OPRD: On-Policy Representation Distillation

온-정책 증류를 출력 공간에서 hidden-state 공간으로 끌어올려 학생-교사 격차를 완전히 해소하는 OPRD(On-Policy Representation Distillation)를 제안한다.

기존 온-정책 증류(OPD)는 출력 분포만 매칭하여 대어휘 샘플링 분산과 LM-head 정보 병목이라는 두 가지 근본적 한계가 있었다. OPRD는 교사의 중간 hidden state를 정렬 대상으로 삼아 분산을 제거하고, 층별·위치별 구조적 정보를 결정론적으로 전달한다. 수학 추론 벤치마크(AIME 2024/2025, AIMO)에서 학생-교사 격차를 완전히 해소했으며, top-k OPD 대비 1.44배 빠른 학습과 최대 54% 적은 GPU 메모리를 달성했다.

핵심 요약

• 출력 공간 매칭에서 hidden-state 공간 매칭으로 증류를 상향 — 기존 모든 OPD 변종의 구조적 한계를 극복
• Qwen ~150K 어휘처럼 대어휘 모델에서 Monte Carlo KL 추정의 분산 문제를 결정론적 MSE로 제거
• LM-head 투영이 폐기하던 교사의 층별·위치별 중간 표현을 직접 학생에 전달
• AIME 2024/2025, AIMO 세 경쟁 수학 벤치마크에서 학생-교사 격차 완전 해소
• 동일 하드웨어에서 top-k OPD 대비 1.44× 빠른 학습, 최대 54% 적은 actor-update 메모리
• 다중 모델 RL 병합, 온-정책 자기 증류(OPSD) 등 확장 시나리오에서 메모리·분산 이점이 특히 큼

논문 링크

외부 연구를 정리한 자료입니다. HDATF가 발표한 논문이나 제품 성능을 측정한 결과는 아닙니다.