DiLaDiff: Distilled Latent-Augmented Diffusion for Language Modeling

Resumen Técnico: DiLaDiff: Difusión Latente Aumentada Destilada para Modelado de Lenguaje

Enunciado del Problema
Los modelos de lenguaje de difusión (DLM) enfrentan una compensación fundamental entre la calidad de muestreo y el rendimiento. Aunque los modelos de difusión discreta (específicamente los Modelos de Lenguaje de Difusión Enmascarada, o MDLM) son efectivos para datos categóricos, intrínsecamente fallan en capturar las correlaciones entre los tokens decodificados. Esta limitación surge porque el denoiser aproxima la posterior como un producto de marginales independientes por token. En consecuencia, generar múltiples tokens en paralelo conduce a oraciones incoherentes, forzando una dependencia de la generación secuencial o una reducción en el paralelismo que limita severamente el rendimiento. Por el contrario, los modelos de difusión continua aplicados a representaciones de texto (por ejemplo, incrustaciones) ofrecen técnicas de muestreo y destilación fundamentadas, pero a menudo luchan con la compensación expresividad-decodabilidad, donde las representaciones contextuales son difíciles de decodificar de nuevo a tokens discretos, o no logran igualar el rendimiento de sus contrapartes discretas.

Metodología
Los autores proponen DiLaDiff (Modelo de Lenguaje de Difusión Latente Aumentada Destilada), un marco híbrido continuo-discreto diseñado para resolver el problema de la correlación de tokens manteniendo un alto rendimiento. La arquitectura consta de tres componentes principales:

1. Autoencoder Latente Aumentado (LaDiff):

   • Codificador: Un codificador basado en transformadores ($E_\phi$) aprende un espacio latente continuo ($z$) a partir de representaciones contextuales (derivadas de los estados ocultos de BERT). Este espacio latente está diseñado para capturar correlaciones globales de tokens e información semántica.
   • Decodificador: Un decodificador discreto ($x_\theta$) se inicializa a partir de un MDLM preentrenado. Genera tokens de texto condicionados a la variable latente $z$.
   • Entrenamiento: El sistema se entrena utilizando una Cota Inferior de Evidencia (ELBO) modificada que combina el autoencoding con la difusión enmascarada. Crucialmente, la distribución posterior se factoriza sobre los tokens condicionados al latente $z$. Dado que $z$ captura las correlaciones, la suposición de independencia condicional se mantiene, permitiendo una generación paralela precisa de tokens.
   • Regularización: Para evitar que el espacio latente se vuelva demasiado disperso para el modelado generativo, los autores emplean heurísticas que fomentan la robustez, incluyendo la normalización por coordenada, el enmascaramiento aleatorio de características y latentes, y la inyección de ruido gaussiano.

2. Aprendizaje del Prior Latente:

   • Dado que la secuencia limpia no está disponible durante la inferencia, se entrena un modelo de difusión continua (Denoiser Latente $z_\psi$) para aprender la distribución a priori $p(z)$ sobre la salida del codificador. Esto permite que el modelo muestree un latente $z$ de alta calidad mediante una EDO de flujo de probabilidad inversa antes de la decodificación.

3. Destilación de Consistencia (DiLaDiff):

   • Para reducir aún más la sobrecarga computacional del paso de difusión continua, los autores aplican la destilación MeanFlow. Se entrena un modelo estudiante ($u_\eta$) para predecir la velocidad promedio de la trayectoria de la EDO entre dos pasos de tiempo ($t$ y $r$), en lugar de la velocidad instantánea.
   • Esta destilación permite generar el latente $z$ en pocos pasos (por ejemplo, 5 pasos) manteniendo un rendimiento cercano al del modelo profesor (LaDiff), que requiere muchos pasos (por ejemplo, 200 pasos).

Contribuciones Clave

• Arquitectura Híbrida: El artículo introduce una receta para entrenar un autoencoder de texto donde el decodificador se inicializa a partir de una línea base de difusión discreta preentrenada. Esto resulta en un espacio latente controlable y regularizado que captura información semántica y correlaciones de tokens.
• Paralelismo Guiado por Latentes: Al aprender un prior latente con difusión continua, el modelo (LaDiff) permite el muestreo de múltiples tokens en paralelo sin sacrificar la calidad del texto, ya que el latente resuelve las dependencias de tokens que los modelos discretos pasan por alto.
• Primera Destilación Exitosa de EDO para Latentes de Texto: Los autores presentan DiLaDiff, el primer intento exitoso de destilar trayectorias de EDO para latentes de texto contextuales. Esto reduce el costo de generación de latentes a niveles insignificantes en comparación con la decodificación discreta.

Resultados
Los experimentos se realizaron en el conjunto de datos OpenWebText con una longitud de secuencia de 1024 tokens.

• Rendimiento vs. Línea Base: Incluso sin destilación, LaDiff supera la línea base de difusión enmascarada (MDLM). En la generación incondicional, LaDiff logra una mejora absoluta de 28.5 en la Perplejidad Generativa (GenPPL) y una mejora absoluta de 0.08 en MAUVE en comparación con MDLM con 1024 pasos discretos.
• Rendimiento: LaDiff acelera la inferencia por un factor de 7× con un tamaño de lote de 32 en comparación con la configuración óptima de MDLM.
• Eficiencia de Destilación: DiLaDiff logra un rendimiento cercano a su profesor LaDiff utilizando solo 5 pasos de difusión latente (en comparación con 200 para el profesor). La sobrecarga computacional del componente de difusión latente cae a aproximadamente el 5% del tiempo total de pared, haciendo que la generación de latentes sea insignificante en comparación con el costo de decodificación discreta.
• Generación en Pocos Pasos: DiLaDiff rivaliza con los métodos de vanguardia de pocos pasos (como MDLM+SDTT y DUO+DCD) sin destilar explícitamente el decodificador discreto, demostrando la eficacia de la aumentación latente.
• Diversidad Semántica: El análisis muestra que el espacio latente preserva la información semántica (verificado mediante BERTScore en latentes ruidosos) y permite una generación diversa, como lo demuestra la capacidad del decodificador para mantener la entropía incluso a temperaturas más bajas.

Significado y Afirmaciones
El artículo afirma que DiLaDiff cierra con éxito la brecha entre la expresividad de las representaciones continuas y la robustez de la decodificación discreta. Al capturar correlaciones globales en un espacio latente, el modelo supera la limitación intrínseca de los modelos de difusión discreta respecto a la generación paralela de tokens. Los autores enfatizan que su enfoque permite la generación de texto de alta calidad con una latencia de inferencia significativamente reducida.

El artículo se mantiene modesto respecto a sus limitaciones:

• DiLaDiff aún no iguala completamente el rendimiento de su modelo profesor (LaDiff) en el régimen de pocos pasos, lo que sugiere que se necesita más trabajo en la destilación.
• El decodificador discreto en sí mismo no se destila; los autores señalan que combinar su destilación latente con la destilación del decodificador discreto es una extensión natural futura para lograr un rendimiento de vanguardia real en pocos pasos.
• Las estrategias de regularización y los diseños de programación de difusión son actualmente heurísticos, aunque los autores presentan un intento más fundamentado en los apéndices.

En resumen, DiLaDiff demuestra que un enfoque híbrido continuo-discreto, aumentado por la destilación latente, puede acelerar significativamente la inferencia de modelado de lenguaje manteniendo o mejorando la calidad de generación, resolviendo efectivamente la compensación rendimiento-calidad inherente a los actuales modelos de lenguaje de difusión.