Resumen Técnico: Muestreo de Posterior Basado en Difusión Probable para Problemas Inversos Lineales mediante DDIM

Planteamiento del Problema

El artículo aborda el desafío de resolver problemas inversos lineales ruidosos de la forma $Y = AX_0 + \varepsilon$, donde $X_0 \in \mathbb{R}^d$ es una señal desconocida, $A \in \mathbb{R}^{n \times d}$ es un operador de medición lineal conocido, y $\varepsilon \sim \mathcal{N}(0, \sigma^2 I_n)$ es ruido gaussiano. El objetivo es generar muestras de la distribución posterior bayesiana $p_{X_0|Y}$, condicionada a la observación $y$.

Si bien los modelos de difusión han emergido como potentes priors basados en datos para capturar distribuciones de señales complejas, los métodos existentes para el muestreo de la posterior suelen sufrir de una de estas dos limitaciones:

1. Falta de garantías teóricas: Muchos enfoques heurísticos (por ejemplo, métodos basados en proyecciones o correcciones de puntuación aproximadas) carecen de una prueba rigurosa de que sus muestras generadas converjan a la verdadera posterior.
2. Sobrecarga computacional: Los métodos que intentan aproximar la puntuación (score) de la posterior o que utilizan Monte Carlo Secuencial (SMC) a menudo incurren en costos computacionales prohibitivos, como requerir el cálculo de gradientes de funciones de puntuación o un gran número de partículas.

Metodología: Posterior-DDIM

Los autores proponen Poster-DDIM, un algoritmo simple y eficiente que realiza el muestreo de la posterior adaptando el muestreador estándar de Modelos de Difusión Implícitos de Denoising (DDIM). La innovación central reside en una partición basada en la descomposición en valores singulares (SVD) del espacio de la señal, donde las actualizaciones se adaptan a la fuerza relativa de la relación señal-ruido (SNR) de la observación frente a la SNR de la difusión.

1. Partición Guiada por la SNR

Utilizando la SVD de la matriz de medición $A = U\Sigma V^\top$, el problema se desacopla en problemas de inferencia escalar a lo largo de cada dirección singular $v_s$. Los autores definen un umbral de tiempo $\tau_s$ para cada dirección donde la SNR de la difusión ($\alpha_t/\sigma_t$) es igual a la SNR de la observación ($\Sigma_{ss}/\sigma$). Basándose en esto, las direcciones singulares se particionan en cualquier tiempo $t$ en tres grupos:

• Direcciones dominadas por la medición ($S^{meas}_t$): Direcciones donde la SNR de la observación supera la SNR de la difusión ($t \ge \tau_s$).
• Direcciones observadas dominadas por el prior ($S^{prior}_t$): Direcciones donde la SNR de la observación es menor que la SNR de la difusión ($t < \tau_s$), pero la dirección aún es observada.
• Direcciones no observadas ($S^c$): Direcciones en el espacio nulo de $A$ donde no existe información de medición.

2. Reglas de Actualización por Coordenadas

El algoritmo aplica distintas reglas de actualización a estos grupos, requiriendo solo modificaciones ligeras a la actualización estándar de DDIM:

• Para $S^{meas}_t$ (Dominadas por la medición): El algoritmo omite el prior aprendido e inyecta directamente una versión ruidosa calibrada de la observación. Específicamente, genera la coordenada $v_s^\top \hat{X}_{t_{i+1}}$ utilizando el proceso de propagación auxiliar definido por la medición $y$, asegurando que la distribución marginal coincida con el proceso de difusión hacia adelante proyectado sobre esa dirección.
• Para $S^{prior}_t$ (Observadas dominadas por el prior): El algoritmo sigue la dinámica estándar de DDIM impulsada por el predictor de datos aprendido $\mu_t(\cdot)$. En el caso general donde los valores singulares no son idénticos, se añade un término de corrección al predictor para dar cuenta de los distintos valores singulares, asegurando que la actualización sea consistente con la posterior condicional.
• Para $S^c$ (No observadas): Estas coordenadas se muestrean enteramente a partir del prior de difusión aprendido condicionado a las coordenadas observadas. Para fomentar una exploración suficiente del subespacio no observado, se utiliza un parámetro de estocasticidad $\eta_1$ más grande, haciendo que la actualización sea más similar a Langevin.

3. Parámetros de Estocasticidad

El método utiliza dos parámetros de estocasticidad, $\eta_0$ y $\eta_1$, que controlan la inyección de ruido en las actualizaciones de DDIM para $S^{prior}_t$ y $S^c$, respectivamente.

• $\eta_0$ se mantiene pequeño (a menudo 0) para mantener una actualización de tipo determinista para las direcciones observadas donde el prior es la guía principal.
• $\eta_1$ se establece para ser suficientemente grande para asegurar que el subespacio no observado se explore eficazmente, ya que su distribución está determinada únicamente por el prior.

Principales Contribuciones

1. Simplicidad Algorítmica: El método propuesto reduce el muestreo de la posterior a una secuencia de actualizaciones de DDIM simples y por coordenadas. Incorpora explícitamente el modelo de medición sin requerir pasos de proyección complejos, cálculos de gradiente de la verosimilitud o el reentrenamiento del modelo de difusión.
2. Garantías Teóricas: El artículo establece la corrección de la posterior asintótica. Bajo el supuesto de un prior de difusión exacto y una discretización temporal suficientemente fina, los autores demuestran que la distribución de las muestras generadas por Posterior-DDIM converge a la verdadera posterior bayesiana $p_{X_0|Y}$ en distancia de variación total. Esto se cumple tanto para el caso especial de valores singulares idénticos como para el caso general con valores singulares arbitrarios.
3. Desempeño Empírico: Experimentos extensos en tareas de restauración de imágenes (inpainting, super-resolución, desenfoque gaussiano y compresión de sensores) demuestran que Posterior-DDIM supera a los muestreadores de difusión existentes (incluyendo DPS, DDRM, DDNM+ y ProjDiff) en la mayoría de las métricas de evaluación (PSRE, SSIM, LPIPS y FID).

Resultados Experimentales

Los autores evaluaron el algoritmo en los conjuntos de datos CelebA e ImageNet. Hallazgos clave incluyen:

• Superioridad: Posterior-DDIM logró el mejor desempeño en al menos tres de las cuatro métricas en la mayoría de las tareas y en al menos dos métricas en cada tarea considerada.
• Robustez: El método superó consistentemente a los modelos de referencia a través de diferentes niveles de ruido y tipos de degradación.
• Eficiencia: El algoritmo requiere solo 100 evaluaciones de función (NFE) para alcanzar un desempeño estable, comparable o incluso más rápido que otros métodos que a menudo requieren significativamente más pasos (por ejemplo, 1000 NFE para DPS).
• Sensibilidad a Hiperparámetros: Los experimentos mostraron que establecer $\eta_0 = 0$ (actualización determinista para las direcciones dominadas por el prior) y un $\eta_1$ suficientemente grande (por ejemplo, 16 o 32) produce resultados óptimos.

Significado y Reivindicaciones

El artículo afirma proporcionar un marco basado en principios y demostrable para integrar los priors de difusión con modelos de medición lineales. Al convertir el complejo problema del muestreo de la posterior en actualizaciones de DDIM simples y por dirección, los autores cierran la breancia entre el rigor teórico y la eficiencia práctica.

La importancia de este trabajo se destaca por su capacidad para:

• Ofrecer soporte teórico riguroso para los muestreadores de tipo DDIM en problemas inversos, abordando la falta de garantías de convergencia en muchos métodos heurísticos.
• Proporcionar una solución de fácil implementación que preserva la eficiencia computacional de los muestreadores de difusión estándar mientras maneja explícitamente el modelo de medición.
• Demostrar que la consistencia de la posterior puede lograrse sin las prohibiciones computacionales de Monte Carlo Secuencial o las aproximaciones heurísticas de los métodos basados en proyecciones.

Los autores señalan que, si bien el presente trabajo se centra en problemas inversos lineales, extender este marco a entornos no lineales y caracterizar las tasas de convergencia son direcciones naturales para investigaciones futuras.