DSL-LLaDA: Scaling Continuous Denoising to 8B Masked Diffusion LMs

Resumen Técnico: DSL-LLaDA: Escalamiento del Denoising Continuo a Modelos de Lenguaje de Difusión Enmascada de 8B

Planteamiento del Problema
Los Modelos de Lenguaje de Difusión Discreta (DLM), como LLaDA, generan texto mediante decodificación paralela iterativa. Aunque estos modelos han escalado a miles de millones de parámetros, su calidad de inferencia se degrada drásticamente bajo presupuestos de pasos bajos (generación de pocos pasos). El problema central es un compromiso entre la longitud de la salida y la calidad: con un número fijo de pasadas hacia adelante, el desenmascarado iterativo estándar produce salidas cortas y de alta calidad, o bien textos largos, repetitivos e incoherentes. Este fallo proviene de la factorización de la transición inversa, donde cada posición de token se decodifica independientemente de su distribución marginal condicionada al contexto, ignorando las correlaciones entre tokens. Cuando los presupuestos de pasos son bajos, cada paso debe comprometerse con muchas posiciones simultáneamente, amplificando estos errores y conduciendo a una terminación prematura o bucles de repetición.

Un denoising continuo en el espacio de embeddings ofrece una solución teórica al evolucionar todas las posiciones de forma conjunta para evitar la factorización en el espacio de tokens, pero construir un modelo de este tipo desde cero a gran escala sigue siendo un problema de diseño abierto debido a las complejidades en la geometría de los embeddings, el esquema de ruido y los objetivos de entrenamiento.

Metodología
Los autores proponen DSL-LLaDA, un método para adaptar ligeramente un DLM preentrenado existente (específicamente LLaDA-8B-Instruct) para soportar el denoising continuo en el espacio de embeddings, en lugar de entrenar un modelo continuo desde cero.

1. Adaptación de Localización Estocástica Discreta (DSL):

   • El método reemplaza la entrada binaria de "limpio o enmascarado" de los DLM estándar por estados continuos ruidosos. Para un token limpio $x_i$, se muestrea una Relación Señal-Ruido (SNR) $\gamma_i$, y se construye un estado ruidoso $z_i$ como una mezcla del embedding del token y ruido Gaussiano.
   • Un conversor softmax mapea estos estados DSL continuos de vuelta al espacio de embedding de entrada del backbone. Este conversor produce una "máscara suave" donde un SNR bajo corresponde a un estado similar a una máscara, un SNR medio a tokens aleatorios poco fiables, y un SNR alto a embeddings de tokens claros.
   • El backbone es sometido a un preentrenamiento continuo durante solo 1,000 pasos en FineWeb-Edu. El objetivo de entrenamiento sigue siendo la entropía cruzada de tokens estándar, pero los estados de entrada son ahora mezclas continuas en lugar de máscaras binarias. Crucialmente, el conversor permite que el modelo procese estas entradas sin cambios arquitectónicos, ya que la invariancia al SNR del denoiser óptimo elimina la necesidad de entradas explícitas de tiempo (timestep).

2. Inferencia Continua (Muestreador SDE):

   • Durante la inferencia, el modelo mantiene un estado continuo para todas las posiciones. Utiliza un muestreador basado en Ecuaciones Diferenciales Estocásticas (SDE) (pasos Heun) para evolucionar el estado desde un SNR bajo (alta incertidumbre) hacia un SNR alto.
   • El modelo predice una distribución de probabilidad sobre los tokens, la cual se convierte en un promedio ponderado de los embeddings de los tokens (usando los $M$ tokens más probables). Esta estimación de embedding actualiza el estado continuo.
   • El compromiso de tokens rígidos (hard token commitment) se pospone hasta el paso final. Esto permite que el modelo refine conjuntamente todas las posiciones, preservando las correlaciones inter-token durante todo el proceso de generación.

3. Generación Condicional:

   • Para tareas como la sumarización, los tokens de contexto se fijan a sus embeddings limpios, mientras que la región de respuesta experimenta un denoising continuo. El backbone atiende bidireccionalmente, condicionando el denoising de la respuesta al contexto limpio.

Contribuciones Clave

1. Conversión Ligera: Se demostró que un DLM enmascarado preentrenado de 8B puede convertirse para soportar el denoising en el espacio continuo mediante una adaptación mínima de 1,000 pasos usando DSL, evitando la necesidad de entrenar un modelo continuo desde cero.
2. Evitación de Modos de Fallo de Bajo NFE: Se mostró que la inferencia SDE continua evita eficazmente la terminación prematura y los bucles de repetición característicos del desenmascarado iterativo con presupuestos de pasos bajos (NFE $\le$ 16).
3. Robustez Selectiva al Estado Ruidoso: Se identificó que la adaptación induce una robustez donde el modelo puede corregir tokens corruptos mientras preserva los tokens limpios (corrección selectiva), un comportamiento ausente en las líneas base discretas de la misma computación.

Resultados

• Sumarización Zero-Shot: En cuatro benchmarks (XSum, CNN/DailyMail, arXiv, PubMed), DSL-LLaDA-SDE logró los mejores puntajes ROUGE-1 con presupuestos de pasos bajos ($\le$ 16 pasadas hacia adelante) sin heurísticas de control de longitud. En contraste, las líneas base discretas truncaban las salidas prematuramente o sufrían altas tasas de repetición.
• Repetición y Longitud: En la generación de texto abierto, DSL-LLaDA-SDE mantuvo tasas de repetición por debajo del 9% en todos los presupuestos de pasos, mientras que las líneas base discretas superaron el 60% de repetición en bajo NFE o se truncaron a longitudes muy cortas.
• Corrección de Errores: El modelo adaptado demostró "corrección selectiva", recuperando tokens corruptos mientras preservaba más del 98% de los tokens limpios. En comparación, las líneas base discretas con corrupción de tokens aleatorios (XDLM) dañaron una parte significativa de los tokens limpios, y el masked diffusion estándar (MDM-CPT) mostró una capacidad de corrección insignificante.
• Experimentos de Control: Los experimentos de control utilizando el entrenamiento de difusión enmascarada estándar con la misma computación (MDM-CPT) y corrupción de tokens aleatorios (XDLM) fallaron en replicar los beneficios del denoising continuo o la robustez selectiva, confirmando que las mejoras derivan de la exposición al ruido continuo y no solo de la computación adicional.

Significancia y Reivindicaciones
El artículo afirma que DSL-LLaDA proporciona un procedimiento de conversión de propósito general para los DLM enmascarados, permitiéndoles aprovechar los beneficios del denoising continuo (evolución conjunta de tokens, compromiso diferido) con un overhead de entrenamiento mínimo. Los resultados sugieren que la adaptación al ruido continuo puede resolver el compromiso entre longitud y calidad inherente al desenmascarado iterativo con presupuestos de pasos bajos.

Los autores mantienen la modestia respecto a las implicaciones más amplias:

• La evaluación se limita principalmente a la generación anclada (sumarización y seguimiento de instrucciones). La escritura creativa abierta todavía muestra tasas más altas de degeneración (24% frente a $\le$4% en sumarización).
• El método está validado actualmente solo en la escala LLaDA-8B; el rendimiento en otras arquitecturas o escalas mayores no ha sido probado.
• Aunque SDE supera a los métodos discretos en bajo NFE, el LLaDA original iguala o supera el rendimiento de SDE en tareas de referencia larga en altos NFE (donde la terminación prematura no es un problema).
• Los autores no reclaman paridad con los modelos autorregresivos (AR) de la misma escala, señalando que los modelos AR siguen liderando en tareas de referencia larga con altos presupuestos de pasos.
• La evaluación humana se identifica como un trabajo futuro.