PaperlayerView exact revision on arXiv

Source-linked AI summary

Trust Region Policy Distillation

Zhengpeng Xie, Li Lyna Zhang, Zeke Xie, Mao Yang

arXiv:2607.04751v1cs.LGcs.AI

TL;DR

Standard On-Policy Distillation is fragile and inefficient because teacher–student disagreements destabilize optimization and prevent effective data reuse. TOP-D constructs a proximal teacher and applies internal trust-region updates, improving mathematical reasoning while adding no computational overhead.

  • Problem

    Standard OPD faces optimization fragility from teacher–student disagreements and strictly on-policy training limits its sample efficiency.

  • Method

    TOP-D interpolates teacher and student distributions to construct a proximal teacher, smooths the distillation reward, and uses internal trust-region iterations.

  • Results

    25.84% absolute improvement in avg@32 accuracy over standard OPD on AIME24, with consistent dominance over RLVR and OPD baselines.

  • Takeaways & Limitations

    TOP-D provides a reliable, sample-efficient distillation paradigm with zero additional computational overhead across mathematical reasoning benchmarks.

  • Takeaways & Limitations

    The gradient-variance analysis assumes a finite vocabulary and uniformly bounded student-policy score-function gradients.

Abstract

from arXiv · show

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.

1 Introduction

Standard OPD is optimization-fragile because teacher–student capacity gaps make its logarithmic probability-ratio signals vulnerable to disagreement. TOP-D addresses this by constructing a proximal teacher that smooths rewards, bounds gradient variance, preserves compatibility and overhead, and improves mathematical reasoning performance.

  • Motivation: Standard OPD is vulnerable to optimization fragility because teacher–student capacity gaps amplify disagreements in its logarithmic probability-ratio objective.The paper identifies this capacity gap and disagreement sensitivity as the fundamental bottleneck of standard OPD.
  • Method: TOP-D dynamically interpolates target-teacher and student distributions to construct a proximal teacher that transforms unbounded distillation rewards into smooth, strictly controlled signals.This construction is designed to systematically resolve standard OPD’s inherent optimization instability.
  • Efficiency: TOP-D adds zero computational overhead because its practical objective is an algebraic transformation of the standard token-level OPD reward.The method is presented as a plug-and-play module compatible with distributed training infrastructures and high-throughput inference engines.
  • Method: TOP-D smooths unbounded reward signals and safely breaks standard OPD’s strict on-policy data-reuse barrier.This positions the method as a plug-and-play distillation paradigm for more flexible training.
  • Theory: TOP-D inherently bounds gradient variance, while convergence analysis and a monotonic improvement bound jointly support stable training and reduced asymptotic error.The theoretical framework forms a closed-loop analysis of training stability and convergence.
  • Results: 25.84% absolute improvement in avg@32 accuracy over standard OPD on AIME24 was achieved by TOP-D, which consistently dominated RLVR and OPD baselines.The empirical evaluation focuses on mathematical reasoning tasks.

2 Preliminaries

The paper formulates autoregressive language modeling and OPD through token-level distillation rewards, then recasts the process as an undiscounted deterministic MDP for theoretical analysis. This setup exposes length bias and high variance in cumulative rewards while defining the value and advantage quantities used for policy-improvement guarantees.

  • Autoregressive modeling: An autoregressive language model πθ generates variable-length responses y token by token from prompts x sampled from Dx.The response is represented as y = (y0, y1, . . . , yT), with T = |y|.
  • On-policy distillation: OPD minimizes a reverse KL divergence to a typically stronger model π∗, yielding a standard policy-gradient objective.The supplied preliminaries identify π∗ as the stronger reference model but do not display the full divergence equation.
  • Reward variance: Using cumulative rewards k=t rk introduces inherent length bias and high variance, motivating the practical use of immediate rewards rk.The passage presents immediate-reward use as a common practical approach.
  • MDP formulation: For theoretical proofs, autoregressive generation is modeled as an undiscounted deterministic MDP with prompt-prefix states, vocabulary actions, and an absorbing terminal state.The MDP uses γ = 1 and defines Qπ, Vπ, Aπ, and dπ to support policy-improvement analysis.
  • Distillation reward: The token-level distillation signal is log π∗(a|s) / π(a|s), interpreted as the immediate reward for each generated action.This reward is also defined at the absorbing state as zero to prevent infinite accumulation after termination.

3 Methodology

TOP-D addresses OPD’s optimization fragility and low sample efficiency through an external proximal teacher and internal trust region iterations. The methodology also uses token-level normalized advantages and combines these components into the complete TOP-D algorithm.

  • Core design: TOP-D constructs a proximal teacher to stabilize training, then uses trust region iterations for the student to approximate it.The methodology presents practical implementation after these two core components.
  • External proximal teacher: Near-zero teacher probabilities can make the OPD reward diverge to negative infinity, destabilizing policy-gradient optimization.The reward is defined through the probability ratio between teacher and student policies.
  • External proximal teacher: The proximal-teacher reward ˜rk is strictly bounded from below as the probability ratio ρk approaches zero, unlike the original OPD reward.When α = 1, the proposed reward reduces to the original OPD reward.
  • Internal trust region iterations: Internal trust region iterations decouple behavior and target policies so previously generated trajectories can be reused instead of discarded after each update.Standard OPD’s strictly on-policy formulation causes notoriously low sample efficiency because it lacks data reuse.
  • Advantage normalization: TOP-D applies token-level advantage normalization across responses for each prompt to better exploit dense reward signals and avoid overly short or excessively long responses.The token-level reward combines immediate reward with a length-normalized future return.
  • Complete algorithm: The external proximal teacher and internal trust region iterations are synthesized into the complete TOP-D framework summarized in Algorithm 1.Together, they address OPD’s optimization fragility and low sample efficiency through a unified design.

4 Theoretical Analysis

The theoretical analysis establishes that TOP-D stabilizes optimization by uniformly bounding gradient variance through its proximal teacher and interpolation coefficient α. It also proves global convergence and monotonic internal trust-region improvement, showing how reduced single-step error closes the asymptotic convergence gap.

  • Gradient Variance Control: TOP-D’s gradient estimator has uniformly bounded variance under a bounded score-function assumption, preventing divergence across the probability space.The result applies for any interpolation coefficient α ∈ (0, 1).
  • Gradient Variance Control: The interpolation coefficient α directly controls variance: variance vanishes as α → 0+, while the bound diverges as α → 1−, recovering standard OPD.TOP-D uses a hard safety bound for heavy penalties and linear dampening for high rewards.
  • Global Convergence: TOP-D’s proximal-teacher iteration admits a global convergence bound that incorporates practical optimization errors at each global step.The operator is T(π) = απ* + (1 − α)π, with updates πk+1 = T(πk) + ϵk.
  • Global Convergence: The asymptotic convergence gap is bounded by ϵ∞/α, so surpassing the precision ceiling while keeping α small requires minimizing the single-step error ∥ϵk∥1.The bound also provides an error-forgetting mechanism that isolates early high-variance updates.
  • Monotonic Improvement: Optimizing the monotonic-improvement lower bound produces decreasing reverse-KL divergences and guarantees continuous reduction of the single-step error through internal trust-region iterations.If the final expected divergence satisfies the stated δ condition, Pinsker’s and Jensen’s inequalities give ∥ϵk∥1 ≤ δ, closing the convergence gap.

5 Experiments

Experiments evaluate TOP-D on mathematical reasoning benchmarks across multiple student-teacher scales and compare it with post-training baselines. TOP-D consistently outperforms the evaluated baselines, while ablations show that its external proximal teacher and internal trust-region iterations are important for stable training.

  • Datasets, models, and baselines: TOP-D uses DAPO-Math-17k for training and evaluates post-trained models on the AIME and AMC mathematical reasoning competitions.The experiments use Qwen3-1.7B-Base and Qwen3-8B-Base students, with Qwen3-14B and Qwen3-30B-A3B-Instruct-2507 teacher policies.
  • Baseline comparison: TOP-D consistently outperforms evaluated baselines across mathematical benchmarks and student model scales.The comparison includes GRPO, DAPO, and standard OPD.
  • Baseline comparison: 50.42% avg@32 accuracy on AIME24 is achieved by TOP-D with Qwen3-8B-Base and the Qwen3-30B-A3B-Instruct-2507 teacher.This is a 25.84% absolute improvement over standard OPD, whose accuracy is 24.58%, and TOP-D exceeds DAPO by 17.5% absolute accuracy on the same benchmark.
  • Baseline comparison: For Qwen3-1.7B-Base, standard OPD suffers severe degradation, whereas TOP-D pushes AIME24 accuracy past 20% and maintains a substantial advantage.The passage attributes OPD’s fragility to the unbounded logarithmic probability ratio under a severe teacher-student capacity gap.
  • Ablation analysis: Setting α = 1.0 removes the external proximal teacher and causes highly unstable training dynamics with significant performance degradation.Disabling internal trust-region iterations forces a strictly on-policy regime and leads to a drastic drop in sample efficiency.

6 Conclusion

TOP-D is introduced as a principled post-training paradigm that addresses OPD’s optimization instability and sample inefficiency. It bounds gradient variance with an external proximal teacher and supports monotonic improvement through safe off-policy data reuse.

  • Conclusion: TOP-D addresses the optimization instability and sample inefficiency of standard On-Policy Distillation (OPD).The paper presents it as a principled post-training paradigm.
  • Conclusion: TOP-D bounds gradient variance through an external proximal teacher.This is the theoretical mechanism used to stabilize training.
  • Conclusion: TOP-D guarantees monotonic improvement through safe off-policy data reuse within internal trust region iterations.The guarantee is tied to the use of internal trust region iterations.
  • Conclusion: TOP-D achieves massive performance gains on mathematical reasoning benchmarks.The supplied passage reports this empirical outcome without specifying numerical results.

A More Results · Computational Resources

The paper reports broader benchmark comparisons for Qwen3-1.7B-Base and describes TOP-D’s experimental computational setup and resource-efficiency motivation.

  • A More Results: Table 4 compares Qwen3-1.7B-Base performance across various mathematical benchmarks.The provided passage identifies the model and scope of the comparison but does not include the benchmark values.
  • Computational Resources: Primary experiments included TOP-D training and reproduction of baseline methods.These experiments were conducted on a high-performance computing cluster.
  • Computational Resources: The computing cluster was equipped with NVIDIA H200 GPUs.The passage describes this hardware as supporting all primary experiments.
  • Computational Resources: The standard distributed training setup used 4 compute nodes.This setup was intended to accelerate the research iteration cycle.
  • Computational Resources: Each compute node contained 8 H200 GPUs.The node-level configuration is specified as part of the standard distributed training setup.
  • Computational Resources: The standard distributed setup amounted to 32 GPUs in total.This total follows directly from the stated 4-node, 8-GPU-per-node configuration.
  • Computational Resources: The authors emphasize that TOP-D is highly resource-efficient and accessible.The passage presents this as a property of the proposed method despite the cluster-based experiments.

Limitations · B Hyperparameters · C Proofs

The paper’s current validation is limited by computational and time budgets: experiments reach only student models up to 8B parameters and roughly 200–400 update steps. Table 5 provides detailed hyperparameters, while the supplied passages contain no proof details.

  • Limitations: Validation currently covers student models only up to the 8B-parameter scale.The authors attribute this boundary to resource limitations and have not investigated scaling beyond it.
  • Limitations: The study has not yet examined TOP-D’s scaling behavior on massive-scale foundation models.The passage specifically identifies scaling the student model beyond the current range as unexplored.
  • C Proofs: The supplied passages do not include substantive content from the proofs section.No proof statement, derivation, or theorem detail is provided in the input.
  • Limitations: No empirical performance saturation was observed within the evaluated optimization window.The authors therefore suggest that the reported results remain below their potential upper bounds.
  • B Hyperparameters: Table 5 presents detailed hyperparameters for the different methods.The supplied passage identifies the table but does not provide its individual hyperparameter values.

C.1 Proof of Theorem 4.2

Under Assumption 4.1 and any α ∈(0, 1), the TOP-D gradient estimator has variance bounded by a uniform absolute limit across the entire probability space. The proof establishes this by uniformly bounding the squared reward terms over the finite vocabulary and applying the variance inequality.

  • Theorem statement: Theorem 4.2 states that TOP-D gradient-estimator variance does not diverge and remains strictly bounded for every α ∈(0, 1).The bound is uniform across the entire probability space.
  • Gradient variance: Applying Var(g) ≤E to the TOP-D estimator and taking total expectation yields an absolutely bounded exact gradient variance.The conclusion uses Assumption 4.1 and the finiteness of M, |V|, α, and C∗.
  • Positive-reward region: For q ≥ p, the positive-reward term is strictly bounded by h(p, q) ≤C∗α, where C∗ is a universal mathematical constant.The bound follows from maximizing the reduced univariate function at the unique stationary point.
  • Conditional expectation: Because the vocabulary is finite, uniformly bounded per-token terms imply an absolutely bounded conditional expectation regardless of context.The finite vocabulary size is denoted |V|.

C.2 Proof of Theorem 4.4

The proof bounds the student’s expected L1 distance to the optimal teacher under iterative proximal updates with optimization error. It derives a contractive recurrence, unrolls it into a finite-step global bound, and obtains an asymptotic error bound.

  • Theorem setup: Theorem 4.4 analyzes πk+1 = T(πk) + ϵk for any initial student policy π0, with ϵk representing practical optimization error.The proof targets the expected L1 distance between πk+1 and the optimal target teacher policy π∗.
  • Single-step recurrence: The proximal teacher contracts the student-to-target distance by factor 1 − α, while the update contributes the step-k optimization error through the triangle inequality.The operator is T(πk) = απ∗ + (1 − α)πk, yielding d(T(πk), π∗) = (1 − α)d(πk, π∗).
  • Finite-step bound: Repeated substitution of the single-step recurrence produces a geometric pattern and a closed-form finite-step global bound for step k+1.The recurrence is unrolled forward from k = 0 by repeatedly substituting earlier bounds.
  • Asymptotic convergence: Under α ∈ (0, 1) and bounded asymptotic optimization error lim supk→∞∥ϵk∥1 ≤ ϵ∞, the initial-distance term vanishes and the limit supremum is bounded by ϵ∞/α.The asymptotic argument splits the accumulated error sum, bounds its tail geometrically, and lets δ → 0.

C.3 Proof of Theorem 4.9

Theorem 4.9 establishes a lower bound on the target policy’s return using a surrogate objective anchored to the base policy’s state distribution. The proof bounds the approximation error through state-distribution shift and local expected advantage drift, yielding a total-variation penalty.

  • Theorem 4.9: Theorem 4.9 lower-bounds η(˜π) by ζπ(˜π) minus a penalty proportional to ξ, Tmax, ℓπ, and expected statewise total variation.The surrogate is ζπ(˜π) = η(π) + ℓπ E[Aπ(s, a)], while ξ is the maximum absolute local expected advantage drift.
  • Objective approximation error: The proof isolates the absolute difference between the true return η(˜π) and surrogate ζπ(˜π), then decouples global state-distribution shift from local advantage drift.Hölder’s Inequality provides the decoupling, with ξ capturing the local term.
  • State-distribution shift: The initial-state marginal distributions coincide exactly because prompt-induced initialization is policy-independent, giving |P(s1 = s | ˜π) − P(s1 = s | π)| = 0.This removes the first generation step from the state-distribution shift accumulation.
  • State-distribution shift: For later steps, a telescoping sequence of hybrid trajectory distributions reduces each adjacent difference to expected L1 divergence between the policies at one branching decision.The reduction relies on deterministic autoregressive transitions and the triangle inequality.
  • Conclusion: Substituting the measure-shift bound into the approximation error yields the stated lower bound on η(˜π).The absolute objective error is bounded by 2ξTmaxℓπ E[DTV(˜π(· | s), π(· | s))].
Loading 2607.04751v1…