🤖 machine learning

OPD+: Rethinking the Advantage Design for On-Policy Distillation

This paper introduces OPD+, a corrected on-policy distillation framework that mathematically proves the bias caused by standard stop-gradient operations in advantage estimation and proposes a generic f-divergence-based optimization method that improves performance on reasoning and tool-use benchmarks.

Original authors: Hanyang Zhao, Haoxian Chen, Han Lin, Genta Indra Winata, David Yao, Wenpin Tang

Published 2026-06-02

📖 4 min read☕ Coffee break read

CC BY 4.0

Original authors: Hanyang Zhao, Haoxian Chen, Han Lin, Genta Indra Winata, David Yao, Wenpin Tang

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are trying to teach a young apprentice (the Student Model) how to solve complex puzzles by watching a master chef (the Teacher Model). The apprentice watches the master cook, tries to copy the moves, and gets feedback on how close they came to the master's recipe. This process is called On-Policy Distillation (OPD).

For a long time, researchers believed there was one specific "recipe" for giving feedback that worked best: a method called Reverse KL. It was like saying, "The only way to learn is to measure the difference exactly this way."

However, the authors of this paper, OPD+, discovered a hidden flaw in how this feedback was being calculated. They found that the way the feedback was being delivered was actually broken, leading to confusion and poor learning, especially when trying different types of feedback recipes.

Here is the breakdown of their discovery using simple analogies:

1. The Broken Feedback Loop (The "Stop Gradient" Problem)

In the old method, when the apprentice tried to learn, the teacher would say, "Here is the score of your move," but then immediately freeze that score so it couldn't change. They did this to keep things stable, like a teacher pointing at a static chart and saying, "This is the target; don't worry about how the chart was drawn."

The paper argues that this is mathematically wrong.

The Analogy: Imagine you are learning to shoot arrows. The teacher says, "Your arrow landed 5 inches to the left." But then, the teacher says, "Don't worry about how I calculated that 5 inches; just treat it as a fixed fact."
The Problem: The "5 inches" actually depends on how you are holding the bow (the student's current state). If you change your grip, the distance changes. By freezing the calculation, the teacher is giving you a lie. You are trying to adjust your grip based on a number that refuses to update with your new grip. This leads to biased estimates—you think you are improving, but you are actually moving in the wrong direction.

2. The New Solution: OPD+

The authors propose OPD+, which is essentially "fixing the math" so the feedback actually updates as the student learns.

The Fix: Instead of freezing the score, OPD+ says, "Here is the score, and here is exactly how that score changes if you move your hand." It adds a small correction term to the feedback.
The Result: It's like the teacher now saying, "Your arrow landed 5 inches left. But remember, if you tighten your grip, that distance will shrink by 1 inch." This tiny correction makes the learning process honest and accurate.

3. The Surprise: Other Recipes Work Too!

For years, everyone thought Reverse KL was the only good way to measure the difference between the student and the teacher. They thought other methods (like Forward KL or JSD) were useless and would cause the student to "collapse" (stop learning entirely).

The paper shows that these other methods weren't actually bad; they were just misused because of the broken feedback loop (the "Stop Gradient" issue).

The Analogy: It's like everyone thought a specific type of hammer was the only tool that could drive a nail. They tried using a screwdriver, and it broke the wood, so they concluded screwdrivers were useless.
The Discovery: The authors fixed the way the screwdriver was held (the math). Suddenly, the screwdriver worked perfectly! In fact, in some cases (like the JSD method), the screwdriver worked better than the hammer, allowing the student to solve harder math problems and use tools more effectively than before.

4. What They Tested

They tested this on two very difficult tasks:

Math Reasoning: Solving high-level competition math problems (like AIME and HMMT).
Tool Use: Teaching the AI to use external tools (like calculators or search engines) to solve problems.

The Outcome:

The old method (with the broken math) caused some learning styles to fail completely (0% accuracy).
The new method (OPD+) fixed these failures.
Even for the "standard" method (Reverse KL), the new math made the student learn faster and reach a higher peak of performance.

Summary

The paper claims that the way we currently teach AI models to copy each other has a fundamental math error. By fixing a single line of code to stop "freezing" the feedback, we can:

Make the learning process mathematically honest.
Unlock the potential of different learning strategies that were previously thought to be broken.
Get better results on hard tasks like math and tool usage, even with the same models we already have.

In short: Don't freeze the feedback. Let the math update, and the AI learns better.

Technical Summary: OPD+: Rethinking the Advantage Design for On-Policy Distillation

1. Problem Statement

On-policy distillation (OPD) is a technique used to transfer capabilities from capable teacher language models to base student models by optimizing a dense, token-level supervision signal derived from student-generated rollouts. While existing OPD implementations (e.g., Tinker) often utilize the Kullback-Leibler (KL) divergence as the primary objective, they employ a specific design choice: stopping the gradient on the reward term (the log-ratio difference between student and teacher likelihoods).

This design is motivated by the desire for numerical stability, drawing parallels to target networks in Deep Q-Networks (DQN). However, the authors identify a critical theoretical flaw in this approach:

Biased Gradient Estimates: For general $f$ -divergences, the reward function depends on the student policy parameters $\theta$ . Stopping the gradient on this reward term ignores the direct dependency of the reward on $\theta$ , leading to mathematically incorrect gradient estimates.
Performance Degradation: This bias causes catastrophic training collapse for divergence functions other than Reverse KL (e.g., Forward KL and Jensen-Shannon Divergence), rendering them untrainable under standard OPD settings. Even for Reverse KL, while the bias theoretically vanishes due to the score identity, empirical results suggest that the stop-gradient design still suboptimalizes performance.

2. Methodology: The OPD+ Framework

The paper proposes OPD+, a corrected optimization framework that re-evaluates the advantage design for on-policy distillation based on $f$ -divergence.

2.1 Theoretical Foundation

The authors formulate the OPD objective as minimizing the $f$ -divergence between the student policy $p_\theta$ and the teacher policy $q$ . The objective function is defined as:
$J_f(\theta) = \mathbb{E}_{x \sim X, y \sim p_\theta(\cdot|x)} \left[ \sum_{n=1}^{L_y} f(u_{\theta,n}) \right]$
where $u_{\theta,n} = q(y_n | y_{<n}, x) / p_\theta(y_n | y_{<n}, x)$ is the token-level density ratio, and $f$ is a convex function with $f(1)=0$ .

2.2 The Correction

The paper derives the true gradient of this objective. Standard OPD uses a "stop-gradient" policy gradient estimator where the reward $r = -f(u_{\theta,n})$ is treated as fixed. The authors show that the true gradient contains two terms:

The standard score-function term: $\mathbb{E}[r \nabla_\theta \log p_\theta]$ .
A direct reward-gradient term: $\mathbb{E}[\nabla_\theta r]$ .

For a general $f$ -divergence, the direct reward-gradient term does not vanish. Specifically, $\nabla_\theta r = u_{\theta,n} f'(u_{\theta,n}) \nabla_\theta \log p_\theta$ .
Consequently, the correct advantage function is not simply $-f(u)$ , but a weighted term:
$w_f(u) = -f(u) + u f'(u)$

2.3 Special Case: Reverse KL

For Reverse KL, $f(u) = -\log u$ . The correction term becomes:
$w(u) = -(-\log u) + u(-1/u) = \log u + 1$
The additional $+1$ term is a constant baseline. Since $\mathbb{E}[\nabla_\theta \log p_\theta] = 0$ (score identity), this term theoretically vanishes in expectation. However, the authors argue that in practice, this constant baseline acts as a variance-reduction mechanism, improving convergence speed and stability.

2.4 Implementation

OPD+ requires only a single-line code change in existing OPD implementations: replacing the raw negative divergence reward $-f(u)$ with the corrected weight $w_f(u)$ . This framework supports various $f$ -divergences, including Forward KL, Reverse KL, and Jensen-Shannon Divergence (JSD).

3. Key Contributions

Generic Optimization Framework: The paper provides a unified framework for on-policy distillation based on general $f$ -divergences, deriving the correct gradient and loss functions that account for reward dependency on policy parameters.
Theoretical Proof of Bias: It mathematically proves that the standard stop-gradient operation leads to biased gradient estimates for general $f$ -divergences, invalidating the common practice of treating the reward as fixed.
OPD+ Algorithm: The authors introduce OPD+, the corrected algorithm that eliminates the training collapse observed in Forward KL and JSD and improves performance even for Reverse KL.
Empirical Validation: The work challenges the prevailing belief that Reverse KL is the superior or only viable distance metric for on-policy distillation, demonstrating that other divergences can outperform it when optimized correctly.

4. Experimental Results

The authors evaluated OPD+ on the QWEN3 model family (1.7B and 8B parameters) across two distinct domains: Competition Mathematics (DeepMath, AIME 24/25, HMMT 25) and Agentic Tool Use.

4.1 Rescue of "Untrainable" Divergences

Under the standard stop-gradient design (OPD), Forward KL and JSD exhibited catastrophic failure, dropping to 0% accuracy on tool-use tasks and failing to match the SFT-initialized baseline on math tasks.

OPD+ Correction: By applying the corrected advantage $w(u)$ , OPD+ completely eliminated this failure mode.
JSD Performance: JSD+ not only stabilized but surpassed the standard Reverse KL baseline. On the challenging AIME 25 benchmark, JSD+ achieved 50.42% accuracy compared to Reverse KL's 48.54%. On the Tool Use task, JSD+ reached 60.29%, outperforming the Reverse KL baseline of 59.97%.

4.2 Improvement on Reverse KL

Even within the domain where Reverse KL is traditionally dominant, OPD+ yielded strict empirical improvements:

AIME 24: Reverse KL+ improved from 64.27% to 64.69%.
Tool Use: Reverse KL+ improved from 59.97% to 61.44%.
Convergence: Training dynamics (Figure 1) show that OPD+ variants accelerate early convergence and prevent performance plateaus or degradation seen in standard baselines.

4.3 Self-Distillation (OPSD)

The method was also validated in an On-Policy Self-Distillation (OPSD) setting, where the model acts as both student and teacher (with the teacher conditioned on privileged ground-truth solutions). OPD+ consistently outperformed the vanilla stop-gradient versions across all divergence types, with the corrected Reverse KL showing the best performance among divergences.

5. Significance and Claims

The paper claims that the historical underperformance of alternative $f$ -divergences (Forward KL, JSD) in on-policy distillation is not an intrinsic limitation of the divergences themselves, but rather a side-effect of a mis-specified gradient coefficient caused by the stop-gradient design.

By correcting the advantage function to be mathematically faithful to the $f$ -divergence objective, OPD+:

Expands the Design Space: It enables the use of a wider variety of statistical distances for distillation, moving beyond the industry standard of Reverse KL.
Improves Performance: It establishes new state-of-the-art performance peaks on reasoning and tool-use benchmarks.
Challenges Conventions: It questions the common belief that Reverse KL is inherently superior due to stability, suggesting that the stability was an artifact of the optimization method rather than the divergence choice.

The authors conclude that their work provides a more mathematically correct and empirically effective algorithm, potentially triggering further research into the design of advantage functions for on-policy distillation methods. They acknowledge limitations regarding model scaling analysis and the comparison between single-token likelihood and logit distillation, noting these as areas for future work.