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

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.

한 줄 요약

β-OPSD는 정책 최적화 기반의 자가 교사 학습(SEL)을 효율적으로 근사하는 새로운 자가 교사 학습 알고리즘으로, 수학적 추론 성능을 향상시킨다.

핵심 기여도

β-OPSD는 β=1일 때 기존 OPSD와 동일하며, β를 조절하여 정책 최적화와 자가 교사 학습 간 균형을 조정할 수 있다.
β-OPSD의 최적 정책은 기하학적 보간으로 표현되며, 이를 토큰 수준 로짓 혼합을 통해 효율적으로 근사한다.
Return-to-go 크레딧 할당을 도입하여 토큰 수준 업데이트를 시퀀스 수준 목적과 정렬시킨다.
AIME 2024, AIME 2025, HMMT 2025 데이터셋에서 Qwen3-1.7B 모델에서 최대 9.16%의 성능 향상을 보인다.

핵심 아이디어

기존 OPSD는 β=1인 KL 정규화 정책 최적화의 특별한 경우로, β를 1로 고정함으로써 정책의 변화를 제어할 수 없다는 한계가 있다. β-OPSD는 β를 조절 가능한 정규화 매개변수로 도입하여, 기존 정책(참조 정책)과 교사 정책 간의 균형을 조정할 수 있도록 한다. 이는 정책 최적화의 해를 토큰 수준 로짓 혼합을 통해 효율적으로 근사하는 방식으로, 정책 최적화의 비용을 줄인다. 또한, Return-to-go 크레딧 할당을 통해 토큰 수준 업데이트가 시퀀스 수준 목적과 정렬되도록 한다. 이는 기존 OPSD의 단점을 해결하고, 더 안정적이고 효과적인 학습을 가능하게 한다.

기술적 접근법

β-OPSD는 KL 정규화 정책 최적화의 β 매개변수를 조절하여, 참조 정책과 교사 정책 간의 기하학적 보간을 정의한다.
최적 정책은 토큰 수준 로짓 혼합을 통해 효율적으로 근사하며, 이는 정책 최적화의 비용을 줄인다.
Return-to-go 크레딧 할당은 토큰 수준 업데이트가 시퀀스 수준 목적과 정렬되도록 하며, 향후 토큰의 분포 불일치를 고려한다.
학습 시 β 값은 일정한 스케줄에 따라 조절되며, 이는 학습 과정에서 부드럽게 타겟을 전환하는 효과를 준다.
하이퍼파라미터로는 할인 계수 γ=0.99, 토큰 수준 로짓 혼합 계수 η∈{0.2, 0.5, 0.8}가 사용된다.

주요 결과

Qwen3-1.7B 모델에서 β-OPSD는 AIME 2024, AIME 2025, HMMT 2025 데이터셋에서 각각 9.16%, 5.27%, 2.78%의 성능 향상을 보이며, 평균 성능은 5.74% 향상된다.
Qwen3-4B와 Qwen3-8B 모델에서도 각각 1.76%, 1.66%의 성능 향상을 기록한다.
β-OPSD는 기존 OPSD, SFT, GRPO 기반 방법보다 평균적으로 더 높은 성능을 보인다.
토큰 수준 로짓 혼합과 Return-to-go 크레딧 할당의 개별 기여도를 분리한 아블레이션 실험에서도 두 요소 모두 중요한 역할을 한다.

의의 및 한계

β-OPSD는 정책 최적화와 자가 교사 학습 간의 관계를 명확히 정의하고, 이를 기반으로 효율적인 학습 알고리즘을 제안함으로써, 기존 OPSD의 불안정성을 개선한다. 특히, β 매개변수를 조절함으로써 학습의 보수성과 진보성을 유연하게 조정할 수 있다. 또한, Return-to-go 크레딧 할당을 통해 토큰 수준 업데이트가 시퀀스 수준 목적과 정렬되도록 하여, 기존 OPSD의 단점을 보완한다. 그러나 β-OPSD는 교사-학생 분포 간의 격차가 큰 경우 성능 향상이 제한될 수 있으며, 이는 특히 대규모 모델에서 나타날 수 있다. 또한, β 스케줄링과 로짓 혼합 계수 η의 선택은 성능에 민감하게 반응하므로, 이에 대한 최적화가 필요하다.

실용적 활용

β-OPSD는 수학적 추론, 논리적 추론, 프로그래밍 등 시퀀스 기반 추론이 필요한 자연어 처리 및 인공지능 분야에서 활용 가능하다. 특히, 대규모 언어 모델의 학습 안정성과 추론 성능 향상을 필요로 하는 산업 및 연구 환경에서 유용하다. 예를 들어, 교육, 의료, 법률 분야에서 정확한 추론이 요구되는 시스템 개발에 적용할 수 있다.