Resumen Técnico: Ascenso Dual para Difusión en Problemas Inversos

Enunciado del Problema

Los problemas inversos consisten en estimar una señal o imagen desconocida $x$ a partir de observaciones $y$ ruidosas y corruptas, gobernadas por un operador directo $A(\cdot)$. Estos problemas son ubicuos en campos que van desde la astrofísica y la imagen médica hasta la determinación de la estructura de proteínas. Suelen estar mal planteados, lo que hace necesario el uso de priores para restringir el espacio de soluciones.

Aunque los enfoques modernos utilizan modelos de difusión preentrenados como priores potentes, los métodos existentes para resolver problemas de Máxima A Posteriori (MAP) o para muestrear desde la posterior enfrentan limitaciones significativas. Muchos dependen de aproximaciones computacionales para combinar la verosimilitud del modelo de observación con el prior de difusión. Estas aproximaciones a menudo conducen a muestras inexactas, calidad de reconstrucción subóptima o una falla en representar fielmente las observaciones (es decir, residuos altos), lo que puede manifestarse como "alucinaciones" en la solución generada. Además, muchos métodos de última generación luchan contra niveles altos de ruido en las mediciones y sufren de ineficiencia computacional debido a la acumulación de errores o a la necesidad de pasos suplementarios de Cadenas de Markov Monte Carlo (MCMC).

Metodología: Ascenso Dual para Difusión (DDiff)

Los autores proponen DDiff, un nuevo marco para resolver problemas MAP utilizando modelos de difusión preentrenados. A diferencia de los métodos que intentan muestrear de manera demostrable desde la posterior, DDiff se centra en derivar una estrategia de optimización que resuelva el problema MAP de manera precisa y eficiente. El método se inspira en el Método de Direcciones Alternas de Multiplicadores (ADMM), pero introduce una estructura novedosa de ascenso dual adaptada específicamente para modelos de difusión.

Componentes Algorítmicos Principales

El marco DDiff itera a través de tres actualizaciones principales en cada paso, alternando entre las variables primales ($x, z$) y la variable dual ($u$):

1. Actualización de x (Fidelidad de Datos): En lugar de resolver exactamente un problema de mínimos cuadrados con restricciones, DDiff emplea un enfoque de ADMM linealizado. Realiza un solo paso de descenso de gradiente sobre el término de fidelidad de datos $-\log p(y|x)$, lo que permite que el método sea directamente aplicable a cualquier modelo directo diferenciable $A(\cdot)$, ya sea lineal o no lineal.
   $$x \leftarrow v - \gamma \nabla_v \|y - A(v)\|_2^2, \quad \text{donde } v = z - u$$

2. Actualización de z (Eliminación de Ruido): Este paso reemplaza la aplicación ingenua de un modelo de difusión como eliminador de ruido. Los autores observan que los modelos de puntuación estándar $s_\theta(x, t)$ son precisos solo en la variedad de difusión en el tiempo $t$. Para abordar el hecho de que la variable $x+u$ generalmente no se encuentra en esta variedad, DDiff construye una variable recursiva $x_t$ que ancla el proceso de eliminación de ruido. La actualización implica un paso de difusión inversa al estilo DDIM donde la "$x_0$ predicha" es la estimación primal actual $x$, y la variable dual $u$ se reescala y se suma para mantener la coherencia con el nivel de la señal.
   $$z \leftarrow \frac{1}{\sqrt{\bar{\alpha}_t}} (x_t + (1-\bar{\alpha}_t)s_\theta(x_t, t))$$
   donde $x_t$ se actualiza mediante un paso DDIM que incorpora la variable dual $u$.

3. Actualización Dual: La variable dual $u$ se actualiza para imponer la restricción $x=z$ acumulando el error residual entre las variables primales:
   $$u \leftarrow u + x - z$$

Extensión a Espacio Latente

El marco se extiende a LatentDDiff para su uso con Modelos de Difusión Latente (LDM). Esta variante opera sobre representaciones latentes compactas para reducir la demanda de memoria y cómputo. Soporta dos modos para imponer la consistencia de datos:

• Latent-DC: Imponer la consistencia directamente en el espacio latente, requiriendo retropropagación a través del decodificador.
• Pixel-DC: Decodifica la variable latente al espacio de píxeles, aplica la corrección de consistencia de datos y vuelve a codificar. Los autores señalan que Pixel-DC produce reconstrucciones de calidad consistentemente superior.

Contribuciones Clave

1. Marco de Ascenso Dual para Priors de Difusión: El artículo introduce una estrategia de optimización novedosa que integra el ascenso dual (inspirado en ADMM) con modelos de difusión. Este enfoque evita las aproximaciones comunes en los métodos existentes de muestreo de posterior.
2. Demostración de Convergencia a Punto Fijo: Los autores establecen la primera demostración de convergencia a punto fijo para un método de optimización de posterior basado en difusión. Demuestran que, bajo suposiciones acotadas de eliminador de ruido y gradiente, las secuencias de variables primales y duales convergen a un punto fijo $(x^\star, z^\star, u^\star)$, incluso sin asumir un prior convexo.
3. Mejora en Robustez y Fidelidad: El método está diseñado para producir soluciones que modelan las observaciones de manera más fiel, evidenciado por residuos más cercanos a cero en comparación con los métodos de última generación. Esto reduce la probabilidad de alucinaciones introducidas por el prior generativo.
4. Eficiencia: DDiff logra tiempos de inferencia más rápidos al eliminar la necesidad de pasos de equilibrado MCMC suplementarios (requeridos por métodos como DAPS) y reducir el número de evaluaciones de funciones neuronales (NFE) necesarias para obtener resultados de alta calidad.

Resultados Experimentales

Los autores evaluaron DDiff en una amplia gama de problemas inversos lineales y no lineales, incluyendo superresolución, inpainting, desenfoque gaussiano y por movimiento, recuperación de fase, desenfoque no lineal e imagen de alto rango dinámico (HDR). Los experimentos se realizaron en los conjuntos de datos FFHQ e ImageNet.

• Rendimiento Cuantitativo: DDiff superó a las líneas base de última generación (incluyendo DAPS, DMPlug, DCDP, DPS, DiffPIR y DDRM) en la mayoría de las tareas. Logró puntuaciones superiores en Relación Señal-Ruido de Pico (PSNR), Índice de Similitud Estructural (SSIM) y Similitud de Parches de Imagen Perceptiva Aprendida (LPIPS). Cabe destacar que logró los errores residuales más bajos, indicando una mejor consistencia de datos.
• Robustez al Ruido: El método demostró una robustez significativamente mayor ante niveles altos de ruido en las mediciones. Mientras que algunas líneas base mostraron una degradación rápida del PSNR a medida que aumentaba el ruido, DDiff mantuvo un rendimiento estable. En escenarios de alto ruido (por ejemplo, recuperación de fase con $\sigma=0.3$), DDiff recuperó estructuras globales coherentes y semánticas donde otros métodos produjeron artefactos corruptos.
• Eficiencia: En términos de eficiencia temporal, DDiff logró una calidad perceptiva superior con menos NFE en comparación con DAPS. Requirió aproximadamente la mitad del tiempo de muestreo para tareas lineales y un sexto del tiempo para tareas no lineales para lograr resultados comparables o superiores.
• Estudios de Ablación: Los experimentos confirmaron que tanto la actualización de la variable dual como el paso específico de adición de ruido (Ecuación 12 en el artículo) son críticos. Eliminar la variable dual (emulando DiffPIR) o el paso de adición de ruido resultó en caídas significativas del rendimiento, particularmente en términos de artefactos de alta frecuencia y consistencia de datos.

Significado y Afirmaciones

El artículo afirma que DDiff representa un avance significativo en la resolución de problemas inversos con priores de difusión, al pasar de aproximaciones de muestreo de posterior a un marco de optimización robusto de ascenso dual. Los autores enfatizan que su método:

• Modela fielmente las observaciones: Al minimizar los residuos, reduce el efecto de "alucinación" a menudo visto en solucionadores de problemas inversos generativos.
• Es computacionalmente eficiente: Elimina el cuello de botella de los pasos MCMC y reduce la carga computacional de los cálculos de gradiente.
• Está fundamentado teóricamente: La demostración de convergencia proporcionada ofrece una garantía teórica para la estabilidad del método, una característica que falta en muchos enfoques de optimización basados en difusión anteriores.

Los autores concluyen que DDiff es una solución práctica para problemas inversos a gran escala y complejos, particularmente en escenarios donde el ruido de medición es prominente, como en la reconstrucción de TC de dosis reducida. Se sugiere trabajo futuro para extender el marco a datos de mayor dimensión, como volúmenes 3D y video.