Resumen Técnico: Problemas Inversos con Modelos de Difusión mediante Pérdida de Búsqueda de Modos

Definición del Problema

El artículo aborda el desafío de resolver problemas inversos utilizando modelos de difusión incondicionales preentrenados. Un problema inverso consiste en estimar una muestra de datos original $x_0$ dado una medición parcialmente degradada $y$, gobernada por la ecuación $y = A(x_0) + \eta$, donde $A$ es un operador de degradación conocido y $\eta$ es ruido gaussiano.

Aunque los modelos de difusión preentrenados pueden resolver estos problemas de manera zero-shot (plug-and-play), los enfoques existentes enfrentan limitaciones significativas:

1. Métodos de Muestreo de la Posterior: Estos buscan extraer muestras proporcionales a la probabilidad posterior $p(x_0|y)$. Sin embargo, a menudo dependen de aproximaciones de modelado (por ejemplo, aproximaciones gaussianas del score condicional) o métodos de Monte-Carlo secuenciales, lo que puede ser computacionalmente demandante y puede no adherirse estrictamente a la verdadera posterior sin validación de verdad fundamental (ground-truth).
2. Métodos de Estimación MAP: Estos buscan encontrar la estimación de Máxima A Posteriori (MAP) (el modo de la posterior). Los solvers existentes suelen utilizar aproximaciones heurísticas, dependen de la división de variables (por ejemplo, HQS, ADMM) o requieren modelos de consistencia. Estos enfoques frecuentemente involucran computaciones de gradiente costosas y suposiciones de modelado que pueden no cumplirse para todos los operadores lineales.

Los autores argumentan que muchas tareas de restauración de imágenes priorizan la generación de una imagen plausible consistente con las mediciones, en lugar de cubrir toda la distribución posterior. Esto se alinea más estrechamente con la estimación MAP, motivando la necesidad de una estrategia MAP más principled y libre de aproximaciones.

Metodología: Pérdida de Búsqueda de Modo Variacional (VML)

La contribución central es la Pérdida de Búsqueda de Modo Variacional (VML), una nueva función objetivo diseñada para dirigir las muestras hacia la estimación MAP durante el proceso de difusión inversa.

Fundamento Teórico

Los autores derivan la VML desde una perspectiva variacional minimizando la divergencia de Kullback-Leibler (KL) entre la posterior de la difusión $p(x_0|x_t)$ y la posterior de la medición $p(x_0|y)$ en cada paso de la difusión inversa $t$:
$$\text{VML}(x_t, t) = D_{KL}(p(x_0|x_t) \parallel p(x_0|y))$$

La intuición es que a medida que $t \to 0$, la distribución $p(x_0|x_t)$ se vuelve cada vez más concentrada alrededor de $x_t$. Al encontrar un $x_t^*$ óptimo que minimice esta divergencia KL en cada paso, el proceso guía la muestra hacia los modos de la posterior $p(x_0|y)$.

Derivación de Forma Cerrada para Operadores Lineales

Un logro teórico clave es la derivación de una expresión de forma cerrada para la VML en el caso de problemas inversos lineales ($A(x) = Hx$) sin ninguna aproximación de modelado. La expresión completa de VML incluye:

1. Un término de consistencia de la medición: $\|y - HD(x_t, t)\|^2 / (2\sigma_y^2)$.
2. Términos de prioridad (prior): Que involucran al denoiser $D(x_t, t)$, el esquema de ruido $\sigma_t$ y la covarianza de la posterior $Cov[x_0|x_t]$.

Objetivos Simplificados y Reponderados

La VML completa contiene términos de orden superior que involucran la matriz de covarianza $Cov[x_0|x_t]$, los cuales son computacionalmente costosos de calcular y diferenciar. Los autores demuestran que estos términos de orden superior convergen uniformemente a cero (o a una constante) cuando $t \to 0$. En consecuencia, proponen una VML Simplificada ($VML_S$) que descarta estos términos, reteniendo solo la consistencia de la medición y los términos de prioridad.

Para mejorar aún más la estabilidad de la optimización, introducen una VML simplificada reponderada ($VML_S^\rho$):
$$VML_S^\rho(x_t, t) = \frac{\|y - HD(x_t, t)\|^2}{2\sigma_y^2} - \rho_t \log p(x_t) - \rho_t \frac{\|D(x_t, t) - x_t\|^2}{2\sigma_t^2} + C$$
donde $\rho_t$ es un esquema de peso de la prioridad tal que $\rho_t \to 1$ cuando $t \to 0$. Esta formulación permite una computación de gradiente eficiente utilizando diferenciación automática estándar.

Algoritmos

Basándose en la minimización de VML, los autores proponen:

1. VML-MAP: Un algoritmo que utiliza el descenso de gradiente para minimizar $VML_S$ en cada paso de la difusión inversa.
2. VML-MAP$_{pre}$: Una variante precondicionada para operadores lineales mal condicionados. Utiliza la Descomposición en Valores Singulares (SVD) de la matriz de degradación $H$ para construir un precondicionador $P$, acelerando la convergencia para tareas difíciles como el desenfoque (deblurring).
3. VML-MAP$^\tau$: Una estrategia para la restauración de imágenes con ruido. Para evitar artefactos agudos causados por desafíos de optimización en niveles bajos de ruido, el algoritmo minimiza la VML solo hasta un umbral de tiempo de difusión $\tau$, y luego cambia al muestreo DDIM estándar para generar la imagen final.
4. LatentVML-MAP: Una extensión a los Modelos de Difusión Latente (LDMs), aproximando la VML en el espacio latente.

Principales Contribuciones

1. Nueva Función de Pérdida: Introducción de VML, que teóricamente dirige las muestras hacia las estimaciones MAP minimizando la divergencia KL entre las posteriores de la difusión y la medición.
2. Derivación Libre de Aproximaciones: Para problemas inversos lineales, la VML se deriva en forma cerrada sin aproximaciones de modelado heurísticas, proporcionando una base teórica formal para la estimación MAP con modelos de difusión.
3. Análisis de Convergencia: Prueba formal de que los objetivos VML simplificados y reponderados convergen al log-posterior negativo (el objetivo MAP) cuando $t \to 0$ bajo supuestos moderados.
4. Algoritmos Prácticos: Desarrollo de VML-MAP y su variante precondicionada, ofreciendo una solución empíricamente efectiva para diversas tareas de restauración de imágenes.
5. Validación Extensa: Demostración de un rendimiento superior sobre las líneas base existentes (por ejemplo, DDRM, $\Pi$GDM, MAPGA, DAPS) a través de múltiples conjuntos de datos (ImageNet, FFHQ, CelebA) y tareas (inpainting, super-resolución, deblurring).

Resultados Experimentales

Los autores evaluaron sus métodos en tareas de restauración de imágenes utilizando modelos preentrenados en ImageNet256, FFHQ256 y CelebA.

• Métricas: El artículo enfatiza las métricas perceptuales (LPIPS y FID) sobre las métricas de píxel a píxel (PSNR, SSIM), argumentando que estas últimas a menudo fallan en capturar la calidad perceptual de las reconstrucciones plausibles en problemas inversos mal planteados.
• Rendimiento:
  • VML-MAP superó consistentemente a los métodos de muestreo de la posterior (DDRM, $\Pi$GDM) y otros estimadores MAP (MAPGA, DAPS) en términos de LPIPS y FID en tareas de inpainting, super-resolución y deblurring.
  • VML-MAP$_{pre}$ mostró mejoras significativas sobre la versión no precondicionada para tareas mal condicionadas (por ejemplo, deblurring), logrando puntuaciones de LPIPS y FID más bajas.
  • Eficiencia Computacional: VML-MAP logró una mejor calidad perceptual con un menor tiempo de ejecución (menos evaluaciones de la función neuronal) en comparación con los métodos que requieren muchos pasos de difusión inversa.
• Restauración con Ruido: La estrategia VML-MAP$^\tau$ manejó con éxito las mediciones con ruido, produciendo reconstrucciones de alta calidad donde el VML-MAP estándar presentaba dificultades con artefactos.
• Modelos Latentes: LatentVML-MAP demostró efectividad en CelebA, aunque los autores señalan que la optimización sigue siendo más desafiante en el espacio latente debido a la no linealidad del decodificador.

Significado y Reivindicaciones

El artículo afirma que la VML representa la primera modelación libre de aproximaciones de la estimación MAP para problemas inversos lineales aprovechando únicamente un modelo de difusión incondicional preentrenado, respaldada por un análisis teórico formal.

Los autores enfatizan que su enfoque difiere fundamentalmente del muestreo de la posterior (que apunta a la cobertura de la distribución) y de los solvers MAP existentes (que a menudo dependen de heurísticas o modelos de consistencia). Al minimizar la VML, el método apunta efectivamente a los "modos" de la posterior, lo cual, según los autores, es el comportamiento deseado para generar restauraciones de imagen plausibles y de alta calidad.

El trabajo es modesto respecto a sus limitaciones, reconociendo que:

• La efectividad depende de la disponibilidad de optimizadores prácticamente eficientes para el objetivo VML.
• Persisten desafíos de optimización para operadores no lineales o mal condicionados, aunque el precondicionamiento ayuda.
• Se necesita trabajo futuro para desarrollar optimizadores de orden superior que mitiguen aún más los problemas de optimización sin costos computacionales prohibitivos.

El artículo concluye que VML-MAP proporciona un marco teóricamente fundamentado, basado en principios y empíricamente efectivo para resolver problemas inversos con modelos de difusión, cerrando la brecha entre la estimación MAP teórica y la restauración de imágenes práctica.