SR-OPSD: Self-Referenced On-Policy Self-Distillation

2026-08-10Machine Learning

Machine LearningArtificial Intelligence
AI summary

The authors study a way to improve learning in language models by giving them more detailed feedback for each token they generate during training, instead of just sparse rewards. They point out problems with previous methods that use a moving target for learning, which can be unstable. To fix this, they propose SR-OPSD, a new approach that carefully balances how the model adjusts itself toward a mix of its own past behavior and a reference policy, using a mathematical tool called Rényi divergence. Their experiments show that this method works well across various complex tasks like math reasoning and coding. Overall, the authors offer a more stable and effective way for models to learn from their own outputs.

On-policy self-distillationReinforcement learningToken-level supervisionSelf-teacher policyRényi divergenceProjection geometryVariational characterizationDistributional concentrationLarge language modelsCoding generation
Authors
Zhuo Sun, Entong Li, Yanlong Zhao, Xiaoyuan Cheng, Wenxuan Yuan, Kaiyu Li, Che Liu, Huihang Liu, Harrison Bo Hua Zhu, Li Zeng
Abstract
On-policy self-distillation (OPSD) converts feedback into dense token-level supervision on trajectories generated by the policy to be optimized, providing a useful complement to reinforcement learning with sparse outcome rewards. However, the self-teacher policy used in OPSD is typically a stop-gradient or exponential-moving-average copy of the policy conditioned on additional context information, and thus co-evolves with both the student policy and its on-policy context distribution. Directly matching such a moving target with a fixed projection objective can lead to unstable optimization or excessive distributional concentration. This nature of OPSD motivates the proposed \emph{Self-Referenced On-Policy Self-Distillation (SR-OPSD)}. At fixed student-generated contexts, a token-level variational characterization identifies the effective distillation target as a geometric interpolation between the self-teacher policy and a reference policy. Meanwhile, we use the Rényi divergence family to generalize the projection geometry. This formulation separates \emph{where} the adaptive target is placed from \emph{how} the student is projected toward it: the interpolation coefficient controls underlying target, while the Rényi order controls the projection geometry and its sensitivity to token-level density ratios. Extensive experiments across scientific evaluation, mathematical reasoning, and coding generation tasks with multiple large language models show that SR-OPSD achieves the state-of-the-art or competitive performance across various settings.