DPMambaIR: All-in-One Image Restoration via Degradation-Aware Prompt State Space Model

Resumen Técnico: DPMambaIR

Definición del Problema

La restauración de imágenes "All-in-One" (todo en uno) tiene como objetivo recuperar imágenes de alta calidad a partir de diversas entradas degradadas utilizando un único modelo unificado, abordando las limitaciones prácticas de diseñar modelos dedicados para cada tipo específico de degradación (por ejemplo, ruido, desenfoque, lluvia, baja iluminación). Los enfoques existentes generalmente se dividen en dos categorías: Mezcla de Expertos (MoE) y métodos basados en Prompts. Sin embargo, estos sufren de desventajas significativas:

• Los métodos basados en MoE suelen incurrir en altos costos computacionales y carecen de un modelado físico explícito de las degradaciones, lo que limita la adaptabilidad a escenarios no vistos.
• Los métodos basados en Prompts dependen típicamente de priors de grano grueso y estáticos (etiquetas de clase discretas o tokens semánticos fijos) que no logran capturar la naturaleza continua y la variabilidad espacial de las degradaciones del mundo real. Además, suelen utilizar un paradigma de fusión aditiva donde las pistas de degradación simplemente se anexan a las características visuales, dejando la arquitectura de la red subyacente gobernada por parámetros fijos que no pueden adaptarse dinámicamente a las demandas de restauración heterogéneas.

Además, la optimización conjunta de degradaciones heterogéneas suele inducir un sesgo implícito hacia las estructuras de baja frecuencia, lo que dificulta la recuperación de detalles finos de alta frecuencia.

Metodología

Los autores proponen DPMambaIR, un nuevo marco de trabajo "All-in-One" que cambia de la fusión de características convencional a la modulación dinámica de parámetros explícitos. La arquitectura se construye sobre una red en forma de U basada en Mamba con un diseño de codificador-decodificador asimétrico y consta de tres componentes principales:

1. Extractor de Degradación

Para capturar características de degradación finas y continuas, los autores diseñan un extractor de degradación pre-entrenado.

• Mecanismo: En lugar de depender de etiquetas discretas, el extractor utiliza un objetivo de reconstrucción para la regresión de embeddings de degradación continuos. Emplea la Convolución de Diferencia Central (CDC) para identificar pistas basadas en gradientes (bordes, texturas) y ramas multiescala para percibir degradaciones en diferentes bandas de frecuencia.
• Entrenamiento: El extractor se pre-entrena mediante una tarea de reconstrucción autosupervisada donde mapea una imagen degradada a un embedding de baja dimensión ($E_d$). Un decodificador intenta entonces reconstruir la imagen degradada usando $E_d$ y un embedding de contenido derivado de la verdad de campo (ground truth) limpia. Esto asegura que $E_d$ capture la información de degradación pura de forma independiente al contenido de la imagen.
• Reparametrización: Durante la inferencia, una técnica de reparametrización fusiona la estructura de múltiples ramas en una representación compacta para equilibrar la precisión y la eficiencia.

2. Modelo de Espacio de Estados de Prompt Consciente de la Degradación (DP-SSM)

La innovación central es el DP-SSM, que integra los embeddings de degradación extraídos directamente en la dinámica del Modelo de Espacio de Estados (SSM).

• Modulación Dinámica: A diferencia de los SSM estándar donde los parámetros son fijos o dependen de la entrada, el DP-SSM utiliza el embedding de degradación para modular dinámicamente el paso de discretización ($\Delta$), la matriz de entrada ($B$) y la matriz de salida ($C$).
  • Modulación de $\Delta$: El modelo aprende a asignar valores de $\Delta$ más pequeños para degradaciones de alta frecuencia (por ejemplo, ruido), creando un estado de alta inercia que actúa como un filtro de paso bajo para suprimir el ruido. Por el contrario, se generan valores de $\Delta$ más grandes para degradaciones estructurales o de baja frecuencia (por ejemplo, desenfoque, baja iluminación), entrando en un régimen de alta ganancia para amplificar señales débiles.
  • Modulación de $B$ y $C$: $B$ controla qué tan fuertemente las observaciones externas impulsan las transiciones de estado, mientras que $C$ inyecta priors de degradación globales en la lectura de características.
• Eficiencia: Este enfoque evita el alto costo computacional del escaneo multidireccional (utilizado para mitigar las limitaciones causales en los SSM estándar) al utilizar en su lugar la matriz $C$ modulada para inyectar contexto global, manteniendo una complejidad lineal.

3. Bloque de Mejora de Alta Frecuencia (HEB)

Para contrarrestar el sesgo de baja frecuencia inherente a la optimización conjunta, se introduce un HEB ligero.

• Función: Extrae explícitamente los componentes de alta frecuencia restando las características de baja frecuencia promediadas globalmente de las características de entrada y sumándolas de nuevo. Esto complementa la modulación dinámica del DP-SSM para asegurar la restauración de texturas locales y bordes.

Contribuciones Clave

1. Extractor de Degradación de Grano Fino: Un módulo basado en regresión que captura características de degradación continuas y detalladas (tipo e intensidad) en lugar de depender de etiquetas de clase discretas.
2. Modelo de Espacio de Estados de Prompt Consciente de la Degradación (DP-SSM): Un mecanismo novedoso que modula dinámicamente los parámetros del SSM ($\Delta, B, C$) basándose en los priors de degradación. Esto permite que el modelo adapte físicamente sus dinámicas de evolución (por ejemplo, cambiando entre filtrado de paso bajo y amplificación de señal) a características de degradación específicas.
3. Bloque de Mejora de Alta Frecuencia (HEB): Un módulo complementario diseñado para aliviar el sesgo de baja frecuencia en el aprendizaje multitarea, asegurando la recuperación de detalles finos.
4. Rendimiento de Vanguardia (SOTA): Experimentos extensos demuestran que DPMambaIR logra un rendimiento superior en siete tipos de degradación, tanto en configuraciones "All-in-One" como en tareas de degradación específicas.

Resultados Experimentales

El método fue evaluado en un conjunto de datos mixto que contiene siete tipos de degradación (neblina, lluvia, nieve, baja iluminación, desenfoque, ruido, JPEG) y benchmarks específicos para tareas individuales.

• Rendimiento All-in-One: En el conjunto de datos mixto, DPMambaIR alcanzó un PSNR de 27.69 dB y un SSIM de 0.893, superando a los métodos All-in-One existentes como AdaIR (27.26 dB) y OneRestore (26.64 dB). Logró el mejor rendimiento en las siete subtareas, equilibrando eficazmente los conflictos multitarea.
• Rendimiento Específico de la Degradación: El modelo también logró resultados competitivos o superiores en tareas específicas en comparación con métodos dedicados SOTA:
  • Eliminación de Lluvia (Deraining): 32.30 dB PSNR (superando a AdaIR por 0.84 dB).
  • Mejora de Baja Iluminación: 24.20 dB PSNR.
  • Eliminación de Desenfoque (Deblurring): 31.42 dB PSNR.
  • Eliminación de Neblina (Dehazing): 30.89 dB PSNR.
• Eficiencia: DPMambaIR mantiene una carga computacional comparable (Parámetros: 29.35M, GFLOPs: 153.33) a otros métodos SOTA mientras ofrece un rendimiento significativamente mayor.
• Generalización (OOD): El modelo demostró una generalización robusta en tipos de degradación no vistos (gotas de lluvia y pixelación) sin ajuste fino, alcanzando un rendimiento de primer nivel. Manejó con éxito degradaciones compuestas (por ejemplo, pixelación + baja iluminación) donde otros métodos fallaron.
• Limitaciones: Los autores señalan un límite de falla con gotas de lluvia pesadas y opacas que ocultan completamente el fondo, atribuyéndolo a la necesidad de alucinación generativa (inpainting) en lugar de pistas de restauración.

Significado

El artículo afirma que DPMambaIR representa un avance significativo en la restauración de imágenes All-in-One al ir más allá de los paradigmas de fusión aditiva estática. Al introducir un mecanismo de prompt consciente de la degradación que reparametriza dinámicamente el Modelo de Espacio de Estados, el marco logra una adaptabilidad física a diversas dinámicas de degradación que antes era inalcanzable. El trabajo subraya la importancia del modelado de degradación de grano fino y la modulación dinámica de características para crear soluciones de restauración unificadas, robustas y eficientes. Los resultados validan que modelar explícitamente las características de la degradación e integrarlas en la dinámica central del sistema puede resolver los conflictos inherentes al aprendizaje multitarea, ofreciendo una dirección prometedora para futuros marcos de restauración de imágenes unificados.