Improving Diffusion Posterior Samplers with Lagged Temporal Corrections for Image Restoration

Resumen Técnico: Mejora de Muestreadores Posteriores de Difusión con Correcciones Temporales Retrasadas (LAMP)

Enunciado del Problema

El muestreo posterior basado en difusión (PS) ha surgido como un marco líder para resolver problemas inversos de imagen (por ejemplo, desenfoque, superresolución, inpainting). Estos métodos combinan un prior de difusión aprendido con restricciones de medición para muestrear la distribución posterior $p(x_0|y)$, donde $y = Kx_0 + e$. Las formulaciones estándar de PS, como DiffPIR y DDRM, operan actualizando iterativamente la estimación actual $x_t$ mediante un paso de eliminación de ruido del modelo de difusión y una corrección de consistencia de datos $D_t$.

Sin embargo, los autores identifican una limitación estructural fundamental en estos enfoques estándar: dependen de estimaciones de consistencia de datos instantáneas y paso a paso ($D_t$) que se recalculan en cada iteración inversa. Este diseño introduce variabilidad temporal e inestabilidad en la dinámica inversa. Desde la perspectiva de los sistemas dinámicos, las actualizaciones estándar de PS corresponden a una discretización de primer orden de la dinámica de difusión más un término de corrección residual. Este término residual actúa como una señal de forzamiento que varía en el tiempo, perturbando la trayectoria y provocando que se desvíe de la dinámica continua suave y bien definida asumida por solucionadores de alto orden (como DPM-Solver) utilizados en la generación incondicional. En consecuencia, la falta de memoria temporal en los métodos estándar de PS conduce a actualizaciones inconsistentes y estabilidad subóptima, particularmente en regímenes mal planteados.

Metodología: LAMP

El artículo propone LAMP (Muestreo Posterior Multiescalón Retrasado), un esquema modular de tipo "plug-in" que aborda estas inestabilidades introduciendo una corrección temporal retrasada.

1. Reinterpretación Dinámica: Los autores reinterpretan la actualización estándar de PS como una discretización de primer orden de una dinámica de difusión forzada. Demuestran que la actualización $x^{PS}_{t-\Delta t}$ puede descomponerse en un paso exponencial de primer orden impulsado por un objetivo consciente de la medición $D_t$ y un término de corrección residual proporcional a la discrepancia $(D_t - \hat{x}_{0|t})$.
2. Discretización de Segundo Orden: Para mejorar la estabilidad, LAMP emplea una discretización de segundo orden de la dinámica subyacente. A diferencia de los solucionadores de alto orden estándar que asumen una EDO suave, LAMP tiene en cuenta explícitamente el término de forzamiento dependiente del residual inherente a los problemas inversos.
3. Mecanismo de Corrección Retrasada: La innovación central es la introducción de una corrección temporal basada en la variación de estimaciones consecutivas. La actualización de LAMP combina la discretización de segundo orden con la corrección residual:
   $$x^{LAMP}_{t-\Delta t} = x^{PS}_{t-\Delta t} + \alpha_{t-\Delta t} A_1(h) \gamma \frac{D_t - D_{t+\Delta t}}{h_{prev}}$$
   Aquí, el término $(D_t - D_{t+\Delta t})/h_{prev}$ actúa como una derivada temporal discreta. Al elegir un coeficiente negativo $\gamma$, el método aplica una corrección retrasada, arrastrando la estimación actual ligeramente hacia las anteriores en lugar de extrapolar. Esto suaviza la trayectoria y contrarresta la variabilidad inducida por las correcciones instantáneas.
4. Preservación Estructural: Crucialmente, los autores demuestran (Proposición 1) que LAMP preserva la estructura estándar de muestreo posterior. La actualización puede reescribirse en la forma canónica $x_{t-\Delta t} = \alpha_{t-\Delta t} \tilde{D}_t + \sigma_{t-\Delta t} \epsilon_\Theta(x_t, t)$, donde $\tilde{D}_t$ es una estimación filtrada temporalmente definida como una combinación convexa de las estimaciones de consistencia de datos actuales y anteriores: $\tilde{D}_t = (1-\beta_t)D_t + \beta_t D_{t+\Delta t}$. Esto asegura que LAMP sea una extensión modular compatible con cualquier columna vertebral de PS existente (por ejemplo, DiffPIR, DDRM) sin modificar el eliminador de ruido ni el procedimiento de entrenamiento.

Análisis Teórico

El artículo proporciona un análisis de riesgo de un paso (Proposición 2) que caracteriza la compensación entre sesgo y varianza de la corrección retrasada. El análisis muestra que LAMP mejora la transición inversa cuando la reducción en la varianza del error de estimación (debida al suavizado temporal) supera el sesgo introducido por el retraso (debido a la deriva del objetivo posterior). Los autores demuestran que esta condición se satisface típicamente en las etapas posteriores del proceso de difusión, lo que motiva una estrategia híbrida donde LAMP se activa después de una breve fase de calentamiento.

Contribuciones Clave

• Identificación de la Incompatibilidad Estructural: Los autores destacan que las interpretaciones estándar basadas en EDO de los solucionadores de difusión no logran capturar los términos de forzamiento dependientes del residual presentes en el muestreo posterior de problemas inversos, lo que conduce a inestabilidad temporal.
• Marco LAMP: Proponen LAMP, un esquema corregido de segundo orden que integra una corrección temporal retrasada en la actualización de muestreo posterior.
• Modularidad y Compatibilidad: Se demuestra que LAMP es un muestreador posterior válido que puede implementarse como un "plug-in" sobre cualquier método de PS existente (por ejemplo, DiffPIR, DDRM) sin reentrenamiento ni cambios en el modelo directo.
• Justificación Teórica: Un análisis de error local establece las condiciones bajo las cuales la corrección retrasada reduce el riesgo, específicamente cuando los objetivos posteriores adyacentes evolucionan suavemente en relación con el ruido de estimación.

Resultados Experimentales

Los autores evalúan LAMP en problemas inversos lineales estándar (desenfoque gaussiano, desenfoque por movimiento y superresolución) en tres conjuntos de datos: FFHQ, ImageNet y CelebA. Comparan versiones mejoradas con LAMP (DiffPIR-LAMP, DDRM-LAMP) contra sus líneas base y variantes ingenuas de alto orden (1M, 2M).

• Rendimiento: LAMP mejora consistentemente la calidad de reconstrucción sobre líneas base sólidas. En el escenario ruidoso ($\sigma_y = 0.05$), DDRM-LAMP logra el mejor rendimiento general, mostrando ganancias consistentes en PSNR, SSIM y LPIPS en todas las tareas. DiffPIR-LAMP también muestra mejoras, particularmente en desenfoque y superresolución.
• Eficiencia: Estas mejoras se logran sin aumentar el número de evaluaciones de eliminación de ruido (NFE), manteniendo el costo computacional de los métodos base (por ejemplo, 20 NFE para DDRM, 100 NFE para DiffPIR).
• Comparación con Alto Orden Ingenuo: Las discretizaciones ingenuas de primer y segundo orden de EDO (sin la corrección residual específica) degradan consistentemente el rendimiento, confirmando que la corrección residual retrasada propuesta es el factor clave para la mejora.
• Resultados Cualitativos: Las comparaciones visuales muestran que LAMP produce texturas más suaves, bordes más nítidos y una mejor preservación de detalles finos (por ejemplo, textura del cabello en CelebA, escamas en imágenes de iguanas) en comparación con las líneas base, reduciendo la granulosidad y los artefactos de desenfoque.

Significado y Afirmaciones

El artículo afirma que el filtrado temporal de estimaciones conscientes de la medición es un mecanismo crítico, previamente poco explorado, para estabilizar solucionadores inversos basados en difusión. Al replantear el muestreo posterior como un sistema dinámico con un término de forzamiento dependiente del residual, los autores demuestran que una simple corrección retrasada puede mejorar significativamente la calidad y la estabilidad de la reconstrucción.

El significado radica en la modularidad del enfoque: LAMP no requiere nuevos priores, reentrenamiento o estrategias de condicionamiento complejas. En su lugar, ofrece una mejora de propósito general y bajo costo para muestreadores posteriores de última generación existentes. Los autores notan modestamente que LAMP hereda las limitaciones de sus métodos base (por ejemplo, no puede compensar incompatibilidades fundamentales del modelo) y que su efectividad depende del comportamiento temporal de la secuencia de corrección del muestreador subyacente. Sin embargo, los resultados sugieren que para líneas base estables, LAMP proporciona una mejora consistente y perceptiblemente significativa en la fidelidad de los datos, particularmente para estructuras de escala fina.