Revisitando a Destilação On-Policy: Modos de Falha Empíricos e Correções Simples
Revisiting On-Policy Distillation: Empirical Failure Modes and Simple Fixes
March 26, 2026
Autores: Yuqian Fu, Haohuan Huang, Kaiwen Jiang, Yuanheng Zhu, Dongbin Zhao
cs.AI
Resumo
A destilação on-policy (OPD) é atraente para o pós-treinamento de grandes modelos de linguagem (LLMs) porque avalia o feedback do professor em rollouts gerados pelo aluno, em vez de em trajetórias fixas do professor. No entanto, em cenários de longo horizonte, a variante comum com token amostrado é frágil: ela reduz o casamento de distribuição a um sinal de um único token e torna-se progressivamente menos confiável à medida que os rollouts se afastam dos prefixos que o professor costuma visitar. Revisitamos a OPD pelos lados do estimador e da implementação. Teoricamente, a OPD em nível de token é tendenciosa em relação à divergência reversa de KL em nível de sequência, mas possui um limite de variância no pior caso muito mais restrito; nosso estudo em um ambiente controlado mostra a mesma compensação empiricamente, com um acoplamento mais forte de recompensa futura produzindo maior variância no gradiente e aprendizado menos estável. Empiricamente, identificamos três modos de falha da OPD com token amostrado: um sinal desequilibrado de um único token, orientação não confiável do professor sobre prefixos gerados pelo aluno e distorções causadas por incompatibilidade de tokenizer ou de tokens especiais. Abordamos essas questões com o casamento do suporte local top-K do professor, implementado como uma divergência reversa de KL truncada com amostragem de rollout top-p e mascaramento de tokens especiais. Tanto no raciocínio matemático de tarefa única quanto no treinamento multitarefa (agência mais matemática), esse objetivo resulta em uma otimização mais estável e em melhor desempenho downstream do que a OPD com token amostrado.
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.