본문으로 건너뛰기

[논문리뷰] Trust-Region Behavior Blending for On-Policy Distillation

링크: 논문 PDF로 바로 열기

메타데이터

저자: Daniil Plyusov, Alexey Gorbatovski, Alexey Malakhov, Nikita Balagansky, Boris Shaposhnikov, Daria Korotyshova, Daniil Gavrilov


1. Key Terms & Definitions (핵심 용어 및 정의)

  • On-policy Distillation (OPD): 학생 모델이 스스로 생성한 trajectory(prefix)를 기반으로 학습하고, 교사 모델의 예측값을 매칭하여 지식을 증류하는 학습 기법입니다.
  • Trust-Region Behavior Blending (TRB): 학생 모델의 초기 학습 불안정성을 해결하기 위해, 학생 모델 중심의 KL trust region 내에서 교사 모델과 가장 가까운 행동 정책(behavior policy)을 생성하여 샘플링하는 방법론입니다.
  • Reverse-KL OPD Loss: 생성형 모델의 분포 학습에 적합한 Reverse-KL divergence를 기반으로 한 목적 함수로, 본 논문에서는 행동 정책만 변경하고 이 목적 함수는 고정하여 사용합니다.
  • Annealed Warmup: 학습 초기에는 교사 모델의 비중을 높이고, 학습이 진행됨에 따라 교사 모델에 대한 의존도(KL budget)를 점진적으로 줄여 최종적으로 학생 모델의 순수 샘플링으로 회귀하는 스케줄링 전략입니다.

2. Motivation & Problem Statement (연구 배경 및 문제 정의)

본 논문은 OPD 초기 단계에서 발생하는 학습 불안정성과 낮은 품질의 데이터 생성 문제를 해결하고자 합니다. 기존 OPD는 학생 모델이 학습 초기에 낮은 품질의 trajectory를 생성하면, 교사 모델의 지도(supervision)가 비효율적인 영역에 집중되는 한계가 있습니다 [Figure 1]. 이러한 '학생-교사 간의 간극'으로 인해 학습이 지연되거나 성능이 정체되는 문제가 발생합니다. 저자들은 기존의 강제적인 토큰 교체나 타겟 분포 변경 방식 대신, 샘플링 분포 자체를 교사 모델 쪽으로 유도하면서도 학생 모델로부터 과도하게 벗어나지 않도록 제어하는 체계적인 접근 방식의 필요성을 제시합니다.

3. Method & Key Results (제안 방법론 및 핵심 결과)

저자들은 Trust-Region behavior Blending (TRB)를 제안하여 학습 초기에 교사 모델의 유도(guidance)를 최적화합니다. 구체적으로, 매 generation 시점에서 학생 정책과 교사 정책 사이의 KL constraint를 설정하고, 이 제약 조건 내에서 교사 정책과 가장 유사한 행동 정책(behavior policy)을 Closed-Form Solution으로 도출합니다 [Figure 1]. 학습이 진행됨에 따라 KL budget($\varepsilon$)을 선형적으로 0으로 수렴시키는 Annealed Warmup 기법을 적용하여, 안정적인 초기 학습을 보장하는 동시에 최종적으로는 순수 학생 모델의 분포로 수렴하도록 설계했습니다. 실험 결과, TRBQwen3-1.7B-BaseQwen3-0.6B-Base 설정 모두에서 Vanilla OPD 및 다양한 기존 교사 유도 기법(Veto, SKD 등) 대비 가장 우수한 pass@1 평균 성능을 기록했습니다 [Table 1]. 특히, 고정된 KL budget을 사용하는 Fixed-$\varepsilon$ blending보다 Annealed Warmup을 적용한 TRB가 일관되게 높은 성능을 보여, 학습 단계별로 적절한 유도 전략이 중요함을 입증했습니다 [Figure 2].

4. Conclusion & Impact (결론 및 시사점)

본 논문은 OPD의 학습 초기 단계에서 발생하는 분포 불일치 문제를 Trust-Region 기반의 행동 정책 제어를 통해 효과적으로 완화했습니다. TRB는 모델의 아키텍처 변경 없이 학습 초기 샘플링 과정을 교사 모델 쪽으로 정교하게 유도함으로써, OPD의 효율성과 최종 성능을 동시에 개선했습니다. 이 연구는 대규모 언어 모델의 지식 증류 과정에서 데이터 생성 품질과 학습의 안정성을 조절하는 핵심적인 방법론을 제시하며, 향후 더 복잡한 추론 태스크 및 고성능 학생 모델의 효율적 정렬(Alignment) 연구에 중요한 토대가 될 것으로 기대됩니다.


Part 2: 중요 Figure 정보

Figure 1: TRB 프레임워크 개요

Figure 1 — TRB 프레임워크 개요

Figure 2: 학습 초기 비교 그래프

Figure 2 — 학습 초기 비교 그래프

Table 1: 성능 비교 벤치마크

Table 1 — 성능 비교 벤치마크

⚠️ 알림: 이 리뷰는 AI로 작성되었습니다.

댓글

관련 포스트

Review 의 다른글