PaperlayerView exact revision on arXiv

Source-linked AI summary

Trust Region On-Policy Distillation

Xingrun Xing, Haoqing Wang, Boyan Gao, Ziheng Li, Yehui Tang

arXiv:2606.01249v3cs.LGcs.CL

TL;DR

On-policy distillation can become unstable when student-generated tokens fall outside regions where teacher supervision is reliable. TrOPD uses trust-region learning, outlier estimation, and off-policy guidance, and consistently outperforms OPD baselines across evaluated tasks.

  • Problem

    Existing on-policy distillation methods lack reliable supervision when teacher and student distributions diverge, causing unstable policy-gradient estimation.

  • Method

    TrOPD partitions student tokens by teacher-verifiable reliability, estimates outliers with top-k forward KL, and adds imitation from teacher-generated trajectories.

  • Results

    Across mathematical and general-domain tasks, TrOPD outperforms OPD; its FKL variant exceeds OPD by 3.06 points on average.

  • Takeaways & Limitations

    TrOPD’s results highlight supervision reliability as important for stable on-policy reasoning distillation and training capable small reasoning models.

  • Takeaways & Limitations

    The study lacks practical deployment and application studies and focuses on OPD-based post-training, which may constrain resulting reasoning performance.

Abstract

from arXiv · show

On-Policy Distillation (OPD) is a fundamental technique for efficient post-training of large language models (LLMs), with broad applications in agent learning, multi-task enhancement, and model compression. However, OPD training becomes unstable when the teacher and student distributions differ substantially, as teacher supervision on student-generated tokens may yield unreliable policy gradients and even cause optimization failure. This work addresses reliable on-policy token-level supervision through credit assignment strategies, and proposes Trust Region On-Policy Distillation, TrOPD. It features the following characteristics: 1) Trust-Region On-Policy Learning: TrOPD performs OPD only in regions where the teacher provides reliable supervision, mitigating the optimization difficulty of the K1 reverse-KL estimator under distribution mismatch. 2) Outlier Estimation: For outlier regions, we explore gradient clipping, masking, and forward-KL estimation to reduce the adverse effects of unreliable supervision. 3) Off-Policy Guidance: The student continues generation from teacher prefixes and uses forward KL to imitate off-policy guidance, encouraging on-policy exploration toward reliable regions. Experiments show that TrOPD consistently outperforms SoTA OPD baselines, including OPD, EOPD, and REOPOLD, across mathematical reasoning, code generation, and general-domain benchmarks.

1 Introduction

Large reasoning models achieve strong results but incur substantial inference costs, motivating small reasoning models and on-policy distillation. This work identifies unreliable supervision as a central OPD failure mode and proposes TrOPD with a benchmark and trust-region supervision strategy.

  • Motivation: Large reasoning models improve mathematics, code generation, and agent-task performance through scaled test-time reasoning, but their inference costs motivate small reasoning models.The introduction frames small reasoning models as a response to the substantial inference costs of large reasoning models.
  • On-Policy Distillation: On-Policy Distillation trains on student-generated trajectories to mitigate exposure bias and efficiently support small reasoning models.OPD addresses the mismatch between teacher-generated training trajectories and student-generated inference trajectories, especially in long-chain-of-thought reasoning.
  • Challenges: Existing OPD methods become unstable when teacher and student distributions diverge, producing unreliable policy gradients that can cause training collapse.Reasoning-oriented OPD also faces prohibitive memory overhead from full-vocabulary supervision.
  • Benchmark: The work establishes a unified benchmark spanning mathematics, code generation, and STEM reasoning while comparing diverse OPD strategies and memory-efficient KL estimators.The benchmark studies reliable OPD through multi-domain evaluation, unified method comparisons, and K1 and top-k KL estimation for long-response distillation.
  • TrOPD: TrOPD partitions student-generated tokens by supervision reliability using teacher–student decoding agreement and applies top-k forward-KL estimation to outliers.This design preserves informative reward signals while avoiding unreliable supervision in outlier regions.

2 Related Works

Prior work has advanced LLM reasoning through prompting, reinforcement learning, supervised fine-tuning, and on-policy distillation. Knowledge distillation has progressed from off-policy, sequence-level supervision to full-vocabulary on-policy methods, including k1-based objectives for reasoning-model post-training.

  • Reasoning Language Models: Reasoning capabilities in LLMs have been elicited through prompts and acquired through reinforcement learning, supervised fine-tuning, and on-policy distillation.Reasoning is also increasingly integrated with agentic and multimodal capabilities.
  • Knowledge Distillation: Knowledge distillation was introduced to efficiently train compact models and later applied to generative language models through sequence-level, off-policy teacher responses.Sequence-level knowledge distillation performs supervised fine-tuning on teacher-generated responses.
  • Knowledge Distillation: Full-vocabulary on-policy distillation methods, including MiniLLM and GKD, were developed to mitigate exposure bias.These methods include work cited as Ko et al., 2024; 2025; and Xu et al., 2025.
  • Knowledge Distillation: For reasoning models, KL objectives based on the k1 estimator have been applied to improve reasoning performance during post-training.The cited works include Lu et al. (2025) and Yang et al. (2026).

3 Problem Formulation

The section formulates OPD as student-generated, reverse-KL-based distillation with a policy-gradient objective, while identifying memory costs and optimization failures caused by distribution mismatch and low-quality student generations.

  • 3.1 Distillation for Language Models: OPD trains on student-generated outputs rather than teacher-generated outputs to mitigate exposure bias.This distinguishes on-policy distillation from off-policy distillation.
  • 3.1 Distillation for Language Models: Reverse-KL OPD samples trajectories from the student policy and rewards sequences assigned high probability by the teacher.Its expectation over the student distribution makes it mode-seeking: low-teacher-probability student outputs are strongly penalized, while unexplored teacher modes receive little direct penalty.
  • 3.2 OPD for Reasoning Models: Token-level reverse-KL OPD methods optimize the KL reward with policy gradients, but full-vocabulary computation requires O(n*k) memory.Here, n denotes sequence length and k denotes vocabulary size.
  • 3.2 OPD for Reasoning Models: Traditional instruction-model distillation methods such as GKD and speculative KD compute full-vocabulary KL divergence for stable optimization.This contrasts with the memory burden of token-level reverse-KL OPD.
  • 3.2 OPD for Reasoning Models: Longer reasoning sequences make memory consumption a major distillation bottleneck, motivating the K1 estimator as an unbiased KL-divergence estimate.Reasoning-oriented models scale performance by extending output sequence length.
  • 3.2 OPD for Reasoning Models: When teacher–student discrepancy is large, extremely low teacher probabilities on student trajectories create significant policy-gradient outliers.In low-confidence regions, Ex∼PS[πT(x)] ≈ 0 and the K1-based gradient can become extremely negative, with πS(x) → −∞, destabilizing optimization and limiting final performance.
  • 3.2 OPD for Reasoning Models: Because OPD trains exclusively on student-sampled trajectories, low-quality student generations restrict the optimization space and reduce access to informative supervision.This is especially problematic for challenging problems, where the student may fail to generate high-quality responses.

4 Trust Region Distillation

Section 4 shows that conventional FKL, entropy filtering, and reward clipping provide inadequate or inconsistent solutions under distribution mismatch, motivating TrOPD. TrOPD combines adaptive trust-region optimization, outlier-specific FKL estimation, and off-policy teacher-prefix guidance to stabilize on-policy distillation.

  • Token Filtering and Reward Clipping: Entropy-aware token selection can degrade OPD performance, while reward clipping provides improvements that become marginal across settings.These results suggest ordinary tokens may still provide informative teacher supervision and should not always be discarded.
  • Divergence Evaluation: Conventional FKL alone fails under constrained vocabulary, while combining FKL with RKL better addresses large distributional mismatch.Top-k FKL is a biased approximation of full-vocabulary FKL and can distort policy gradients when applied to all sampled tokens.
  • Trust Region Distillation: TrOPD optimizes only within regions where token-level policy gradients are reliable, using trust-region learning to suppress unreliable supervision.Within the trust region, TrOPD uses the K1 RKL estimator for student-sampled tokens.
  • Outlier Handling: Mask Outlier and Clip Outlier both outperform vanilla OPD, with Mask Outlier preserving higher policy entropy by eliminating unreliable gradients.Masking removes the token-level advantage when reward magnitude exceeds the threshold, whereas clipping limits the reward.
  • Adaptive and Off-Policy Guidance: TrOPD adaptively defines trust regions from student and teacher policies, adds auxiliary FKL in outlier regions, and guides student continuation from teacher-generated prefixes.The off-policy trajectory uses a teacher-generated prefix followed by a student-generated continuation, with the prefix length annealed to zero during training.

5 Experimental Results

TrOPD consistently improves on-policy distillation across single-domain and multi-domain reasoning settings, with gains supported by trust-region variants, outlier estimation, and off-policy guidance. Ablations show FKL outlier estimation and off-policy guidance are effective, while combining TrOPD with AOPD yields further improvement.

  • Benchmark Evaluation: Experiments evaluate distilled models on mathematics, STEM, instruction following, and code generation benchmarks, including AIME, AMC, GPQA Diamond, MMLU-Redux v2, IFBench, and LiveCodeBench v6.Mathematical results average 32 evaluation runs.
  • Single-Domain Distillation: TrOPD improves over OPD by +3.06 points on mathematical reasoning and +2.63 points on general-domain tasks in single-domain distillation.Evaluation includes AIME 2024, AIME 2025, AMC 2023, and out-of-domain tasks.
  • Multi-Domain Distillation: TrOPD achieves average gains of +4.62 and +3.44 points over OPD for DeepSeek-Qwen2.5-1.5B and Qwen3-SFT-1.7B, respectively.The multi-domain evaluations cover mathematics and code-generation tasks across different teacher–student configurations.
  • Ablation Studies: FKL applied only to outlier regions outperforms masking and clipping, while off-policy guidance further improves all three outlier-estimation variants.TrOPD Mask, TrOPD Clip, and TrOPD FKL outperform OPD by 2.00, 1.94, and 3.06 points on average, respectively.
  • Comparison with AOPD: TrOPD outperforms AOPD, and combining them raises the average score from 40.63 to 41.67.The result suggests that AOPD and TrOPD provide complementary optimization strategies.

6 Conclusion

The paper proposes TrOPD as a reliable, stable framework for reasoning-oriented on-policy distillation. It combines trust-region optimization, outlier estimation, and off-policy guidance to suppress unreliable gradients, preserve informative supervision, and encourage exploration toward teacher-supported trajectories.

  • TrOPD is proposed as a reliable and stable framework for reasoning-oriented on-policy distillation.The framework targets improved reliability in OPD training.
  • Trust-region optimization and outlier estimation suppress unreliable policy gradients while preserving informative supervision.These mechanisms address supervision reliability during on-policy reasoning distillation.
  • Off-policy guidance encourages exploration toward trajectories supported by the teacher.This provides additional guidance beyond the student’s on-policy trajectories.
  • Extensive multi-domain results highlight the importance of supervision reliability in on-policy reasoning distillation.The conclusion reports evidence across multiple domains without specifying individual benchmark results.

Limitations

The work lacks practical deployment and application studies on small reasoning models. Its focus on OPD-based post-training with two specific models may constrain the upper bound of reasoning performance.

  • Limitations: The primary limitation is the lack of practical deployment and application studies on small reasoning models.
  • Limitations: High-performing small reasoning models may require mid-training to further improve post-training reasoning capabilities in real-world scenarios.
  • Limitations: The study primarily uses OPD-based post-training on DeepSeek-Qwen2.5-1.5B and Qwen3-SFT-1.7B, potentially constraining resulting reasoning performance.

Appendix · A Training Details of Teacher Model Qwen3-Nemotron-4B

Qwen3-Nemotron-4B is trained from Qwen3-4B-Base through supervised finetuning and reinforcement learning with verifiable rewards. The appendix specifies the data mixtures, optimization settings, and rollout configuration for both stages.

  • A Training Details of Teacher Model Qwen3-Nemotron-4B: The pipeline starts from Qwen3-4B-Base and combines supervised finetuning with reinforcement learning using verifiable rewards.SFT uses publicly available Nemotron 3 Nano1 datasets, removes entries without a messages field, combines multiple domains, and upsamples smaller datasets to match reported domain ratios.
  • A Training Details of Teacher Model Qwen3-Nemotron-4B: RLVR uses the publicly available Nemotron 3 Nano2 training blend spanning math, coding, science, and instruction following.The mixture contains 22,056 math, 19,169 coding, 19,670 science, and 16,575 instruction-following samples.
  • A Training Details of Teacher Model Qwen3-Nemotron-4B: SFT uses Adam with a learning rate of 5 × 10−5, weight decay of 0.1, and warmup over 10% of total training steps.The SFT batch size is 512, with an average response length of approximately 7K tokens.
  • A Training Details of Teacher Model Qwen3-Nemotron-4B: SFT and RLVR use batch sizes of 512 and 128, respectively, reflecting distinct training configurations.SFT averages approximately 7K-token responses, while RLVR updates parameters every 2048 rollouts.
  • A Training Details of Teacher Model Qwen3-Nemotron-4B: RLVR employs GRPO with group size 16 and masked importance sampling to improve consistency between training and inference.RLVR uses batch size 128 and updates model parameters every 2048 rollouts.
  • A Training Details of Teacher Model Qwen3-Nemotron-4B: RLVR caps maximum generation length at 32K tokens and sets sampling temperature to 1.0 to encourage exploration.Further details on the multi-task reinforcement-learning recipe are provided in Wang et al. (2026a).
Loading 2606.01249v3…