On-Policy or Off-Policy Learning? A Systematic Study of Distillation Dynamics
On-policy 학습이 항상 우월하다는 통념과 달리, KL 발산 방향과 학습률이 모델 성능 및 망각 제어에 더 결정적인 역할을 함을 규명함.
논문이 다루는 내용
기존에는 On-policy 학습이 파라미터 업데이트의 희소성을 높이고 일반화 성능을 개선한다고 알려져 왔으나, 기존 연구들은 여러 변수가 혼재되어 있어 Rollout policy의 순수 효과를 분리하기 어려웠습니다. 본 연구는 강력한 교사 모델에서 약한 학생 모델로의 증류(Distillation) 환경에서 Rollout policy, KL 방향, 학습률을 독립적으로 변화시키며 그 영향을 분석했습니다. 실험 결과, Rollout policy 자체보다는 토큰 수준의 KL 방향이 성능과 출력 범위를 결정하며, 학습률이 망각과 업데이트 희소성을 조절하는 핵심 요소임을 발견했습니다. 특히 Forward KL는 Rollout policy 변화에 강건한 반면, Reverse KL은 학생 모델의 생성 데이터(On-policy)에 더 민감하게 반응했습니다. 결론적으로 On-policy 학습의 이점은 목적 함수와 하이퍼파라미터 설정에 따라 달라질 수 있음을 보여줍니다.
핵심 결과
-
Rollout policy(On-policy vs Off-policy)가 성능에 미치는 영향은 KL 방향에 따라 상이함
-
Forward KL는 Rollout policy 변화에 매우 강건하며 안정적인 성능을 유지함
-
Learning rate는 모델의 망각(forgetting)과 파라미터 업데이트 희소성을 결정하는 핵심 변수임
실무에서 볼 만한 점
LLM 미세 조정(Fine-tuning) 시 무조건적인 On-policy 학습보다는, 사용하는 KL divergence 방향과 학습률 설정이 성능과 망각 방지에 더 중요함을 시사합니다.
읽을 때 확인할 점
-
KL 방향(Forward vs Reverse)에 따른 모델 성능 및 출력 분포 변화 비교 실험
-
학습률(Learning Rate) 변화가 기존 지식 망각(Catastrophic Forgetting)에 미치는 영향 측정
-
Rollout policy를 변경하며 KL gradient의 안정성 테스트