PowerOPD: Stabilizing On-Policy Distillation with Bounded Power Transformation

2026-06-15

AI Summary

This work addresses the high gradient variance in on-policy distillation (OPD) caused by unbounded log-ratio rewards, which leads to training instability, low sample efficiency, and degraded performance. To resolve this, we propose PowerOPD, the first method to integrate the Box–Cox power transformation into on-policy distillation, yielding a natively bounded and sign-consistent reward function that naturally generalizes the log-ratio reward as α → 0, thereby fundamentally suppressing gradient variance. Combining single-sample Monte Carlo estimation with a large language model distillation framework, PowerOPD achieves up to 6.37 and 5.71 absolute improvements in Avg@8 and Pass@8, respectively, across six mathematical reasoning benchmarks and four Qwen3 teacher–student model pairs, while reducing training time by 59.2%, peak GPU memory usage by 23.1%, and gradient norm by over three orders of magnitude.

Abstract

Standard on-policy distillation (OPD) for large language models estimates the reverse-KL objective using student-sampled tokens, yielding an unbiased single-sample Monte Carlo estimator that avoids vocabulary-wide computation. However, we show that this estimator suffers from severe training pathologies in practice: sample inefficiency, unstable generation dynamics, and a substantial performance gap compared to exact full-vocabulary OPD. Reward-level diagnosis traces these pathologies to the log-ratio reward, which is unbounded by construction, producing extremely high-variance gradients concentrated at early positions and persisting throughout training; standard post-hoc scaling fail as they operate only after this distortion occurs. To solve this problem, we propose PowerOPD: a family of natively bounded, sign-consistent rewards from the Box-Cox power transformation, parameterized by alpha > 0, of which the log-ratio is the degenerate alpha -> 0 limit. Across six mathematical reasoning benchmarks and four Qwen3 teacher-student pairs, PowerOPD achieves benchmark-averaged Avg@8/Pass@8 gains of up to +6.37/+5.71 over vanilla OPD, +3.01/+3.54 over post-hoc stabilization, and +2.59/+8.90 over full-vocabulary OPD, while reducing wall-clock time by 59.2% and peak GPU memory by 23.1%. Larger alpha generally improves accuracy, consistently shortens responses, and keeps gradient norms more than 3,000x smaller than vanilla OPD.

Problem

• on-policy distillation
• log-ratio reward
• gradient variance
• training instability
• sample inefficiency

Innovation

• On-Policy Distillation
• Power Transformation
• Bounded Reward
• Gradient Variance Reduction
• Large Language Model Compression