Trust Region On-Policy Distillation

Recent years have seen exponential growth in large language models (LLMs) such as GPT, BERT, and T5, driven by their impressive capabilities in understanding and generating complex language. Knowledge distillation emerged as a key technique to compress these models, making deployment feasible in resource-constrained environments. Early methods like sequence-level distillation (Kim & Rush, 2016) focused on off-policy training, where a student mimics teacher outputs. However, these approaches suffer from exposure bias and inefficiency in reasoning tasks, especially with long outputs. On-policy distillation (Lu et al., 2025) addresses some issues but faces instability when teacher-student distributions diverge, leading to unreliable gradients. Long-sequence reasoning further complicates training due to memory and computational constraints, necessitating efficient divergence estimation techniques and stability mechanisms. Recent efforts incorporate KL estimators like K1 and top-k to handle these challenges, but instability persists, especially in complex reasoning scenarios. This paper builds on these foundations, proposing a trust-region framework to improve stability and performance.

The main challenge is the divergence between teacher and student distributions during on-policy distillation, which causes unstable gradients and training collapse, especially in reasoning tasks with long outputs. Existing methods struggle to balance exploration and exploitation, often amplifying errors due to unreliable supervision signals. Memory limitations restrict full-vocabulary KL computation, forcing approximations that may introduce bias. Consequently, training instability hampers the model’s ability to learn complex reasoning skills efficiently. Addressing these issues is crucial for developing small, resource-efficient models capable of high-level reasoning, which remains an open problem in the field.

This work introduces a trust-region based framework (TrOPD) that dynamically partitions tokens into reliable and outlier regions based on decoding agreement, effectively controlling the divergence estimation process. It combines K1-based reverse KL within trusted regions with top-k forward KL in outliers, balancing stability and informativeness. The framework also employs gradient clipping and masking to suppress outlier gradients, and incorporates off-policy prefix guidance to promote exploration. Unlike prior methods, TrOPD adaptively adjusts trust regions during training, ensuring stable convergence even under large distribution mismatches. These innovations enable efficient, robust training of reasoning models with long sequences, addressing core limitations of existing approaches.

• �� Compute decoding agreement ratio between teacher and student to define trust regions (trusted vs. outlier tokens).
• �� In trusted regions, optimize using K1 estimator for reverse KL (RKL), ensuring unbiased, memory-efficient divergence estimation.
• �� In outlier regions, apply top-k KL estimation and auxiliary forward KL (FKL) to recover informative supervision.
• �� Use gradient clipping and masking to suppress unreliable gradients from outliers.
• �� Incorporate off-policy prefix guidance, sampling trajectories from teacher to guide student exploration.
• �� Dynamically adjust trust regions based on real-time agreement metrics, refining the division during training.
• �� Combine all components into a unified objective that balances on-policy and off-policy learning within trust regions.

Experiments conducted on benchmarks like AIME, AMC, GPQA, and LiveCodeBench, with models trained for 200 steps at learning rate 5×10^-6. The support set for KL estimators was fixed at k=64. Comparisons included OPD, EOPD, REOPOLD, and variants of TrOPD with different outlier handling strategies. Ablation studies examined the impact of trust-region partitioning, estimator choices, and outlier suppression techniques. Results demonstrated that TrOPD outperforms baselines in accuracy, stability, and robustness, especially in long-sequence reasoning tasks. Gradient norm analysis confirmed reduced outlier influence, and performance gains were consistent across multiple domains, validating the framework’s effectiveness.

TrOPD achieved up to +6.18 points improvement over SOTA baselines across reasoning benchmarks. In mathematical reasoning, accuracy increased from 29.16 to 32.50 points; in code generation, from 29.80 to 35.98; and in multi-domain tasks, from 21.66 to 27.60. Gradient norm analysis showed a 20% reduction, indicating enhanced stability. Ablation results confirmed that trust-region partitioning and multi-estimator KL are critical for these improvements. The method also demonstrated superior exploration capabilities, maintaining higher entropy during training, which contributed to better generalization.

TrOPD can be applied to resource-efficient training of small reasoning models, enabling deployment in edge devices and real-time systems. It is suitable for domains requiring reliable long-sequence inference, such as automated reasoning, scientific computation, and intelligent tutoring. The framework also facilitates model compression and transfer learning, making it valuable for industry applications where stability and robustness are critical. Long-term, integrating adaptive trust-region mechanisms and multimodal data could further expand its utility.

The reliance on decoding agreement ratios for trust-region partitioning may misjudge trustworthiness under extreme distribution shifts. Additional computational overhead from multiple estimators and dynamic trust adjustments could limit scalability. Effectiveness on ultra-long sequences and highly complex reasoning tasks remains to be validated, requiring further optimization of estimators and trust mechanisms. Future work should focus on reducing complexity and enhancing robustness under diverse scenarios.