기술 요약: OISD – 언어 모델의 온-폴리시 내부 자기 증류

1. 문제 제기

최근 대규모 언어 모델 (LLM) 을 위한 강화 학습 (RL) 후훈련 접근법, 특히 검증 가능한 보상을 활용한 그룹 상대적 정책 최적화 (GRPO) 를 사용하는 방법들은 수학적 추론을 크게 발전시켰습니다. 그러나 이러한 방법들은 주로 최종 출력 중심입니다. 이들은 최종 답변의 정확성과 같은 희소하고 결과 수준의 보상을 사용하여 최종 행동 정책을 최적화하지만, 중간 계층에 인코딩된 예측 신호, 추론 과정, 그리고 내부 표현을 크게 간과합니다.

지식 증류 (KD) 는 토큰 수준의 밀집된 감독을 제공할 수 있는 메커니즘을 제공하지만, 기존 패러다임들은 다음과 같은 특정 한계를 겪고 있습니다:

• 오프-폴리시 증류: 학생이 고정된 교사 모델이 생성한 궤적이 아닌 자신의 생성물로 훈련되기 때문에 분포 불일치를 초래합니다.
• 온-폴리시 증류: 불일치를 줄이지만 일반적으로 외부 교사 모델에 의존합니다.
• 온-폴리시 자기 증류: 최근 방법들은 외부 교사를 제거하지만, 모델을 감독하기 위해 외부에서 주입된 특권 정보(예: 특권 프롬프트, 정답 솔루션, 또는 보조 피드백) 에 의존합니다.

여기에는 외부 특권 데이터나 여러 행동 정책을 요구하지 않고 모델의 자신 내부 계산을 감독의 원천으로 활용하는 데 있어 간극이 존재합니다. 또한, 이전 연구는 추론 표현이 계층을 가로질러 체계적으로 진화한다는 것 (초기 계층의 광범위한 후보 예측에서 깊은 계층의 수렴된 결정으로) 을 시사하며, 이는 중간 계층이 추론의 일관성을 개선하기 위해 더 깊은 계층에 의해 안내될 수 있음을 나타냅니다.

2. 방법론: OISD 프레임워크

저자들은 OISD(온-폴리시 내부 자기 증류) 를 제안합니다. 이는 모델의 최종 계층을 "분리된 내부 교사"로, 선택된 중간 계층을 "보조 내부 학생"으로 취급하는 프레임워크입니다. 이는 롤아웃 및 GRPO 최적화 과정 중 단일 순전파 (forward pass) 내에서 발생합니다.

핵심 아키텍처

• 단일 행동 정책: 최종 계층 ($L$) 은 롤아웃 생성 및 GRPO 목적 함수 최적화를 위한 유일한 정책으로 남습니다.
• 내부 교사 - 학생: 동일한 순전파 동안, 선택된 중간 계층 ($l < L$) 의 은닉 상태는 언임베딩 행렬을 사용하여 토큰 공간으로 투영됩니다 (logit lens 프레임워크를 따름). 이는 "교사"(최종 계층) 와 동일한 문맥에서 파생된 임시 "학생" 분포를 생성합니다.
• 외부 특권 없음: 감독 신호는 동일한 온-폴리시 롤아웃에서 모델의 자체 내부 상태로부터 완전히 파생되며, 외부 정답 솔루션이나 특권 프롬프트가 필요하지 않습니다.

정렬 목표

OISD 는 최종 계층에서 중간 계층으로 예측 신호를 전달하기 위해 두 가지 보완적인 내부 정렬 메커니즘을 사용합니다:

1. 로그잇 정렬 ("어떻게 생각할 것인가"):

   • 목표: 고수준 추론 행동 및 예측 신념을 전달합니다.
   • 메커니즘: 중간 계층 ($p^l_t$) 과 최종 계층 ($q^L_t$) 의 온도 스케일링된 토큰 분포 간에 제센 - 샤논 (JS) 발산을 계산합니다.
   • 가중치: 손실은 서명된 GRPO 이득 ($A_t$) 으로 가중됩니다. 높은 이득을 가진 토큰은 중간 계층을 최종 계층의 분포 쪽으로 끌어당기는 반면, 낮은 이득을 가진 토큰은 나쁜 연속을 강화하지 않도록 밀어냅니다.
   • 공식: $\mathcal{L}^{(l)}_{\text{think}} = \mathbb{E}_t [ A_t \cdot \text{JS}(p^{l,\tau}_t, \text{sg}(q^{L,\tau}_t)) ]$

2. 어텐션 정렬 ("어디를 볼 것인가"):

   • 목표: 일관된 어텐션 패턴 및 증거 라우팅 행동을 강제합니다.
   • 메커니즘: 샘플링된 인과 키 집합 (메모리 관리를 위해 로컬 윈도우 및 글로벌 스트라이드 위치를 결합) 에 대한 중간 계층과 최종 계층의 어텐션 가중치 간에 JS 발산을 계산합니다.
   • 가중치: 로그잇 정렬과 유사하게, 손실은 서명된 이득으로 가중됩니다. 양의 이득은 중간 계층이 최종 계층과 동일한 증거에 주의를 기울이도록 장려하고, 음의 이득은 낮은 보상 패턴에 대한 합의를 줄입니다.
   • 공식: $\mathcal{L}^{(l)}_{\text{attn}} = \mathbb{E}_t [ A_t \cdot \frac{1}{H} \sum_{h=1}^H \text{JS}(\tilde{a}^l_{t,h}, \text{sg}(\tilde{a}^L_{t,h})) ]$

훈련 목적 함수

총 손실은 표준 GRPO 손실과 두 가지 내부 정렬 손실을 결합합니다:
$$\mathcal{L}_{\text{OISD}} = \mathcal{L}_{\text{GRPO}} + \lambda_{\text{think}} \mathcal{L}_{\text{think}} + \lambda_{\text{attn}} \mathcal{L}_{\text{attn}}$$
여기서 $\lambda_{\text{think}}$와 $\lambda_{\text{attn}}$은 증류의 강도를 조절합니다.

3. 주요 기여

1. 새로운 패러다임: 단일 행동 정책을 유지하면서 모델의 자체 내부 계산을 감독에 활용하는 추론 RL 패러다임인 온-폴리시 내부 자기 증류의 도입.
2. OISD 프레임워크: 외부 특권 정보 없이 온-폴리시 예측 신호를 전달하기 위해 분리된 교사를 최종 계층으로, 학생을 중간 계층으로 사용하는 구체적인 구현.
3. 이중 정렬 메커니즘: 두 가지 보완적 목적 함수의 개발:
   • 로그잇 정렬: 추론 과정 ("어떻게 생각할 것인가") 을 감독합니다.
   • 어텐션 정렬: 증거 라우팅 ("어디를 볼 것인가") 을 감독합니다.
4. 실증적 검증: 네 가지 수학적 추론 벤치마크에서 강력한 추론 RL 베이스라인 (GRPO 및 BuPO 포함) 대비 상당하고 일관된 개선 효과 입증.

4. 실험 결과

저자들은 Qwen3-4B, Qwen3-8B, OctoThinker 모델을 AMC23, MATH500, AIME24, AIME25의 네 가지 벤치마크에 대해 평가했습니다.

• 성능 향상:
  • Qwen3-4B에서 OISD 는 평균 점수 64.75를 달성하여 GRPO(55.08) 보다 9.67 점, BuPO(58.51) 보다 6.24 점 더 높은 성과를 거두었습니다.
  • AIME 벤치마크에서 개선이 특히 두드러졌습니다 (예: AIME24 에서 GRPO 대비 +14.31 점).
  • Qwen3-8B에서 OISD 는 66.72의 최고 전체 평균을 달성하여 GRPO(64.23) 와 BuPO(66.36) 를 능가했습니다.
• Pass@K 분석: OISD 는 모든 $K$값에서 GRPO 보다 일관되게 높은 Pass@K 비율을 달성하여, 평균 정확도뿐만 아니라 샘플링된 솔루션 집합의 전반적인 품질이 개선되었음을 나타냅니다.
• 애블레이션 연구:
  • 로그잇 및 어텐션 정렬 모두 긍정적인 기여를 했습니다. 로그잇 정렬은 AIME24 에서 가장 큰 개선을 보였으며, 어텐션 정렬은 평균적으로 약간 더 좋은 성능을 보였습니다.
  • 두 가지를 결합했을 때 가장 좋은 결과가 나왔으며, 이는 그들의 보완적 성격을 확인시켜 줍니다.
• OPSD 와의 비교: OISD(64.75) 는 특권 정답 솔루션에 의존하는 최첨단 온-폴리시 자기 증류 방법인 OPSD(56.63) 보다 현저히 우수했습니다. 이는 OISD 의 순수 내부 감독의 효과성을 강조합니다.
• 정성적 분석:
  • 훈련 역학: OISD 는 안정화되기 전에 꾸준한 보상 향상과 초기 엔트로피 증가 (더 넓은 탐색을 시사) 를 보였습니다.
  • 어텐션 시각화: 훈련 중 OISD 의 중간 계층은 점차 중요한 증거 토큰 (예: 끝점, 상한) 에 대한 최종 계층의 어텐션과 정렬되었습니다.
  • 로그잇 - 렌즈 시각화: OISD 는 종종 단편적인 중간 예측을 생성한 GRPO 에 비해 추론에 중요한 토큰에서 중간 및 최종 예측 간의 더 빠른 정렬을 보였습니다.

5. 중요성 및 주장

이 논문은 OISD 가 온-폴리시 내부 자기 증류 조사를 시작한다고 주장합니다. 그 주요 중요성은 다음과 같습니다:

• 외부 의존성 제거: 훈련 중 외부 교사, 특권 프롬프트, 또는 정답 솔루션 없이 밀집된 내부 감독을 제공합니다.
• 정책 일관성: 롤아웃 생성 및 최적화를 위해 단일 행동 정책을 유지하여 오프-폴리시 방법의 분포 불일치 문제를 피합니다.
• 향상된 추론: 중간 표현을 최종 계층의 추론 및 어텐션 패턴과 정렬함으로써, OISD 는 모델이 더 강력하고 전역적으로 일관된 중간 표현을 학습하도록 장려합니다.

저자들은 검증 가능한 답변을 가진 수학적 추론에 초점을 맞추고 최적의 계층 선택에 대한 심층 연구가 부족하다는 한계를 지적합니다. 향후 작업은 더 복잡한 추론 작업으로의 일반화 및 다중 계층 내부 감독을 탐구해야 한다고 제안합니다.