ReadPaper Blog
Trust Region Policy Distillation
Trust Region Policy Distillation introduces TOP-D, a plug-and-play alternative to standard On-Policy Distillation for post-training large language models. The paper addresses OPD’s unstable, high-variance optimization by replacing direct teacher supervision with a dynamically constructed proximal teacher and trust-region-style updates, yielding stronger theoretical stability guarantees and better reported performance on AIME mathematical reasoning benchmarks.
Source: Trust Region Policy Distillation
Why OPD Gets Shaky
The paper studies why On-Policy Distillation, a popular post-training paradigm for large language models, can be attractive in theory yet fragile in practice. In OPD, a student policy learns from a stronger teacher through token-level signals, with the immediate reward defined as the logarithmic probability ratio between the teacher policy and the student policy. This gives OPD dense rewards while preserving on-policy sampling, which the paper contrasts with Supervised Fine-Tuning and Reinforcement Learning with Verifiable Rewards. The authors argue that this combination should help avoid catastrophic forgetting and improve sample efficiency, but only if the optimization dynamics remain stable. Their central motivation is therefore to keep the benefits of OPD while removing the instability caused by forcing a weaker student to directly match a stronger teacher at every token.
The Hidden Problem
The paper identifies the core failure mode of standard OPD as the unbounded logarithmic probability difference between teacher and student. When the teacher assigns very low probability to a token generated by the student, the reward term log(π*/πθ) can diverge toward negative infinity, creating high-variance gradients and fragile policy updates. The authors frame this as a capacity-gap problem: the student and target teacher may disagree sharply, and direct distillation turns those disagreements into extreme optimization signals. They note that existing remedies such as teacher-student mixed sampling, reward clipping, top-p sampling, off-policy cold starts, and full-vocabulary supervision are largely empirical heuristics. The paper’s contribution is to replace these ad hoc stabilizers with a formal mechanism that bounds the reward signal and supports theoretical analysis.
Meet the Proximal Teacher
TOP-D stabilizes OPD by constructing an external proximal teacher that interpolates between the target teacher π* and the current student πθ. The proximal teacher is defined in probability space as α π* + (1 − α) πθ, where α controls how strongly the target teacher influences the update. This changes the token-level reward from log ρk to log(α ρk + 1 − α), where ρk is the teacher-to-student probability ratio for the sampled token. The paper emphasizes that this transformed reward is strictly lower-bounded by log(1 − α), preventing the negative divergence that destabilizes standard OPD. It also explains why interpolation in probability space is essential: interpolation in log-probability space would merely scale the original OPD reward and would not remove the unboundedness.
Trust Region Steps
After defining the proximal teacher, the paper adds internal trust region iterations so the student can approximate this safer target while improving sample efficiency. Instead of discarding trajectories after every policy update, TOP-D decouples the behavior policy πθold from the target policy πθ and uses a clipped importance-ratio objective inspired by trust region reinforcement learning. The algorithm samples groups of responses for each prompt, computes TOP-D rewards using the algebraic transformation log(α π*/πθold + 1 − α), forms token-level returns, and normalizes token-level advantages across responses. The authors claim a theoretical closed loop: the external proximal teacher bounds gradient variance, the internal trust region iterations support monotonic policy improvement, and the convergence analysis controls the asymptotic optimization gap. A practical implication is that the proximal teacher need not be explicitly instantiated, so TOP-D reduces to a simple transformation of the OPD reward and introduces zero additional computational overhead.
Did It Work?
The paper evaluates TOP-D on mathematical reasoning tasks using Qwen3-8B-Base as the student-side base model and Qwen3-30B-A3B-Instruct-2507 as the teacher model. On AIME24, the reported avg@32 accuracy rises from 24.58 for OPD to 50.42 for TOP-D, while RLVR reaches 30.10 and the base model reaches 9.38. On AIME25, TOP-D reports 34.06 avg@32 accuracy compared with 23.33 for OPD, 22.08 for RLVR, and 8.02 for the base model. On AIME26, TOP-D reaches 44.06 compared with 25.42 for OPD, 21.67 for RLVR, and 6.15 for the base model. The authors interpret these results as evidence that bounding the distillation reward and reusing samples through trust-region updates improves training stability, sample efficiency, and final reasoning performance without adding computational cost.
