VIPaint: Image Inpainting with Pre-Trained Diffusion Models via Variational Inference

1. Planteamiento del Problema

El artículo aborda el desafío del inpainting de imágenes (rellenar regiones faltantes o enmascaradas de una imagen) utilizando Modelos de Difusión Preentrenados (DM), particularmente Modelos de Difusión Latente (LDM) como Stable Diffusion.

• La Dificultad Central: Aunque los DM sobresalen en generar imágenes novedosas a partir de ruido, condicionarlos sobre imágenes corruptas u observadas parcialmente (inpainting) no es trivial. El objetivo es muestrear de la verdadera distribución posterior condicional $p(x|y)$, donde $y$ es la imagen observada (enmascarada) y $x$ es la imagen completa.
• Limitaciones de los Métodos Existentes:
  • Métodos Heurísticos/Iterativos: Enfoques como RePaint, CoPaint y DPS (Muestreo Posterior de Difusión) a menudo dependen de refinamiento iterativo o heurísticas de "viaje en el tiempo". Estos frecuentemente fallan en capturar la estructura global cuando se enmascaran grandes regiones, lo que conduce a resultados incoherentes o poco realistas.
  • Métodos Variacionales (REDdiff): Los intentos previos de inferencia variacional (por ejemplo, REDdiff) a menudo optimizan una sola imagen a nivel libre de ruido ($t=0$) utilizando una aproximación gaussiana simple. Esto falla en capturar la incertidumbre, es propenso a la inestabilidad de la optimización (especialmente en regímenes de bajo ruido) y lucha para generalizar a LDM.
  • Requisitos de Entrenamiento: Muchos métodos de última generación requieren ajustar fino el modelo generativo en conjuntos de datos o máscaras específicos, lo cual es computacionalmente costoso y evita la aplicación "zero-shot" a nuevas tareas.

2. Metodología: VIPaint

Los autores proponen VIPaint, un algoritmo de inferencia variacional (VI) jerárquica que trata el modelo de difusión preentrenado como un prior. En lugar de optimizar una sola imagen, VIPaint optimiza una distribución sobre variables latentes en múltiples niveles de ruido.

A. Posterior Variacional Jerárquico

VIPaint define una aproximación no gaussiana y markoviana de la verdadera posterior de difusión. Introduce un conjunto de $K$ "puntos clave" (pasos de tiempo intermedios) $\{h(1), \dots, h(K)\}$ dentro de la trayectoria de difusión inversa (donde $h(K)$ es ruido alto y $h(1)$ es ruido bajo).

La distribución variacional se define como:
$$q_\lambda(z_{h(1):h(K)}) = q_\lambda(z_{h(K)}) \prod_{i=1}^{K-1} q_\lambda(z_{h(i)} | z_{h(i+1)})$$

• Estructura:
  • Ruido Alto ($h(K)$): Modelado como un Gaussiano factorizado.
  • Ruido Intermedio/Bajo: Modelado como un Gaussiano donde la media es una combinación convexa de la predicción de difusión del prior (del modelo preentrenado) y un parámetro variacional aprendido específico de la imagen.
  • Varianza: La varianza se aprende y es específica por dimensión, permitiendo al modelo aumentar dinámicamente la incertidumbre en regiones enmascaradas mientras mantiene certeza en regiones observadas.

B. La Función Objetivo

VIPaint optimiza una cota inferior variacional sobre la verosimilitud marginal de la observación $y$. La función de pérdida consta de tres términos:

1. Pérdida de Reconstrucción: Asegura que las muestras generadas en el nivel de ruido más bajo ($h(1)$) sean consistentes con los píxeles observados $y$. Esto se aproxima utilizando una predicción de eliminación de ruido en un solo paso.
2. Pérdida del Prior: Regulariza el nivel de ruido más alto ($h(K)$) para alinearse con la semántica global codificada por el prior de difusión en el régimen de alto ruido.
3. Pérdida Jerárquica: Asegura la consistencia entre puntos clave adyacentes en la jerarquía, capturando detalles y estructura de nivel medio de la imagen.

C. Optimización y Muestreo

• Optimización: Los parámetros variacionales $\lambda$ (medias, varianzas y pesos de mezcla) se optimizan por imagen utilizando descenso de gradiente (Adam). Crucialmente, esto se realiza sin recocido (se utiliza muestreo de tiempo aleatorio para la pérdida), lo que proporciona una cota imparcial y evita la inestabilidad observada en métodos que recocen de $T$ a $0$.
• Muestreo: Después de la optimización, se generan reconstrucciones diversas mediante:
  1. Muestreo Ancestral: Muestreo desde la posterior jerárquica optimizada desde $h(K)$ hasta $h(1)$.
  2. Refinamiento: Refinar further la muestra desde $h(1)$ hasta $t=0$ utilizando la red de eliminación de ruido preentrenada, guiada por el gradiente de verosimilitud (similar a DPS pero comenzando desde una posterior mejor inicializada).

3. Contribuciones Clave

1. Aproximación No Gaussiana Jerárquica: A diferencia de los métodos VI anteriores que aproximan la posterior en un solo punto ($t=0$), VIPaint modela la posterior a través de una jerarquía de niveles de ruido. Esto le permite capturar simultáneamente tanto semántica global (en ruido alto) como detalles locales (en ruido bajo).
2. Adaptación Zero-Shot: VIPaint no requiere ajustar fino el modelo de difusión preentrenado. Adapta el modelo a nuevas máscaras e imágenes optimizando parámetros variacionales para cada consulta, haciéndolo aplicable a cualquier DM preentrenado (basado en píxeles o latente).
3. Estabilidad y Modelado de Incertidumbre: Al evitar el régimen de bajo ruido durante la fase principal de optimización (donde las redes de eliminación de ruido son no suaves e inestables) y modelar explícitamente la varianza, VIPaint produce imputaciones diversas y plausibles que reflejan la incertidumbre de los datos faltantes.
4. Generalización a LDM: El método está diseñado explícitamente para funcionar con Modelos de Difusión Latente, permitiendo inpainting de alta resolución (por ejemplo, 1024x1024) con un costo computacional menor que los enfoques basados en píxeles.

4. Resultados Experimentales

Los autores evaluaron VIPaint en ImageNet64, ImageNet256, LSUN-Churches y Stable Diffusion 3.5 (condicionado por texto).

• Rendimiento Cuantitativo:

  • Máscaras Grandes: VIPaint supera significativamente a las líneas base (DPS, RePaint, CoPaint, REDdiff) en regiones enmascaradas grandes (40-80% de la imagen) en términos de PSNR, SSIM y LPIPS.
  • Modelos Condicionados por Texto: En Stable Diffusion 3.5, VIPaint logra un PSNR superior (26.69 vs 25.22 para DPS) y LPIPS (0.050 vs 0.061), produciendo imágenes globalmente coherentes que preservan detalles finos.
  • Otros Problemas Inversos: El método también muestra un rendimiento sólido en desenfoque y superresolución, superando líneas base especializadas como PSLD y ReSample.

• Rendimiento Cualitativo:

  • Los resultados visuales muestran que VIPaint genera estructuras globalmente coherentes (por ejemplo, formas de objetos consistentes e iluminación) donde otros métodos producen artefactos borrosos o incoherentes.
  • Genera con éxito múltiples completaciones diversas y plausibles para la misma imagen enmascarada, demostrando su capacidad para modelar la naturaleza multimodal de la posterior.

• Eficiencia:

  • VIPaint es computacionalmente eficiente. Para $K=2$, requiere solo ~150 llamadas a la red de eliminación de ruido para la optimización (en comparación con más de 1000 para métodos basados en muestreo) y es más rápido que muchas líneas base mientras ofrece mayor calidad.

5. Significado

VIPaint representa un avance significativo en problemas inversos con modelos de difusión. Al cambiar desde el muestreo heurístico o la optimización de un solo punto hacia un marco de inferencia variacional jerárquica principista, resuelve el problema crítico de la incoherencia global en el inpainting con máscaras grandes.

Su capacidad para funcionar zero-shot con modelos fundacionales de última generación (como Stable Diffusion 3.5) sin reentrenamiento lo convierte en una herramienta altamente práctica para tareas reales de edición, restauración y completado de imágenes. Cierra la brecha entre el rigor teórico de la inferencia variacional y las capacidades prácticas de la IA generativa moderna.