Escaping the KL Agreement Trap in On-Policy Distillation

기술 요약: 온폴리시 증류에서의 KL 합의 함정 탈출 (Escaping the KL Agreement Trap in On-Policy Distillation)

1. 문제 정의

온폴리시 증류(On-Policy Distillation, OPD)는 학생 모델이 생성한 롤아웃(rollout)에 대해 토큰 단위의 조밀한 감독(supervision)을 제공함으로써, 거대 교사 모델로부터 작은 학생 모델로 능력을 전이하는 강력한 방법으로 부상했습니다. 그러나 본 논문은 **낮은 KL 합의 함정(low-KL agreement trap)**이라 명명된 결정적인 실패 모드를 식별합니다.

표준 OPD에서 학생과 교사 사이의 낮은 역(reverse) Kullback-Leibler(KL) 발산은 일반적으로 진전이나 정렬(alignment)의 신호로 해석됩니다. 저자들은 학생의 롤아웃이 진행됨에 따라, 모델이 "손상된 접두사(corrupted prefix)"(예: 초기 추론 오류 또는 반복적인 퇴행 패턴)로 빠져들 수 있음을 관찰했습니다. 이렇게 저하된 상태에 조건화되면, 교사 모델은 국소적으로는 그럴듯하지만 정보가 없는 연속(continuation)을 생성할 수 있으며, 결과적으로 낮은 KL 발산을 초래합니다. 이는 **퇴행적 합의(degenerate agreement)**를 만듭니다. 즉, 교사와 학생이 서로 일치하기는 하지만, 교사가 교정적인 가이드를 제공하는 대신 단순히 학생의 오류를 따르고 있기 때문에 신호가 약해지는 현상입니다.

본 논문은 기존의 OPD 방식들이 전체 롤아웃을 유용한 감독 신호로 취급한다고 주장합니다. 결과적으로, 이러한 방식은 감독 신호가 희석되거나 해로울 수 있는 함정 진입 이후의 토큰들에 대해 계산 자원을 낭비하며, 잠재적으로 최적화 과정을 잘못된 방향으로 유도할 수 있습니다.

2. 방법론: KAT (KL Agreement Trap Termination)

이를 해결하기 위해 저자들은 근본적인 OPD 목적 함수를 변경하거나 보조 모델을 필요로 하지 않으면서, 이러한 퇴행적 합의 영역을 탐지하고 필터링하도록 설계된 온라인 롤아웃 종료 규칙인 KAT를 제안합니다.

핵심 메커니즘:
KAT는 OPD 과정에서 이미 계산된 역 KL 발산($d_t$)을 재사용합니다. 단일 토큰의 KL 임계값을 설정하는 대신, KAT는 $W$개 토큰의 윈도우 동안 평균 KL을 측정하는 슬라이딩 윈도우 통계량($z_t$)을 모니터링합니다.

적응형 임계값 설정 (Adaptive Thresholding):
KL 발산의 규모가 훈련 과정 전반에 걸쳐 변한다는 점을 인식하여, KAT는 동적이고 데이터 기반인 임계값($s_k$)을 채택합니다.

• 웜업 단계 (Warmup Phase): 초기 $K$번의 옵티마이저 스텝 동안, KAT는 완료된 롤아웃으로부터 얻은 최소 슬라이딩 윈도우 KL 값($m(\hat{y})$)들을 FIFO 버퍼에 수집합니다.
• 보정 (Calibration): 웜업 이후, 중단 임계값은 버퍼의 최근 최솟값들에 대한 $(1-\eta)$ 분위수(quantile)로 설정됩니다. 이를 통해 규칙이 진화하는 학생 정책과 KL 규모에 적응할 수 있도록 합니다.

종료 로직 (Termination Logic):
슬라이딩 윈도우 KL 통계량이 $T$개의 연속된 윈도우 동안 적응형 임계값 미만으로 유지되면 롤아웃이 온라인으로 종료됩니다. 결정적으로, KAT는 KL이 자연스럽게 낮아지는 롤아웃 초기 단계에서 조기 종료되는 것을 방지하기 위해 면제 길이($L_0$)를 포함합니다. 트리거가 확인되면, 롤아웃은 첫 번째 낮은 KL 윈도우의 시작점에서 잘리며(truncate), 이후의 토큰은 더 이상 생성되거나 감독에 사용되지 않습니다.

3. 주요 기여

본 논문은 세 가지 주요 기여를 합니다:

1. 낮은 KL 합의 함정의 식별: 저자들은 낮은 KL이 유익한 정렬이 아닌, 손상된 접두사에 대한 "퇴행적 합의"를 나타내는 특정 실패 모드를 규명했습니다.
2. 감독 품질 분석: 그래디언트 기하학 분석(주요 업데이트 부공간으로 그래디언트 투영)과 어휘 시각화를 통해 다음을 입증했습니다:
   • "합의(agreement)" 및 "합의 후(post-agreement)" 단계의 그래디언트는 주요 파라미터 업데이트 방향과 정렬이 잘 되지 않습니다.
   • 이 영역에서 교사의 어휘는 추론 관련 토큰에서 국소적인 표면 파편(surface fragments)이나 일반적인 연속 문구로 변화합니다.
   • 이러한 영역에서는 KL이 약간 증가하더라도 감독의 품질이 회복되지 않습니다.
3. KAT 제안: 정보가 없는 접미사를 필터링하는 적응형 온라인 종료 규칙을 제안합니다. 이는 유익한 사전 함정(pre-trap) 접두사에 업데이트를 집중함으로써 학습 효율성과 정확도를 개선하며, 토큰당 $O(1)$의 기록 작업만을 요구합니다.

4. 실험 결과

저자들은 Qwen3-8B 교사를 사용하여 두 가지 학생 규모(Qwen3-0.7B/1.6B 및 Qwen3-4B)에 대해 네 가지 수학적 추론 벤치마크(AMC, MATH500, MinervaMath, AIME24)에서 KAT를 평가했습니다.

성능 향상:

• 정확도: KAT는 표준 OPD와 비교하여 평균 avg@k 정확도를 2.66%, pass@k를 3.43% 향상시켰습니다.
• 효율성: 평균 롤아웃 길이가 59.73% 감소했습니다 (약 1480개에서 약 596개 토큰으로).
• 학습 비용: 4B 학생 모델의 경우, KAT는 표준 OPD 대비 학습 시간을 2.4배 단축(26.28시간에서 11.00시간)시켰고, 연산 사용량을 59.4% 줄였습니다.

비교 분석:

• 고정 접두사(Fixed Prefix) 대비: KAT는 고정 길이 절단(예: 1024 또는 2048 토큰)보다 우수한 성능을 보였으며, 이는 절단 지점의 품질이 길이 자체보다 더 중요하다는 것을 입증했습니다. 고정 접두사는 유용한 추론을 잘라내거나 무익한 접미사를 남기는 경우가 많습니다.
• 무작위 종료(Random Termination) 대비: KAT는 무작위 절단보다 성능이 현저히 높았으며, 이는 낮은 KL 신호가 단순한 길이 지표가 아니라 약한 감독을 나타내는 특정한 지표임을 확인시켜 줍니다.
• 어블레이션(Ablation): KAT가 수행하는 것처럼 훈련에서 "합의" 및 "합의 후" 단계를 제거했을 때, 해당 단계들을 포함했을 때보다 더 높은 정확도를 얻었으며, 이는 해당 영역이 약하거나 해로운 신호를 제공한다는 점을 검증합니다.

5. 의의 및 주장

본 논문은 KAT가 감독 신호를 활용하는 방식을 근본적으로 재고함으로써, 더 효과적이고 효율적인 온폴리시 증류 접근법을 제공한다고 주장합니다. 낮은 KL이 항상 양호한 것은 아니라는 점을 식별함으로써, KAT는 정보가 없는 토큰으로 인해 학습 신호가 희석되는 것을 방지합니다.

저자들은 KAT가 새로운 손실 함수나 보조 모델을 도입하지 않고도 기존의 OPD 목적 함수를 변경하지 않고 이러한 이득을 달성한다는 점을 강조합니다. KAT는 기존의 역 KL 계산을 활용하여 온라인으로 적응적인 결정을 내립니다. 결론적으로 KAT는 특히 추론 경로가 복구 불가능한 상태로 쉽게 퇴행할 수 있는 수학적 추론 작업에서, OPD를 더 효과적(높은 정확도)이고 더 효율적(낮은 연산량)으로 만듭니다.

한계점 명시:
저자들은 본 연구가 수학적 추론에 집중되어 있음을 겸허히 언급했습니다. 코드 생성, 개방형 지시 이행(open-ended instruction following), 또는 다회차 대화(multi-turn interaction)와 같은 더 넓은 도메인에서의 종료 규칙 동작은 향후 연구 과제로 남아 있습니다. 또한, 하이퍼파라미터가 직관적이기는 하지만, 다양한 환경에서 이를 튜닝하기 위한 보다 체계적인 가이드라인이 제공된다면 사용성을 더욱 높일 수 있을 것입니다.