기술 요약: 기둥인가 걸림돌인가? 온-정책 증류에서의 Rock Token 해독

1. 문제 제기

온-정책 증류 (On-Policy Distillation, OPD) 는 학생 모델이 자신의 정책에서 샘플링된 궤적을 통해 우월한 교사 모델과 정렬함으로써 토큰 수준의 추론 정제를 가능하게 하여, 대형 언어 모델 (LLM) 의 사후 학습 파이프라인의 기둥이 되었습니다. 최근 검증 가능한 보상과 강화 학습 (RLVR) 에 대한 연구는 소수의 중요한 토큰 (예: 고엔트로피 '분기 토큰') 이 추론 향상을 불균형적으로 주도한다는 점을 규명했지만, OPD 에 대한 유사한 토큰 수준의 이해는 아직 탐구되지 않았습니다.

OPD 에서 목적 함수는 학생과 교사 정책 간의 토큰 단위 역 KL 발산을 최소화합니다. 직관적으로 훈련이 수렴함에 따라 고손실 토큰 (학생 - 교사 불일치를 나타냄) 은 감소해야 합니다. 그러나 실증적 관찰은 역설을 드러냅니다. 훈련이 명백한 포화 상태에 도달한 후에도 상당한 부분의 토큰이 지속적으로 높은 손실을 보입니다. 이 연구가 다루는 핵심 질문은 이러한 지속적 고손실 토큰이 정책 정렬에 필수적인 '기둥 (Cornerstones)'인지, 아니면 불필요한 최적화 중복을 초래하는 단순한 '걸림돌 (Stumbling Blocks)'인지 여부입니다.

2. 방법론

2.1 Rock Token 의 정의 및 탐지

저자들은 Rock Token을 긴 궤적에서의 고립된 스파이크가 아니라 유사한 문맥 전반에 걸쳐 안정적인 지속적 고손실 행동을 보이는 토큰 유형으로 정의합니다. 탐지 과정은 세 단계로 구성됩니다:

1. 집합적 부담 계산: 학생이 생성한 궤적에서 각 토큰 유형 $v$의 평균 토큰 단위 손실 ($\bar{\ell}_v$) 과 빈도 ($Freq(v)$) 의 곱으로 정의된 'Rock 점수' $R(v)$를 계산합니다.
2. 지속성 필터링: OPD 훈련 전후 모두 고손실로 남아 있는 발생 사례를 식별합니다.
3. 문맥 일관성 필터링: 고손실 발생 사례가 유사한 지역 문맥에서 반복되는지 확인합니다 (문맥 윈도우 유사성 지수 사용). 이는 고립된 위치적 노이즈와 진정한 구조적 불일치를 구분합니다.

최종 지표인 문맥 인식 Rock 점수 $R_{ctx}(v)$는 집합적 부담과 문맥 일관성을 결합합니다. 선택 안정성 (재현성) 과 손실 커버리지 (코퍼스 수준 KL 발산의 약 60% 포착) 를 균형 있게 맞추기 위해 컷오프 $K=100$을 선택했습니다.

2.2 기울기 기하학 및 지속성 분석

Rock Token 이 유용한 학습 신호를 제공하는지 여부를 해결하기 위해 저자들은 토큰 단위 기울기의 기하학을 분석합니다:

• 크기 대 빈도: Rock Token 의 발생 단위 기울기 크기를 드문 고-KL 토큰과 비교합니다.
• 정렬: Rock Token 기울기와 빈도 균형 전역 하강 방향 ($G_{balanced}$) 간의 코사인 정렬을 측정합니다.
• 정체: 초기 및 후기 훈련 체크포인트 간의 토큰 단위 KL 발산 변화를 추적하여 손실이 실제로 감소하는지 확인합니다.

2.3 인과적 개입: 녹아웃 실험

이러한 토큰의 기능적 필요성을 판단하기 위해 저자들은 추론 시 녹아웃 (knockout) 절차를 적용합니다. 그들은 각 디코딩 단계에서 후보 토큰 $v$의 로그를 $-\infty$로 강제 설정하여 정책 $\pi_{\theta}^{\setminus v}$를 형성합니다. 인과적 영향은 벤치마크 정확도 변화 ($\Delta_B(v)$) 로 측정됩니다. 토큰은 다음과 같이 분류됩니다:

• 강력한 기둥 (Strong Pillars): 제거 시 정확도가 현저히 저하됨.
• 강력한 걸림돌 (Strong Stumbling Blocks): 제거 시 정확도가 향상됨.
• 중립: 제거 시 영향이 미미함.

2.4 선택적 증류 (기울기 동결)

Rock Token 이 훈련 과정에 미치는 영향을 평가하기 위해 저자들은 문맥 인식 토큰 재가중치 메커니즘을 도입합니다. OPD 훈련 중 Rock Token 과 관련된 지역 고발산 윈도우의 기울기를 동결 (가중치 $\lambda=0$ 설정) 합니다. 이 'Just Not Train' 전략을 통해 이러한 토큰이 필수적인 기능적 앵커인지 아니면 비생산적인 최적화 부담인지 분리하여 확인할 수 있습니다.

3. 주요 발견

3.1 정체성과 보편성

• 구성: Rock Token 은 주로 수학적/LaTeX 구분자, 마크다운/공백 구조, 담화 표지자 (예: "So", "Wait"), 그리고 숫자와 같은 문법적 및 구조적 발판입니다. 이들은 내용-bearing 단어가 아니라 추론의 경계를 표시합니다.
• 보편성: Rock Token 은 어휘의 약 6% 만을 차지함에도 불구하고 생성된 출력에서 **전체 토큰 빈도의 약 18%**를 차지합니다.

3.2 기울기 역설

• 정렬: 노이즈라는 가정과 달리, Rock Token 은 전역 최적화 목표와 비자명한 양의 방향 정렬을 보입니다. 그들은 올바른 최적화 경로를 가리키는 기울기를 제공합니다.
• 정체: 이러한 정렬에도 불구하고 Rock Token 은 훈련 내내 정체됩니다. 드문 고-KL 토큰이 상당한 감소를 보이는 것과 달리, 이들의 KL 발산은 유의미하게 감소하지 않습니다 (중앙값 $\Delta KL \approx 0$).
• 원인: 저자들은 적응형 옵티마이저 (Adam 등) 가 두꺼운 꼬리 분포를 가진 클래스 불균형으로 인해 이러한 보편적 토큰에 대한 유효 학습률을 억제하여, 올바른 방향 신호가 존재함에도 불구하고 학생을 기존 구조적 습관에 가두는 것으로 가정합니다.

3.3 인과적 필요성 (기둥 대 걸림돌)

• 기둥의 희소성: Rock Token 중 추론 성능을 저하시키는 '강력한 기둥'으로 작용하는 것은 극히 일부 (1.5%–3.5%) 입니다.
• 걸림돌의 부재: '강력한 걸림돌'은 발견되지 않았습니다. Rock Token 을 제거해도 성능이 향상된 적은 없습니다.
• 직교성: 토큰의 인과적 중요성 (기둥성) 은 엔트로피, 빈도, 잔여 KL 과 같은 표준 지표와 직교합니다. 기존 중요성 분류 체계는 어떤 토큰이 인과적으로 필요한지 예측하지 못합니다.
• 버퍼 존: 학생 모델은 전체 고손실 세트에 의존하기보다는 몇 가지 중요한 의미적 앵커를 보호하기 위해 지속적 잔여의 '버퍼 존'을 유지하는 것으로 보입니다.

3.4 훈련 효율성

• 선택적 동결: 훈련 중 Rock Token (및 해당 지역 윈도우) 의 기울기를 동결함으로써 저자들은 1.4 배에서 1.7 배의 훈련 월클락 속도 증가를 달성했습니다.
• 성능 영향: 이 가속화는 하류 추론 성능 (AIME, HMMT, MATH 벤치마크에서 측정) 에서 미미한 손실로 이어졌습니다. 이는 Rock Token 에 할당된 상당한 최적화 대역폭이 불필요함을 시사합니다.

4. 의의 및 주장

이 논문은 온-정책 증류에서 모든 고손실 토큰을 정렬하는 것이 학습에 필수불가결하다는 통념에 도전합니다. 저자들은 다음과 같이 주장합니다:

1. 숨겨진 동역학의 식별: 그들은 RLVR 에서 관찰된 '분기 토큰'과 구별되는 OPD 최적화 동역학의 숨겨진 측면을 드러내는 지속적 고손실 현상인 'Rock Token'을 식별합니다.
2. 현재 파이프라인의 비효율성: OPD 의 상당한 부분의 훈련 신호는 모델의 궁극적 추론 성능에 기대보다 훨씬 낮은 기여를 제공합니다.
3. 효율성 - 성능 프론티어: 이 연구는 이러한 '걸림돌'을 전략적으로 우회 (기울기 동결을 통해) 함으로써 정렬 과정을 크게 간소화할 수 있음을 보여줍니다. 이는 추론 무결성을 훼손하지 않으면서 대규모 모델 증류를 위한 더 효율적인 패러다임을 제공합니다.

이 논문은 Rock Token 이 완전히 폐기 가능한 것은 아니라는 점 (소수 하위 집합이 필수 기둥으로 작용함) 을 인정하면서도, 비균일 기울기 할당은 계산적 중복을 효과적으로 우회하여 온-정책 정렬의 효율성 - 성능 프론티어를 최적화할 수 있다고 결론지었습니다.

5. 한계

저자들은 그들의 실증적 평가가 주로 경쟁적 수학 추론에 의존하며, Rock Token 의 역할은 개방형 생성이나 코딩 작업에서는 다를 수 있다고 지적합니다. 또한, 이 연구는 특정 30B 에서 4B 로의 증류 페어링을 활용하므로, 서로 다른 아키텍처 패밀리나 규모에서의 일반화 가능성은 여전히 열린 질문으로 남아 있습니다. 마지막으로, 완화 전략은 엄격한 이진 기울기 동결 마스크를 사용하며, 저자들은 향후 연구에서 동적 재가중치 체계로 이를 정교화할 수 있음을 제안합니다.