Resumen Técnico: La Trayectoria como Docente: Flujo Discreto de Correspondencia en Pocos Pasos mediante Destilación Navegada por Energía

Enunciado del Problema

Los modelos de Flujo Discreto de Correspondencia (DFM) generan texto transformando iterativamente una secuencia fuente (por ejemplo, tokens aleatorios uniformes o tokens de máscara) en lenguaje coherente. Aunque la calidad coincide con la de los LLM autoregresivos, el proceso es computacionalmente prohibitivo, requiriendo a menudo cientos o miles de pasadas hacia adelante (pasos de refinamiento) por generación.

Las técnicas de destilación, como el Flujo Discreto de Correspondencia en Pocos Pasos (FS-DFM), intentan comprimir este proceso entrenando a un modelo estudiante para reproducir la salida iterativa del docente en significativamente menos pasos. Sin embargo, los estudiantes destilados a menudo rinden menos que el docente, incluso cuando el estudiante tiene capacidad suficiente. La explicación estándar atribuye esto a limitaciones en la capacidad del modelo.

Este artículo argumenta que el cuello de botella no es la capacidad del estudiante, sino la calidad de la trayectoria de entrenamiento. En la destilación DFM estándar (específicamente usando estimadores RK-4), los estados intermedios se construyen mediante "saltos estocásticos ciegos". En cada punto medio de la trayectoria, los logits de velocidad se convierten en tokens concretos mediante un proceso estocástico sin ninguna evaluación de la calidad de la secuencia. Una sola decisión estocástica deficiente en un punto medio temprano se propaga a través de los pasos subsiguientes, corrompiendo las evaluaciones de velocidad y el objetivo final. El estudiante imita fielmente estos objetivos "desviados", aprendiendo los errores acumulados en lugar de la trayectoria óptima.

Metodología: Flujo Discreto de Correspondencia con Forma de Trayectoria (TS-DFM)

Los autores proponen TS-DFM, un marco que reemplaza los saltos estocásticos ciegos con navegación guiada durante la fase de entrenamiento. Esta "configuración de navegación" ocurre exclusivamente durante el entrenamiento; la inferencia permanece sin cambios.

1. La Brújula de Energía

Se introduce un modelo de energía escalar y ligero ($E_\phi$) para evaluar la calidad de los estados de flujo parcialmente revelados.

• Arquitectura: Un transformador de 6 bloques (90 millones de parámetros) entrenado en FineWeb-Edu.
• Entrada: Evalúa secuencias de tokens crudas sin condicionamiento temporal, valorando la calidad intrínseca de la secuencia.
• Entrenamiento: La brújula se entrena con un objetivo de contraste de ruido utilizando "negativos conscientes de la generación". Estos negativos se construyen mediante casi-errores de velocidad: retroceder a lo largo del flujo para crear un estado menos revelado y luego ejecutar el docente congelado hacia adelante para reconstruirlo. Esto crea secuencias plausibles pero corruptas que la brújula debe distinguir de los estados de flujo reales.
• Rendimiento: La brújula alcanza una precisión del 98,5% al 99,8% en distinguir estados de flujo reales de corrupciones relevantes para la generación.

2. Navegación de Secuencia a Token

Para evitar el costo intratable de evaluar cada posición de token individualmente, TS-DFM emplea una política de dos fases en cada punto medio RK-4:

• Fase 1 (Nivel de Secuencia): Generar $K$ continuaciones candidatas diversas desde el estado actual utilizando transiciones CTMC con temperatura diversa. La brújula de energía evalúa todos los candidatos en una sola pasada hacia adelante por lotes, y se selecciona el candidato de menor energía. Esto asegura que la trayectoria entre en una región globalmente de baja energía.
• Fase 2 (Nivel de Token): El candidato seleccionado se refina a nivel de token utilizando la propia confianza de la red de velocidad como proxy de calidad. Las posiciones donde el candidato no coincide con la predicción de alta confianza de la velocidad son candidatas para reemplazo. Esta fase no incurre en costo adicional de inferencia del modelo de energía.
• Salvaguarda de Energía: Una verificación final de energía asegura que la secuencia refinada no aumente la energía más allá de una pequeña tolerancia ($\epsilon_{safe}$), evitando que los refinamientos locales degraden la calidad global.

3. Activación por Umbral de Tiempo

La navegación está controlada por un umbral de tiempo $\tau$. En tiempos de flujo tempranos ($t < \tau$), la secuencia está dominada por ruido fuente, haciendo ineficaz la discriminación basada en energía. La navegación se activa solo cuando ha surgido suficiente estructura lingüística ($t \ge \tau$), asegurando que la brújula opere sobre estados ricos en señal.

Contribuciones Clave

1. Reencuadre del Cuello de Botella: El artículo identifica la calidad de la trayectoria, y no la capacidad del estudiante, como la limitación principal en la destilación DFM en pocos pasos. Demuestra que los errores acumulados de los saltos estocásticos ciegos degradan la señal de supervisión.
2. Configuración de Navegación: Un marco novedoso que configura las trayectorias de entrenamiento reemplazando los saltos ciegos con una selección guiada mediante una brújula de energía.
3. Política de Secuencia a Token: Un enfoque eficiente en dos fases que combina la selección global de secuencias con el refinamiento local de tokens, evitando la necesidad de evaluaciones de energía por token.
4. Sobrecarga Solo de Entrenamiento: Todos los mecanismos de configuración se aplican únicamente durante el entrenamiento. El procedimiento de inferencia, la arquitectura y la velocidad del modelo estudiante permanecen idénticos a los modelos destilados estándar.

Resultados Experimentales

Los experimentos se realizaron en modelos de lenguaje de 170 millones y 1,3 mil millones de parámetros utilizando distribuciones fuente Uniforme y de Máscara.

• Rendimiento vs. Docente: En el modelo de 170M con fuente Uniforme, el estudiante TS-DFM en 8 pasos alcanza una perplejidad de GPT-2 de 56,1. Esto representa una reducción del 32% comparado con el docente de 1.024 pasos (82,8 PPL) y una reducción del 36% comparado con la línea base FS-DFM (87,6 PPL).
• Velocidad: El modelo TS-DFM de 8 pasos es 128 veces más rápido que el docente de 1.024 pasos mientras logra una mejor perplejidad.
• Fuente de Máscara: Las mejoras son aún más pronunciadas para la fuente de Máscara, donde TS-DFM logra una mejora de 5,6 veces sobre la línea base FS-DFM en 8 pasos.
• Estado del Arte: TS-DFM supera a los métodos existentes de generación discreta en pocos pasos (incluyendo Duo y SDTT) en todos los evaluadores (GPT-2, LLaMA-2, LLaMA-3), a pesar de usar modelos más pequeños y menos datos de entrenamiento que algunos competidores.
• Escalado: A 1,3 mil millones de parámetros, la brecha entre TS-DFM y FS-DFM se amplía (reducción del 41% en PPL en 8 pasos), y el estudiante configurado supera a su docente de 1.024 pasos en solo 8 pasos.
• Tareas de Razonamiento: El método se transfiere al razonamiento matemático (GSM8K) a escala de 1,3 mil millones, donde los estudiantes TS-DFM de 4 y 8 pasos superan la precisión del docente de pasos completos.

Significado y Afirmaciones

El artículo afirma que la calidad de la trayectoria es el cuello de botella fundamental en el flujo discreto de correspondencia en pocos pasos. Al tratar los estados intermedios como objetivos para la optimización de calidad en lugar de aceptarlos como fijos por el ruido estocástico, TS-DFM rompe el techo de rendimiento de los métodos de destilación actuales.

Los autores enfatizan que este enfoque es ortogonal al mecanismo de destilación en sí mismo; puede aplicarse sobre cualquier marco de destilación basado en trayectorias. El método logra una perplejidad del estado del arte sin aumentar el costo de inferencia, dependiendo únicamente de una sobrecarga de entrenamiento única (aproximadamente 2,2 veces el costo del FS-DFM estándar). Los resultados sugieren que para la generación discreta, el "docente" (el proceso de construcción de la trayectoria) es a menudo el factor limitante, y mejorar la trayectoria del docente produce mayores retornos que simplemente aumentar la capacidad del estudiante.