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

The Core Problem and Solution

Vanilla on-policy self-distillation (OPSD)—where a language model learns from its own better outputs—works in principle but is fragile in practice, requiring heavy tuning to succeed reliably. The authors identify why: standard OPSD is mathematically equivalent to a policy optimization problem where the regularization strength is locked at β=1. By treating β as a tunable parameter, they recover a principled family of algorithms that balances fidelity to a reference policy against learning from a stronger teacher model. The key insight is that the optimal policy under this formulation is a geometric interpolation—a learnable blend—between the reference and teacher policies.

How β-OPSD Works

The authors derive a closed-form solution: the student should mix the reference policy and privileged teacher along a path parameterized by β. Rather than solving this via expensive reinforcement learning, they convert the mathematical solution into a distillation target. Concretely, at the token level, they interpolate logits (the model's raw output scores before probabilities) from both policies, weighted by β. A student model then learns to match this target through standard supervised learning. This sidesteps the high variance and computational cost of policy optimization while preserving its benefits.

The method adds return-to-go credit assignment—a technique that weights each token's update by the cumulative reward from that point forward in a sequence. This keeps token-level updates aligned with the ultimate goal of better reasoning, without sacrificing the simplicity that makes OPSD efficient.

Why This Matters

Mathematical reasoning benchmarks (the paper tests on such domains) demand reliable improvement methods. Vanilla OPSD sometimes helps, sometimes hurts, and its behavior is hard to predict. β-OPSD provides practitioners with a control knob: β trades off staying close to a safe reference model against absorbing lessons from a stronger teacher. Empirically, the authors show consistent wins over vanilla OPSD in both stability and final reasoning performance.

The deeper contribution is conceptual. The authors prove that a family of self-distillation algorithms—varying β—all solve valid optimization problems. This transforms OPSD from a black-box heuristic into a point on a principled Pareto frontier. The efficient implementation (mixing logits rather than running RL) means practitioners can explore this frontier without reimplementing expensive optimization loops.

The result bridges two worlds: the theoretical cleanliness of policy optimization (which says what the optimal solution should be) and the practical efficiency of distillation (which says how to approximate it cheaply). For builders working with reasoning models, this offers a path to more robust, interpretable improvements—and a way to understand why tuning one knob (β) stabilizes training where vanilla approaches fail.