Stochastic Generative Plug-and-Play Priors

1. Planteamiento del Problema

El artículo aborda los problemas inversos de imágenes, donde el objetivo es recuperar una imagen desconocida $x$ a partir de mediciones incompletas y ruidosas $y$ (modeladas como $y = A(x) + e$).

• Limitaciones de los métodos PnP actuales: Los métodos Plug-and-Play (PnP) tradicionales utilizan un desruido (denoiser) preentrenado como un prior dentro de algoritmos de optimización. Sin embargo, estos métodos enfrentan dos problemas críticos en problemas inversos severamente mal planteados (como la inpainting con máscaras grandes):
  1. Desajuste de distribución: Las iteraciones intermedias de PnP contienen artefactos estructurados, no ruido gaussiano aditivo puro. Los desruidores convencionales (entrenados para bajos niveles de ruido) fallan cuando se aplican fuera de su régimen de entrenamiento.
  2. Puntos de silla estrictos: Los métodos deterministas tienden a quedar atrapados en puntos de silla estrictos en paisajes de optimización no convexos, limitando la calidad de la reconstrucción.
• La brecha: Aunque los Modelos de Difusión Basados en Puntuación (SBDM) han demostrado un gran rendimiento generativo al ser entrenados en un amplio rango de niveles de ruido, no existía un marco sistemático para utilizarlos directamente como priores dentro de los algoritmos PnP sin recurrir a la muestreo de difusión inversa en cada paso de optimización.

2. Metodología Propuesta

Los autores proponen un marco unificado que conecta la optimización PnP clásica con los modelos de difusión, introduciendo el marco SGPnP (Stochastic Generative PnP).

A. Interpretación Basada en Puntuación (Score-based)

Establecen un vínculo matemático directo entre las iteraciones PnP y el desruido basado en puntuación (score). Utilizan la fórmula de Tweedie para adaptar los modelos de difusión preentrenados (tanto de tipo Variance-Exploding -VE- como Variance-Preserving -VP-) para que funcionen como desruidores condicionados por el nivel de ruido $\sigma$ dentro de algoritmos PnP estándar (como ADMM, HQS/DPIR o PGM).

• Ajuste de Parámetros: Desarrollan una estrategia para mapear el nivel de ruido continuo $\sigma$ utilizado en PnP a los pasos de tiempo discretos $t$ de los modelos de difusión preentrenados mediante interpolación e inversión numérica.

B. Marco SGPnP (Inyección de Ruido Estocástica)

La contribución central es el SGPnP, que modifica el paso de actualización del prior en PnP:

1. Re-ruido explícito: Antes de aplicar el desruidor SBDM, se inyecta ruido gaussiano ($\sigma_{inject}$) a la iteración intermedia.
2. Desacoplamiento de niveles de ruido: A diferencia de enfoques anteriores, SGPnP permite que el nivel de ruido inyectado ($\sigma_{inject}$) sea diferente del nivel de condicionamiento del desruidor ($\sigma_{cond}$). Esto es crucial porque las iteraciones intermedias de PnP sufren corrupción por el ruido de medición y los artefactos del operador de medición, no solo por el ruido inyectado.
3. Fórmula de actualización:
   $$x_{k+1} = D_\theta(z_k + \sigma_{inject, k} n; \sigma_{cond, k})$$
   Donde $z_k$ es el resultado del paso de consistencia de datos y $n \sim \mathcal{N}(0, I)$.

3. Contribuciones Clave

1. Interpretación Teórica de PnP: Proporcionan una interpretación basada en puntuación que justifica el uso directo de SBDMs preentrenados como desruidores en PnP, sin necesidad de pasos de muestreo de difusión inversa en cada iteración.
2. Marco SGPnP: Introducen un marco que inyecta ruido para alinear las iteraciones intermedias con la distribución de entrenamiento del SBDM y, simultáneamente, introduce estocasticidad para mejorar la optimización.
3. Garantías Teóricas:
   • Escape de Puntos de Silla: Demuestran que la inyección de ruido induce una optimización sobre un objetivo suavizado por Gaussiano. Bajo ciertas condiciones, esto garantiza que el algoritmo escape de puntos de silla estrictos (strict saddle points), algo que los métodos deterministas no pueden asegurar.
   • Convergencia: Bajo un programa de ruido enfriado (annealed noise schedule), demuestran que las iteraciones convergen a un punto estacionario del objetivo original (no suavizado), recuperando así un punto estacionario del estimador MAP.

4. Resultados Experimentales

Los autores evaluaron el método en tareas desafiantes: reconstrucción de MRI multicapa (fastMRI) y inpainting de imágenes naturales (FFHQ) con máscaras grandes.

• Comparación Cuantitativa: SGPnP superó consistentemente a los métodos PnP deterministas (DPIR, PnP-ADMM), a los métodos PnP estocásticos anteriores (SNORE) y a solucionadores basados en difusión pura (DPS, DiffPIR).
  • En tareas de inpainting con máscaras grandes (donde la información faltante es crítica), SGPnP logró mejoras significativas en métricas como PSNR, SSIM y LPIPS. Por ejemplo, en inpainting de FFHQ, SGPnP alcanzó un PSNR de 25.21 frente a 18.17 de DPIR.
• Análisis de Ablación:
  • Prior Basado en Puntuación: Reemplazar el prior CNN (DRUNet) por un SBDM adaptado mejoró el rendimiento incluso en la versión determinista (SDPnP).
  • Inyección de Ruido: La comparación entre la versión determinista (SDPnP) y la estocástica (SGPnP) con el mismo prior mostró que la inyección de ruido es esencial para resolver problemas severamente mal planteados, permitiendo "alucinar" contenido semántico realista en regiones grandes faltantes.
  • Desacoplamiento: Mostrar que usar niveles de ruido diferentes para la inyección y el condicionamiento ($\sigma_{inject} \neq \sigma_{cond}$) mejora la estabilidad y la reconstrucción.

5. Significado e Impacto

Este trabajo es significativo porque:

• Cierra la brecha entre optimización y generación: Permite utilizar la potencia de los modelos generativos modernos (SBDMs) dentro de esquemas de optimización clásicos, evitando la necesidad de recalcular gradientes de verosimilitud costosos en cada paso.
• Resuelve problemas mal planteados: Proporciona una solución robusta para problemas donde la información es escasa (como máscaras grandes en inpainting), superando las limitaciones de los priores tradicionales que solo funcionan bien en regímenes de bajo ruido.
• Fundamento Teórico: Ofrece por primera vez garantías teóricas sobre la capacidad de los métodos PnP estocásticos para escapar de puntos de silla estrictos, validando empíricamente por qué la inyección de ruido mejora la reconstrucción en escenarios difíciles.

En resumen, SGPnP representa un avance fundamental al transformar los métodos PnP de enfoques deterministas limitados a un marco generativo estocástico robusto, capaz de aprovechar modelos de difusión preentrenados para tareas de restauración de imágenes de alta complejidad.