기술 요약: TRACE - 토큰 라우팅 자기 온-정책 정렬을 통한 핵심 영역 증류

1. 문제 제기

본 논문은 수학 추론 작업에서 검증 가능한 보상 (RLVR) 을 활용한 장기-지평 (long-horizon) 강화학습에 적용될 때의 **온-정책 자기 증류 (Self-OPD)**의 한계를 다룹니다. Self-OPD 방법들 (예: SDPO, SRPO) 은 정책 스냅샷을 교사 (teacher) 로 사용하여 학생 (student) 을 안내함으로써 희소한 보상 신호를 밀집화하지만, 저자들은 장기 훈련에서 다음과 같은 치명적인 실패 모드를 식별합니다:

• 세분성 불일치 (Granularity Mismatch): 기존 방법들은 응답 내 모든 토큰에 걸쳐 지식 증류 (KL) 손실을 적용합니다. 그러나 분석에 따르면, 올바른 롤아웃 (rollout) 에서의 확률 질량의 약 97% 는 이미 학생과 교사 간에 공유되고 있어, 대부분의 KL 기울기는 불필요합니다.
• "균일 증류 세 (Uniform Distillation Tax)": 모든 토큰에 KL 을 적용하면 모델이 수정이 필요 없는 위치들에 기울기를 소모하게 됩니다. 이는 다음과 같은 결과를 초래합니다:
  • 특권 정보 누출 (Privileged Information Leakage): 학생은 추론 시점에 존재하지 않는 특권 컨텍스트 (예: 검증된 추적) 에 존재하는 상관관계를 흡수하여, 분포 외 (OOD) 성능 저하를 유발합니다.
  • 엔트로피 상승 및 추론 붕괴: 저자들은 장기 훈련에서 "3 가지 증상 붕괴"를 관찰합니다: 응답 길이가 급격히 단축되고, 토큰별 액터 엔트로피가 상승하며 (종종 기준 대비 4 배 이상), 검증 정확도가 크게 하락합니다.
• 핵심 질문: 올바른 증류 단위는 전체 응답이 아니라, 학생의 분포가 실제로 수정을 필요로 하는 결정적 추론 토큰의 소규모 부분집합입니다.

2. 방법론: TRACE

저자들은 TRACE(Critical rEasoning 을 위한 토큰 라우팅 정렬) 를 제안합니다. 이는 균일하게 적용하는 대신 특정 토큰 구간에만 증류 신호를 라우팅하는 방법입니다.

핵심 구성 요소

1. 토큰 라우팅 자기 Self-OPD: 전역 KL 손실 대신 TRACE 는 특권 주석자 (privileged annotator)($\pi_A$) 를 사용하여 롤아웃 내 특정 구간을 식별합니다:

   • 핵심 구간 ($K_y$): 올바른 롤아웃 (보상=1) 의 결정적 추론 단계.
   • 오류 구간 ($E_y$): 잘못된 롤아웃 (보상=0) 의 근본 원인 추론 오류.
   • 비-구간 ($N_y$): 나머지 모든 토큰.
   • 제약: 특권 정보 노출을 제한하기 위해 총 구간 커버리지는 $\alpha = 0.25$(토큰의 25%) 로 제한됩니다.

2. 대략적 진단 라우팅 (Coarse Diagnostic Routing):

   • 주석자는 대략적인 유형 레이블(예: case_split, missed_case) 을 생성하지만, 구간 텍스트나 전체 롤아웃 내용은 공개하지 않습니다.
   • 교사($\pi_T$) 는 문제와 이 대략적 레이블을 개인 진단 접두사로 받습니다. 이는 학생이 생성한 접두사에 대해 인과적으로 로그 확률 (logits) 을 계산하지만, 구체적인 구간 위치나 내용은 보지 않습니다.
   • 학생($\pi_S$) 은 KL 항이 마스킹된 구간에만 적용되는 라우팅 손실을 사용하여 훈련됩니다.

3. 비대칭 발산 라우팅 (Asymmetric Divergence Routing):

   • 핵심 구간에서의 순방향 KL (FKL): 올바른 롤아웃에 적용됩니다. FKL ($KL(\pi_T \parallel \pi_S)$) 은 학생의 현재 확률 질량과 관계없이, 학생이 과소 할당한 토큰에 "상승 (lift)"을 제공합니다.
   • 오류 구간에서의 역방향 KL (RKL) (선택적): 잘못된 롤아웃에 적용됩니다. RKL ($KL(\pi_S \parallel \pi_T)$) 은 교사가 선호하지 않지만 학생이 과도하게 확신하는 토큰을 억제합니다.
   • 비-구간에서의 KL 부재: 이러한 토큰들은 표준 GRPO(그룹 상대 정책 최적화) 손실에만 의존합니다.

4. 유한 노출 일정 (Finite Exposure Schedule):

   • KL 가중치 ($\lambda_k$) 는 짧은 워밍업 (예: 40 스텝) 후 0 으로 감쇠 (annealed) 됩니다.
   • 이는 훈련 기간 동안 누적 특권 기울기 노출이 유한하게 유지되도록 하여, OOD 성능을 저하시키는 상관관계의 장기적 축적을 방지합니다.

3. 이론적 분석

논문은 라우팅 전략에 대한 이론적 근거를 제공합니다:

• 비대칭 포인트별 압력 (Corollary 2):
  • FKL은 과소 할당된 토큰 (학생 확률 $p \ll$ 교사 확률 $q$) 에 효과적입니다. 왜냐하면 기울기 압력은 간격 $(q-p)$에 비례하며 $p \to 0$으로 갈 때 소멸하지 않기 때문입니다.
  • RKL은 과도하게 확신하는 잘못된 토큰 (where $p \gg q$) 에 효과적입니다. 왜냐하면 기울기가 학생의 과도한 확신 ($p$) 에 비례하기 때문입니다.
• 누적 노출 상한 (Proposition 3): 구간 마스킹 (대역폭을 $\alpha$로 제한) 과 KL 감쇠 (기간 제한) 를 결합함으로써, TRACE 는 누적 특권 기울기 노출을 $O(\alpha \Lambda^2)$로 제한하여, 영구적인 모든-토큰 KL 기준선에서 관찰되는 무한한 성장을 방지합니다.
• 신호 vs. 누출: 이 방법은 핵심 구간에서의 정렬 신호를 극대화하면서, 위장 상관관계를 유발하는 특권 정보 누출 위험을 최소화합니다.

4. 주요 기여

1. 진단: 모든-토큰 Self-OPD 의 "균일 증류 세"를 식별하여, 엔트로피 상승, 인식적 억제, 특권 기울기 누출을 세분성 불일치와 연결합니다.
2. 방법론 및 이론: 토큰 클래스 (핵심, 오류, 비-구간) 에 기반하여 $\{FKL, RKL, \emptyset\}$을 라우팅하고 특권 채널을 감쇠시키는 TRACE 를 도입합니다. 이론은 강력한 베이스 영역에서 핵심 구간에 FKL 이 우세한 이유와, 약한 베이스 영역에서 오류 구간에 최적의 모서리가 RKL 로 이동하는 이유를 설명합니다.
3. 실증적 증거:
   • 성능: Qwen3-8B 에서 TRACE 는 GRPO 대비 5 개 벤치마크 평균을 2.76 퍼센트 포인트 향상시킵니다.
   • 안정성: SDPO/SRPO 와 달리, TRACE 는 기준 모델이 저하되는 GPQA-Diamond 에서 베이스 모델의 OOD 성능을 유지합니다.
   • 자기 주석: 이 방법은 온라인 자기 주석자(학생 정책 자체) 와 함께 작동하여, 강력한 외부 API 주석자로부터의 이득의 약 69% 를 회복합니다. 이는 이득이 외부 능력 때문이 아니라 라우팅 메커니즘에서 비롯됨을 증명합니다.
   • 크로스-스케일 적응: 약한 베이스 (Qwen3-1.7B) 에서 지배적인 라우팅 모서리는 FKL-on-Key 에서 RKL-on-Error 로 자연스럽게 반전되어, 이론적 예측과 일치합니다.

5. 결과

• Qwen3-8B (강력한 베이스): TRACE(FKL on Key Spans) 는 5 개 벤치마크 (MATH-500, AIME 24/25, AMC 23, GPQA-Diamond) 에서 평균 정확도 **81.51%**를 달성하여 GRPO(78.75%) 와 가장 강력한 이전 Self-OPD 기준선인 RLSD(78.99%) 를 능가합니다. 특히, GRPO 와 모든-토큰 기준선이 크게 하락하는 반면, 베이스 모델의 GPQA-Diamond 점수 (58.33% vs 베이스 58.27%) 를 유지합니다.
• Qwen3-1.7B (약한 베이스): TRACE(RKL on Error Spans) 는 평균 **60.16%**를 달성하여 베이스 모델 (58.77%) 을 초과한 유일한 훈련된 방법입니다. 모든-토큰 기준선 (SDPO/SRPO) 은 붕괴하여 GRPO 보다 11~19 점 낮아집니다.
• 애블레이션 연구:
  • 마스킹 위치 지정: TRACE 의 특정 구간 마스킹을 무작위 25% 마스킹이나 모든-토큰 KL 로 대체하면 성능이 하락하여, 어떤 토큰이 선택되느냐가 양보다 더 중요함을 확인합니다.
  • 주석자 품질: 능동적으로 훈련 중인 학생을 자신의 주석자로 사용하면 +1.90 점의 향상을 보여주며, 이 방법이 외부 "초인간" 추론 능력을 가져오는 것에 의존하지 않음을 입증합니다.

6. 중요성 및 주장

본 논문은 TRACE 가 도메인 특화 해결책이 아니라 장기-지평 Self-OPD 의 불안정성에 대한 손실-기하학적 (loss-geometric) 해결책을 제공한다고 주장합니다. 그 중요성은 다음과 같습니다:

• 증류와 불안정성의 분리: "특권" 신호를 핵심 구간으로 제한하고 이를 감쇠시킴으로써, TRACE 는 SDPO 와 SRPO 에서 관찰된 엔트로피 상승 및 추론 붕괴를 피합니다.
• 적응형 라우팅: 최적의 증류 방향 (FKL 대 RKL) 이 고정된 하이퍼파라미터가 아니라 베이스 모델의 오류 프로파일 (과소 할당 대 과도한 확신) 에 의존함을 보여줍니다.
• 실용적 타당성: 외부 교사나 비싼 쿼리당 평가 기준 없이도 고품질 증류가 가능함을 입증하며, 대신 자가 생성 가능한 경량 대략적 진단 주석에 의존합니다.

저자들은 TRACE 가 고정된 레시피가 아니라 "코너 라우팅 프레임워크"라고 강조하며, 특정 라우팅 동작 (FKL 대 RKL) 은 베이스 모델의 잔여 오류 모드에 따라 선택되어야 한다고 주장합니다.