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

Enunciado del Problema

La difusión continua se ha convertido en un marco estándar para la generación de alta dimensión en imágenes, video y audio, sin embargo, ha quedado consistentemente rezagada frente a los modelos de difusión discreta enmascarada (MDM) en la generación de lenguaje. Los modelos de lenguaje de difusión continua existentes operan típicamente sobre incrustaciones de tokens ruidosas utilizando ruido gaussiano de Preservación de Varianza (VP) o Explosión de Varianza (VE). Los autores argumentan que el cuello de botella no es la continuidad en sí misma, sino la elección de la representación.

En la difusión continua estándar, el proceso de eliminación de ruido depende de un régimen de ruido indexado por pasos de tiempo. El modelo debe recibir una etiqueta de paso de tiempo o nivel de ruido para calibrar cuánta evidencia contiene la incrustación ruidosa. En consecuencia, los eliminadores de ruido en esta representación son altamente sensibles a la condicionamiento temporal para distinguir entre estados enmascarados (sin información) e inciertos (SNR finita). Esta dependencia del condicionamiento temporal explícito impide que una sola red soporte naturalmente diversos caminos de inferencia, como el muestreo autoregresivo de orden aleatorio o la decodificación híbrida continua-discreta, sin modelos específicos para el muestreador o reentrenamiento.

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 eliminador de ruido óptimo de Bayes sea invariante a la relación señal-ruido (SNR) nominal.

1. Geometría de Tokens en Esfera Unitaria

La innovación central es restringir todas las incrustaciones de tokens limpios para que se encuentren en la esfera unitaria ($\|x_i\|_2 = 1$). Bajo un canal de localización estocástica donde la observación ruidosa es $z_t = t x + \sqrt{t}\epsilon$, la factorización del posterior de Bayes incluye un término proporcional a $t \|x\|_2^2$. Dado que las incrustaciones tienen norma unitaria, $\|x\|_2^2$ es constante para cualquier oración de una longitud dada. Este término se cancela durante la normalización, haciendo que el posterior y, por lo tanto, el eliminador de ruido óptimo, sean independientes de la SNR nominal $t$.

El eliminador de ruido óptimo se convierte en una función del estado ruidoso $z$ únicamente:
$$\hat{x}(z) = \mathbb{E}[x | z] = \frac{\mathbb{E}_P[x e^{x \cdot z}]}{\mathbb{E}_P[e^{x \cdot z}]}$$
Esto permite que una sola red agnóstica al tiempo sirva a toda la familia de caminos de SNR por token.

2. Caminos de SNR por Token

Dado que el eliminador de ruido no requiere un paso de tiempo global, DSL permite que cada posición de token evolucione bajo su propio camino de SNR arbitrario $\gamma_i(t)$. Esto unifica varias estrategias de inferencia bajo un solo modelo:

• Autoregresivo de Orden Aleatorio (ROAR): Revelar tokens uno por uno corresponde a impulsar $\gamma_i$ específicos de 0 a $\infty$ mientras se mantienen los demás en 0.
• Refinamiento Enmascarado: Volver a enmascarar tokens corresponde a restablecer sus $\gamma_i$ a 0.
• Difusión Continua: Aumentar conjuntamente todos los $\gamma_i$ corresponde a la eliminación de ruido continua estándar.

3. Objetivo de Entrenamiento y Arquitectura

Para entrenar un solo modelo para estos caminos diversos, DSL emplea una estrategia de muestreo de SNR de soporte mixto:

• Rama Continua: Muestra SNRs de una distribución log-normal para cubrir estados de SNR finita.
• Rama de Punto Final: Muestra de una distribución estilo ROAR concentrada cerca de los extremos enmascarado ($\gamma=0$) y limpio ($\gamma=\infty$).

Ambas ramas se entrenan utilizando una sola pérdida de entropía cruzada (CE) a nivel de token. Los autores justifican esto mostrando que minimizar la CE en el posterior categórico de tokens es un sustituto para el Error Cuadrático Medio (MSE) exacto de la integral de camino y la Negativa Log-Verosimilitud (NLL), evitando al mismo tiempo patologías asociadas con la regresión directa en el espacio de incrustaciones (como colapso de incrustaciones o memorización).

La arquitectura incluye un convertidor de vista posterior que mapea la incrustación ruidosa $z_i$ a una representación acotada de mezcla de tokens $m_i$. Esto asegura que la entrada al núcleo Transformer permanezca estable e interpretable como una distribución de probabilidad sobre el vocabulario, independientemente del nivel de SNR.

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 eliminador de ruido óptimo de Bayes es una función del estado ruidoso únicamente, eliminando la necesidad de condicionamiento temporal explícito.
2. Una Receta Práctica: El artículo propone una estrategia de muestreo de SNR de soporte mixto que entrena un único eliminador de ruido tanto en estados de punto final (enmascarar/revelar) como en estados de SNR finita, acoplado con un convertidor que une la geometría posterior continua con los núcleos Transformer estándar.
3. Evidencia a través de Caminos: Los autores demuestran que un único punto de control DSL puede soportar refinamiento enmascarado, muestreo autoregresivo de orden aleatorio y muestreo híbrido continuo-luego-discreto sin reentrenamiento.

Resultados Experimentales

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

• Generación OpenWebText: DSL mejora sustancialmente la fidelidad distribucional (medida por MAUVE) en todos los presupuestos de pasos ($T=128$ a $T=1024$) en comparación con el MDLM de referencia.
  • Con un decodificador de refinamiento enmascarado (ReMDM), DSL logra un MAUVE de 0.722 en $T=1024$, superando a métodos anteriores del estado del arte como PRISM-loop (0.527) y ADLM (0.699).
  • Incluso sin refinamiento iterativo (usando un ROAR ingenuo o un decodificador MDLM vanilla), el punto de control ajustado fino con DSL supera significativamente a la referencia, lo que indica que la mejora proviene del posterior aprendido y no solo del programa de muestreo.
• Muestreo Híbrido: Un muestreador híbrido que realiza una etapa corta de eliminación de ruido continua seguida de refinamiento enmascarado discreto logra una generación coherente en tan solo 48 pasos totales (MAUVE 0.662), superando a las referencias sin refinamiento en conteos de pasos mucho más altos.
• Verosimilitud Text8: En la referencia Text8, DSL logra un Bits por Carácter (BPC) de $\le 1.45$. Si bien esto es un límite superior (ELBO) y no directamente comparable a la NLL exacta de los modelos autoregresivos, se acerca al rango de los modelos de difusión discreta enmascarada (MDLM $\le 1.40$) y mejora sobre otras referencias de difusión continua (Plaid $\le 1.48$).

Significado y Afirmaciones

El artículo afirma que la brecha entre la difusión continua y la difusión discreta en el lenguaje no se debe a la naturaleza de la continuidad, sino a la representación del régimen de ruido. Al parametrizar la corrupción de tal manera que el estado ruidoso determine únicamente el posterior del token limpio, DSL unifica el comportamiento de "revelación de punto final" de la difusión enmascarada con las capacidades de "refinamiento continuo" de los modelos de difusión.

La afirmación central es que la revelación de punto final, el refinamiento enmascarado y la eliminación de ruido continua/híbrida pueden ser servidos por un único estimador posterior en lugar de requerir modelos separados específicos para el muestreador. Esto sugiere un nuevo paradigma para la generación no autoregresiva donde el modelo aprende un mapa posterior agnóstico al tiempo, y el camino de inferencia (AR, enmascarado o híbrido) se determina dinámicamente en el momento del muestreo.

Los autores permanecen modestos respecto a las limitaciones, señalando que su evaluación se limita a OWT y Text8, y que los muestreadores reportados (familia ReMDM e híbrido) no son políticas de decodificación exhaustivamente optimizadas, sino más bien pruebas de concepto que demuestran la flexibilidad del posterior DSL. Las direcciones de trabajo futuro incluyen programas de remapeo optimizados, modelos más grandes y generación condicional.