Zhengpeng Xie, Li Lyna Zhang, Zeke Xie, Mao Yang
6 min
On-Policy Distillation (OPD) is a popular method for aligning large language models, but it is notoriously unstable due to the unbounded nature of the logarithmic probability ratio between the teacher and student. This paper investigates how to stabilize these optimization dynamics and improve sample efficiency without introducing significant computational costs.
The authors introduce Trust Region Policy Distillation (TOP-D), which replaces the standard distillation objective with a stabilized version. The core innovation is the construction of an external proximal teacher—a dynamic interpolation between the target teacher and the current student policy. This construction mathematically transforms the unbounded distillation reward into a smooth, lower-bounded signal. Furthermore, the authors incorporate internal trust region iterations, which allow for off-policy data reuse, effectively decoupling the behavior policy from the target policy to improve sample efficiency.
TOP-D provides a rigorous theoretical framework that guarantees bounded gradient variance and monotonic policy improvement. Empirically, the method demonstrates significant performance gains on mathematical reasoning tasks. For instance, using a Qwen3-8B-Base student, TOP-D achieved a 50.42% accuracy on the AIME24 benchmark, representing a 25.84% absolute improvement over standard OPD. The authors show that TOP-D consistently outperforms existing reinforcement learning baselines (like RLVR and GRPO) across multiple model scales while maintaining zero additional computational overhead.
Standard distillation techniques often suffer from fragility that requires extensive engineering heuristics to mitigate. By providing a principled, theoretically grounded alternative that is natively compatible with existing distributed training infrastructures, TOP-D offers a robust and efficient solution for post-training alignment. Its ability to deliver substantial performance improvements without increasing training time or memory requirements makes it a highly practical tool for researchers and practitioners.
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.
Alex: So it's more of a mathematical adjustment than a new component entirely.
Sam: Exactly. Think of it like adjusting the volume knob on a speaker rather than building a new sound system. The signal is the same — you're just preventing it from ever going to eleven.
Alex: You also mentioned something called token-level advantage normalisation. That sounds very specific. What does it mean in practice?
Sam: Most training systems look at an entire response and assign it one overall score — like giving a whole essay a single letter grade. The problem is the student has no idea which sentences were strong and which were weak. Token-level advantage normalisation changes that. It looks at each individual word — or "token" — and asks: compared to the average quality of everything else the model produced for this prompt, was this particular word a good choice or a poor one? It gives the model a much more detailed map of its own performance.
Alex: So instead of one broad grade, the model gets specific feedback on every word it chose.
Sam: Right. And that granularity matters. It also helps with response length. By tracking the value of each word relative to the whole sequence, the system learns to balance immediate word choices with the longer-term goal of completing the task well — which naturally discourages responses that are either too short to be useful or unnecessarily padded.
Alex: There's also something called an importance sampling ratio in the paper. What role does that play?
Sam: As the student model improves through training, its behaviour changes. The importance sampling ratio is a way of tracking how much the model's current choices differ from the choices it was making earlier. It acts as a safety check. If the model has changed significantly, we scale down the size of the update we apply — because applying a large correction based on outdated information could undo progress. It keeps changes small and grounded in what the model is actually doing right now.
Alex: So it stops the model from overcorrecting based on stale data.
Sam: Precisely. Together with the proximal teacher, this creates what the researchers call a "trust region" — a boundary that limits how much the model's behaviour can shift in any single training step. Small, safe steps rather than large, unpredictable lurches.
Alex: And all of this adds up to what the paper calls a monotonic improvement bound?
Sam: That's the theoretical centrepiece. A monotonic improvement bound is a mathematical guarantee that, as long as the training follows these rules, the model's performance will trend consistently upward. It won't randomly collapse or oscillate. That kind of guarantee is relatively rare in this field — most training processes involve a degree of hoping for the best.
Alex: Does the guarantee hold up in practice? Does the model actually perform better?
Sam: The paper reports meaningful gains on mathematical reasoning tasks, which are considered a reliable test of a model's ability to think through multi-step problems. Performance improved compared to standard distillation methods, and the training process was substantially more stable. The two outcomes — stability and accuracy — appear to reinforce each other. A model that trains without crashing has more opportunity to actually learn.
Alex: That feels like the key insight, really. Stability isn't just a nice property — it's what makes consistent improvement possible in the first place.
Sam: That's a good way to put it. The authors frame TOP-D not as a replacement for existing methods, but as a principled correction to a known weakness. By building reliability into the mathematics of the training process itself, rather than hoping the system stays on track, they've produced something that is both theoretically grounded and straightforward to apply.
Alex: A more stable foundation for teaching these models — and one that doesn't cost anything extra to use. Thanks for walking us through it, Sam, and thanks to everyone listening to ResearchPod.