On-Policy Distillation
Adventures

Matching the complete KL divergence for LLMs requires summing over the entire vocabulary at every token position, which is computationally expensive when the vocabulary is large. Instead, previous works estimate the KL from the single token actually sampled at each position24. This reduces the problem to a token-level RL-like objective, where the advantage is the per-token log-probability difference between teacher and student:

Here $A_t$ is the per-token advantage: it measures how much more likely the teacher considers token $y_t$ compared to the student, serving as a dense signal that drives the student toward the teacher's distribution1. We keep all updates strictly on-policy to keep this estimator unbiased.

Throughout the rest of this post, we use the token-level ratio \( r_t := \frac{\pi_\theta(y_t \mid y_{< t}, x)}{\pi_{\rm T}(y_t \mid y_{< t}, x)} \). Under this convention, the reverse-KL reward used in Algorithm 1 is \( A_t = \log \pi_{\rm T}(y_t \mid y_{< t}, x) - \log \pi_\theta(y_t \mid y_{< t}, x) = -\log r_t \). We call \(A_t\) either the advantage or the policy-gradient weight in the remainder of this post.

Naive Pseudo Code

We implement the naive baseline using the k1 estimator16, the single-sample estimator of reverse KL, as the per-token advantage. The pseudo code is shown in Algorithm 1.

for batch in train_loader:
    y, logp = student.rollout(batch)
    logp_t = teacher.log_prob(batch, y)

    A = logp_t - logp                   # per-token k1
    loss = -(A.detach() * logp).mean()  # standard token-level REINFORCE

    student.update(loss)

Experiment Settings

For this naive baseline, we follow Algorithm 1 directly. Unless otherwise stated, all runs use Qwen/Qwen3-4B-Base as the student and Qwen/Qwen3-4B as the teacher, and we train for one epoch on the DeepMath-103K prompts. For evaluation, we use both mean@8 and best@8 as metrics, averaged over: BeyondAIME, AIME25, AIME24, HMMT25, BRUMO25. All other hyperparameters use the defaults listed in the Hyperparameter Analysis section. We implemented our experiments using the veRL framework.

Naive Results

We plot the Train KL results, which in the case of OPD correspond to the training cost/negative reward (left), and the aggregate eval performance (right) below. We can immediately see that there is a stability issue with the naive implementation: the KL estimator drops sharply after around 120 training steps (even going negative), which is also directly accompanied by a collapse in eval performance.

Given the instability observed in the previous section, we first establish a stable baseline before comparing different design choices. Without this, apparent improvements could simply reflect training instability rather than genuine gains from the proposed method or settings.

The Instability Problem

We investigated two angles to try to stabilize the runs:

• Training-inference mismatch correction313: the rollout engine and learner backend can assign different token probabilities even when they share the same parameters (due to differences in CUDA kernels, batching strategies, reduce ordering, etc.), making the updates effectively off-policy. We test two importance-sampling correction techniques to mitigate this mismatch:
  • Truncated-IS (TIS): clamp the per-token importance ratio of rollout and learner to a fixed range $[\epsilon_{\rm lo}, \epsilon_{\rm hi}]$. Tokens whose ratio falls outside the range still contribute to the loss, but with a bounded importance weight.
  • Masked-IS (MIS) or IcePop12: rather than clamping, completely mask out any token whose importance ratio falls outside the allowed range. Masked tokens contribute no gradient at all, preventing extreme-ratio tokens from corrupting the update.
• Adjusting Learning rate: high learning rates can destabilize gradient descent through “overshooting.”25 Our naive baseline used a learning rate of 1e-5, following the Thinking Machines blog, which is substantially higher than common RLVR defaults (around 1e-6). We therefore conducted an ablation run with learning rate=1e-6.

Unless otherwise stated, we use the same rollout-ratio thresholds for both methods: \(\epsilon_{\rm lo}=0.5\) and \(\epsilon_{\rm hi}=2.0\). For TIS this means clamping to \([0.5, 2.0]\), while for MIS this means masking tokens outside \([0.5, 2.0]\).

After stabilizing the OPD experiments in the previous section with MIS rollout correction, we now take a closer look at how some key hyperparameters could affect the eval performance from OPD training.

Hyperparameters

We refer to default as the baseline setting used for all runs throughout this blog: the Masked-IS configuration from §3 with the hyperparameters listed below (click to expand).

Hyperparameter Sweeps

We select key hyperparameters for OPD and sweep them across values to compare them against the default run. Note that moving a second slider automatically resets the others back to default hyperparameters. We leave the effects of hyperparameter combinations to future work.

Naive k1 vs k3

We start with a small pilot run: keep the setup fixed and change the per-token advantage form, from k1 (default) to k316, which is a lower variance estimator for the reverse KL. The two plots below are the train KL and average mean@8 of eval benchmarks over training. For k3, we observe a slower increase in eval performance and unstable KL which follows by an eventual collapse.

Derivation of Divergence Advantages

Using the raw divergence value directly as a REINFORCE weight drops the reward-gradient correction and gives a biased gradient.

Compact derivation. We start from an f-divergence objective. Here \(D_f\) is generated by a convex function \(f:\mathbb{R}_{+}\!\to\!\mathbb{R}\) with \(f(1)=0\). The default reverse KL used in the main experiments is the special case \(f(r)=r\log r\): \[ D_f(\pi_\theta \,\|\, \pi_{\rm T}) = \mathbb{E}_{y\sim \pi_{\rm T}}\!\left[f(r_\theta(y))\right], \qquad r_\theta(y):=\frac{\pi_\theta(y)}{\pi_{\rm T}(y)}. \] Since the teacher policy \(\pi_{\rm T}\) is fixed with respect to \(\theta\), only \(\pi_\theta\) contributes to the derivative inside \(r_\theta(y)\). The key identity is \[ \nabla_\theta r_\theta(y) = \frac{\nabla_\theta \pi_\theta(y)}{\pi_{\rm T}(y)} = r_\theta(y)\,\nabla_\theta \log \pi_\theta(y). \] Plugging this in gives the gradient:

This shows that the correct policy-gradient weight is \(f'(r_\theta(y))\). The full derivation below reproduces the same result step by step and shows how it relates to the alternative form \(h(r_\theta(y))+r_\theta(y)h'(r_\theta(y))\), where \(h(r):=f(r)/r\) (click to expand).

Divergence Types

With the properly derived advantage for general f-divergences, we examine three representative divergences (i.e., Jensen, Hellinger, Tsallis) and illustrate their corresponding advantage form, \( -f'(r) \), below. We also plot the distribution of importance ratios \( r_\theta \) observed in actual responses before any training (Qwen3-4B-Base vs. Qwen3-4B), which can reveal which weighting regimes dominate the early training dynamics.

Divergence Shaping Results

Using the derived policy-gradient weight -f'(r), we compare the average mean@8 of multiple f-divergences (Jensen, Hellinger, Tsallis) against the default (k1) in the bar chart below.

Instead of a single teacher policy, we can define $\pi_{\rm T}$ as an ensemble or a mixture of teachers. For example, we can have multiple teacher models of different sizes or from different model families, and combine their outputs to define the reference distribution. This raises the question: how does the choice of aggregation strategy affect eval performance? In this section, we use Qwen3-4B,Qwen3-8B,Qwen3-14B,Qwen3-32B as ensemble of teachers, and leave teachers of different model families for future work.

Ensemble Strategies

We explore several strategies for aggregating the outputs of teacher models:

• Arithmetic mean: $\pi_{\rm T}=\frac{1}{K}\sum_k \pi_k$ (probability-space averaging).
• Geometric mean: $\pi_{\rm T}\propto \prod_k \pi_k^{w_k}$ (logit-space averaging).
• Hard switch: pick a single teacher per query sample. As training progresses, we switch to the next larger teacher, similar to a form of curriculum learning.
• Soft switch: the reference distribution is a linearly weighted combination of “neighboring” teachers. We gradually shift the weights from smaller to larger teachers over time.
• Control-token routing: condition the teacher choice on a control prefix token before generation. We split each batch evenly across teachers (e.g., 25% each for Qwen3-4B, Qwen3-8B, Qwen3-14B, and Qwen3-32B).
• Random: pick a teacher at random for each query sample.

The animation below shows how each ensemble strategy computes the policy-gradient advantage token by token. The student first generates an on-policy response on the left. Some methods depend on training progress, represented as a ratio in [0,1], where 0 indicates the start of training and 1 indicates its completion (illustrated in the left panel). In the middle, the chosen teachers evaluate the sampled token, where $W$ denotes the weights used to linearly combine the teacher advantages. Finally, the aggregation method combines these evaluations into a single reference score on the right.

Switch between aggregation methods using the tabs below to see how the routing and weighting change. Press the replay button to start the animation from the beginning, or step button to go through it token by token to see the intermediate values.

Ensemble Results

We train each ensemble strategy with other hyperparameters fixed to the same default (no ensemble) setup and compare eval performance curves over training.

Recent work on RLVR19 shows that high-entropy "forking tokens," roughly 20% of tokens in a chain of thought, drive most of the learning signal. Motivated by this finding, we explore whether this insight transfers to on-policy distillation: can we improve training by selectively filtering which tokens receive gradient updates?

Filtering Criteria

We test three token-level scoring functions to rank tokens, then only apply gradient updates to tokens within a chosen percentile range:

• KL divergence: focus on tokens where the student-teacher disagreement is smallest/largest.
• Student entropy: focus on tokens where the student is most certain/uncertain, analogous to the "forking tokens" identified in RLVR.
• Teacher entropy: focus on tokens where the teacher itself is certain/uncertain.

Formally, let $s_t$ be the token-level score under the chosen criterion and $m_t \in \{0,1\}$ indicate whether token $t$ falls inside the selected percentile range. The filtered advantage is:

This concentrates the gradient budget on selected tokens based on heuristic rules. We sweep the percentile ranges (0-30, 30-70, 70-100) for each criterion. The interactive visualization below shows which tokens are selected under each criterion and range for a sample sequence. Move the P_MIN and P_MAX sliders to see how the tokens are filtered, and choose the filtering criterion at the top left to explore different strategies.

The table below summarizes all token-filtering runs, showing both peak and end-of-training outcomes so we can compare raw gains versus training stability.

Conclusion

Through empirical exploration of on-policy distillation, we come with several key take-aways.

• Stability is the first bottleneck. The naive baseline can run into stability issues and make the k1 estimator values negative. Masked importance sampling (MIS / IcePop) resolves this and is a prerequisite for all subsequent experiments.
• Hyperparameters are high-leverage. Max response length (higher the better), rollout temperature (best at 1.0), and number of epochs (more the better) highly affect eval performance.
• Divergence choice has reward-shaping effects. Replacing KL with alternative f-divergences (Hellinger, Tsallis) produces gains (+0.3% mean@8), suggesting that the shape of the reward function does influence optimization.
• Curriculum-based teacher scheduling is the biggest win. Hard switch between a weaker and stronger teacher yields +1.3% mean@8 and +3.6% best@8 over default by end of training.
• Token filtering improves training stability. Concentrating gradient updates on high-KL tokens (top 30%) improves end-of-training mean@8 by +1.8% and achieves the best best@8 of 37.50% (+1.7% over default).

Limitations

• Domain scope is narrow. The experiments focus on DeepMath-103K prompts and math evaluations, so results may not transfer directly to coding, tool use, or broader instruction-following tasks.
• Model-family/size scope is limited. Most comparisons use Qwen3-4B teacher-student variants, so cross-family and student size behavior remains untested.
• Statistical confidence is limited. The runs are primarily exploratory comparisons and do not yet include multi-seed confidence intervals.

Future Work

• SFT-then-OPD: run a supervised fine-tuning warmup to bring the student closer to the teacher policy. This may be particularly beneficial for the control-token ensemble method, since the control-token prefix introduces distribution shift.
• More diverse settings: we only investigated under mathematical reasoning and need to expand to coding, tool-use, and agentic scenarios.
• Self-context distillation: recent work on self-distillation20212223 shows that a single model can act as both teacher and student by conditioning on privileged information (verified reasoning traces, environment feedback, or demonstrations). On-policy distillation is a key step in these frameworks. Investigating how conditioning on privileged context interacts with the techniques explored in this blog (divergence shaping, token filtering, curriculum scheduling) is a promising direction for future work.

@misc{zhao2026opd,
  title = {On-Policy Distillation Adventures},
  author = {Zhao, Andrew},
  journal = {Andrew Zhao's Blogs},
  year = {2026},
  url = {https://andrewzh112.github.io/opd_adventures}
}

Papers and References

1. Lu, K., & Thinking Machines Lab. (2025). On-Policy Distillation . Technical blog post.
2. Schulman, J., & Thinking Machines Lab. (2025). LoRA Without Regret . Technical blog post.
3. Yao, F., et al. (2025). Your Efficient RL Framework Secretly Brings You Off-Policy RL Training . Technical blog post.
4. Anthropic. (2025). Claude 3.7 Sonnet and Claude Code . Product announcement blog post.
5. OpenAI. (2025). Introducing Codex . Product announcement blog post.
6. Agarwal, R., Vieillard, N., et al. (2024). On-Policy Distillation of Language Models: Learning from Self-Generated Mistakes . ICLR 2024.
7. Chu, T., Zhai, Y., et al. (2025). SFT Memorizes, RL Generalizes: A Comparative Study of Foundation Model Post-training . ICML 2025.
8. Shenfeld, I., Pari, J., et al. (2025). RL's Razor: Why Online Reinforcement Learning Forgets Less . ICLR 2025.
9. Hinton, G., Vinyals, O., & Dean, J. (2015). Distilling the Knowledge in a Neural Network . arXiv preprint.
10. Shao, Z., Wang, P., Zhu, Q., et al. (2024). DeepSeekMath: Pushing the Limits of Mathematical Reasoning in Open Language Models . arXiv preprint.
11. Shah, V., Obando-Ceron, J., Jain, V., et al. (2025). A Comedy of Estimators: On KL Regularization in RL Training of LLMs . arXiv preprint.
12. Ling Team, et al. (2025). Every Step Evolves: Scaling Reinforcement Learning for Trillion-Scale Thinking Model . arXiv preprint.
13. Liu, J., Li, Y., et al. (2025). When Speed Kills Stability: Demystifying RL Collapse from the Training-Inference Mismatch . Technical blog post.
14. Qwen Team. (2025). Qwen3 Technical Report . arXiv preprint.
15. Hu, J., et al. (2025). BroRL: Scaling Reinforcement Learning via Broadened Exploration . arXiv preprint.
16. Schulman, J. (2020). Approximating KL Divergence . Blog post.
17. Chen, L., Lu, K., et al. (2021). Decision Transformer: Reinforcement Learning via Sequence Modeling . NeurIPS 2021.
18. Schmidhuber, J. (2019). Reinforcement Learning Upside Down: Don't Predict Rewards -- Just Map Them to Actions . arXiv preprint.
19. Wang, S., et al. (2025). Beyond the 80/20 Rule: High-Entropy Minority Tokens Drive Effective Reinforcement Learning for LLM Reasoning . NeurIPS 2025.
20. Zhao, S., et al. (2026). Self-Distilled Reasoner: On-Policy Self-Distillation for Large Language Models . arXiv preprint.
21. Hübotter, J., et al. (2026). Reinforcement Learning via Self-Distillation . arXiv preprint.
22. Shenfeld, I., et al. (2026). Self-Distillation Enables Continual Learning . arXiv preprint.
23. Shi, T., et al. (2026). Experiential Reinforcement Learning . arXiv preprint.
24. Ouyang, L., Wu, J., Jiang, X., et al. (2022). Training language models to follow instructions with human feedback . NeurIPS 2022.
25. Goodfellow, I., Bengio, Y., & Courville, A. (2016). Deep Learning . MIT Press.

Select runs from the exported logs. The default view starts with Default.