Google Gemma 2 Release Notes
Google has released Gemma 2, a new family of open-weight large language models (LLMs) based on Google DeepMind's Gemini. The release includes four models: base and instruction-tuned versions of both a 9 billion (9B) and a 27 billion (27B) parameter model, designed to provide state-of-the-art performance in a permissive license that allows for commercial use and redistribution.
Model Specifications and Training
Gemma 2 is available in two primary sizes with a context length of 8,192 tokens and uses Rotary Position Embedding (RoPE). The models were trained on significantly more data than the first Gemma iteration: the 27B model was trained on 13 trillion tokens, and the 9B model on 8 trillion tokens of web data, code, and mathematics.
- gemma-2-9b: Base 9B model
- gemma-2-9b-it: Instruction-tuned 9B model
- gemma-2-27b: Base 27B model
- gemma-2-27b-it: Instruction-tuned 27B model
Technical Innovations
Gemma 2 introduces four key technical advances to improve generation quality and training stability:
Sliding Window Attention
Gemma 2 employs a hybrid attention mechanism. It interleaves sliding window attention (covering 4,096 tokens) with full quadratic global attention (covering 8,192 tokens) every other layer. This design aims to maintain high quality in long-context scenarios while gaining the efficiency benefits of sliding window attention.
Logit Soft-Capping
To prevent logits from growing excessively large and stabilizing training, Gemma 2 uses soft-capping. Logits are scaled to a fixed range using the formula logits → soft_cap * tanh(logits/soft_cap).
- Attention layers: Capped at 50.0
- Final layer: Capped at 30.0
Note that soft-capping is currently incompatible with Flash Attention/SDPA during training; therefore, eager attention is recommended for stable fine-tuning.
Knowledge Distillation
The 9B model was pre-trained using knowledge distillation, where a larger teacher model provides a richer signal for the student model to learn from. For post-training, the team used "on-policy distillation." In this process, the student generates completions from SFT prompts, and the KL divergence between the teacher's and student's logits is minimized to reduce the train-inference mismatch.
Model Merging via WARP
Gemma 2 utilizes a merging technique called WARP to combine models in three stages:
- Exponential Moving Average (EMA): Applied during RL fine-tuning.
- Spherical Linear intERPolation (SLERP): Applied after RL fine-tuning of multiple policies.
- Linear Interpolation Towards Initialization (LITI): Applied after the SLERP stage.
Performance Evaluation
According to the Technical Report, Gemma 2 shows competitive performance against other open models.
Comparison of Large Models
| Benchmark | Llama 3 (70B) | Qwen 1.5 (32B) | Gemma 2 (27B) |
|---|---|---|---|
| MMLU | 79.2 | 74.3 | 75.2 |
| GSM8K | 76.9 | 61.1 | 75.1 |
| ARC-c | 68.8 | 63.6 | 71.4 |
| HellaSwag | 88.0 | 85.0 | 86.4 |
| Winogrande | 85.3 | 81.5 | 83.7 |
Comparison of Small Models
| Benchmark | Mistral (7B) | Llama 3 (8B) | Gemma (8B) | Gemma 2 (9B) |
|---|---|---|---|---|
| MMLU | 62.5 | 66.6 | 64.4 | 71.3 |
| GSM8K | 34.5 | 45.7 | 50.9 | 62.3 |
| ARC-C | 60.5 | 59.2 | 61.1 | 68.4 |
| HellaSwag | 83.0 | 82.0 | 82.3 | 81.9 |
| Winogrande | 78.5 | 78.5 | 79.0 | 80.6 |
Implementation and Deployment
Prompting
The instruction-tuned versions use a specific conversation structure that must be exactly reproduced for effectiveness:
<start_of_turn>user [Prompt]<end_of_turn> <start_of_turn>model [Response]<end_of_turn>
Hardware Requirements and Quantization
- Gemma 2 9B: Requires approximately 18 GB of RAM.
- Gemma 2 27B: Requires approximately 56 GB of RAM in bfloat16.
To reduce memory footprints, the models can be loaded in 4-bit or 8-bit mode. The 27B version in 4-bit takes approximately 18 GB of memory. It is critical to use bfloat16 for the 27B instruction-tuned model, as float16 can produce erratic outputs.
Fine-Tuning
Fine-tuning can be achieved using the Hugging Face TRL library with QLoRA. When targeting linear layers for PEFT, it is recommended to target attention blocks (q_proj, k_proj, v_proj, o_proj) rather than MLP layers, as the MLP layers are sparse and do not interact well with PEFT.