Resumen Técnico: Difusión de Conjuntos (Set Diffusion)

Planteamiento del Problema

Los modelos de difusión discretos han surgido como una alternativa competitiva a los modelos de lenguaje autorregresivos (AR), ofreciendo velocidades de generación más rápidas y un control mejorado. Sin embargo, los enfoques de difusión existentes enfrentan limitaciones significativas:

1. Restricciones de Longitud Fija y Bidireccionalidad: La mayoría de los modelos de difusión discretos dependen de arquitecturas que requieren un contexto bidireccional completo, lo que impide el uso de caché de Clave-Valor (KV) y restringe la generación a longitudes fijas.
2. Rigidez de la Difusión por Bloques: Aunque la Difusión por Bloques (BD3LM) cierra parcialmente la brecha al generar bloques secuenciales de tamaño fijo de izquierda a derecha, sufre de inflexibilidad. No puede soportar inserciones en posiciones arbitrarias (infilling) ni la generación de longitud variable dentro de un bloque, y las actualizaciones de la caché KV se retrasan hasta que se decodifica un bloque entero.
3. Limitaciones de Compromiso (Trade-off): Los métodos actuales interpolan entre la autorregresión y la difusión principalmente variando un tamaño de bloque fijo, lo que restringe la flexibilidad de decodificación y el paralelismo.

El desafío central es desarrollar un modelo de lenguaje que interpole entre la autorregresión y la difusión no solo mediante el tamaño del bloque, sino variando el orden en el que se generan los conjuntos de tokens, permitiendo así la generación de longitud flexible, el uso de caché KV y la decodificación de orden arbitrario, manteniendo una alta calidad de muestreo.

Metodología

Los autores introducen Set Diffusion (Difusión de Conjuntos), un marco que factoriza el proceso de generación sobre conjuntos de tokens de posición y longitud flexibles, en lugar de bloques fijos.

1. Modelo Probabilístico

La difusión de conjuntos define una distribución de probabilidad autorregresiva sobre una secuencia de conjuntos de tokens $\sigma = (\sigma_1, \dots, \sigma_N)$, donde cada $\sigma_n$ es un subconjunto de posiciones de tokens. La verosimilitud se factoriza como:
$$p_\theta(x) = \sum_{\sigma} \prod_{n=1}^N \pi(\sigma_n | x_{<\sigma_n}) \cdot p_\theta(x_{\sigma_n} | x_{<\sigma_n})$$
Aquí, $\pi$ es una política de orden que selecciona el siguiente conjunto de posiciones, y $p_\theta$ modela la verosimilitud condicional de los tokens en ese conjunto utilizando difusión. El modelo optimiza un Límite Inferior de Evidencia Negativa (NELBO) que suma las pérdidas de difusión sobre estos conjuntos.

2. Arquitectura Set-Causal y Caché KV

Una innovación arquitectónica clave es el transformer set-causal. A diferencia de los modelos de difusión bidireccionales estándar, esta arquitectura soporta la caché KV después de cada paso de inferencia.

• Mecanismo: Cuando se genera un conjunto de tokens $x_{\sigma_n}$, sus vectores de clave (key) y valor (value) se almacenan en caché y se reutilizan para condicionar las predicciones subsecuentes.
• Eficiencia de Entrenamiento: Los autores proponen un algoritmo de entrenamiento vectorizado que procesa el contexto limpio y las consultas ruidosas simultáneamente en una sola pasada hacia adelante utilizando una máscara de atención especializada. Esta máscara asegura que los tokens ruidosos en un conjunto atiendan solo a otros tokens ruidosos en el mismo conjunto y a todos los tokens limpios en los conjuntos precedentes.

3. Interpolación de Ordenamientos de Tokens

Para controlar el compromiso entre la calidad de la muestra (estrechez del límite) y la eficiencia de la inferencia (paralelismo), los autores introducen programas de tiempo de revelación dependientes de la posición.

• Programas de Desplazamiento de Posición (Position-Offset Schedules): Cada posición de token $\ell$ recibe un tiempo de revelación aleatorio $R_\ell$ extraído de una distribución definida por una función de distribución acumulativa (CDF) $\alpha^\ell_\tau$.
• Instanciación de Ventana Deslizante (SW-SetDLM): Los autores proponen una instanciación específica donde los tokens tienen un intervalo de generación activo $[a_\ell, a_\ell + w]$. El tamaño de la ventana $w$ controla el sesgo de "izquierda a derecha":
  • Un $w$ pequeño (aproximándose a $1/L$) concentra el orden en la generación de izquierda a derecha, recuperando el comportamiento AR y estrechando el ELBO.
  • Un $w$ grande (aproximándose a 1) permite la generación agnóstica al orden, recuperando la difusión de máscara estándar.
  • Los valores intermedios permiten la decodificación de ventana deslizante, permitiendo al modelo predecir conjuntos de tokens de longitud flexible en posiciones arbitrarias, soportando tareas como el infilling.

4. Inferencia

Durante la inferencia, el modelo muestrea posiciones candidatas basadas en el tiempo de ordenamiento actual y el programa de desplazamiento de posición. Luego, denoise (elimina el ruido) de estas posiciones en paralelo utilizando la caché KV. Los tokens aceptados se consolidan y la caché se actualiza inmediatamente, permitiendo la generación de longitud variable y una decodificación eficiente.

Contribuciones Clave

1. Marco de Difusión de Conjuntos (Set Diffusion Framework): Una nueva clase de modelos de lenguaje que factoriza la generación sobre conjuntos de tokens de posición y longitud flexibles, generalizando la autorregresión, la difusión agnóstica al orden y la difusión por bloques.
2. Arquitectura Set-Causal: Una arquitectura transformer que soporta actualizaciones de caché KV después de cada paso de inferencia, permitiendo la generación eficiente de longitud variable.
3. Interpolación de Orden: Un método principled para interpolar entre AR y difusión ajustando los programas de ordenamiento dependientes de la posición (específicamente los programas de ventana deslizante) en lugar de tamaños de bloque fijos.
4. Rendimiento de Vanguardia (State-of-the-Art): Demostración de que la difusión de conjuntos logra compromisos de velocidad-calidad superiores en razonamiento matemático, resumen y generación incondicional comparado con modelos de difusión previos, y supera significativamente a la difusión por bloques en tareas de infilling.

Resultados Experimentales

Los autores evaluaron Set Diffusion (específicamente SW-SetDLM) en varios benchmarks:

• Razonamiento Matemático (GSM8K): La difusión de conjuntos logra una frontera de Pareto de velocidad-precisión estrictamente mejor que la Difusión por Bloques. Por ejemplo, con un tamaño de ventana de decodificación $S \le 8$, alcanza una precisión pass@1 del 66.41% con un rendimiento de $\ge 60.42$ tokens/seg, superando a la Difusión por Bloques ($S=8$) que logra 63.53% de precisión a $\ge 55.39$ tokens/sec.
• Infilling (ROCStories): La difusión de conjuntos supera sustancialmente a la Difusión por Bloques en tareas de infilling (llenar 1 o 3 oraciones intermedias), logrando puntuaciones ROUGE más altas con hasta un ~25% de decodificación más rápida. La difusión por bloques tiene dificultades aquí debido a su estructura rígida de bloques de izquierda a derecha.
• Resumen (CNN/DailyMail): La difusión de conjuntos logra puntuaciones ROUGE competitivas en comparación con las líneas base de difusión y hasta un 10% más rápido de decodificación que la Difusión por Bloques. También supera a una línea base AR, que los autores señalan que tiende al sobreajuste en este conjunto de datos.
• Generación Incondicional y Verosimilitud (Likelihood): En OpenWebText y LM1B, la difusión de conjuntos mejora los puntajes MAUVE sobre la Difusión por Bloques mientras decodifica más rápido. En LM1B, logra la mejor perplejidad entre los modelos de difusión (27.83 PPL para $S \le 8$) manteniendo un mayor rendimiento (throughput) que los modelos de difusión previos.

Significado y Reivindicaciones

El artículo afirma que la Difusión de Conjuntos representa un avance significativo al desacoplar la interpolación entre la autorregresión y la difusión de las estructuras de bloques rígidas. Al interpolar mediante el ordenamiento de tokens, el marco ofrece:

• Mayor Flexibilidad: Soporte para inserciones en posiciones arbitrarias y generación de longitud variable, que son críticas para tareas como el infilling.
• Eficiencia: La capacidad de actualizar las cachés KV después de cada paso, una característica previamente no disponible en los modelos de difusión, lo que conduce a una inferencia más rápida.
• Rendimiento: Un compromiso superior entre la calidad de la muestra y la velocidad de inferencia en comparación con los modelos de lenguaje de difusión existentes.

Los autores posicionan a la Difusión de Conjuntos como un marco unificador que recupera la AR, la difusión agnóstica al orden y la difusión por bloques como casos especiales, proporcionando un mecanismo ajustable para equilibrar el paralelismo y la calidad de la muestra según los requisitos específicos de la tarea de inferencia.