Resumen Técnico: Alinear e Invertir

Enunciado del Problema

Los problemas inversos buscan recuperar una señal limpia $x_0$ a partir de mediciones ruidosas o degradadas $y$. Aunque los modelos preentrenados basados en difusión y flujos se han convertido en priores potentes para estas tareas, a menudo tienen dificultades para producir imágenes altamente detalladas en escenarios que involucran degradaciones severas (por ejemplo, super-resolución extrema, desenfoque por movimiento) o escenas complejas. Una limitación específica surge con los Modelos de Difusión Latente (LDM), donde las no linealidades de la arquitectura codificador-decodificador pueden introducir artefactos durante el proceso de muestreo inverso. Además, los métodos existentes a menudo requieren un gran número de pasos de discretización para lograr reconstrucciones de alta calidad, y carecen de restricciones semánticas fuertes durante la inferencia para guiar la trayectoria hacia soluciones perceptualmente realistas.

Metodología

El artículo propone Alinear e Invertir (Align & Invert), un marco que integra Alineación de Representaciones (REPA) en el proceso de inferencia de solucionadores basados en difusión y flujos para problemas inversos.

Concepto Central

El método impone una alineación entre las representaciones internas del modelo generativo (difusión/flujos) y las características de un codificador visual preentrenado y auto-supervisado (específicamente DINOv2). Si bien REPA se utilizaba anteriormente como regularizador de entrenamiento para mejorar la calidad generativa, este trabajo lo adapta para la inferencia en problemas inversos donde las señales de verdad fundamental no están disponibles.

Implementación Técnica

1. Construcción de un Proxy: Dado que la verdad fundamental $x_0$ es desconocida, el método construye una representación proxy $c_{proxy}$ para suplir las características objetivo faltantes.

   • Inicialmente, $c_{proxy}$ se deriva de las características DINOv2 de la medición disponible $y$.
   • A medida que avanza el proceso de difusión inversa, el proxy se actualiza gradualmente a las características DINOv2 de la estimación actual desruidada del modelo, $f_{DINOv2}(D(E[z_0 | z_t]))$.
   • Esta estrategia se basa en la robustez observada de las características DINOv2 frente a degradaciones comunes de imágenes.

2. Distribución Inclinada: El proceso de muestreo se guía mediante una distribución inclinada que incorpora el término de alineación como una recompensa:
   $$\tilde{p}(x_t | y) \propto p(y | x_t) p_t(x_t) \exp(\lambda \cdot \text{REPA}(\bar{x}_0, h_t))$$
   Aquí, $p(y | x_t)$ impone la consistencia de la medición, $p_t(x_t)$ es el prior incondicional, y el término REPA fomenta la alineación entre los estados ocultos intermedios del modelo de difusión $h_t$ y las características proxy de DINOv2.

3. Integración del Algoritmo: El enfoque se integra en solucionadores existentes de última generación (por ejemplo, Latent DPS, ReSample) añadiendo un paso de gradiente basado en el objetivo REPA durante el bucle de muestreo inverso. Este paso alinea el estado de difusión latente con las características semánticas del objetivo.

4. Adaptabilidad: Para modelos que no fueron preentrenados originalmente con REPA, el marco aprende una cabeza de proyección ligera post-hoc ($g_\phi$) para mapear las representaciones internas del modelo de difusión al espacio de características de DINOv2. Este mapeo se optimiza manteniendo el modelo de difusión congelado, lo que hace que la adaptación sea computacionalmente económica.

Perspectivas Teóricas

El artículo proporciona dos resultados teóricos principales que caracterizan el regularizador REPA:

1. Minimización de la Divergencia: Se demuestra que maximizar el objetivo REPA es equivalente a minimizar la Discrepancia de la Media Máxima (MMD) entre las distribuciones de características DINOv2 de la imagen reconstruida y la verdad fundamental, hasta términos que accountan por el desalineamiento de características y los errores de aproximación del proxy.
2. Contracción del Estado Latente: Bajo suposiciones de regularidad respecto a los jacobianos del codificador de difusión y la cabeza de proyección, una sola actualización de REPA induce una contracción en el espacio de estados de difusión. Específicamente, dirige el estado latente actual $z_t$ más cerca de la representación latente de la imagen limpia $z^*$, siempre que el estado actual se encuentre dentro de una vecindad local de la solución.

Resultados Experimentales

Los autores evaluaron el método en cuatro problemas inversos: Super-resolución (4×, 8×, 16×), Inpainting de cajas, Desenfoque Gaussiano y Desenfoque por movimiento, utilizando los conjuntos de datos ImageNet y FFHQ.

• Calidad Perceptiva: La integración de REPA mejoró consistentemente las métricas perceptivas (LPIPS y FID) en todas las tareas y solucionadores (Latent DPS, ReSample, Latent DAPS). Por ejemplo, en Super-resolución 4× en ImageNet, Latent DPS + REPA redujo el LPIPS de 0.238 a 0.217 y el FID de 123.82 a 86.88.
• Métricas de Píxel: Las métricas estándar a nivel de píxel (PSNR, SSIM) permanecieron comparables a las líneas base, reflejando el conocido compromiso entre percepción y distorsión donde REPA prioriza el realismo semántico sobre la precisión perfecta a nivel de píxel.
• Eficiencia: Un hallazgo significativo es que REPA permite a los solucionadores alcanzar el mismo nivel de calidad perceptiva con significativamente menos pasos de discretización. En super-resolución y desenfoque gaussiano, el método logró puntuaciones LPIPS comparables con aproximadamente 4 veces menos pasos que la línea base, lo que conduce a un tiempo de ejecución general reducido a pesar del cómputo adicional por paso.
• Robustez: El método permaneció efectivo bajo degradaciones severas (por ejemplo, escalado 16×, desenfoque de alta intensidad) y cuando se aplicó a modelos preentrenados independientemente a través de la cabeza de proyección aprendida.

Significado y Afirmaciones

El artículo afirma que la alineación de representaciones constituye una herramienta efectiva para mejorar los solucionadores de problemas inversos generativos en tiempo de inferencia.

• Guía Semántica: Al alinear las representaciones internas del modelo con características visuales robustas (DINOv2), el método proporciona un prior semántico fuerte que guía la reconstrucción hacia soluciones perceptivamente significativas, incluso sin supervisión de verdad fundamental.
• Justificación Teórica: El trabajo aclara por qué funciona la alineación, mostrando que actúa como un mecanismo para minimizar la divergencia en el espacio de características y contraer el estado latente hacia la solución de imagen limpia.
• Aplicabilidad General: El marco es general, compatible con diversos solucionadores basados en difusión/flujos, y adaptable a modelos que no fueron entrenados originalmente con alineación mediante un mapeo ligero post-hoc.
• Eficiencia: Aborda el cuello de botella computacional del muestreo de difusión acelerando la convergencia, permitiendo una reconstrucción de alta calidad con menos pasos de muestreo.

Los autores concluyen que este enfoque cierra la brecha entre las capacidades generativas de los modelos de difusión y la comprensión semántica de los codificadores auto-supervisados, ofreciendo una manera principista de mejorar la calidad de la reconstrucción en configuraciones de problemas inversos desafiantes.