Gemma 2: Improving Open Language Models at a Practical Size

In plain words

Gemma 2 is a family of open language models in three sizes: about 2, 9 and 27 billion parameters. The authors keep the basic Transformer design but change a few parts of it, and they train the two smaller models differently. Instead of learning only from the single correct next word, the small models learn from the full set of word probabilities produced by a larger model. That gives each training step richer feedback.

The architecture changes are small. Attention alternates between a local layer that looks only at nearby text and a global layer that looks at the whole context. Attention scores are squashed by a smooth cap so they cannot grow without bound, and several query heads share one set of key and value heads to save memory when generating text. The authors report that the models beat other open models of similar size and sometimes compete with models about twice as large. They release the weights.

TL;DR

Gemma 2 (2B, 9B, 27B) is a decoder-only Transformer with alternating local sliding-window (4096) and global attention layers, logit soft-capping and grouped-query attention with two query heads per key/value head. The 2B and 9B models are trained by distilling a larger teacher's token distributions; the 27B model is trained from scratch. Author-reported results place the 27B model near LLaMA-3 70B on the HuggingFace suite and the 9B model at 70.2 average on eight benchmarks. Use this card when you need the architecture details or want to know where the paper is thin on training recipe.

Problem

Small open models were improving mainly by training on more tokens, and the authors note that this gain scales only logarithmically with data size §1. They ask whether small models can do better per training step by receiving a richer training signal, and whether a few known architecture changes help at this size §1. The distillation ablations use a 2B model and smaller models (200M to 1B parameters) §5 tab. 6 tab. 7; the architecture ablations use 9B §5 tab. 8 tab. 9 tab. 10.

Method

Architecture §2 tab. 1. Decoder-only Transformer. The main changes from Gemma 1 are:

1. Local and global attention alternate, one layer each, with the local sliding window set to 4096 tokens and the global span set to 8192 §2.
2. Logit soft-capping in every attention layer and in the final output layer, with soft_cap = 50.0 for attention and 30.0 for the final layer §2.
3. RMSNorm applied to the input and output of both the attention and the feed-forward sub-layers (pre- and post-norm) §2.
4. Grouped-query attention with num_groups = 2, chosen from ablations that showed faster inference with no clear loss §2 §5 tab. 8.
5. Rotary position embeddings and the approximated GeGLU activation, both carried over from Gemma 1 §2.
6. A 256k-entry SentencePiece vocabulary shared by input and output embeddings (tied embeddings) §3.1 tab. 1.

# Gemma 2 attention for one layer  [§2]
# Input:  queries (B, T, H, d); keys, values (B, S, K, d), K = H / 2
# Output: attended values (B, T, H, d)

def attention_layer(queries, keys, values, is_local):
    # Scale queries by a constant; the 27B value differs (see Pitfalls) [§2]
    queries = queries * query_scale
    # Every two query heads share one key/value head  [§2] [tab. 1]
    logits = grouped_scores(queries, keys)            # (B, T, H, S)
    # Soft-cap attention logits to (-50, 50)  [§2]
    logits = 50.0 * tanh(logits / 50.0)
    # Causal mask; local layers also drop keys beyond 4096 tokens  [§2]
    mask = causal_mask(T, S)
    if is_local:
        mask = mask and within_window(T, S, 4096)
    logits = where(mask, logits, -inf)
    # Normalize over keys and mix the values  [§2]
    weights = softmax(logits, axis=-1)
    return grouped_mix(weights, values)               # (B, T, H, d)

# Layer schedule  [§2] [code: gemma/gm/nn/_gemma.py#L60-L64]
# The paper says local and global layers alternate; the code starts local.
def layer_types(num_layers):
    # Even-indexed layers local, odd-indexed layers global  [code]
    return ["local" if i % 2 == 0 else "global" for i in range(num_layers)]

Knowledge distillation §3.2. For the 2B and 9B models the student is trained to match the teacher's distribution over the next token given the same context, not the one-hot label. The objective is the cross-entropy between the teacher and student distributions, summed over the vocabulary §3.2.

# Distillation loss for one next-token position  [§3.2]
# Input:  teacher_logits (V,), student_logits (V,); V = 256128
# Output: scalar loss

def distillation_loss(teacher_logits, student_logits):
    # Teacher probability of every vocabulary token  [§3.2]
    teacher_probs = softmax(teacher_logits)
    # Student log-probability of every vocabulary token  [§3.2]
    student_log_probs = log_softmax(student_logits)
    # Cross-entropy: minimize -sum_x P_T(x) log P_S(x)  [§3.2]
    return -sum(teacher_probs * student_log_probs)

Pre-training data and budget §3.1 §3.3 tab. 3. Tokens: 13T for 27B, 8T for 9B, 2T for 2B, mostly English web, code and science text §3.1. Data mixture was chosen by ablations similar to Gemini 1.0 §3.1. Data filtering matches Gemma 1 and includes decontamination of evaluation sets §3.1. Training used TPUv5e (2B), TPUv4 (9B) and TPUv5p (27B) with the sharding given in tab. 3. The optimizer state is sharded in a way the authors describe as similar to ZeRO-3 §3.3. The paper reports 1247.61 tCO2eq for pre-training §3.4.

Post-training §4. Supervised fine-tuning on synthetic and human prompt and response pairs, with some distillation from the teacher on the student's own outputs; then RLHF with a reward model that is "an order of magnitude larger than the policy"; then averaging of models produced by several hyperparameter settings §4. Data filtering removes personal information, unsafe or toxic outputs, self-identification mistakes and duplicates §4. Control tokens are the Gemma 1 set, with a new turn format: the model ends each turn with <end_of_turn> then <eos> §4 tab. 4 tab. 5.

Hyperparameters

Results

Author-reported. Benchmark setups are listed in the tables cited. Baselines come from the authors' own evaluations, the HuggingFace leaderboard or the LLaMA-3 report, as noted in the paper tab. 12 tab. 13.

Selected extra numbers. Pre-trained vs instruction-tuned (MMLU): 2B 52.2 to 56.1, 9B 71.3 to 72.3, 27B 75.2 to 76.2 tab. 17. MBPP: 30.2 to 36.6, 52.4 to 59.2, 62.6 to 67.4 tab. 17. Safety benchmarks (lower is better for RealToxicity, CrowS-Pairs, Toxigen; higher is better for others) are in tab. 18. Offensive-security CTF: 27B solves 34 of 76 InterCode-CTF tasks, 1 of 13 internal CTFs and 0 of 13 Hack the Box tasks tab. 19. None of this is independently verified.

Paper vs code

Searched: the paper links only google-deepmind/dangerous-capability-evaluations §8.4 and a Gemma repository is not linked. A GitHub search of the google-deepmind organization found google-deepmind/gemma, the official Gemma library (Apache-2.0). It is the JAX inference and model code at commit 9e1f74131ba9d45e5fe6ae7bb3941d9f2abb39c3, read below. The repository has since gained Gemma 3 and Gemma 4 code, which we do not use here. The repository contains model definitions and sampling only. No training, distillation or RLHF code is present, so those parts can be checked only against the paper.

Model parameters match tab. 1 for head counts, KV heads, head size, softcaps and window: Gemma 2 2B uses 8 heads, 4 KV heads, head size 256 [code: gemma/gm/nn/_gemma.py#L52-L70]; 9B uses 16/8/256 [code: gemma/gm/nn/_gemma.py#L81-L100]; 27B uses 32/16/128 [code: gemma/gm/nn/_gemma.py#L111-L130]. Layer counts are 26, 42 and 46 [code: gemma/gm/nn/_gemma.py#L29-L31].

Differences and conventions, each with both citations:

• Feed-forward width: paper gives 18432 / 28672 / 73728 tab. 1; code sets hidden_dim to 9216 / 14336 / 36864 _gemma.py _gemma.py _gemma.py. The gating weight has a leading dimension of 2 _modules.py, so the paper's figure is twice the per-branch width. Both describe the same model.
• 27B query scaling: paper does not say. Code uses (embed_dim / num_heads)^-0.5 for 27B, i.e. 144^-0.5, not head_dim^-0.5 used for 2B and 9B _gemma.py _config.py. The 27B head size is 128 tab. 1, so the two choices differ.
• Layer order: paper says "every other layer" with no starting point §2. Code alternates local first, then global _gemma.py.
• Attention soft-cap and mask order: code applies the cap to the raw scores before masking _modules.py. The paper does not state an order.
• Sliding-window mask: code's create_sliding_mask uses two bounds, keys with position greater than pos - W and less than pos + W _modules.py. The paper describes only a window of 4096. Under a causal mask the upper bound adds no restriction (inference, not checked by running the code).
• Final soft-cap: paper states 30.0 §2; code divides by 30, applies tanh and multiplies back _transformer.py.
• Attention soft-cap in code is 50.0 for all three sizes [code: gemma/gm/nn/_gemma.py#L68] _gemma.py _gemma.py, matching §2.
• Text vs table: §6.1 says "Qwen1.5 34B" §6.1 while the table caption says 32B tab. 12. §6.2 writes "Gemma 2.6B" for the 2B model §6.2. Use the table values and treat these as typos.

Pitfalls

• Feed-forward width is fused. Copy the paper's width as the gate+up width, so the per-branch hidden size is half of it. Doubling it by mistake doubles the feed-forward parameters.
• 27B query scale. Use 1/sqrt(d_model / num_heads) for 27B and 1/sqrt(head_dim) for 2B and 9B, as the code does. The paper does not state this, so an implementer following only the paper would pick the 2B/9B rule.
• Layer ordering is unstated in the paper. Code puts a local layer first. Weights are indexed by layer, so the ordering decides which parameters apply to which layer. Follow the code, not the paper's wording.
• Distillation details are missing. The paper gives only the cross-entropy over the full vocabulary §3.2. Temperature, top-k truncation, whether teacher logits are capped, and what the student's training data is are not stated.
• 2B teacher is unnamed. The only teacher size given is 7B in the ablation §5 and 27B for 9B by ratio §5. Do not assume the teacher for 2B.
• Training hyperparameters are absent. Learning rate, batch size, optimizer settings, warmup, schedule and sequence packing are not stated anywhere in the paper. The appendix is empty in the source. There is no training code in the repository, so these cannot be recovered.
• Post-training values are absent. SFT, RLHF and model-averaging settings are described as tuned but no values are given §4.
• GQA naming. The paper calls the setting num_groups = 2 §2. In code the KV head count is half the query head count in all three models _gemma.py, which means two query heads per KV head.
• Caps are not ablated. The paper sets the attention cap to 50 and the final cap to 30 §2 but does not report an ablation that removes or changes either one, so their effect on quality is not stated.
• Window at inference. The window can be shrunk at inference with small perplexity change, but this was measured only on the 9B model and on one validation set tab. 10. Results for other sizes and tasks are not stated.
• Baseline protocols differ. Several baselines are copied from the HuggingFace leaderboard or from LLaMA-3's report, and the paper notes a +3% to +4% change in some baseline scores under LLaMA-3's evaluation protocol tab. 13 (footnotes). Cross-model comparisons mix protocols.
• Not uniform gains. 2B HumanEval is lower than Gemma 1 2B (20.1 vs 22.0) tab. 13. Gains are not on every benchmark.
• Arena and human studies. Elo numbers are from the Chatbot Arena and the side-by-side studies use the authors' own prompt sets tab. 14 tab. 15. Treat them as author-reported.
• Memorization. The paper uses a narrow definition: a model can be prompted to generate near-copies of training examples §7. The test is 50-token memorization given a 50-token prompt, with exact and approximate (10% edit distance) matching §7. The personal-data check uses an automated tool that the authors say produces false positives, so the rates may be overestimates §7.
• CBRN. The CBRN section reports only qualitative findings ("knowledge in these domains is low") and gives no numbers §8.4. Do not cite a CBRN number from this card.

Lineage

• Builds on: Gemma 1 (cited as gemmateam2024gemma in the paper, not in this wiki); local sliding-window attention and global attention, both cited to Beltagy et al., 2020 (Longformer) §2; grouped-query attention, Ainslie et al., 2023 §2; knowledge distillation, Hinton et al., 2015 §3.2; RMSNorm §2; RoPE §2; GeGLU §2; logit capping, Bello et al., 2016 §2.
• Compared against: Qwen1.5 32B and LLaMA-3 70B (HF suite) tab. 12; LLaMA-3 8B, Mistral 7B and Gemma 1 2B/7B tab. 13; GPT-4o and other chat models in the Arena tab. 14 and in human studies tab. 15.
• Superseded by: not stated in the wiki.