Resumen Técnico: FlowADMM: ADMM Plug-and-play con Priors de Renoise-Denoise Basados en Flujos

Enunciado del Problema

Los problemas inversos en imagen, como la eliminación de ruido, la desenfoque, la super-resolución y la inpainting, suelen estar mal planteados y se formulan como la recuperación de una imagen $x$ a partir de mediciones degradadas $y = Ax + \nu$. Aunque el aprendizaje profundo ofrece mapeos directos potentes, a menudo carecen de interpretabilidad y garantías teóricas. Los enfoques variacionales clásicos formulan esto como la minimización de un término de fidelidad de datos más un término de regularización, a menudo resuelto mediante algoritmos proximales como la División Forward-Backward (FBS).

Los métodos Plug-and-Play (PnP) mejoran esto reemplazando el operador proximal de un regularizador diseñado a mano con un eliminador de ruido aprendido. Recientemente, los flujos generativos y los modelos de difusión han proporcionado familias dependientes del tiempo de eliminadores de ruido $\{D_t\}$ adaptados a niveles de ruido específicos. Sin embargo, surge un desafío clave: durante la reconstrucción iterativa, las iteraciones intermedias a menudo se desvían de la distribución de ruido $p_t$ para la cual se entrenó el eliminador de ruido. Los métodos PnP basados en flujos existentes (por ejemplo, PnP-Flow, Flower) abordan esto incorporando pasos explícitos de renoise-denoise, donde la iteración actual se perturba hacia el nivel de ruido apropiado antes de eliminar el ruido. Estos métodos suelen promediar sobre múltiples muestras de ruido para estabilizar el proceso.

A pesar de su éxito empírico, los métodos PnP de renoise-denoise existentes se formulan como algoritmos estocásticos, lo que dificulta el análisis teórico de su comportamiento de convergencia. Específicamente, el operador determinista inducido implícitamente al promediar la salida del eliminador de ruido sobre la distribución de ruido latente no ha sido caracterizado, dejando las propiedades de convergencia de estos métodos poco comprendidas.

Metodología

El Operador Medio de Renoise-Denoise

Los autores identifican y formalizan el operador determinista subyacente a los métodos PnP estocásticos basados en flujos. Dada una red de velocidad de flujo preentrenada $v_\theta$ que define un eliminador de ruido MMSE dependiente del tiempo $D_t = I + (1-t)v_\theta^t$, el paso estocástico estándar implica re-ruidar una iteración $x$ con ruido $\epsilon \sim \mathcal{N}(0, I)$ para obtener $tx + (1-t)\epsilon$, seguido de la eliminación de ruido.

El artículo define el operador medio de renoise-denoise $\bar{S}_t(x)$ como la esperanza de este proceso sobre la distribución de ruido:
$$\bar{S}_t(x) = \mathbb{E}_{\epsilon \sim \mathcal{N}(0, I)} [D_t(tx + (1-t)\epsilon)]$$
Este operador representa el límite determinista de las actualizaciones estocásticas utilizadas por métodos como PnP-Flow. Los autores analizan su forma residual $R_t(x) = \bar{S}_t(x) - x$, interpretándola como un campo vectorial que proporciona correcciones locales hacia la variedad de datos. Bajo suposiciones ideales de MMSE, los puntos en la variedad de datos son puntos fijos de este operador.

Algoritmo FlowADMM

Basándose en esta idea, los autores proponen FlowADMM, que integra el operador medio de renoise-denoise en el marco del Método de Direcciones Alternas de Multiplicadores (ADMM). A diferencia de FBS, ADMM permite múltiples operadores de prior y maneja términos de datos no diferenciables sin restricciones de tamaño de paso dependientes del operador.

El algoritmo resuelve el problema restringido $\min_{x,z} F_y(x) + R(z)$ sujeto a $x=z$ minimizando el Lagrangiano aumentado. Las actualizaciones son:

1. Consistencia de Datos: $x^{k+1} = \text{prox}_{\tau F_y}(z^k - u^k)$
2. Actualización de Prior: $z^{k+1} = \bar{S}_{t_k}(x^{k+1} + u^k)$
3. Actualización Dual: $u^{k+1} = u^k + x^{k+1} - z^{k+1}$

Dado que la esperanza en $\bar{S}_t$ es intratable, se aproxima mediante muestreo de Monte Carlo. Los autores proponen una estrategia de muestreo dependiente del paso de tiempo: usar menos muestras en las iteraciones tempranas (donde se realizan correcciones gruesas y el operador es menos sensible al ruido) y aumentar el número de muestras en las iteraciones posteriores (donde se refinan los detalles finos y la precisión es crítica).

Análisis de Convergencia

El artículo establece garantías de convergencia para FlowADMM bajo condiciones de Lipschitz débiles en la red de flujo subyacente:

• Programación de Tiempo Fija: Para un paso de tiempo constante $t$, la convergencia está garantizada si la constante de Lipschitz $L_v$ de la red de flujo satisface $L_v < \frac{1}{1-t}$. Esta condición es significativamente más débil que la contracción estricta ($L < 1$) requerida por muchos otros métodos PnP, especialmente cuando $t \to 1$.
• Programaciones de Tiempo No Estacionarias: El análisis se extiende a programas de tiempo $(t_k)_k$ que convergen a un límite $t_{max} < 1$. Los autores prueban la convergencia a un punto fijo del operador límite $T_{t_{max}}$ siempre que la secuencia de operadores se aproxime al límite suficientemente rápido y la red de flujo sea Lipschitz continua con respecto al tiempo $t$.

Contribuciones Clave

1. Caracterización Teórica: Los autores identifican y analizan el operador medio determinista de renoise-denoise subyacente a los métodos PnP estocásticos basados en flujos, caracterizando su estructura residual y su comportamiento de Lipschitz.
2. Propuesta de Algoritmo: Proponen FlowADMM, un algoritmo PnP que integra este operador medio directamente en el marco ADMM.
3. Garantías de Convergencia: Derivan criterios de convergencia para FlowADMM bajo suposiciones débiles de Lipschitz, cubriendo tanto programas de pasos de tiempo constantes como no constantes (variables en el tiempo).
4. Estrategia de Muestreo: Proponen una estrategia de muestreo Monte Carlo de grueso a fino que aumenta el número de muestras en las iteraciones posteriores, motivada por el comportamiento del operador medio.
5. Rendimiento Empírico: Demuestran un rendimiento de vanguardia en múltiples problemas inversos mientras requieren menos evaluaciones de consistencia de datos que los enfoques anteriores de promediado estocástico.

Resultados Experimentales

Los autores evaluaron FlowADMM en los conjuntos de datos CelebA y AFHQ-Cat en cinco tareas: eliminación de ruido, desenfoque, super-resolución, inpainting aleatorio e inpainting en caja.

• Rendimiento: FlowADMM logró resultados de vanguardia en métricas de distorsión (PSNR y SSIM) en 7 de 10 tareas en comparación con métodos PnP recientes basados en flujos y difusión (incluyendo PnP-Flow, Flower, OT-ODE, D-Flow y DiffPIR). Se observaron mejoras notables en las tareas de super-resolución e inpainting.
• Eficiencia: Aunque iguala el número total de evaluaciones de flujo (pasadas de red hacia adelante) del competidor líder Flower5-OT, FlowADMM requirió significativamente menos evaluaciones del operador proximal del término de datos. Esto se debe a que Flower5-OT promedia trayectorias completas de reconstrucción, mientras que FlowADMM promedia solo el paso de prior dentro de cada iteración.
• Estudios de Ablación: Los experimentos confirmaron que el programa de muestreo "pesado al final" (aumentando muestras en iteraciones posteriores) mejora significativamente el rendimiento sobre el muestreo constante. Además, aplicar este programa a PnP-Flow mejoró sus resultados, pero FlowADMM aún lo superó, lo que sugiere que la propia formulación ADMM proporciona una ventaja.
• Limitaciones: El análisis de convergencia depende de constantes de Lipschitz de la red preentrenada que no se penalizan explícitamente durante el entrenamiento. Además, las garantías teóricas se aplican al operador medio determinista, mientras que la implementación práctica utiliza aproximaciones estocásticas de Monte Carlo.

Significado y Afirmaciones

El artículo afirma que, al formalizar explícitamente el operador determinista subyacente a los métodos estocásticos de renoise-denoise, cierra la brecha entre el éxito empírico y la comprensión teórica en PnP basado en flujos. El trabajo proporciona:

• Convergencia Demostrable: A diferencia de los métodos PnP estocásticos anteriores, FlowADMM cuenta con garantías de convergencia bajo condiciones de Lipschitz realistas en la red de flujo.
• Eficiencia Algorítmica: La integración en ADMM permite un uso más eficiente de los recursos computacionales (menos pasos de consistencia de datos) en comparación con los métodos de promediado de trayectorias.
• Perspectiva Teórica de Operadores: El análisis sugiere que el comportamiento de "grueso a fino" del operador medio puede aprovecharse mediante programas de muestreo adaptativos, un principio que beneficia tanto a los enfoques basados en ADMM como a los basados en FBS.

Los autores concluyen que FlowADMM ofrece un marco robusto y fundamentado teóricamente para resolver problemas inversos con priores basados en flujos, aunque reconocen que el trabajo futuro podría centrarse en imponer restricciones de Lipschitz durante el entrenamiento para fortalecer las garantías de estabilidad.