Riesaminare la distillazione on-policy: modalità di fallimento empiriche e soluzioni semplici
Revisiting On-Policy Distillation: Empirical Failure Modes and Simple Fixes
March 26, 2026
Autori: Yuqian Fu, Haohuan Huang, Kaiwen Jiang, Yuanheng Zhu, Dongbin Zhao
cs.AI
Abstract
La distillazione on-policy (OPD) è interessante per il post-addestramento di grandi modelli linguistici (LLM) perché valuta il feedback del docente sui rollout generati dallo studente, anziché su tracce fisse del docente. Tuttavia, in contesti a lungo orizzonte, la variante comune basata sul token campionato è fragile: riduce l'adattamento della distribuzione a un segnale a singolo token e diventa progressivamente inaffidabile man mano che i rollout si allontanano dai prefissi che il docente visita comunemente. Rivediamo l'OPD dal punto di vista dell'estimatore e dell'implementazione. Teoricamente, l'OPD a livello di token è distorto rispetto alla KL inversa a livello di sequenza, ma ha un limite di varianza nel caso peggiore molto più stretto; il nostro studio giocattolo mostra empiricamente lo stesso compromesso, con un accoppiamento più forte della ricompensa futura che produce una maggiore varianza del gradiente e un apprendimento meno stabile. Empiricamente, identifichiamo tre modalità di fallimento dell'OPD a token campionato: un segnale a singolo token sbilanciato, una guida inaffidabile del docente sui prefissi generati dallo studente e distorsioni causate da mismatch del tokenizer o dei token speciali. Affrontiamo questi problemi con l'adattamento del supporto locale top-K del docente, implementato come KL inversa troncata con campionamento dei rollout top-p e mascheramento dei token speciali. In contesti di ragionamento matematico a singolo compito e di addestramento multi-compito (agente+matematica), questo obiettivo produce un'ottimizzazione più stabile e prestazioni migliori a valle rispetto all'OPD a token campionato.
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.