$β$-OPSD: Deriving with Policy Optimization, Training with Self-Distillation

2026-07-30

AI Summary

This work addresses the optimization instability and heavy reliance on extensive hyperparameter tuning that plague conventional On-Policy Self-Distillation (OPSD) when enhancing reasoning language models. The authors embed OPSD within a policy optimization framework regularized by a KL divergence term with tunable weight β, revealing standard OPSD as a special case where β = 1. They propose a generalized formulation with adjustable β and derive an optimal policy via geometric interpolation. Their approach employs a closed-form distillation target constructed from token-level logits that blend those of a mixture teacher and a reference policy, and further incorporates a Return-to-go credit assignment mechanism. This unified framework seamlessly integrates policy optimization with efficient distillation, achieving substantially improved training stability and downstream reasoning performance over the original OPSD on mathematical reasoning benchmarks.

Abstract

On-policy self-distillation (OPSD) is a promising approach to improve reasoning language models, but it remains brittle in practice: making it work reliably often requires substantial engineering effort. We identify a structural source of this difficulty: vanilla OPSD is precisely the $β=1$ member of a broader policy-optimization family, where $β$ weights the KL penalty anchoring the student to a reference policy. This equivalence turns $β$ from an implicit value fixed at one into a controllable regularization parameter, yielding a more general formulation that trades off proximity to a reference policy against privileged teacher guidance. We introduce $β$-OPSD and derive its optimal policy as a geometric interpolation between the reference policy and the privileged teacher. Directly optimizing this objective with reinforcement learning, however, would be costly and high-variance. Rather than optimize the RL objective directly, we turn its closed-form solution into a distillation target. Each value of $β$ selects a target along the reference-to-teacher path, which we implement efficiently by mixing their token-level logits. In this way, inexpensive distillation approximates the solution of expensive policy optimization. Return-to-go credit assignment further aligns token updates with the sequence-level objective while retaining the simplicity of OPSD. Experiments on mathematical reasoning benchmarks show that $β$-OPSD consistently outperforms vanilla OPSD, improving optimization stability and downstream reasoning performance. Our results provide a principled route from self-distillation to policy optimization and back without sacrificing the efficiency that makes OPSD practical.

Problem

• on-policy self-distillation
• policy optimization
• KL regularization
• reasoning language models
• distillation stability

Innovation

• policy optimization
• self-distillation
• KL regularization
• geometric interpolation
• return-to-go credit assignment