Trust Region Policy Distillation

Q: What is the main contribution of TOP-D?

TOP-D introduces a proximal teacher mechanism that interpolates between the target teacher and the current student distribution in probability space, transforming the volatile log-ratio reward of standard On-Policy Distillation (OPD) into a strictly lower-bounded signal and enabling internal trust-region iterations for stable, sample-efficient training.

Q: What problem does TOP-D address?

TOP-D addresses the fragility of standard On-Policy Distillation (OPD) for LLMs, where the logarithmic probability ratio between teacher and student can diverge to negative infinity when the policies disagree, causing gradient variance to explode and training to become unstable.

Q: Why does standard OPD produce unstable training?

In standard OPD, a near-zero teacher probability yields an unbounded negative reward because the log-ratio can become arbitrarily large, causing gradient variance to explode and making optimization fragile.

Q: How does the proximal teacher work?

The proximal teacher is a surrogate target distribution that interpolates between the target teacher and the current student in probability space, guaranteeing a hard lower bound of log(1−α) on the reward regardless of how small the teacher probability becomes, which prevents unbounded negative rewards.

Q: How does TOP-D differ from simply adding a constant offset to the OPD reward?

Adding a constant shifts all rewards uniformly but does not bound the reward when the teacher probability approaches zero; TOP-D's interpolation reshapes the probability distribution itself, guaranteeing a hard lower bound of log(1−α) regardless of how small the teacher probability becomes.

Q: What are the internal trust-region iterations in TOP-D and why are they used?

Internal trust-region iterations decouple the behavior policy from the target policy, allowing the student to learn from off-policy data and breaking the strict on-policy data-reuse barrier of standard distillation, which maximizes sample efficiency and ensures monotonic improvement while minimizing single-step optimization error.

Q: What theoretical guarantees does TOP-D provide?

The paper proves two core guarantees: a bound on the variance of the gradient estimator (with a max term capturing two regimes—a safety regime and a damping regime) and a convergence result showing that iterative policy updates under the trust-region operator guarantee monotonic improvement of the true objective.

Q: What are the key empirical results of TOP-D?

TOP-D achieves a 25.84% absolute improvement in accuracy on mathematical reasoning benchmarks compared to standard OPD, and consistently outperforms OPD across all six evaluated math benchmarks.

Q: What benchmarks and datasets were used to evaluate TOP-D?

TOP-D was evaluated on six mathematical reasoning benchmarks; the paper does not specify the individual benchmark names beyond describing them as math benchmarks, and comparisons are made against standard OPD as well as GRPO and DAPO baselines.

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

Trust Region Policy Distillation (TOP-D) stabilizes LLM distillation by bounding reward variance and enabling off-policy reuse.

How can we stabilize on-policy distillation for language models by replacing direct teacher supervision with a dynamically constructed proximal teacher?

Standard On-Policy Distillation (OPD) for LLMs is notoriously fragile because the logarithmic probability ratio between teacher and student can diverge to negative infinity, causing gradient variance to explode. TOP-D resolves this by dynamically constructing a proximal teacher that interpolates between the target teacher and the current student, effectively smoothing the distillation reward into a strictly lower-bounded signal. This stabilization allows for internal trust-region iterations that decouple behavior and target policies, delivering a 25.84% absolute improvement in accuracy on mathematical reasoning benchmarks compared to standard OPD.

Paper Primer

The core mechanism of TOP-D is the proximal teacher: a surrogate target distribution that prevents the reward signal from becoming unbounded. By interpolating the teacher and student probability distributions, the method transforms the volatile log-ratio reward into a smooth, bounded signal that acts as a variance controller.

To maximize sample efficiency, TOP-D incorporates internal trust-region iterations that allow the model to learn from off-policy data. This approach breaks the strict on-policy data-reuse barrier of standard distillation, ensuring that the student policy improves monotonically while minimizing the single-step optimization error.

TOP-D significantly outperforms standard OPD and RLVR baselines on complex mathematical reasoning tasks.

Evaluation on the AIME24 benchmark using a Qwen3-8B-Base student model. 25.84% absolute improvement in avg@32 accuracy over standard OPD.

The proximal teacher construction provides a rigorous theoretical bound on gradient variance.

Theorem 4.2 establishes that the variance of the TOP-D gradient estimator is strictly bounded by a uniform absolute limit, preventing the divergence seen in standard OPD. Variance is bounded by a function of the interpolation coefficient $\alpha$, where $\alpha \to 0$ causes variance to vanish.

Why is the proximal teacher constructed in probability space rather than log-probability space?

Interpolating in probability space yields a smooth, lower-bounded reward signal that prevents the gradient variance explosion inherent to standard OPD, whereas log-probability interpolation merely scales the original, unbounded reward.

Does this method introduce significant computational overhead during training?

No. The proximal teacher is a conceptual construction that reduces to a simple algebraic transformation of the standard token-level reward, making TOP-D a plug-and-play module with zero additional computational overhead.

Introduction and Motivation

We expose OPD’s instability and propose TOP‑D to stabilize distillation via a proximal teacher.

Standard On‑Policy Distillation (OPD) forces a student to match the teacher’s output distribution directly, which makes the gradient variance explode because the logarithmic probability ratio can become arbitrarily large when the policies disagree. This high‑variance signal leads to unstable training and fragile optimization. TOP‑D mitigates the problem by dynamically constructing a proximal teacher $\tilde{\pi}_{\text{teacher}}^{(t)}$ that interpolates the target teacher with the current student $\pi_{\text{student}}^{(t)}$, thereby keeping the reward signal smooth, bounded, and computationally free.

**Figure 1.** While standard OPD forces direct teacher supervision and suffers from unstable optimization, TOP-D ensures stable training by dynamically constructing a proximal teacher at every iteration.

The core problem is that direct teacher supervision in OPD yields unstable, high‑variance gradients.

On-Policy Distillation Background

Defines on‑policy distillation and formalizes the RL view of autoregressive language modeling.

OPD trains a student language model to imitate a stronger teacher by matching the teacher’s token‑level distribution on the same prompts.

Generating a sequence can be seen as a deterministic MDP: each generated token is an action, the state holds the prompt and the already generated prefix, and the reward measures how much the student’s token probability improves toward the teacher.

**Table 1.** Comparison of different paradigms.

The TOP-D Algorithm

TOP‑D builds a proximal teacher and trust‑region updates to stabilize distillation.

Standard OPD forces the student to match the teacher’s raw distribution, so a near‑zero teacher probability yields an unbounded negative reward and makes the gradient explode.

Instead of demanding exact imitation, we blend the teacher with the current student to create a “nearby” target that never pushes the student into a fatal loss.

Compute the raw ratio $\rho = 0.01/0.4 = 0.025$.

Blend: $\alpha\rho + 1-\alpha = 0.5\times0.025 + 0.5 = 0.5125$.

Take the log: $\tilde{r}= \log(0.5125)\approx -0.669$.

Compare to the raw OPD reward $\log\rho = \log(0.025)\approx -3.688$, which is far more negative.

The proximal teacher caps the penalty at $\log(1-\alpha) = \log(0.5)\approx -0.693$, preventing catastrophic gradients.

How does this differ from simply adding a constant offset to the OPD reward?

Adding a constant would shift all rewards uniformly but would not bound the reward when $\rho\to0$. The interpolation reshapes the probability distribution itself, guaranteeing a hard lower bound $\log(1-\alpha)$ regardless of how small $\pi^{*}$ becomes.

We keep a stale “behavior” policy for data collection while repeatedly improving a fresh “target” policy inside a trust region, so each update reuses the same trajectories.

Clip $p$ to $1+\epsilon=1.2$, so the clipped term becomes $1.2\times0.5=0.6$.

Take the minimum with the unclipped advantage $0.5$, yielding $0.5$ for this token.

Repeat for all tokens, sum, divide by $|y|=3$, then average over the two prompts.

The resulting surrogate loss is a modest positive value, indicating a safe improvement direction.

Even though the raw ratio $p=1.3$ suggests a large policy shift, clipping forces the update to stay within the trusted $20\%$ band.

Why not use the PPO objective directly on the sequence‑level return?

Sequence‑level returns are sparse and hide token‑wise signal. By normalizing advantages per token, TOP‑D exploits the dense reward $\tilde{r}_k$, yielding finer‑grained gradients while still benefiting from PPO’s trust‑region clipping.

After computing raw token rewards, we standardize them within each prompt group so that advantages reflect relative quality rather than absolute scale.

Compute $\tilde{R}_4 = \tilde{r}_4 = 0.3$ (no future tokens).

Compute $\tilde{R}_3 = \tilde{r}_3 + \frac{1}{1}\tilde{r}_4 = 0.0 + 0.3 = 0.3$.

Compute $\tilde{R}_2 = -0.1 + \frac{1}{2}(\tilde{r}_3+\tilde{r}_4) = -0.1 + 0.15 = 0.05$.

Compute $\tilde{R}_1 = 0.2 + \frac{1}{3}(\tilde{r}_2+\tilde{r}_3+\tilde{r}_4) = 0.2 + 0.067 \approx 0.267$.

Collect all $\tilde{R}$ values from the three responses, compute $\mu \approx 0.2$, $\sigma \approx 0.08$, then normalize each to obtain $\hat{A}$.

Normalization compresses the spread of returns, ensuring that no single long response dominates the gradient.

**Figure 3.** Reward curves of $r_k$ and $\tilde{r}_k$.

**Figure 4.** In contrast to sequence-level advantage normalization, we perform token-level normalization across the responses generated for a given prompt to better leverage the dense reward signal.

TOP‑D training loop (Algorithm 1)

Theoretical Foundations

We prove TOP‑D’s variance control, global convergence, and monotonic improvement.

Standard OPD suffers from exploding gradient variance, which makes training unstable. TOP‑D introduces a proximal teacher to keep variance in check.

TOP‑D scales the teacher’s reward by an interpolation coefficient $\alpha$, turning the raw variance into a controllable quantity that never exceeds a fixed ceiling.

Compute the logarithmic term: $\log(1-0.3)=\log(0.7)\approx -0.357$, square gives $0.127$.

Compute the linear term: $C^{*}\alpha = 1\times0.3 = 0.3$.

Take the max: $\max\{0.127,0.3\}=0.3$.

Apply the bound: $M^{2}|V|\times0.3 = 2^{2}\times5\times0.3 = 4\times5\times0.3 = 6$.

The bound tells us that, with these modest settings, the TOP‑D gradient variance cannot exceed 6, a concrete guarantee absent in vanilla OPD.

Why isn’t the variance bound simply $M^{2}|V|$ without the max term?

Because TOP‑D’s reward scaling has two regimes: a safety regime (log term) that dominates when $\alpha$ is near 1, and a damping regime (linear term) that dominates for smaller $\alpha$. The max captures the worst‑case of the two.

Having secured per‑step variance, we turn to the overall learning dynamics of the proximal teacher.

For any initial policy $\pi_{0}\in\Pi$, the iterates $\pi_{k+1}=T(\pi_{k})+\varepsilon_{k}$ satisfy $d(\pi_{k+1},\pi^{*}) \le (1-\alpha)^{k+1} d(\pi_{0},\pi^{*}) + \sum_{i=0}^{k} (1-\alpha)^{k-i}\,\|\varepsilon_{i}\|_{1}$, and if $\limsup_{k\to\infty}\|\varepsilon_{k}\|_{1}\le\varepsilon_{\infty}$ then $\limsup_{k\to\infty} d(\pi_{k+1},\pi^{*}) \le \varepsilon_{\infty}/\alpha$.

This convergence result reveals that the remaining bottleneck is the single‑step error $\|\varepsilon_{k}\|_{1}$, motivating internal trust‑region iterations.

We now prove that TOP‑D’s internal trust‑region updates guarantee monotonic improvement of the true objective.

Because language models generate sequences of varying length, we measure a policy’s average output length to normalize visitation counts.

For any policies $\tilde{\pi}$ and $\pi$, $\eta(\tilde{\pi}) \ge \zeta_{\pi}(\tilde{\pi}) - \frac{2\xi $T_{\max}$}{\ell_{\pi}} \, \mathbb{E}_{s\sim $d_{\pi}$}\!\bigl[ $D_{\mathrm{TV}$}\bigl(\tilde{\pi}(\cdot\mid s),\pi(\cdot\mid s)\bigr) \bigr]$, where $\zeta_{\pi}(\tilde{\pi}) = \eta(\pi) + \ell_{\pi}\,\mathbb{E}_{s\sim d_{\pi}}\!\bigl[\,\mathbb{E}_{a\sim \tilde{\pi}(\cdot\mid s)}[A_{\pi}(s,a)]\,\bigr]$ and $\xi = \max_{s}\bigl| \mathbb{E}_{a\sim \tilde{\pi}(\cdot\mid s)}[A_{\pi}(s,a)] \bigr|$.

Empirical Results

TOP‑D delivers large, stable gains on math benchmarks, confirming the trust‑region distillation idea.

Recall that standard on‑policy distillation (OPD) directly matches teacher outputs, which leads to high‑variance gradients. TOP‑D replaces this with a proximal teacher that keeps updates inside a trust region, stabilising training.

TOP‑D yields a large absolute gain over standard OPD on the hardest benchmark.

On AIME24 with the Qwen3‑8B‑Base student, TOP‑D reaches 50.42 % avg@32 accuracy versus 24.58 % for OPD, a +25.84 % improvement.

**Table 3.** Performance comparison of the Qwen3-1.7B-Base across AIME benchmarks.

**Figure 5.** Learning curves for the ablation study on the Qwen3-1.7B-Base. We isolate the effects of the external proximal teacher ($\alpha = 1.0$) and the internal trust region iterations (w/o off-policy).

**Table 4.** Performance comparison of the Qwen3-1.7B-Base across various mathematical benchmarks.

TOP‑D consistently outperforms OPD across all evaluated math benchmarks.

Extended Benchmarks

Additional benchmark gains and resource analysis for TOP-D on the Qwen3-1.7B-Base.

TOP‑D improves over OPD by up to +11.35 points on the AIME24 benchmark.

Table 4 shows OPD 8.96 vs TOP‑D 20.31, a gain of +11.35.

Across the six mathematical benchmarks, TOP‑D consistently outperforms the standard OPD baseline, delivering larger scores on every task.

All primary experiments were run on a high‑performance cluster with NVIDIA H200 GPUs; the standard setup used 4 nodes × 8 GPUs (32 GPUs total), but the entire TOP‑D pipeline can be reproduced on a single 8‑GPU node.

Current validations are limited to student models up to 8 B parameters and to short training windows of roughly 200–400 update steps, leaving open the question of scaling TOP‑D to larger foundations and longer horizons.

Theoretical Proofs I

We prove variance boundedness of TOP‑D gradients and convergence of the policy updates.

This section establishes the two core theoretical guarantees for TOP‑D. First we bound the variance of the gradient estimator, then we show that the iterative policy updates converge under the trust‑region operator.

Variance is bounded by the second moment.

It suffices to bound $\mathbb{E}\|\tilde g_k\|^{2}$.

Bound the second moment of the TOP‑D reward.

Case 1 ($q<p$) yields $h(p,q)<(\log(1-\alpha))^{2}$.

Case 2 ($q\ge p$) yields $h(p,q)\le p(\log\alpha+1-\alpha)^{2}$.

Optimizing $p(\log\alpha+1-\alpha)^{2}$ over $p\in(0,1)$ yields $C^{*}\alpha$.

Combine the two cases to bound $h(p,q)$.

Sum over the finite vocabulary.

Apply Assumption 4.1 to the gradient term.

Final variance bound.

Single‑step recurrence from the update rule.

Proximal teacher contracts towards $\pi^{*}$.

Recurrence inequality.

Unroll the recurrence.

Take the limit superior.

Theoretical Proofs II

We prove TOP‑D’s global convergence bound and the surrogate‑objective guarantee of Theorem 4.9.

Apply the recurrence for k=1 to bound $d(\pi_2,\pi^{*})$.

Iterate the substitution to obtain a closed‑form finite‑step bound.

Take the limit supremum to derive the asymptotic convergence bound.

Express the performance difference between the true return $\eta(\tilde\pi)$ and the surrogate $\zeta_{\pi}(\tilde\pi)$.

Bound the state‑distribution difference $|d_{\tilde\pi}(s)-d_{\pi}(s)|$.

Combine the two bounds to obtain the final surrogate‑objective guarantee.

Implementation Details

Lists the hyperparameter settings for each baseline and our TOP‑D method.

Table 5 enumerates the hyperparameter configuration used for GRPO, DAPO, OPD, and our TOP‑D method.

Questions & answers

What is the main contribution of TOP-D?

TOP-D introduces a proximal teacher mechanism that interpolates between the target teacher and the current student distribution in probability space, transforming the volatile log-ratio reward of standard On-Policy Distillation (OPD) into a strictly lower-bounded signal and enabling internal trust-region iterations for stable, sample-efficient training.

What problem does TOP-D address?

TOP-D addresses the fragility of standard On-Policy Distillation (OPD) for LLMs, where the logarithmic probability ratio between teacher and student can diverge to negative infinity when the policies disagree, causing gradient variance to explode and training to become unstable.

Why does standard OPD produce unstable training?

In standard OPD, a near-zero teacher probability yields an unbounded negative reward because the log-ratio can become arbitrarily large, causing gradient variance to explode and making optimization fragile.

How does the proximal teacher work?

The proximal teacher is a surrogate target distribution that interpolates between the target teacher and the current student in probability space, guaranteeing a hard lower bound of log(1−α) on the reward regardless of how small the teacher probability becomes, which prevents unbounded negative rewards.

Why is the proximal teacher constructed in probability space rather than log-probability space?

Interpolating in probability space yields a smooth, lower-bounded reward signal that prevents gradient variance explosion, whereas interpolating in log-probability space merely scales the original unbounded reward without providing a hard lower bound.

How does TOP-D differ from simply adding a constant offset to the OPD reward?

Adding a constant shifts all rewards uniformly but does not bound the reward when the teacher probability approaches zero; TOP-D's interpolation reshapes the probability distribution itself, guaranteeing a hard lower bound of log(1−α) regardless of how small the teacher probability becomes.

What are the internal trust-region iterations in TOP-D and why are they used?

Internal trust-region iterations decouple the behavior policy from the target policy, allowing the student to learn from off-policy data and breaking the strict on-policy data-reuse barrier of standard distillation, which maximizes sample efficiency and ensures monotonic improvement while minimizing single-step optimization error.

What theoretical guarantees does TOP-D provide?

The paper proves two core guarantees: a bound on the variance of the gradient estimator (with a max term capturing two regimes—a safety regime and a damping regime) and a convergence result showing that iterative policy updates under the trust-region operator guarantee monotonic improvement of the true objective.

What are the key empirical results of TOP-D?

TOP-D achieves a 25.84% absolute improvement in accuracy on mathematical reasoning benchmarks compared to standard OPD, and consistently outperforms OPD across all six evaluated math benchmarks.

What benchmarks and datasets were used to evaluate TOP-D?

TOP-D was evaluated on six mathematical reasoning benchmarks; the paper does not specify the individual benchmark names beyond describing them as math benchmarks, and comparisons are made against standard OPD as well as GRPO and DAPO baselines.

Does TOP-D introduce significant computational overhead?

Why does TOP-D use token-level advantages rather than sequence-level returns?

Sequence-level returns are sparse and hide token-wise signal; by normalizing advantages per token, TOP-D exploits the dense reward signal, yielding finer-grained gradients while still benefiting from PPO's trust-region clipping.

What are the limitations of TOP-D as acknowledged in the paper?

Current validations are limited to student models up to 8 billion parameters and short training windows of roughly 200–400 update steps, leaving open the question of whether TOP-D scales to larger foundation models and longer training horizons.

How does TOP-D compare to related methods such as GRPO and DAPO?

The paper includes GRPO and DAPO as baseline comparisons in its hyperparameter configuration table and empirical evaluation, and TOP-D consistently outperforms standard OPD across all six math benchmarks; the paper does not report specific numerical comparisons against GRPO and DAPO beyond this.

How can TOP-D be reproduced, and what hardware is required?

Primary experiments were run on a high-performance cluster with NVIDIA H200 GPUs using 4 nodes × 8 GPUs (32 GPUs total), but the paper states the entire TOP-D pipeline can be reproduced on a single 8-GPU node; hyperparameter configurations for GRPO, DAPO, OPD, and TOP-D are enumerated in Table 5.

Who authored TOP-D, and where was it published?

The paper is available on arXiv at https://arxiv.org/abs/2607.04751; the paper does not specify author names or a venue in the provided text.

Key terms

On-Policy Distillation (OPD): A training method where a student language model is directly supervised to match the output probability distribution of a teacher model using data generated by the student itself.
TOP-D (Trust Region Policy Distillation): The method introduced in this paper, which stabilizes on-policy distillation by replacing the raw teacher distribution with a proximal teacher and using internal trust-region iterations.
Proximal teacher: A surrogate target distribution constructed by interpolating between the target teacher and the current student in probability space, ensuring the distillation reward signal has a hard lower bound.
Log-ratio reward: The token-level reward in standard OPD computed as the logarithm of the ratio of teacher probability to student probability, which can diverge to negative infinity when the teacher assigns near-zero probability.
Gradient variance explosion: A training instability where the variance of gradient estimates becomes arbitrarily large, causing erratic parameter updates and preventing convergence.
Trust-region iteration: An optimization step that constrains policy updates to remain within a bounded region of the current policy, ensuring stable and monotonically improving learning.
Interpolation parameter α (alpha): A scalar in TOP-D that controls the degree of mixing between the teacher and student distributions when constructing the proximal teacher, with the reward lower-bounded by log(1−α).
Behavior policy: The policy used to collect training data (rollouts), which in off-policy settings may differ from the target policy being optimized.
Target policy: The policy whose parameters are being updated during training, as distinct from the behavior policy that generated the data.
PPO (Proximal Policy Optimization): A reinforcement learning algorithm that uses clipping to constrain policy updates within a trust region, which TOP-D leverages for its token-level advantage normalization.
GRPO (Group Relative Policy Optimization): A policy optimization baseline method included in the paper's comparative hyperparameter configuration and empirical evaluation.
DAPO: A policy optimization baseline method included alongside GRPO and OPD in the paper's hyperparameter configuration table and empirical comparisons.
Single-step optimization error (ε_k): The residual error introduced in a single policy update step, which the paper identifies as the remaining bottleneck motivating internal trust-region iterations.
Token-level reward: A dense reward signal assigned at each individual token position in a sequence, as opposed to a single sparse reward given for the entire generated sequence.
Monotonic improvement: A theoretical property guaranteeing that each successive policy update does not decrease the value of the true objective function.

Read the original paper

Open the simplified reader on Paperglide

Browse all simplified papers