Technical Summary: PowerOPD

Problem Statement

Standard On-Policy Distillation (OPD) for Large Language Models (LLMs) aims to minimize the reverse-KL divergence between a teacher and a student policy using trajectories generated by the student itself. To avoid the computational expense of full-vocabulary computation, modern implementations estimate the objective using a single-sample Monte Carlo estimator based on student-sampled tokens.

The authors identify that this "vanilla OPD" approach suffers from severe training pathologies in practice, including:

1. Sample Inefficiency: Validation accuracy often drops initially and fails to recover for hundreds of steps despite dense token-level supervision.
2. Unstable Generation Dynamics: Response lengths exhibit large oscillations before stabilizing, indicating the student repeatedly enters unstable generation regimes.
3. Performance Gap: Vanilla OPD converges to a policy significantly weaker than full-vocabulary OPD (e.g., an 8.19-point accuracy gap on MATH-500).

Root Cause Diagnosis:
The paper traces these failures to the log-ratio reward ($r_t = \log \frac{\pi_T(o_t|c_t)}{\pi_\theta(o_t|c_t)}$). This reward is unbounded by construction:

• High Variance: The log-ratio can reach extreme negative values (near -50), allowing single rare tokens to dominate gradient updates.
• Early-Position Extremes: Massive reward magnitudes disproportionately affect the early positions of student rollouts, destabilizing the prefix distribution and causing cascading errors.
• Persistence: These extreme values do not decay during training, injecting instability throughout the optimization process.
• Failure of Post-Hoc Fixes: Standard reinforcement learning stabilization techniques (clipping, tanh compression, Z-score normalization) fail because they operate after the unbounded log-ratio has already been computed, leaving the underlying distortion intact.

Methodology: PowerOPD

To address the unbounded nature of the reward, the authors propose PowerOPD, which reframes the reward design as learning a principled probability-to-reward mapping rather than stabilizing a pathological log-ratio post-hoc.

Core Design Principles

The authors define two necessary properties for a robust OPD reward:

1. Boundedness (P1): The reward function must be bounded on $[0, 1]^2$ to prevent rare Monte Carlo events from inducing catastrophic gradient updates.
2. Sign Consistency (P2): The reward sign must correctly align with the teacher-student probability gap (positive if the teacher assigns higher probability, negative otherwise).

The Box-Cox Transformation

The authors retain the "transform-then-subtract" structure of the standard reward ($h(\pi_T) - h(\pi_\theta)$) but replace the unbounded logarithmic function $h(x) = \log x$ with the Box-Cox power transformation:
$$h_\alpha(x) = \frac{x^\alpha - 1}{\alpha}, \quad \alpha > 0$$

This yields the PowerOPD reward:
$$r^\alpha_t = \text{sg}[\pi_T(o_t | c_t)^\alpha - \pi_\theta(o_t | c_t)^\alpha]$$
where $\text{sg}$ denotes stop-gradient.

• Boundedness: For any $\alpha > 0$, the reward is strictly bounded within $[-1, 1]$.
• Sign Consistency: Since $x^\alpha$ is strictly increasing for $\alpha > 0$, the sign of the reward correctly reflects the probability gap.
• Degenerate Limit: The standard log-ratio reward is recovered as the degenerate limit $\alpha \to 0$.

Mechanism of Improvement

The exponent $\alpha$ controls the reward sensitivity profile:

• Smaller $\alpha$: Retains higher sensitivity to low-probability tokens, resembling the log-ratio's emphasis on rare events but remaining bounded.
• Larger $\alpha$: Attenuates differences for low-probability tokens. It suppresses rewards for tokens that both models assign low probability (unreliable measurements) and focuses learning on tokens where at least one model assigns substantial probability. This filters out "high-leverage" but statistically unreliable policy-gradient terms.

Key Contributions

1. Diagnosis of OPD Pathologies: The paper provides a rigorous empirical and theoretical diagnosis showing that vanilla OPD's training instability stems from the unbounded log-ratio reward, not a lack of post-hoc stabilization.
2. PowerOPD Framework: The introduction of a family of natively bounded, sign-consistent rewards derived from the Box-Cox transformation, parameterized by $\alpha$.
3. Theoretical Insight: Demonstrating that the log-ratio is a degenerate limit of a broader class of bounded transformations and that stability must be enforced at the probability-to-reward mapping level, not after the fact.
4. Efficiency Gains: Showing that a bounded sampled-token reward can outperform full-vocabulary distillation while avoiding vocabulary-wide computation.

Experimental Results

The authors evaluated PowerOPD on six mathematical reasoning benchmarks (GSM8K, MATH-500, AMC23, AIME24/25, Minerva, Olympiad-Bench) using four Qwen3 teacher-student pairs (0.6B/1.7B students; 4B/8B teachers).

Performance Gains:

• vs. Vanilla OPD: PowerOPD achieved benchmark-averaged gains of up to +6.37 (Avg@8) and +5.71 (Pass@8). On specific benchmarks like AMC23, gains reached +16.75 Avg@8.
• vs. Post-Hoc Stabilization: Outperformed the best post-hoc baselines (Clip, Tanh, Z-Score) by +3.01 (Avg@8) and +3.54 (Pass@8) on average.
• vs. Full-Vocabulary OPD: Surpassed full-vocabulary OPD by +2.59 (Avg@8) and +8.90 (Pass@8) on average, despite using only sampled tokens.

Efficiency:

• Wall-Clock Time: Reduced by 59.2% compared to full-vocabulary OPD.
• Peak GPU Memory: Reduced by 23.1%.
• Training Dynamics: PowerOPD eliminated the initial accuracy drop and response length oscillations seen in vanilla OPD. Gradient norms were kept >3,000× smaller than the initial spike of vanilla OPD.

Scaling with $\alpha$:

• Larger $\alpha$ generally improved accuracy, shortened response lengths (reducing length inflation common in full-vocabulary OPD), and further stabilized training dynamics.

Significance and Claims

The paper claims that PowerOPD resolves the fundamental optimization difficulties of on-policy distillation by addressing the root cause: the unbounded nature of the log-ratio reward. By enforcing boundedness and sign consistency at the probability-to-reward mapping level, PowerOPD enables:

1. Stable Training: Elimination of pathological training dynamics (accuracy drops, length oscillations).
2. Superior Performance: Achieving higher accuracy than both vanilla OPD and full-vocabulary OPD.
3. Computational Efficiency: Matching or exceeding the performance of full-vocabulary distillation while reducing computational costs by nearly 60%.

The authors note that their experiments are primarily conducted on mathematical reasoning benchmarks and Qwen3 models, serving as a controlled testbed for reward design. They suggest future work could validate these benefits on code generation, instruction following, and multilingual reasoning. The method is presented as a plug-and-play replacement for the reward function in standard OPD pipelines without modifying rollout generation or teacher scoring.