기술 요약: 모든 불일치가 학습 가능한 것은 아님: 온-정책 증류에서의 토큰 가르침 가능성

문제 제기

온-정책 증류 (OPD) 는 교차 정책 (teacher policy) 의 토큰 수준 감독을 사용하여 학생 언어 모델을 자체 롤아웃 (rollouts) 으로 훈련합니다. 최근의 선택적 OPD 방법들은 높은 엔트로피나 높은 교사 - 학생 불일치 (KL 발산) 와 같은 "주목할 만한" (salient) 토큰을 우선시함으로써 효율성을 높이려 시도하지만, 이러한 접근법들은 높은 불일치가 곧 높은 학습 가치를 의미한다는 가정에 의존합니다.

본 논문은 이러한 가정에 치명적인 결함이 있음을 지적합니다. 즉, 원시 KL 불일치는 학습에 대한 거친 대리 지표일 뿐입니다. 이는 질적으로 다른 두 가지 시나리오를 혼동합니다:

1. 학습 가능한 불일치 (Learnable Disagreement): 교사가 학생의 현재 상위-K 예측 지지 (predictive support) 내에 이미 있는 토큰에 교정 확률 질량을 할당하는 경우.
2. 불일치 (Incompatible Disagreement): 교사가 학생의 현재 지지 범위에서 멀리 떨어진 토큰에 확률 질량을 할당하는 경우. 이는 큰 KL 격차를 생성하지만, 학생은 이러한 교정을 쉽게 흡수할 수 없어 약하거나 부정적인 국부적 학습 효과를 초래합니다.

핵심 문제는 어떤 토큰 수준의 교사 신호가 실제로 학습 가능한지 (즉, 국부적 개선을 유도할 수 있는지) versus 단순히 주목할 만하거나 불일치하는지를 구분하는 것입니다.

방법론

1. 고정 컨텍스트 진단

롤아웃 노이즈와 하류 컨텍스트 변화로부터 특정 토큰의 학습 가치를 분리하기 위해, 저자들은 **고정 컨텍스트 진단 (fixed-context diagnostic)**을 도입합니다.

• 프로토콜: 학생이 생성한 접두사 (prefixes) 를 수집하여 컨텍스트 뱅크로 고정합니다. 초기 (업데이트 전) 및 훈련된 학생 모델 모두 이러한 고정 컨텍스트에 대해 동일한 교사 분포에 대해 재점수 (rescored) 됩니다.
• 지표: **고정 컨텍스트 토큰 이득 (Fixed-Context Token Gain, $G^{\text{fix}}_{i,t}$)**은 동일한 접두사에서 훈련 후 교사와 학생 간의 KL 발산 감소량으로 정의됩니다.
• 목표: 이를 통해 저자들은 샘플링 분산을 통제하면서 업데이트 전 신호 (예: 원시 KL) 와 실제 국부적 개선을 상관관계 분석할 수 있습니다.

2. 토큰 가르침 가능성의 형식화

이 진단을 사용하여 저자들은 **국부적 지지 호환성 (local support compatibility)**에 기반하여 불일치를 두 가지 구성 요소로 분해합니다:

• 국부적 지지 ($S^S_t(K)$): 학생의 상위-K 토큰 집합.
• 호환성 질량 ($C_t$): 교사가 학생의 국부적 지지 $S^S_t(K)$에 할당하는 총 확률 질량.
• 분해:
  • 학습 가능한 불일치 ($D^L_t$): 교사의 교정이 학생의 지지와 정렬된 높은 불일치 (높은 $C_t$).
  • 불일치 ($D^I_t$): 교사의 교정이 학생의 지지에서 벗어난 높은 불일치 (낮은 $C_t$).

**가르침 가능성 점수 (Teachability Score)**는 정규화된 불일치와 호환성의 곱으로 정의됩니다:
$$s^{\text{teach}}_t = \tilde{D}_t \cdot \tilde{C}_t$$
여기서 $\tilde{D}_t$는 정규화된 불일치이고 $\tilde{C}_t$는 정규화된 호환성 질량입니다.

3. 가르침 가능성 인식 OPD (TA-OPD)

$D^L_t$가 원시 KL 보다 고정 컨텍스트 이득을 더 잘 예측한다는 발견에 기반하여, 저자들은 TA-OPD를 제안합니다. 이는 경량 토큰 위치 선택 방법입니다:

• 선택: 주어진 예산 (예: 토큰의 5% 또는 10%) 에 대해, TA-OPD 는 가장 높은 가르침 가능성 점수 ($s^{\text{teach}}_t$) 를 가진 상위-N 위치를 선택합니다.
• 구현: 이는 표준 OPD 중 사용 가능한 교사 및 학생 토큰 확률 (특히 상위-K 로그 확률) 만 필요합니다. 보상 모델, 검증기, 또는 추가 레이블은 필요하지 않습니다.
• 대조: 엔트로피나 원시 발산 (TIP) 을 기반으로 선택하는 이전 방법들과 달리, TA-OPD 는 명시적으로 "지지 정렬"된 교정을 필터링합니다.

주요 기여

1. 고정 컨텍스트 진단: 롤아웃 노이즈를 혼동하지 않고 업데이트 전 토큰 신호를 국부적 OPD 개선과 연결하는, 동일한 컨텍스트 KL 감소를 측정하는 방법.
2. 토큰 가르침 가능성 형식화: 원시 KL 불일치가 학습 가능하고 불일치하는 신호를 혼동한다는 증명. 본 논문은 지지 정렬된 불일치가 유용한 학습의 신뢰할 수 있는 예측 변수임을 입증합니다.
3. TA-OPD 알고리즘: 외부 도구 없이 효율성을 개선하면서 고타임 가능성 위치에 대해서만 OPD 손실을 적용하는 실용적이고 경량인 선택 방법.
4. 실증적 검증: 주목할 만한 토큰을 선택하는 것보다 가르침 가능성을 선택하는 것이 토큰 예산을 크게 줄임에도 불구하고 더 나은 하류 성능을 산출한다는 광범위한 실험.

실험 결과

저자들은 Qwen3 (4B, 8B, 14B) 및 Qwen2.5 (3B) 모델을 포함한 네 가지 교사 - 학생 설정에서 TA-OPD 를 평가했으며, AIME, GPQA-Diamond, HumanEval, MATH-500 과 같은 벤치마크를 사용했습니다.

• 전체 OPD 대비 성능: 많은 설정에서 TA-OPD 는 토큰의 5% 만 사용하면서도 전체 토큰 OPD 와 일치하거나 이를 능가했습니다. 예를 들어, Qwen3-4B $\to$ 1.7B 설정에서 TA-OPD 는 10% 예산을 사용했음에도 불구하고 평균 점수 44.89를 기록하여 전체 OPD 의 42.37을 상회했으며, 특정 지표에서도 전체 OPD 를 능가했습니다.
• 베이스라인 대비 비교: TA-OPD 는 엔트로피 기반, 발산 기반, 그리고 TIP (OPD 내 토큰 중요도) 베이스라인을 일관되게 능가했습니다.
  • 크로스-백본 설정 (DeepSeek-R1-Distill-Qwen-14B $\to$ Qwen2.5-3B) 에서 TA-OPD 는 평균 점수를 30.62로 향상시켜 베이스 모델 (29.98) 과 전체 OPD (28.76) 를 크게 능가했습니다.
• 예산 민감도: 결과는 OPD 성능이 예산에 의해 제한되기보다는 토큰 품질에 의해 제한됨을 나타냅니다. 10% 를 초과하여 토큰 예산을 늘리는 것은 종종 추가 이득을 가져오지 않았으며 때로는 성능을 저하시켰는데, 이는 모든 추가 감독이 유익한 것은 아님을 재확인합니다.
• 강건성: 긍정적인 "고지지 이득 격차" (높은 가르침 가능성과 낮은 가르침 가능성 토큰 간의 성능 차이) 는 다양한 컨텍스트 (수학, GSM8K-COT), 교사 규모 (8B, 14B), 그리고 지지 크기 ($K=8, 16, 32$) 에서 일관되게 유지되었습니다.

의미 및 주장

본 논문은 선택적 온-정책 증류를 "주목할 만한" 또는 "불확실한" 토큰을 선택하는 문제가 아니라, 국부적으로 학습 가능한 교사 신호를 식별하는 문제로 재정의합니다.

• 핵심 주장: OPD 의 이득은 교사의 교정과 학생의 현재 예측 지지 간의 호환성에 달려 있습니다. 높은 KL 발산만으로는 부족하며, 교사는 학생이 이미 그럴듯하다고 고려하는 토큰에 확률 질량을 할당해야 합니다.
• 실제적 영향: TA-OPD 는 (불일치 토큰을 필터링함으로써) 증류의 계산 비용을 극적으로 줄이면서도 모델 성능을 유지하거나 개선하는 방법을 제공합니다. 이는 보상 모델이나 검증기의 오버헤드 없이 달성되므로 표준 훈련 파이프라인에 적용 가능합니다.
• 한계: 저자들은 분석이 수학적 추론과 Qwen 계열 모델에 초점을 맞추고 있음을 지적합니다. 진단은 모델에 무관하지만, 대화, 코드, 비-Qwen 백본에 대한 광범위한 검증을 통해 일반성을 확인해야 합니다. 또한, "토큰 예산"은 벽시계 속도 향상 (wall-clock speedup) 이 아니라 감독 밀도를 의미하며, 이 방법은 트랜스포머 연산을 가지치기하는 것이 아니라 위치를 선택합니다.

결론적으로, 이 연구는 효과적인 증류가 호환성 인식 감독 전략을 필요로 함을 시사하며, 학생은 교정적이면서도 국부적으로 도달 가능한 교사 신호로부터 가장 잘 학습합니다.