On-Policy Distillation by Thinking Machines Lab

***Non-members can read for free***

Kevin Lu in collaboration with others at Thinking Machines introduced on-policy distillation, a method that efficiently transfers knowledge and behavioral alignment within large language models, enabling continual learning, personalization, and performance recovery without the cost of reinforcement learning.

As large language models (LLMs) grow more capable, training them efficiently has become a balancing act between compute, data, and human alignment. Traditional knowledge distillation compresses large models but struggles with adaptation. Reinforcement learning from human feedback (RLHF) helps models align with complex goals, but it’s costly and unstable.

Source Image

Between these extremes lies a method that merges efficiency with adaptability: On-Policy Distillation (OPD), a framework where a model learns not from frozen data, but from its own behavior, refined by a teacher model in real time.

Why On-Policy Distillation Matters

Classical distillation is off-policy; the student learns from teacher data generated long before its own updates. Over time, this causes distributional mismatch: the student is evaluated on behavior it no longer follows.

On-policy distillation eliminates this gap. Here, the student actively generates samples using its current policy (π_θ), the teacher provides token-wise feedback, and updates happen immediately. This makes learning both dense (every token supervised) and adaptive (reflecting current behavior).

Source Image

For example, to solve this math problem:

Source Image

On-policy training through reinforcement learning (RL) allows a student model to learn by grading its output, or “rollout,” on its correctness, either by a human or a teacher model . While training on its own samples helps the student directly avoid mistakes, RL’s sparse feedback teaches only whether the final answer is right or wrong, not the specific location of the error, making it inefficient for some applications.

Source Image

Off-policy training utilizes supervised fine-tuning (SFT) with labeled examples, often from a proficient teacher model. Distillation trains a student model to emulate the teacher’s outputs by training on the teacher’s full sequence of generated tokens.

Source Image

The drawbacks of off-policy training include the student learning from the teacher’s frequent contexts, leading to compounding errors when the student deviates, especially in long sequences. Additionally, off-policy distillation can result in the student imitating the teacher’s style and confidence without learning its factual accuracy.

On-policy distillation for learning chess provides fine-grained feedback on a player’s moves, unlike standard reinforcement learning (playing without a coach). This method gives delayed, holistic feedback or off-policy distillation that involves observing grandmaster moves in scenarios a novice won’t face. It further combines on-policy relevance with a dense reward signal, akin to a teacher grading each individual move.

Source Image | blunders (red), mistakes (orange), inaccuracies (yellow), or brilliant (blue)

Best of both worlds

On-policy distillation represents an important step toward refining how large language models (LLMs) learn and reason from their own behaviors. Unlike traditional distillation techniques, which rely on static datasets or teacher-student imitation, on-policy distillation introduces a dynamic, self-referential learning process. It allows the model to iteratively evolve its reasoning under its current policy, ensuring that each new generation of knowledge is shaped by what the model presently understands not what it once learned.

On-policy distillation samples a student model’s output trajectory and uses a high-performing teacher to provide dense, token-by-token feedback. This approach is a powerful and computationally efficient method for post-training, combining the interactive learning of on-policy training with the detailed guidance of a teacher model.

Source Image

Using an approach similar to DAGGER and process reward modeling, on-policy distillation method replicates the Qwen3 team’s findings. The work demonstrates that equivalent reasoning performance can be achieved at a significantly lower cost than traditional reinforcement learning.

The Core Mechanism: Reverse KL Optimization

At the mathematical heart of on-policy distillation lies the Reverse Kullback-Leibler (KL) divergence, which measures how the student’s probability distribution diverges from the teacher’s but under the student’s own samples.

Source

Using reverse Kullback-Leibler (KL) divergence as a reward function trains a student model to emulate a teacher’s behavior by focusing on mode-seeking behavior. This approach learns to prefer outputs that the teacher also finds probable. A low KL value directly indicates desirable, teacher-like behavior and avoids generating lower-quality alternatives.

An example from SimpleBench shows a student model, Qwen3–4B-Instruct-2507, incorrectly solving a problem mathematically without considering that ice melts in a frying pan. A teacher model, Qwen3–235B-A22B-Instruct-2507, penalized the tokens that led the student model astray but not the final incorrect answer.

Source Image

Simplified On-Policy Distillation Pseudocode

An on-policy distillation method based on the Tinker platform trains a smaller student model to imitate a larger teacher model using reinforcement learning (RL). The process involves:

Initializing a teacher client

Sampling trajectories from the student

Computing the reward based on reverse KL divergence

Training the student model with an RL importance-sampling loss function

This pseudocode represents the training loop: student rollouts, teacher evaluation, reverse KL loss, and gradient update forming a self-correcting cycle.

The reverse-KL objective acts like a continuous distillation engine, compressing reasoning and behavioral signals without external reward models.

Source Image

Experimental Evaluation: The Qwen3 Study

Using Qwen3–32B as a teacher, mathematical reasoning was distilled into the Qwen3–8B-Base model, and this experiment can be reproduced with the Tinker cookbook.

Off-policy distillation

In an experiment using off-policy distillation, a Qwen3–8B-Base model was supervised fine-tuned on 400k prompts from the OpenThoughts-3 dataset. The training achieved a 60% score on the AIME’24 math benchmark, demonstrating that performance increases log-linearly with continued training.

Source Image

Improving a fine-tuned model’s AIME’24 score from 60% to 70% is exceptionally difficult, as performance gains become logarithmically more expensive, requiring significant effort from post-training techniques like PPO, DPO, or online reinforcement learning. Fine-tuning an 8B model with approximately 2 million off-policy distillation prompts is estimated to reach 70% performance on AIME’24, though this relies on the scaling law holding.

Reinforcement learning

In large-scale model training, this idea marks a fundamental shift. The concept bridges reinforcement learning (RL) and supervised fine-tuning(SFT), turning the learning loop itself into a continuous feedback mechanism. Instead of relying on fixed data or past teacher signals, the model interacts with its environment (or task), gathers new samples aligned with its own reasoning distribution, and then distills this experience into a refined version of itself. This is learning from one’s own updated thought process: a mechanism much closer to human reflection.

The Qwen3 technical report highlights a performance of 67.6% on a benchmark after 17,920 GPU hours of RL, noting that this is computationally comparable to training on 2M off-policy distillation prompts, depending on the SFT training assumptions.

Source Image

Qwen team achieved a high score of 74.4 on the AIME’24 benchmark using a cost-effective method called on-policy distillation.

On-policy distillation

Unlike standard distillation, which depends on off-policy data (teacher-generated samples that the student later imitates), on-policy distillation keeps training synchronized with the model’s current behavior. The student generates responses using its own policy, and the teacher model evaluates these outputs at every token. This allows training signals to stay relevant throughout the learning process rather than drifting away from the student’s distribution.

On-policy distillation trains a student model on sequences it generates dynamically, with a teacher providing supervision on those self-generated samples. This rapidly improves model performance, achieving a 70% AIME’24 score from an SFT checkpoint.

Source Image

On-policy distillation is a significantly more compute-efficient fine-tuning method than SFT, particularly for LoRA models, reducing the performance gap between LoRA and full fine-tuning from 13% to 6% at rank 32.

When comparing the computational costs of different methods, the ratio of costs for training, sampling, and calculating log-probabilities can vary drastically depending on the specific implementation. Using FLOPs (floating-point operations) as a cost metric has limitations, particularly as it can significantly overestimate the practical cost of operations that benefit from parallel processing, such as calculating log-probabilities on GPUs.

Source Image

SFT yields a 9x cost reduction, which increases to 18x in GPU hours because teacher model calculations can be parallelized. When off-policy training is used, the cost savings are more significant because the “teacher” model’s costs aren’t included. The total cost reduction grows to 30x when accounting for all teacher model costs in off-policy distillation, particularly for new tasks without a pre-existing data.

Distillation for Personalization

Beyond improving general performance, distillation can tailor models to specific needs such as conversational tone, output formatting, or cost-aware reasoning. However, teaching both new knowledge and behavioral nuance simultaneously has always been difficult.

Lightweight finetunes often fail to meet this dual goal, forcing practitioners to rely on complex post-training stacks with proprietary reward models; a luxury only frontier labs can afford.

This is where on-policy distillation stands out. By allowing models to learn from their own generated data, it aligns domain expertise with desirable behaviors. The same mechanism extends naturally to continual learning where models update after deployment without losing prior skills.

Training an internal assistant

Consider the goal of building an internal company assistant, a model that combines deep understanding of internal documents with assistant-like responsiveness. Achieving both simultaneously is challenging because domain learning often disrupts previously learned patterns.

The two objectives are:

Knowledge recall:

Correctly answering questions using internal company data (Internal QA)

Instruction following:

Maintaining conversational competence (IF-eval)

Training on new knowledge degrades learned behavior

The researchers begin with Qwen3–8B, already fine-tuned for reasoning and instruction-following through reinforcement learning. Prior work shows RL activates only small subnetworks within the model; a fragile structure easily disturbed by later finetuning.

To reduce catastrophic forgetting during mid-training, one might usually mix in “background data” from the pretraining distribution. When that’s unavailable, an alternative is to sample on-policy background data.

Here, prompts from Tulu3, a large instruction dataset, are resampled using Qwen3–8B itself. This on-policy background acts as a forward KL regularizer, anchoring the model’s behavioral style. Using the smaller Qwen3–8B for sampling preserves chat capabilities better than the larger Qwen3–32B.

Source Image

When they fine-tune Qwen3–8B on different mixes of internal document data and chat data, increasing document proportion improves knowledge but harms IF-eval. Even with 30% chat data, behavior still deteriorates.

Source Image

Neither longer training nor LoRA (Low-Rank Adaptation) fully resolves this issue. LoRA limits catastrophic forgetting but learns less and cannot preserve IF-eval scores.

Source Image

On-policy distillation recovers post-training behavior

To restore instruction-following after finetuning on internal data, they use on-policy distillation, training with an earlier version of Qwen3–8B as the teacher.

This approach is highly effective: after a 70–30 finetune (documents–chat data), on-policy distillation recovers almost full IF-eval performance without losing knowledge. It also yields positive transfer; improved chat skills slightly enhance internal QA scores.

Source Image

In essence, the model acts as its own reward model, rewarding high-probability behaviors akin to inverse reinforcement learning. This allows any instruction-tuned model to self-distill if it provides a function like compute_logprobs.

Discussion

Dense Supervision Greatly Improves Compute Efficiency

Reinforcement learning and on-policy distillation both optimize behavior using reverse Kullback-Leibler (KL) divergence, which prunes undesirable actions from the base policy. The distinction lies in reward density. RL offers sparse rewards: roughly O(1) bits per episode while on-policy distillation provides O(N) bits per episode, where N equals the number of generated tokens. This dense signal drastically enhances learning efficiency.

Experimental Setup:

Begin with Qwen3–8B-Base, an unmodified model.

Train it with RL on DeepMath, matching the earlier LoRA Without Regret setup using a LoRA rank of 128.

Use this RL-trained model as the teacher, and perform on-policy distillation back into the base model.

Source Image

The results are striking. Starting from the same initialization, on-policy distillation replicated the RL-trained behavior 7–10× faster, requiring 50–100× less compute overall. The reverse KL loss converged to near zero, and the AIME’24 score matched the teacher’s in under 10 gradient steps, compared to 70 in RL.

Several insights support this efficiency gap.

1. First, while RL must train at full evaluation context lengths to avoid format penalties, distillation remains stable even at shorter contexts, as it lacks sharp reward cutoffs.

2. Second, when initialized from strong SFT checkpoints, on-policy distillation functions effectively with smaller batch sizes, as the dense reward signal reduces gradient noise.

Despite RL’s difficulty in scaling with process supervision, distillation demonstrates that dense token-level feedback can improve training efficiency by an order of magnitude.

Distillation can effectively reuse training data for data efficiency

Collecting diverse prompts for training can be expensive and slow. Thus, a practical advantage of on-policy distillation is its ability to reuse prompts across multiple epochs without overfitting.

Unlike RL, which tends to memorize final answers when re-exposed to the same prompt, on-policy distillation learns the full output distribution of the teacher via reverse KL, enabling learning across repeated samples of identical data.

Source Image

Using Qwen3–8B-Base, the authors trained on a single randomly chosen prompt from DeepMath for 20 steps, each with 256 rollouts (5120 total sequences). Despite repeated exposure to the same example, the model nearly matched the teacher’s AIME’24 performance.

This shows that multi-epoch distillation on even one example can effectively transfer the teacher’s policy. Their default setup, also used for personalization uses 64 prompts per batch with 4 samples per prompt, but even the single-prompt setup sustained strong generalization.

This property highlights on-policy distillation’s data efficiency, showing that models can retain quality even when training data is limited; a vital capability for real-world fine-tuning where large, diverse datasets are unavailable.

RL Searches in the Space of Semantic Strategies

A major takeaway from these results is that reinforcement learning and distillation serve different roles. While distillation compresses and transfers learned reasoning, RL acts as a search mechanism across the landscape of semantic strategies.

In this view, pre-training explores a high-dimensional parameter space through stochastic gradient descent, a costly and exhaustive process. RL, in contrast, explores semantic strategy space, iteratively refining behaviors by trial and error rather than weight updates. Each rollout in RL represents a hypothesis test of behavior, not a gradient descent step.

Once RL discovers a successful strategy, on-policy distillation efficiently encodes it, without replicating the costly exploration process. This distinction explains why distillation converges so much faster. It focuses solely on imitating final strategies, not retracing all intermediate ones.

Researchers may spend years exploring hypotheses before reaching a result, yet once found, teaching that concept to others is far easier. Distillation, then, becomes the language of transmission, efficiently conveying high-level knowledge without redoing the original exploration.

On-Policy Learning as a Tool for Continual Learning

Beyond efficiency, the research team highlights on-policy distillation as a framework for continual learning allowing a model to integrate new behaviors without forgetting old ones.

Previous work observed that on-policy RL forgets less than off-policy training, but RL’s sparse supervision limits its ability to absorb new knowledge. Conversely, SFT and off-policy distillation may teach new information but tend to degrade behavioral quality.

Source Image

To investigate, the authors constructed a dataset using Tulu3 prompts sampled from Qwen3–32B at temperature 1.0, ensuring a KL of zero between data and model. Running SFT on these self-generated samples surprisingly led to performance degradation on instruction-following tasks even with practical learning rates.

The explanation lies in finite-batch divergence: although expected KL is zero, each batch slightly differs, creating small but accumulating drifts in policy that make the model effectively go off-policy over time. By contrast, on-policy distillation remains anchored to a fixed teacher, ensuring stable convergence and preventing this drift as an essential property for continual or self-improving models.

Conclusion

This investigation reveals that on-policy distillation blends the behavioral reliability of RL with the efficiency and scalability of dense supervision. Across experiments from mathematical reasoning to personalization, the technique consistently delivers frontier-level performance at a fraction of RL’s compute cost.

By leveraging student-generated samples and teacher-guided dense feedback, on-policy distillation emerges as a practical post-training recipe for modern LLMs. Its efficiency, adaptability, and continual learning capacity position it as a key enabler of scalable model alignment.

At Thinking Machines, the vision is clear: to create AI models that combine frontier performance with adaptability and personalization. On-policy distillation stands as a crucial step toward that future: a practical, elegant path to keeping the intelligent systems both current and aligned.

References:

[1] Lu, Kevin and Thinking Machines Lab, “On-Policy Distillation”, Thinking Machines Lab: Connectionism, Oct 2025.