Bridging Reasoning Trajectories in On-Policy Distillation via Near-Future Guidance

Resumen Técnico: Puente de Trayectorias de Razonamiento en la Destilación On-Policy mediante Guía de Futuro Cercano

Declaración del Problema
La Destilación On-Policy (OPD, por sus siglas en inglés) es un pilar fundamental en el post-entrenamiento moderno de Modelos de Lenguaje de Gran Escala (LLMs), diseñado para mejorar el razonamiento mediante el entrenamiento de un modelo estudiante sobre trayectorias muestreadas de su propia política bajo la supervisión de un profesor. La OPD estándar opera mediante un mecanismo de "muestreo de trayectoria pero aprendizaje de tokens": identifica desviaciones a través de tokens de alta pérdida y las repara utilizando una corrección de divergencia reverse-KL local.

Los autores identifican un desajuste fundamental en este enfoque, argumentando que la supervisión a nivel de token está estructuralmente desalineada con los requisitos de la corrección de razonamiento a nivel de trayectoria. Esta desalineación se manifiesta en dos modos críticos de falla:

1. Falsas Alarmas Ruidosas: Los tokens de alta pérdida no siempre corresponden a una divergencia genuina de razonamiento. El análisis empírico revela que aproximadamente el 30% de los tokens de alta pérdida caen en un régimen de baja divergencia, lo que representa discrepancias de forma superficial (por ejemplo, preferencias estilísticas como "y" frente a "luego") en lugar de bifurcaciones lógicas. Tratar estos casos como errores críticos proporciona una supervisión ruidosa que puede interferir activamente con la corrección del razonamiento.
2. La Trampa del Aprendizaje Token por Token: Incluso cuando se identifican puntos de divergencia genuinos, la corrección reverse-KL de un solo token suele ser insuficiente para redirigir toda la trayectoria de razonamiento. Las fallas de razonamiento suelen evolucionar como derivas distributivas de corto horizonte. Corregir un token localmente no garantiza que el modelo se alineará con la trayectoria futura del profesor; el modelo a menudo vuelve a derivar hacia ramas erróneas inmediatamente después de la corrección local.

Metodología: OPD Consciente de la Trayectoria (TOPD)
Para abordar estas limitaciones, los autores proponen la Destilación On-Policy Consciente de la Trayectoria (TOPD), la cual cierra la brecha entre la corrección de tokens locales y la evolución de la trayectoria global incorporando información de la trayectoria de futuro cercano. El método consta de dos etapas primarias:

1. Detección de Divergencia:

   • Para las posiciones candidatas de alta pérdida, el sistema retrocede al token precedente para establecer un prefijo compartido.
   • Tanto el profesor como el estudiante generan continuaciones de ventana corta (longitud $K=50$) desde este prefijo.
   • El sistema mide la discrepancia a nivel de trayectoria utilizando la distancia de Transporte Óptimo (OT) entre las continuaciones del profesor y del estudiante.
   • Los puntos divergentes reales se identifican como posiciones que exhiben tanto una alta pérdida a nivel de token como una alta divergencia de OT en la ventana corta. Esto filtra las discrepancias superficiales benignas.

2. Inyección de Señal del Profesor a Nivel de Trayectoria:

   • En lugar de optimizar solo el token actual, TOPD extiende la supervisión a toda la trayectoria de la ventana corta.
   • Utilizando el plan de transporte OT ($\gamma$), el método construye un "objetivo suave consciente de la trayectoria" para cada posición futura del estudiante. Este objetivo alinea los tokens futuros del estudiante con la trayectoria del profesor basándose en el mapeo de transporte óptimo.
   • El estudiante es entrenado para coincidir con estos objetivos suaves alineados con OT a través de la ventana, aprendiendo efectivamente cómo dirigir su trayectoria de razonamiento cercana hacia la trayectoria del profesor, no solo cómo corregir el token actual.
   • El objetivo final combina la pérdida de la OPD estándar con un término de pérdida de trayectoria ($L_{TOPD} = L_{OPD} + \lambda L_{traj}$).

Contribuciones Clave

• Identificación del Desajuste: Este trabajo establece que la supervisión a nivel de token es un "proxy impreciso" para la corrección de razonamiento a nivel de trayectoria, fallando al distinguir entre bifurcaciones lógicas genuinas y variaciones superficiales benignas.
• Caracterización Empírica de los Modos de Falla: Los autores proporcionan evidencia de que ~30% de los tokens de alta pérdida son falsas alarmas de baja divergencia y demuestran la "trampa del aprendizaje token por token", donde las correcciones locales fallan en restaurar la trayectoria de razonamiento general.
• Principio de Guía a Nivel de Trayectoria: El artículo introduce un principio donde la información de la trayectoria de futuro cercano se distribuye a través de múltiples tokens. Medi el uso de la alineación de OT de ventana corta, TOPD inyecta las discrepancias de la trayectoria profesor-estudiante directamente en la función de pérdida, permitiendo que el estudiante aprenda cambios de trayectoria en lugar de ediciones de tokens aislados.

Resultados Experimentales
Los experimentos se realizaron utilizando modelos Qwen3 (destilando un profesor de 30B en un estudiante de 4B) en bancos de pruebas matemáticos competitivos (AIME24, AIME25, HMMT25-Feb).

• Ganancias de Rendimiento: La OPD estándar mejoró la precisión promedio del modelo de inicio en caliente (warm-start) del 38.9% al 47.8%. TOPD aumentó esto aún más hasta el 52.2%.
• Bancos de Pruebas Específicos: TOPD mejoró la precisión de AIME24 del 60.0% al 63.3% y de AIME25 del 46.7% al 53.3%.
• Estudios de Ablación:
  • Suprimir los tokens de alta pérdida y baja OT (falsas alarmas) en la OPD estándar produjo una ganancia modesta (+0.4%), confirmando que estos tokens son ruidosos.
  • Suprimir los tokens de alta OT y alta pérdida (divergencias reales) causó una caída significativa en el rendimiento, confirmando su necesidad.
  • El análisis de la corrección de trayectoria local mostró que, mientras que la OPD estándar redujo parcialmente la distancia de OT de ventana corta (de 0.538 a 0.350), TOPD logró una reducción sustancial (a 0.204), demostrando una capacidad superior para redirigir las trayectorias de futuro cercano.
• Eficiencia: TOPD introduce un costo computacional modesto, siendo el tiempo de entrenamiento 1.41$\times$ mayor que el de la OPD estándar, ya que el OT se calcula solo en las ventanas seleccionadas de alta pérdida.

Significado y Reivindicaciones
El artículo sostiene que las fallas de razonamiento en la OPD son fundamentalmente fenómenos a nivel de trayectoria en lugar de simples desajustes de tokens aislados. Al pasar de la edición de tokens punto a punto a una alineación consciente de las transiciones globales, TOPD permite una alineación de razonamiento más efectiva. Los autores posicionan este trabajo como un paso hacia objetivos de destilación más conscientes de la trayectoria, argumentando que la corrección exitosa del razonamiento requiere que el estudiante aprenda no solo qué token corregir, sino cómo navegar la trayectoria de razonamiento subsiguiente para permanecer consistente con el profesor.

Limitaciones
Los autores reconocen que TOPD incurre en costos computacionales adicionales debido a los cálculos de OT, se enfoca principalmente en la divergencia de corto horizonte (posiblemente omitiendo errores de contexto largo) y actualmente se valida principalmente en tareas de razonamiento matemático en lugar de código o tareas de lenguaje abierto. Además, la distancia de OT mide la proximidad de la trayectoria, pero puede no capturar la equivalencia semántica donde diferentes caminos conducen a la misma respuesta correcta.