Diffusion Language Model Parallel Decoding via Product-of-Experts Bridge

Resumen Técnico: Decodificación Paralela de Modelos de Lenguaje de Difusión mediante un Puente de Producto de Expertos

Declaración del Problema
Los Modelos de Lenguaje de Difusión (DLM) ofrecen una alternativa prometedora a los modelos autorregresivos (AR) al permitir la decodificación paralela, lo que teóricamente permite generar secuencias completas en una fracción del tiempo requerido por los modelos AR secuenciales. Sin embargo, los DLM actualmente sufren una brecha de calidad significativa en comparación con los modelos AR robustos. Esta degradación surge de los supuestos de independencia condicional requeridos para la decodificación paralela, donde los tokens generados dentro de un mismo paso se modelan de forma independiente en lugar de conjunta.

Los intentos recientes de cerrar esta brecha utilizando métodos de Monte Carlo —específicamente el muestreo de rechazo (RS) y el muestreo de importancia (IS)— han enfrentado limitaciones fundamentales. En estas configuraciones, un DLM actúa como la distribución de propuesta ($p_D$) y un modelo AR más fuerte actúa como el objetivo ($p_{AR}$). Debido al gran desajuste distributivo entre la propuesta paralela del DLM y el objetivo secuencial del AR:

1. El Muestreo de Rechazo sufre de rechazos frecuentes, causando que el proceso de decodificación colapse hacia una generación casi secuencial con un bajo rendimiento (throughput).
2. El Muestreo de Importancia requiere un número imprácticamente grande de partículas candidatas para encontrar una muestra de alta calidad; con presupuestos limitados, a menudo selecciona la continuación del DLM "menos implausible" en lugar de una muestra verdaderamente fiel al AR.

Metodología: PoE-Bridge
El artículo introduce PoE-Bridge, un marco de decodificación que mitiga el desajuste entre la propuesta y el objetivo mediante la introducción de una distribución intermedia "puente" construida como un Producto de Expertos (PoE) del DLM propuesto y el objetivo AR.

1. La Distribución Puente PoE
En lugar de corregir directamente los borradores del DLM hacia el objetivo AR, PoE-Bridge descompone la corrección en dos etapas a través de una distribución intermedia $p_{PoE}$. El puente se define como una interpolación a nivel de token:
$$p_{PoE}(x_i | x_{<i}) \propto \tilde{p}_D(x_i | x_{<i})^w \cdot p_{AR}(x_i | x_{<i})^{1-w}$$
donde $w \in [0, 1]$ controla la posición del puente. Cuando $w$ está cerca de 1, el puente se asemeja a la propuesta del DLM; cuando está cerca de 0, se aproxima al objetivo AR. Esta definición a nivel de token asegura la tratabilidad tanto para el muestreo como para la evaluación de la verosimilitud (likelihood).

2. Proceso de Decodificación de Dos Etapas
El marco opera en dos etapas paralelizables:

• Etapa 1: Muestreo de Rechazo Especulativo (DLM $\to$ PoE). El DLM redacta múltiples continuaciones de tokens en paralelo. Estos borradores se verifican contra la distribución PoE en lugar de contra el objetivo AR. Debido a que el PoE está más cerca de la propuesta del DLM que el objetivo AR, la tasa de aceptación es significativamente mayor. Los tokens aceptados se muestrean de la distribución puente $p_{PoE}$.
• Etapa 2: Muestreo de Importancia (PoE $\to$ AR). Los candidatos aceptados de la Etapa 1 son luego reponderados mediante muestreo de importancia para corregirlos hacia el objetivo AR final $p_{AR}$. Dado que los candidatos ya están más cerca de la distribución AR (habiendo pasado por el puente PoE), los pesos de importancia son más estables y se necesitan menos candidatos para lograr muestras de alta fidelidad en comparación con el IS directo desde el DLM.

3. Mejoras Prácticas
Para hacer que el marco sea eficiente en entornos de inferencia realistas, los autores proponen dos técnicas clave:

• Muestreo de Temperatura Mixta: Para evitar el colapso de modo cuando se utiliza un número limitado de candidatos ($K$), el marco muestrea candidatos de una familia de distribuciones PoE temperadas con variaciones de temperatura. Esto estratifica la entropía de la propuesta, asegurando una cobertura diversa de la distribución objetivo y mejorando el tamaño de muestra efectivo.
• Ventana de Rechazo Elástica: En lugar de verificar un sufijo fijo y potencialmente largo, el marco ajusta adaptativamente el alcance de la verificación basándose en las longitudes de aceptación observadas. Esto reduce el cómputo desperdiciado en tokens que probablemente serán rechazados, lo cual es particularmente beneficioso al procesar múltiples candidatos en paralelo.

Contribuciones Clave

• Nuevo Marco de Decodificación: PoE-Bridge es el primer método que combina con éxito la decodificación paralela de DLM con la fidelidad de nivel AR mediante la introducción de un puente PoE intermedio, dividiendo efectivamente una tarea de corrección difícil en dos pasos manejables.
• PoE Tratable a Nivel de Token: Los autores definen una interpolación PoE a nivel de token que soporta la verificación especulativa eficiente y la ponderación de importancia, evitando la intratabilidad del muestreo PoE a nivel de secuencia.
• Mecanismos de Eficiencia: La introducción del muestreo de temperatura mixta y las ventanas de rechazo elásticas aborda los cuellos de botella prácticos en diversidad y sobrecarga de verificación.

Resultados Experimentales
Evaluado en desafíos de razonamiento matemático (GSM8K, MATH) y codificación (HumanEval, MBPP), PoE-Bridge demuestra:

• Precisión: Recupera al menos el 95% del rendimiento del modelo AR objetivo, superando significativamente la decodificación estándar de DLM y el muestreo especulativo ingenuo.
• Velocidad: Logra una aceleración de hasta 5$\times$ sobre los enfoques de decodificación DLM estándar y 2$\times$ el rendimiento (throughput) de la decodificación AR.
• Escalabilidad: El uso de muestreo de temperatura mixta permite que la precisión escale consistentemente con el número de candidatos paralelos ($K$), mientras que el muestio de temperatura uniforme ofrece rendimientos decrecientes.
• Sensibilidad a Hiperparámetros: Se encontró que un peso PoE intermedio ($w \approx 0.3$) proporciona el equilibrio óptimo entre la tasa de aceptación (rendimiento) y la fidelidad al objetivo AR.

Significado y Reivindicaciones
El artículo afirma que PoE-Bridge logra un rendimiento de "lo mejor de ambos mundos", rompiendo el convencional compromiso entre calidad y eficiencia en los modelos de lenguaje de difusión. Al aprovechar una distribución puente intermedia, el método permite que los DLM produzcan generaciones tan fieles como la decodificación AR, preservando al mismo tiempo las ventajas de la decodificación paralela que reducen la latencia. Los autores aseguran que este enfoque hace que la corrección de Monte Carlo sea efectiva bajo presupuestos de cómputo realistas y limitados, ofreciendo una solución práctica para la generación de texto de alta calidad sin sacrificar los beneficios de velocidad de los modelos de difusión. El trabajo se presenta como un método de inferencia que no requiere el reentrenamiento de modelos ni la modificación de parámetros, apoyándose en cambio en la sinergia entre una propuesta de DLM rápida y un verificador AR robusto.