Most teams doing knowledge distillation do the same thing. Run a large teacher model on a dataset. Save the outputs. Fine-tune the small student model on those outputs. Done!
That is off-policy distillation. And it has a flaw that nobody talks about when they share their benchmark numbers: you are training the student to navigate the teacher’s world, not its own. At inference time, the student operates in its own world, the one it was never trained for.
Thinking Machines Lab published this blogpost by thinking machines lab showing they trained a 0.5B model to match a 1.5B teacher’s performance on math reasoning at 9–30x lower compute than RL, just by switching from off-policy to on-policy. The training loop fits in 10 lines of PyTorch.
This post explains why the switch matters and exactly what changes in the code.
Table of Contents
- What Off-Policy Distillation Actually Is
- The Compounding Error Problem
- What On-Policy Distillation Changes
- The Loss Function: Reverse KL
- The Code: 10 Lines That Change Everything
- The Numbers from Thinking Machines
- Run It Yourself: Colab Demo
- Summary
What Off-Policy Distillation Actually Is
When people say “distillation” they almost always mean this:
# Step 1: run teacher on your dataset and save outputs
teacher_outputs = [teacher_model.generate(prompt) for prompt in dataset]
# Step 2: fine-tune student on those outputs (standard SFT)
student_model.train_on(teacher_outputs) # cross-entropy loss
This is off-policy distillation. The word “off-policy” means: the data you train on was generated by a different policy (model) than the one you are training. The teacher generated the trajectories. The student just imitates them.
The loss is standard cross-entropy, at each token position, minimise $-\log P_\text{student}(\text{correct token})$ where “correct” means whatever the teacher wrote:
\[\mathcal{L}_\text{SFT} = -\frac{1}{N}\sum_{t=1}^{N} \log P_\theta(x_t \mid x_1, \ldots, x_{t-1})\]This works. Many excellent models are built this way. But it has a structural flaw that gets worse as sequences get longer.
The Compounding Error Problem
Take a concrete problem: solve 3x + 7 = 22.
The teacher’s training data looks like this:
Step 1: subtract 7 from both sides → 3x = 15
Step 2: divide both sides by 3 → x = 5 ✓
The student trains on hundreds of examples like this. It learns: when you see a linear equation, subtract the constant first, then divide. Good so far.
Now deploy the student and give it the same problem. At step 1, the student is a 0.5B model with limited capacity. It makes a different move:
Step 1: multiply both sides by 3 → 9x + 21 = 66 ✗
This is not in the training data. The teacher never wrote this. Now what? The student has to continue from 9x + 21 = 66, a state it has zero training signal for. It was only ever taught how to continue from correct intermediate states. It has never seen what to do after it goes wrong.
So it keeps generating. It produces the next token, and the next, and each one is conditioned on a context that is further and further from anything it trained on. Every mistake makes the next token harder to get right. Errors compound.
The student went off track at step one. It took a path the teacher never explored. Everything downstream is uncharted. The student has no training signal for how to recover, it just keeps generating in a distribution it was never trained on.
This is called compounding error or exposure bias. It is the fundamental structural problem with off-policy training. Each mistake shifts the student further from the distribution it trained on, and errors accumulate.
Off-policy distillation trains the student to be good in the teacher’s world. But at inference time, the student lives in its own world, made of its own mistakes.
What On-Policy Distillation Changes
The fix is simple in concept. Instead of generating training data from the teacher, generate it from the student. Then use the teacher only to grade it.
# On-policy distillation — the key change
# Step 1: student generates its OWN trajectories (on-policy)
student_outputs = student_model.generate(prompts, do_sample=True)
# Step 2: teacher scores every token of the student's output
teacher_log_probs = teacher_model.score_tokens(student_outputs)
# Step 3: train student to be more like teacher ON ITS OWN OUTPUTS
loss = reverse_kl(student_outputs, teacher_log_probs)
The student now trains on trajectories it actually generates, including the wrong paths and the states it reaches after mistakes. The teacher grades every token: “that word was a mistake, that word was good, that transition was wrong.” The student learns to correct itself in exactly the situations it will face at inference.
The training data distribution now matches the inference distribution. The compounding error problem disappears.
| Off-Policy Distillation | On-Policy Distillation | |
|---|---|---|
| Training data from | Teacher’s generations | Student’s own generations |
| Teacher’s role | Source of examples | Token-level grader |
| Trains in | Teacher’s distribution | Student’s own distribution |
| Compounding error | Yes, grows with sequence length | No, trains on its own mistakes |
| Reward density | Dense (every token) | Dense (every token) |
The Loss Function: Reverse KL
The per-token grading from the teacher is formalised as the reverse KL divergence.
There are two ways to measure the gap between two distributions.
Forward KL : sample from the teacher $P$, penalise the student $Q$ for missing teacher mass:
\[KL(P \| Q) = \sum_x P(x) \log \frac{P(x)}{Q(x)}\]This is mean-seeking. It forces the student to cover everything the teacher covers, even if the student spreads thin.
Reverse KL : sample from the student $Q$, penalise the student for going where the teacher would not:
\[KL(Q \| P) = \mathbb{E}_{x \sim Q}\left[\log Q(x) - \log P(x)\right]\]This is mode-seeking. The student commits to one correct reasoning style rather than hedging. And crucially, the expectation is over $Q$, meaning the tokens the student actually generates. This is what makes it on-policy.
Per-token, the penalty is:
\[\text{penalty at token } t = \log P_\text{student}(x_t) - \log P_\text{teacher}(x_t)\]High when the student was confident about a token the teacher would never write. Near zero when both agree. Exactly zero when student equals teacher.
Reverse KL also decomposes as:
\[KL(Q \| P) = \underbrace{H(Q, P)}_{\text{cross-entropy}} - \underbrace{H(Q)}_{\text{student entropy}}\]The student entropy term acts as a regulariser, it prevents the student from collapsing to a single deterministic output and losing diversity.
The Code: 10 Lines That Change Everything
Here is the full on-policy distillation training step in raw PyTorch.
# 1. Rollout: student generates on-policy (no grad just sampling)
student_model.eval()
with torch.no_grad():
full_out = student_model.generate(
prompt_ids, max_new_tokens=150,
do_sample=True, temperature=0.9,
pad_token_id=tokenizer.eos_token_id
)
response_ids = full_out[:, prompt_ids.shape[1]:] # strip prompt tokens
# 2. Student log-probs WITH gradients (single parallel forward pass)
student_model.train()
student_logprobs = compute_token_logprobs(student_model, prompt_ids, response_ids)
# 3. Teacher log-probs WITHOUT gradients (frozen, one forward pass)
with torch.no_grad():
teacher_logprobs = compute_token_logprobs(teacher_model, prompt_ids, response_ids)
# 4. Reverse KL, student learns from its own mistakes
reverse_kl = student_logprobs - teacher_logprobs.detach() # [B, seq_len]
loss = (reverse_kl * mask).sum() / mask.sum()
loss.backward()
optimizer.step()
compute_token_logprobs does one forward pass and picks out the log-probability for each response token:
def compute_token_logprobs(model, prompt_ids, response_ids):
full_ids = torch.cat([prompt_ids, response_ids], dim=1)
prompt_len = prompt_ids.shape[1]
resp_len = response_ids.shape[1]
logits = model(input_ids=full_ids).logits # [B, L, V]
# logit[i] predicts token[i+1] so response token j uses logit[prompt_len+j-1]
resp_logits = logits[:, prompt_len - 1 : prompt_len + resp_len - 1, :]
log_probs = F.log_softmax(resp_logits, dim=-1) # [B, resp_len, V]
return log_probs.gather(-1, response_ids.unsqueeze(-1)).squeeze(-1) # [B, resp_len]
Notice the two student forward passes:
- The first is inside
generate()autoregressive, token by token, no grad. It determines what tokens to produce. - The second is inside
compute_token_logprobsparallel over the full sequence, with grad. It computes the log-probabilities the optimizer differentiates through. Sampling is not differentiable, so you need this second pass to get gradients.
The Numbers from Thinking Machines
Thinking Machines applied this on AIME’24 competition math using Qwen3-8B as student and Qwen3-32B as teacher.
| Method | AIME’24 | GPU Hours |
|---|---|---|
| Off-policy SFT (400K prompts) | 60% | — |
| Reinforcement learning (Qwen3 report) | 67.6% | 17,920 |
| On-policy distillation | 70% | ~1,800 |
On-policy distillation beat RL’s score at one-tenth the compute. The reason is information density:
\[\text{RL signal per episode} = O(1) \qquad \text{On-policy distillation signal per episode} = O(N)\]where $N$ is the number of tokens. At 150 tokens per rollout that is a 150× difference in learning signal per episode. Fewer episodes needed. Less compute.
They also found you can train on a single prompt repeatedly. RL overfits, it memorises the answer. On-policy distillation keeps sampling diverse reasoning chains from the student and grading them, converging to the teacher’s distribution on that problem without overfitting.
The continual learning result is equally striking. After fine-tuning Qwen3-8B on company documents, instruction-following (IF-eval) dropped from 85% to 45%. Running on-policy distillation from the original frozen Qwen3-8B as teacher, with no domain data involved, recovered IF-eval to 83%.
This unlocks a practical continual learning loop: fine-tune on new knowledge, then on-policy distill from the previous version to recover behaviour. Repeat indefinitely.
Run It Yourself: Colab Demo
I built a self-contained notebook that runs this exact loop on GSM8K grade-school math with Qwen2.5-0.5B as student and Qwen2.5-1.5B as teacher. Runs on a Colab T4 free tier ~35 minutes, ~9 GB VRAM.
The notebook includes:
- Baseline accuracy on 200 GSM8K test examples before training
- The full training loop from this post, step by step
- Accuracy after 250 steps of on-policy distillation
- A per-token KL visualisation: which exact tokens in a wrong answer the teacher most strongly disagreed with
Summary
In this post, we learnt that Off-policy distillation (standard SFT on teacher outputs) trains the student in the teacher’s world. The student never sees its own mistakes. Compounding error means every wrong token at inference takes the student further from anything it was trained on.
On the other hand, On-policy distillation generates rollouts from the student and uses the teacher only to score them token by token. The training distribution matches the inference distribution. The compounding error problem disappears.
If on-policy distillation copies what RL learns at 1/10th the compute, the question is how much of RL’s expensive search is genuinely necessary versus just being the expensive prerequisite for cheap distillation to finish the job.
Did you find this post useful? I am curious to hear from you in comments below.
Comments