기술 요약: 온-폴리시 증류(On-Policy Distillation)의 위치 편향성에 대하여

문제 정의

온-폴리시 증류(On-Policy Distillation, OPD)는 학생 모델이 생성한 자체 롤아웃(rollout)으로부터 교사 모델의 밀도 높은 토큰 수준 감독을 제공함으로써 대규모 언어 모델(LLM)을 위한 강화 학습(RL)의 학습 효율성을 높입니다. 그러나 표준적인 OPD 목적 함수는 일반적으로 모든 토큰에 동일한 가치를 부여한다고 가정하며, 토큰 수준의 손실(주로 KL 발산)을 균등하게 평균화합니다.

저자들은 이 접근 방식에서 결정적인 **위치 편향(Position Bias)**을 식별했습니다. 학생의 롤아웃이 길어질수록 생성된 궤적(trajectory)은 교사의 분포로부터 점점 더 멀어지게 됩니다. 결과적으로, 시퀀스의 후반부로 갈수록 교사의 감독 품질이 저하됩니다. 경험적 관찰 결과는 다음과 같습니다:

1. 감독 가치의 비대칭성: 처음 30%의 토큰(접두사, prefix)만을 사용하여 학습하는 것이 전체 토큰을 사용하는 것과 비슷하거나 더 나은 성능을 보인 반면, 마지막 30%의 토큰(접미사, suffix)을 사용하는 것은 미미한 학습 이득만을 제공했습니다.
2. 지속적인 괴리: 수렴 후에도 교사와 학생 사이의 KL 발산은 여전히 유의미하게 유지되는데, 이는 OPD가 정책 공간의 좁은 국소 영역 내에서만 학생을 최적화하고 있음을 시사합니다.
3. 비효율성: 모든 토큰에 균등한 가중치를 부여하는 것은 학생이 이미 교사에 대해 분포 외(Out-of-Distribution, OOD) 영역으로 이탈한 후반부 토큰에 한정된 그래디언트 업데이트 예산을 낭비하게 만들며, 이로 인해 감독 신호가 노이즈가 심하거나 도움이 되지 않게 됩니다.

방법론: 중요도 가중치 적용 온-폴리시 증류 (IW-OPD)

이러한 편향을 해결하기 위해 저자들은 학생과 교사 분포 사이의 누적된 불일치에 따라 토큰 가중치를 동적으로 조정하는 방법인 **중요도 가중치 적용 온-폴리시 증류(Importance-Weighted On-Policy Distillation, IW-OPD)**를 제안합니다.

이론적 기초

저자들은 OPD를 KL 발산 예산 $\rho$로 정의된 신뢰 영역 내에서 학생 정책이 업데이트되는 제약 최적화 문제(국소 투영)로 프레임화합니다.

• 최적 정책 도출: $D_{KL}(q \| \pi_T)$를 최소화하면서 $D_{KL}(q \| \pi_\theta) \leq \rho$라는 제약을 만족하는 문제를 풀면 폐쇄형 최적 정책 $q^\star$를 얻을 수 있습니다.
• 우도 비율 재가중치(Likelihood Ratio Reweighting): 솔루션 $q^\star$는 학생 정책 $\pi_\theta$를 우도 비율 $(\pi_T / \pi_\theta)^\alpha$로 재가중합니다. 이는 교사의 확률이 학생보다 현저히 높은 토큰(높은 호환성)은 더 높은 가중치를 받아야 하는 반면, 학생이 이탈한 토큰(낮은 비율)은 가중치를 낮추어야 함을 의미합니다.

실제 구현

원시 우도 비율을 직접 사용하는 것은 긴 시퀀스에 걸친 확률의 지수적 누적으로 인해 수치적으로 불안정할 수 있습니다. 따라서 IW-OPD는 안정적인 **부호 없는 누적 접두사 불일치(unsigned cumulative prefix discrepancy)**를 중요도 가중치의 대리 지표로 도입합니다.

1. 누적 불일치: 이 방법은 원시 비율 대신, 시퀀스의 시작부터 현재 토큰 $t$까지의 절대적인 토큰 수준 어드밴티지($|A^{OPD}_k|$)를 합산합니다. 이는 교사가 선호하는 경로로부터의 총 "이탈(drift)"을 포착합니다.
2. 정규화: 누적된 불일치는 각 샘플 내에서 $[0, 1]$ 범위로 표준화되어 안정적인 스케일링을 보장합니다.
3. 보간(Interpolation): 이탈한 토큰에 대한 감독을 완전히 제거하는 것을 방지하기 위해, 최종 가중치는 표준 OPD 신호(가중치 1)와 불일치로부터 도출된 중요도 가중치 사이의 보간법을 사용합니다.
   • 공식: $\tilde{r}^{IW-OPD}_t = 1 + \gamma \cdot (1 - \frac{\sum_{k<t} |A^{OPD}_k|}{\sum_{k<T} |A^{OPD}_k|})$.
   • 여기서 $\gamma$는 중요도 가중치 적용의 기여도를 조절합니다.

이 접근 방식은 자연스럽게 초기 토큰(접두사가 교사와 호환되는 부분)에는 가중치를 높이고, 후반부 토큰(누적 이탈이 큰 부분)에는 가중치를 낮추어, 가장 정보가 풍부한 궤적 부분으로 그래디언트 예산을 효과적으로 재배분합니다.

주요 기여

1. 위치 편향성 식별: 본 논문은 분포 이탈로 인해 초기 토큰이 후기 토큰보다 훨씬 더 가치 있는 감독을 제공한다는 위치 편향성을 경험적으로 입증하고 이론적으로 설명합니다.
2. 제약 최적화 관점: 저자들은 OPD를 제약된 국소 투영으로 보는 원칙적인 이해를 제공하며, 최적의 업데이트 전략이 본질적으로 비율 기반 재가중치를 필요로 한다는 것을 도출했습니다.
3. IW-OPD 알고리즘: 누적 접두사 불일치를 사용하여 토큰 중요도를 추정하는 실용적이고 안정적인 알고리즘인 IW-OPD를 제안합니다. 이는 표준 OPD 외에 추가적인 교사 평가를 요구하지 않습니다.
4. 경험적 검증: 다양한 모델 규모와 교사-학생 구성에 대해 IW-OPD가 표준 OPD보다 더 빠르게 수렴하고 더 높은 최종 성능을 달여냄을 광범위한 실험을 통해 보여줍니다.

실험 결과

저자들은 수학적 추론(AIME-2024, AIME-2025, HMMT-2025) 및 코드 생성(HumanEval+, MBPP+) 벤치마크에서 Qwen3 모델(0.6B, 1.7B, 4B, 30B, 235B)을 교사 및 학생으로 사용하여 IW-OPD를 평가했습니다.

• 성능 향상: IW-OPD는 일관되게 표준 OPD보다 우수한 성능을 보입니다.
  • Qwen3-30B 교사를 사용하는 Qwen3-1.7B 학생 설정에서, IW-OPD는 10단계에서 AIME-2025 정확도를 +6.9 포인트, 수렴 시 +1.7 포인트 향상시켰습니다.
  • 모든 학생 규모(0.6B ~ 30B)와 교사 설정에 걸쳐 IW-OPD는 평균 성능 점수를 개선했습니다.
• 샘플 효율성: IW-OPD는 현저히 빠르게 수렴합니다. 여러 경우에서 10단계에서의 IW-OPD 체크포인트가 표준 OPD의 최종 성능과 일치하거나 이를 상회했습니다.
• 스케일링 효과:
  • 강력한 교사: IW-OPD는 강력한 교사의 효율성을 높입니다. 표준 OPD는 초기 단계에서 더 큰 교사-학생 간 격차로 인해 어려움을 겪는 반면(더 강력한 교사의 이점을 얻기 위해 더 많은 단계가 필요함), IW-OPD는 즉각적으로 강력한 교사의 신호를 활용합니다.
  • 작은 학생: IW-OPD의 상대적 이득은 학생 모델이 작아질수록 커집니다(예: 4B 학생의 +1.8% 대비 0.6B 학생의 +9.5% 상대적 이득). 이는 이 방법이 이탈이 더 심한 교차 규모 증류(cross-scale distillation)에서 특히 효과적임을 나타냅니다.

의의 및 주장

본 논문은 효과적인 온-폴리시 감독은 단순히 토큰 수준의 불일치뿐만 아니라 궤적의 문맥(context)을 고려해야 한다고 주장합니다. 교사의 감독이 학생이 생성한 궤적을 따라 항상 균등하게 신뢰할 수 있는 것은 아니라는 점을 인식함으로써, IW-OPD는 유한한 그래디언트 예산을 "교사와 호환되는" 접두사에 재배분합니다.

저자들은 학생이 교사의 분포에서 이탈한 토근의 가중치를 낮추는 이 간단한 수정이 온-폴리시 증류의 샘플 효율성과 최종 성능을 모두 크게 향상시킨다고 결론짓습니다. 이러한 통찰은 향후 온-폴리시 방법들이 궤적 내의 모든 토큰을 독립적이고 동일하게 가치 있는 것으로 취급하기보다, 롤아웃의 누적된 이력(접두사)을 고려하여 신용 할당이나 감독 가중치를 부여해야 함을 시사합니다.