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

Este artículo identifica modos de fallo en la destilación en línea (OPD) basada en tokens muestreados y propone una solución simple mediante el emparejamiento de soporte local top-K del profesor, logrando una optimización más estable y un mejor rendimiento en tareas de razonamiento matemático y agentes.

Lo esencial

Imagina que estás aprendiendo a tocar el piano. Tienes un maestro (un modelo de IA muy inteligente) y tú eres el estudiante (un modelo más pequeño).

El problema que este paper intenta resolver es cómo el estudiante debe aprender de sus propios errores cuando ya no está tocando las mismas canciones que el maestro, sino creando sus propias melodías.

Aquí tienes la explicación sencilla, paso a paso:

1. El Problema: "El Maestro se pierde en tu viaje"

En el método antiguo (llamado On-Policy Distillation o "distilación en línea"), el estudiante toca una nota, y el maestro le dice: "Bien" o "Mal" solo sobre esa única nota.

• La analogía: Imagina que estás caminando por un bosque nuevo. El maestro te guía desde una torre, pero solo puede ver el árbol que estás tocando en ese preciso segundo. Si te desvías un poco del camino que él conoce, empieza a gritar cosas extrañas.
• El fallo: Si el estudiante se equivoca y empieza a caminar por un sendero que el maestro nunca ha visto, el maestro se confunde. A veces, el estudiante empieza a repetir una frase sin sentido (como un disco rayado) o a divagar, y el maestro, al no entender el contexto nuevo, sigue diciendo "¡Bien hecho!" en cada repetición. El estudiante aprende a repetir tonterías porque el maestro le da "premios" falsos.

2. La Solución: "Mirar el mapa, no solo la nota"

Los autores dicen: "No basta con mirar solo la nota que acabas de tocar. Debemos mirar un pequeño grupo de notas que el maestro considera probables en este momento".

• La nueva idea (Emparejamiento de soporte local): En lugar de preguntar al maestro solo sobre la nota que tocaste, le preguntamos: "Maestro, de las 10 notas que tú crees que podrían sonar bien aquí, ¿cuáles son las mejores?".
• Cómo funciona:
  1. El estudiante toca una nota.
  2. El maestro mira sus "top 10" de notas posibles para ese momento.
  3. Comparamos la distribución del estudiante con la del maestro dentro de ese grupo de 10.
  4. Si el estudiante se aleja de ese grupo de 10, recibe una corrección suave pero clara. Si se queda dentro, se ajusta finamente.

3. ¿Por qué es mejor? (La analogía del GPS)

• El método viejo (Token a Token): Es como tener un GPS que solo te dice "Gira a la derecha" cuando ya estás en la esquina. Si te equivocas y vas a la izquierda, el GPS se rompe y te sigue diciendo "Gira a la derecha" aunque estés en el río.
• El método nuevo (Soporte Local): Es como tener un GPS que te dice: "Estás en la zona de las calles A, B y C. Deberías estar en la calle B". Si te vas a la calle D (que no existe en el mapa del maestro), el GPS te corrige inmediatamente porque sabe que no estás en ninguna de las opciones válidas.

4. Los Tres Enemigos que derrotaron

El paper identifica tres cosas que hacían fallar al método viejo:

1. La señal desequilibrada: El maestro solo premiaba o castigaba una sola nota, lo que era como intentar aprender a cocinar probando solo un grano de sal.
2. El maestro ciego: Cuando el estudiante se iba por caminos raros, el maestro ya no sabía qué era bueno y empezaba a dar consejos contradictorios.
3. El problema de la traducción: A veces, el estudiante y el maestro hablan "idiomas" diferentes (usaban diferentes formas de escribir las palabras). El maestro pensaba que el estudiante había dicho una tontería cuando en realidad había dicho lo correcto, pero escrito de otra forma. El nuevo método ignora estos errores de traducción.

5. El Resultado

Al usar este nuevo método (que llaman "Emparejamiento de soporte local"):

• El estudiante aprende más rápido y de forma más estable.
• No se queda atascado en bucles de repetición (como decir "espera, espera, espera" una y otra vez).
• Funciona mejor en tareas difíciles de matemáticas y en tareas donde el agente debe tomar decisiones complejas (como un robot que debe mover muebles).

En resumen:
El paper dice que para enseñar a una IA a pensar por sí misma, no debemos castigarla o premiarla por cada palabra individual que dice. En su lugar, debemos preguntarle al maestro: "¿Qué conjunto de palabras probables deberíamos estar considerando aquí?". Al comparar al estudiante con ese grupo de opciones válidas, evitamos que la IA se pierda en sus propios delirios y aprende a pensar de forma más lógica y creativa.

Resumen técnico

Resumen Técnico: Revisión de la Destilación On-Policy (OPD)

1. El Problema: Fragilidad de la OPD en Entornos de Largo Horizonte

La Destilación On-Policy (OPD) es un enfoque atractivo para el entrenamiento posterior (post-training) de Modelos de Lenguaje Grande (LLM), donde un modelo "estudiante" se entrena sobre sus propias secuencias generadas (rollouts) mientras recibe retroalimentación de un modelo "maestro" más fuerte. A diferencia de la destilación off-policy (que usa trazas fijas del maestro), la OPD permite al estudiante explorar regiones del espacio de estados que el maestro podría no visitar.

Sin embargo, el artículo identifica que la implementación estándar de OPD, basada en la comparación de tokens muestreados (sampled-token), es extremadamente frágil en escenarios de largo horizonte (razonamiento complejo, agentes). Los autores identifican tres modos de fallo críticos:

1. Señal de un solo token desequilibrada: La actualización se basa únicamente en la razón de log-probabilidades del token muestreado. Esto convierte una discrepancia de distribución en una señal de aprendizaje altamente desequilibrada, donde la mayoría de los tokens reciben recompensas negativas y la optimización se vuelve sensible a continuaciones locales favorables pero semánticamente irrelevantes (ej. marcadores de vacilación).
2. Guía no fiable en prefijos generados por el estudiante: A medida que el estudiante se desvía de las trazas típicas del maestro, el maestro puede asignar alta probabilidad a tokens que parecen plausibles localmente pero que pertenecen a trayectorias degradadas (ej. bucles de repetición, reinicios de razonamiento). La OPD estándar falla al penalizar estos comportamientos.
3. Distorsión por desajuste de tokenización: Si el maestro y el estudiante usan esquemas de tokenización diferentes, un token generado por el estudiante podría no corresponder a un token natural para el maestro, generando penalizaciones falsas (falsos negativos) aunque el contenido semántico sea correcto.

2. Metodología: Emparejamiento de Soporte Local Top-K del Maestro

Para abordar estos problemas, los autores proponen un nuevo objetivo de entrenamiento llamado Teacher Top-K Local Support Matching (Emparejamiento de Soporte Local Top-K del Maestro).

• Concepto Central: En lugar de comparar al estudiante y al maestro solo en el token muestreado, la comparación se realiza sobre un conjunto de soporte local definido por el maestro.
• Implementación (Truncated Reverse-KL):
  • En cada paso de prefijo, se define un conjunto de soporte $S$ que contiene los $K$ tokens con mayor probabilidad según la distribución del maestro ($TopK_q$).
  • Se renormalizan las distribuciones del estudiante y del maestro dentro de este subconjunto restringido.
  • El objetivo de pérdida es una divergencia de Kullback-Leibler (KL) inversa truncada calculada sobre este soporte local, en lugar de sobre todo el vocabulario o un solo token.
• Estabilización Práctica:
  • Muestreo Top-p: Se utilizan rollouts generados con muestreo top-p (ej. $p=0.9$) para mantener las trayectorias dentro de regiones donde la señal del maestro es informativa, evitando tokens de probabilidad extremadamente baja.
  • Máscara de Tokens Especiales: Se enmascaran tokens especiales problemáticos (como marcadores de fin de secuencia o tokens de pensamiento) para mitigar el ruido causado por diferencias de tokenización entre modelos.

3. Contribuciones Clave

1. Análisis de Compensación (Bias-Variance): Los autores demuestran teórica y empíricamente que la OPD a nivel de token es sesgada respecto al objetivo de nivel de secuencia (Reverse-KL completo), pero tiene un límite de varianza mucho más estricto ($O(T^2)$ frente a $O(T^4)$ para secuencias largas). Esto explica por qué la OPD de token es estable pero inexacta, mientras que la de secuencia es precisa pero inestable.
2. Identificación de Modos de Fallo: Documentan empíricamente cómo la señal de un solo token falla en escenarios de largo horizonte, incluyendo la incapacidad de penalizar bucles de repetición y la sensibilidad a artefactos de tokenización.
3. Propuesta de Solución Efectiva: Introducen el emparejamiento de soporte local Top-K, que logra un equilibrio: mantiene la eficiencia computacional de la OPD local pero reemplaza la estimación de un solo punto por una comparación de distribución local, reduciendo la varianza y mejorando la robustez.

4. Resultados Experimentales

Los autores evaluaron su método en dos configuraciones principales utilizando Qwen2.5-7B como estudiante:

• Razonamiento Matemático (Tarea Única):
  • En benchmarks como Math500, AIME24/25 y OlympiadBench, la OPD con soporte local superó consistentemente a la OPD de token muestreado.
  • La versión con soporte local alcanzó un promedio de 41.5, superando a la OPD estándar (36.4) y a la OPD estándar con máscaras (40.7). Esto demuestra que la mejora no es solo por corregir errores de tokenización, sino por una señal de destilación local más fuerte.
• Entrenamiento Multi-tarea (Agente + Matemáticas):
  • En un entorno que alterna entre tareas de agentes (ALFWorld) y matemáticas, el método propuesto mejoró significativamente el rendimiento matemático (de 36.6 a 41.7 en promedio) sin sacrificar el rendimiento en tareas de agentes (manteniéndose competitivo o mejorando a 97.7 en ALFWorld).
• Dinámica de Entrenamiento:
  • El método propuesto mostró menor varianza en los gradientes, menos cortes (clipping) en la actualización de políticas y una alineación más rápida entre estudiante y maestro en comparación con la OPD estándar.
  • Las curvas de aprendizaje mostraron mejoras estables a lo largo del entrenamiento, no solo al final.

5. Significado e Impacto

Este trabajo es fundamental para el futuro del entrenamiento de LLMs en tareas de razonamiento complejo y agentes autónomos.

• Validación de la OPD: Confirma que la supervisión on-policy es viable y necesaria para el largo horizonte, pero requiere un diseño de objetivo más sofisticado que la simple comparación de tokens.
• Solución Práctica: Ofrece una solución computacionalmente eficiente (no requiere calcular KL sobre todo el vocabulario) que mitiga los problemas de inestabilidad y "reward hacking" (engaño a la recompensa) comunes en la destilación on-policy.
• Dirección Futura: Sugiere que la combinación de objetivos locales robustos (como el soporte Top-K) con un control estricto de la deriva de los rollouts (drift) es la vía para cerrar la brecha entre la coincidencia con el maestro y el éxito real en la tarea.

En resumen, el artículo demuestra que reemplazar la supervisión de un solo token por una comparación de distribución local sobre el soporte del maestro es una "solución simple" pero poderosa que estabiliza el entrenamiento de LLMs en escenarios complejos de largo horizonte.