Your Pre-trained Diffusion Model Secretly Knows Restoration

1. El Problema

El objetivo de la Restauración Todo-en-Uno (AiOR) es recuperar imágenes o videos dañados por diversas degradaciones (niebla, lluvia, nieve, desenfoque de movimiento, baja iluminación) utilizando un único modelo unificado.

Aunque los modelos de difusión latente preentrenados (como Stable Diffusion, FLUX y WAN) poseen priors (conocimientos previos) fuertes sobre la distribución de imágenes naturales, las metodologías actuales para aprovecharlos en AiOR presentan limitaciones:

• Ajuste Fino (Fine-tuning): Requiere un entrenamiento costoso y puede sesgar el modelo hacia los datos de entrenamiento, reduciendo su robustez ante degradaciones complejas o del mundo real.
• Módulos de Control (ControlNet): Introducen ramas adicionales que modifican características intermedias, lo que puede debilitar los priors preentrenados si el módulo de control no es robusto.
• Métodos de Edición Basados en Texto: Intentar usar prompts de lenguaje natural (ej. "eliminar niebla") o optimización de tokens de texto (como Textual Inversion) para tareas de restauración de bajo nivel ha demostrado ser ineficaz. Los modelos tienden a suprimir el ruido añadido sin eliminar la degradación subyacente, o simplemente regresan a la imagen de entrada degradada.

La pregunta central: ¿Contienen los priors de los grandes modelos de difusión conocimiento relevante para la restauración y, de ser así, cómo podemos acceder a él sin romper la estructura del modelo?

2. Metodología Propuesta

Los autores proponen un marco de trabajo que desbloquea la capacidad de restauración inherente de los modelos de difusión congelados mediante el aprendizaje de prompts en el espacio de incrustaciones (embedding space) y una formulación de puente de difusión (diffusion bridge).

A. Optimización en el Espacio de Incrustaciones (No en Tokens)

El hallazgo clave es que la optimización de tokens de texto (entrada del codificador de texto) es insuficiente. En su lugar, los autores optimizan directamente las incrustaciones de contexto en la salida del codificador de texto (antes de entrar al modelo de difusión).

• Mecanismo: Se mantiene el modelo de difusión (backbone) y el codificador de texto congelados. Se aprende un vector de contexto $p$ que guía al modelo desde un estado ruidoso y degradado hacia una imagen limpia.
• Resultado: A diferencia de los prompts de texto, esta optimización directa en el espacio de incrustaciones logra que el modelo mapee consistentemente las entradas degradadas a objetivos limpios.

B. Resolución del Desajuste de Trayectoria (Trajectory Mismatch)

Un problema crítico identificado es la inestabilidad al entrenar prompts de manera ingenua.

• El Problema: Si se entrena el prompt usando un proceso forward donde el estado ruidoso $z_t$ se construye anclado a la imagen degradada ($z_{deg}$), el prompt aprende a manejar estados que nunca encontrará durante la inferencia real (donde el estado debe transitar suavemente desde lo degradado hacia lo limpio). Esto causa un desajuste entre entrenamiento e inferencia.
• La Solución: Puente de Difusión Monótono (EBR): Los autores proponen entrenar los prompts utilizando una formulación de puente de difusión (basada en la técnica EBR - Energy-Oriented Diffusion Bridge).
  • En lugar de anclarse solo a $z_{deg}$, los estados intermedios de entrenamiento se construyen como una interpolación monótona entre la imagen degradada ($z_{deg}$) y la limpia ($z_{clean}$), añadiendo ruido progresivo.
  • Esto asegura que el prompt vea durante el entrenamiento los mismos tipos de estados intermedios que encontrará durante el muestreo inverso (inferencia), alineando las dinámicas de entrenamiento y prueba.

C. Diseño de Prompts

• Se aprende un prompt específico para cada tipo de degradación (niebla, lluvia, etc.).
• Para degradaciones mixtas, los prompts se combinan promediando sus velocidades predichas durante el muestreo.
• Se utiliza una inyección residual en el contexto de atención para mayor eficiencia de parámetros.

3. Contribuciones Clave

1. Descubrimiento de Capacidad Oculta: Demostraron que los modelos de difusión preentrenados poseen conocimiento de restauración que es inaccesible mediante prompts de lenguaje natural o optimización de tokens, pero que se desbloquea totalmente mediante la optimización directa de incrustaciones de contexto.
2. Identificación y Solución del Desajuste de Trayectoria: Identificaron que el entrenamiento ingenuo de prompts falla debido a la discrepancia entre los estados de entrenamiento y los de inferencia. Propusieron el uso de un puente de difusión (EBR) para alinear estas trayectorias, permitiendo un entrenamiento estable.
3. Eficiencia y Rendimiento: Implementaron el marco en modelos masivos como WAN (video, 1.3B parámetros) y FLUX (imagen, 12B parámetros). Lograron un rendimiento competitivo con el estado del arte (SOTA) utilizando solo prompts ligeros aprendidos, sin necesidad de ajustar el backbone ni añadir módulos de control pesados.

4. Resultados Experimentales

Los experimentos se realizaron en tareas de restauración de imágenes y video con diversos tipos de degradaciones (in-distribution, out-of-distribution, mixtas y no vistas).

• Rendimiento Cuantitativo:
  • En Imágenes (FLUX): El método superó o igualó a enfoques SOTA como DFPIR, AutoDIR y PixWizard en métricas perceptuales (CLIPIQA, MUSIQ) y de distorsión (PSNR, SSIM), especialmente en escenarios de degradación no vistos (OOD) y mixtos.
  • En Video (WAN): Superó a métodos específicos de video como ViWS-Net y AverNet, mostrando una fuerte generalización temporal gracias a los priors preentrenados de WAN.
• Rendimiento Cualitativo: Las visualizaciones muestran que el método restaura detalles finos y elimina degradaciones complejas (niebla densa, lluvia fuerte) mejor que los métodos basados en ajuste fino, evitando artefactos comunes como la pérdida de textura o la sobre-suavización.
• Ablaciones:
  • La optimización en espacio de incrustaciones superó significativamente a la optimización de tokens.
  • El uso del puente EBR fue crucial; el entrenamiento ingenuo y el uso de puentes DDBM estándar produjeron resultados inferiores o inestables.
  • El nivel de ruido inicial ($T_0$) óptimo se encontró en 0.4, equilibrando la corrección de degradaciones severas y la preservación de detalles.

5. Significado e Impacto

Este trabajo cambia el paradigma de cómo se utilizan los grandes modelos de difusión para tareas de visión de bajo nivel:

• Eficiencia: Elimina la necesidad de costosos procesos de ajuste fino (fine-tuning) o la adición de arquitecturas complejas (ControlNet), reduciendo drásticamente los parámetros entrenables.
• Generalización: Al mantener los priors originales del modelo intactos, el método demuestra una robustez superior ante degradaciones del mundo real y no vistas, algo difícil de lograr con modelos ajustados a datasets específicos.
• Versatilidad: Proporciona una vía para convertir modelos generativos de propósito general en herramientas de restauración de alta calidad mediante un mecanismo de condicionamiento ligero y directo.

En resumen, el artículo demuestra que la "inteligencia" para restaurar ya existe dentro de los modelos de difusión, pero requiere la interfaz correcta (optimización de incrustaciones + alineación de trayectorias) para ser activada.