再考：オンライン蒸留の実証的失敗モードと簡易修正

Revisiting On-Policy Distillation: Empirical Failure Modes and Simple Fixes

March 26, 2026

著者: Yuqian Fu, Haohuan Huang, Kaiwen Jiang, Yuanheng Zhu, Dongbin Zhao

cs.AI

要旨

オン方針蒸留（OPD）は、教師モデルの固定軌跡ではなく、生徒モデルが生成したロールアウトに対する教師フィードバックを評価するため、大規模言語モデル（LLM）の事後学習において魅力的な手法である。しかし、長期的な設定では、一般的なサンプリングトークン方式は脆弱である。これは分布マッチングを単一トークンの信号に還元し、ロールアウトが教師モデルが頻繁に訪れる接頭部から逸脱するにつれて信頼性が低下する。本論文では、推定量と実装の両面からOPDを再検討する。理論的には、トークンレベルOPDはシーケンスレベルの逆KLダイバージェンスに対してバイアスを持つが、最悪ケースの分散限界は遥かに厳密である。玩具実験では、将来報酬の結合が強いほど勾配分散が増大し学習が不安定になるという、同じトレードオフが実証的に観察された。実証的には、サンプリングトークンOPDの3つの失敗モードを特定する：不均衡な単一トークン信号、生徒生成接頭部に対する信頼性の低い教師指導、トークナイザーまたは特殊トークンの不一致による歪みである。これらの問題を解決するため、教師モデルのトップKローカルサポートマッチングを提案する。これは、トップpロールアウトサンプリングと特殊トークンマスキングを伴う切断逆KLとして実装される。単一タスクの数学推論とマルチタスクのエージェント・数学複合学習において、この目的関数はサンプリングトークンOPDよりも安定した最適化と優れた下流性能を達成する。

English

On-policy distillation (OPD) is appealing for large language model (LLM) post-training because it evaluates teacher feedback on student-generated rollouts rather than fixed teacher traces. In long-horizon settings, however, the common sampled-token variant is fragile: it reduces distribution matching to a one-token signal and becomes increasingly unreliable as rollouts drift away from prefixes the teacher commonly visits. We revisit OPD from the estimator and implementation sides. Theoretically, token-level OPD is biased relative to sequence-level reverse-KL, but it has a much tighter worst-case variance bound; our toy study shows the same tradeoff empirically, with stronger future-reward coupling producing higher gradient variance and less stable learning. Empirically, we identify three failure modes of sampled-token OPD: an imbalanced one-token signal, unreliable teacher guidance on student-generated prefixes, and distortions caused by tokenizer or special-token mismatch. We address these issues with teacher top-K local support matching, implemented as truncated reverse-KL with top-p rollout sampling and special-token masking. Across single-task math reasoning and multi-task agentic-plus-math training, this objective yields more stable optimization and better downstream performance than sampled-token OPD.