On-Policy Distillation
for LLMs

Chapter 1: Exposure Bias

The intuition from Chapter 0 has a precise mathematical formulation. Let's derive exactly WHY errors compound — and quantify how much worse off-policy is than on-policy.

The off-policy training objective

Standard knowledge distillation trains the student to match the teacher's output distribution, conditioned on teacher-generated prefixes:

The critical detail: y_<t comes from the teacher. The student never conditions on its own generated tokens during training. At inference, however, it must condition on its own previous outputs.

Distribution mismatch

This creates a train-test distribution mismatch. During training, the student sees the distribution of prefixes p_T(y_<t). During inference, it sees p_θ(y_<t). These two distributions diverge more with each token generated.

Let ε be the per-step error probability — the chance that the student generates a token that deviates from the teacher's distribution at any given step. After T steps:

This is the DAgger theorem (Ross et al., 2011), originally proven for imitation learning in robotics. It applies directly to autoregressive language modeling.

Worked example: 10-step proof

Suppose your student has 95% per-step accuracy (ε = 0.05). For a 10-step reasoning chain:

Independent errors (best case): P(all correct) = 0.95¹⁰ = 60%. Already losing 40% of trajectories.

With compounding (off-policy): Each error makes subsequent errors more likely because the student enters unseen territory. The actual success rate drops well below 60% — closer to the O(εT²) = O(0.05 × 100) = O(5) bound, meaning errors dominate completely.

With on-policy training: The student has practiced recovery. Errors at step 3 don't cascade because the student knows what to do from its own imperfect states. Success follows the O(εT) = O(0.05 × 10) = O(0.5) bound — much more controlled.

The on-policy training objective

On-policy distillation replaces teacher prefixes with student-generated (or mixed) prefixes:

Here π_mix is a mixture policy — typically λ · p_θ + (1-λ) · p_data — that blends student rollouts with ground-truth sequences. The divergence D_f can be forward KL, reverse KL, JSD, or any f-divergence (Chapter 2 explores which to choose).

Chapter 2: The f-Divergence Framework

Now we know WHY on-policy helps. The next question: what should the student optimize when training on its own outputs? The choice of divergence measure fundamentally shapes what the student learns.

The f-divergence family

An f-divergence is a general way to measure how different two probability distributions are. Every f-divergence has the form:

where f is a convex function with f(1) = 0. Different choices of f give different divergences — and each one tells the student to prioritize different aspects of the teacher distribution.

Forward KL: mode-covering

Forward KL (f(u) = u log u) measures D_KL(P_teacher || Q_student). The expectation is over the student's distribution. This means: wherever the teacher has probability mass, the student MUST also have probability mass (or the divergence explodes).

Result: the student covers all modes of the teacher. If the teacher assigns probability to two different valid answers, the student spreads probability over both. But it also assigns probability to the space between the modes — hallucinating intermediate outputs that neither mode produces.

Reverse KL: mode-seeking

Reverse KL (f(u) = -log u) measures D_KL(Q_student || P_teacher). The expectation is over the teacher's distribution. This means: wherever the student has probability mass, the teacher must also (or the divergence explodes).

Result: the student seeks one mode and concentrates on it. It won't hallucinate between modes because it never places mass where the teacher doesn't. But it drops entire modes — if the teacher has two valid answers, the student picks one and ignores the other.

JSD and α-divergence: interpolations

The Jensen-Shannon Divergence (JSD) is a symmetric average of forward and reverse KL, computed through a mixture M = (P+Q)/2. It's bounded in [0, log 2], which prevents gradient explosion.

The α-divergence family provides a continuous interpolation parameter α between forward KL (α → 1) and reverse KL (α → 0). Setting α = 0.5 gives a symmetric divergence similar to JSD.

Which divergence for which task?

Chapter 3: Fixed Divergence Methods

Three papers crystallized the core approaches to on-policy distillation. Each picks a different point in the design space — and each reveals a different tradeoff.

GKD: Generalized Knowledge Distillation (Agarwal et al., 2024)

GKD is the simplest framework. Its key insight: make on-policy distillation modular. Separate the sampling policy (who generates the prefixes?) from the divergence (what does the student optimize?).

The sampling policy is a mixture:

When λ = 0, it's standard off-policy (train on ground-truth prefixes). When λ = 1, it's fully on-policy (train only on student's own generations). Values in between blend both — the student sees some perfect prefixes and some of its own imperfect ones.

GKD is divergence-agnostic: you can plug in forward KL, reverse KL, JSD, or any f-divergence. The sampling policy and the divergence are orthogonal choices.

MiniLLM: Reverse KL via REINFORCE (Gu et al., 2024)

MiniLLM goes all-in on reverse KL. The problem: reverse KL D_KL(p_θ || p_T) requires computing the gradient through the student's own sampling process (because the expectation is over p_θ). You can't just backpropagate through discrete token sampling.

MiniLLM's solution: treat it as a reinforcement learning problem. The "reward" for generating token y_t is:

This is the log-probability ratio: how much more the teacher likes this token than the student does. High reward = the teacher strongly approves of something the student isn't yet confident about. The REINFORCE estimator uses this reward to compute policy gradients.

The catch: REINFORCE has high variance. The gradient estimates are noisy, requiring large batch sizes, careful baseline subtraction, and many training steps to converge. Training is unstable — the loss oscillates wildly before settling.

DistiLLM: Skewed KL for stability (Ko et al., 2024)

DistiLLM solves MiniLLM's instability problem with a clever trick: instead of minimizing D_KL(p_θ || p_T) directly, minimize D_KL(p_θ || p̃) where p̃ is a skewed mixture:

Why does this help? The gradient explosion in MiniLLM happens when the student assigns near-zero probability to a token the teacher strongly prefers. The ratio p_T/p_θ → ∞, and the gradient explodes.

By mixing in the student's own distribution, DistiLLM ensures p̃ never diverges too far from p_θ. Even when p_T(y) is high and p_θ(y) is near zero, the mixture p̃(y) = α · p_T(y) + (1-α) · 0 ≈ α · p_T(y), and the ratio p̃/p_θ grows more slowly than p_T/p_θ. Gradient explosions are tamed.

The tradeoff triangle

python
# GKD: simple mixture sampling + any divergence
def gkd_loss(student, teacher, x, lam=0.5, div='jsd'):
    # Sample prefix from mixture policy
    if random() < lam:
        prefix = student.generate(x)   # on-policy
    else:
        prefix = ground_truth(x)       # off-policy

    p_t = teacher.forward(x, prefix)   # teacher logits
    p_s = student.forward(x, prefix)   # student logits
    return compute_divergence(p_t, p_s, div)

# MiniLLM: reverse KL via REINFORCE
def minillm_loss(student, teacher, x):
    y = student.generate(x)            # full on-policy rollout
    r = teacher.logprob(y, x) - student.logprob(y, x)
    baseline = moving_avg(r)          # variance reduction
    return -((r - baseline) * student.logprob(y, x)).mean()

# DistiLLM: skewed mixture target
def distillm_loss(student, teacher, x, alpha=0.9):
    y = student.generate(x)
    p_t = teacher.forward(x, y)
    p_s = student.forward(x, y)
    p_mix = alpha * p_t + (1 - alpha) * p_s  # skewed target
    return kl_div(p_s, p_mix)          # bounded gradient!

Chapter 4: Adaptive & RL-Augmented Objectives

Fixed divergences are one-size-fits-all. But what if different tokens need different treatment? A confident token where teacher and student agree needs less gradient pressure than a confused token at a reasoning branch point.

Consider a math proof: tokens like "therefore" are near-deterministic — both teacher and student assign 95%+ probability. But choosing between "integration by parts" vs. "substitution" is the crux. A single divergence treats both identically. That's wasteful at best, harmful at worst.

Adaptive Divergences

ToDi (Token-level Divergence routing) scores each token by its log-ratio r_t = log p_θ(y_t) − log p_T(y_t). If r_t > 0 (student overestimates), use FKL to pull it down. If r_t < 0 (student underestimates), use RKL to push it up. The routing is per-token and parameter-free.

AKL (Adaptive KL) blends FKL and RKL based on the head/tail gap of the teacher's distribution. Peaked teacher → RKL. Flat teacher → FKL. The weight is a smooth function of teacher entropy.

EOPD interpolates continuously: α_t = σ(H(p_T) − τ). High teacher entropy → more FKL. Low entropy → more RKL.

The RL Equivalence

G-OPD proves a startling result: on-policy distillation IS reinforcement learning.

Left side: minimize KL from student to teacher. Right side: maximize reward (teacher log-prob) plus entropy (exploration). Identical. Every RL trick — baselines, advantage estimation, PPO clipping — applies directly to distillation.

Chapter 5: Signal Sources — Where Does Feedback Come From?

We've discussed WHAT to optimize. Now: where does the teaching signal come from? The answer determines everything — compute budget, quality ceiling, deployment constraints.

Think of it as a spectrum. At one end, full access to the teacher's brain (every logit). At the other end, just a thumbs-up or thumbs-down on complete outputs.

White-Box: Full Logit Access

You run the teacher and get its full output distribution at every token position. Richest signal — you know exactly how much probability the teacher assigns to every possible next token. Cost: full teacher forward pass for every sample (~140 GB for a 70B model).

Cross-family challenge: Different tokenizers mean logits aren't aligned. DSKD and ULD solve this via shared vocabulary projection or distribution-level matching.

Black-Box: API-Only Access

Only generated text or scalar scores. Reality when distilling from GPT-4 or Claude. GAD trains a discriminator. OVD uses the teacher to SELECT among student candidates.

Self-Distillation: No Teacher At All

The model teaches itself by exploiting asymmetries in its own capabilities:

• Privileged information: Model is better WITH extra context. Distill that advantage: p(y|x, hints) → p(y|x).
• Self-play: Current model vs. previous version. DPO until convergence.
• External feedback: Verifier (math checker, code runner) scores outputs. GRPO optimization.

python
# White-box: exact KL at every token (dense signal)
loss = kl_div(teacher_logits, student_logits)  # [B, T, V]

# Black-box: only outcome reward (sparse signal)
reward = teacher_api.score(student_output)  # scalar
loss = -reward * log_prob(student_output)   # REINFORCE

# Self-distillation: model teaches itself
good = model.generate(prompt, with_hints=True)
loss = kl_div(model(prompt).logits, model(prompt+hints).logits)

Chapter 6: Self-Distillation Deep Dive

The most surprising finding: you don't always need a teacher. A model can improve itself by exploiting asymmetries in its own capabilities.

Privileged Information (PI)

OPSD: Condition on ground-truth during training. The model with answers generates better reasoning chains. Distill: p(reasoning | x, answer) → p(reasoning | x). The model learns to reason well WITHOUT seeing the answer.

CRISP: Concise prompt as PI — saves 57% tokens at inference while maintaining quality.

GATES: Full document as PI → distill to model with only a summary.

Self-Play

SPIN: At iteration k, generate y_k. DPO with "chosen" = reference, "rejected" = y_k. Update. Repeat. Converges when p_θ = p_ref (distributions indistinguishable, DPO signal = zero).

π-Play: Multi-agent variant — several models co-evolve, playing against each other rather than a fixed reference.

External Feedback (Verifier-Driven)

SD-ZERO: Generate N solutions, verifier gives binary reward (correct/incorrect), optimize with GRPO. No teacher at all — just trial and error with automated grading. This is essentially what DeepSeek-R1-Zero does.

Chapter 7: Training Dynamics & Stabilization

On-policy training is inherently unstable. The data distribution shifts with every gradient step. The student explores bad regions and gets noisy gradients. How do Qwen3 and DeepSeek-V4 make this actually work?

Token Weighting: Not All Tokens Are Equal

TIP (Token Importance Profiling) scores each token on a 2D quadrant: teacher entropy (is the teacher confident?) × student-teacher divergence (does the student disagree?). The golden tokens: teacher is confident but student is confused. These get highest weight.

SCOPE uses surprise weighting: tokens where student assigns very low probability but teacher assigns high get upweighted by their log-ratio.

Curriculum (PACED)

Start with easy samples (low teacher perplexity), gradually increase difficulty. Beta-kernel sampling from the difficulty frontier — the difficulty slider auto-adjusts as the student improves.

Compute Optimization

On-policy is 3-8× more expensive (generate + score + train). Solutions:

The Flawed Prefix Trap

If the student generates a terrible prefix (off-topic, degenerate repetition), teacher feedback on it is meaningless noise. Detection: perplexity threshold, length filter, reward floor.

Chapter 8: OPD Explorer

Let's put it all together. You're going to run an on-policy distillation pipeline and see every design choice interact in real time.

The teacher has a bimodal distribution — two valid answer modes. The student starts broad and must learn to match. Pick a method, set sequence length, and watch.

At short sequences (T<10), all methods work equally well. As T grows, compounding error separates them: off-policy O(εT²) explodes, on-policy O(εT) stays manageable. This is why every frontier lab uses on-policy for reasoning models.

Chapter 9: Connections & Open Problems

On-Policy Distillation sits at the intersection of three converging fields: knowledge distillation, reinforcement learning, and imitation learning.

The Unified View

The convergence is not just conceptual — it's infrastructural. OPD and RLHF need the same pipeline: rollout generation, scoring, filtering, gradient updates. They share 90% of code.

Core Equations (Cheat Sheet)

Open Problems

1. Distillation scaling laws. How does quality scale with teacher size, student size, rollout budget? No clean scaling laws exist yet.
2. Agent-level distillation. Multi-step tool-use chains require credit assignment across steps — fundamentally harder than single-turn.
3. Uncertainty-aware feedback. Teacher should abstain on OOD student prefixes rather than hallucinate confident noise.
4. Distillation + RL unification. They share 90% infrastructure. GKD and DPO are almost identical — both optimize divergences with different signal sources.

Industrial Adoption

• Qwen3: Off-policy warmup → GKD on-policy → GRPO exploration
• DeepSeek-V4: Multi-teacher R-KL OPD replacing their mixed RL stage entirely
• Gemma 2: KD in pre-training with on-policy logit distillation
• MiMo-V2: Multi-teacher + SCOPE token weighting → self-play refinement