Resumen Técnico: Entrenamiento Posterior de Equilibrio de Trayectoria para Modelos de Lenguaje Difusivos

Problema: Bloqueo de Trayectoria en RL Difusivo

Los Modelos de Lenguaje Difusivos (dLLMs) ofrecen una alternativa prometedora a los modelos autoregresivos para el razonamiento y la generación de código. Sin embargo, los métodos existentes de entrenamiento posterior para dLLMs adaptan en gran medida objetivos de Aprendizaje por Refuerzo (RL) de maximización de recompensas (por ejemplo, PPO, GRPO) diseñados para modelos autoregresivos. Los autores identifican un modo de fallo central en esta adaptación denominado bloqueo de trayectoria.

En los modelos difusivos, la recompensa se asigna únicamente a la finalización desruidada final ($y$), mientras que la distribución del modelo se induce sobre toda la trayectoria de desruido ($\tau$). Debido a que las recompensas terminales son "indiferentes a la trayectoria" (es decir, cualquier trayectoria que alcance la misma $y$ correcta recibe la misma recompensa), las actualizaciones de gradiente de política muestreadas tienden a reforzar las trayectorias específicas que ocurren a muestrearse temprano en el entrenamiento. Esto crea un bucle de retroalimentación auto-reforzante donde la masa de probabilidad colapsa sobre un conjunto estrecho de rutas de desruido. En consecuencia, incluso si un modelo aprende a producir una respuesta correcta, pierde la capacidad de generar un conjunto diverso de soluciones válidas (modos) bajo muestreo repetido. Los autores formalizan esto mediante el Teorema 1, demostrando que la cobertura de modos terminales está acotada superiormente por la diversidad de trayectorias; por lo tanto, el bloqueo de trayectoria limita directamente la cobertura de soluciones del modelo.

Metodología: TraFL (Equilibrio de Flujo de Trayectoria)

Para abordar el bloqueo de trayectoria, los autores proponen TraFL, un objetivo de entrenamiento posterior fundamentado en el principio de equilibrio de trayectoria de las Redes de Flujo Generativo (GFlowNets). En lugar de maximizar la recompensa directamente, TraFL entrena la política para igualar una distribución objetivo inclinada por recompensa anclada a un modelo de referencia congelado.

El Objetivo

La distribución objetivo $p^*(\tau | x)$ se define como:
$$p^*(\tau | x) = \frac{1}{Z(x)} p_{\text{ref}}(\tau | x) \exp(\beta r(x, y))$$
donde:

• $p_{\text{ref}}$ es un modelo de referencia congelado.
• $r(x, y)$ es la recompensa terminal.
• $\beta$ controla la nitidez de la recompensa.
• $Z(x)$ es un término de normalización dependiente del prompt (función de partición).

El objetivo de entrenamiento minimiza el residuo cuadrado entre el log-probabilidad del modelo actual y el objetivo:
$$\delta(\tau, x) = \log p_\theta(\tau | x) - \log p_{\text{ref}}(\tau | x) - \beta r(x, y) + \log Z(x)$$
$$L_{\text{TraFL}} = \mathbb{E}_{\tau \sim p_\theta} [\delta(\tau, x)^2]$$

Implementación Práctica para Difusión

Dado que los dLLMs no exponen log-probabilidades de trayectoria exactas, los autores introducen dos componentes clave para hacer TraFL práctico:

1. Sustituto de Nivel de Secuencia Compatible con Difusión: Derivan un sustituto de log-probabilidad tratable ($\log \hat{p}_\theta(y|x)$) basado en una cota inferior de reconstrución enmascarada. Esto implica muestrear un nivel de enmascaramiento $l$, corromper $l$ tokens y promediar las log-probabilidades de los tokens reconstruidos. Se utiliza un factor $1/l$ para normalizar la puntuación independientemente del número de tokens enmascarados.
2. Normalización Aprendida Dependiente del Prompt: A diferencia de los métodos que estiman $Z(x)$ mediante softmax sobre un lote (lo cual puede inducir concentración), TraFL aprende una función de partición $Z_\phi(x)$ mediante una pequeña MLP. Este término recibe señales de gradiente en cada paso independientemente de qué finalizaciones se muestreen, evitando el fallo de normalización "local al lote" observado en otros enfoques de coincidencia de distribuciones.

Contribuciones Clave

1. Identificación del Bloqueo de Trayectoria: El artículo identifica y demuestra formalmente que los objetivos de maximización de recompensas en dLLMs conducen a un colapso de la diversidad de trayectorias, lo que a su vez limita la cobertura de modos de solución correctos distintos.
2. Objetivo TraFL: La propuesta de un objetivo de equilibrio de trayectoria que ancla la política a un modelo de referencia, separando el papel de la recompensa (inclinando las finalizaciones terminales) del modelo de referencia (asignando masa de probabilidad sobre trayectorias plausibles).
3. Sustituto Práctico y Normalización: El desarrollo de un sustituto de nivel de secuencia compatible con difusión y un término de normalización aprendido que permite un entrenamiento estable sin probabilidades de trayectoria exactas.
4. Superioridad Empírica: TraFL es el único método de entrenamiento posterior evaluado que mejora consistentemente sobre el modelo base en todas las configuraciones de longitud de benchmark (GSM8K, MATH-500, HumanEval, MBPP) y presupuestos de muestreo.

Resultados

Los autores evalúan TraFL en benchmarks de razonamiento matemático y generación de código frente al modelo base (LLaDA-8B-Instruct) y baselines fuertes (ESPO, JUSTGRPO).

• Rendimiento In-Distribution: TraFL mejora Pass@k en todos los presupuestos de muestreo ($k \in \{1, 2, 4, 8, 16\}$) y temperaturas. Cabe destacar que mejora Pass@1 (precisión de muestra única) mientras también escala mejor con presupuestos de muestreo aumentados, a diferencia de los baselines que a menudo sacrifican la precisión de muestra única para obtener ganancias de múltiples muestras.
• Generalización Fuera de Muestra: Las ganancias se transfieren a benchmarks fuera de distribución. En Minerva Math, TraFL se mantiene por encima del modelo base. En LiveCodeBench, TraFL es el método más fuerte en todas las divisiones de dificultad (Fácil, Medio, Difícil), superando significativamente a los baselines en problemas medios y difíciles.
• Análisis de Diversidad: Utilizando un protocolo de LLM como juez, los autores muestran que las ganancias de TraFL están asociadas con una cobertura más amplia de estrategias de solución distintas (por ejemplo, diferentes enfoques de prueba o descomposiciones algorítmicas) en lugar de simplemente afinar la puntuación de un solo modo. Los baselines como ESPO, aunque mejoran la precisión, no muestran ganancias significativas de diversidad sobre el modelo base, lo cual es consistente con la hipótesis de bloqueo de trayectoria.

Significado y Afirmaciones

El artículo afirma que TraFL proporciona una solución principial al fallo de "búsqueda de modos" del entrenamiento posterior estándar por RL para modelos difusivos. Al cambiar de la maximización de recompensas al equilibrio de trayectoria, el método preserva la diversidad de rutas de desruido necesaria para cubrir múltiples soluciones válidas.

Los autores declaran que TraFL es el único método evaluado que mejora sobre el modelo base en cada configuración de longitud de benchmark probada. Subrayan que las mejoras son robustas ante aumentos en el presupuesto de muestreo y se generalizan a conjuntos de datos retenidos, lo que sugiere que el objetivo de equilibrio de trayectoria desacopla efectivamente la señal de recompensa del mecanismo de asignación de rutas. El trabajo aboga por objetivos de coincidencia de distribuciones sobre la maximización pura de recompensas para el entrenamiento posterior de dLLM, ofreciendo un marco práctico que evita los modos de fallo de concentración inherentes a los enfoques estándar de gradiente de política.

El artículo se mantiene modesto respecto a las limitaciones, señalando que el método depende de un sustituto aproximado de nivel de secuencia en lugar de probabilidades de trayectoria exactas y utiliza recompensas binarias gruesas (coincidencia exacta o aprobación de prueba), lo cual puede limitar la eficiencia de aprendizaje en escenarios con éxitos raros.