Resumen Técnico: Destilación de Refinamiento de Trayectoria (TRD)

1. Declaración del Problema: Fallo de Prefijo en la Destilación On-Policy

El artículo identifica una limitación estructural fundamental en la Destilación On-Policy (OPD) y la Autodestilación On-Policy (OPSD) para Modelos de Lenguaje Grande (LLMs), denominada fallo de prefijo (prefix failure).

En la OPD/OPSD estándar, un modelo estudiante genera un despliegue (trayectoria) $y_o$, y un modelo profesor proporciona una supervisión densa por cada token basada en esa trayectoria específica. Los autores argumentan que cuando la trayectoria del estudiante toma un camino de razonamiento incorrecto (un "prefijo fallido"), la señal de supervisión del profesor se ve estructuralmente comprometida:

• Mezcla Bimodal del Profesor: Cuando un prefijo es incorrecto, la distribución del profesor se convierte en una mezcla de dos modos: uno que continúa el prefijo fallido (para consistencia de la secuencia) y otro que pivota hacia la solución correcta.
• Gradientes Fragmentados: Las funciones de pérdida estándar por token (divergencia KL) evalúan la supervisión a lo largo de la trayectoria fallida congelada. En consecuencia, las actualizaciones de gradiente resultan fragmentadas. El profesor recomienda repetidamente un token de corrección (por ejemplo, "Espera" o "En realidad") condicionado al contexto erróneo (el prefijo fallado), en lugar de condicionarse al desarrollo del camino de corrección.
• Ineficacia de las Correcciones a Nivel de Token: Los remedios existentes (por ejemplo, recorte de pérdida, truncamiento top-K, reponderación) operan a nivel de pérdida de token. Intentan seleccionar tokens "informativos" o suprimir aquellos "ruidosos", pero no logran abordar la causa raíz: el propio prefijo fallido permanece estructuralmente intacto. Los pares de supervisión entregados al estudiante divergen del camino de corrección ideal, impidiendo que el modelo aprenda a recuperarse de los errores de manera efectiva.

2. Metodología: Destilación de Refinamiento de Trayectoria (TRD)

Para mitigar el fallo de prefijo, los autores proponen la Destilación de Refinamiento de Trayectoria (TRD), un método que desplaza la intervención del nivel de token al nivel de trayectoria.

Mecanismo Central

En lugar de utilizar la trayectoria bruta del estudiante $y_o$ directamente para la destilación, la TRD introduce un paso de refinamiento antes del cálculo de la pérdida:

1. Muestreo: El modelo estudiante $\pi_\theta$ genera un despliegue on-policy bruto $y_o \sim \pi_\theta(\cdot | x)$.
2. Refinamiento: El modelo profesor $\pi_T$ (o el propio estudiante en OPSD, condicionado por información privilegiada como la solución de referencia $y^*$) genera una trayectoria refinada $y_r$.
   • En OPD: $y_r \sim \pi_T(\cdot | x, y_o)$.
   • En OPSD: $y_r \sim \text{sg}[\pi_\theta(\cdot | x, y_o, y^*)]$.
   • Crucialmente, se le pide al profesor que revise el despliegue del estudiante guiado por la solución de referencia, "reescribiendo" efectivamente las porciones erróneas del prefijo mientras mantiene los patrones de razonamiento del estudiante donde estos son válidos.
3. Destilación: La trayectoria refinada $y_r$ se utiliza como el objetivo de supervisión para la pérdida estándar de emparejamiento de distribuciones (KL directa con vocabulario completo).

Justificación Teórica

La TRD aborda el problema del "gradiente fragmentado" asegurando que el contexto de supervisión crezca a lo largo del camino de corrección en lugar del prefijo fallido.

• Recuperación del Gradiente Ideal: Al supervisar a lo largo de $y_r$, el contexto para cada token incluye el prefijo corregido, permitiendo que el profesor proporcione una señal de corrección coherente y de múltiples pasos ($g_{ideal}$) en lugar de una señal repetitiva y desajustada al contexto ($g_{frag}$).
• Exploración: Incluso cuando $y_o$ es correcto, la TRD expone al estudiante a derivaciones válidas alternativas sugeridas por el profesor, ampliando la distribución de caminos de razonamiento correctos que el estudiante aprende.

Algoritmo

El proceso se resume en el Algoritmo 1:

1. Muestrear $y_o \sim \pi_\theta(\cdot | x)$.
2. Generar $y_r \sim \pi_T(\cdot | x, y_o)$.
3. Actualizar $\theta$ utilizando la divergencia KL entre $\pi_\theta$ y la distribución del profesor sobre $y_r$.

3. Contribuciones Clave

1. Identificación del Fallo de Prefijo: El artículo formaliza el "fallo de prefijo" como una causa estructural de degradación en la destilación on-policy, explicando por qué las intervenciones a nivel de token fallan al intentar recuperar los errores de razonamiento.
2. Corrección a Nivel de Trayectoria: Se propone la TRD como el primer diseño de optimización a nivel de trayectoria que apunta explícitamente a la recuperación de prefijos fallidos respetando la restricción de soporte on-policy, sin un costo computacional prohibitivo.
3. Mejora en Regímenes Duales: La TRD mejora el rendimiento en dos regímenes:
   • Corrección de Errores: Mitiga el fallo de prefijo corrigiendo los prefijos problemáticos antes de la destilación.
   • Exploración: Amplía la cobertura de razonamiento al exponer a los estudiantes a derivaciones válidas alternativas incluso cuando el despliegue original era correcto.
4. Aplicabilidad a la Autodestilación: El método se extiende naturalmente a la OPSD, donde el estudiante actúa como su propio profesor condicionado por información privilegiada (la verdad fundamental), eliminando la necesidad de un modelo profesor separado.

4. Resultados Experimentales

Los autores evaluaron la TRD en cinco benchmarks de matemáticas competitivas (AIME24/25, HMMT25, BeyondAIME, AMOBench) y benchmarks de generación de código (HumanEval+, MBPP+, LiveCodeBench) utilizando modelos Qwen3 (1.7B, 4B, 8B).

Hallazgos Clave

• Rendimiento Superior: La TRD supera consistentemente a los modelos base previos (Forward/Reverse KL, con/sin recorte/truncamiento) tanto en entornos OPD como OPSD.
• Ganancias Significativas en Tareas Difíciles: Las mejoras son más pronunciadas en los benchmarks más complejos.
  • AMOBench: La TRD logró una mejora relativa de ~50% en Pass@16 para Qwen3-8B bajo la configuración OPSD en comparación con el modelo base.
  • AIME24/25: La TRD mostró ganancias consistentes, particularmente para modelos más pequeños (por ejemplo, +4.6% en AIME24 para Qwen3-1.7B).
• Estabilidad: A diferencia de algunos modelos base que degradan el rendimiento en ciertos benchmarks (por ejemplo, Reverse KL bajando el rendimiento en AIME24), la TRD nunca cayó por debajo del rendimiento del modelo base.
• Eficiencia: En la configuración OSD, las trayectorias refinadas $y_r$ fueron significativamente más cortas (longitud mediana comprimida en ~9x) que los despliegues brutos $y_o$, reduciendo el tiempo de entrenamiento en aproximadamente un 60% a pesar del paso adicional de muestreo.
• Generación de Código: La TRD igualó o superó a los modelos base en HumanEval+ y MBPP+, aunque las ganancias en LiveCodeBench fueron limitadas, sugiriendo que la capacidad de refinamiento del profesor puede estar limitada por la dificultad de la tarea.

5. Significado y Reivindicaciones

El artículo afirma que la TRD representa un cambio necesario en la estrategia de post-entrenamiento para los LLMs. Los autores argumentan que:

• Las intervenciones a nivel de token son insuficientes: La naturaleza estructural del fallo de prefijo requiere correcciones a nivel de trayectoria.
• Se preserva el soporte on-policy: La TRD logra estas correcciones sin violar la restricción on-policy (a diferencia del uso de un experto fijo), lo que la hace compatible con los paradigmas de aprendizaje computacionalmente eficientes estándar.
• La auto-mejora es viable: El éxito de la TRD en el entorno OPSD demuestra que los modelos pueden refinar efectivamente sus propios caminos de razonamiento utilizando información privilegiada (verdad fundamental) sin requerir un modelo profesor separado y más grande, ofreciendo una vía ligera para la auto-mejora.

El trabajo sugiere que los futuros métodos de destilación deben considerar la integridad estructural de la trayectoria de razonamiento, no solo el alineamiento estadístico de los tokens individuales.