On-policy distillation
Watch on YouTube →
Overview
Mohit Iyyer connects supervised fine-tuning to reinforcement learning by defining policies, rollouts, rewards, and the difference between off-policy data and model-generated on-policy data. He explains exposure bias, soft-target distillation with teacher probabilities, and how forward versus reverse KL leads to different training behavior, including mode covering and mode seeking.
Key takeaways
- Off-policy training exposes a language model to ground-truth prefixes, while inference conditions on the model’s own outputs; this mismatch creates exposure bias and lets early mistakes compound.
- Soft-target distillation preserves information discarded by a sampled teacher response: matching a distribution such as 80% Claude, 15% GPT, and 5% Gemini is richer than treating one sampled token as 100% correct.
- Standard next-token negative log-likelihood is a special case of forward KL when the target is one-hot, because the target distribution’s entropy is zero.
- Reverse-KL on-policy distillation trains on student-generated rollouts and uses teacher token probabilities to score them, so it requires teacher logits or probabilities rather than only black-box API responses.
- Forward KL tends to cover the teacher’s modes, while reverse KL tends to select modes where the teacher assigns high probability; these objectives can produce different levels of response diversity.
- Expected rewards over enormous sequence spaces are estimated by sampling rollouts rather than enumerating them, and methods such as GRPO can use multiple samples per prompt.
Chapters
- Mohit Iyyer frames on-policy distillation as a bridge between supervised fine-tuning and reinforcement learning.
- Full fine-tuning updates every model parameter, while LoRA adapts low-rank matrices in linear layers with fewer trainable parameters.
- The announced homework covers PyTorch and transformer implementation, with a later assignment planned around post-training.
- In standard distillation, a stronger teacher such as GPT or Claude supplies responses that a student learns to imitate.
- In reinforcement learning terminology, the policy is the language model currently being trained; policy notation commonly uses π.
- As gradient descent changes the model, its policy distribution changes too.
- A rollout is a complete response sampled token by token from a policy after receiving a prompt, ending at the end-of-sequence token.
- A reward scores how good a rollout is, such as assigning a higher score to “Paris” than “Washington, D.C.” for a France-capital question.
- Dense rewards provide a score per token, while sparse rewards provide one score for the entire sequence, as in checking generated Python code with unit tests.
- Off-policy training uses responses from external sources such as humans, teacher models, or pretraining text; SFT and vanilla distillation are examples.
- On-policy training generates responses from the model being trained, then computes a learning signal on those generations.
- The driving analogy distinguishes watching an instructor’s video from driving yourself while receiving interactive corrections.
- During pretraining and SFT, the model predicts each next token while conditioning on a human-written ground-truth prefix.
- At inference, the model conditions on its own previous tokens, so an early mistake can alter and compound later predictions.
- This training-test mismatch is known as exposure bias and is a classic issue in imitation learning.
- On-policy feedback can tell a model to avoid a poor decision made during its own decoding, unlike training only on gold prefixes.
- Applying rollout-based learning during pretraining is computationally expensive because each sampled token requires another transformer forward pass.
- Autoregressive rollouts lose much of the parallelism available to standard next-token pretraining, limiting large-scale use.
- Geoffrey Hinton introduced distillation in 2014 for classification, a setting that illustrates the value of teacher probability distributions.
- A teacher distribution such as Claude 80%, GPT 15%, and Gemini 5% carries more information than a single sampled answer.
- Training only on the sampled token treats it as a one-hot target, whereas soft-target distillation teaches the student to match the teacher’s full distribution.
- Soft-target cross-entropy compares the teacher distribution P with the student distribution Q across the vocabulary at each token position.
- Commercial APIs such as Claude and GPT generally return sampled responses rather than their logits or full token probabilities.
- Repeated black-box sampling can approximate a distribution, but it becomes expensive for long prompts such as an entire code repository.
- The Kullback-Leibler divergence measures how different two distributions are, with P denoting the teacher and Q the student.
- KL divergence is nonnegative and reaches zero when P and Q are identical.
- KL is asymmetric: KL(P‖Q) and KL(Q‖P) weight mismatches differently, producing the forward and reverse forms.
- Forward KL weights each possible outcome by its probability under the teacher distribution P.
- A student’s mismatch on a teacher-likely token matters more than its mismatch on a token the teacher considers unlikely.
- Because the weighting comes from P, forward KL is not interchangeable with the reverse divergence, which weights outcomes under Q.
- Forward KL decomposes into cross-entropy from P to Q minus the entropy of P.
- For one-hot targets in pretraining, SFT, and vanilla distillation, the teacher entropy is zero.
- Consequently, standard negative log-likelihood on the correct next token is a special case of forward KL minimization.
- An expectation is a probability-weighted average, such as averaging a rollout’s reward R(X) under a model distribution.
- Enumerating every possible 2,000-token sequence is intractable, so expected rewards are estimated from sampled rollouts.
- Some reinforcement-learning methods, including GRPO, sample multiple rollouts per prompt; estimating an expectation requires samples, not direct access to the distribution.
- Forward KL takes its expectation over teacher-generated outcomes, while reverse KL takes its expectation over student-generated outcomes.
- In reverse-KL distillation, the teacher scores the student’s sampled tokens using teacher probabilities, so access to teacher probabilities is required.
- The teacher therefore acts as a scoring function for the student’s own rollout, rather than supplying the reference response as in conventional SFT.
- When a teacher has two modes but a limited student can represent only one, forward KL encourages the student to cover probability across both teacher modes.
- Reverse KL penalizes probability placed where the teacher has little probability, so the student may concentrate on one mode and ignore the other.
- These mode-covering and mode-seeking behaviors affect response diversity and help explain why post-training can narrow a model’s outputs.
Transcript excerpt
0:06 All right. Let's get started. So, today we are going to bridge between supervised fine-tuning, which we talked about in the previous couple classes, and reinforcement learning, which we will spend quite a bit of time on. So, today we'll be introducing a bunch of notation, and we will also be focusing a lot of our time on a new class new algorithm for this semester called on-policy distillation, which is very common these days in training frontier models. So, we will start this topic and likely not finish today. We'll have to finish on Monday, and then we will start talking about
0:55 reinforcement learning. But, the two things are related enough that after we cover this variant of distillation, you should be well positioned to understand how reinforcement learning differs from the distillation setup. Also, we'll try on the you know, in the next few weeks to keep unifying all these algorithms together as much as possible so that you understand their pros, cons, implementation, what they're actually optimizing. Okay, so before we start some logistics, your homework is finally released. So, please look into that, especially if you don't have that much experience coding
1:44 in PyTorch. Um, it is going to be a useful homework to understand some of the low-level implementation behind transformers. Once you guys submit that homework, we'll have a follow-up homework that has you kind of doing SFT and possibly even some RL on so so like more of a post training focused homework. Um, and we're just waiting for one person to finish their exam on Friday and we'll post the results and discuss them on Monday. Okay, great. So, um, last time um, we talked about SFT and instruction
2:39 tuning. And now this can either be implemented via full model tuning where we compute the DLD data for all of
This is the start of the transcript. Save this video free for the whole transcript, searchable and time-coded.
Related notes
Summary, takeaways, and chapters were generated by AI from the video's transcript and may contain errors. The video belongs to its creator, Mohit Iyyer.