기술 요약: 온폴리시리 셀프 디스틸레이션을 위한 특권적 은닉 흐름 (Privileged Hidden Flow)

문제 정의

온폴리시 셀프 디스틸레이션(On-policy self-distillation, OPSD)은 추론 모델을 학습시키는 기법으로, 학생 모델이 자신의 롤아웃(rollout)을 샘플링하고, "특권적(privileged)" 교사로부터 감독을 받는 방식이다. 여기서 교사는 동일한 아키텍처를 가지되, 검증된 참조 솔루션과 같은 추가적인 정보로 조건화된 모델이며, 학생은 오직 문제 프롬프트만을 본다.

기존의 OPSD 목적 함수는 학생과 특권적 교사 사이의 출력 분포(다음 토큰 확률)를 일치시키는 데에만 전적으로 집중한다. 이는 최종 예측을 정렬하는 데에는 효과적이지만, 그 예측을 생성하는 내부 연산을 감독하는 데에는 실패한다. 구체적으로, 특권적 컨텍스트는 네트워크 레이어를 통과하는 중간 은닉 상태(hidden states)와 이 상태들의 궤적(trajectory)에 영향을 미친다. 현재의 방법들은 최종 출력 헤드만을 제약하며, 모델의 내부 표현이 단계별로 어떻게 진화하는지를 나타내는 "은닉 흐름(hidden flow)"은 제약하지 않는다. 본 논문은 이 내부 프로세스를 증류하는 것이 온폴리시 학습 설정의 근본적인 틀을 변경하지 않으면서도 추론 성능을 향상시킬 수 있다고 상정한다.

방법론: 특권적 은닉 흐름 (Privileged Hidden Flow, PHF)

저자들은 은닉 상태 자체를 매칭하는 대신, 은닉 상태의 **이동(movement)**을 증류하는 방법인 **특권적 은격 흐름(PHF)**을 제안한다.

핵심 메커니즘

PHF는 학생의 은닉 벡터 $h^S_{\ell,t}$가 동일한 토큰 위치에서 교사의 벡터 $h^T_{\ell,t}$와 일치하도록 강제하는 (지점별 매칭, pointwise matching) 방식 대신, 토큰 위치 간의 **전이(transition)**를 매칭한다.

1. 전이 정의: 선택된 생성 위치의 집합 $W = \{w_1, \dots, w_m\}$에 대해, 로컬 은닉 전이는 인접한 위치 간의 차이로 정의된다: $\Delta h_{\ell,i} = h_{\ell, w_{i+1}} - h_{\ell, w_i}$.
2. 정규화: 이러한 전이들은 절대적인 크기가 아닌 움직임의 기하학적 구조에 집중하기 위해 단위 방향으로 정규화된다.
3. 손실 구성 요소:
   • 방향 손실 ($L_{dir}$): 각 레이어에서 학생의 전이 방향을 교사의 전이 방향과 정렬한다.
   • 기하학 손실 ($L_{geo}$): 전이 방향의 **그람 행렬(Gram matrix)**을 매칭한다. 이는 절대적인 좌표가 아니라 궤적의 관계적 구조(예: 어떤 단계들이 평행하거나, 직교하거나, 반대되는지)를 포착한다.
   • 인접 레이어 관계 ($L_{adj}$): 인접한 레이어($\ell$ 및 $\ell+1$) 간의 전이 방향 관계를 계산하여, 깊이에 따른 움직임의 변환이 교사와 일관되게 이루어지도록 한다.

구조적 불변성

PHF의 주요 이론적 기여는 훈련 중에 은닉 좌표가 이동하는 "표현 드리프트(representational drift)" 문제를 해결하는 불변성 속성이다:

• 오프셋 불변성 (Offset Invariance): 전체 궤적에 상수 오프셋이 더해지더라도 손실 값은 변하지 않는다.
• 스케일 불변성 (Scale Invariance): 전이(정규화 전)에 양의 스칼라 곱이 적용되어도 손실 값은 불변하다.
• 직교 불변성 (Orthogonal Invariance): 기하학적 항은 학생과 교사의 전이 집합에 대한 독립적인 직교 변환에 대해 불변하다.

이러한 속성들은 PHF가 학생과 교사가 고차원 은닉 공간 내에서 정확히 동일한 지점에 위치할 필요 없이, 추론 과정의 형태와 방향을 감독할 수 있게 한다.

학습 레시피

• 교사 소스: 특권적 교사는 참조 솔루션으로 조건화된 온라인 학생의 지수 이동 평균(EMA)이다.
• 집계: 이 방법은 특정 서브셋을 선택하는 대신 모든 레이어에 걸쳐 전이 손실을 집계하며, 로컬 전이 손실과 인접 레이어 관계를 결함한다.
• 목적 함수: 최종 손실은 $L_{PHF} = L_{OPSD} + \alpha L_{Flow}$이며, 여기서 $L_{Flow}$는 로컬 및 인접 손실을 결합한 것이다. 계수 $\alpha$는 모든 스케일에서 0.05로 고정된다.
• 제약 사항: 메모리 관리를 위해 128개의 생성 위치를 가진 고정 윈도우를 사용하며, 보상 모델, 정답 라우팅(correctness routing) 또는 추가적인 롤아웃을 도입하지 않는다.

주요 기여

1. 특권적 전이-운송 목적 함수: 단순한 지점별 상태 매칭을 넘어, 정규화된 토큰 간 은닉 전이와 그 궤적의 기하학을 매칭하는 새로운 손실 함수를 도입했다.
2. 이론적 규명: 운송 손실(transport loss)이 궤적 오프셋 및 직교 변환에 대해 정확한 대수적 불변성을 가짐을 증명하여, 표준 은닉 상태 매칭과 차별화했다.
3. 전 레이어 집계: 전이 손실을 모든 레이어에 걸쳐 평균 내는 것이 수동적인 레이어 선택 없이도 일관된 개선을 가져온다는 것을 입증하여, 깊이에 따른 평균화 방식을 지원한다.
4. 경험적 검증: 공식 OPSD 수학적 추론 프로토콜 하에서 세 가지 모델 규모(Qwen3-1.7B, 4B, 8B)에 대해 검증을 완료했다.

실험 결과

저자들은 AIME 2024, AIME 2025, HMMT 2025 벤치마크에서 Average@12 지표(문제당 12개 샘의 평균 정답률)를 사용하여 PHF를 평가했다.

• 성능 향상: 고정된 100단계 학습 일정 하에서, PHF는 재현된 OPSD 베이스라인 대비 Average@12 수치를 개선했다:
  • Qwen3-1.7B: +2.2 포인트 (40.83 $\to$ 43.06)
  • Qwen3-4B: +1.5 포인트 (61.89 $\to$ 63.39)
  • Qwen3-8B: +1.7 포인트 (64.26 $\to$ 65.94)
• 일관성: 이득은 모든 규모와 거의 모든 개별 벤치마크 셀에서 양수로 나타났다.
• 절제 연구 (Ablation Studies):
  • 방향 vs. 기하학: 방향 전용 또는 기하학 전용 변형 모두 모든 규모에서 전체 PHF 레시피를 따라잡지 못했으며, 이는 두 구성 요소가 모두 필요함을 시사한다.
  • 레이어 선택: 직접 선택한 작은 서브셋을 매칭하는 "선택된 레이어" 대조군이 전 레이어 집계보다 성능이 낮았으며, 이는 깊이 전체에 걸친 평균화 결정이 타당함을 뒷받침한다.
  • 지점별 매칭: 논문은 지점별 은닉 상태 매칭이 경쟁력 있는 경계 제어(boundary control) 방식이지만, 전이 기반 접근 방식보다 표현 드리프트에 덜 견고하다고 언급했다.

의의 및 주장

본 논문은 PHF가 특권적 정보를 더 구조적으로 활용함으로써 OPSD 프레임워크에 완만하지만 의미 있는 개선을 제공한다고 주장한다.

• 프로세스 수준의 감독: PHF의 주요 의의는 특권적 참조가 (모델이 자신의 롤아웃을 통해 이동하는 방식과 같은) 특정한 "은닉 프로세스"를 유도하며, 이를 증류할 수 있음을 보여준 데 있다. PHF는 학생이 교사의 절대적인 은닉 좌표를 맞출 필요 없이 이 프로세스를 포착한다.
• 배포 호환성: 테스트 시점에 보상 모델이나 외부 검증기를 요구하는 방법들과 달리, PHF는 기존의 OPSD 배포 설정을 유지한다. 즉, 평가되는 모델은 온라인 학생이며 추론 시에는 특권적 정보가 필요하지 않다.
• 효율성: PHF는 스칼라 손실 계수와 은닉 상태 저장 윈도우만을 추가하여, 추가적인 롤아웃이나 복잡한 라우팅 메커니즘에 따른 계산 오버헤드를 피한다.
• 주장의 범위: 저자들은 자신들의 증거가 고정된 레시피, 특정 모델 제품군(Qwen3), 그리고 수학적 추론 작업으로 제한된다는 점을 명시한다. 이들은 은닉 흐름 매칭이 추론의 모든 측면(예: 증명의 질 또는 보정/calibration)을 개선한다거나 보편적인 해결책이라고 주장하지 않으며, 단지 온폴리시 셀프 디스틸레이션에 적용했을 때 일관되고 긍정적인 종합적 변화를 제공한다는 점을 강조한다.