[논문리뷰] On-Policy or Off-Policy Learning? A Systematic Study of Distillation Dynamics

링크: 논문 PDF로 바로 열기

메타데이터

저자: Julianna Piskorz, Antonin Berthon, Mihaela van der Schaar

---

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

• OnPD (On-Policy Distillation): 모델이 자신의 현재 정책($\pi_{S}^{\theta}$)을 사용하여 궤적(rollout)을 생성하고, 해당 궤적에서 Teacher 모델의 토큰 수준 지도를 받아 학습하는 방식입니다.
• OffPD (Off-Policy Distillation): 고정된 Teacher 모델의 정책($\pi_{T}$)을 사용하여 생성된 데이터셋 혹은 궤적을 학습 데이터로 활용하는 방식입니다.
• KL Direction (Forward vs. Reverse): Student와 Teacher 분포 간의 불일치를 측정하는 방향성입니다. Forward-KL은 Teacher의 분포를 중심으로 모드를 포괄(mode-covering)하며, Reverse-KL은 Student의 분포를 중심으로 특정 모드를 탐색(mode-seeking)합니다.
• Rollout-Policy Spectrum: $\lambda$ 파라미터를 통해 Student와 Teacher 사이의 샘플링 비율을 연속적으로 조절하여 rollout 방식을 변경하는 실험적 프레임워크입니다.

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

본 논문은 Post-training 과정에서 널리 활용되는 On-policy와 Off-policy 학습 방법론 사이의 성능 차이를 유발하는 핵심 요인을 규명합니다. 기존 연구들은 On-policy 학습이 catastrophic forgetting 감소, 파라미터 업데이트의 sparsity 향상, 일반화 능력 강화에 효과적이라고 주장해 왔습니다. 하지만 기존 비교 실험들은 학습 목표, 데이터 생성 방식, 최적화 하이퍼파라미터 등 여러 요인을 동시에 변경하기 때문에 Rollout Policy 자체의 효과를 고립시켜 분석하기 어렵습니다. 저자들은 Strong-to-Weak Distillation이라는 통제된 환경을 통해 Rollout Policy가 실제로 성능에 미치는 독립적인 영향을 검증하고자 합니다 [Figure 1].

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

본 연구는 Llama 3 및 Qwen2.5 모델 패밀리를 활용하여 Rollout Policy, KL direction, Learning rate를 독립적으로 변인 통제하며 실험을 수행합니다. 제안된 프레임워크는 Rollout Policy를 $\lambda$ 파라미터로 제어하는 연속적 스펙트럼상에서 학습을 진행하며, Forward-KL과 Reverse-KL에 따른 최적화 거동을 분석합니다 [Figure 3].

주요 정량적 결과는 다음과 같습니다:

• Performance: Rollout Policy 자체는 최종 성능에 일관된 이점을 제공하지 않으며, 오히려 KL direction이 성능과 Output Coverage를 결정하는 결정적 요인임을 확인했습니다.
• Robustness: Forward-KL은 Rollout Policy 변화에 매우 강건(robust)하여 전체 스펙트럼에서 안정적인 성능을 보인 반면, Reverse-KL은 Rollout Policy에 극도로 민감하며 Student가 선호하는 궤적에서 더 높은 성능을 보였습니다.
• Forgetting & Sparsity: Catastrophic forgetting과 파라미터 업데이트의 Sparsity는 Rollout Policy가 아닌 Learning rate에 의해 지배됩니다. 높은 Learning rate는 성능과 상관없이 공통적으로 더 높은 forgetting과 낮은 Sparsity를 야기합니다 [Figure 2].

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

본 논문은 On-policy Rollout이 절대적으로 우월하다는 통념에 의문을 제기하며, 그 가치가 목적 함수와 최적화 설정에 따라 달라짐을 입증했습니다. 연구 결과는 Forward-KL을 사용할 경우 Off-policy 전략만으로도 충분한 성능을 달성할 수 있음을 시사하며, 이는 계산 비용이 높은 On-policy 학습의 무분별한 사용을 재고하게 합니다. 향후 언어 모델의 효율적 Post-training 레시피를 설계함에 있어 Rollout Policy 보다는 KL Objective 및 최적화 하이퍼파라미터의 정밀한 튜닝이 우선시되어야 한다는 중요한 학술적 근거를 제시합니다.

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