Rethinking On-Policy Distillation I

On-policy distillation(OPD)은 학생이 스스로 생성한 rollout 위에서 teacher의 token별 log-probability를 dense reward로 쓰는 post-training 기법이다. Qwen3·MiMo·GLM-5가 채택했고 Thinking Machines Lab이 RL 대비 적은 비용으로 재현했지만, 왜 token 수준 신호가 학생을 올바른 방향으로 옮기는지, 언제 실패하는지는 거의 연구되지 않았다. 이 논문은 현상(§3) → 기제(§4) → 처방(§5) → 비용(§6) 순으로 OPD를 체계적으로 분석한다.

OPD는 학생 $\pi_\theta$가 생성한 rollout $\hat y$ 위에서 sequence-level reverse KL을 최소화하며, 이는 token-level KL의 합으로 정확히 분해된다.

여기서 $p_t,q_t$는 학생 prefix $\hat y_{<t}$에서의 student·teacher 다음 token 분포다. 구현은 이 per-token KL을 어떻게 추정하느냐로 갈린다.

학습 내내 추적하는 진단 지표는 세 가지다.

설정. 학생은 Qwen3-1.7B-Base. teacher는 Qwen3-4B(Non-thinking)와, Qwen3-4B-Base에 GRPO zero-RL을 적용한 Qwen3-4B-Base-GRPO. 데이터는 DAPO-Math-17K, 평가는 AIME 2024·AIME 2025·AMC 2023의 avg@16.

결과. 두 teacher의 benchmark 점수는 비슷한데, base 학생과 thinking pattern이 가까운 GRPO teacher에서 distillation이 일관되게 더 좋다. GRPO teacher는 초기 overlap ratio가 더 높다. 학습 후반에 두 run의 overlap 곡선이 수렴해도 성능 격차는 남는다. 초기의 pattern mismatch가 만든 손실은 나중에 회복되지 않는다.

Thinking pattern이 맞고 점수도 높은데 OPD가 실패하는 경우가 있다. 두 model family에서 같은 pipeline의 teacher와 추가 RL을 거친 teacher를 대조한다.

두 family 모두 같은 pipeline의 teacher는 개선이 제한적이고, RL을 거친 teacher는 모든 benchmark에서 훨씬 큰 이득을 낸다. gap recovery rate $(\mathrm{Acc}_{\text{after}}-\mathrm{Acc}_{\text{before}})/(\mathrm{Acc}_{\text{teacher}}-\mathrm{Acc}_{\text{before}})$도 훨씬 높다. RL teacher는 같은 base checkpoint에서 나왔으므로 thinking pattern은 유지된 채, RL로 얻은 새 능력만 추가된 것이다. 같은 데이터와 recipe로 학습한 두 모델은 각자의 규모에서 비슷한 분포로 수렴해 teacher가 넘겨줄 신호가 거의 없다.

두 조건을 동시에 검증하는 실험이다. JustRL-1.5B는 R1-Distill-1.5B에 RL을 적용해 만든 모델이다. 이제 JustRL-1.5B를 학생으로, 자기 pre-RL checkpoint인 R1-Distill-1.5B를 teacher로 거꾸로 증류한다. 비교군으로 같은 family의 R1-Distill-7B도 teacher로 쓴다. 7B는 JustRL-1.5B보다 점수가 약간 높고, 1.5B는 훨씬 낮다.

여기서 세 가지 결론이 나온다.

• OPD는 thinking pattern을 배운다. teacher의 pattern을 적극적으로 획득해 학생 자신의 것을 덮어쓴다. 그래서 pattern이 너무 다르면 학습이 안 된다.
• Benchmark 점수는 OPD 결과를 예측하지 못한다. 점수가 더 높은 teacher가 개선은커녕 퇴행을 일으킬 수 있다. 학습 dynamics는 점수와 완전히 분리될 수 있다.
• 높은 점수 ≠ 새 지식. 규모만 다른 같은 family의 두 모델은 학생에게 같은 효과를 낸다. 점수 차이는 같은 데이터에 대한 적합도 차이일 뿐이다.

설정. 학생 R1-Distill-1.5B에 JustRL-1.5B(성공)와 R1-Distill-7B(실패) teacher를 붙이고 세 지표를 추적한다. 두 teacher의 수학 성능은 비슷하며 7B가 약간 높다.

결과. JustRL-1.5B로부터는 teacher와의 격차 80% 이상을 회복하고, R1-Distill-7B로부터는 전혀 개선되지 않는다. dynamics가 갈림길을 보여준다.

성공 설정에서 loss가 덮는 token 집합만 바꾼 세 변형을 비교한다. (i) Student Top-$k$: 학생 top-$k$ 전체 $S_t^{(p)}$, (ii) Overlap Top-$k$: 교집합 $S_t^{(p)}\cap S_t^{(q)}$, (iii) Non-Overlap Top-$k$: 대칭차 $S_t^{(p)}\triangle S_t^{(q)}$. $k=16$.

• Overlap Top-$k$는 세 benchmark 모두에서 Student Top-$k$의 이득을 거의 전부 회복한다. Non-Overlap Top-$k$는 일관되게 약하다.
• Student Top-$k$와 Overlap Top-$k$의 overlap-token advantage 곡선은 거의 구별되지 않는다. 학생 전용 token은 확률 질량이 매우 작기 때문이다.
• 자기 강화 dynamics: Student·Overlap Top-$k$는 overlap ratio를 약 72%에서 91% 이상으로 꾸준히 올리는 반면, Non-Overlap Top-$k$는 먼저 떨어졌다가 부분 회복한다. token이 공유 고확률 영역에 들어와 teacher가 선호하면 reverse KL이 그 token에 질량을 더 모으고, 경쟁하는 non-overlap token을 학생 top-$k$ 밖으로 밀어낸다. overlap 영역은 최적화에도 불구하고가 아니라 최적화 때문에 자란다.

새 지식은 teacher 고유 속성이라 바꿀 수 없지만, thinking pattern 격차는 학습 설계로 줄일 수 있다. 첫 번째 처방은 off-policy cold start다. teacher가 생성한 rollout으로 학생을 먼저 SFT해 pattern을 가깝게 만든 뒤 표준 OPD를 이어 간다.

설정. 학생 Qwen3-1.7B-Base, teacher Qwen3-4B(Non-thinking). OpenThoughts3-1.2M의 수학 subset에서 teacher가 200K 응답을 생성하고 이를 SFT해 Qwen3-1.7B-SFT를 만든다. 이어서 SFT prompt와 중복 제거한 나머지 약 30K prompt로 OPD. 대조군은 같은 teacher·같은 OPD prompt로 Base에서 바로 시작하는 pure OPD.

Off-policy distillation이 초기 pattern mismatch를 줄여, OPD가 시작되는 순간부터 teacher의 token 수준 감독을 활용할 수 있게 만든다.

두 번째 처방은 데이터 쪽에서 정렬을 높인다. teacher의 policy는 post-training에서 본 prompt에 의해 형성되므로, teacher-aligned prompt로 OPD하면 감독이 더 효과적이다. 두 granularity에서 실험한다.

Prompt content 실험의 미묘한 점은 학생이 더 적지만 더 강하게 공유되는 token에 질량을 집중한다는 것이다. overlap 집합은 작아지지만 고확률 token에 대한 유효 정렬은 강해진다. 동시에 학생 entropy가 크게 낮아진다. teacher의 post-training prompt만 쓰면 policy entropy를 과도하게 줄여 탐색 능력을 해칠 수 있으므로, 실전에서는 teacher-aligned prompt와 그 밖의 prompt를 섞는 전략이 더 견고하다.

OPD의 매력은 모든 token이 reward를 받는 dense supervision이다. 그러나 이는 teacher의 token 수준 reward가 학생 방문 state에서 신뢰할 만하다는 가정에 기댄다.

최대 응답 길이를 0.5K부터 15K까지 바꾸면, 0.5K·1K는 감독 token이 너무 적어 sample-efficient하지 않고, 3K·7K가 가장 강하며, 10K·15K는 정체하거나 하락한다. 10K·15K는 후반에 overlap ratio가 급락하고 학생 entropy와 gradient norm이 튀는 late-stage collapse를 보인다.

DAPO-Math-17K에서 2K prompt를 뽑아 16K token을 넘는 학생 rollout을 여러 위치에서 자르고 teacher가 이어 쓰게 하면, teacher의 정확도 우위는 1K prefix에서 +0.37, 16K prefix에서 +0.02로 단조 감소한다. dense reward는 중간 길이의 추론에는 효과적이지만 학생 prefix가 teacher에게 익숙한 state에서 멀어질수록 신뢰도가 떨어진다. 긴 chain-of-thought나 agentic multi-turn 같은 long-horizon 설정으로 OPD가 깔끔하게 확장되지 않을 수 있다.

실패한 7B teacher의 reward가 근본적으로 무정보인지 확인하기 위해 학생 rollout별 sequence mean reward $\bar r(y)$가 정답·오답을 얼마나 구분하는지 본다. 두 teacher 모두 정답 rollout에 더 높은 reward를 주며 AUROC는 0.73(JustRL-1.5B)과 0.75(R1-Distill-7B)로 비슷하다. 실패한 teacher의 전역 신호가 더 약하지 않다.

그렇다면 왜 실패하는가. 7B teacher에서는 후반에 overlap-token advantage 크기가 더 큰데도 gradient norm은 계속 작다. 가설은 anisotropy다. 7B의 per-token advantage가 개별적으로는 크지만 sequence 안 위치마다 방향이 제각각이라 gradient로 합쳐질 때 서로 상쇄된다. 호환되는 pattern의 JustRL-1.5B는 더 일관된 token 부분집합에 advantage를 집중해, 작은 신호라도 reverse KL의 mode-seeking이 증폭할 수 있는 한 방향을 가리킨다. 이 가설은 직접 검증되지 않았고 per-token gradient의 방향 구조 분석이 필요하다.

학생 R1-Distill-1.5B, teacher JustRL-1.5B에서 Top-$k$ OPD($k\in\{1,4,16,64\}$)와 sampled-token OPD를 비교하면 sampled-token은 Top-$k$ 평균과 대등하다. 유일하게 나쁜 설정은 Top-1이며 $k$를 4 이상으로 키워도 이득은 미미하다. :::figure assets/fig-15.png Figure 15. Top-$k$ OPD의 support 크기 $k$ 효과 (avg@16). sampled-token OPD가 Top-4·16·64와 대등하고 Top-1만 일관되게 낮다. :::

Top-1은 항상 argmax token만 골라 reward를 한 mode에 집중시키고, 작은 policy 변화로 1위 token이 바뀌어 신호가 불안정해진다. sampled-token은 매 step 학생 분포에 비례해 다른 token을 뽑으므로 고확률 영역을 unbiased하게 덮는다. 문제는 token 수가 아니라 편향된 mode-집중 선택 규칙이다.

• Off-policy KD의 한계: Kim & Rush(2016)의 sequence-level distillation 계열은 teacher 생성 sequence로 학습하고 추론은 학생 분포에서 하므로 exposure bias가 누적된다. 이 mismatch가 OPD의 동기다.
• OPD 계보: MiniLLM(reverse KL을 policy gradient로), GKD(on/off-policy와 divergence를 잇는 통합 틀), Yang et al. 2026(OPD를 dense KL-constrained RL의 특수 경우로 형식화). 산업 pipeline과 self-distillation으로 확장됐지만 언제·왜 실패하는지의 체계적 분석은 없었다.
• Capacity gap: Cho & Hariharan, teacher assistant, distillation scaling law의 U자형 regime, long CoT의 learnability gap 등은 대부분 off-policy KD에 관한 것이다. OPD에서의 distillability는 미개척이었다.

• 모든 실험이 수학 benchmark다. code·open-ended 영역에서도 같은 조건과 기제가 성립하는지는 열려 있다.
• "새 지식" 조건은 pre-training corpus 차이에 기대는데, cross-family 증류는 tokenizer·architecture 차이와 뒤섞이고 통제된 pre-training ablation은 너무 비싸다.
• Self-distillation에서는 pattern 일관성이 보장되고 novelty가 privileged information에서 오므로, 같은 통찰을 확장하는 것이 자연스러운 다음 단계다.
• Long-horizon 한계는 짧은 구간의 dense token 감독과 긴 horizon의 sparse outcome reward를 결합하는 hybrid, 감독 horizon을 점진적으로 늘리는 curriculum을 시사한다.