Revisión de la Destilación On-Policy: Modos de Falla Empíricos y Soluciones Sencillas

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

Resumen

La destilación bajo política (OPD) es atractiva para el post-entrenamiento de modelos de lenguaje grandes (LLM) porque evalúa la retroalimentación del profesor sobre las trayectorias generadas por el estudiante, en lugar de trazas fijas del profesor. Sin embargo, en entornos de horizonte largo, la variante común de token muestreado es frágil: reduce el emparejamiento de distribuciones a una señal de un solo token y se vuelve cada vez más poco fiable a medida que las trayectorias se alejan de los prefijos que el profesor visita comúnmente. Revisamos la OPD desde los aspectos del estimador y la implementación. Teóricamente, la OPD a nivel de token está sesgada en relación con la KL inversa a nivel de secuencia, pero tiene un límite de varianza en el peor caso mucho más ajustado; nuestro estudio de juguete muestra la misma disyuntiva empíricamente, donde un acoplamiento más fuerte de la recompensa futura produce una mayor varianza del gradiente y un aprendizaje menos estable. Empíricamente, identificamos tres modos de fallo de la OPD con token muestreado: una señal desequilibrada de un solo token, una guía del profesor poco fiable sobre prefijos generados por el estudiante y distorsiones causadas por el tokenizador o la falta de coincidencia de tokens especiales. Abordamos estos problemas con el emparejamiento del soporte local top-K del profesor, implementado como KL inversa truncada con muestreo de trayectorias top-p y enmascaramiento de tokens especiales. En tareas únicas de razonamiento matemático y entrenamiento multitarea agéntico-matemático, este objetivo produce una optimización más estable y un mejor rendimiento final que la OPD con token muestreado.

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.