Trust Region Policy Distillation

Big goals are hard to achieve all at once; breaking them into small steps is wiser. We present Trust Region Policy Distillation (TOP-D), which transforms the notoriously unstable, high-variance On-Policy Distillation (OPD) into a stable training paradigm by dynamically constructing a proximal teacher. Theoretically, we establish a rigorous framework demonstrating that TOP-D inherently controls gradient variance. By providing a formal global convergence analysis alongside a monotonic improvement bound, we mathematically formalize the reliability and stability of the overall training dynamics. Empirically, TOP-D dramatically enhances training stability, sample efficiency, and final performance on mathematical reasoning tasks. More importantly, TOP-D introduces zero additional computational overhead, positioning itself as a promising alternative to the well-established OPD paradigm.

한 줄 요약

TOP-D는 기존 OPD의 불안정성을 해결하며 계산 오버헤드 없이 수학적 추론 성능을 25.84% 개선한 정책 디스틸레이션 방법이다.

핵심 기여도

외부 근사 교사(proximal teacher)를 동적으로 구성하여 기존 OPD의 무제한 로그 확률 차이 문제를 해결.
내부 신뢰 영역 반복(trust region iterations)을 도입하여 단계적 정책 개선을 보장.
AIME24 데이터셋에서 avg@32 정확도 25.84% 절대 개선을 달성.
추가적인 계산 오버헤드 없이 기존 분산 학습 인프라와 호환 가능.

핵심 아이디어

기존 OPD는 교사-학생 정책 간 로그 확률 비율로 인해 학습 불안정성이 발생한다. TOP-D는 이 문제를 해결하기 위해 **외부 근사 교사**를 동적으로 구성하여, 기존의 무제한(distillation reward) 신호를 **smooth하고 하한이 있는 신호**로 변환한다. 이는 수학적으로 **gradient variance를 엄격히 제어**함으로써 학습 안정성을 확보한다. 또한, **내부 신뢰 영역 반복**을 통해 현재 근사 교사에 더 가까운 정책을 학습함으로써 단계적 개선을 보장한다. 이 두 가지 기법은 **수학적 이론적 보장**과 **실제 학습 성능 향상**을 동시에 달성한다.

기술적 접근법

**외부 근사 교사(proximal teacher)**: 교사 정책과 학생 정책의 확률 분포를 보간하여 구성.
**내부 신뢰 영역 반복(trust region iterations)**: 단계적 정책 개선을 보장하며 샘플 효율성을 향상.
**수학적 변형**: 기존 OPD의 토큰 수준 보상에 **대수적 변형(algebraic transformation)**을 적용하여 계산 오버헤드 없이 구현 가능.
**하이퍼파라미터**: 보간 계수 α ∈ {0.1, 0.2, 0.3}에서 안정성과 성능에 미미한 차이 발생.

주요 결과

AIME24 데이터셋에서 avg@32 정확도 25.84% 절대 개선 (기존 OPD 대비).
RLVR 및 OPD 기반 베이스라인 대비 일관된 성능 우위.
α = 1.0로 설정 시 기존 OPD로 복귀되며, 학습 불안정성과 성능 저하 발생.
내부 신뢰 영역 반복 비활성화 시 샘플 효율성 30% 이상 하락.

의의 및 한계

TOP-D는 기존 OPD의 학습 불안정성과 샘플 비효율성을 이론적으로 해결하며, **수학적 이론적 보장**과 **실용적 성능 향상**을 동시에 달성한 점에서 학술적 의의가 크다. 또한, **추가 계산 오버헤드 없이 기존 인프라와 호환** 가능하다는 점에서 실용적 가치가 높다. 그러나, 보간 계수 α에 대한 민감도 분석은 TOP-D가 모든 상황에서 최적의 성능을 보장하지는 않는다는 한계를 시사한다. 또한, 다른 유형의 추론 작업이나 대규모 모델에서의 일반화 가능성은 추가 실험을 통해 검증이 필요하다.

실용적 활용

TOP-D는 대규모 언어 모델의 후처리 훈련(post-training)에 적용 가능하며, 특히 수학적 추론, 코드 생성, 복잡한 추론 작업에서 샘플 효율성과 성능을 동시에 향상시킬 수 있다. 분산 학습 인프라와 호환되므로, 산업 현장에서 대규모 모델의 정책 정렬(policy alignment) 과정에 유용하게 활용될 수 있다.