SPG: Sandwiched Policy Gradient for Masked Diffusion Language Models

1. El Problema: Alineación de Modelos de Difusión con RL

Los Modelos de Lenguaje de Difusión (dLLMs), como el Modelo de Lenguaje de Difusión enmascarada (MDLM), han surgido como una alternativa eficiente a los modelos autoregresivos (AR) debido a su capacidad para decodificar múltiples tokens en paralelo, reduciendo la latencia de inferencia. Sin embargo, alinear estos modelos con preferencias humanas o recompensas específicas de tareas mediante Aprendizaje por Refuerzo (RL) presenta un desafío fundamental:

• Intratabilidad de la Verosimilitud: A diferencia de los modelos AR, la log-verosimilitud ($\log \pi_\theta(x|c)$) de un dLLM es computacionalmente intratable.
• Sesgo en los Gradientes de Política: Los métodos anteriores intentan eludir este problema utilizando el Límite Inferior de Evidencia (ELBO) como sustituto de la log-verosimilitud real.
  • El ELBO es un límite inferior ($ELBO \le \log \pi_\theta$).
  • Maximizar el ELBO es válido para recompensas positivas.
  • Sin embargo, minimizar el ELBO para recompensas negativas no garantiza la minimización de la log-verosimilitud real. Esto introduce un sesgo significativo en los gradientes de la política, impidiendo que el modelo aprenda efectivamente de los errores (retroalimentación negativa) y limitando el rendimiento de algoritmos avanzados de RL que utilizan recompensas relativas o negativas.

2. Metodología: SPG (Sandwiched Policy Gradient)

Para abordar estas limitaciones, los autores proponen SPG, un nuevo algoritmo de gradiente de política diseñado para dLLMs enmascarados. La idea central es "sándwich" (encerrar) la log-verosimilitud intratable utilizando límites superiores e inferiores de forma condicional según la recompensa.

A. El Objetivo Sándwich

En lugar de usar solo el ELBO, SPG define un objetivo de optimización que utiliza:

1. Para secuencias con ventaja positiva ($A_j \ge 0$): Se maximiza el ELBO (Límite Inferior de Evidencia). Esto asegura que las buenas respuestas se vuelvan más probables.
2. Para secuencias con ventaja negativa ($A_j < 0$): Se minimiza un Límite Superior de Evidencia (EUBO) tractable. Minimizar un límite superior garantiza que la probabilidad real de la secuencia disminuya, penalizando efectivamente las respuestas malas.

La función de objetivo se formula como:
$$J_{SPG}(\theta) = \mathbb{E} \left[ \frac{1}{g} \sum_{j=1}^g \left( \mathbb{1}_{A_j \ge 0} \cdot A_j \mathcal{L}_{ELBO} + \mathbb{1}_{A_j < 0} \cdot A_j \mathcal{L}_{EUBO} \right) \right]$$

B. Derivación del Límite Superior (EUBO)

Los autores derivan un límite superior tractable basado en la desigualdad de Rényi. Para un parámetro $\beta \ge 1$, el EUBO se expresa como:
$$\mathcal{L}_{EUBO} \approx \frac{1}{\beta} \sum_{i=1}^n \log \mathbb{E}_{t, z_t} \left[ w(t) \cdot \mathbb{1}(z_{t,i} = m) \cdot \pi_\theta^\beta(x_i | z_t) \right]$$
A diferencia del ELBO, donde el logaritmo está dentro de la esperanza, en el EUBO el logaritmo está fuera. Esto permite que, al minimizarlo, se reduzca la probabilidad de las secuencias indeseadas de manera más robusta.

C. Estrategia de Enmascaramiento por Bloques (Block-Wise Masking)

Para estimar estos límites mediante Monte Carlo de manera estable y eficiente, SPG introduce una estrategia de enmascaramiento específica:

• En lugar de enmascarar tokens aleatoriamente, la secuencia se divide en bloques.
• Se selecciona un bloque aleatorio para enmascarar, manteniendo los bloques anteriores limpios y los posteriores completamente enmascarados.
• Esto alinea mejor la distribución de datos durante el despliegue de la política (rollout) y la optimización, imitando la estrategia de generación semi-autoregresiva utilizada en modelos como LLaDA.

D. Mezcla de Límites (Mixture)

Para mitigar el sesgo y la varianza en la estimación del EUBO (que puede ser inestable con pocas muestras), SPG utiliza una mezcla convexa del ELBO y el EUBO para las trazas de ventaja negativa:
$$\mathcal{L}_{Mix} = \omega \cdot \mathcal{L}_{EUBO} + (1-\omega) \cdot \mathcal{L}_{ELBO}$$
Los autores demuestran teóricamente (Proposición 1) que existe un coeficiente de mezcla óptimo $\omega^*$ que reduce estrictamente la varianza del gradiente en comparación con usar solo uno de los límites, logrando un equilibrio entre la estabilidad del ELBO y la corrección fuerte del EUBO.

3. Contribuciones Clave

1. Algoritmo SPG: Un nuevo algoritmo de gradiente de política que reduce el sesgo en la optimización de dLLMs al utilizar límites variacionales sándwich basados en la recompensa.
2. Técnica de Enmascaramiento por Bloques: Una estrategia de estimación de Monte Carlo que mejora la estabilidad del objetivo de entrenamiento y la eficiencia de la optimización de la política.
3. Resultados de Estado del Arte (SOTA): Demostración empírica de que SPG supera significativamente a los algoritmos de RL existentes para dLLMs en tareas de razonamiento matemático, lógico y de codificación.

4. Resultados Experimentales

Los experimentos se realizaron sobre el modelo base LLaDA-8B-Instruct en cuatro benchmarks de razonamiento y dos de codificación. SPG se comparó contra métodos baselines como D1, WD1, UniGRPO y LLaDA-1.5.

Rendimiento en Benchmarks de Razonamiento (Longitud de generación 256):

• GSM8K: Mejora del 3.6% sobre el estado del arte anterior (86.1% vs 82.5%).
• MATH500: Mejora del 2.6% (40.0% vs 37.4%).
• Countdown: Mejora masiva del 18.4% (70.7% vs 52.3%).
• Sudoku: Mejora espectacular del 27.0% (94.0% vs 67.0%).

Rendimiento en Codificación:

• HumanEval: Mejora del 1.9%.
• MBPP: Mejora del 4.7%.

Análisis Adicional:

• Convergencia: SPG muestra una convergencia más rápida y estable en las dinámicas de recompensa en comparación con los baselines.
• Robustez: El modelo entrenado con SPG mantiene su superioridad bajo diferentes estrategias de inferencia (desenmascaramiento aleatorio vs. basado en confianza, tamaños de bloque variables), demostrando una gran generalización.
• Eficiencia: El costo computacional por actualización de gradiente es comparable al de los métodos baselines (como D1 y UniGRPO), ya que el costo principal sigue siendo el paso de despliegue (rollout).

5. Significado e Impacto

Este trabajo es significativo porque resuelve una barrera teórica fundamental en la aplicación de Aprendizaje por Refuerzo a los Modelos de Lenguaje de Difusión. Al demostrar que el uso exclusivo del ELBO es insuficiente y potencialmente perjudicial para las recompensas negativas, SPG establece un nuevo estándar para la alineación de dLLMs.

La capacidad de SPG para mejorar drásticamente el rendimiento en tareas de razonamiento lógico y matemático (donde la corrección de errores es crítica) sugiere que los dLLMs, cuando se entrenan correctamente con RL, pueden competir e incluso superar a los modelos autoregresivos en capacidades de razonamiento complejo, aprovechando al mismo tiempo sus ventajas de inferencia paralela. Esto abre la puerta a una nueva generación de modelos de lenguaje más eficientes y capaces.