Tail-Aware Top-$k$ On-Policy Distillation

기술 요약: 꼬리 인지형 Top-k 온-폴리시 증류 (Tail-Aware Top-k On-Policy Distillation, TA-OPD)

문제 정의
온-폴리시 증류(On-policy distillation, OPD)는 학생 언어 모델이 자신의 궤적(trajectory)을 따라 교사 모델의 다음 토큰 분포와 정렬되도록 학습함으로써, 오프-폴리시 증류에서 발생하는 훈련-추론 간의 불일치 문제를 해결합니다. 널리 사용되는 방식인 **정규화된 Top-k OPD(Normalized Top-k OPD)**는 교사의 상위 $k$개 토큰에 국한된 역 쿨백-라이블러(reverse KL) 발산을 최소화합니다.

본 논문은 이 정규화된 목적 함수가 꼬리 확률(tail probability)(상위 $k$개 토큰 이외의 총 확률 질량)에 관한 정보를 버린다는 치명적인 결함을 지적합니다. 이론적 분석과 경험적 증거에 따르면, 정규화된 목적 함수를 최소화하는 것은 학생의 꼬리 확률과 엔트로피를 지속적으로 증가시키는 체계적인 편향을 생성합니다. 학생의 꼬리 확률이 커짐에 따라, 학생은 교사의 top-$k$ 집합 외부의 토큰을 점점 더 많이 샘플링하게 되며, 이는 교사의 감독이 신뢰할 수 없는 접두사(prefix)로 표류하게 만듭니다. 이러한 현상은 특히 학생과 교사 사이의 능력 격차가 클 때 다운스트림 정확도의 저하를 초కు합니다.

방법론: 꼬리 인지형 Top-k OPD (TA-OPD)
꼬리 확률 신호의 손실을 해결하기 위해, 저자들은 명시적인 **꼬리 토큰($v_{tail}$)**을 도입하여 누락된 정보를 복원하는 방법인 TA-OPD를 제안합니다.

1. 확장된 토큰 집합: 확률을 오직 상위 $k$개 토큰에 대해서만 정규화하는 대신, TA-OPD는 $S^+_t = S^k_t \cup \{v_{tail}\}$로 정의된 확장된 집합을 사용합니다. 여기서 $S^k_t$는 교사의 상위 $k$개 토큰을 나타냅니다.
2. 확률 할당:
   • $S^k_t$에 속한 토큰들은 원래의 확률 $p_t(v)$를 유지합니다.
   • 꼬리 토큰 $v_{tail}$에는 전체 꼬리 확률이 할당됩니다: $p_t(v_{tail}) = p^{tail}_t = \sum_{v \notin S^k_t} p_t(v)$.
   • 마찬가지로, 교사의 분포 또한 $q^{tail}_t$로 확장됩니다.
3. 손실 함수: 이 방법은 확장된 집합에 대해 역 KL 발산을 최소화합니다:
   $$\ell^{TA}_t = \sum_{v \in S^k_t} p_t(v) \log \frac{p_t(v)}{q_t(v)} + p^{tail}_t \log \frac{p^{tail}_t}{q^{tail}_t}$$
   이 공식은 학생의 꼬리 확률을 교사의 꼬리 확률 쪽으로 유도함으로써, 학생과 교사의 꼬리 확률을 명시적으로 비교합니다.

이론적 특성

• 그래디언트 역학(Gradient Dynamics): top-$k$ 집합 외부의 로짓(logits)에 대해 제로 그래디언트를 제공하는 정규화된 목적 함수와 달리, TA-OPD는 꼬리 로그-비(log-ratio)를 통해 꼬리 로짓을 업데이트합니다. 또한, top-$k$ 로짓에 대한 그래디언트가 더 이상 합이 0이 되지 않으므로, 학생이 꼬리에 과도한 가중치를 둘 때 top-$k$ 로짓을 꼬리 로짓에 비해 명시적으로 높일 수 있습니다.
• 하한(Lower Bound): TA-OPD 손실은 전체 어휘 역 KL 발산의 타이트한 하한(tight lower bound)임이 증명되었습니다. 그 차이는 꼬리 확률에 비례하는 비음수 잔차 항입니다. 실제로 TA-OPD가 꼬리 확률을 0으로 유도함에 따라 이 잔차는 사라지며, 목적 함수는 타이트한 근사치가 됩니다.
• 편향 없는 변형 (SC-TA-OPD): 샘플링된 토큰의 확률을 사용할 수 있는 경우, 저자들은 **샘플 교정 TA-OPD(Sample-Corrected TA-OPD)**를 도출했습니다. 이 변형은 샘플링된 토큰에 기반한 편향 교정 항을 추가하여, 값과 그래디언트 모두에서 전체 어휘 역 KL 발산의 **불편 추정치(unbiased estimate)**를 제공합니다.

주요 기여

1. 꼬리 확률 증가 식별: 본 논문은 정규화된 top-k OPD가 학생의 꼬리 확률과 엔트로피를 지속적으로 증가시켜 성능 저하를 초래한다는 것을 이론적 및 경험적으로 입증했습니다.
2. TA-OPD 제안: 꼬리 토큰을 통해 꼬리 확률 신호를 복원함으로써, top-k 정규화와 관련된 엔트로피 증가를 방지하는 새로운 증류 방법을 제안했습니다.
3. 이론적 보장: TA-OPD가 전체 어휘 목적 함수의 타이트한 하한임을 증명하고, 불편 추정치인 샘플 교정 변형을 도출했습니다.
4. 경험적 우월성: TA-OPD가 특히 학생과 교사 사이의 능력 격차가 큰 시나리오에서 베이스라인보다 우수한 성능을 보임을 광범위한 실험을 통해 입증했습니다.

실험 결과
저자들은 세 가지 학생-교사 쌍(예: OpenThinker3-7B로부터 증류된 Qwen2.5-7B-Instruct)을 사용하여 수학적 추론 벤치마크(MATH500, Minerva, OlympiadBench, AMC, AIME24/25)와 분포 외(OOD) 작업(ARC-c, MMLU-Pro)에서 TA-OPD를 평가했습니다.

• 성능 향상: TA-OPD는 일관되게 가장 높은 평균 정확도를 달림이 확인되었습니다. 주목할 만한 개선 사항은 다음과 같습니다:
  • Llama-3.1-8B의 수학 벤치마크 평균에서 +8.05 포인트 상승 (정규화된 top-k OPD 대비).
  • Llama-3.1-8B의 MATH500에서 +19.53 포인트 상승 (31.30%에서 50.83%로 개선).
  • OOD 성능에서의 상당한 이득 (예: Qwen2.5-7B-Instruct의 ARC-c 정확도가 23.07%에서 76.30%로 개선).
• 훈련 역학: TA-OPD는 정규화된 top-k OPD가 엔트로피를 ~6까지, 꼬리 확률을 ~0.6까지 높인 것과 달리, 학생의 엔트로피를 1.5 미만으로, 꼬리 확률을 0 근처로 유지하며 훈련을 안정화했습니다.
• 능력 격차 민감도: 이 방법은 학생과 교사 사이의 능력 격차가 클 때 가장 효과적입니다. 격차가 작을 때는 성능 향상이 줄어듭니다.
• 파라미터 민감도: TA-OPD는 작은 $k$(예: $k=2$)에서도 효과적이며, $k \ge 2$인 경우 $k$ 값에 무관하게 안정적인 성능을 보입니다.

의의 및 주장
본 논문은 TA-OPD가 꼬리 확률 신호를 보존함으로써 기존의 top-k 증류 방식의 근본적인 결함을 해결한다고 주장하며, 이는 학생이 교사와 정렬 상태를 유지하는 데 매우 중요합니다. 이 방법은 교사의 top-$k$ 확률 외에 추가적인 쿼리를 요구하지 않으며 구현이 간단한 실용적인 개선책으로 제시됩니다. 저자들은 TA-OP입니다 큰 격차가 있는 시나리오에서 상당한 이점을 제공하지만, $k$가 증가하거나 능력 격차가 좁아질수록 그 이점이 감소한다는 점을 언급했습니다. 본 연구는 계산상의 제약으로 인해 최대 8B 파라미터 규모의 모델까지 제한됩니다.