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

왜 중요한가

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

핵심 기여

근접 교사 보정을 통해 로그 보상의 발산을 하한으로 묶는 학습 보상 변환

교사 확률 π^*와 학생 확률 π_θ를 α로 보간하여 근접 교사 π^* = απ^* + (1−α)π_θ를 구성하고 보상을 r̃_k = log(α ρ_k + 1 − α)로 재정의하였다. 이 변환은 ρ_k→0일 때 r̃_k의 하한 log(1−α)을 보장하므로 음의 극단적 보상으로 인한 분산 폭주를 방지한다. 결과적으로 보상 신호의 안정성이 확보되어 학습 과정에서 발생하는 최적화 취약점을 완화했다.

내부 trust region 반복으로 단일 스텝 최적화 오차 축소와 샘플 재사용 보장

글로벌 스텝 내에서 행동 정책(π_{θ_old})으로 생성한 그룹 응답을 토큰 수준 중요도비와 clip 연산을 사용해 여러 내부 epoch 동안 재활용하였다. 토큰 수준의 정규화된 이점(advantage)과 clip 기반 objective를 통해 내부 반복마다 기대되는 KL 다이버전스를 점진적으로 감소시켜 단일 스텝 오차 ∥ε_k∥_1을 엄격히 줄였다. 이 구조는 엄밀한 monotonic improvement 보장을 도출하여 α를 작게 유지하면서도 수렴 간극을 좁힐 수 있게 했다.

수학적 분석으로 그래디언트 분산 상한과 수렴 경계 도출

정리 4.2에서 보상 변환에 의해 토큰 수준 그래디언트 분산이 상한을 가진다는 수식을 도출하였고, 이는 Var(ĝ_k) ≤ M^2 |V| max{(log(1−α))^2, C^* α} 형태로 표현되었다. 이러한 표현은 α가 분산 조절 인자임을 명확히 하며 α→1^-일 때 OPD의 발산을 재현하고 α→0^+일 때 분산이 소산되는 동작을 수학적으로 정량화했다. 추가로 반복적 proximal 연산을 추상화해 전역 수렴 경계와 비동등식을 제공하여 알고리즘의 안정성 근거를 마련했다.

추론 비용 증가 없이 다양한 모델 규모에서 실험적으로 검증된 실효성

TOP-D는 개념적으로 근접 교사를 구성하지만 실무적으로는 토큰 보상에 대한 대수적 변환으로 구현되어 추가적인 계산 비용을 요구하지 않는다. Qwen3-8B-Base와 Qwen3-1.7B-Base 학생 모델에 대해 AIME 및 AMC 등 수학 벤치마크에서 OPD와 RLVR 계열 기법들을 상대로 일관되게 우수한 성능을 보였다. 특히 Qwen3-8B-Base에 대해 AIME24 avg@32에서 OPD 대비 +25.84%p 절대 향상을 기록하여 안정화 전략의 실효성을 실험적으로 확인했다.

핵심 아이디어 이해하기

출발점과 기존 한계: 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 수와 계산 리소스 사이의 실무적 조정이 필요하다.

실무 활용

TOP-D는 추가 연산 없이 기존 OPD 파이프라인에 플러그인 가능한 형태로 동작하므로 사후 학습(post-training) 워크플로에 직접 적용 가능하다. 학습 과정에서는 행동 정책으로 생성된 그룹 응답을 재활용하므로 샘플 효율이 중요하거나 대규모 교사 모델을 활용해 소형 학생을 향상시키려는 실무 환경에서 특히 유용하다. 공개 코드 저장소가 논문 메타데이터에 없으므로 실무 적용을 위해서는 논문에 기재된 보상 변환과 내부 반복 스킴을 재현해야 한다.

대형 교사 모델로 소형 학생을 안정적으로 증류해 배포 가능한 경량 대화 모델을 제작할 때
수학적 추론이나 긴 문맥을 요구하는 태스크에서 OPD가 불안정해 성능이 저하되는 경우의 대체 파이프라인으로
분산 학습 환경에서 오프라인 롤아웃을 재활용하여 샘플 효율을 개선하고 계산 비용을 억제해야 하는 파이프라인 최적화

코드 공개 여부: 미확인

키워드

TOP-D(Trust Region Policy Distillation)OPD(On-Policy Distillation)reverse KL(역 KL 발산)proximal teacher(근접 교사 분포)trust region(트러스트 리전)token-level normalization(토큰 수준 정규화)

용어 해설

근접 교사 분포(Proximal Teacher): — 교사 모델과 현재 학생 정책의 확률을 α로 보간하여 만든 대체 확률분포로서, 로그 확률비를 직접 사용하지 않고 αρ + 1 − α 형태의 보상을 계산해 로그 보상의 하한을 보장함으로써 극단적 음의 보상으로 인한 분산 폭주를 막는다.
중요도 샘플링(Importance Sampling): — 행동 중요도 비율 p_t = π_θ(y_t|·)/π_{θ_old}(y_t|·)을 통해 다른 정책으로 생성된 토큰을 재사용하여 손실을 추정하는 기법으로, TOP-D에서는 내부 trust region 반복에서 오프라인 롤아웃을 재활용해 샘플 효율을 높인다.
토큰 수준 이점 정규화(Token-level Advantage Normalization): — 동일 프롬프트에 대해 생성된 여러 응답의 토큰별 반환값을 그룹 단위로 평균과 표준편차로 정규화하여 토큰 수준의 이점을 계산함으로써 길이 편향을 줄이고 조밀한 보상 신호를 더 잘 활용하는 정규화 절차이다.
역 KL 발산(Reverse KL): — 학생 분포 π_θ를 교사 분포 π^*에 근사시키기 위해 사용되는 불평등성으로, OPD에서는 기대 로그비(log π^*/π_θ)를 최대화하는 형태로 학습 목표가 정의되어 교사-학생 확률 불일치가 로그 보상으로 증폭될 수 있다.
트러스트 리전(Trust Region): — 정책 업데이트의 크기를 제한하여 최적화의 급격한 변화를 방지하는 기법으로, TOP-D에서는 내부 반복에서 중요도비와 clip 연산을 통해 오프라인 데이터를 안전하게 재사용하면서 단계별 최적화 오차를 줄인다.

코드 예제

text

0: Student π_θ, teacher π^*, dataset 𝒟_x, interpolation α, group size G, internal epochs E
1: while not converged do
2:   Initialize π_{θ_old} ← π_θ
3:   Sample a batch of prompts 𝒳 ⊂ 𝒟_x
4:   for each prompt x ∈ 𝒳 do
5:     Perform rollout to sample G responses {y^i}_{i=1}^G ∼ π_{θ_old}(·|x)
6:     Compute TOP-D rewards  r̃_k^i ← log(α π^*(y_k^i|x,y_{<k}^i)/π_{θ_old}(y_k^i|x,y_{<k}^i) + 1 − α)
7:     Compute token-level returns R̃_k^i ← r̃_k^i + 1/(|y^i| − k) ∑_{j=k+1}^{|y^i|} r̃_j^i
8:     Compute token-level advantages Â_k^i ← (R̃_k^i − μ)/σ
9:   end for
10:  for epoch e = 1 to E do
11:    J(θ) ← (1/∑_i |y^i|) ∑_{i=1}^G ∑_{t=1}^{|y^i|} min[p_t^i ⋅ Â_t^i, clip(p_t^i,1−ε,1+ε) ⋅ Â_t^i]
12:    Update θ by maximizing J(θ)
13:  end for
14: end while

이 코드는 TOP-D의 학습 루프와 근접 교사 보상 계산, 토큰 수준 반환 및 내부 trust region 반복으로 정책을 업데이트하는 핵심 절차를 의사코드 형태로 요약한다.

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

왜 중요한가

핵심 기여

핵심 아이디어 이해하기

방법론

주요 결과

기술 상세

한계점

실무 활용

키워드

용어 해설

코드 예제

관련 토론

댓글

관련 기사

OPD는 정말 교사 지식을 증류하는가

서로 다른 토크나이저로 장문 증명을 증류하는 SimpleOPD

다수결로 정답을 만드는 비지도 Self-Distillation

왜 중요한가

핵심 기여

핵심 아이디어 이해하기

방법론

주요 결과

기술 상세

한계점

실무 활용

키워드

관련 토론

댓글

관련 기사

OPD는 정말 교사 지식을 증류하는가

서로 다른 토크나이저로 장문 증명을 증류하는 SimpleOPD

다수결로 정답을 만드는 비지도 Self-Distillation