[Paper] On-Policy Distillation 재검토: 경험적 실패 모드와 간단한 해결책

Source: arXiv - 2603.25562v1

개요

논문은 **on‑policy distillation (OPD)**를 재조명한다—학생 모델이 자체 텍스트를 생성하고 더 강력한 교사 모델로부터 피드백을 받는 방식으로 대형 언어 모델(LLM)을 미세 조정하는 기술이다. 저자들은 일반적인 “sampled‑token” 버전의 OPD가 장기 지평 과제에서 왜 붕괴되는지를 밝히고, 훈련을 보다 안정적이고 효과적으로 만드는 몇 가지 간단하고 실용적인 해결책을 제안한다.

핵심 기여

• Theoretical insight: 토큰‑level OPD가 원하는 시퀀스‑level reverse‑KL 목표의 편향된 추정량임을 보여주지만, 더 엄격한 최악‑사례 분산 경계를 갖는다.
• Failure‑mode taxonomy: 실무에서 샘플링된 토큰 OPD가 잘못되는 구체적인 세 가지 방식을 식별한다:
  1. 과도하게 불균형한 단일 토큰 학습 신호,
  2. 학생이 생성한 프리픽스에 대해 신뢰성이 떨어지는 교사 피드백, 그리고
  3. 토크나이저 또는 특수‑token 불일치로 인한 왜곡.
• Simple corrective recipe: top‑K local support matching 전략을 도입하며, 이는 top‑p 롤아웃 샘플링과 특수‑token 마스킹을 적용한 잘린 reverse‑KL 손실로 구현된다.
• Empirical validation: 단일 작업 수학 추론 및 다중 작업 “agentic‑plus‑math” 벤치마크에서 새로운 목표가 기존 vanilla sampled‑token OPD에 비해 최적화 곡선을 더 부드럽게 만들고 하위 작업 성능을 향상시킴을 보여준다.

방법론

1. Baseline OPD (sampled‑token): 학생이 토큰을 생성하고, 교사는 동일한 프리픽스가 주어졌을 때 해당 토큰의 확률을 평가하며, 토큰별로 reverse‑KL 손실을 적용한다.
2. Bias‑variance analysis: 저자들은 이 토큰별 추정기가 실제 시퀀스 수준 reverse‑KL(편향)과 동일하지 않지만 분산 하한이 낮아, 학생이 교사의 일반적인 경로에서 크게 벗어날 때 불안정해질 수 있음을 설명한다.
3. Diagnosing failure modes: 제어된 실험을 통해 불안정성의 세 가지 원인을 분리한다.
   • Signal imbalance: 대부분의 그래디언트가 소수의 고확률 토큰에서 발생하고, 나머지는 충분히 학습되지 않는다.
   • Teacher drift: 교사 전용 생성 과정에서 학생이 한 번도 보지 못한 프리픽스에 대한 교사의 확률 분포가 노이즈가 많다.
   • Tokenizer mismatch: 특수 토큰(예: padding, EOS)이나 서브워드 토크나이징 차이로 인해 손실이 관련 없는 차이를 벌한다.
4. Proposed fix:
   • Top‑K local support matching: reverse‑KL 계산을 교사의 top‑K(또는 top‑p) 토큰 집합으로 제한하여, 교사가 자신 있는 영역으로 분포를 효과적으로 잘라낸다.
   • Special‑token masking: 특수 토큰(예: padding, BOS/EOS)과 같이 인위적인 토큰에서 오는 손실 기여도를 0으로 만든다.
   • Rollout sampling: 순수한 그리디 또는 균등 샘플링 대신 top‑p 샘플링을 사용해 전체 롤아웃을 샘플링함으로써, 학생의 경로가 교사의 고확률 영역 안에 머물도록 한다.

결과적인 손실은 truncated reverse‑KL 형태로 표현될 수 있으며, 이는 교사의 로컬하게 지원되는 어휘에 대한 발산만을 벌한다.

결과 및 발견

핵심 관찰

• 새로운 목표는 평균 약 30 % 정도 그라디언트 분산을 감소시켜, “폭발” 업데이트가 감소합니다.
• 성능 향상은 좁은(단일 작업) 및 넓은(다중 작업) 커리큘럼 모두에서 일관됩니다.
• 소거 연구를 통해 각 구성 요소(Top‑p 샘플링, Top‑K 트렁케이션, 특수 토큰 마스킹)가 긍정적으로 기여함을 확인했으며, 이를 제거하면 안정성과 최종 점수가 모두 악화됩니다.

실용적 함의

• 보다 신뢰할 수 있는 파인‑튜닝 파이프라인: 개발자는 잘라낸 역‑KL 레시피를 채택하여 대형 교사 LLM을 더 작고 빠른 학생 모델로 증류할 수 있으며, 긴 프롬프트에서의 재앙적인 발산을 걱정할 필요가 없습니다.
• 컴퓨팅 낭비 감소: 그래디언트 분산이 낮아지면 수렴에 필요한 학습 에포크 수가 줄어들어 대규모 LLM 사후 학습에 대한 클라우드‑GPU 비용을 절감합니다.
• 에이전트형 애플리케이션에 대한 향상된 정렬: 다단계 추론이나 계획이 필요한 작업(예: 코드 생성, 도구 사용 에이전트)은 안정화된 OPD의 혜택을 받아, 교사의 능력을 그대로 유지하면서도 가벼운 어시스턴트를 구축하기가 쉬워집니다.
• 플러그‑앤‑플레이 호환성: 수정 사항은 손실‑함수 수준에서 구현되며 기존 OPD 코드베이스에 약간의 변경(Top‑p 샘플링 및 마스크 추가)만 필요하므로, 전체 학습 스택을 재설계하지 않고도 팀이 쉽게 통합할 수 있습니다.

제한 사항 및 향후 과제

• 편향이 남아 있음: 교사의 분포를 잘라내면 편향이 발생한다; 저자들은 편향과 분산 사이의 트레이드‑오프가 모델 크기에 따라 정량화될 필요가 있다고 언급한다.
• 토크나이저 의존성: 특수 토큰 마스킹이 많은 문제를 완화하지만, 이 접근 방식은 비교적 안정적인 토크나이저를 전제로 한다; 토크나이저가 빠르게 업데이트되면 불일치가 다시 발생할 수 있다.
• 확장성 테스트: 실험은 약 13 B 파라미터까지의 모델에 한정되었으며, 70 B 이상 대형 언어 모델에 적용하면 새로운 실패 모드가 드러날 수 있다.
• 향후 방향: 저자들은 롤아웃 엔트로피에 기반한 적응형 K/p 임계값을 탐색하고, OPD를 강화학습 스타일의 보상 형태와 결합하여 학생 행동을 장기 목표에 더욱 정렬시키는 방안을 제안한다.

저자

• Yuqian Fu
• Haohuan Huang
• Kaiwen Jiang
• Yuanheng Zhu
• Dongbin Zhao

논문 정보

• arXiv ID: 2603.25562v1
• 분류: cs.LG, cs.AI, cs.CL
• 출판일: 2026년 3월 26일
• PDF: PDF 다운로드