Timestep Rescheduling in Diffusion Inversion

Resumen Técnico: Reprogramación de Pasos de Tiempo en la Inversión de Difusión

Planteamiento del Problema
La inversión de difusión es el proceso de mapear una imagen real $z_0$ de vuelta al espacio latente gaussiano ($z_T$) de un modelo de difusión, permitiendo tareas como la reconstrucción y edición de imágenes. Aunque los Modelos de Difusión Implícitos de Denoising (DDIM) proporcionan una formulación determinista y eficiente para esto, el proceso de inversión es inherentemente aproximado. El enfoque estándar consiste en revertir analíticamente el paso de eliminación de ruido, lo que introduce desviaciones que se acumulan con el tiempo, provocando errores de reconstrucción notables y una degradación en la fidelidad de la edición.

Los métodos existentes intentan mitigar estos errores resolviendo problemas de punto fijo para refinar la ecuación de inversión. Sin embargo, estos enfoques pasan por alto en gran medida un factor crítico: la influencia de la programación de los pasos de tiempo (la selección y el espaciamiento de los pasos de difusión) en la fidelidad de la inversión. La mayoría de las canalizaciones utilizan pasos de tiempo distribuidos uniformemente, lo cual, según los autores, es subóptimo. El artículo postula que la distribución de los pasos de tiempo impacta significativamente en la acumulación de errores globales de inversión, aunque esta relación permanece poco explorada en comparación con su papel en el muestreo o el entrenamiento de la difusión.

Metodología: TRDI (Reprogramación de Pasos de Tiempo en la Inversión de Difusión)
Los autores proponen TRDI, un programador de pasos de tiempo no uniforme diseñado para minimizar el error de inversión global sin requerir parámetros adicionales ni sobrecarga computacional durante la etapa de inferencia. La metodología se fundamenta en un análisis teórico del error de inversión y se implementa mediante una estrategia de dos etapas de "grueso a fino" (coarse-to-fine).

1. Perspectiva Teórica:
   Los autores derivan que el error de inversión inducido por un salto de paso de tiempo $\Delta t$ puede reformularse como un problema de punto fijo escalado. La magnitud del error $\delta$ es proporcional a un coeficiente de escala $c_\alpha(t, \Delta t)$ multiplicado por la discrepancia entre las predicciones consecutivas del modelo ($\Delta \epsilon_\theta$).

   • Observación Clave: El coeficiente de escala exhibe una tendencia parabólica respecto al índice del paso de tiempo. Los errores se concentran tanto en los pasos de tiempo pequeños (pasos de difusión tempranos) como en los grandes (pasos tardíos), con un descenso en el medio. Además, los tamaños de paso más grandes ($\Delta t$) inducen inherentemente errores mayores.

2. Reescalado Global (Ajuste Grueso):
   Para abordar la distribución parabólica del error, los autores introducen un mecanismo de reescalado global. Modifican la secuencia original de pasos de tiempo uniformes $t_k$ utilizando una transformación de ley de potencia controlada por un hiperparámetro $\gamma$:
   $$t_k = t_1 + (t_K - t_1) \left( \frac{k-1}{K-1} \right)^\gamma$$

   • Cuando $\gamma > 1$, los pasos de tiempo se estiran hacia los pasos iniciales (asignando más resolución a la región de alto error al principio).
   • Cuando $\gamma < 1$, los pasos de tiempo se vuelven más densos hacia el final.
     Este paso prepara la programación para una optimización de grano fino mediante el ajuste global de la asignación de pasos.

3. Reprogramación Local mediante Programación Dinámica (Ajuste Fino):
   Para minimizar aún más el error acumulado, los autores emplean un enfoque de programación dinámica (DP) de ventana deslizante.

   • Mapa de Coste: Se inicializa un mapa de coste $E[k, t]$ para rastrear el error acumulado mínimo al alcanzar el índice de paso $t$ en el $k$-ésimo paso de inversión.
   • Recursión: Dentro de una ventana de ancho $2d+1$ alrededor del paso actual, el algoritmo busca recursivamente la transición localmente óptima entre pasos de tiempo adyacentes que minimice la suma del error acumulado previo y el coeficiente de error del paso actual.
   • Retroceso (Backtracking): La secuencia óptima de pasos de tiempo se recupera mediante un retroceso desde el paso final hasta el primero.

Todo el proceso de reprogramación se realiza fuera de línea (o como un paso de preprocesamiento) basado en la programación de ruido conocida, lo que significa que añade una sobrecarga de tiempo de ejecución insignificante al proceso real de inversión o edición de la imagen.

Contribuciones Clave

• Formulación Teórica: El artículo deriva una conexión teórica entre los saltos de tiempo y el error de inversión, enmarcando el error como un problema de punto fijo escalado dependiente de la programación de ruido.
• Programador Adaptativo: Propone un marco de reprogramación de grueso a fino (TRDI) que asigna adaptativamente pasos más grandes a regiones de bajo error y pasos más finos a regiones de alto error, equilibrando eficazmente la eficiencia y la precisión.
• Mejora Plug-and-Play: El método está diseñado como una mejora "lista para usar" (off-the-shelf) que no requiere parámetros adicionales, entrenamiento ni una sobrecarga computacional significativa, siendo compatible con las canalizaciones de inversión existentes.

Resultados Experimentales
Los autores validan TRDI mediante experimentos extensos en reconstrucción de imágenes (conjunto de datos MSCOCO) y edición de imágenes (PIE-Bench) utilizando varios modelos de difusión (Stable Diffusion v1.5, SDXL y SDXL Turbo) y métodos de inversión de referencia (DDIM, ReNoise, NPI, GNRI).

• Reconstrucción: En SD v1.5, la integración de TRDI mejoró consistentemente las métricas en todas las bases de comparación. Por ejemplo, con la base DDIM, el PSNR aumentó de 20.07 a 20.21, y el LPIPS disminuyó (mejorando la calidad perceptual) de 193.97 a 187.85. Los resultados cualitativos mostraron correcciones en detalles estructurales como ventanas de trenes y formas de frutas que las bases de comparación no lograron reconstruir con precisión.
• Edición: En el PIE-Bench, TRDI mejoró significativamente la preservación de la estructura y la preservación del fondo, manteniendo una alta similitud CLIP. Notablemente, para modelos de pocos pasos como SDXL Turbo (4 pasos), donde la degradación del fondo es típicamente severa, TRDI logró mejoras sustanciales en la preservación del fondo (por ejemplo, reduciendo la Distancia Estructural en un 17.43% para DDIM).
• Estudios de Ablación: El estudio confirmó que un factor de escala $\gamma \approx 1.05$ y un ancho de ventana $d=8$ proporcionaron el mejor equilibrio entre fidelidad estructural y editabilidad semántica.

Significado y Reivindicaciones
El artículo sostiene que la programación de los pasos de tiempo es un factor crítico, aunque previamente pasado por alto, en la fidelidad de la inversión de difusión. Al revelar la relación parabólica entre el tamaño/índice del paso de tiempo y la acumulación de error, los autores demuestran que una asignación estratégica y no uniforme de los pasos computacionales puede reducir significamente los errores globales de inversión.

Los autores posicionan a TRDI no como un reemplazo para formulaciones de inversión más fuertes (como las iteraciones de punto fijo o la inversión estocástica), sino como un complemento de nivel de programador ligero y complementario. Permite que las canalizaciones de inversión deterministas existentes logren resultados de reconstrucción y edición superiores sin aumentar los costes computacionales ni requerir el reentrenamiento del modelo. El trabajo enfatiza que optimizar cuándo tomar los pasos es tan importante como cómo tomarlos.