Preprint
Knowledge Graphs

Gemma 2

June 1, 2024

0

Citations

0

Influential Citations

Venue

2024

Year

Abstract

Utilizes interleaving local-global attentions and group-query attention, trained with knowledge distillation instead of next token prediction to achieve competitive performance comparable with larger models.

Analysis

Why This Paper Matters

The race for ever-larger language models often overshadows the practical value of smaller, more efficient architectures. Gemma 2, from Google DeepMind, challenges the assumption that size alone determines capability. By introducing a clever mix of local and global attention mechanisms and leveraging knowledge distillation during training, the Gemma 2 models—especially the 27B parameter variant—achieve performance that rivals or even surpasses models two to three times their size. This is a significant step for the open-source AI community, as it lowers the hardware and energy requirements for deploying state-of-the-art language models, making advanced AI more accessible to researchers and practitioners with limited resources.

How It Works

Figure 1

Gemma 2 retains the core design of its predecessor, Gemma 1, including an 8192-token context length, Rotary Position Embeddings (RoPE), and approximated GeGLU non-linearity. However, several key modifications set it apart.

Interleaved Local-Global Attention: Instead of using full global attention in every layer, Gemma 2 alternates between local sliding window attention and global attention. Local layers attend to a window of 4096 tokens, while global layers attend to the full 8192-token context. This hybrid approach reduces computational cost while preserving the ability to capture long-range dependencies.

Group-Query Attention (GQA): Both the 9B and 27B models use GQA with two key-value heads per query group (num_groups=2). This reduces memory bandwidth and speeds up inference compared to standard multi-head attention, especially for autoregressive generation.

Logit Capping: Following Gemini 1.5, the logits in each attention layer and the final layer are capped using a soft cap function: logits ← soft_cap * tanh(logits / soft_cap). For the 9B and 27B models, attention logits are capped at 50.0 and final logits at 30.0. This prevents extreme logit values and stabilizes training.

Normalization: RMSNorm is applied to both the input and output of every transformer sub-layer, including attention and feedforward layers, further enhancing training stability.

Overview of the main model parameters and design choices.

Knowledge Distillation for Pre-training: A standout innovation is the use of knowledge distillation during pre-training for the 9B and 2.6B models. Instead of standard next-token prediction, these models are trained to mimic the output probabilities of the larger 27B teacher model. Since the vocabulary has 256k tokens, only a sampled subset of teacher probabilities is stored to manage memory. This richer training signal helps the smaller models learn more effectively from the same data.

Post-Training Pipeline: The post-training process involves three phases:

  1. Supervised Fine-Tuning (SFT): Fine-tuning on synthetic and human-generated prompt-response pairs, using behavioral cloning and distillation from a larger teacher.
  2. Reinforcement Learning from Human Feedback (RLHF): Using the SFT model as a policy, trained with a reward model on labeled preference data.
  3. Model Merging (Warp): The models from each phase are averaged using a three-stage technique: Exponential Moving Average (EMA) during RL, Spherical Linear Interpolation (SLERP) after RL, and Linear Interpolation Towards Initialization (LITI) after SLERP. This merging improves overall performance.

Results

Figure 6

On pre-training benchmarks, the Gemma 2 27B model outperforms Qwen1.5 32B and is only a few percent behind LLaMA-3 70B, despite being 2.5× smaller and trained on two-thirds less data. The 9B and 2.6B models also lead their size categories.

Figure 7

Post-training evaluations on Chatbot Arena show the Gemma 2 27B instruction-tuned model setting a new state-of-the-art for open-weights models, slightly surpassing the much larger Llama3-70B-Instruct and Nemotron-4-340B-Instruct. The 9B model strongly outperforms all other models in its parameter range.

Evaluation of Gemma 2 9B and 27B Instruction Tuned models on the Chatbot Arena.

Significance

Gemma 2 demonstrates that thoughtful architectural design and training strategies can yield highly competitive models at a fraction of the size and cost of the largest LLMs. By making state-of-the-art performance accessible on consumer-grade hardware, it empowers a wider range of developers and researchers to build upon and deploy advanced language models. The use of knowledge distillation during pre-training is a particularly promising direction, suggesting that the field may shift focus from scaling model size to scaling the quality of training signals. Gemma 2 sets a new benchmark for what is achievable with practical-sized open models.