On-Policy Distillation: A Cheaper Path to Specialized LLMs

Thinking Machines reports that on-policy distillation pushed a Qwen3-8B student to 74.4% on AIME'24 for roughly 1,800 GPU-hours, against the 17,920 GPU-hours the Qwen3 team spent reaching 67.6% with reinforcement learning. That's a higher score for about a tenth of the compute. If you're specializing an open-weight model and you framed the choice as supervised fine-tuning versus RL, this method changes the default answer. It borrows the mechanic that makes RL reliable, then pays for feedback the way SFT does.

Key takeaways

On-policy distillation samples trajectories from the student model, then uses a stronger teacher to score every token by reverse KL. It combines RL's on-policy relevance with the dense per-token signal of supervised fine-tuning.
On Qwen3-8B math reasoning, Thinking Machines reports 74.4% on AIME'24 at about 1,800 GPU-hours, versus 67.6% at 17,920 GPU-hours for the Qwen3 team's RL run. That's a higher score for roughly a tenth of the compute.
The dense signal is the reason it wins. Outcome-reward RL returns on the order of O(1) bits of feedback per episode, while per-token distillation returns on the order of O(N) bits, one for every token generated.
It repairs the regressions that domain fine-tuning causes. A model whose instruction-following fell from 85% to 79% during mid-training recovered to 83% while its internal-knowledge score climbed from 18% to 41%.
The teacher only runs a single forward pass to score the student's own tokens, so the expensive model never has to generate. That asymmetry is where the cost savings come from.
Reverse KL is mode-seeking, so the student converges toward the teacher and won't exceed it on the distilled skill. You need a teacher that already clears your quality bar.

What Is On-Policy Distillation, and How Does It Differ From SFT?

On-policy distillation trains a student model on trajectories it generates itself, and grades each token against a larger teacher's probability for that same token. The student does the talking; the teacher does the marking. Formally, the objective minimizes the reverse KL divergence between the student and teacher next-token distributions at every position the student actually visits.

Standard supervised fine-tuning, in this framing, is off-policy distillation. You collect a corpus of the teacher's completions, then train the student to imitate those transcripts token for token. The signal is dense, one target per token, but it describes states the teacher chose, not states the student will land in once deployed.

That distinction is the whole game. Off-policy training teaches the student to continue the teacher's sentences. On-policy training teaches the student to fix its own. The lineage here runs from sequence-level knowledge distillation to the generalized on-policy formulation, which established that sampling from the student and scoring it against a divergence measure beats imitating a fixed dataset.

Why Does Sampling From the Student Matter So Much?

Because a model in production sees the states its own earlier tokens produced, not the states a teacher would have produced. Train only on teacher transcripts and the student learns to handle inputs it will rarely encounter, while the mistakes it does make compound over a long generation with no correction anywhere in the data.

This is the classic distribution-shift problem. On a short answer it barely registers; on a 4,000-token chain of reasoning, one early misstep drifts the student into territory the SFT corpus never covered, and the errors accumulate. On-policy distillation closes the gap by scoring exactly the tokens the student emits, so every correction lands on a state the model genuinely reaches.

There's a second, quieter reason, and it's about information density. Outcome-reward RL tells the model whether a whole trajectory earned a reward, roughly O(1) bits per episode, with no indication of which step went wrong. Per-token distillation returns a graded signal at every position, on the order of O(N) bits for an N-token rollout. Denser feedback means far fewer rollouts to learn the same behavior.

How Much Compute Does On-Policy Distillation Actually Save?

Roughly an order of magnitude against RL, on the published Qwen3 math-reasoning numbers. The table below sets the three post-training routes side by side, using AIME'24 accuracy for a Qwen3-8B student, drawn from the Qwen3 technical report and Thinking Machines' analysis of it.

Method	Policy	Feedback signal	AIME'24 (Qwen3-8B)	Compute
Off-policy distillation (SFT on teacher transcripts)	Off-policy	Dense, but on the teacher's states	55.0%	Large labeled corpus
Reinforcement learning (outcome reward)	On-policy	Sparse, ~O(1) bits per episode	67.6%	17,920 GPU-hours
On-policy distillation	On-policy	Dense, ~O(N) bits per token	74.4%	~1,800 GPU-hours

The cost asymmetry has a concrete cause. Generating rollouts is the expensive part of any on-policy method, and here the small, cheap student does that generation. The teacher only computes log-probabilities for tokens that already exist, which is a single forward pass with no sampling loop. You pay a large model to grade, not to talk.

Thinking Machines pushes the point further with a data-efficiency result: repeatedly training on a single prompt, 256 rollouts across 20 steps, reached teacher-level AIME performance. Reverse KL learns to approximate the teacher's full next-token distribution rather than memorize one answer, so even a handful of prompts carries a lot of signal.

Can It Recover Skills a Model Loses During Domain Fine-Tuning?

Yes, and that may be its most practical use for teams shipping a private assistant. Fine-tuning an already-aligned model on your own documents reliably injects knowledge and quietly damages behavior the base post-training installed, chat quality and instruction-following most of all.

Thinking Machines' personalization run makes the trade-off legible. Starting from Qwen3-8B, mid-training on internal company documents lifted an internal-knowledge score from 18% to 36%, but instruction-following on IF-eval slid from 85% to 79%. The domain knowledge arrived; the manners degraded.

A short on-policy distillation pass, using the original instruction-tuned model as the teacher, restored IF-eval to 83% and carried the knowledge score up to 41% rather than trading it away. The student practiced on its own outputs while the teacher reminded it how to follow instructions, so the two objectives stopped competing. Treat distillation less as a one-shot specialization step and more as a recovery phase that runs after any mid-training that risks regression.

Where Does On-Policy Distillation Fit Against RL and Prompt Optimization?

Reach for it when a teacher already exists that clears your quality bar on the target skill and you want a smaller or cheaper model to match it in the states it will actually see. That covers most enterprise specialization: a frontier or larger open model can do the task, and you need it to run at a fraction of the serving cost.

Reinforcement learning stays necessary when no teacher can demonstrate the behavior, when you're pushing past the frontier rather than compressing it, or when the objective is a verifiable outcome with no reference distribution to imitate. On-policy distillation is a compression and transfer tool, not a way to discover capabilities that don't exist yet. This is the same fork covered in when reinforcement fine-tuning an agent beats prompting, now with a third, cheaper option on the table.

If you can't train at all, prompt optimization is the lighter lever, and methods like the reflective evolution described in GEPA versus reinforcement learning get you part of the way without touching weights. Distillation is the step up when a prompt can't close the gap and you'll run a training loop. Because it produces a compact specialist, it pairs naturally with the case for small language models in agentic systems, and the student is usually trained with the same LoRA or QLoRA setup you'd use for any parameter-efficient fine-tune.

What Do You Need in Place to Run It?

Three things: a capable teacher, the infrastructure to run its forward pass alongside student generation, and a base model that already gets the task right some of the time. The last one matters more than it sounds. Reverse KL sharpens a distribution the student can already sample from, so if the base model never produces a correct trajectory, there's nothing for the teacher to reinforce.

Operationally, you run a student sampler and a teacher scorer in the same loop. The student generates rollouts, the teacher scores their tokens in one forward pass, and a policy-gradient update follows with the per-token reverse KL as the advantage. Open-source RL trainers now ship on-policy distillation as a first-class recipe, so you're assembling components rather than writing the loss from scratch.

Teacher access is the real gate. Self-distillation, where a model teaches a smaller sibling from the same family, is clean because you control both endpoints and the tokenizers match. Distilling across model families is harder: you need shared or mappable vocabularies to compare log-probabilities token for token, and closed API models generally won't return the full per-token distribution the method depends on.

What Are the Trade-Offs and Failure Modes to Watch?

The hard ceiling is the teacher. Reverse KL is mode-seeking, so the student concentrates on the teacher's high-probability behavior and converges toward it, which means it won't exceed the teacher on the distilled skill. Distillation makes a cheaper copy; it doesn't make a better model than the one you're copying from.

Mode-seeking carries a subtler cost: it narrows output diversity. The student learns the teacher's dominant modes and sheds the long tail, which is fine for a math solver and a real consideration for anything creative or exploratory, where you may want the variety a forward KL or a lighter touch would preserve.

Two operational cautions round it out. First, supervised fine-tuning on a finite batch of the model's own samples drifts off-policy as training proceeds, which is exactly the degradation on-policy distillation exists to avoid, so don't approximate it with a static self-generated dataset and expect the same result. Second, the method assumes your evaluation measures the skill you're transferring. A dense per-token signal faithfully reproduces a teacher's blind spots, so a teacher that's confidently wrong in a corner of the distribution passes that failure straight to the student.