Bridging Reasoning Trajectories in On-Policy Distillation via Near-Future Guidance

기술 요약: 근미래 가이드를 통한 온폴리시 증류(On-Policy Distillation)에서의 추론 궤적 연결

문제 정의
온폴리시 증류(On-Policy Distillation, OPD)는 대규모 언어 모델(LLM)의 추론 능력을 향상시키기 위해 설계된 현대적 사후 학습(post-training)의 핵심 기법으로, 교사(teacher) 모델의 감독 하에 학생(student) 모델의 자체 정책에서 샘플링된 궤적을 통해 학습합니다. 표준적인 OPD는 "궤적은 샘플링하되 토큰은 학습하는(trajectory-sampled but token-learned)" 메커니즘을 기반으로 작동합니다. 즉, 높은 손실(high-loss)을 보이는 토큰을 통해 편차를 식별하고, 로컬 역-KL 발산(reverse-KL divergence) 수정을 통해 이를 복구합니다.

저자들은 이러한 접근 방식에 근본적인 불일치가 존재한다고 지적하며, 토큰 수준의 감독이 궤적 수준의 추론 교정에 필요한 요구사항과 구조적으로 정렬되지 않는다고 주장합니다. 이러한 불일치는 두 가지 결정적인 실패 모드로 나타납니다:

1. 노이즈가 섞인 허위 경보(Noisy False Alarms): 높은 손실을 보이는 토큰이 반드시 실제 추론의 발산(divergence)을 의미하는 것은 아닙니다. 실증적 분석 결과, 높은 손실을 보이는 토큰의 약 30%가 낮은 발산 영역에 속해 있으며, 이는 논리적 분기점이 아닌 표면적 형태의 불일치(예: "and"와 "then" 같은 스타일적 선호도 차이)를 나타냅니다. 이러한 토큰을 결정적인 오류로 취급하는 것은 추론 교정을 오히려 방해할 수 있는 노이즈 섞인 감독을 제공합니다.
2. 토큰 단위 학습의 함정(The Token-by-Token Learning Trap): 실제 발전 지점이 식별되더라도, 단일 토큰 단위의 역-KL 수정만으로는 전체 추론 경로를 재지향하기에 불충분한 경우가 많습니다. 추론 실패는 일반적으로 단기 지평의 분포 드리프트(distributional drift) 형태로 진화합니다. 하나의 토큰을 로컬하게 수정하더라도 모델이 교사의 미래 궤적과 정렬되는 것을 보장할 수 없으며, 모델은 로컬 수정 직후 다시 잘못된 분기로 이탈하는 경우가 빈번합니다.

방법론: 궤적 인식 OPD (Trajectory-aware OPD, TOPD)
이러한 한계를 해결하기 위해 저자들은 근미래 궤적 정보를 통합함으로써 로컬 토큰 수정과 글로벌 궤적 진화 사이의 간극을 메우는 **궤적 인식 OPD (TOPD)**를 제안합니다. 이 방법은 두 가지 주요 단계로 구성됩니다:

1. 발산 탐지(Divergence Detection):

   • 높은 손실 후보 위치에 대해, 시스템은 이전 토큰으로 되돌아가 공유 접두사(shared prefix)를 설정합니다.
   • 교사와 학생 모두 이 접두사로부터 짧은 윈도우 연속체(length $K=50$)를 생성합니다.
   • 시스템은 교사와 학생의 연속체 사이의 최적 운송(Optimal Transport, OT) 거리를 사용하여 궤적 수준의 불일치를 측정합니다.
   • 높은 토큰 수준 손실과 높은 단기 윈도우 OT 발산을 모두 보이는 위치를 실제 발산 지점으로 식별합니다. 이를 통해 무해한 표면적 불일치를 필터링합니다.

2. 궤적 수준의 교사 신호 주입(Trajectory-level Teacher Signal Injection):

   • 현재 토큰만을 최적화하는 대신, TOPD는 감독 범위를 전체 단기 윈도우 궤적으로 확장합니다.
   • OT 운송 계획($\gamma$)을 사용하여, 각 학생의 미래 위치에 대한 "궤적 인식 소프트 타겟(trajectory-aware soft target)"을 구축합니다. 이 타겟은 최적 운송 매핑을 기반으로 학생의 미래 토큰을 교사의 궤적에 정렬시킵니다.
   • 학생 모델은 이 OT 정렬된 소프트 타겟을 윈도우 전반에 걸쳐 일치하도록 학습하며, 이를 통해 단순히 현재 토큰을 수정하는 것이 아니라 자신의 근미래 추론 경로를 교사의 경로로 조종하는 법을 효과적으로 학습합니다.
   • 최종 목적 함수는 표준 OPD 손실과 궤적 손실 항을 결합합니다 ($L_{TOPD} = L_{OPD} + \lambda L_{traj}$).

주요 기여

• 불일치 식별: 본 연구는 토큰 수준의 감독이 궤적 수준의 추론 교정을 위한 "무딘 대리 지표(blunt proxy)"이며, 실제 논리적 분기와 무해한 표면적 변이를 구분하지 못한다는 점을 입증했습니다.
• 실증적 실패 모드 규명: 저자들은 높은 손실을 보이는 토큰의 약 30%가 낮은 발산성을 가진 허위 경보임을 보여주는 증거를 제시하고, 로컬 수정이 전체 추론 경로를 복구하지 못하는 "토큰 단위 학습의 함정"을 입증했습니다.
• 궤적 수준 가이드 원칙: 저자들은 근미래 궤적 정보가 여러 토큰에 걸쳐 분산되어 있다는 원칙을 도입했습니다. 짧은 윈도우 OT 정렬을 사용함으로써, TOPD는 교사-학생 간의 경로 불일치를 손실 함수에 직접 주입하여 학생이 고립된 토큰 편집이 아닌 궤적 변화를 학습할 수 있게 합니다.

실험 결과
Qwen3 모델(30B 교사를 4B 학생에게 증류)을 사용하여 경쟁력 있는 수학 벤치마크(AIME24, AIME25, HMMT25-Feb)에서 실험을 수행했습니다.

• 성능 향상: 표준 OPD는 웜스타트(warm-start) 모델의 평균 정확도를 38.9%에서 47.8%로 개선했습니다. TOPD는 이를 **52.2%**까지 추가로 높였습니다.
• 특정 벤치마크: TOPD는 AIME24 정확도를 60.0%에서 63.3%로, AIME25를 46.7%에서 53.3%로 향상시켰습니다.
• 절제 연구(Ablation Studies):
  • 표준 OPD에서 낮은 OT를 가진 높은 손실 토큰(허위 경보)을 억제했을 때 미미한 성능 향상(+0.4%)이 있었으며, 이는 해당 토큰들이 노이즈임을 확인시켜 줍니다.
  • 높은 OT를 가진 높은 손실 토کن(실제 발산)을 억제했을 때는 상당한 성능 저하가 발생하여, 이들의 필요성을 확인했습니다.
  • 로컬 궤적 수정 분석 결과, 표준 OPD는 단기 윈도우 OT 거리를 부분적으로 감소시켰으나(0.538에서 0.350으로), TOPD는 실질적인 감소(0.204까지)를 달성하여 근미래 궤적을 재지향하는 탁월한 능력을 입증했습니다.
• 효율성: TOPD는 OT가 선택된 높은 손실 윈도우에 대해서만 계산되므로, 표준 OPD 대비 1.41배의 완만한 계산 오버헤드를 발생시킵니다.

의의 및 주장
본 논문은 OPD에서의 추론 실패가 고립된 토큰 불일치가 아니라 근본적으로 궤적 수준의 현상이라고 주장합니다. 점 단위의 토큰 편집에서 글로벌하고 전이 인식적인(transition-aware) 정렬로 전환함으로써, TOPD는 더욱 효과적인 추론 정렬을 가능하게 합니다. 저자들은 이 연구가 더 궤적 인식적인 증류 목적 함수를 향한 단계라고 보고하며, 성공적인 추론 교정을 위해서는 학생이 단순히 어떤 토큰을 수정할 것인가가 아니라, 교사와 일관성을 유지하기 위해 어떻게 후속 추론 경로를 탐색할 것인가를 학습해야 한다고 강조합니다.

한계점
저자들은 TOPD가 OT 계산으로 인해 추가적인 계산 비용이 발생한다는 점, 단기 지평의 발산에 집중하고 있다는 점(잠재적으로 장기 컨텍스트 오류를 놓칠 수 있음), 그리고 현재 주로 수학적 추론 벤치마크에서 검증되었다는 점(코드나 개방형 작업에는 미흡할 수 있음)을 인정합니다. 또한, OT 거리는 궤적의 근접성을 측정하지만, 서로 다른 경로가 동일한 정답에 도달하는 경우의 의미론적 동등성을 포착하지 못할 수 있습니다.