Resumen Técnico: Muestreo de Posterior Inspirado en Langevin Discreto (∆LPS)

Enunciado del Problema

El artículo aborda el desafío de resolver problemas inversos en espacios de estados discretos. Los problemas inversos implican recuperar una señal desconocida $x$ a partir de mediciones indirectas, escasas y ruidosas $y$ (modeladas como $y = A(x) + n$). Si bien los modelos de difusión continuos se han convertido en el estándar para resolver tales problemas mediante muestreo de posterior, sus contrapartes discretas permanecen poco exploradas.

Los métodos existentes para el muestreo de posterior discreto adolecen de tres limitaciones principales:

1. Relajaciones Continuas: Métodos como G2D2 relajan las variables categóricas en espacios continuos para utilizar guía por gradiente, alterando la geometría de la inferencia y arriesgando que los tokens discretos proyectados no se decodifiquen de manera consistente.
2. Actualizaciones Secuenciales: Métodos como SGDD dependen de actualizaciones al estilo de Gibbs, que son secuenciales y lentas para mezclarse en espacios de alta dimensión.
3. Especialización: Métodos como APS están estrechamente acoplados a procesos de corrupción específicos (por ejemplo, difusión enmascarada) y paradigmas de entrenamiento, lo que limita su generalidad.

El objetivo es desarrollar un muestreador de posterior que permanezca totalmente discreto, actualice todas las dimensiones de los tokens en paralelo y sea agnóstico al paradigma de entrenamiento específico del prior de difusión discreto.

Metodología: ∆LPS

Los autores proponen ∆LPS (Muestreador de Posterior Inspirado en Langevin Discreto), un algoritmo libre de entrenamiento que utiliza información de gradiente para guiar transiciones discretas sin abandonar el espacio de estados discreto.

Formulación Central

El problema se formula en el espacio latente discreto $z_0 \in \{1, \dots, K\}^L$ de un modelo codificador-decodificador. El objetivo es muestrear de la posterior $p(z_0|y) \propto p(y|z_0)p_\theta(z_0)$, donde $p_\theta$ es un prior de difusión discreto preentrenado. El muestreador aproxima la transición posterior $p(z_s|z_t, y)$ inferiendo primero un latente limpio $z_0$ consistente con la medición y el prior, y luego añadiéndole ruido hasta el siguiente paso de tiempo.

Propuesta de Langevin Discreto

La innovación clave es adaptar la dinámica de Langevin a espacios discretos. En configuraciones continuas, una actualización de Langevin mueve una muestra $z_0$ en la dirección del gradiente del potencial log-posterior $U(z_0) = \log p(y|z_0) + \log p_\theta(z_0; z_t)$. Dado que $z_0$ es discreto, los autores no añaden un vector de gradiente directamente. En su lugar, utilizan la distribución de propuesta gaussiana de la dinámica de Langevin continua para definir probabilidades de transición sobre candidatos discretos válidos $z'_0$:

$$q(z'_0|z_0; z_t, y) \propto \exp\left( -\frac{1}{4\eta} \|z'_0 - z_0 - \eta \nabla_{z_0} U(z_0; z_t, y)\|_2^2 \right)$$

Al expandir el exponente y normalizar sobre el conjunto finito de tokens válidos, la propuesta se factoriza a través de las coordenadas de los tokens. Esto permite al muestreador calcular una distribución categórica para cada posición de token $\ell$ de forma independiente:
$$q_\ell(z'_0[\ell]) = \text{Cat}\left( \sigma\left( \frac{1}{2} \nabla_{z_0[\ell]} U \cdot (k - z_0[\ell]) - \frac{(k - z_0[\ell])^2}{4\eta} \right) \right)$$

Este mecanismo permite actualizaciones totalmente paralelas a través de las $L$ dimensiones de los tokens, evitando el cuello de botella secuencial de los muestreadores de Gibbs.

Generalización y Precondicionamiento

• Agnosticismo del Prior: El método solo requiere probabilidades predictivas sobre estados categóricos, lo que lo hace compatible con priores de difusión enmascarada (MDLM) y difusión de estado uniforme (Duo).
• Problemas Inversos Ciegos: Para operadores directos no diferenciables o desconocidos, el gradiente de verosimilitud se reemplaza por un sustituto contrastivo aprendido (utilizando codificadores para datos y mediciones) para proporcionar una señal de guía suave.
• Precondicionamiento Adam: Para manejar gradientes mal escalados en espacios latentes de alta dimensión, el gradiente crudo se reemplaza por una dirección precondicionada con Adam. Esto reescala adaptativamente las actualizaciones, amortiguando gradientes grandes ruidosos mientras preserva los consistentemente útiles, mejorando la mezcla y la convergencia.

Contribuciones Clave

1. Muestreo de Posterior Totalmente Discreto: ∆LPS es un muestreador libre de entrenamiento que mantiene cada iteración en el espacio de estados discreto, evitando la optimización de tokens suaves y errores de proyección.
2. Actualizaciones Paralelas Eficientes: Al factorizar la propuesta a través de las coordenadas, el método actualiza todas las dimensiones de los tokens en paralelo, mejorando significativamente la eficiencia sobre los muestreadores secuenciales al estilo de Gibbs.
3. Diseño Agnóstico al Prior: El muestreador no está atado a un proceso de corrupción específico (enmascarado vs. uniforme) y funciona con cualquier prior de difusión discreto preentrenado que proporcione probabilidades predictivas.

Resultados Experimentales

Los autores evalúan ∆LPS en un conjunto diverso de problemas inversos en MNIST, CIFAR-10 y FFHQ, cubriendo configuraciones lineales, no lineales y ciegas.

• Restauración de Imágenes (FFHQ): En tareas que incluyen inpainting, desenfoque gaussiano, superresolución, HDR y desenfoque por movimiento, ∆LPS logra un rendimiento comparable o superior a las líneas base continuas (por ejemplo, DPS, PSLD) y supera a las líneas base discretas existentes (SGDD, G2D2, APS) en calidad perceptual (LPIPS).
• Tareas Discretas (MNIST/CIFAR): El método demuestra robustez a través de niveles de dificultad variables (Fácil, Medio, Difícil). Consistentemente iguala o supera a SGDD y G2D2 en precisión de tokens y PSNR, incluso con presupuestos computacionales significativamente menores (menos evaluaciones de funciones) que la variante "completa" de SGDD.
• Problemas Inversos Ciegos: En una tarea de mapeo espacial (inferir planos de planta a partir de trayectorias de caminata donde el operador directo es desconocido), ∆LPS supera a las líneas base continuas (CoGuide, variantes de DPS) en puntuación F1 e Intersección sobre Unión (IoU).

Significado y Afirmaciones

El artículo afirma que los muestreadores de posterior totalmente discretos e informados por gradientes ofrecen un camino escalable y general para resolver problemas inversos sobre representaciones discretas.

• Escalabilidad: El mecanismo de actualización paralela aborda las limitaciones de velocidad de mezcla de los muestreadores discretos en espacios de alta dimensión.
• Generalidad: El método desacopla el algoritmo de inferencia del paradigma de entrenamiento específico del prior, permitiéndole aprovechar el ecosistema en crecimiento de modelos de difusión discreto (texto, código, imágenes, audio).
• Practicidad: Al evitar relajaciones continuas, ∆LPS asegura que las muestras finales sean tokens discretos válidos, previniendo la "brecha de proyección" donde una aproximación continua falla al decodificarse en una señal discreta consistente.

Los autores concluyen que, a medida que los modelos generativos discretos escalan, los métodos de inferencia principistas como ∆LPS se volverán cada vez más críticos para aplicaciones en diseño biomolecular, música simbólica, finalización de código y otros dominios donde los datos son inherentemente discretos.