On-Policy or Off-Policy Learning? A Systematic Study of Distillation Dynamics

AI summary

arXiv: 2609.35259v1 [cs.LG], 28 Sep 2026 Authors: Julianna Piskorz, Antonin Berthon, Mihaela van der Schaar (University of Cambridge)

Overview

• Research area: Large language model post-training, specifically strong-to-weak knowledge distillation and the on-policy versus off-policy distinction.
• Technical level: Advanced. The paper combines controlled training experiments with a token-level gradient analysis of forward and reverse KL divergence and a continuous rollout-policy interpolation.
• Scope in one sentence: A controlled study that disentangles rollout policy, token-level KL direction, and learning rate in strong-to-weak distillation across two model families and three reasoning domains, finding that rollout policy is not the dominant factor that prior work suggested.

What This Paper Is About

Comparisons between supervised fine-tuning (SFT) and reinforcement learning with verifiable rewards (RLVR) have led to the claim that on-policy learning reduces catastrophic forgetting, produces sparser parameter updates, and improves generalisation, but those comparisons change many training factors at once. The authors use strong-to-weak distillation as a controlled testbed where the data-generating policy can be changed while everything else in the pipeline is held fixed, allowing them to ask whether the reported advantages actually come from the rollout policy. The goal is to isolate what rollout policy, KL direction, and learning rate each contribute to final accuracy, forgetting, update sparsity, output coverage, and generalisation.

Key Contributions

1. A controlled study of distillation. The authors disentangle rollout policy, KL direction, and learning rate in strong-to-weak distillation, and find no consistent advantage from on-policy rollouts in final accuracy, catastrophic forgetting, or parameter-update sparsity.
2. An objective-dependent account of rollout-policy sensitivity. Through gradient analysis and experiments along a continuous student–teacher rollout-policy spectrum, they show that forward KL is robust to the rollout policy whereas reverse KL is substantially more sensitive.
3. A characterisation of when on-policy data helps. Training stability and output coverage are largely driven by KL direction and learning rate determines forgetting and update sparsity, while on-policy rollouts consistently improve generalisation to harder task variants and can reduce incidental teacher-style transfer.
4. Robustness checks on the experimental design. The main conclusions are tested against removing gradient clipping, replacing full-vocabulary KL with sampled KL estimators, and training on a task with much longer reasoning traces.

Main Findings

• No consistent accuracy advantage for on-policy rollouts: across MedReason, Science, and Countdown-3, the best mean accuracies are nearly identical, at 72% for on-policy distillation (OnPD) and 73% for off-policy distillation (OffPD).
• KL direction matters more for accuracy: forward KL achieves 71–73% across all rollout policies and learning rates, whereas reverse KL ranges from 35% to 72% and is substantially more sensitive to the learning rate.
• Learning rate, not rollout policy, governs forgetting: mean out-of-distribution (OOD) performance changes by at most 1.3 percentage points at the lower learning rate (1×10⁻⁵) but drops by 11.2–14.0 points at the higher rate (5×10⁻⁵).
• Update sparsity is also set by learning rate: sparsity ranges from 85.3–89.6% at the lower learning rate compared with 51.9–60.0% at the higher learning rate, and OffPD produces at least as sparse updates as OnPD in every matched comparison.
• Forgetting is not a necessary cost of strong performance: under forward KL, a small learning rate achieves top target-task performance without decreasing prior capabilities, resembling the "cliff" phenomenon previously identified in the SFT setting.
• Gradient analysis explains the asymmetry: the forward-KL derivative with respect to student logits is π_S^θ(v) − π_T^θ(v), nonzero whenever the distributions differ and bounded in [−1, 1], so under bounded student-logit Jacobians the difference between the updates induced by two rollout policies is bounded linearly by the total-variation distance between their prefix distributions. Reverse KL's derivative is weighted by π_S^θ(v) and includes a log-ratio term that can become arbitrarily large, admitting no bound dependent solely on rollout distance and the Jacobian.
• Forward KL is robust across the rollout-policy spectrum: held-out accuracy stays above 80% across the entire spectrum, varying by only 5.2 percentage points at a learning rate of 1×10⁻⁵, even though OnPD with forward KL passes through a regime of largely incoherent, high-entropy rollouts.
• Reverse KL favours student-generated rollouts: reverse KL exhibits large changes in mean performance across λ and considerable variance across seeds, benefiting strongly from student-favoured rollouts (λ > 0) at the lower learning rate and remaining unstable at the higher learning rate, occasionally collapsing.
• Forward KL yields better output coverage: at comparable pass@1 for λ > 0, forward KL leads to significantly higher pass@10 than reverse KL on Countdown-3, with the same effect observed on Qwen2.5.
• On-policy rollouts generalise better to harder tasks: for both KL directions, pass@k on the harder Countdown-4E increases gradually across the spectrum, with on-policy rollouts (λ > 0) consistently giving 10–15% higher pass@k accuracy than off-policy rollouts.
• The on-policy generalisation advantage does not reliably survive RLVR: after 300 RLVR steps on Countdown-4, reverse-KL checkpoints distilled at 1×10⁻⁵ improve quickly but later undergo reward collapse, while OffPD checkpoints distilled with forward KL at 1×10⁻⁵, or reverse KL at 5×10⁻⁵, improve steadily and reach the highest final performance despite starting near 0% accuracy.
• On-policy plus reverse KL reduces incidental teacher-style transfer: when the teacher is instructed to reason in Spanish, OnPD with reverse KL largely preserves the student's English response style, whereas the other configurations exhibit near-complete transfer to Spanish.
• Conclusions survive design changes: with sampled KL estimators, forward KL remains robust, reverse KL remains brittle, and learning rate continues to govern forgetting and sparsity; removing gradient clipping exposes a stronger difference between KL directions but no consistent OnPD–OffPD gap; and on Numina–MATH, where teacher responses average 622 tokens, OnPD with reverse KL achieves the highest MATH-500 accuracy at the lower learning rate, though this is one run per condition and treated as suggestive.

Methodology in Plain English

The authors take a large "teacher" model and a smaller "student" from the same family so that they share a tokenizer and vocabulary, then train the student to match the teacher's next-token probability distributions. The key trick is that the policy used to generate the training text can be changed without touching anything else, so any resulting difference can be attributed to that choice rather than to the objective or the optimiser.

They cross two design choices that are usually tied together. The first is the rollout policy: in off-policy distillation the student trains on fixed teacher-generated responses, and in on-policy distillation it generates its own trajectories and queries the teacher for supervision along them. The second is the token-level KL direction: forward KL is mode-covering, weighting tokens the teacher considers likely, while reverse KL is mode-seeking, weighting tokens the student favours. Because they compute the KL over the whole vocabulary rather than from single sampled tokens, the usual sampling-based reason to pair forward KL with teacher rollouts and reverse KL with student rollouts no longer applies.

For each configuration they vary the learning rate (1×10⁻⁵ or 5×10⁻⁵) and train with full-parameter fine-tuning for 150 optimisation steps, using three random seeds. They measure held-out accuracy on the target task, catastrophic forgetting as the drop in mean score across seven OOD benchmarks, and update sparsity as the fraction of parameters whose absolute change from the initial checkpoint is below 10⁻⁶.

To map the space between the two extremes, they build a rollout-policy spectrum: a policy π_λ formed by blending the log-probabilities of the student and teacher, where λ = −1 recovers standard off-policy distillation, λ = 1 recovers on-policy distillation, λ = 0 is the symmetric midpoint, negative λ gives teacher-favoured rollouts, and positive λ gives student-favoured rollouts, with λ < −1 or λ > 1 extrapolating beyond the two models. Log-ratio clipping and plausibility masking prevent the scheme from amplifying tokens both models consider unlikely.

Why This Matters

• Impact on research: The paper challenges the widely repeated claim that on-policy rollouts are inherently preferable, showing that many of the differences attributed to rollout policy in SFT-versus-RLVR comparisons are better explained by the learning objective and the learning rate. It recommends that on-policy distillation work include off-policy distillation as a standard baseline.
• Real-world applications:
  • Training small, deployable reasoning models for medical, scientific, and arithmetic domains using larger teachers.
  • Reducing post-training compute, since off-policy methods can reuse existing data instead of requiring continual generation throughout training.
  • Designing multi-step post-training pipelines, where distillation checkpoints are used as starting points for subsequent RLVR.
  • Controlling incidental behaviour transfer, such as a student adopting a teacher's language or stylistic quirks, via on-policy rollouts combined with reverse KL.
• Industry relevance: On-policy training is expensive because it requires generating rollouts continuously, while off-policy training can reuse a fixed dataset. If the advertised benefits of on-policy data come from other design choices, practitioners may be paying a substantial cost for little gain, and the paper suggests that tuning the learning rate and choosing the KL direction are higher-leverage decisions.

Future Directions

1. Scaling beyond the studied regime. The authors deliberately bound the study to student models of at most 1.5B parameters and reasoning traces of at most 2,000 tokens; extending the analysis to larger models and longer rollouts is left open.
2. Varying the teacher. The teacher model is held fixed throughout, so whether OnPD–OffPD differences emerge under specific student–teacher combinations remains untested.
3. Explaining the RLVR interaction. The initial generalisation advantage of on-policy checkpoints does not translate into a reliable advantage after RLVR, and the strongest sustained RLVR performance comes from off-policy checkpoints, which the authors say requires further research into how best to combine post-training methods into multi-step pipelines.
4. Confirming the long-rollout result. The Numina–MATH experiment, where on-policy rollouts with reverse KL achieved the highest MATH-500 accuracy at the lower learning rate, contains only one run per condition, so the authors treat the advantage as suggestive and call for more runs.

Target Audience

Researchers and engineers working on LLM post-training, knowledge distillation, and reinforcement learning from verifiable rewards, particularly those designing training recipes that mix supervised fine-tuning, distillation, and RLVR. The paper is also relevant to practitioners who must justify the additional compute cost of on-policy training, and to readers interested in the theory of forward versus reverse KL divergence in sequence-level objectives. Familiarity with KL divergence and language model fine-tuning is assumed.

Authors’ abstract

On-policy learning has been argued to reduce catastrophic forgetting, produce sparser parameter updates, and improve generalisation. However, existing comparisons between supervised fine-tuning and reinforcement learning vary many factors simultaneously, making the contribution of rollout policy difficult to isolate. We study the effect of rollout policy in a controlled strong-to-weak distillation setting, by independently varying rollout policy, token-level KL direction, and learning rate across the Llama3 and Qwen2.5 model families and reasoning tasks spanning scientific, medical, and arithmetic domains. Our analysis reveals a nuanced picture of distillation dynamics in which rollout policy does not necessarily play a central role. Instead, token-level KL direction more clearly shapes task performance and output coverage, while learning rate governs forgetting and update sparsity. Analysis of KL gradients and experiments along a continuous student-teacher rollout-policy spectrum explain this pattern: forward KL is remarkably robust to rollout policy, with its performance stable and strong despite changes to the rollout policy, whereas reverse KL is substantially more sensitive and favours student-generated rollouts. On-policy data nevertheless improves generalisation to harder variants of the Countdown arithmetic task under both KL directions, although this advantage does not reliably persist after subsequent RLVR. Our broader conclusions remain robust to removing gradient clipping, using sampled KL estimators, and training on tasks requiring longer reasoning chains. Overall, our results challenge the view that on-policy rollouts are inherently preferable and show that their value depends critically on the objective, evaluation setting, and optimisation hyperparameters.