KL for a KL: Control Variate Baseline을 이용한 On-Policy Distillation

요약

본 논문은 대규모 언어 모델(LLM)의 추론 영역 훈련에 사용되는 On-Policy Distillation (OPD)의 불안정성 문제를 해결하기 위해 vOPD(Control Variate Baseline을 이용한 On-Policy Distillation)를 제안합니다. 이는 OPD를 정책 경사(policy-gradient RL)로 간주하고, 전통적인 강화 학습(RL) 기법인 Control Variate Baseline을 도입하여 훈련의 안정성을 높입니다. 특히, OPD의 값 함수가 학생 및 교사 모델 간의 토큰당 음의 역 KL 발산이라는 형태로 표현되어 추가 계산 없이도 쉽게 얻을 수 있다는 점을 강조합니다.

핵심 포인트

• On-Policy Distillation (OPD)은 LLM 훈련에서 중요한 패러다임이지만, 높은 기울기 분산으로 인해 구현이 불안정하다는 문제가 있다.
• 제안된 vOPD는 OPD를 정책 경사(policy-gradient RL)로 재해석하고 Control Variate Baseline을 적용하여 훈련의 안정성을 크게 향상시킨다.
• vOPD의 핵심은 OPD의 값 함수가 학생 모델과 교사 모델 간의 토큰당 음의 역 KL 발산이라는 닫힌 형태로 표현되어 계산 효율성이 높다는 것이다.