β-OPSD:以策略最佳化推導,以自我蒸餾訓練
β-OPSD: Deriving with Policy Optimization, Training with Self-Distillation
July 30, 2026
作者: Jiawei Xu, Minghui Liu, Juzheng Zhang, Tom Goldstein, Furong Huang
cs.AI
摘要
在策略自蒸餾(OPSD)是一種有望改善推理語言模型的方法,但在實踐中仍然脆弱:要使其可靠運作,往往需要大量的工程努力。我們辨識出此困難的結構性來源:基本版OPSD正是更廣泛策略最佳化家族中β=1的特例,其中β加權KL懲罰,將學生模型錨定至參考策略。此等價關係將β從固定為一的隱含值轉變為可控的正則化參數,從而產生更一般的表述,可在接近參考策略與特權教師引導之間進行取捨。我們引入β-OPSD,並推導出其最佳策略為參考策略與特權教師之間的幾何插值。然而,直接用強化學習最佳化此目標將成本高昂且具有高變異。我們不直接最佳化RL目標,而是將其閉式解轉化為蒸餾目標。每個β值會沿著從參考到教師的路徑選擇一個目標,我們透過混合兩者的詞元層級logits來有效實現。如此一來,廉價的蒸餾可近似昂貴的策略最佳化之解。剩餘回報信用分配進一步使詞元更新與序列層級目標對齊,同時保留OPSD的簡潔性。在數學推理基準上的實驗顯示,β-OPSD持續優於基本版OPSD,改善了最佳化穩定性與下游推理效能。我們的結果提供了一條從自蒸餾到策略最佳化再返回的原則性途徑,且不犧牲使OPSD實用的效率。
English
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.