기술 요약: $\beta$-OPSD

1. 문제 정의

온-폴리시 자기 증류(on-policy self-distillation, OPSD)는 대규모 언어 모델(LLM)의 추론 능력을 향상시키기 위한 유망한 기술로 부상했습니다. 표준 OPSD에서 학생 모델은 자신의 샘플링된 궤적(trajectory)으로부터 학습하는 동시에, 훈련 중에만 사용 가능한 "특권 교사(privileged teacher)"(정답 솔루션 또는 검증된 흔적에 접근할 수 있는 동일한 아키텍처의 모델)로부터 감독을 받습니다.

그러나 저자들은 일반적인 OPSD가 실제 적용 시 매우 취약하며, 안정적으로 작동시키기 위해 상당한 엔지니어링 노력이 필요하다는 점을 지적합니다. 핵심 문제는 표준 OPSD가 교사를 유일한 최적화 대상으로 취급하여, 학생 모델을 초기 또는 참조 정책(reference policy)에 고정시키는 정규화 강도를 암묵적으로 고정시킨다는 점입니다. 이러한 경직된 공식은 학생 모델이 참조 정책으로부터 얼마나 공격적으로 교사에게 이동할지를 제어할 수 있는 메커니즘이 부족하며, 이는 종종 불안정성이나 하위 최적해(suboptimal convergence)로 이어집니다.

2. 방법론

본 논문은 $\beta$-OPSD를 제안하며, 이는 온-폴리시 자기 증류를 더 넓은 범위의 KL-정규화된 정책 최적화 문제의 가족(family)으로 재정의합니다. 방법론은 세 단계로 진행됩니다: 이론적 타겟 도출, 이를 실용적인 증류 목적 함수로 변환, 그리고 크레딧 할당(credit assignment)의 정교화.

2.1. 이론적 도출: 정책 최적화로서의 $\beta$-OPSD

저자들은 일반적인 OPSD가 정확히 $\beta = 1$인 KL-정규화된 정책 최적화 계열의 일원임을 입증합니다. 이들은 목적 함수 $J_\beta(\theta)$를 다음과 같이 정의합니다:
$$J_\beta(\theta) = \mathbb{E}_{y \sim \pi_\theta} \left[ \log \frac{p_T(y|x,c)}{\pi_{\text{ref}}(y|x)} \right] - \beta D_{KL}(\pi_\theta(\cdot|x) \parallel \pi_{\text{ref}}(\cdot|x))$$
여기서:

• $\pi_\theta$: 학생 정책.
• $p_T$: 특권 교사 정책 (문제 $x$와 특권 정보 $c$에 조건화됨).
• $\pi_{\text{ref}}$: 참조 정책 (예: 초기 학생 모델).
• $\beta$: 조절 가능한 정규화 파라미터.

핵심 통찰: $\beta=1$일 때, 참조 항들이 상쇄되어 교사를 향한 표준 역방향 KL 최소화(reverse KL minimization)를 회복합니다. $\beta > 1$인 경우, 목적 함수는 더 보수적이 되어 참조 정책으로부터의 이탈에 대해 더 강력하게 페널티를 부여합니다.

2.2. 최적 정책 및 기하학적 보간(Geometric Interpolation)

저자들은 이 목적 함수에 대한 폐쇄형 최적 정책 $\pi^\star_\beta$를 도출합니다:
$$\pi^\star_\beta(y|x,c) = \frac{\pi_{\text{ref}}(y|x)^{1-1/\beta} p_T(y|x,c)^{1/\beta}}{Z_\beta(x,c)}$$
이 해법은 참조 정책과 특권 교사 사이의 기하학적 보간을 나타냅니다. 파라미터 $\beta$(특히 그 역수인 $1/\beta$)는 이 경로상의 위치를 결정하는 가중치 역할을 합니다.

• $\beta = 1 \implies$ 타겟은 교사입니다.
• $\beta \to \infty \implies$ 타겟은 참조 정책에 근접합니다.

2.3. 실용적 훈련: 스케줄링된 로짓 보간(Scheduled Logit Interpolation)

시퀀스 레벨의 정규화 상수(normalizer) $Z_\beta$가 다루기 어렵고 그래디언트 분산이 높기 때문에, RL 목적 함수를 직접 최적화하는 것은 계산 비용이 많이 듭니다. 대신, $\beta$-OPSD는 이론적 해법을 효율적인 증류 타겟으로 변환합니다.

• 국소 근사(Local Approximation): 시퀀스 레벨의 정규화 상수는 다루기 어렵기 때문에, 이 방법은 각 디코딩 접두사(prefix) $h_t$에서 참조와 교사의 **로짓(logits)**을 선형적으로 보간하여 최적 정책을 근사합니다:
  $$\tilde{p}_{\beta_k}(\cdot | h_t, c) = \text{softmax}\left( \left(1 - \frac{1}{\beta_k}\right) z_{\text{ref}} + \frac{1}{\beta_k} z_T \right)$$
• 커리큘럼 스케줄링(Curriculum Scheduling): $\beta$를 고정하는 대신, 이 방법은 훈련 단계에 따라 교사 가중치 $w_k = 1/\beta_k$가 $w_{\text{start}}$에서 $w_{\text{end}}$까지 선형적으로 증가하는 스케줄을 채택합니다. 이는 학생에게 더 가까운 타겟에서 시작하여 점진적으로 더 강력한 교사의 가이드를 통합하는 매끄러운 커리큘럼을 생성합니다.

2.4. Return-to-Go 크레딧 할당

표준 토큰 레벨 OPSD 업데이트는 각 토큰의 손실을 독립적으로 처리하는 "근시안적(myopic)"인 특성을 가집니다. 이는 초기 토큰의 선택이 미래의 접두사와 그에 따른 타겟과의 불일치에 어떻게 영향을 미치는지 무시합니다.
이를 해결하기 위해, $\beta$-OPSD는 Return-to-go (RTG) 크레딧 할당을 도입합니다. 토큰 $t$에 대한 손실은 궤적을 따르는 누적 미래 불일치(할인된 수익)에 의해 가중치가 부여됩니다:
$$G_{t, \gamma} = \sum_{s=t}^{T} \gamma^{s-t} \rho_s$$
여기서 $\rho_s$는 단계 $s$에서의 로그 비율 불일치입니다. 이를 통해 초기 토큰이 시퀀스 후반부에 유발하는 분포적 불일치에 대해 책임을 지거나(credit) 비난받도록(blame) 보장하며, 토큰 레벨의 업데이트를 시퀀스 레벨의 목적 함수와 정렬시킵니다.

3. 주요 기여

1. 이론적 통합: 본 논문은 일반적인 OPSD가 더 넓은 KL-정규화된 정책 최적화 계열의 $\beta=1$ 케이스임을 확립하여, 암묵적인 정규화 선택을 명시적이고 제어 가능하게 만들었습니다.
2. $\beta$-OPSD 프레임워크: 최적 정책이 참조와 교사 사이의 기하학적 보간임을 보여주는 일반화된 목적 함수를 도입하였으며, 이를 스케줄링된 로짓 보간을 통해 실용적으로 구현하였습니다.

• 효율적인 훈련 알고리즘: 복잡한 어드밴티지 추정(advantage estimation)과 같은 까다로운 RL 메커니즘을 피하면서, 저렴한 토큰 레벨 증류를 사용하여 값비싼 RL 유도 해법을 근사하는 방법을 제안하였습니다.

4. 시퀀스 레벨 정렬: RTG 크레딧 할당을 통합하여 표준 토큰-KL 업데이트의 근시안적 특성을 교정하고, 시퀀스 레벨 목적 함수에 대한 편향되지 않은 추정치를 제공합니다.

4. 실험 결과

이 방법은 Qwen3 모델(1.7B에서 8B 파라미터 범위)을 사용하여 수학적 추론 벤치마크(AIME 2024, AIME 2025, HMMT 2025)에서 평가되었습니다.

• 성능 향상: $\beta$-OPSD는 vanilla OPSD, SFT, GRPO 베이스라인보다 일관되게 우수한 성능을 보였습니다.
  • Qwen3-1.7B에서, $\beta$-OPSD는 vanilla OPSD 대비 평균 성능(avg@12)을 5.74 퍼센트 포인트 향상시켰으며, AIME 2024에서는 9.16 포인트의 구체적인 이득을 보였습니다.
  • 더 큰 모델(4B 및 8B)에서도 경쟁력 있거나 우월한 성능을 유지하며, 각각 1.76 및 1.66 포인트를 개선했습니다.
• 절제 연구(Ablation Studies):
  • 타겟 분포: 고정된 교사 타겟을 스케줄링된 로짓 보간 타겟으로 교체했을 때 상당한 이득(예: AIME 2024에서 +6.03 포인트)을 얻었으며, 이는 매끄러운 전환의 이점을 확인시켜 줍니다.
  • 크레딧 할당: 로컬 토큰 그래디언트 대신 return-to-go 가중치를 사용했을 때 성능이 향상되었습니다(예: AIME 2025에서 +5.55 포인트). 이는 미래를 고려한 업데이트의 중요성을 입증합니다.
  • 참조 구축: 가장 좋은 성능은 동적 학생(현재 모델)과 고정된 교사(특권 교사)를 사용할 때 달성되었습니다. 이는 진화하는 학생 분포를 추적하면서도 안정적인 감독 앵커를 유지하는 것이 최적임을 시사합니다.

5. 의의 및 주장

본 논문은 "자기 증류에서 정책 최적화로, 그리고 다시 돌아오는 원칙적인 경로"를 제공한다고 주장합니다. 주요 의의는 다음과 같습니다:

• 안정성: $\beta$를 제어 가능한 파라미터로 취급함으로써, $\beta$-OPSD는 vanilla OPSD의 취약성을 해결하고, 참조 정책에 대한 근접성과 교사의 가이드 사이의 절충을 허용하는 매끄러운 최적화를 가능하게 합니다.
• 효율성: 완전한 RL의 계산 비용 없이도 정책 최적화(시퀀스 레벨 목적 함수, KL 정규화)의 이점을 얻으면서, 표준 증류의 효율성을 유지합니다.
• 일반화 가능성: 이 접근 방식은 학생과 교사 사이의 분포 격차를 줄이는 메커니로를 제공하며, 이는 특히 이 격차가 큰 작은 모델들에게 유익합니다.

저자들은 매끄러운 타겟 경로와 궤적 레벨의 크레딧 할당이 온-폴리시 증류을 위한 효과적인 원칙이며, 직접적인 교사 모방에 대한 강력한 대안을 제공한다고 결론짓습니다.