Trust Region Policy Distillation (TOP-D): 보상 분산을 상한으로 묶어 OPD 최적화를 안정화한 방법

On-Policy Distillation(OPD)은 교사와 학생의 확률차를 로그 보상으로 직접 사용하기 때문에 극단적 확률 불일치에서 보상이 발산하여 최적화가 불안정해지는 실무적 병목을 낳았다. TOP-D는 교사 확률과 현재 학생 확률을 선형 보간하여 근접 교사 분포를 구성하고 로그 비율을 αρ+1−α 형태로 변환하여 음의 보상에 대한 하한을 확보함으로써 그래디언트 분산을 수학적으로 상한했다. 이 방식은 추가 연산 비용을 요구하지 않으면서 학습 안정성과 샘플 효율을 동시에 개선하여 고난도 수학 추론 벤치마크에서 의미 있는 성능 이득을 제공했다.

출발점과 기존 한계: Autoregressive LLM의 토큰 생성은 각 시점에 대해 교사와 학생의 확률분포를 비교하는 형태로 OPD가 정의되며, 표준 OPD 목적은 기대 역 KL을 최소화하는 것이다. 이 목적은 로그 확률비 r_k = log(π^*(y_k)/π_θ(y_k))를 즉시 보상으로 사용하므로 교사와 학생 사이에 용량 격차가 존재할 때 로그비가 음의 큰 값으로 발산하여 정책 그래디언트의 분산을 폭발적으로 증가시킨다. 결과적으로 최적화가 불안정해지고 특히 소형 학생 모델에서 성능 저하가 심각하게 나타난다.

해결 원리와 작동 방식: TOP-D는 근본적으로 보상 표현을 바꾸어 이 문제를 완화한다. 구체적으로 확률 공간에서 보간을 수행해 근접 교사 π^* = απ^* + (1−α)π_θ를 정의하고 보상을 r̃_k = log(α ρ_k + 1 − α)로 재정의한다. 이 변환은 ρ_k가 작아질 때 로그값이 log(1−α)로 하한을 갖게 하여 음의 극단적 보상으로 인한 그래디언트 분산 폭주를 수학적으로 차단하고, 보상 신호의 상한·하한 특성을 통해 안정적인 경사 추정이 가능해진다.

달라지는 점과 기대 효과: 보상 변환만으로도 분산을 제어할 수 있으므로 추가적인 샘플링·클리핑과 같은 실무적 휴리스틱에 의존할 필요가 줄어든다. 내부적으로는 행동 정책으로부터 생성한 그룹 응답을 중요도비와 clip 기반 목표로 여러 내부 epoch 동안 재활용해 샘플 효율을 높이며, 이 과정에서 토큰 수준 이점 정규화를 적용해 길이 편향을 제어한다. 이 결합은 최종적으로 안정된 훈련 곡선과 더 높은 최종 성능을 가능하게 하며, α를 조절함으로써 분산-수렴 간의 명확한 절충을 제어할 수 있게 한다.

전체 접근 방식과 핵심 아이디어: TOP-D는 외부 근접 교사 구성과 내부 trust region 반복의 두 축으로 설계되었다. 외부 구성은 교사와 현재 학생의 확률을 α로 보간하여 근접 교사를 개념적으로 정의하고 보상을 r̃_k = log(α ρ_k + 1 − α)로 대수적으로 변환해 보상의 하한을 확보한다. 내부 반복은 행동 정책 π_{θ_old}으로부터 생성한 그룹별 응답을 재활용하되 토큰 수준의 중요도비 p_t^i와 clip(p_t^i,1−ε,1+ε) 연산을 사용해 안전하게 업데이트를 수행한다.

핵심 메커니즘과 알고리즘 세부: 토큰별 즉시 보상 r̃_k와 길이 정규화된 미래 반환 R̃_k를 계산하고, 동일 프롬프트 내 여러 응답을 그룹으로 묶어 평균 μ와 표준편차 σ로 토큰 수준 반환을 정규화해 Â_k^i를 얻는다. 내부 objective는 PPO 스타일의 clipped surrogate를 토큰 수준으로 적용한 형태이며 J(θ)= (1/∑|y^i|) ∑ min[p_t^i⋅Â_t^i, clip(p_t^i,1−ε,1+ε)⋅Â_t^i]으로 계산되어 내부 epoch마다 θ를 갱신한다. 알고리즘 구현은 근접 교사 계산을 보상 변환으로 치환하므로 추가적인 전방/후방 계산 비용을 요구하지 않으며 분산된 트레이닝 인프라와도 호환된다.

학습 전략과 안정성 보장: 이론적으로는 근접 교사 연산을 정책 공간으로의 수축 연산으로 보고 반복적 업데이트를 분석하여 전역 수렴 경계와 단일 스텝 최적화 오차의 역할을 도출했다. 내부 반복을 충분히 수행해 KL 다이버전스를 낮추면 Pinsker 불평등을 통해 ∥ε_k∥1을 원하는 δ 수준으로 제어할 수 있고, 이로써 최종 수렴 간극 ε∞/α를 실무적으로 감소시킬 수 있다. 하이퍼파라미터로는 보간 계수 α, 그룹 크기 G, 내부 epoch 수 E, 클립 계수 ε와 토큰 정규화 방식이 핵심이며 논문은 α=0.1,0.2 범위에서의 강건성을 보고한다.

메인 벤치마크 성과: 논문은 Qwen3-8B-Base 및 Qwen3-1.7B-Base 학생 모델에 대해 AIME24/AIME25/AIME26 등 수학 추론 벤치마크에서 비교를 수행하였다. Qwen3-8B-Base를 Qwen3-30B-A3B-Instruct-2507 교사로 지도했을 때 TOP-D는 AIME24 avg@32에서 50.42%를 기록하여 표준 OPD의 24.58%보다 절대 +25.84%p의 이득을 보였다. Qwen3-1.7B-Base에서도 TOP-D는 AIME24 avg@32에서 20.31%를 기록해 OPD(8.96%) 대비 +11.35%p 향상을 달성했다.

성분별 영향과 ablation 결과: 외부 근접 교사를 α=1.0으로 설정해 제거하면 학습이 불안정해지고 성능이 크게 하락했으며, 내부 오프정책 반복을 제거하면 샘플 효율이 급격히 저하되어 수렴 속도가 느려졌다. α를 0.1, 0.2, 0.3으로 변화시킨 민감도 실험에서는 학습 안정성, 수렴 속도, 최종 성능이 해당 범위에서 소폭 차이를 보이며 전체적으로 방법이 α에 대해 비교적 강건함이 관찰되었다. 이 결과는 외부 보상 변환과 내부 trust region 반복이 상호 보완적으로 작동해야 성능 이득을 확보함을 실험적으로 확인한다.

효율성 및 연산 비용: 근접 교사 구성은 확률 보간의 대수적 표현으로 구현되어 추가적인 추론 또는 교사 샘플링 비용을 요구하지 않는다. 논문은 TOP-D가 계산적으로 플러그인 형태로 동작하므로 분산 학습 인프라와 고처리량 추론 엔진과 호환되며 실무에서의 연산 오버헤드가 없음을 강조했다. 따라서 동일한 계산 예산 내에서 안정성과 성능을 동시에 개선할 수 있는 실용적 장점이 입증되었다.

전체 아키텍처 구조: 알고리즘은 외부 근접 교사 연산과 내부 trust region 반복을 결합한 형태로, 입력 프롬프트 배치에 대해 먼저 행동 정책 π_{θ_old}로 G개의 응답을 롤아웃한다. 각 토큰에 대해 근접 교사 기반 보상 r̃_k = log(α π^*(·)/π_{θ_old}(·) + 1 − α)을 계산하고, 토큰별 길이 정규화된 반환 R̃_k을 구한 뒤 그룹 내 평균 μ와 표준편차 σ로 정규화해 Â_k를 얻는다. 내부 업데이트는 토큰 수준의 clipped surrogate objective를 최적화하는 방식으로 E 내부 epoch 동안 수행되어 정책을 π_θ로 갱신한다.

핵심 메커니즘의 수학적 기반: 근접 교사 보상 변환은 r̃_k = log(α ρ_k + 1 − α)로 주어지며 이 식은 ρ_k→0일 때 r̃_k ≥ log(1−α) 성질을 갖는다. 이를 이용해 토큰 수준 그래디언트의 두 번째 모멘트에 대해 Var(ĝ_k) ≤ M^2 |V| max{(log(1−α))^2, C^* α} 형태의 상한을 도출했고, 이 상한은 α를 통해 분산을 조절할 수 있음을 수학적으로 보장한다. 내부 trust region 반복은 클리핑과 중요도비를 통해 기대 KL을 감소시키며 Pinsker 불평등을 통해 단일 스텝 최적화 오차 ∥ε_k∥_1을 δ 수준으로 제어할 수 있음을 보였다.

Prior work 대비 차별점: 기존 수치적 휴리스틱(보상 클리핑, top-p 샘플링, 혼합 샘플링 등)은 불안정성의 원인을 로그 확률비의 발산으로 규정하지 않았고 이론적 보증이 부족했다. TOP-D는 확률 공간에서의 보간을 통해 보상의 수학적 상한을 도입하고 내부 반복을 통해 최적화 오차를 시스템적으로 축소하는 닫힌 이론 경로를 제공하여 안정성과 전역 수렴 경계 모두에 대해 엄밀한 근거를 제공한다. 또한 근접 교사 구성은 대수적 변환으로 실무 상의 계산 부담을 추가하지 않는 구현적 장점을 가진다.

구현 및 학습 세부사항: 실험에서는 그룹 크기 G=8, 글로벌 배치 512 프롬프트(4096 샘플), 미니배치 32 프롬프트를 사용했고 보간 계수 α는 0.1 또는 0.2를 주로 사용하였다. 토큰 정규화는 그룹 내 R̃_k의 평균 μ와 표준편차 σ로 표준화하였고 내부 objective는 ε에 의한 클리핑을 적용한 PPO 계열 공식으로 구성하였다. 검증 시에는 top-p=0.7, 온도 1.0을 사용했으며 추가 하이퍼파라미터는 부록에 상세히 수록되어 있다.

논문 자체에서 명시한 근본적 절충은 α의 선택과 단일 스텝 최적화 오차 ∥ε_k∥1 사이의 균형이다. 구체적으로 α를 작게 설정하면 안정성은 향상되지만 이론적으로 수렴 간극이 ε∞/α로 증가하므로 단일 스텝 오차를 엄격히 줄이지 못하면 최종 정밀도 제한이 발생할 수 있다. 또한 구현면에서는 내부 trust region 반복을 충분히 수행해야 ∥ε_k∥_1을 작은 δ로 제어할 수 있으므로 내부 epoch 수와 계산 리소스 사이의 실무적 조정이 필요하다.