ArXiv

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

Authors
Jiawei Xu, Minghui Liu, Juzheng Zhang...
Categories
cs.LG
arXiv
https://arxiv.org/abs/2607.28582v1
PDF
https://arxiv.org/pdf/2607.28582v1

Brief

β-OPSD introduces a controllable β that weights the KL anchoring term in on-policy self-distillation, exposing vanilla OPSD as the β=1 special case. The paper derives a closed-form optimal policy as a geometric interpolation between a reference policy and a privileged teacher, and approximates that RL solution efficiently by mixing token-level logits as distillation targets. Adding return-to-go credit assignment, the method reduces RL variance and yields more stable optimization and better performance on mathematical-reasoning benchmarks compared to vanilla OPSD.

Why it matters

Proposes β-OPSD: a generalization of on-policy self-distillation that introduces a scalar β to weight the KL penalty anchoring the student to a reference policy (vanilla OPSD = β=1); the authors derive a closed-form optimal policy that is a geometric interpolation between the reference policy and a privileged teacher, and implement targets by mixing token-level logits.

Key details

  • Avoids costly/high-variance RL by turning the closed-form solution into an inexpensive distillation target and adding return-to-go credit assignment to align token updates with sequence-level returns; experiments on mathematical-reasoning benchmarks (Jiawei Xu et al., arXiv:2607.28582v1, published 2026-07-30) show β-OPSD consistently outperforms vanilla OPSD, improving optimization stability and downstream reasoning performance.
Source evidence

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.