Trust Region Policy Distillation

Abstract: Big goals are hard to achieve all at once; breaking them into small steps is wiser. We present Trust Region Policy Distillation (TOP-D), which transforms the notoriously unstable, high-variance On-Policy Distillation (OPD) into a stable training paradigm by dynamically constructing a proximal teacher. Theoretically, we establish a rigorous framework demonstrating that TOP-D inherently controls gradient variance. By providing a formal global convergence analysis alongside a monotonic improvement bound, we mathematically formalize the reliability and stability of the overall training dynamics. Empirically, TOP-D dramatically enhances training stability, sample efficiency, and final performance on mathematical reasoning tasks. More importantly, TOP-D introduces zero additional computational overhead, positioning itself as a promising alternative to the well-established OPD paradigm.

• The paper introduces TOP-D, which replaces direct teacher supervision with a dynamic proximal teacher to yield lower gradient variance.
• TOP-D incorporates internal trust region iterations, allowing data reuse and ensuring monotonic policy improvement with global convergence guarantees.
• Empirical results show a 25.84% accuracy improvement over OPD, demonstrating enhanced sample efficiency and stability in training LLMs.

Trust Region Policy Distillation: An Authoritative Summary

Motivation and Problem Formulation

Standard On-Policy Distillation (OPD) has become the predominant paradigm for post-training LLMs, leveraging token-level signals from a teacher to drive policy learning. OPD is theoretically appealing due to its dense reward regime and avoidance of catastrophic forgetting, but it is plagued by severe optimization instability and low sample efficiency. The root cause is the unbounded logarithmic probability ratio between teacher and student, which yields high gradient variance and vulnerability to teacher-student capacity mismatches. Previous mitigation efforts—teacher-student mixing, reward clipping, and sampling controls—are heuristics and lack principled theoretical guarantees.

The paper "Trust Region Policy Distillation" (2607.04751) introduces Trust Region Policy Distillation (TOP-D), a framework designed to resolve OPD’s inherent instability. TOP-D constructs a dynamic proximal teacher at every training iteration through probability interpolation, transforming unbounded distillation rewards into smooth, lower-bounded signals and enabling rigorous control of gradient variance. The approach incorporates internal trust region iterations, decoupling data reuse from strict on-policy requirements and producing monotonic policy improvement bounds with formal global convergence guarantees.

Figure 1: TOP-D replaces direct teacher supervision in OPD with a dynamically constructed proximal teacher at every iteration, ensuring stable optimization.

Methodology

Proximal Teacher Construction

TOP-D’s central innovation is the external proximal teacher, defined as the convex interpolation of the teacher and student policies:

$$\tilde{\pi}^*(y_k | x, y_{<k}) = \alpha \pi^*(y_k | x, y_{<k}) + (1 - \alpha) \pi_\theta(y_k | x, y_{<k}),$$

where $\alpha \in (0, 1)$. This design bounds the token-level reward:

$$\tilde{r}_k = \log\left(\alpha \frac{\pi^*(y_k | x, y_{<k})}{\pi_\theta(y_k | x, y_{<k})} + 1 - \alpha\right),$$

which, unlike standard OPD, cannot diverge as $\pi^* \to 0$. This prevents catastrophic penalization for rare tokens and yields robust optimization dynamics.

Internal Trust Region Iterations

TOP-D incorporates trust region iterations to decouple the behavior and target policies. Data generated under the current student (behavior policy) can be reused for multiple internal epochs. Importance sampling corrects for policy evolution, and token-level normalized advantage estimation is employed, leveraging dense rewards and penalizing excessive response lengths.

Figure 2: TOP-D leverages token-level normalization across responses to better harness dense reward signals compared to sequence-level normalization.

Algorithmically, these innovations combine to break OPD’s strict on-policy barrier, achieve high sample efficiency, and maintain optimization stability. The proximal teacher eliminates variance explosions, while trust region data reuse ensures monotonic improvement and tighter convergence bounds.

Theoretical Guarantees

The paper establishes strong mathematical foundations for TOP-D:

• Bounded Gradient Variance: The proximal teacher ensures that the variance of the gradient estimator is strictly bounded by a uniform absolute limit (dependent on $\alpha$ and vocabulary size), preventing divergence in the probability space. As $\alpha \to 1$, TOP-D reduces to OPD with unbounded variance; as $\alpha \to 0$, variance vanishes.
• Global Convergence: The iterative proximal teacher operator yields a geometric convergence rate and introduces an error-forgetting mechanism. The asymptotic convergence gap depends on the minimized single-step optimization error and $\alpha$, motivating internal trust region iterations.
• Monotonic Policy Improvement: The internal trust region updates guarantee monotonic improvement in the reverse KL divergence objective, yielding a formally decreasing sequence of divergences during internal epochs.

Empirical Results

Empirical validation focuses on mathematical reasoning benchmarks (AIME, AMC, MATH, Olympiad) using Qwen3-1.7B and Qwen3-8B students, supervised by stronger teacher models (Qwen3-14B, Qwen3-30B-A3B-Instruct-2507). TOP-D dramatically outperforms OPD and RLVR baselines in both sample efficiency and final accuracy. Notably, TOP-D delivers a 25.84% absolute improvement in avg@32 accuracy on AIME24 relative to OPD, and clear superiority on all evaluated tasks.

Figure 3: TOP-D outperforms OPD and RLVR in avg@32 accuracy across AIME benchmarks with the Qwen3-8B-Base student.

Learning curves from ablation studies demonstrate that both the external proximal teacher and internal trust region iterations are critical. Removing either causes optimization instability or loss of sample efficiency. TOP-D is robust to the choice of $\alpha$, requiring minimal hyperparameter tuning.

Figure 4: Isolating the external proximal teacher ($\alpha=1$) or disabling trust region iterations degrades sample efficiency and optimization stability during training.

Practical and Theoretical Implications

TOP-D demonstrates that embedding trust region concepts into policy distillation addresses core instability, balancing variance reduction with policy improvement guarantees. Empirically, TOP-D achieves these advances with zero additional computational cost and is compatible with existing distributed training infrastructures. The framework generalizes to diverse teacher-student gaps and avoids the fragility of OPD under severe capacity disparities.

Theoretically, the closed-loop mechanism of bounded variance and monotonic improvement formalizes reliable optimization dynamics for post-training large models. This has implications for scalable alignment, robustness under distributional shift, and safe reuse of off-policy data in practical deployment settings.

Future Directions

• Scaling Behavior: While the current validation focuses on student models up to 8B scale due to resource constraints, the theoretical guarantees of variance reduction and monotonic improvement suggest robust performance at much larger scales.
• Extended Training Horizons: No saturation is observed within the limited training steps (~200-400); more extensive training may yield further improvements, and the method’s sample efficiency is expected to increase with longer optimization windows.
• Generalization Beyond Mathematical Reasoning: Applying TOP-D to other reasoning domains, multimodal settings, and structured outputs will examine the breadth of its reliability.

Conclusion

Trust Region Policy Distillation (TOP-D) establishes a principled, theoretically guaranteed, and empirically robust paradigm for on-policy distillation in LLMs. By replacing direct teacher supervision with dynamically constructed proximal teachers and incorporating internal trust region optimization, TOP-D delivers substantial gains in training stability, sample efficiency, and final performance—without any increase in computational demands. Its closed-loop theoretical foundation positions TOP-D as a highly reliable framework for scalable model alignment and post-training.