Resumen Técnico: SONIC (Optimización Espectral del Ruido para Inpainting con Consistencia)

Declaración del Problema

Los enfoques actuales para el inpainting de imágenes utilizando modelos de texto a imagen genéricos (específicamente modelos de difusión de eliminación de ruido y modelos de flujo) enfrentan un compromiso entre versatilidad y rendimiento. Aunque los métodos sin entrenamiento que dependen de la guía o del muestreo posterior permiten utilizar modelos genéricos para problemas inversos sin un ajuste fino costoso, a menudo producen resultados con estructuras de imagen inconsistentes en comparación con los modelos de inpainting especializados y entrenados.

Los autores argumentan que la limitación principal de los métodos existentes sin entrenamiento es la omisión de la muestra de ruido inicial. Si bien está bien establecido que el ruido inicial determina en gran medida la disposición y estructura de la generación final, optimizar directamente este ruido es computacionalmente prohibitivo. Los enfoques estándar requieren la retropropagación a través de toda la cadena de eliminación de ruido (desplegando $T$ pasos), lo cual es inviable para los modelos modernos a gran escala debido a las restricciones de memoria y cómputo. En consecuencia, los métodos existentes suelen establecerse en una inicialización de ruido subóptima, lo que conduce a desalineaciones estructurales en las regiones de inpainting.

Metodología

SONIC propone un nuevo marco de trabajo sin entrenamiento que optimiza la muestra de ruido inicial para reproducir fielmente las regiones no enmascaradas de una imagen de entrada, asegurando así la consistencia estructural en la generación final. El método evita la retropropagación a través del eliminador de ruido empleando dos innovaciones técnicas centrales:

1. Linealización de la Trayectoria de Eliminación de Ruido

Para evitar la necesidad de desplegar la cadena de eliminación de ruido, SONIC aproxima la trayectoria de eliminación de ruido como un camino lineal.

• Formulación: En lugar de tratar el proceso de eliminación de ruido $x_0 = D_T(x_T)$ como una función no lineal compleja que requiere el cálculo de gradientes a través de cada paso, los autores aproximan la trayectoria $x(t)$ como:
  $$\hat{x}(t) \approx \text{stopgrad}(D_T(x_T) - x_T)(1 - \frac{t}{T}) + x_T$$
• Mecanismo: Esta linealización asume que la diferencia entre la imagen final eliminada de ruido y el ruido inicial es constante durante todo el proceso. Esto permite que el objetivo de optimización se formule como una pérdida diferenciable con respecto a $x_T$ sin retropropagar a través del eliminador de ruido $D_T(\cdot)$. La pérdida minimiza la diferencia entre la imagen enmascarada observada $y$ y la versión corrompida de la salida aproximada linealmente.

2. Precondicionamiento Espectral

La optimización directa del ruido inicial en el dominio espacial suele conducir a la inestabilidad, donde diferentes frecuencias espaciales convergen a ritmos distintos, causando parpadeo o artefactos.

• Optimización en el Dominio Espectral: El método transforma el problema de optimización al dominio de la frecuencia. El ruido inicial $x_T$ se representa mediante su transformada de Fourier $X_T$.
• Optimización Adam: La optimización se realiza sobre $X_T$ utilizando el optimizador Adam. Dado que Adam divide las actualizaciones de los parámetros por sus segundos momentos, operar en el dominio espectral actúa efectivamente como un precondicionador. Esto permite que los diferentes componentes de frecuencia converjan a sus ritmos preferidos, estabilizando el proceso de optimización y evitando que el ruido se desvíe de la variedad aceptable.
• Restricciones Espaciales: Para evitar que las regiones no observadas (enmascaradas) del ruido se alejen de la distribución Gaussiana esperada por el modelo, las actualizaciones de los gradientes son enmascaradas. Solo se aplican los gradientes correspondientes a los píxeles observados (no enmascarados).

Detalles de Implementación

• Eficiencia en el Espacio Latente: El método opera enteramente en el espacio latente. Para imágenes enmascaradas, los autores emplean una estrategia simple de llenado por vecindad más cercana antes de codificar al espacio latente, considerando esto suficiente para un rendimiento robusto.
• Flujo de Trabajo: El proceso consiste en optimizar el ruido inicial $X_T$ durante un número fijo de iteraciones (por ejemplo, 20) utilizando la pérdida linealizada y el precondicionamiento espectral. El ruido optimizado resultante se introduce luego en un pipeline de inpainting estándar sin entrenamiento (por ejemplo, BLD-SD3.5) para generar la imagen final.

Contribuciones Clave

1. Optimización del Ruido Inicial sin Retropropagación: El artículo introduce una estrategia de linealización que permite la optimización de la muestra de ruido inicial para problemas inversos sin el costo computacional de desplegar la cadena de eliminación de ruido.
2. Precondicionamiento Espectral: Los autores demuestran que optimizar en el dominio de la frecuencia con Adam proporciona una convergencia significativamente más estable que la optimización espacial, abordando las tasas de convergencia dependientes de la frecuencia inherentes a los modelos de difusión/flujo.
3. Estado del Arte Sin Entrenamiento: El método logra un rendimiento superior comparado con los baselines de entrenamiento gratuito (FLAIR, FlowChef, FlowDPS, BLD-SD3.5) y resultados competitivos frente a modelos especializados entrenados como BrushNet.

Resultados Experimentales

Los autores evaluaron SONIC en tres conjuntos de datos estándar: FFHQ (rostros), DIV2K (imágenes naturales) y BrushBench (máscaras complejas).

• Rendimiento Cuantitativo: SONIC superó a los métodos de vanguardia sin entrenamiento (FLAIR, FlowChef, FlowDPS, BLD-SD3.5) y al modelo especializado entrenado BrushNet en métricas perceptuales clave: SSIM, LPIPS y FID.
  • En FFHQ, SONIC logró un LPIPS de 0.121 y un FID de 12.883, significativamente mejor que FLAIR (0.266 LPIPS, 18.982 FID) y BrushNet (0.237 LPIPS, 17.950 FID).
  • En DIV2K y BrushBench, se observaron tendencias similares, mostrando SONIC la mejor consistencia estructural y similitud perceptual.
• Resultados Cualitativos: Las comparaciones visuales mostraron que SONIC produce regiones de inpainting que se alinean mejor con la estructura global, la iluminación y la disposición de la imagen de entrada en comparación con los baselines, que a menudo sufrían de desalineación estructural o artefactos borrosos.
• Estudios de Ablación:
  • Eliminar el precondicionamiento espectral (optimización en el dominio espacial) condujo a un rendimiento significativamente peor (LPIPS y FID más altos), confirmando la necesidad de la estabilización en el dominio de la frecuencia.
  • Eliminar las restricciones espaciales causó que la optimización divergiera, introduciendo artefactos visibles y cambios de color.
  • El uso de píxeles de verdad fundamental (ground-truth) para el codificador (límite superior) produjo resultados casi idénticos al simple llenado por vecindad más cercana, validando la robustez de la estrategia del espacio latente.

Significado y Reivindicaciones

El artículo afirma que SONIC aborda una brecha crítica en la aplicación de modelos generativos genéricos a problemas inversos. Al identificar la muestra de ruido inicial como el "ingrediente faltante" y proporcionar un método práctico y estable para optimizarla, SONIC permite que los modelos genéricos alcancen el rendimiento previamente reservado para redes de inpainting especializadas y entrenadas.

Los autores enfatizan que su método no requiere entrenamiento, evitando los costos y los problemas de generalización asociados con el ajuste fino de modelos. Posicionan a SONIC no solo como una herramienta de inpainting, sino como una prueba de concepto de que la linealización y el precondicionamiento espectral pueden desbloquear el potencial de los modelos de flujo genéricos para una gama más amplia de problemas inversos, como la superresolución y el desenfoque (deblurring), sin la necesidad de un muestreo posterior complejo o retropropagación.