Resumen Técnico: BlockGen: Modelado de Secuencias por Bloques Flexible con Muestreadores Híbridos

Planteamiento del Problema

Los Transformers autorregresivos (AR) dominan el modelado de secuencias, pero sufren de limitaciones de generación secuencial (rendimiento/latencia) y restricciones de atención causal que dificultan las tareas de razonamiento que requieren contexto bidireccional. Los modelos de difusión discreta ofrecen una alternativa al refinar iterativamente secuencias ruidosas, permitiendo actualizaciones paralelas. Sin embargo, comparaciones recientes entre los Modelos de Difusión de Máscara (MDM) y los Modelos de Difusión de Estado Uniforme (USDM) han arrojado perspectivas contradictorias dependiendo del régimen de muestreo:

1. Muestreo Ancestral: Los USDM suelen superar a los MDM en tareas de downstream a pesar de tener una menor perplejidad, lo que se atribuye a la capacidad de los USDM para autocorregir errores (los tokens pueden transicionar entre cualquier estado), mientras que los MDM fijan los tokens una vez que dejan de estar enmascarados.
2. Generación por Bloques: El trabajo previo compara mayoritariamente la difusión de secuencia completa. No está claro si la ventaja de los USDM se mantiene cuando los tokens se generan bloque a bloque (semi-autorregresivamente), un entorno que permite el uso de caché KV y una inferencia más rápida.
3. Correctores Informados: Las comparaciones existentes suelen utilizar muestreadores predictor-corrector (PC) "no informados" que reinyectan ruido en posiciones aleatorias. Se desconoce cómo los correctores "informados" (que apuntan a errores probables) afectan el ranking entre MDM y USDM, particularmente sin modelos de verificación auxiliares.

Metodología: BlockGen

Los autores presentan BlockGen, un modelo de secuencia por bloques entrenado sobre una mezcla de tamaños de bloque, instanciado ya sea con difusión enmascarada o uniforme dentro de los bloques.

1. Formulación de la Mezcla y Entrenamiento

BlockGen define la densidad de la secuencia como una mezcla sobre $M$ tamaños de bloque $S = \{s_1, \dots, s_M\}$:
$$p_{\theta}^{\text{BlockGen}}(x) = \sum_{i=1}^M \gamma_i p_{\theta}^{(s_i)}(x)$$
donde $p_{\theta}^{(s_i)}$ es un factor de densidad válido factorizado en bloques de tamaño $s_i$.

• Denoiser Compartido: En lugar de entrenar modelos separados, un único denoiser compartido $\hat{x}_\theta$ toma como entrada el tamaño de bloque $L'$, ajustando los patrones de atención en consecuencia.
• Límites de Verosimilitud: El artículo deriva dos límites para la densidad de la mezcla. Se utiliza un límite log-sum-exp más ajustado para la evaluación, mientras que un límite de verosimilitud de mezcla más laxo (una esperanza sobre los tamaños de bloque) se utiliza para el entrenamiento para mantener la eficiencia computacional.
• Muestreo Estratificado: Para reducir la varianza durante el entrenamiento multi-GPU, los tamaños de bloque se seleccionan mediante muestreo estratificado en lugar de muestreo i.i.d.

2. Predictor-Corrector Informado por AR (ARPC)

Una contribución clave es el muestreador ARPC, que aprovecha la mezcla de BlockGen para realizar correcciones informadas sin un verificador auxiliar.

• Mecanismo: Durante la generación, tras los pasos de difusión ancestral estándar, el muestreador ejecuta un paso de "propuesta limpia". Luego realiza una segunda pasada hacia adelante en modo AR ($L'=1$) utilizando el mismo denoiser para puntuar los tokens propuestos.
• Puntuación: Se identifican los tokens con baja log-verosimilitud AR como "improbables" y se vuelven a enrarecer (corregir) a un nivel de ruido anterior.
• Ventaja: Esto permite que el modelo utilice sus propias predicciones AR para guiar las correcciones de difusión, intercalando efectivamente los paradigmas AR y de difusión dentro de un solo modelo.

3. Atención a Nivel de Bloque

BlockGen utiliza atención de causalidad de bloque: los tokens dentro de un bloque atienden bidireccionalmente, mientras que el bloque atiende causalmente a todos los bloques precedentes. Esto permite el uso de caché KV a través de los bloques decodificados, preservando los beneficios de eficiencia de la generación semi-autorregresiva.

Resultados Clave

1. Razonamiento Matemático (GSM8K)

Evaluado en GSM8K utilizando modelos entrenados en TinyGSM:

• Muestreo Ancestral: En el entorno de bloque a bloque, los USDM superan a los MDM, especialmente en el régimen de pocos pasos (bajo NFE). Esto confirma que la ventaja de la secuencia completa de los USDM se traslada a la generación por bloques.
• Muestreo ARPC: La brecha de rendimiento entre los USDM y los MDM se estrecha y se invierte con un alto NFE. Bajo ARPC con bloques de 16 tokens, los MDM superan ligeramente a los USDM en precisión de GSM8K.
• ARPC vs. Baselines: ARPC supera consistentemente tanto al muestreo ancestral como al EIPC (Entropy-Informed PC) en la mayoría de los presupuestos de NFE. ARPC con MDM se acerca al rendimiento del muestreo AR con decodificación codiciosa (greedy).

2. Modelado de Lenguaje (OpenWebText & LM1B)

• Verosimilitud: Entrenar sobre una mezcla de tamaños de bloque mejora significamente la perplejidad en comparación con modelos de bloque fijo. BlockGen con difusión enmascarada logra una perplejidad de validación (Val. PPL) de 17.5 en OpenWebText (frente a 16.7 para AR y 21.6 para el mejor BDM de bloque fijo), reduciendo la brecha entre los modelos de difusión y los autorregresivos.
• Calidad de la Muestra:
  • Bajo muestreo ancestral, los USDM logran una menor Perplejidad Generativa (Gen. PPL) que los MDM.
  • Bajo ARPC, la tendencia se invierte con un NFE más alto (ej. 64 NFE/bloque), donde el MDM-ARPC logra una frontera de Gen. PPL mejor que el uniforme-ARPC.
  • ARPC reduce la brecha con AR, pero no supera a AR en Gen. PPL dentro del régimen práctico de baja temperatura.

Significancia y Reivindicaciones

El artículo afirma resolver dos preguntas abiertas sobre la difusión discreta:

1. Validez por Bloques: La ventaja de la difusión de estado uniforme sobre la difusión enmascarada observada en la generación de secuencia completa se traslada a la generación bloque a bloque bajo muestreo ancestral.
2. Corrección Informada: El ranking entre MDM y USDM depende del régimen. Mientras que los USDM sobresalen con muestreadores ancestrales simples o no informados, los correctores informados (específicamente ARPC) pueden cerrar esta brecha e incluso invertir el ranking, favoreciendo la difusión enmascarada con presupuestos de cómputo (NFE) más altos.

Los autores enfatizan que BlockGen proporciona un marco flexible que interpola entre AR y difusión pura, permitiendo que un solo modelo actúe tanto como muestreador por bloques como verificador AR. Señalan que, aunque sus afirmaciones empíricas son robustas dentro de los presupuestos y muestreadores probados, no establecen si estos intercambios persisten a la escala de los grandes modelos de lenguaje (LLM), ya que sus experimentos se limitan a arquitecturas de 170 millones de parámetros. El trabajo sugiere que la elección del paradigma de difusión (enmascarada vs. uniforme) no es universalmente superior, sino que depende fuertemente de la estrategia de muestreo y el presupuesto de cómputo.