Resumen Técnico: Localización Estocástica Discreta para Generación No Autoregresiva

1. Declaración del Problema

Los modelos de difusión continua se han convertido en un estándar para la generación de alta dimensión en imágenes, video y audio. Sin embargo, su aplicación a la generación de lenguaje mediante difusión continua sobre incrustaciones de tokens ha rendido consistentemente menos que los Modelos de Difusión Discreta Enmascarada (MDM). Los modelos de lenguaje de difusión continua existentes suelen quedarse atrás respecto a los MDM en fidelidad distribucional (medida por métricas como MAUVE).

Los autores argumentan que el cuello de botella no es la continuidad en sí misma, sino la representación utilizada. Los enfoques estándar (Preservación de Varianza o Explosión de Varianza) representan el límite de información cero como ruido gaussiano isotrópico en lugar de un estado de token tipo máscara. Además, en estos marcos, el desruidizador debe recibir una etiqueta de paso de tiempo o nivel de ruido para calibrar la cantidad de evidencia en la incrustación ruidosa. Esto hace que los desruidizadores sean altamente sensibles a la condición temporal, requiriendo un manejo distinto para estados enmascarados versus inciertos. En consecuencia, los modelos de difusión continua han fallado en aprovechar la supervisión más rica proporcionada por el continuo de estados finitos de Relación Señal-Ruido (SNR).

2. Metodología: Localización Estocástica Discreta (DSL)

El artículo introduce la Localización Estocástica Discreta (DSL), un marco de estado continuo diseñado para hacer que el proceso de desruido sea invariante al SNR nominal.

Decisiones de Diseño Central

1. Incrustaciones de Token en Esfera Unitaria: DSL restringe todas las incrustaciones de tokens limpios para que se encuentren en la esfera unitaria ($\|x_i\|_2 = 1$). Esta es una restricción estructural sobre la geometría, no un regularizador de pérdida.
2. Canal de Localización Estocástica: El proceso de corrupción se define como $z_t = t x + \sqrt{t} \epsilon$, donde $t$ representa el SNR. Bajo la restricción de esfera unitaria, el término que involucra la norma de $x$ en el posterior de Bayes se vuelve constante y se cancela durante la normalización.
3. Desruidizador Invariante al SNR: Un resultado teórico clave es que, bajo esta geometría, el desruidizador óptimo de Bayes (la media posterior) depende únicamente del estado de token ruidoso $z$, no del SNR nominal $t$.
   $$\hat{x}(z, t) = \hat{x}(z) = \frac{\mathbb{E}_P[x e^{x \cdot z}]}{\mathbb{E}_P[e^{x \cdot z}]}$$
   Esto permite que una sola red, ajena al tiempo, sirva como desruidizador para todos los regímenes de SNR (desde completamente enmascarado hasta casi limpio).

Entrenamiento y Arquitectura

• Entrenamiento de Soporte Mixto: Para aprovechar tanto la verosimilitud de integral de trayectoria continua como la verosimilitud del punto final de la difusión enmascarada, DSL emplea un objetivo de entrenamiento de soporte mixto. Muestrea SNRs de dos ramas:
  • Trayectoria Continua: SNRs extraídos de una distribución log-normal (estados de SNR finito).
  • Punto Final ROAR (Autoregresivo de Orden Aleatorio): SNRs extraídos de valores extremos (cerca de 0 para enmascarado, cerca de $\infty$ para revelado).
    Ambas ramas se entrenan utilizando una única pérdida de entropía cruzada sobre el posterior del token, que actúa como un sustituto del Error Cuadrático Medio (MSE) de la integral de trayectoria, evitando patologías de optimización en el espacio de incrustaciones (por ejemplo, colapso de incrustaciones).
• Convertidor de Vista Posterior: Dado que el estado ruidoso $z$ vive en un espacio de incrustaciones continuo pero el núcleo (Transformer/DiT) espera entradas tipo token, un convertidor mapea $z$ a una representación de "mezcla de tokens". Esto mapea el estado ruidoso a una combinación acotada de incrustaciones de norma unitaria, preservando la geometría posterior (la dirección controla la preferencia de token, la norma controla la confianza) mientras garantiza una autoatención estable.

Flexibilidad de Inferencia

Dado que el desruidizador es ajeno al tiempo, DSL admite trayectorias de SNR por token arbitrarias dentro de una sola ejecución de generación:

• Refinamiento Enmascarado: Remascaramiento iterativo estándar (ReMDM).
• Autoregresivo de Orden Aleatorio (ROAR): Revelar tokens uno por uno en orden aleatorio sin remascaramiento.
• Muestreo Híbrido: Una etapa de desruido continua seguida de una etapa de refinamiento enmascarado discreta.

3. Contribuciones Clave

1. Un Marco de Estado Continuo para Generación Discreta: DSL introduce un canal de localización estocástica sobre incrustaciones en esfera unitaria donde el desruidizador óptimo de Bayes es una función del estado ruidoso únicamente. Esto permite que una sola red ajena al tiempo soporte toda una familia de trayectorias de SNR por token.
2. Una Receta de Entrenamiento Práctica: Los autores proponen una estrategia de muestreo de SNR de soporte mixto que entrena al mismo desruidizador tanto en estados de punto final (enmascarado/revelado) como de SNR finito. Utilizan un convertidor de mezcla de tokens para alinear la vista posterior continua con arquitecturas Transformer estándar.
3. Evidencia a través de Trayectorias: Un único punto de control de DSL demuestra una mejora en la generación de pocos pasos en OpenWebText bajo refinamiento enmascarado, soporta muestreo AR de orden aleatorio sin reentrenamiento y logra una verosimilitud competitiva en Text8.

4. Resultados Experimentales

Los autores evaluaron DSL ajustando fino un punto de control de MDLM preentrenado en OpenWebText (OWT) y Text8.

• Generación en OpenWebText (MAUVE):
  • Sin Refinamiento: Incluso sin remascaramiento iterativo, el punto de control ajustado con DSL superó significativamente al MDLM de referencia. En $T=1024$, MAUVE mejoró de 0.042 (MDLM) a 0.495 (DSL-FT + MDLM) y 0.551 (DSL-FT + ROAR).
  • Refinamiento Iterativo (ReMDM): Cuando se combina con decodificadores de la familia ReMDM, DSL logró resultados de vanguardia entre las referencias sin referencia. La mejor configuración (bucle ReMDM con $\eta_{cap}$ principista) alcanzó un MAUVE de 0.722 en $T=1024$, superando a métodos anteriores como PRISM-loop (0.527) y ADLM (0.699).
  • Eficiencia de Pocos Pasos: DSL mostró un rendimiento sólido con presupuestos bajos de pasos. En $T=128$, DSL-FT + bucle ReMDM alcanzó 0.639 MAUVE, en comparación con 0.015 para el MDLM de referencia.
• Muestreo Híbrido: Un muestreador híbrido (desruido continuo seguido de refinamiento discreto) alcanzó un MAUVE de 0.662 en solo 48 pasos totales, superando el resultado de MDLM sin refinamiento en $T=1024$.
• Verosimilitud en Text8: DSL logró Bits por Carácter (BPC) de $\le$ 1.45, acercándose al rango de modelos de difusión discreta enmascarada (MDLM $\le$ 1.40, MD4 $\le$ 1.37) y mejorando sobre las referencias de difusión continua anteriores (Plaid $\le$ 1.48).

5. Significado y Afirmaciones

El artículo afirma que la ventaja principal de DSL es la unificación de las trayectorias de inferencia. Al parametrizar la difusión continua como desruido de estado posterior sobre tokens ruidosos (en lugar de desruido indexado por paso de tiempo), DSL permite:

• Que el revelado de punto final, el refinamiento enmascarado y el desruido continuo/híbrido sean atendidos por un único estimador posterior.
• Compatibilidad de puntos de control: Un único modelo puede ser consultado a lo largo de múltiples trayectorias de SNR (por ejemplo, AR de orden aleatorio, refinamiento enmascarado, híbrido) sin reentrenamiento o destilación.
• Mejora de la Fidelidad Distribucional: Los estados finitos de SNR continuos proporcionan una supervisión más rica, lo que conduce a puntuaciones MAUVE sustancialmente mejores en todos los presupuestos de pasos en comparación con la difusión continua estándar o las referencias de difusión enmascarada.

Los autores concluyen que el ingrediente faltante en la difusión continua para el lenguaje no era la continuidad, sino la elección de la representación. Al adoptar incrustaciones en esfera unitaria y localización estocástica, cierran la brecha entre la difusión continua y discreta, permitiendo que un único modelo soporte diversas estrategias de generación con alta eficiencia y calidad.

Limitaciones: La evaluación se limita actualmente a OpenWebText y Text8. Los muestreadores reportados (familia ReMDM, híbrido) no son políticas de decodificación exhaustivamente optimizadas, sino que sirven para demostrar la flexibilidad del posterior DSL subyacente. Las direcciones de trabajo futuro incluyen la optimización de programas de remascaramiento, reglas de transferencia híbrida y la exploración de generación condicional.