기술 요약: 온-폴리시 증류(On-Policy Distillation)에서의 감독 충실도 저하(Supervision Fidelity Decay) 극복

문제: 감독 충실도 저하 (Supervision Fidelity Decay, SFD)

본 논문은 학생 모델이 교사 모델으로부터 제공되는 토큰 수준의 피드백을 사용하여 자신이 생성한 궤적(trajectory)으로부터 학습하는 패러다임인 **온-폴리시 증류(On-Policy Distillation, OPD)**의 핵심적인 병목 현상을 식별합니다. OPD는 추론 능력을 전이하는 데 효과적이었으나, 저자들은 생성 길이(예: 16k 토큰에서 성능 정체, 39k 토ques에서 하락)가 길어짐에 따라 성능이 크게 저하되는 현상을 관찰했습니다.

저자들은 이를 감독 충실도 저하(SFD) 때문이라고 설명합니다. 학생이 더 긴 시퀀스를 생성할수록, 그 출력은 점진적으로 교사의 훈련 분포(Out-of-Distribution 또는 OOD)에서 벗어나게 됩니다. 이러한 드리프트(drift)는 교사의 다음 토큰 분포를 덜 확신하게 만들고 변별력을 떨어뜨립니다.

• 메커니즘: 표준 역-KL(reverse-KL) 증류에서 학습 신호(advantage)는 학생의 로그 확률과 교사의 로그 확률 사이의 차이에 의존합니다. SFD로 인해 교사의 분포가 확산(고엔트로피)되면, 교사의 로그 확률의 분산이 붕괴됩니다.
• 결과: 교사의 변별적 기여도가 사라집니다. 학습 신호는 사실상 학생의 기존 모드(mode)를 강화하기만 하는 "학생 전용" 신호로 축소됩니다. 이는 드리프트가 누적되는 자기 강화 피드백 루프를 생성하여, 유효 학습 수평선을 축소시키고 추론 길이에 대한 한계를 설정합니다.

방법론: Lookahead Group Reward (LGR)

SFD를 해결하기 위해, 저자들은 현재의 국소적 일치(local matching)가 아닌 미래의 교사 확신도를 최적화함으로써 감독 품질을 능동적으로 유지하는 방법인 **Lookahead Group Reward (LGR)**를 제안합니다.

핵심 통찰

교사의 현재 위치 $t$에서의 국소적 감독이 붕괴되었더라도(즉, 분포가 균등에 가까워지더라도), $t$에서의 서로 다른 토큰 선택이 다음 위치($t+1$)에서의 교사 확신도에 유의미하게 다른 영향을 미칠 수 있습니다. 본 논문은 **한 단계 앞선 변별력(one-step-ahead discriminability)**이 국소적 감독이 실패한 상황에서도 생존할 수 있음을 이론적으로 증명합니다. 다음 단계에서 교사의 확신도를 극대화하는 토큰을 선택함으로써, 학생은 드리프트 속도를 늦출 수 있습니다.

알고리즘 설계

1. 확신도 보상 (Confidence Reward): 각 위치 $t$에서, 이 방법은 학생의 상위-$K$ 후보 토큰들에 대해 교사의 최대 다음 토큰 확률($\max_v \pi_T(v|x_{<t}, x_t)$)을 평가합니다. 이 값은 후보 토근이 다음 단계에서 고품질의 감독을 제공할 수 있는 능력을 얼마나 잘 보존하는지를 나타내는 보상 신호 역할을 합니다.
2. 그룹 정규화 (Group Normalization): 다양한 작업과 위치에 걸친 절대적 확신도 수준의 변동을 처리하기 위해, 원시 보상은 (GRPO와 유사하게) 상위-$K$ 후보 그룹 내에서 정규화됩니다. 이는 절대적인 스케일 노이즈를 제거하면서 후보들의 상대적 순위를 유지합니다.
3. 엔트로피 트리거 활성화 (Entropy-Triggered Activation): 모든 토큰에 대해 lookahead 보상을 계산하는 것은 계산 비용이 많이 듭니다. 저자들은 SFD와 lookahead 가이드의 필요성이 가장 중요한 지점이 학생의 생성 엔트로피가 높을 때(학생이 불확실할 때)임을 관찰했습니다. 따라서 엔트로피가 임계값 $\tau$를 초과하는 위치에서만 확신도 보상을 계산합니다. 이를 통해 계산 오버헤드를 4~7배 줄입니다.
4. 트리-어텐션 메커니즘 (Tree-Attention Mechanism): 매 후보마다 교사의 prefill을 다시 실행하지 않고 여러 위치의 $K$개 후보를 효율적으로 평가하기 위해, 저자들은 트리-어텐션 메커니즘을 사용합니다. 이는 학생의 궤적을 메인 브랜치로 하고 후보 토큰들을 브랜치로 추가하는 확장된 시퀀스를 구축합니다. 희소 어텐션 마스크(sparse attention mask)를 통해 교사가 단 한 번의 순방향 패스(forward pass)로 모든 후보를 처리할 수 있게 하여, 나이브한 구현 대비 이론적으로 최대 1000배의 속도 향상을 달면합니다.

최종 손실 함수는 표준 역-KL advantage와 하이퍼파라미터 $\gamma$로 가중치가 부여된 lookahead 확신도 보상을 결합합니다.

주요 기여

1. SFD 식별: 본 논문은 감독 충실도 저하를 온-폴리시 역-KL 증류의 근본적인 실패 모드로 공식적으로 규정하며, 이것이 어떻게 그래디언트를 자기 강화적 드리프트 루프로 붕괴시켜 추론 길이를 제한하는지 증명합니다.
2. Lookahead Group Reward (LGR): 저자들은 한 단계 앞을 내다보는 원칙적인 해결책을 제안합니다. 저하된 국소 분포를 맞추려고 노력하는 대신, LGR은 미래의 교사 확신도를 극대화하는 토큰을 선택하여 학생이 감독이 여전히 신뢰할 수 있는 상태로 향하도록 유도합니다.
3. 효율적인 구현: 엔트로피 트리거 트리-어텐션 메커니즘 설계는 계산 집약적인 lookahead 과정을 실용적으로 만들어, 과도한 비용 없이 긴 컨텍스트(최대 39k 토큰)에서 학습을 가능하게 합니다.

실험 결과

이 방법은 수학적 추론(AIME-24, AIME-25, AIME-26, HMMT-25, HMMT-26)과 코드 생성(LiveCodeBench v6)을 포함하는 6개의 벤치마크에서 평가되었습니다.

• 성능 향상: 7B 학생 모델에 대해, LGR은 9k 토큰 기준 표준 OPD 대비 mean@8 점수를 평균 2.57 포인트 개선했습니다.
• Long-Context 스케일링: LGR의 이점은 생성 길이에 따라 증가합니다. 39k 토큰의 AIME-26 벤치마크에서 LGR은 OPD 대비 +4.92 포인트의 이득을 달성했습니다. 반면, 표준 OPD의 성능은 이 길이에서 크게 저하되었습니다.
• 베이스라인과의 비교: LGR은 JSD, REOPOLD 및 표준 OPD 변형을 포함한 다른 증류 방법들보다 우수한 성능을 보였습니다. 특히, (단순히) 훈련을 안정화할 뿐(REOPOLD와 같이) 감독 품질 저하를 해결하지 못하는 방법들은 LGR의 성능을 따라잡지 못했습니다.
• 절제 연구 (Ablation Studies):
  • Max-p vs. Entropy: 교사의 최대 확률(Max-p)을 확신도 지표로 사용하는 것이 엔트로피나 퍼플렉서티(perplexity)를 사용하는 것보다 우수했으며, 이는 Max-p가 교사의 분포 내 매니폴드(in-distribution manifold)에 대한 직접적인 대리 지표라는 이론적 발견과 일치합니다.
  • 가중치 민감도: 확신도 보상 가중치 $\gamma$는 1.0에서 1.5 사이의 범위에서 최적의 성능을 보였습니다. 가중치가 지나치게 높으면 보상 해킹(reward hacking, 반복적인 시퀀스 생성)이 발생했습니다.

의의 및 주장

본 논문은 LGR이 증류에서의 장기 호라이즌 추론 문제에 대해 구조적으로 다른 솔루션을 제공한다고 주장합니다. 약한 감독 영역을 수동적으로 피하거나(예: 절단) 단순히 최적화를 안정화하는 전략과 달리, LGR은 학생을 교사가 높은 감독 충실도를 유지하는 궤적으로 능동적으로 유도합니다.

저자들은 자신들의 연구가 긴 추론 체인에서 발생하는 성능 저하의 근본 원인, 즉 학생의 드리프트가 교사의 확신도에 미치는 누적 효과를 다루고 있음을 강조합니다. "높은 다음 위치 확신도는 더 변별력 있는 미래 신호를 의미한다"는 통찰을 실행에 옮김으로써, LGR은 감독 저하의 악순환을 끊고 모델이 표준 온-폴리시 증류보다 훨씬 더 긴 추론 체인(최대 39k 토큰)까지 효과적으로 확장할 수 있도록 합니다.