On-Policy Distillation

TL;DR

Thinking Machines’ Tinker-team write-up of on-policy distillation for LLM post-training: sample trajectories from the student, then use a larger teacher to score each token via per-token reverse KL against the teacher’s distribution. The recipe combines RL’s on-policy relevance (student learns in the states it actually visits) with distillation’s dense reward (every token graded, not just the final answer), at a fraction of RL’s compute. They replicate Qwen3’s reported result — Qwen3-8B-Base → 74.4% AIME’24 — and add a new use case: recovering instruction-following after mid-training catastrophically forgets it, by distilling from an earlier checkpoint of the same model on Tulu3 prompts.

Key claims

On-policy distillation = on-policy sampling (student rollouts) + dense reward (per-token reverse KL to teacher); SFT is off-policy + dense, RL is on-policy + sparse [§Best of both worlds, table].
Reverse KL is “unhackable” in the sense that low KL always corresponds to high probability of teacher-desirable behavior; it is mode-seeking (learns one teacher behavior rather than spreading across suboptimal options) and reduces exposure bias [§Loss function: reverse KL].
Implementation is a one-line modification on top of an RL training loop that already uses KL regularization — swap the regularizer model for the teacher; reward computation needs only a single teacher forward pass (logprobs) per trajectory [§Implementation, §Pseudocode].
On a Qwen3-8B-Base student initialized from 400k off-policy SFT on OpenThoughts-3 (60% AIME’24), on-policy distillation with a Qwen3-32B teacher reaches 70% AIME’24 in ~150 steps (~77K prompts × 4 samples) [§On-policy distillation, table].
Compute-efficiency: on-policy distillation is 9× cheaper than the SFT-2M extrapolated baseline when the SFT dataset is reused (teacher FLOPs amortized), and ~30× cheaper when the teacher must also be sampled fresh [§On-policy distillation, table; CE column].
Replicates Qwen3’s reported 74.4% AIME’24 / 63.3% GPQA-Diamond at ~1,800 GPU hours vs 17,920 GPU hours for RL on the same SFT initialization — i.e. ~10× cheaper than RL at higher quality [§Reinforcement learning, table from Qwen3 Technical Report Table 21].
Mid-training Qwen3-8B on internal company documents catastrophically degrades instruction-following: IF-eval drops from 85% (base Qwen3-8B) to 45% at 100% document data; even mixing 30% Tulu3 chat data only recovers to 79%, and LoRA does not save it [§Training on new knowledge degrades learned behavior].
On-policy distillation with an earlier version of the same model as teacher on Tulu3 prompts restores IF-eval to 83% after document mid-training, while preserving — and even slightly improving — internal-QA knowledge recall (36% → 41%). The student is using its own pre-mid-training checkpoint as a reward model [§On-policy distillation recovers post-training behavior, table].
The approach generalizes to continual learning: alternating mid-training (knowledge injection) with on-policy distillation against an earlier checkpoint (behavior recovery) provides a phased recipe that can keep small expert models up-to-date without losing post-trained behaviors [§On-policy distillation recovers post-training behavior].
Conceptual framing: any instruction-tuned open-weight model can be used as a reward model in this sense — compute_logprobs is the only API needed; this is a connection to DPO’s “language model is secretly a reward model” framing and to inverse RL [§On-policy distillation recovers post-training behavior].

Method

The student samples rollouts on a task prompt set, exactly as in RL. For each token in each rollout, the per-token advantage is set to the negative reverse KL between student and teacher distributions conditioned on the same prefix, i.e. $-\text{KL}(\pi_\theta(\cdot \mid x_{1..t}) \,\|\, \pi_\text{teacher}(\cdot \mid x_{1..t}))$ . In practice this is computed without materializing full distributions: the student’s logprobs come for free from the sampling step (already needed for importance sampling), and the teacher returns logprobs for the sampled tokens via one forward pass. The negative reverse KL is then fed as per-token advantage into a vanilla policy-gradient importance-sampling loss. Discount factor is set to zero — each step optimizes only the immediate next token. No separate reward model, no trajectory completion is required (training can use partial rollouts), and the teacher need not be queried for full distributions.

For the reasoning experiment, both teacher (Qwen3-32B) and student (Qwen3-8B-Base) are off-the-shelf Tinker-supported models, initialized from an SFT checkpoint on 400k OpenThoughts-3 prompts. For the internal-assistant experiment, the recipe is staged: (1) mid-train Qwen3-8B on a 70/30 mix of internal documents and Tulu3-sampled chat background data; (2) on-policy distill the mid-trained model on Tulu3 prompts with the original Qwen3-8B (the pre-mid-training checkpoint) as teacher.

Results

AIME’24 / GPQA-Diamond reasoning: SFT-400K baseline 60% / —; + on-policy distillation reaches 74.4% / 63.3% at ~1,800 GPU hours, vs RL’s 67.6% / 61.3% at 17,920 GPU hours on top of the same SFT init [§Reinforcement learning, table; Qwen3 Technical Report Table 21].
Compute-efficiency vs SFT extrapolation: on-policy distillation uses ~8.4 × 10¹⁹ teacher FLOPs and ~8.2 × 10¹⁹ student FLOPs to reach 70% AIME’24, vs an extrapolated 3.4 × 10²¹ / 1.5 × 10²¹ for SFT-2M at the same target — 9× cheaper with an existing SFT dataset, ~30× cheaper if the teacher must also generate fresh data [§On-policy distillation, table].
Internal-assistant recovery: Qwen3-8B IF-eval 85% / internal-QA 18% → mid-train (70% doc) 79% / 36% → + distill 83% / 41% — recovers almost full IF-eval while improving knowledge recall by 5 points [§On-policy distillation recovers post-training behavior, table].
100% document mid-training pushes internal-QA to 43% but craters IF-eval to 45%, and no mix ratio of background chat data preserves the original 85% IF-eval — distillation, not data mixing, is the lever that restores behavior [§Training on new knowledge degrades learned behavior, plot].

Why it’s interesting

Lands the dense-reward-but-on-policy combination as a third post-training axis alongside SFT (off-policy + dense) and reasoning RL (on-policy + sparse). The wiki’s Reasoning RL page has been accumulating evidence that GRPO + shaped rewards is the current default, but the open question of whether RL is the right primitive at all — explicitly raised in History May Repeat Itself: RSI Seen from a Previous AI Era — gets a concrete partial answer here: at ~10× lower compute than RL, on-policy distillation matches or beats RL on Qwen3-8B reasoning benchmarks. The continual-learning recipe (distill against an earlier checkpoint of the same model to recover lost behaviors) is the conceptually fresh contribution, and is closely related to Smaller, Weaker, Yet Better: Training LLM Reasoners via Compute-Optimal Sampling‘s W2S-I framing — both papers find that “the teacher” can be a weaker or earlier version of the student family without losing the educational signal. Sits between the wiki’s Diffusion Distillation cluster (which is about multi-step → few-step diffusion distillation, conceptually similar reverse-KL mechanics but a different problem) and Distillation Scaling Laws (which is LLM pretraining-time off-policy distillation scaling laws); the on-policy axis is currently uncovered by either.