Resumen Técnico: Restauración de Imágenes Basada en Modelos de Difusión mediante el Aprendizaje Interactivo de Determinismo y Estocasticidad

Planteamiento del Problema

La restauración de imágenes (IR) tiene como objetivo reconstruir la información visual latente a partir de observaciones degradadas, una tarea complicada por las significativas discrepancias distribucionales entre diferentes tipos de degradación (por ejemplo, ruido, lluvia, sombras, baja iluminación) en el espacio de características. Los métodos de aprendizaje profundo existentes, incluyendo CNN, Transformers y GAN, a menudo luchan por generalizar a través de diversas tareas de IR sin rediseños específicos para cada tarea. Si bien los modelos de difusión ofrecen alta fidelidad e interpretabilidad, los enfoques actuales enfrentan dos limitaciones principales:

1. Ineficiencia en la IR: Los modelos de difusión estándar comienzan el proceso inverso desde ruido puro, lo cual es ineficiente para tareas de IR donde la entrada ya contiene información estructural.
2. Modelado Estático de Determinismo y Estocasticidad: Los intentos previos de unificar la difusión determinista (basada en residuos) y la estocástica (basada en ruido), como el Modelo de Difusión de Denoising de Residuos (RDDM), emplean U-Nets independientes para cada trayectoria. Esta separación estática no logra capturar las distintas distribuciones de probabilidad de diferentes tipos de degradación y puede introducir una diversidad excesiva que obstaculiza la reconstrucción precisa.

Metodología

Los autores proponen un nuevo marco de difusión que emplea el Aprendizaje Interactivo de Determinismo y Estocasticidad (ILDS) para abordar estas limitaciones. La arquitectura central utiliza una red de doble trayectoria basada en dos U-Nets simples.

1. Construcción del Proceso de Difusión hacia Adelante (Forward)

A diferencia de la difusión estándar que solo añade ruido, el proceso de difusión propuesto incorpora tanto residuos (determinismo) como ruido (estocasticidad). La transición se define como:
$$x_t = x_{t-1} + \alpha_t x_{res} + \beta_t \epsilon_{t-1}$$
donde $x_{res} = x_{in} - x_0$ representa el residuo de la degradación, y $\epsilon$ es el ruido gaussiano. Esto permite que la trayectoria hacia adelante preserve las estructuras relacionadas con la degradación, permitiendo que el proceso inverso comience desde estados más cercanos al objetivo que el ruido puro, reduciendo así los pasos de muestreo requeridos.

2. Mecanismo de Aprendizaje Interactivo

Para superar las limitaciones del modelado independiente en RDDM, los autores introducen una estrategia de aprendizaje interactivo dentro de la arquitectura U-Net:

• Predicción de Doble Trayectoria: Se utilizan dos U-Nets para predecir el ruido ($\epsilon_\theta$) y el residuo ($x_{res}^\theta$) respectivamente.
• Fusión de Características: En lugar de operar de forma aislada, los codificadores de las dos U-Nets comparten características. Específicamente, después del primer bloque de submuestreo (downsampling), las características de ambas trayectorias se concatenan y se procesan mediante convoluciones de $1 \times 1$ ($\tilde{h}_{res}$ y $\tilde{h}_\epsilon$).
• Modelado Dinámico Dependiente del Tiempo: Las entradas de los pasos temporales (timesteps) para las dos trayectorias son moduladas dinámicamente por parámetros que controlan la acumulación temporal ($\bar{\alpha}_t$) y la intensidad del ruido ($\bar{\beta}_t$). Esto permite al modelo ajustar adaptativamente la contribución de los componentes deterministas y estocásticos basándose en la distribución de probabilidad de características específica del tipo de degradación.
• Contexto Compartido: Las características interactuadas proporcionan información contextual compartida entre trayectoras durante la fase de decodificación, asegurando que la guía determinista y la diversidad estocástica se influyan mutuamente.

3. Entrenamiento y Generación

El modelo se entrena para minimizar la divergencia de Kullback-Leibler (KL) entre la posterior real y la transición inversa aprendida. La función de pérdida final es la suma de los errores cuadráticos para las predicciones de ruido y residuo:
$$L = L_{res}(\theta) + L_{\epsilon}(\theta)$$
Notablemente, estos componentes no están ponderados, ya que el mecanismo interactivo facilita un equilibrio autoadaptativo durante el entrenamiento. Durante la inferencia, el modelo genera imágenes de alta calidad mediante el muestreo iterativo de la distribución condicional, proceso que se vuelve totalmente determinista cuando el factor de escala de ruido $\eta = 0$.

Principales Contribuciones

1. Nuevo Marco de Difusión: El diseño de un marco que modela interactivamente los aspectos deterministas y estocásticos de la degradación de imágenes, permitiendo que el proceso de difusión se adapte a tareas que requieren alta fidelidad (precisión) o diversidad.
2. Modelado Dinámico Dependiente del Tiempo: La incorporación de parámetros de acumulación temporal e intensidad de ruido para modular dinámicamente las entradas de los pasos temporales para las rutas de difusión duales, mejorando la generalización en diversos escenarios de IR.
3. Arquitectura de Doble Trayectoria Eficiente: La implementación de la difusión de doble trayectoria utilizando solo dos U-Nets simples con características de codificador compartidas. Este diseño logra un rendimiento competitivo con un sobrecoste de parámetros mínimo y una alta escalabilidad.

Resultados Experimentales

El método propuesto fue evaluado en cuatro tareas representativas de IR: Reducción de Ruido (Denoising), Eliminación de Sombras, Mejora de Baja Iluminación y Eliminación de Lluvia (Deraining).

• Rendimiento: El método logró un rendimiento superior o competitivo en todas las tareas.
  • Reducción de Ruido (SIDD): Alcanzó el mejor PSNR (41.67 dB) y SSIM (0.9385), superando a RDDM y otros modelos de referencia.
  • Eliminación de Sombras (ISTD): Obtuvo el mejor PSNR (29.29 dB) y MAE (8.776), con el segundo mejor SSIM.
  • Eliminación de Lluvia (Raindrop): Logró el mejor PSNR (31.30 dB) y SSIM (0.9460), eliminando con éxito las trazas de lluvia donde otros métodos dejaban artefactos.
  • Mejora de Baja Iluminación (LOL): Alcanzó el segundo mejor rendimiento (PSNR 24.04, SSIM 0.9118), situándose ligeramente por detrás de LLFormer en métricas, pero manteniendo una fuerte preservación estructural.
• Eficiencia: El aprendizaje interactivo mejoró significamente la eficiencia de muestreo. El modelo requirió solo de 1 a 4 pasos de muestreo para lograr resultados de alta calidad, en comparación con cientos o miles para otros métodos basados en difusión (por ejemplo, DiffusionVI requiere 1000 pasos).
• Estudios de Ablación:
  • La eliminación del módulo de interacción resultó en caídas de rendimiento en todas las tareas, confirmando el beneficio de la fusión de características.
  • El uso de modelado dinámico dependiente del tiempo ($\bar{\alpha}_t, \bar{\beta}_t$) superó al modelado de paso temporal estático ($t$).
  • Los pasos de muestreo óptimos variaron según la tarea (1 para baja iluminación, 2 para reducción de ruido, 3 para eliminación de sombras, 4 para eliminación de lluvia).

Significado y Reivindicaciones

El artículo afirma que el marco propuesto ofrece una solución unificada para la restauración de imágenes multitarea que equilibra la precisión de la reconstrucción y la diversidad. Al integrar el modelado determinista y estocástico mediante el aprendizaje interactivo, el enfoque proporciona una mayor interpretabilidad que los métodos basados en GAN y una mejor generalización que las CNN o Transformers específicos de cada tarea.

Los autores reconocen limitaciones, particularmente en tareas de mejora de baja iluminación con grandes gradientes de iluminación donde la restauración del brillo puede ser insuficiente en comparación con modelos de referencia especializados. También señalan la dependencia actual de datos emparejados (limpios-degradados) para el entrenamiento supervisado. El trabajo futuro se dirige hacia la optimización de las ecuaciones de difusión para tareas sensibles a la iluminación y el desarrollo de estrategias de entrenamiento no supervisadas para reducir la dependencia de datos.

En general, el estudio demuestra que una arquitectura de difusión de doble trayectoria simple e interactiva puede manejar eficazmente la heterogeneidad de los procesos de degradación de imágenes, ofreciendo una alternativa escalable y eficiente para diversos desafíos de restauración de imágenes.