💬 NLP

Mask-Aware Policy Gradients for Diffusion Language Models

Este artículo presenta Mask-Aware Policy Gradients, un marco de aprendizaje por refuerzo que formaliza la generación de modelos de lenguaje de difusión con máscara como un proceso de decisión de dos etapas para optimizar conjuntamente la selección de tokens y el enmascaramiento de posiciones, superando así la intratabilidad de la verosimilitud de logaritmo y logrando un rendimiento de vanguardia en evaluaciones de razonamiento matemático y codificación.

Autores originales: Haran Raajesh, Kulin Shah, Adam Klivans, Philipp Krähenbühl

Publicado 2026-07-17

📖 3 min de lectura☕ Lectura para el café

CC BY 4.0

Autores originales: Haran Raajesh, Kulin Shah, Adam Klivans, Philipp Krähenbühl

Artículo original bajo licencia CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ Esta es una explicación generada por IA del artículo a continuación. No ha sido escrita ni avalada por los autores. Para mayor precisión técnica, consulte el artículo original. Leer descargo de responsabilidad completo

Imagina que estás intentando enseñarle a un robot cómo escribir una historia o resolver un problema matemático. Durante mucho tiempo, la mejor manera de hacer esto era hacer que el robot escribiera una palabra a la vez, como un humano tecleando en un teclado. Esto se llama generación "autorregresiva". Es fiable, pero también es lenta, como un caracol cruzando una autopista. Recientemente, los científicos descubrieron una forma más rápida llamada "Difusión". En lugar de escribir desde cero, el robot comienza con una página llena de máscaras en blanco (como [MASK]) y las va completando gradualmente, revelando palabras una por una hasta que la oración está completa. Es como un rompecabezas donde comienzas con una imagen cubierta y lentamente levantas las cubiertas para ver la imagen.

Sin embargo, enseñar a estos robots de "difusión" rápidos a volverse más inteligentes mediante recompensas (una técnica llamada Aprendizaje por Refuerzo) ha sido una pesadilla. En el antiguo método "palabra por palabra", es fácil calcular exactamente qué tan probable era que el robot hiciera una elección específica. Pero en el método de "difusión", el robot toma dos decisiones en cada paso: decide qué palabra poner en un espacio en blanco y también decide qué espacio en blanco revelar después. Los métodos anteriores intentaron enseñar al robot mirando solo la primera elección (la palabra) e ignorando la segunda elección (el orden de revelación). Es como intentar enseñarle a un chef a cocinar una comida perfecta criticando únicamente los ingredientes que eligió, mientras se ignora por completo el orden en el que picó y revolvió.

Este artículo, presentado en la conferencia COLM 2026, sugiere que nos ha estado faltando una parte enorme del rompecabezas. Los autores, Haran Raajesh, Kulin Shah y su equipo de la Universidad de Texas en Austin, argumentan que para dominar verdaderamente el arte de la difusión, tienes que recompensar al robot tanto por elegir las palabras correctas como por elegir el orden correcto para revelarlas. Desarrollaron un nuevo método llamado "Gradientes de Política Conscientes de la Máscara" (Mask-Aware Policy Gradients). En lugar de solo adivinar la siguiente palabra, su sistema trata la decisión de "qué máscara levantar después" como una parte crucial de la estrategia del robot. Al descomponer matemáticamente el proceso de aprendizaje en estas dos partes distintas, descubrieron que el robot aprende mucho más rápido y comete menos errores.

Los resultados son bastante impresionantes. Cuando probaron su nuevo método en problemas matemáticos difíciles y desafíos de programación, el robot mejoró significativamente al resolverlos. En un famoso examen de matemáticas llamado GSM8K, su método alcanzó una precisión del 87.1%, y en un examen de programación llamado MBPP, alcanzó el 53.4%. Estos números son más altos que cualquier método anterior que haya intentado enseñar modelos de difusión utilizando recompensas. El artículo muestra que, al prestar atención a las decisiones de "enmascaramiento" —el orden en el que el robot revela sus pensamientos—, podemos desbloquear un nuevo nivel de inteligencia en estos modelos de IA rápidos y flexibles, haciéndolos no solo más rápidos, sino también más inteligentes.

Resumen Técnico: Gradientes de Política Conscientes de la Máscara para Modelos de Lenguaje de Difusión

Planteamiento del Problema

El aprendizaje por refuerzo (RL) ha mejorado con éxito las capacidades de razonamiento en los Grandes Modelos de Lenguaje (LLMs) autorregresivos mediante la optimización de gradientes de política basados en log-verosimilitudes tratables. Sin embargo, extender estos métodos a los Modelos de Lenguaje de Difusión con Máscara (MDLMs) sigue siendo un desafío. A diferencia de los modelos autorregresivos, que generan tokens secuencialmente de izquierda a derecha, los MDLMs generan texto desmascarando iterativamente los tokens de una secuencia completamente enmascarada.

La dificultad central reside en la intratabilidad de la estimación de la log-verosimilitud para los MDLMs. Los enfoques existentes aproximan esta verosimilitud utilizando variantes de la Cota Inferior de Evidencia (ELBO) o modelando únicamente el componente de predicción de tokens de la trayectoria de generación. Estos métodos ignoran un aspecto crítico de la generación de los MDLMs: el orden en el que se desmascaran las posiciones. En los MDLMs, la trayectoria de generación implica dos decisiones distintas en cada paso: (1) predecir qué token colocar en las posiciones actualmente enmascaradas, y (2) decidir cuáles de esas posiciones mantener desmascaradas (y cuáles volver a enmascarar). Los métodos previos que ignoran la segunda decisión no logran utilizar toda la señal disponible en el proceso de generación para la optimización de la política.

Metodología

Los autores formalizan la generación de MDLM como un Proceso de Decisión de Markov (MDP) de dos etapas de acción. En cada paso de eliminación de ruido $t$ , el modelo realiza:

Predicción de Tokens: Predecir los tokens para todas las posiciones actualmente enmascaradas.
Selección de Posición (Remascarado): Decidir qué subconjunto de posiciones reveladas mostrar y cuáles devolver al estado [MASK].

Para permitir la optimización del gradiente de la política sobre este proceso, los autores introducen el Remascarado Probabilístico. La inferencia estándar de MDLM suele utilizar una estrategia determinista y codiciosa (greedy) de top- $K$ para seleccionar las posiciones a desmascarar basándose en puntuaciones de confianza. Esta selección codiciosa es no diferenciable, lo que impide la optimización basada en gradientes del orden de desmascarado. El método propuesto reemplaza esto con una variante probabilística:

En lugar de seleccionar determinísticamente las posiciones top- $K$ , el modelo muestrea un subconjunto de posiciones $U_t$ de una distribución Plackett–Luce.
Las probabilidades para esta distribución se derivan directamente de los propios log-verosimilitudes (logits) del modelo, escalados por un parámetro de temperatura $\tau$ .
Esta formulación crea una distribución diferenciable sobre los subconjuntos de posiciones, permitiendo que la log-probabilidad de la decisión de desmascarado se incluya en la verosimilitud de la trayectoria.

Descomposición del Gradiente de la Política

Al modelar la verosimilitud completa de la trayectoria $\pi_\theta(\hat{z}|c)$ como un producto de las probabilidades de predicción de tokens y las probabilidades de desmascarado, el gradiente de la política se descompone naturalmente en dos términos:
$\nabla_\theta J(\theta) = \mathbb{E}_{\hat{z}} \left[ R(c, z_T) \sum_{t=1}^T \left( \underbrace{\nabla_\theta \log \pi_\theta(\hat{z}_t | c, z_{t-1})}_{\text{Gradiente de Token}} + \underbrace{\nabla_\theta \log p_{\text{unmask}}(U_t | \hat{z}_t, z_{t-1}, c)}_{\text{Gradiente de Máscara}} \right) \right]$

Término de Token: Refuerza la corrección de los tokens predichos.
Término de Máscara: Refuerza la selección del orden de desmascarado.
Eficiencia: El término de máscara se calcula a partir de los logits existentes del modelo utilizados para la predicción de tokens, lo que requiere ningún parámetro adicional, cambios arquitectónicos o pasadas hacia adelante extra.

Los autores implementan esto utilizando el algoritmo Optimización de Política de Secuencia de Grupo (GSPO), recortando las razones de importancia para los términos de token y posición por separado. Para manejar el costo computacional de evaluar trayectorias completas, emplean la aproximación StepMerge, agrupando los pasos de eliminación de ruido en segmentos para reducir el número de pasadas hacia adelante preservando la estructura de la trayectoria.

Contribuciones Clave

Formalización de MDLM como un MDP de dos etapas: El artículo proporciona un marco teórico que muestra que la generación de MDLM involucra tanto la selección de tokens como la selección de posiciones, y que el gradiente de la política se descompone de esa manera.
Remascarado Probabilístico: La derivación de una distribución de desmascarado diferenciable (Plackett–Luce) basada en los propios logits del modelo, permitiendo la optimización del orden de desmascarado sin modificaciones arquitectónicas.
Perspectiva Teórica: Demostración de que ignorar el componente de posición del gradiente de la política puede resultar en la pérdida de direcciones que mejoran el retorno esperado, incluso cuando las probabilidades de los tokens permanecen inalteradas.
Rendimiento de Vanguardia (State-of-the-Art): Validación empírica que muestra que la optimización conjunta de ambos términos produce mejoras consistentes sobre los métodos basados en trayectorias y en ELBO existentes.

Resultados Experimentales

El método fue evaluado en tareas de razonamiento matemático (GSM8K, MATH500) y generación de código (HumanEval, MBPP), utilizando LLaDA-8B-Instruct como modelo base.

Ganancias de Rendimiento: El método propuesto alcanzó nuevos resultados de vanguardia en todos los benchmarks.
- GSM8K: 87.1% de precisión (con longitud de generación 128), una mejora de +2.5% sobre el anterior mejor resultado (StepMerge).
- MATH500: 53.4% de precisión, una mejora de +4.0%.
- HumanEval: 47.1% de precisión, una mejora de +2.2%.
- MBPP: 53.4% de precisión, una mejora de +2.5%.
Estudios de Ablación:
- Tamaño de Bloque: La brecha de rendimiento entre el método propuesto y los baselines se amplió a medida que aumentaba el tamaño del bloque de inferencia (de 32 a la secuencia completa), lo que sugiere que el beneficio de optimizar la selección de posición crece con la complejidad del espacio de decisión de desmascarado.
- Eficiencia: El método convergió más rápido en tiempo de ejecución real en comparación con el fuerte baseline basado en ELBO (SPG), alcanzando la precisión final de SPG en ~15 horas frente a ~18 horas, y logrando una precisión final más alta (81.0% vs 78.5% en GSM8K).
- Generalización: El enfoque se generalizó a un modelo base diferente (Dream-7B) y tareas de planificación (Sudoku, Countdown), superando consistentemente a los baselines.
Comparación con DCoLT: Al compararse con LLaDOU (que utiliza una cabeza separada para la selección de posición), el método propuesto logró una precisión comparable o mejor con costos de entrenamiento significativamente menores (~160 horas de GPU vs ~800), ya que no requiere módulos adicionales.

Significado y Reivindicaciones

El artículo sostiene que la flexibilidad de los MDLMs conlleva una estructura de trayectoria más rica que la de los modelos autorregresivos, específicamente la decisión de cuándo desmascarar posiciones. Al tratar esta decisión como parte de la política y optimizarla mediante un marco probabilístico diferenciable, los autores demuestran que la optimización conjunta de las predicciones de tokens y la selección de posiciones es crucial para maximizar el rendimiento de los modelos de lenguaje de difusión en entornos de RL.

Los autores enfatizan que su enfoque es prácticamente eficiente, ya que aprovecha los logits existentes del modelo sin requerir nuevos parámetros o cambios arquitectónicos, lo que lo hace compatible con infraestructuras de entrenamiento de MDLM estándar como StepMerge. Los resultados sugieren que el "gradiente de máscara" proporciona una señal de entrenamiento significativa que los métodos previos, centrados únicamente en las verosimilitudes de los tokens, no lograron capturar.