One-sentence Summary

Researchers from HKUST-GZ and Microsoft introduce Trust Region Policy Distillation (TOP-D), which stabilizes on-policy distillation by dynamically constructing a proximal teacher, formally controlling gradient variance through global convergence and monotonic improvement bounds, and dramatically enhancing training stability, sample efficiency, and final performance on mathematical reasoning tasks without any additional computational overhead.

Key Contributions

• Trust Region Policy Distillation (TOP-D) is a plug-and-play distillation paradigm that smooths the unbounded reward signal and safely relaxes the strict on-policy data-reuse constraint of standard On-Policy Distillation.
• A rigorous theoretical framework is established where an external proximal teacher inherently bounds gradient variance to stabilize training, and a convergence analysis with a monotonic improvement bound forms a closed-loop that systematically minimizes asymptotic error.
• Experiments on mathematical reasoning benchmarks show that TOP-D delivers a 25.84% absolute improvement in avg@32 accuracy over standard OPD on AIME24 and consistently outperforms both OPD and RLVR baselines, all without additional computational overhead.

Introduction

On-Policy Distillation (OPD) has become a widely adopted post-training paradigm for large language models, where a student model learns from token-level log-probability ratios of a stronger teacher. This formulation behaves like dense-reward reinforcement learning, preserving supervised fine-tuning knowledge while offering better sample efficiency than sparse verifiable reward methods. However, standard OPD suffers from optimization fragility because the student-teacher capacity gap leads to unbounded logarithmic differences in the reward signal, causing high variance. Existing fixes, such as reward clipping and mixed sampling, are empirical heuristics that lack theoretical guarantees. The authors introduce Trust Region Policy Distillation (TOP-D), which dynamically interpolates the teacher and student distributions to create a proximal teacher, turning the unbounded reward into a smooth, strictly lower-bounded signal that analytically prevents variance explosion. Paired with internal trust region iterations for safe off-policy data reuse, TOP-D guarantees monotonic policy improvement and tightens the convergence gap, all while reducing to a simple algebraic reward transformation with no additional computational overhead.

Method

The authors propose Trust Region Policy Distillation (TOP-D), a unified algorithmic framework designed to resolve the optimization fragility and low sample efficiency inherent in standard On-Policy Distillation (OPD). To establish a rigorous theoretical foundation, the authors cast autoregressive language modeling as a deterministic Markov Decision Process (MDP). In this formulation, the state space represents the prompt and prefix, while the action space encompasses the vocabulary and the EOS token. The immediate reward is defined as the distillation signal, specifically the log probability ratio between the target teacher and the student policy.

A fundamental challenge in standard OPD arises from the capacity gap between the student and the target teacher. When the teacher assigns a near-zero probability to a generated token, the immediate reward diverges to negative infinity, destabilizing the policy gradient. To address this instability, the authors construct an external proximal teacher, which serves as a localized surrogate target distribution. This proximal teacher interpolates between the target teacher and the current student policy in the probability space: $\tilde{\pi}^* (y_k \mid x, y_{<k}) = \alpha \pi^* (y_k \mid x, y_{<k}) + (1 - \alpha) \pi_\theta (y_k \mid x, y_{<k})$π~∗(yk​∣x,y<k​)=απ∗(yk​∣x,y<k​)+(1−α)πθ​(yk​∣x,y<k​) Consequently, the token-level reward is inherently smoothed and strictly bounded from below, preventing catastrophic divergence. The reward behavior with respect to the probability ratio is visualized below.

Building upon the stabilized reward signal, the authors introduce internal trust region iterations to overcome the notoriously low sample efficiency of strictly on-policy methods. By decoupling the behavior policy from the target policy, the framework inherits the high sample efficiency of modern trust region algorithms. The optimization objective employs an importance sampling ratio with a clipping mechanism to restrict destructively large policy updates: $J(\theta) = \mathbb{E} \left[ \frac{1}{\sum_{i=1}^G |y^i|} \sum_{i=1}^G \sum_{t=1}^{|y^i|} \left\{ \min \left[ p_t^i \cdot \hat{A}_t^i, \operatorname{clip}(p_t^i, 1-\epsilon, 1+\epsilon) \cdot \hat{A}_t^i \right] \right\} \right]$J(θ)=E[∑i=1G​∣yi∣1​∑i=1G​∑t=1∣yi∣​{min[pti​⋅A^ti​,clip(pti​,1−ϵ,1+ϵ)⋅A^ti​]}] To fully exploit the dense reward signals provided by the proximal teacher, the authors adopt a fine-grained token-level normalized advantage. For each prompt, the advantage relies on a token-level reward that combines the immediate smoothed reward with a length-normalized future return: $\tilde{R}_k^i = \tilde{r}_k^i + \frac{1}{|y^i| - k} \sum_{j=k+1}^{|y^i|} \tilde{r}_j^i$R~ki​=r~ki​+∣yi∣−k1​∑j=k+1∣yi∣​r~ji​ This formulation effectively prevents the model from generating overly short or excessively long responses. The token-level normalization process across responses is illustrated as follows.

The TOP-D framework is supported by a comprehensive theoretical analysis. The authors prove that the proximal teacher inherently bounds the gradient variance, with the interpolation coefficient acting as a direct variance controller. Furthermore, they derive a global convergence bound that reveals an inherent error forgetting mechanism, isolating early high-variance updates. Finally, to minimize the single-step optimization error and tightly close the asymptotic convergence gap, the authors establish a monotonic improvement guarantee for the internal trust region iterations. This theoretical synergy ensures that the reverse KL divergence objective monotonically improves during the internal updates, rigorously guaranteeing the continuous minimization of the optimization error.

Experiment

The evaluation involves training on the DAPO-Math-17k dataset and benchmarking on AIME and AMC competitions, using Qwen3-1.7B/8B base students supervised by Qwen3-14B/30B teachers, and comparing TOP-D against post-training baselines including GRPO, DAPO, and standard OPD. The results show that TOP-D consistently outperforms all baselines, with the most dramatic gains over standard OPD, whose unbounded reward variance causes instability and performance collapse when the teacher-student capacity gap is large. An ablation study confirms that both the external proximal teacher and the internal trust region iterations are essential for training stability and sample efficiency, while the method remains robust to the interpolation coefficient.

Standard On-Policy Distillation (OPD) uses on-policy sampling and dense rewards but suffers from low training stability and only weak theoretical guarantees, in contrast to Supervised Fine-Tuning (off-policy, dense, high stability) and RL with Verifiable Rewards (on-policy, sparse, low stability). Trust Region Policy Distillation (TOP-D) retains on-policy dense rewards while achieving high stability and strong theoretical guarantees by bounding the distillation reward with a proximal teacher and trust region iterations, leading to large performance gains over standard OPD, especially when the teacher–student capacity gap is large. TOP-D transforms the unbounded logarithmic probability ratio of standard OPD into a bounded reward, dramatically improving training stability and yielding over 25 absolute percentage points higher accuracy on AIME24 for an 8B student model. Both the external proximal teacher and the internal trust region iterations are indispensable: removing the proximal teacher causes unstable training, while disabling trust region iterations drastically reduces sample efficiency.

TOP-D substantially outperforms all baselines on mathematical benchmarks. For the Qwen3-8B-Base model, it more than doubles the AIME24 accuracy of standard OPD and surpasses the best RLVR method by a large margin. When scaling down to a smaller 1.7B student, TOP-D remains robust while standard OPD collapses, highlighting the method's stability across model sizes. TOP-D achieves a 25.84 percentage point absolute improvement over standard OPD on AIME24 for the Qwen3-8B-Base student, reaching 50.42% accuracy. TOP-D beats the strongest RLVR baseline DAPO by 17.5 percentage points on the same AIME24 benchmark, confirming its advantage over reinforcement learning approaches.

For the Qwen3-1.7B-Base student, standard OPD distillation falls behind RLVR methods, achieving only 8.96% on AIME24 compared to DAPO's 12.29%. TOP-D overcomes this fragility, lifting AIME24 accuracy past 20% and surpassing standard OPD by over 10 percentage points on both AIME24 and AIME25. This gap highlights the instability of unbounded distillation when the teacher-student capacity difference is severe. Standard OPD with the 30B teacher reaches 8.96% on AIME24, trailing the RLVR baseline DAPO at 12.29%. TOP-D pushes AIME24 accuracy past 20% for the 1.7B student, a gain of more than 10 absolute percentage points over standard OPD.

The experiments evaluate policy distillation methods on mathematical reasoning benchmarks using Qwen3 student models of varying sizes. Standard On-Policy Distillation suffers from instability and weak performance, especially with a large teacher–student capacity gap, while Trust Region Policy Distillation (TOP-D) introduces a bounded reward and trust region iterations to achieve high stability and strong theoretical guarantees. TOP-D dramatically outperforms standard OPD and reinforcement learning baselines, proving robust even when the student model is much smaller, and both the proximal teacher and internal trust region steps are essential for its success.