추론 접두어로부터의 빠르고 효과적인 온라인 증류

Fast and Effective On-policy Distillation from Reasoning Prefixes

온라인 증류(OPD)는 학생 모델에서 샘플링된 경로를 사용하고 토큰 수준에서 교사 모델을 통해 이를 지도하며, 검증 가능한 최종 보상에만 의존하지 않고 오프라인 증류보다 더 나은 일반화 성능을 제공할 수 있습니다. 그러나 OPD는 훈련 중에 학생 정책을 실시간으로 샘플링해야 하므로, 특히 긴 응답의 경우 훈련 비용이 크게 증가합니다. 초기 분석 결과, OPD 과정에서 훈련 신호가 종종 각 출력의 접두부에 집중되며, 짧은 교사 모델이 생성한 접두어만으로도 학생 모델이 올바른 답변을 생성하는 데 크게 도움이 될 수 있다는 것을 알 수 있습니다. 이러한 관찰을 바탕으로, 우리는 간단하면서도 효과적인 OPD의 수정 방법을 제안합니다. 즉, 증류 목표를 학생 모델이 생성한 출력의 접두부에만 적용하고, 각 샘플링을 조기에 종료합니다. AI-for-Math 및 다양한 도메인 벤치마크에서의 실험 결과, 온-라인 접두어 증류는 전체 OPD와 동등한 성능을 보이면서 훈련 FLOP을 2배에서 47배까지 줄입니다.

On-policy distillation (OPD), which samples trajectories from the student model and supervises them with a teacher at the token level, avoids relying solely on verifiable terminal rewards and can yield better generalization than off-policy distillation. However, OPD requires expensive on-the-fly sampling of the student policy during training, which substantially increases training cost, especially for long responses. Our initial analysis shows that, during OPD, training signals are often concentrated in the prefix of each output, and that even a short teacher-generated prefix can significantly help the student produce the correct answer. Motivated by these observations, we propose a simple yet effective modification of OPD: we apply the distillation objective only to prefixes of student-generated outputs and terminate each sampling early during distillation. Experiments on a suite of AI-for-Math and out-of-domain benchmarks show that on-policy prefix distillation matches the performance of full OPD while reducing training FLOP by 2x-47x.