On-policy distillation: small models learning from their own mistakes

Ordinary distillation trains a small model on text the big model wrote. Then, at test time, the small model has to continue from its own text, which it never practised on. On-policy distillation fixes this mismatch. The student writes, the teacher grades every token of what the student wrote, and the student learns from that dense feedback. This guide follows the video's path from soft targets to exposure bias, reinforcement learning, the KL-direction question and the applications.

Video by Jia-Bin Huang · Watch on YouTube

Going deeper: everything below is the full version, with some maths. On a first read, skim the headings and the "Why it matters for robotics" box, and come back for the rest later.

1. Watch

Chapter guide

The video's description doesn't list references. The sources below are the original papers behind each chapter's topic. The three application titles match the sections of Thinking Machines Lab's "On-Policy Distillation" post, which makes it the natural companion reading.

2. The ideas, explained

Soft targets and dark knowledge

Hinton, Vinyals and Dean (2015) trained a small model to match a large model's (or an ensemble's) output probabilities instead of only the hard labels. The teacher's logits z are softened with a temperature T so that small probabilities become visible, and the student is trained on a mix of the soft-target loss and the usual hard-label loss:

p_T(i) = exp(z_i / T) / Σ_j exp(z_j / T)          (teacher, softened)
L = α · T² · KL(p_T^teacher ‖ p_T^student)  +  (1 − α) · CE(y, p_1^student)

The T² factor keeps the soft-target gradients on a similar scale as T changes. The information in the wrong-answer probabilities (an image of a 2 that looks a bit like a 3 and a bit like a 7) is often called dark knowledge. It tells the student how the teacher generalises, which hard labels can't.

From one token to many: token-level and sequence-level distillation

For a language model, the "label" is a whole sequence. Two classic options:

• Token-level (word-level) distillation: on a fixed sequence (ground truth or teacher-generated), match the teacher's next-token distribution at every position.
• Sequence-level distillation (Kim and Rush, 2016): generate outputs with the teacher (they used beam search) and train the student on them with ordinary cross-entropy. This approximates matching the teacher's distribution over whole sequences, and their best student ran 10× faster than its teacher with little loss in quality.

token-level:     L = Σ_t KL( p_teacher(· | y_<t, x) ‖ q_θ(· | y_<t, x) ),   y fixed
sequence-level:  L = − Σ_t log q_θ(ŷ_t | ŷ_<t, x),   ŷ ~ teacher

Supervised fine-tuning on outputs written by a stronger model (as with many open reasoning datasets) is sequence-level distillation. Both forms are off-policy: the prefixes y_<t come from the data or the teacher, never from the student.

Exposure bias

At training time the student always conditions on a "good" prefix. At inference it conditions on its own, possibly flawed, prefix. One early slip puts it in a state it never trained on, where it is more likely to slip again. Bengio et al. (2015) proposed scheduled sampling, gradually feeding the model its own predictions during training. The same problem was formalised years earlier in imitation learning by Ross, Gordon and Bagnell (2011): a policy trained only on expert states can see its errors grow quadratically with the horizon T, and their algorithm DAgger fixes this by running the learner's policy, asking the expert to label the states the learner actually visits, and training on the growing dataset. On-policy distillation is DAgger for language models, with the teacher's full distribution as the label.

behaviour cloning:   error can grow like  O(ε · T²)
DAgger (on-policy):  error grows like     O(ε · T)       (Ross et al., 2011)

On-policy reinforcement learning

RL also trains on the student's own samples, so it doesn't suffer exposure bias. The catch is feedback density: a verifier says "correct" or "wrong" once per sequence, so a 2,000-token solution gets one number, with no indication of which step went wrong. Thinking Machines' post puts it as RL teaching O(1) bits per episode, against O(N) bits for distillation, where N is the number of tokens. The policy-gradient machinery (REINFORCE, baselines, PPO, GRPO) is covered in Video 06 and 30u30 Day 29.

On-policy distillation

Agarwal et al. (Google DeepMind, 2023) introduced Generalized Knowledge Distillation (GKD; the arXiv title is "On-Policy Distillation of Language Models: Learning from Self-Generated Mistakes"). The student is trained on its own sampled outputs, with the teacher's token-level distribution as the target, and the divergence can be chosen freely:

GKD:  L(θ) = (1 − λ) · E_{(x,y) ~ data}[ D(p_T ‖ q_θ)(y | x) ]
           +      λ  · E_{x ~ X, y ~ q_θ(·|x)}[ D(p_T ‖ q_θ)(y | x) ]

D(p_T ‖ q_θ)(y|x) = (1/|y|) Σ_t D( p_T(· | y_<t, x) ‖ q_θ(· | y_<t, x) )
D ∈ { forward KL, reverse KL, generalised JSD(β) }

λ is the fraction of student-generated data; λ = 1 is fully on-policy. Gradients aren't taken through the sampling step, which keeps training stable and cheap. The paper reports gains on summarisation, translation and arithmetic reasoning, and shows the same loss can be added to RL fine-tuning. MiniLLM (Gu, Dong, Wei and Huang, 2023) reached a related recipe from the reverse-KL side, deriving an on-policy policy-gradient optimisation of sequence-level reverse KL.

Thinking Machines Lab's post (Kevin Lu and colleagues, October 2025) uses the simplest version: per-token reverse KL on student samples, implemented as RL with the reward set to the negative KL and a discount of zero:

sample y ~ π_θ(· | x)                                   (student rollout)
advantage_t = −( log π_θ(y_t | y_<t, x) − log π_teacher(y_t | y_<t, x) )
update π_θ with the usual importance-sampled policy-gradient loss

Only one teacher forward pass per rollout is needed, rollouts can be partial, and no reward model is required. The post describes it as a small change to an RL loop with a KL penalty: swap the reference model for the teacher.

Why the KL direction matters

forward KL   KL(p ‖ q) = E_{y~p}[ log p(y) − log q(y) ]   samples from the teacher
reverse KL   KL(q ‖ p) = E_{y~q}[ log q(y) − log p(y) ]   samples from the student

Forward KL punishes the student heavily wherever the teacher has mass and the student doesn't, so the student spreads out to cover everything the teacher might do: mode-covering. A small student that can't represent all the teacher's modes ends up putting mass in the gaps between them, which for text means bland or incoherent mixtures. Reverse KL punishes the student wherever it has mass and the teacher doesn't, so it concentrates on what it can do well: mode-seeking. That suits a low-capacity student, and the expectation is over the student's own samples, which is what makes it naturally on-policy. The trade-off is lower diversity; GKD's generalised JSD(β) interpolates between the two. Thinking Machines adds that per-token reverse KL is hard to "hack", since low KL to the teacher means behaviour the teacher itself finds likely.

Applications and empirical results

All numbers in this section are as reported by Thinking Machines Lab's post, which in turn draws on the Qwen3 technical report for one comparison. They are the authors' measurements in specific setups, not general laws.

• Reasoning. The post quotes the Qwen3 report: from the same off-policy-distilled starting point, RL reached 67.6% on AIME'24 using 17,920 GPU hours, while on-policy distillation reached 74.4% using 1,800 GPU hours. In its own replication (Qwen3-8B-Base student, starting from a 400k-prompt SFT checkpoint at 60%), on-policy distillation reached about 70% in roughly 150 steps, which the post estimates as 9–30× cheaper than continuing SFT to the same score.
• Personalisation. Mid-training Qwen3-8B on internal company documents raised an internal-knowledge score but lowered instruction following (IF-eval), and no data mix fully preserved it. Running on-policy distillation afterwards, with the original Qwen3-8B as teacher on general chat prompts, brought IF-eval from 79% back to 83% (the original was 85%) while keeping the new knowledge.
• Data reuse. Training repeatedly on a single maths prompt (20 steps of 256 rollouts) approximately matched the teacher's AIME'24 score. With RL, repeated epochs on the same prompt tend to memorise the final answer; the distillation target is a whole distribution, so it doesn't collapse the same way.
• Dense versus sparse. Distilling an RL-trained teacher back into its own base model recovered the teacher's performance in about 7–10× fewer gradient steps than the RL run took, which the post estimates as 50–100× less compute overall.

A practical caveat from the same post: on-policy distillation works best when the student already puts some probability on what the teacher would do, which is why the experiments start from an SFT or mid-trained checkpoint. Reverse KL can only sharpen within the student's existing support; SFT (forward KL) is what adds new support.

On-policy self-distillation

The last two chapters cover self-distillation, where the teacher is the same model (or a copy of it) rather than a bigger one. The video's description doesn't name the specific papers, so this guide stays general. The common idea is that some version of the model has an advantage the student version lacks: an earlier checkpoint that still has a skill, or the same weights given extra context the student won't have at test time. The student is trained on its own samples to match that teacher token by token. Thinking Machines' personalisation result is one concrete case for continual learning: the earlier Qwen3-8B acted as teacher to restore what fine-tuning had erased, and the post suggests alternating "learn new data" and "distil old behaviour back" phases. It also reports that plain SFT on a model's own samples still degraded instruction following in its tests, which it takes as evidence that the on-policy, teacher-graded signal is what matters, not just training on self-generated text.

The two-by-two

3. Where the course goes deeper

Why it matters for robotics

Self-check

References

The video's description lists no references. Primary sources for this guide:

• On-Policy Distillation of Language Models: Learning from Self-Generated Mistakes (GKD) (Agarwal, Vieillard, Zhou, Stanczyk et al., 2023)
• On-Policy Distillation (Kevin Lu and others, Thinking Machines Lab, October 2025)
• MiniLLM: On-Policy Distillation of Large Language Models (Gu, Dong, Wei, Huang, 2023)
• Distilling the Knowledge in a Neural Network (Hinton, Vinyals, Dean, 2015)
• Sequence-Level Knowledge Distillation (Kim, Rush, 2016)
• Scheduled Sampling for Sequence Prediction with Recurrent Neural Networks (Bengio, Vinyals, Jaitly, Shazeer, 2015)
• A Reduction of Imitation Learning and Structured Prediction to No-Regret Online Learning (DAgger) (Ross, Gordon, Bagnell, 2011)
• Qwen3 Technical Report (Qwen Team, 2025), source of the RL vs on-policy distillation comparison quoted above