Resumen Técnico: Muestreo MAP Local para Modelos de Difusión

Enunciado del Problema

El Muestreo Posterior de Difusión (DPS) ha surgido como un marco potente para resolver problemas inversos mediante el muestreo de la distribución posterior $p(x_0 | y)$, ofreciendo cuantificación de incertidumbre fundamentada y multimodalidad. Sin embargo, muchas aplicaciones prácticas en imagenología y ciencia priorizan la estimación puntual precisa sobre el muestreo completo de la distribución. El estimador de Máxima A Posteriori (MAP), definido como $\arg \max p(x_0 | y)$, sigue siendo un objetivo de reconstrucción estándar.

Los métodos existentes basados en optimización (por ejemplo, Resample, DiffPIR, DCDP) intentan mejorar el rendimiento de la reconstrucción alternando entre pasos de eliminación de ruido y pasos de optimización. No obstante, la conexión teórica entre estos procedimientos iterativos, la solución MAP global y el muestreo posterior permanece poco clara. No está establecido si estos métodos convergen al MAP global o si son consistentes con DPS. Además, la naturaleza específica de los pasos de optimización en estos métodos, ya sea que resuelvan un objetivo global o local, requiere aclaración.

Metodología: Muestreo MAP Local (LMAPS)

Los autores introducen Muestreo MAP Local (LMAPS), un nuevo marco de inferencia que resuelve iterativamente subproblemas de máxima a posteriori locales a lo largo de la trayectoria de difusión.

1. Marco Teórico

LMAPS opera resolviendo una secuencia de problemas MAP locales en lugar de uno global único. En cada paso de difusión $t$, dado un estado ruidoso $x_t$ y una medición $y$, el método busca:
$$x^*_0(t, x_t, y) := \arg \max_{x_0} p(x_0 | x_t, y)$$
Esto es seguido por un paso de transición estándar DDIM:
$$x_{t-\Delta t} = g(x^*_0, x_t, \epsilon)$$

Distinciones con Métodos Existentes:

• vs. MAP Global: El MAP global busca el modo de la posterior completa $p(x_0 | y)$. LMAPS busca el modo de la posterior condicional $p(x_0 | x_t, y)$ en cada paso. Los autores argumentan que la secuencia de maximizadores locales no necesariamente converge al maximizador global de $p(x_0 | y)$.
• vs. DPS: DPS utiliza la media condicional $E[x_0 | x_t, y]$ en la actualización. LMAPS reemplaza esta media por el modo condicional. Estos son equivalentes solo si la posterior condicional es Gaussiana; de lo contrario, LMAPS introduce un sesgo de búsqueda de modos, evitando regiones de baja densidad entre modos.
• vs. DMAP: Aunque DMAP también utiliza una perspectiva MAP local, optimiza en el espacio de variables de difusión latente ($x^*_{t-\Delta t} = \arg \max p(x_{t-\Delta t} | x_t, y)$). LMAPS plantea el problema directamente en el espacio de la señal limpia ($x^*_0$), proporcionando una visión complementaria que conecta explícitamente con la verosimilitud de la medición $p(y|x_0)$.

2. Algoritmo Práctico y Aproximaciones

Para hacer el problema MAP local tratable para problemas inversos no lineales generales, los autores proponen aproximaciones y reformulaciones específicas:

• Aproximación de Covarianza: Los autores aproximan la covarianza condicional $\Sigma_{0|t}$ del prior $p(x_0 | x_t)$ utilizando una forma isotrópica motivada por una suposición de prior Gaussiano. A medida que el nivel de ruido $\sigma_t \to 0$, la covarianza posterior se vuelve asintóticamente isotrópica. Aproximan $\Sigma_{0|t} \approx \frac{1}{SNR} I$, donde $SNR = \alpha_t^2 / \sigma_t^2$.

• Reformulación del Objetivo: El objetivo MAP local se formula como la minimización de una suma ponderada de la distancia al prior y la fidelidad de los datos:
  $$x^*_0 = \arg \min_{x_0} \left\{ \frac{SNR}{k} \|x_0 - m_{0|t}\|^2 + \frac{1}{\sigma_y^2} \|y - H(x_0)\|^2 \right\}$$
  Para mejorar la estabilidad numérica y la interpretabilidad, esto se reformula en una combinación convexa invariante a la escala:
  $$x^*_0 = \arg \min_{x_0} \left\{ (1 - \mu_t) \frac{1}{2} \|x_0 - m_{0|t}\|^2 + \mu_t k_2 \|y - H(x_0)\|^2 \right\}$$
  donde $\mu_t = \sigma_t^2 / (\sigma_t^2 + k_1^2)$. Esta formulación permite un recocido automático: a medida que $t$ disminuye ($\sigma_t \to 0$), el objetivo se desplaza de estar impulsado por la medición a estar impulsado por el prior.

• Lineal vs. No Lineal: Para problemas inversos lineales, la solución MAP local puede derivarse en forma cerrada, recuperando la Difusión Proyectada de Momentos de Tweedie (TMPD) como un caso especial. Para problemas no lineales, la solución se obtiene mediante descenso de gradiente.

Contribuciones Clave

1. Unificación Teórica: El artículo formula LMAPS para unificar los métodos de problemas inversos basados en optimización y la Difusión Proyectada de Momentos de Tweedie (TMPD) bajo un único marco probabilístico. Aclara que los métodos basados en optimización resuelven inherentemente problemas MAP locales en lugar de MAP global o muestreo completo de la posterior.
2. Avances Metodológicos: Los autores introducen una aproximación de covarianza motivada por la Gaussiana y una función objetivo reformulada que mejora la estabilidad numérica y la interpretabilidad mediante un recocido automático.
3. Validación Empírica: El marco se valida en 10 tareas de restauración de imágenes (incluyendo operadores lineales, no lineales y no diferenciables) y 3 problemas inversos científicos.

Resultados Experimentales

Los autores evaluaron LMAPS en los conjuntos de datos FFHQ e ImageNet para restauración de imágenes y en benchmarks para problemas inversos científicos (Dispersión Inversa Lineal, MRI de Muestreo Comprimido e Imagenología de Agujeros Negros).

• Restauración de Imágenes: LMAPS logró un rendimiento de vanguardia en 43 de 60 casos en las métricas PSNR, SSIM y LPIPS en FFHQ e ImageNet. Consistentemente superó la línea base DMAP, mostrando mejoras de más de 2 dB en PSNR en tareas como desenfoque por movimiento, restauración JPEG y cuantización.
• Problemas Inversos Científicos: LMAPS alcanzó consistentemente el PSNR más alto en todas las tareas científicas. Notablemente, logró ganancias superiores a 1.5 dB en PSNR en tres tareas de Dispersión Inversa Lineal (LIS) en comparación con métodos anteriores.
• Eficiencia: LMAPS demostró una eficiencia superior en comparación con DAPS, logrando mejores resultados con menos recursos computacionales (por ejemplo, menos pasos de optimización por paso de difusión).

Significado y Afirmaciones

El artículo afirma que LMAPS proporciona una interpretación probabilística fundamentada para los métodos basados en optimización, cerrando la brecha entre el muestreo posterior (DPS) y la estimación puntual (MAP). Al resolver explícitamente subproblemas MAP locales, el método ofrece una visión unificada que explica el comportamiento de los solucionadores basados en optimización existentes.

Los autores enfatizan que, aunque LMAPS no muestrea de la distribución posterior completa (y por lo tanto no captura la multimodalidad de la misma manera que DPS), está diseñado específicamente para escenarios donde la estimación puntual precisa es el objetivo principal. El marco ofrece una alternativa robusta que equilibra la fidelidad de la reconstrucción con los priores estructurales de los modelos de difusión, logrando resultados de vanguardia tanto en restauración estándar de imágenes como en tareas desafiantes de reconstrucción científica. El trabajo sugiere que la investigación futura podría explorar aún más la utilidad de la estimación MAP global con priores de difusión, lo cual sigue siendo un desafío abierto.