Papers·어제
On-Policy Representation Distillation (OPRD) — hidden state 정렬로 AIME 2024/2025 gap 0으로
Shenzhi Yang 팀이 기존 on-policy distillation(OPD)의 출력 공간 한계를 보완해, student와 teacher의 hidden state를 같은 rollout 상에서 정렬하는 OPRD를 제안했습니다. 이론적으로 Monte Carlo sampling variance를 제거하고, 실험적으로 AIME 2024/2025와 AIMO에서 student-teacher gap을 0으로 만들었으며, top-k OPD 대비 1.44배 빠른 학습과 54% 적은 메모리를 사용합니다. 단, Qwen의 ~150k 토큰 어휘를 가정한 실험으로, 작은 vocab 모델에서는 이득이 줄어들 가능성이 있습니다.
Shenzhi Yang 팀이 출력 공간 KL distillation의 sampling variance와 black-box 한계를 극복하는 hidden-state 정렬 방법 OPRD를 공개했습니다.
핵심 결론
- 성능 — AIME 2024/2025, AIMO에서 student-teacher gap을 0으로 — output-space OPD는 teacher 아래에서 정체.
- 효율 — top-k OPD 대비 학습 1.44배 빠르고, 메모리 54% 절감.
방법
- 아이디어 — 같은 rollout 상에서 student와 teacher의 selected layer hidden state를 정렬, LM head를 거치지 않음.
- 이론적으로 Monte Carlo sampling variance를 제거하고, per-layer 구조 정보를 활용합니다.
한계·조건
- 어휘 — Qwen의 ~150k 토큰처럼 큰 vocab에서 효과가 두드러짐 — 작은 vocab에서는 이득이 줄어들 수 있음.
- 코드 — GitHub 공개 (https://github.com/ShenzhiYang2000/OPRD).
편집자 한 줄
hidden-state distillation이 output-space의 한계를 깔끔하게 우회한 점이 인상적입니다. 다만 layer 선택 기준이 실험마다 달라질 수 있어 일반화는 더 봐야 할 듯.
- #distillation
- #representation-learning
- #on-policy
- #qwen
Shenzhi Yang