Пересмотр дистилляции на основе стратегии: эмпирические проблемы и простые решения
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 с семплированием токенов: несбалансированный однотокенный сигнал, ненадежное руководство учителя на сгенерированных студентом префиксах и искажения, вызванные несовпадением токенизаторов или специальных токенов. Мы решаем эти проблемы с помощью сопоставления локального носителя учителя по топ-K, реализованного как усеченная обратная KL-дивергенция с семплированием разверток по top-p и маскированием специальных токенов. В рамках однозадачного математического рассуждения и многозадачного обучения (агентское поведение плюс математика) этот подход обеспечивает более стабильную оптимизацию и лучшую итоговую производительность по сравнению с 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.