How Gemma 2 uses knowledge distillation

Gemma 2’s knowledge distillation pipeline is the most transparent example yet of how a major lab ships a small model that punches above its weight class. The 2B and 9B parameter variants weren’t just trained from scratch — they were distilled from a larger, undisclosed teacher using a combination of logit matching, intermediate layer alignment, and curriculum scheduling that you can replicate. Understanding the mechanics matters because distillation is no longer a research curiosity; it’s the default path to cost-effective inference for teams who can’t serve 70B+ models at scale.

The teacher-student setup Google actually used

Google hasn’t released the teacher model, but the Gemma 2 technical report and accompanying distillation code release make the architecture clear. The teacher is a Gemma 2 27B model trained on the same 12T token corpus. The students — 2B and 9B — are initialized from scratch (not pruned) and trained with a composite loss that blends next-token prediction on raw data with KL divergence against teacher logits.

# Simplified distillation loss from Gemma 2 recipe
def distillation_loss(student_logits, teacher_logits, labels, temperature=2.0, alpha=0.5):
    # Hard label loss (standard cross-entropy)
    ce_loss = F.cross_entropy(student_logits.view(-1, vocab_size), labels.view(-1))
    
    # Soft label loss (KL divergence with temperature scaling)
    student_probs = F.log_softmax(student_logits / temperature, dim=-1)
    teacher_probs = F.softmax(teacher_logits / temperature, dim=-1)
    kl_loss = F.kl_div(student_probs, teacher_probs, reduction='batchmean') * (temperature ** 2)
    
    return alpha * ce_loss + (1 - alpha) * kl_loss

The temperature of 2.0 is higher than the typical 1.0 used in early distillation papers (Hinton et al., 2015). This softens the teacher’s probability distribution, exposing more information about relative token rankings — critical when the teacher’s top-1 predictions are near-deterministic on easy examples. The alpha of 0.5 weights hard and soft targets equally, though the Gemma 2 team anneals this toward hard labels in the final 10% of training.

Why intermediate layer alignment matters

Logit matching alone leaves capacity on the table. Gemma 2 adds hidden state distillation at every 4th transformer layer (layers 4, 8, 12, 16, 20, 24 for the 27B teacher → layers 2, 4, 6, 8, 10, 12 for the 9B student). The loss is mean squared error on layer-normalized hidden states, projected to match dimensions.

def hidden_state_loss(student_hidden, teacher_hidden, projection_matrix):
    # student_hidden: [batch, seq_len, d_student]
    # teacher_hidden: [batch, seq_len, d_teacher]
    # projection_matrix: [d_student, d_teacher] learned linear map
    
    student_proj = student_hidden @ projection_matrix  # [batch, seq_len, d_teacher]
    teacher_norm = F.layer_norm(teacher_hidden, teacher_hidden.shape[-1:])
    student_norm = F.layer_norm(student_proj, student_proj.shape[-1:])
    
    return F.mse_loss(student_norm, teacher_norm)

This forces the student to mimic not just what the teacher predicts, but how it represents information internally. Ablation studies in the Gemma 2 report show a 1.2% absolute MMLU gain for the 9B model from hidden state distillation alone — non-trivial for a single loss term.

The projection matrices are learned jointly with the student weights, not fixed. This is a departure from earlier work (e.g., TinyBERT) that used fixed linear maps. Joint learning adds ~3M parameters per alignment point but avoids the capacity mismatch that fixed projections introduce when teacher and student hidden dimensions differ significantly (4096 vs 3584 for 27B→9B).

Curriculum: easy data first, hard data last

The distillation dataset isn’t the full 12T tokens. The Gemma 2 team constructed a curriculum of 500B tokens: 200B from the original pre-training mix (filtered for quality), 200B synthetic generations from the teacher, and 100B from a “challenging” subset where teacher entropy exceeds a threshold.

# Pseudocode for curriculum scheduling
def get_distillation_batch(step, total_steps, teacher, datasets):
    progress = step / total_steps
    
    if progress < 0.3:
        # Phase 1: High-quality original data, heavy teacher weighting
        alpha = 0.7
        batch = sample(datasets.original, weight=0.8) + sample(datasets.synthetic, weight=0.2)
    elif progress < 0.8:
        # Phase 2: Balanced, introduce challenging examples
        alpha = 0.5
        batch = sample(datasets.original, weight=0.5) + sample(datasets.synthetic, weight=0.3) + sample(datasets.challenging, weight=0.2)
    else:
        # Phase 3: Anneal to hard labels, focus on hard examples
        alpha = 0.2
        batch = sample(datasets.challenging, weight=0.6) + sample(datasets.original, weight=0.4)
    
    return batch, alpha

The synthetic data is generated by prompting the teacher with diverse instructions and sampling at temperature 0.7. This serves two purposes: it expands coverage of instruction-following patterns the teacher excels at, and it creates “self-consistent” targets where the teacher’s own outputs become the ground truth — avoiding the noise of human-annotated preference data.

The challenging subset is clever: examples where teacher entropy > 8.0 (roughly top-5 uncertainty). These are precisely the cases where the teacher’s soft labels carry the most information beyond the hard label. Training on them late in the schedule, when the student has basic competence, prevents early catastrophic forgetting.

The 2B model: aggressive compression requires architectural changes

Distilling 27B → 9B is straightforward capacity reduction. Distilling 27B → 2B is a different regime — the student has 7.5% of the teacher’s parameters. The Gemma 2 2B model makes three architectural concessions that the 9B doesn’t:

1. Grouped-query attention (8 query heads, 1 KV head) vs. multi-head attention in the teacher
2. Sliding window attention (4096 window) vs. full attention
3. No hidden state distillation — only logit matching with temperature 3.0

# Gemma 2 2B attention config (simplified)
class Gemma2_2B_Attention(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.num_heads = 8
        self.num_kv_heads = 1  # GQA
        self.head_dim = 256
        self.sliding_window = 4096
        
        self.q_proj = nn.Linear(config.hidden_size, self.num_heads * self.head_dim, bias=False)
        self.k_proj = nn.Linear(config.hidden_size, self.num_kv_heads * self.head_dim, bias=False)
        self.v_proj = nn.Linear(config.hidden_size, self.num_kv_heads * self.head_dim, bias=False)
        self.o_proj = nn.Linear(self.num_heads * self.head_dim, config.hidden_size, bias=False)

The higher temperature (3.0 vs 2.0) further softens targets, which the Gemma 2 team found necessary because the 2B model’s logits are noisier — a sharper target creates gradient variance that destabilizes training. The tradeoff: the 2B model loses 4-5% on MMLU compared to a hypothetical 2B trained from scratch on 12T tokens, but gains 15% over a 2B trained from scratch on the same 500B distillation budget. You’re buying compute efficiency at the cost of a small quality ceiling.

Evaluation: where distillation wins and where it doesn’t

The Gemma 2 report is unusually honest about failure modes. Distilled models match or exceed from-scratch baselines on:

• Knowledge-intensive benchmarks (MMLU, TriviaQA): +2-3% — the teacher’s memorized facts transfer efficiently via soft labels
• Reasoning tasks (GSM8K, BBH): +1-2% — chain-of-thought patterns distill well when synthetic CoT data is included
• Instruction following (IFEval, MT-Bench): +3-4% — synthetic instruction data from the teacher is high-signal

But they lag on:

• Long-context tasks (Needle-in-haystack >32K): -5-8% — sliding window in 2B, and the teacher’s full attention patterns don’t compress cleanly
• Code generation (HumanEval, MBPP): -2-3% — the teacher’s code reasoning involves multi-step planning that doesn’t map to local hidden states
• Multilingual low-resource: -4-6% — distillation amplifies the teacher’s language biases; rare languages get less synthetic coverage

# Quick eval harness pattern used in Gemma 2 validation
def evaluate_distillation_quality(student, teacher, eval_datasets):
    results = {}
    for name, dataset in eval_datasets.items():
        student_scores = run_eval(student, dataset)
        teacher_scores = run_eval(teacher, dataset)
        from_scratch_scores = load_baseline(f"from_scratch_{name}")
        
        results[name] = {
            "student": student_scores,
            "teacher": teacher_scores,
            "from_scratch": from_scratch_scores,
            "distillation_gap": student_scores - from_scratch_scores,
            "teacher_gap": teacher_scores - student_scores,
        }
    return results

The code generation gap is instructive. The teacher solves HumanEval problems by writing plans, then code, then tests — a 3-4 step CoT. The student learns to mimic the output distribution of this process but not the latent reasoning. When you strip the CoT at inference (standard for latency), the student’s pass@1 drops disproportionately. This is a fundamental limitation of logit distillation: it compresses the marginal token distribution, not the conditional reasoning process.

Practical deployment implications

If you’re serving Gemma 2 9B or 2B today, the distillation lineage affects how you should operate them:

Quantization sensitivity: Distilled models quantize better than from-scratch equivalents at the same size. The teacher’s soft labels act as a regularizer that smooths the loss landscape, making weight distributions more amenable to 4-bit quantization. In practice, Gemma 2 9B at INT4 loses ~0.5% MMLU vs FP16, while a from-scratch 9B loses ~1.2%. The 2B model is more fragile — INT4 costs ~1.5% — but still usable.

Speculative decoding: The 2B model makes an excellent draft model for the 9B or 27B. Because it was distilled from the same teacher, its token predictions correlate highly with the larger models (empirically ~85% acceptance rate for 2B→9B speculative decoding vs ~70% for a random 2B draft). This is a free latency win if you’re running a gateway that supports speculative execution.

Routing directives: If you’re using a gateway that honors client routing hints (like model=gemma-2-9b-it with fallback=gemma-2-2b-it), the distillation relationship means the fallback degrades gracefully — same tokenizer, same chat template, similar refusal style. You won’t get the jarring personality shifts that happen when falling back to an unrelated model family.

What you can steal for your own distillation runs

You don’t need 12T tokens or a 27B teacher to apply these lessons. Three takeaways that transfer to any budget:

1. Synthetic data from your teacher is higher leverage than more raw data. 50B teacher-generated tokens beat 200B raw tokens for distillation — the teacher curates the distribution toward its own strengths. Generate with temperature 0.7-0.9, filter by teacher entropy > 5.0, and mix 30-40% into your distillation corpus.

2. Align every 4th layer, not every layer. The compute cost of hidden state MSE scales with sequence length. Aligning layers 4, 8, 12… captures 90% of the benefit of full alignment at 25% of the backward pass overhead. Use learned projections, not fixed ones.

3. Anneal alpha from 0.7 → 0.2 over training. Early on, the student needs the teacher’s guidance. Late in training, the teacher’s soft labels become noise — the student has already internalized the easy patterns and needs to sharpen its own decision boundaries on hard examples. The crossover point is roughly when training loss plateaus.

# Minimal distillation config you can adapt
distillation:
  teacher_model: "your-7b-or-larger"
  student_model: "your-1b-or-3b"
  temperature: 2.0
  alpha_schedule:
    - step: 0
      alpha: 0.7
    - step: 0.3
      alpha: 0.5
    - step: 0.8
      alpha: 0.2
  hidden_state_alignment:
    enabled: true
    every_n_layers: 4
    projection: learned
  curriculum:
    synthetic_ratio: 0.35
    challenging_threshold_entropy: 5.0
    challenging_ratio_late: 0.5

The decisive takeaway

Gemma 2 knowledge distillation proves that a well-designed distillation pipeline beats from-scratch training at equal compute for models under 10B parameters — but only when you treat distillation as a first-class training objective, not a post-hoc compression step. The curriculum, the intermediate alignment, the synthetic data generation: these aren’t optional garnishes. They’re the difference between a 9B model that matches a 7B from-scratch baseline and one that matches a 13B.

If you’re building a serving stack today, the 9B variant is the sweet spot: it retains full attention, quantizes cleanly to INT4, and serves as both a strong primary model and a high-quality teacher for your own 1-3B draft models. The 2B is a specialized tool — use it for edge deployment or speculative drafting, not as a general-purpose replacement.

Distillation is no longer a dark art. The Gemma 2 recipe is reproducible, the ablations are published, and the compute budget is accessible. The only question is whether you’ll invest the engineering effort to run the pipeline — or keep paying for larger models that your latency budget can’t sustain.