# On-Policy Distillation for LLMs, Simply Explained

#106 Jun 15 Training

https://www.aibutsimple.com/p/on-policy-distillation-for-llms-simply-explained

Large language models are impressive. They summarize articles, translate languages, solve math problems, and write code.

However, one of the key problems surrounding them is that they are expensive. Serving a model with billions of parameters costs real money at real scale.

This has made model compression a critical topic in modern AI research.

One of the most effective compression strategies is [knowledge distillation](https://newsletter.aibutsimple.com/p/llm-knowledge-distillation-simply-explained): train a smaller student model to replicate the behavior of a larger teacher model. The student runs faster, uses less memory, and can be deployed cheaply.

This is simple in principle but harder in practice.

---

### The Problem With LLM Distillation

Here is the core issue with how distillation is done today.

During training, the student model learns from a fixed dataset of sequences. These sequences are either from the ground-truth data (teacher forcing) or generated by the teacher in advance.

During inference, the student generates its own sequences from scratch, one token at a time.

The problem, however, is that these are two completely different situations.

The sequences the student encounters while training look nothing like the sequences it will generate when deployed. More formally, the distribution of contexts seen during training differs from the distribution encountered at inference.

This gap is called the train-inference distribution mismatch, and it is a well-known problem in imitation learning research.

![](https://media.beehiiv.com/cdn-cgi/image/fit=scale-down,width=1200,format=auto,onerror=redirect,quality=80/uploads/asset/file/8b78b381-5df8-421f-8d16-8e3e33acd112/Train-Inference_Mismatch.png)

Here is why the mismatch matters. In auto-regressive models, every token depends on all the previous tokens, meaning that the student is conditioned on its own outputs.

If the student makes a slightly wrong prediction early in a sequence, that error feeds forward into the next token, and the next, and the next. This way, the errors compound.

The student was trained on clean, expert-level sequences. It has never learned how to recover from its own mistakes.

This is the central problem that [Agarwal et al. (2024)](https://arxiv.org/pdf/2306.13649v3?utm_source=newsletter.aibutsimple.com&utm_medium=referral&utm_campaign=on-policy-distillation-for-llms-simply-explained), published at ICLR, set out to fix.

---

### What You’ll Learn

1. Knowledge distillation from the ground up
2. Adding “on-policy” to distillation
3. Generalized knowledge distillation (GKD)
4. Combining GKD with reinforcement learning
5. How on-policy distillation outperforms
6. Why this all matters

### What’s Helpful to Know

- Auto-Regressive Generation
  - A model that produces one token at a time, conditioning each new token on everything generated so far. They are the core generation method of language models like [GPT](https://openai.com/index/introducing-gpt-5-5/?utm_source=newsletter.aibutsimple.com&utm_medium=referral&utm_campaign=on-policy-distillation-for-llms-simply-explained) and [T5](https://huggingface.co/docs/transformers/en/model_doc/t5?utm_source=newsletter.aibutsimple.com&utm_medium=referral&utm_campaign=on-policy-distillation-for-llms-simply-explained).
- [KL Divergence](https://newsletter.aibutsimple.com/p/kl-divergence-simply-explained)
  - A measure of how different two probability distributions are. For discrete distributions P and Q:
    ![](https://media.beehiiv.com/cdn-cgi/image/fit=scale-down,width=1200,format=auto,onerror=redirect,quality=80/uploads/asset/file/d7d79b23-e6f0-46f4-ac7a-ae62b2d707ca/image.png)
  - KL divergence is not symmetric. DKL(P || Q) is not DKL(Q || P). This asymmetry turns out to matter a great deal for distillation.
- Forward vs. Reverse KL
  - DKL(P || Q) is called the forward KL. Minimizing it forces Q to cover all the regions where P has mass. It is mode-covering.
  - DKL(Q || P) is called the reverse KL. Minimizing it forces Q to concentrate only where P is large. It is mode-seeking.

![](https://media.beehiiv.com/cdn-cgi/image/fit=scale-down,width=1200,format=auto,onerror=redirect,quality=80/uploads/asset/file/a92ba1b9-2bb2-406c-9ded-764e6127897f/Forward_vs_Reverse_KL.png)

- Jensen-Shannon Divergence (JSD)
  - A symmetric and bounded divergence that interpolates between forward and reverse KL. For a parameter β between (0,1):
    ![](https://media.beehiiv.com/cdn-cgi/image/fit=scale-down,width=1200,format=auto,onerror=redirect,quality=80/uploads/asset/file/ca2d3714-516e-4ff2-aa8f-86b2462c4b07/image.png)
  - Here, M = βP + (1-β)Q. When β approaches 0, JSD behaves like forward KL. When β approaches 1, it behaves like reverse KL.

![](https://media.beehiiv.com/cdn-cgi/image/fit=scale-down,width=1200,format=auto,onerror=redirect,quality=80/uploads/asset/file/47cd0296-1c63-4687-9cbf-7979ae6d48a4/JSD_Interpolation_for_OPD__1_.png)

- Imitation Learning
  - A family of techniques where an agent learns a policy by mimicking an expert. The key challenge is that mistakes during deployment can compound in ways the expert never demonstrated.

---

### What Is Knowledge Distillation?

You have a large teacher model. It performs well but is expensive to run.

You want a small student model. It should perform nearly as well but at a fraction of the cost.

Knowledge distillation (KD) trains the student not just to predict the right answer, but to match the teacher's entire probability distribution over possible next tokens.

This is a much richer training signal than simply predicting the correct label.

Think of it this way. A teacher does not just tell a student, "the answer is 4." A good teacher explains why 4 is correct and also explains why 3 and 5 are close but wrong.

Token-level distributions from the teacher carry exactly that kind of nuanced information. And the student learns to not do something the teacher wouldn't.

There are two dominant approaches in the existing literature.

Supervised KD trains the student to minimize the forward KL divergence between the teacher's distributions and its own over a fixed dataset (X, Y):

![](https://media.beehiiv.com/cdn-cgi/image/fit=scale-down,width=1200,format=auto,onerror=redirect,quality=80/uploads/asset/file/f98c423e-8527-4fbd-be24-692c01ea205c/image.png)

Here, pT is the teacher, and psθ is the student with learnable parameters. The divergence is computed token by token and averaged over the sequence.

Sequence-level KD (SeqKD) generates a large set of sequences from the teacher, then trains the student on those sequences using maximum likelihood. It is essentially supervised fine-tuning on teacher-generated data.

Both approaches share the same flaw. The student trains on a fixed set of sequences it did not produce itself, so the distribution mismatch remains.

---

### The On-Policy Insight

The paper draws a direct connection to imitation learning, which studies how an agent can learn to mimic an expert.

The naive approach in imitation learning is behavioral cloning: record what the expert does, then train the agent to copy it. This works when the agent stays close to the expert's trajectory.

But when the agent makes a mistake and drifts into a new state, it has no training signal to rely on; it has simply never been in this situation before.

The solution is on-policy data collection. Let the agent act in the environment, observe the states it actually ends up in, and query the expert for the right action at those states.

![](https://media.beehiiv.com/cdn-cgi/image/fit=scale-down,width=1200,format=auto,onerror=redirect,quality=80/uploads/asset/file/2235c380-26e9-4aa1-9040-9d9bd197808c/GKD_and_RL__3_.png)

The agent learns from its own mistakes, corrected by expert feedback.

This is exactly the idea the paper brings to language model distillation.

Instead of training the student on fixed sequences, let the student generate its own output sequences.

Then, ask the teacher to evaluate those sequences, token by token, and train the student on this feedback.

The on-policy loss captures this directly. Given an input x, the student samples a full output sequence y~ps(·|x) and then imitates the teacher's token-level distributions along that sequence:

![](https://media.beehiiv.com/cdn-cgi/image/fit=scale-down,width=1200,format=auto,onerror=redirect,quality=80/uploads/asset/file/a8a91401-10fd-4725-9631-0945ceae08f1/image.png)

The gradient does not flow back through the student's sampling distribution ps(|x). This is a deliberate choice: backpropagating through sampling introduces high variance and instability. Treating the sampled sequences as fixed makes training stable and computationally efficient.

The student now trains on the sequences it is actually likely to generate. The distribution seen at training time matches the distribution encountered at inference. The mismatch is resolved.

---

If you enjoy our content and custom visuals, consider sharing this newsletter with others or [upgrading](http://www.aibutsimple.com/pro?utm_source=newsletter.aibutsimple.com&utm_medium=referral&utm_campaign=on-policy-distillation-for-llms-simply-explained) so we can keep doing what we do.

### [Code AI Projects With AI, But Simple PRO](https://www.aibutsimple.com/pro?utm_source=newsletter.aibutsimple.com&utm_medium=referral&utm_campaign=on-policy-distillation-for-llms-simply-explained)

![](https://media.beehiiv.com/cdn-cgi/image/fit=scale-down,width=1200,format=auto,onerror=redirect,quality=80/uploads/asset/file/963eba00-ab31-4912-8c1f-1470a52f9ecc/aibutsimple_academy__2_.png)

Supporting us by purchasing a PRO membership will get you:

- Exclusive Hands-on Tutorials With Code
- Real-World Interview Questions and Cases (From Real Companies)
- Full Math Explanations and Beautiful Visualizations
- Ad-Free Reading for an Uninterrupted Learning Experience

[Stay Ahead in AI With PRO](https://www.aibutsimple.com/pro?utm_source=newsletter.aibutsimple.com&utm_medium=referral&utm_campaign=on-policy-distillation-for-llms-simply-explained)

---

### Generalized Knowledge Distillation (GKD)

The paper formalizes the on-policy addition into Generalized Knowledge Distillation (GKD). GKD has 2 main tunable components that combine both on-policy and off-policy distillation.

The first is the student data fraction λ, which is between 0 and 1.

- When λ = 1, the student trains entirely on its own generated sequences.
- When λ = 0, it trains on a fixed dataset.
- Any value in between mixes the two.

Our previously mentioned supervised KD corresponds to λ = 0, while pure on-policy KD corresponds to λ = 1.

![](https://media.beehiiv.com/cdn-cgi/image/fit=scale-down,width=1200,format=auto,onerror=redirect,quality=80/uploads/asset/file/333be65a-be01-426a-baff-a19729bc17c2/Generalalized_Knowledge_Distillation__GKD__Algorithm__1_.png)

The second is the choice of divergence D. The full GKD objective is a linear (convex) interpolation between on and off-policy objectives:

![](https://media.beehiiv.com/cdn-cgi/image/fit=scale-down,width=1200,format=auto,onerror=redirect,quality=80/uploads/asset/file/ca3c5165-06fe-4f79-8201-9e06f40d967f/image.png)

The divergence D between teacher and student distributions is computed token by token (token-level divergence) and averaged over the sequence length.

![](https://media.beehiiv.com/cdn-cgi/image/fit=scale-down,width=1200,format=auto,onerror=redirect,quality=80/uploads/asset/file/6a9d6dde-9f5e-4cd3-862b-568735317563/image.png)

This is flexible, as we now have options for both our divergence and data fraction:

- Plug in forward KL and λ=0, and you recover supervised KD.
- Plug in forward KL and λ=1, and you get pure on-policy KD.
- Plug in Jensen-Shannon Divergence (JSD) or reverse KL, and you get new variants that can outperform both.

![](https://media.beehiiv.com/cdn-cgi/image/fit=scale-down,width=1200,format=auto,onerror=redirect,quality=80/uploads/asset/file/12b1510c-a575-4fda-b30a-c0e09beefd28/Token-level_teacher_feedback__3_.png)

Here, we should ask ourselves: why does the divergence choice matter?

Forward KL is mode-covering. It forces the student to assign probability mass everywhere the teacher does.

When the student has far fewer parameters than the teacher, it cannot faithfully represent every mode of the teacher's distribution. It spreads mass thin across unlikely tokens. This can hurt generation quality, especially at a high temperature.

Reverse KL is mode-seeking. It forces the student to concentrate on the regions where the teacher assigns high probability.

The student learns to generate the teacher's most likely outputs well, but at the cost of diversity. It ignores low-probability regions entirely.

JSD interpolates between them, and the paper finds that the optimal choice is task-dependent: mode-seeking divergences work best when evaluation uses temperature sampling, while forward KL performs comparably with greedy decoding.

---

### Combining GKD with Reinforcement Learning

One of the more interesting contributions is showing that on-policy GKD integrates naturally into RL fine-tuning.

RL fine-tuning of language models, particularly [RLHF](https://newsletter.aibutsimple.com/p/reinforcement-learning-for-large-language-models-simply-explained) and RLAIF, has become standard.

The model is trained to maximize a reward signal while being regularized (also typically using KL divergence) to stay close to a reference policy, typically its own supervised fine-tuned checkpoint.

In GKD, instead of regularizing toward the student's own initial checkpoint, we will regularize toward the teacher, creating a combined objective of:

![](https://media.beehiiv.com/cdn-cgi/image/fit=scale-down,width=1200,format=auto,onerror=redirect,quality=80/uploads/asset/file/65efe569-bb6e-428d-9ee1-747a5df7b48d/image.png)

The first term is the RL objective, which represents the expected reward r(y) over student-generated sequences we should maximize.

The second term is the on-policy distillation loss that incentivizes the student to be close to the teacher while generating.

Note that this objective is very similar to the GKD objective. We have simply swapped supervised KD with the RL reward. The interpolation parameter is simply given a different Greek symbol (α instead of λ).

- Same idea: α = 1 means pure distillation, and α = 0 is pure RL.

Both terms operate on student-generated sequences, so they share the same sampling process.

But what about compute? The additional computational overhead compared to standard RL fine-tuning is small: you only need to run the teacher to get its log-probabilities, not to generate full sequences.

![](https://media.beehiiv.com/cdn-cgi/image/fit=scale-down,width=1200,format=auto,onerror=redirect,quality=80/uploads/asset/file/a5980fd7-3935-4545-9d96-21880a3a6922/GKD_and_RL__2_.png)

The paper tests this on summarization with factual consistency as the reward. GKD + RL improves factual consistency substantially over the teacher model while also improving summarization quality. Neither distillation alone nor RL alone achieves both at the same time.

---

### Experimental Results

The paper tests GKD across 3 tasks:

- Abstractive summarization (XSum)
- Machine translation (WMT English to German)
- Arithmetic reasoning (GSM8K)

It also includes a task-agnostic instruction-tuning experiment on the FLAN dataset.

Students are T5-small (77M), T5-base (250M), and T5-large (800M). The teacher in most experiments is T5-XL (approximately 3B parameters), between 4 and 38 times larger than the students.

![](https://media.beehiiv.com/cdn-cgi/image/fit=scale-down,width=1200,format=auto,onerror=redirect,quality=80/uploads/asset/file/7f88a38c-2518-4ecf-b5bb-90481a80d659/image.png)

On summarization (XSum), on-policy GKD consistently outperforms supervised KD, SeqKD, and more across all student sizes. GKD also achieves 2.1x the performance gain of supervised KD over the base fine-tuned student. GKD with T5-Large exceeds the few-shot performance of PaLM (540B), a model roughly 7,000 times larger.

In terms of data efficiency, on-policy GKD trained on just 5% of XSum, without any ground-truth summaries, outperforms supervised KD and ImitKD trained on the full dataset.

On translation (WMT), on-policy GKD with JSD divergences outperforms purely supervised off-policy variants across all student sizes. Using entirely student-generated sequences consistently performs best.

In reasoning (GSM8K), on-policy GKD substantially outperforms all baselines across all student sizes, achieving 90% relative improvement over the base fine-tuned student compared to significantly smaller gains from supervised distillation.

On the FLAN instruction-tuning dataset, on-policy GKD with reverse KL achieves 2% absolute accuracy improvement on MMLU and 1% on BBH over supervised KD, on held-out benchmarks.

---

### Why GKD Matters

Distillation is one of the most practically important techniques in applied machine learning. It makes a capable research model deployable at scale.

GKD addresses the distribution mismatch with no architectural changes, empirically showing consistent improvements over supervised distillation across tasks, student sizes, and evaluation setups.

The paper provides practical guidance: start by tweaking λ, use forward KL for greedy evaluation tasks, and use reverse KL or JSD (β=0.9) for temperature sampling tasks.

The integration with RL fine-tuning is particularly valuable (and it works empirically). Standard RLHF already applies a KL penalty against the initial model, but replacing it with a stronger teacher is a nice improvement.

---

### What to Take Away

The main idea is that fixed training data creates a gap between what the student sees during training and what it generates at inference.

Letting the student generate its own training sequences, with the teacher providing token-level feedback on those sequences, closes that gap.

The student learns better by learning from its own mistakes.

The key equation is the GKD objective. It mixes fixed data and student-generated data through the parameter λ, and it generalizes over any divergence D through a clean token-averaged formulation.

For anyone working on compressing or aligning language models, GKD is a well-tested and easy-to-integrate upgrade to the standard distillation toolkit.

---

Here’s a special thanks to our [biggest supporters](https://www.aibutsimple.com/upgrade?utm_source=newsletter.aibutsimple.com&utm_medium=referral&utm_campaign=on-policy-distillation-for-llms-simply-explained):

- Sushant Waidande
- Sai Krishna Pashapu (AI Developer, BTech, [LinkedIn](https://www.linkedin.com/in/sai-krishna-pashapu-997358283/?utm_source=newsletter.aibutsimple.com&utm_medium=referral&utm_campaign=on-policy-distillation-for-llms-simply-explained), [GitHub](https://github.com/saikrishna-pashapu?utm_source=newsletter.aibutsimple.com&utm_medium=referral&utm_campaign=on-policy-distillation-for-llms-simply-explained))
- Bob Boncella
- Demet Yazici
- Imran Khan

---

If you enjoy our content, consider supporting us so we can keep doing what we do. Please share this with a friend!

Want to reach 8000+ ML engineers? Advertising Inquiries:

[WORK WITH US 📈](https://tally.so/r/A7Dgdk?utm_source=newsletter.aibutsimple.com&utm_medium=referral&utm_campaign=on-policy-distillation-for-llms-simply-explained)

Feedback or inquiries? Send us an email at [edwin@aibutsimple.com](mailto:edwin@aibutsimple.com).

---

That’s it for this week’s issue of AI, but simple. See you next week!

—AI, but simple team