기술 요약: 온-폴리 추론 증류를 위한 증거 탐색 (Finding the Evidence for On-Policy Reasoning Distillation)

1. 문제 정의

온-폴리 증류(On-Policy Distillation, OPD)는 교사 모델(Teacher model)로부터 조밀한 토큰 수준의 감독(supervision)을 제공함으로써 소형 모델에게 추론을 가르치는 지배적인 패러다임으로 부상했습니다. 그러나 현재의 OPD 방식은 결정적인 한계에 직면해 있습니다. 즉, 추론 체인 내의 서로 다른 지식 유형을 구분하지 못한다는 점입니다. 모든 토큰에 대해 학습을 진행하면 학습 신호가 희석되는 반면, 최근의 엔트로피 기반 선택 방식(예: Wang et al., 2026)은 "결정 토큰(decision tokens)"(학생 모델이 불확실성을 느끼는 지점)만을 포착할 뿐, "증거 토큰(evidence tokens)"은 놓치게 됩니다.

본 논문은 기존 방식들이 가진 구조적 사각지대를 식별합니다:

• 결정 토큰 (Decision Tokens): 학생 모델이 진정한 불확실성을 마주하는 위치(예: 분기점, 전략적 선택). 이들은 높은 학생 엔트로피(student entropy)에 의해 자연스럽게 식별됩니다.
• 증거 토큰 (Evidence Tokens): 결정을 정당화하는 중간 단계(예: 유도 과정, 변수 대입). 이들은 종종 학생 모델이 확신을 가지고 있으나 틀린(confident yet wrong) 상태에서 나타납니다. 학생 모델이 확신을 가지고 있기 때문에 엔트로피는 낮으며, 이로 인해 엔트로피 기반 선택기(selector)의 눈에 보이지 않게 됩니다. 결과적으로 실질적인 추론 지식이 전달되지 못하며, 이는 학생 모델이 언제 분기해야 하는지는 알지만, 그 분기를 어떻게 올바르게 실행해야 하는지는 배우지 못하는 결과를 초래합니다.

2. 방법론: DEAR (Decision-Evidence Aware Reasoning Distillation)

저자들은 결정 토큰과 증거 토큰을 모두 회복하여 OPD 학습에 활용하기 위해 설계된 2단계 토큰 선택 메커니즘인 DEAR를 제안합니다.

1단계: 결정 식별 (Decision Identification)

이 방법은 먼저 **학생 엔트로피($H^S_t$)**를 사용하여 결정 토큰을 식별합니다. 엔트로피가 가장 높은 상위 $p\%$의 토큰을 결정 집합($D$)으로 선택합니다. 이 토큰들은 학생 모델이 선택을 내려야 하는 추론 과정의 "골격(skeleton)"을 표시합니다.

씨: 증거 발견 (Evidence Discovery)

엔트로피가 놓치는 증거 토큰을 찾기 위해, DEAR는 결정과 그 주변 문맥 사이의 관계를 활용합니다:

1. 문맥적 유사성 (Contextual Similarity): 방법론은 가장 깊은 트랜스포머 레이어에서 학생의 은닉 상태($h^L_t$)를 추출합니다. 모든 비결정 토큰에 대해, 가장 가까운 결정 앵커(결정 집합 $D$에 속한 임의의 토큰)와의 **코사인 유사도(cosine similarity)**를 계산합니다. 높은 유사도는 해당 토큰이 동일한 추론 문맥에 참여하고 있음(즉, 결정의 근거가 되는 유도 과정의 일부임)을 나타냅니다.
2. 지식 격차 부스팅 (Knowledge Gap Boosting): 문맥적으로 유사하다고 해서 모든 토큰을 수정할 필요는 없습니다. 지식이 부족한 토큰을 우선순위에 두기 위해, 이 방법은 교사-학생 발산(teacher–student divergence)($\delta_t = |\log \pi_\theta(y_t) - \log \pi_T(y_t)|$)을 통합합니다.
3. 스코어링 (Scoring): 비결정 토큰에 대한 **증거 돌출도 점수(Evidence Salience Score, $s_j$)**를 다음과 같이 계산합니다:
   $$s_j = \hat{a}_j \times (1 + \hat{\delta}_j)$$
   여기서 $\hat{a}_j$는 노이즈를 걸러내는 게이트 역할을 하는 정규화된 코사인 유사도이며, $\hat{\delta}_j$는 지식 격차가 큰 토큰을 강조하는 정규화된 발산값입니다.
4. 선택 (Selection): 가장 높은 점수를 가진 상위 $q\%$의 비결정 토큰이 증거 집합($E$)을 형성합니다.

최종 학습 목적 함수 (Final Training Objective)

최종 학습 세트 $S$는 결정과 증거의 합집합($S = D \cup E$)입니다. OPD 손실(loss)은 오직 선택된 이 토큰들에 대해서만 계산되어, 모델이 분기 전략과 실질적인 중간 단계 모두를 학습하도록 보장합니다.

3. 주요 기여

1. 추론 체인의 분해: 본 논문은 추론 체인이 서로 다른 발견 메커니즘을 요구하는 두 가지 뚜렷한 지식 유형(결정: 높은 엔트로피, 증거: 낮은 엔트로피 및 높은 발산)으로 구성되어 있음을 경험적으로 입증합니다.
2. 발견 메커니즘: 은닉 상태 코사인 유사도(토큰을 결정과 연결)와 발산 부스팅(지식 격차를 우선순위화)을 결합하여 "확신에 찬 오류(confidently wrong)" 증거 토큰을 발견하는 방법을 도입합니다. 이는 엔트로피 기반 선택으로는 이러한 토큰을 찾는 것이 구조적으로 불가능하다는 점을 극복합니다.
3. 성능 향상: DEAR는 다양한 벤치마크에서 표준 OPD 및 오프라인 지식 증류(offline knowledge distillation)보다 일관되게 우수한 성능을 보입니다.

4. 실험 결과

저자들은 세 가지 학생-교사 구성(모델 패밀리 및 크기 변화)에 대해 6개의 수학적 추론 벤치마크와 3개의 코드 생성 벤치마크에서 DEAR를 평가했습니다.

• 수학적 추론 (Mathematical Reasoning):
  • DEAR는 경쟁 수준의 수학 문제(예: AIME)에서 표준 OPD 대비 최대 **2.5 퍼센트 포인트(pp)**의 개선을 달em했습니다.
  • 성능 향상은 지식 격차가 가장 큰 "교차 패밀리(cross-family)" 증류(설정 B: Qwen2.5-1.5B 학생, Qwen3-4B 교사)에서 가장 두드러졌습니다.
  • AIME 2024(설정 A)에서 DEAR는 4.58%의 정확도를 기록하여, 표준 OPD의 2.08% 대비 2.2배의 향상을 보였습니다.
• 코드 생성 (Code Generation):
  • DEAR는 MBPP에서 +5.74pp, HumanEval에서 +2.13pp, APPS에서 +2.56pp의 통과율(pass rate)을 개선했습니다.
  • 특히, 표준 OPD는 MBPP에서 기본 모델보다 성능이 저하되었으나, DEAR는 이를 해결함으로써 코드 생성 시 노이즈를 효과적으로 필터링함을 보여주었습니다.
• 신호 분석 (Signal Analysis):
  • 엔트로피 전용 선택은 전체 그래디언트 질량(gradient mass)의 약 39%만을 포착했습니다.
  • DE의 증거 토큰 포함은 그래디언트 질량 포착을 **약 76%**로 높여, 유효한 학습 신호를 거의 두 배로 늘렸습니다.
  • 분석 결과, 증거 토큰은 "낮은 엔트로피, 높은 발산" 영역에 집중되어 있으며, 이는 엔트로피 기반 방식이 체계적으로 무시하는 영역임을 확인했습니다.

5. 의의 및 주장

본 논문은 **완전한 추론 증류를 위해서는 단순히 어디에서 결정할 것인가뿐만 아니라, 그 결정을 정당화하는 지식을 전달해야 한다**고 주장합니다.

• 엔트로피를 넘어서: 이 연구는 높은 엔트로피를 가진 토큰만이 추론 능력을 운반하는 유일한 매개체라는 기존의 가정을 반박합니다. "확신에 찬 오류"가 중요한, 전달되지 못한 지식을 포함하고 있음을 입증합니다.
• 구조적 통찰: 증거 토큰이 은닉 상태 공간에서 결정 주변에 군집한다는 점을 식별함으로써, 본 논문은 이전의 선택 알고리즘에서는 보이지 않았던 실질적인 추론 단계를 복구하는 메커니즘을 제공합니다.
• 확장성: 이 이득은 문제의 난이도에 따라 확장됩니다. 증거 체인이 긴 어려운 문제일수록 증거 발견 메커니즘으로부터 가장 큰 혜hen을 얻습니다.

저자들은 결정-증거 구조가 수학 및 코드 분야에서 주로 검증되었다는 점과, 현재의 토큰 분류가 인과적 개입(causal intervention)보다는 표면적인 휴리스틱에 의존하고 있다는 점을 언급하며 신중한 입장을 유지합니다. 또한, 향-후 연구로서 학습이 진행됨에 따라 선택 비율($p, q$)을 적응적으로 스케줄링하는 방안을 제시했습니다.