The "aha" moment with On Policy Distillation
The promise of knowledge distillation is quite interesting. It is a machine learning technique that lets you distill, or transfer, knowledge from a more capable teacher network into a less capable student network.
On-Policy Distillation (OPD) is a variant of knowledge distillation that has recently gained popularity in the large language model (LLM) literature, because it can train new LLMs using dense token-level feedback from stronger LLMs.
In simple terms, OPD allows AI labs to use their existing LLMs to train knowledge directly into their newer models.
In this article, we will discuss various aspects of On-Policy Distillation - the main design problems, algorithmic branches, when to use it, and how it differs from other LLM training techniques.
This article is also available as a video 👇🏼
First, an overview of various LLM training techniques
Before looking into OPD, let's get some context about other popular LLM training algorithms first.
Self-supervised Pretraining
What you need: A massive dataset of raw text, code, documents, etc. You don't need manually labeled answers, preferences, rewards, or a teacher model.
When you use it: This is the initial training stage. You're creating the base model, so often you start from scratch, or continue from a pretrained-base to do continued pretraining on a specialized domain.
Goal of training: Learn general-purpose language modeling capabilities: syntax, knowledge, patterns, concepts, code, etc from raw text.
##
Supervised Finetuning (SFT)
Now you already have a pretrained model. It is great at next-token prediction, but it is not capable of generating useful answers to users' question. You need to teach the model basic behaviours.
What you need: A dataset of input → desired output examples. These outputs might be written by humans or generated by a stronger model.
Question:
Solve 2x + 4 = 10.
Desired response:
Subtract 4 from both sides...
2x = 6
x = 3.When you use SFT:Post-training, usually one of the first stages after pretraining.
Goal: Teach the base model what kind of responses we actually want it to produce. That can mean:
- following instructions,
- chatting,
- producing particular formats,
- learning first steps of how to format tool calls,
- behaving as an assistant.
Training is still essentially next-token prediction.
SFT does not train the model to recover from its own generated sequences. It is an "off-policy" technique meaning we tell the student model to follow the exact output token distribution of the dataset/external model. The student studies the teacher's exact answers step-by-step.
SFT can train new behaviors stably, but the student never learns how to recover when it makes its own mistake, because it only sees the teacher's perfect path
Direct Preference Optimization (DPO)
DPO changes the type of supervision. You show the model a paired dataset of good responses and bad responses. And train to align the model to prefer the good responses over the bad ones.What you need: A preference dataset containing something like
Prompt
│
├── "Sure I will help you with your issue" ← preferred
│
└── "I don't feel like helping today since I am busy" ← rejectedSo unlike SFT, you don't necessarily tell the model:
"This is exactly what you should say."Instead you tell it:
"Between these two responses, A is better than B."The preference can come from humans, AI judges, or another ranking mechanism.When you use it::Post-training/alignment, typically after SFT.Goal: Increase the probability of behavior people prefer while decreasing the probability of undesirable alternatives.
Reinforcement Learning with Verifiable Rewards (RLVR)
This is quite different. Instead of giving the model the correct response to imitate, you give it a task and some mechanism capable of checking whether it succeeded.
What you need: Task / Prompts + Verifier or Environment
Verifier environments can be created for tasks where the answer can be easily checked against the ground truth, such as math, coding, logical puzzles. The model generates its own solutions, receives rewards, and an RL algorithm such as GRPO/PPO/ updates the model.
When you use it::Post-training, frequently after some SFT/cold-start training.
Goal: Make the model discover strategies that result in successful outcomes, rather than merely imitate existing solutions.
This is different from SFT also because RL is on-policy. Which means that the learning examples comes from the responses generated by the learning model itself - not external sources.
Problem:
What is x...?
Verifier:
Does final_answer == 42?On Policy Distillation (OPD)
Now we get much closer to our actual topic.
In the more traditional "off-policy" distillation, teacher generates solutions --> save them as dataset --> student trains on them SFT style.
It is like SFT, but instead of a fixed dataset, the target outputs come from a separate teacher neural network.
In "on-policy distillation", we flip the order. The student generates solutions, and the teacher provides dense feedback about how the student did. We train the student on the teacher's feedback.
We call it on-policy since the training examples came from the student itself - not the teacher! The crucial part is that the student generates the trajectory.
What you need:
- a student model you are training
- a stronger teacher model
- a dataset of prompts/problems
Goal: Transfer the teacher's capability into the student specifically at the states the student itself encounters during generation.
In this article, we will build on the above intuition by discussing about on-policy vs off-policy (OPD vs SFT), dense feedback vs sparse feedback (OPD vs RL), as well as the excatly how OPD models are trained.
##
How OPD works, algorithmically!
To properly appreciate the power of On-Policy Distillation, we must understand it's advantages (and disadvantages) compared to SFT and RLVR.
In Supervised Finetuning, we give the student a correct answer written by a human or larger model. Every token becomes a training target, so the supervision is dense. However the answers come from the teacher's path, not the student's.
In Reinforcement Learning with Verifiable Rewards, we let the student generate several responses (GRPO) and assign a positive/negative advantage to each attempt. Here the supervision is sparse (we can't do token-level feedback) but the answers do come from the student's path.
In On Policy Distillation, we try to get dense token-level supervision (like SFT) and get on-policy rollouts i.e. we train on the student's path (like RL).
Let's understand these differences by seeing how OPD works
Step 1: Let's say we want to train a small student model to solve a mathematics problem, Given this math prompt, the student generates a full attempt. It may solve the problem, take a detour, or fail.
Step 2: We send the student's token sequence to the (more capable) teacher model.
Step 3: The teacher scores the log-probabilities of every token the student generated. This is often called the "teacher's distribution over the student's trajectory" (this terminology may sound confusing, but read that back a few more times until it makes sense!)
Step 4: For each token, find the difference between the teacher generated logprob and the student generated logprob.
Note: This specific variant of loss calculation is called "Sampled Token / Reverse KL" style OPD. There are other variants of losses which are discussed later on in the article.
Step 5: Using policy optimization, we update the student's weights so that its probability distribution moves closer to the teacher's distribution along that exact path.
Step 6: Repeat. The updated student produces new answers, the teacher scores them, and the training distribution moves with the student.
From the above algorithm, two things should be clear:
1. Why is OPD an on-policy method?
The training trajectories come from the current student policy. And these trajectories move as the student grows and improves.
2. Why does OPD give dense feedback to the student?
Since the teacher generates the log-probs for each token that the studen produces, we call it "dense" feedback. A "sparse" feedback, on the other hand, is when a verifier provides a single advantage/reward score for the entire sequence of token
On-Policy vs Off-Policy (OPD vs SFT)
Imagine you want to learn chess from a Grandmaster. The off-policy way is to just watch the grandmaster play a bunch of chess games and try to learn from how they play. The on-policy way is to play games yourself and the grandmaster give you feedback on how you played. In the first approach, you learn by imitating. In the second approach, you learn by correcting your own mistakes using the grandmaster's feedback on every move you made.
In simple terms, off-policy means we tell the student model to follow the exact output token distribution of the teacher.
On-policy means we let the student model generate its own sequences, and the teacher model simply tells it how good (or bad) those sequences are.
This is one of the main reasons why on-policy methods are so popular.
When AI labs deploy models out in public and let it run in the wild, it will generate entire sequences on its own, without any external supervision. With only off-policy training, the LLM generally develops blind spots in its language and might falter in scenarios that it did not see in training.
Off-policy methods can train new behaviors stably, but the student is primarily trained on externally provided trajectories rather than on the states produced by its own mistakes, because it only sees the teacher's perfect path.
On the other hand, on-policy methods allow LLMs to gain experience by generating their own trajectories, and learning from their mistakes through external teacher feedback. It makes the model more well rounded as the student learns how to fix its own errors during generation.
##
Dense vs Sparse Feedback (OPD vs RL)
Many popular RL training methods like REINFORCE, Proximal Policy Optimization (PPO), or Group Relative Policy Optimization (GRPO) are also on-policy, i.e. they work by generating tokens from the model and then using a verifier/environment to reward the final answer.
The main difference between RL and OPD is about the granularity of supervision:
Note that SFT also gives dense token-level feedback to the model since we have access to the target token sequences for every prompt via the external dataset (or an off-policy teacher).
This gives us a useful map:
##
The problem with On-Policy Training
Upon initial viewing, the concept of on-policy training might feel super attractive.
However, this comes with a huge caveat.
In on-policy training, the teacher only provides feedback on the tokens the student generated.
In off-policy training, the teacher can shift the student's trajectories to literally anything.
Meaning if the student never generated certain situations during training, it will not receive direct training signal on those states unless they are reached or introduced externally. The student model can only learn about trajectories it can reach. It does not learn about trajectories completely outside its reachable distribution / behavior space.
This is why when we are teaching a model completely new behaviours, we avoid starting directly with on-policy training. Some off-policy training as a warm-up is useful to show the student training examples that externally shape it's behavior, response format, or instruction-following capabilities.
Only when the model is already somewhat capable in the task, on-policy training (through OPD or RL) can further optimize its performance in the task.
Here is Dr. Omar Khattab explaining the big gotcha of OPD:
##
The problem with Teacher Supervision
The main advantage of RLVR over OPD is that RLVR relies on an undeniable source of truth (the verifier) as its source of feedback.
In OPD, if the teacher model has weird biases, they will shape the capabilities and behavior of the student model. Since the teacher model is ultimately a neural network, it will always have its own limitations and blind spots. We cannot assume it is an oracle.
An external verifier however can be trusted as an oracle. If you are training a neural network on math problems, the ground truth solution to each math question can be computed deterministically - so you can always judge whether the final answer the model came up with is correct or not. The reasoning trajectory the model followed to arrive at the answer is completely left up to the magic of exploration and optimization. You do not bias its intermediate tokens, you only judge the final outcome.
RLVR is typically slower and more expensive, but it can optimize against an external source of truth without inheriting a teacher model's exact preferences.
OPD is faster, cheaper, but has to rely on the capability of the teacher model.
In the next few sections, let's continue our exploration of OPD and the various design choices that researchers use to combat the problems mentioned above.
##
Sampled-token vs full-vocabulary training
So, there are (at least) two different ways the teacher can provide this token-level supervision. Let's understand these two ways using a simple example.
Suppose the student has generated the prefix, "The capital of France is," and then samples the token "Paris".
The first way is called Sampled-TokenTraining. Here we only inspect the single token that the student generated - we compare the probability the student assigned to "Paris" against the probability the teacher assigned to "Paris." That gives us one scalar training signal for this position.
This is cheap and scalable because we never need the teacher's full output distribution. But the con is that the signal is also pretty narrow. We only learn what the teacher thought about the token "Paris," but nothing about the alternatives the student could have selected.
The second way is called Full-VocabularyTraining. At the same prefix, we take the difference between the teacher's and student's probability assigned to every possible next token in the vocabulary. Then we calculate the divergence between those two distributions.
That is a much richer signal because the student can learn all the alternatives the teacher thought were possible, but it requires considerably more computation and memory at every generated position.
There is a third middle-ground too called Top-K distillation. As you might have guessed, instead of comparing one sampled token or the entire vocabulary, we instead compare the top-k likely candidate tokens according to the student generation.
##
Forward KL and Reverse KL (Mass Covering vs Mode Seeking)
Forward KL and Reverse KL are two ways to implement the distillation loss function. And the way this is implemented pretty much changes how the student model behaves.
Imagine the teacher knows two equally valid ways to solve the same math problem. For example, one path is algebraic and the other geometric. Think of the teacher distribution as two islands of good answers.
When is one useful vs when it backfires?
In theory, Forward KL makes the student put probabilities everywhere that the teacher assigns probabilities to. So the student can, in theory, learn both solutions to the math question we stated above, by learning to place mass on both the teacher's islands.
However, if the student model is too small and can't physically fit all that knowledge into its weights, Forward KL can severely backfire. This results in something called as "mass-covering" behavior.
Basically the student can just outright ignore the two islands, and decide to put a bunch of probability in the ocean between them. Trying to learn too much, but not able to actually learn any of the correct ways.
With Reverse KL, the student gets a lower loss if the tokens the student actually produced was favoured by the teacher. Reverse KL does not heavily penalize the student for failing to cover teacher modes on which the student itself places essentially no mass. It just checks if the student generated tokens had high probability according to the teacher.
The result of Reverse KL is "mode-seeking" behavior. Here the LM settles on learning just one solution to the math problem. Basically build a full city on one of the islands, and outright abandon the other one.
For a math or logical problem, learning just one solution is often enough to train a good model. So mode-seeking behavior actually helps producing a stronger student model.
However, for certain domains such as creative writing where world knowledge and diversity are super important, being mode-seeking often results in an inferior model.
##
On-policy Self-Distillation
Now here is where it gets really science fiction-y. Can we actually train a student model by using that same student model as the teacher as well?
The short answer is yes. We can use the same model to act as it's own teacher, but with one caveat.
When acting as the teacher, the model is given additional context, or privileged information, that the student itself does not receive.
For example, lets say we asked our student policy network to answer the question "What is 6 times 7?" And let's say that after generating a reasoning trajectory, the student arrived at the answer 41. A clearly wrong answer.
To turn that same model into a more capable teacher policy, we evaluate the student's trajectory again, but this time give the model some additional information. For example:
"A model answered 41 to this question before with the following reasoning: {$reasoning}. Carefully reconsider the reasoning and determine where it went wrong. Then regenerate the answer."
Just this switch in the context changes the LM's role and often this generates a better response that the student model can learn from using distillation.
The important part is that the trajectory still came from the unprivileged student. The privileged teacher now provides token-level supervision along that student-generated trajectory, and we distill this improved distribution back into the student!
Another common approach is to use a RL-like verifier-in-the-loop. Especially useful for verifiable tasks such as math problems. You can use an external calculator to check if the language model's original answer was correct or not. And then write the prompt as follows:
"A model answered 41 to this question with this reasoning: {$reasoning}. An external verifier confirmed that this answer was incorrect, and should have been 42 instead. Please regenerate your reasoning to the original question."
These new hints alters the behavior of the model. The student policy then learns to generate the same tokens without the hint. Having a verifier in the loop makes it easier for us to confidently check if the outputs are correct.These hints are called privileged information.Privileged Information is context that will not be available during deployment. We only show it during training so the student can generate better rollouts, as if it's a stronger model. The student network then uses those log-probabilities as if it came from a stronger teacher.
##
A quick note on "Privileged Illusion"
Privileged Illusion is a fairly new term, but the problem has always existed with Self-Distillation.
A teacher may generate better results not because it has learned a better capability, but because it can see information the student will never receive at inference time. Meaning that the teacher generates its response using leaked information that will never be available to the student at inference time, no matter how smart the student gets.
Example: The prompt asked the model "What is the current time?" and the privileged information used a clock and injected "It is 5:30 PM" into the prompt. This external information will never be available to the model, so there is no point training on it.
If we distil that distribution blindly, the student learns to imitate the privileged outcome without learning a reusable way to produce it. You can read more about PI in this recent DOPD paper, or this X article.
##
The gotchas of OPD
If you have been enjoying this article, consider checking out the 30 minute companion YouTube video too where I explain these concepts along with additional material in greater detail.
The winning features of On-Policy Distillation should be apparent by now, but let's close out this article by explaining some of its failure modes.
Thanks for reading! If you liked my writing style, here's where you can find more content on OPD:
My YouTube: https://www.youtube.com/@avb_fj
Thinking Machines Blog: https://thinkingmachines.ai/blog/on-policy-distillation/
MiniLLM Paper: https://arxiv.org/abs/2306.08543
GKD Paper: https://arxiv.org/abs/2306.13649
Self Distillation Paper: https://arxiv.org/pdf/2601.18734
A survey paper on recent OPD advancements: https://arxiv.org/pdf/2604.00626
A great GIT repo containing various papers and article links: https://github.com/chrisliu298/awesome-on-policy-distillation
Will Brown's article: https://x.com/willcb/status/2050038277454143918
