OPRD: 온폴리시 표현 증류

OPRD: On-Policy Representation Distillation

온폴리시 증류(OPD)는 학생 모델을 다음 토큰 확률을 매칭하여 출력 공간에서만 학습하도록 합니다. 이러한 출력-중심 접근 방식은 다음과 같은 두 가지 한계점을 가지고 있습니다 (1) 대규모 어휘(예: Qwen의 약 150,000개의 토큰)에 대한 몬테카를로 KL 추정으로 인한 샘플링 분산이 학습 전체 과정에서 지속되고, (2) 교사 모델을 블랙박스로 취급하여 LM 헤드 이후의 모든 중간 은닉 상태를 버립니다. 우리는 온폴리시 표현 증류(OPRD)를 제안합니다. OPRD는 동일한 배치 내에서 선택된 레이어에 걸쳐 학생 및 교사 모델의 표현을 정렬하여 증류를 은닉 상태 공간으로 확장하며, LM 헤드를 완전히 우회합니다. 이론적으로, OPRD는 샘플링 분산을 제거하고 각 레이어별로 더 풍부한 구조적 정보를 제공합니다. 실험적으로, OPRD는 AIME 2024/2025 및 AIMO 데이터셋에서 학생-교사 모델 간의 성능 격차를 줄이는 반면, 출력 공간 OPD 기반 모델은 교사 모델 수준에 도달하지 못합니다. 또한, OPRD는 top-k OPD보다 1.44배 더 빠르게 학습되고 메모리 사용량도 54% 적습니다. 코드: https://github.com/ShenzhiYang2000/OPRD.

On-policy distillation (OPD) supervises the student only in output space by matching next-token probabilities. This output-only paradigm has two limits: (1) sampling variance from Monte Carlo KL estimates over large vocabularies (e.g., Qwen's ~150k tokens) persists throughout training, and (2) it treats the teacher as a black-box, discarding all intermediate hidden states after the LM head. We propose On-Policy Representation Distillation (OPRD), which lifts distillation into hidden-state space by aligning student and teacher representations across selected layers on the same rollouts, bypassing the LM head entirely. Theoretically, OPRD eliminates sampling variance and provides richer per-layer structural information. Empirically, OPRD closes the student-teacher gap on AIME 2024/2025 and AIMO, while output-space OPD baselines plateau below the teacher. OPRD also trains 1.44x faster and uses 54% less memory than top-k OPD. Code: https://github.com/ShenzhiYang2000/OPRD.