Issue #106 Training Edwin Dong, Anurag Shinde
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: 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.
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), published at ICLR, set out to fix.
What You’ll Learn
Knowledge distillation from the ground up
Adding “on-policy” to distillation
Generalized knowledge distillation (GKD)
Combining GKD with reinforcement learning
How on-policy distillation outperforms
Why this all matters
