Free Lunch for Stabilizing Rectified Flow Inversion

1. El Problema: Inestabilidad en la Inversión de Rectified Flow

Los modelos generativos basados en Rectified Flow (RF) han emergido como alternativas de alto rendimiento a los modelos de difusión tradicionales, ofreciendo generación de alta calidad y la capacidad de realizar tareas de inversión (mapear datos observados de vuelta a su representación latente) sin necesidad de entrenamiento adicional. Sin embargo, los métodos de inversión existentes (como la inversión basada en RF "vanilla") sufren de un problema crítico:

• Errores de Aproximación Acumulativos: Durante el proceso de inversión, la latente no es accesible exactamente, por lo que se debe aproximar utilizando los datos observados en cada paso de tiempo. Esto introduce errores de estimación en cada paso.
• Inestabilidad del Campo de Velocidad: Estos errores se acumulan a lo largo de los pasos de tiempo, perturbando el campo de velocidad que debería ser casi constante en un RF ideal.
• Consecuencias: Esta inestabilidad provoca que la trayectoria de inversión se desvíe hacia regiones de baja densidad de la distribución de datos, resultando en una degradación significativa en la calidad de la reconstrucción de imágenes y en la fidelidad de las tareas de edición (como la preservación de la estructura y el contenido original).

2. Metodología Propuesta

Los autores proponen dos métodos principales, ambos libres de entrenamiento (training-free) y diseñados para corregir el campo de velocidad perturbado:

A. Inversión Proximal-Media (PMI)

Destinado a estabilizar el proceso de inversión y reconstrucción.

• Concepto Central: En lugar de seguir la predicción de velocidad pura (que contiene errores), PMI guía la velocidad actual hacia un promedio móvil (running average) de las velocidades predichas desde el inicio hasta el paso actual.
• Mecanismo de Corrección:
  1. Calcula el promedio de las velocidades históricas.
  2. Aplica una actualización de tipo proximal que empuja la velocidad actual hacia la dirección de este promedio global.
  3. Restricción Teórica: Para evitar que la trayectoria corregida se desvíe hacia regiones de baja densidad (donde la reconstrucción fallaría), se restringe cada paso de corrección dentro de una Gaussian esférica. El radio de esta esfera se deriva teóricamente basándose en un análisis de inestabilidad, asegurando que los puntos de muestreo permanezcan en regiones de alta densidad.
• Eficiencia: No requiere evaluaciones adicionales de la función neuronal (NFE) ni entrenamiento, y se integra fácilmente en cualquier solucionador ODE (Euler, Heun, RF-Solver, FireFlow).

B. Mimic-CFG (Edición)

Destinado a equilibrar la consistencia estructural y la capacidad de edición en tareas de modificación de imágenes.

• Inspiración: Se basa en la estructura de la Guía Libre de Clasificador (CFG), pero aplicada a la velocidad en lugar de la puntuación de log-verosimilitud.
• Mecanismo:
  1. Calcula el promedio de las velocidades pasadas durante la fase de edición.
  2. Proyecta la velocidad predicha actual sobre la dirección de este promedio histórico.
  3. Realiza una interpolación lineal entre la velocidad predicha original y su proyección.
• Objetivo: Esto permite mantener la estructura de la imagen original (preservación del fondo) mientras se permite suficiente desviación para lograr el cambio semántico deseado, evitando tanto la inestabilidad excesiva como la falta de editabilidad.

3. Contribuciones Clave

1. Propuesta de PMI: Un método libre de entrenamiento que corrige el campo de velocidad perturbado mediante una corrección de gradiente hacia un promedio móvil, reduciendo efectivamente los errores de inversión en modelos RF.
2. Propuesta de Mimic-CFG: Una estrategia ligera de corrección de velocidad para tareas de edición que imita la estructura de CFG, interpolando entre la velocidad actual y su proyección sobre la dirección promedio histórica.
3. Fundamentación Teórica: Derivación de un radio de corrección basado en el análisis de inestabilidad de ODEs de alta dimensión, garantizando que las trayectorias corregidas permanezcan en regiones de alta densidad de probabilidad.
4. Rendimiento SOTA: Establecimiento de un nuevo estado del arte (SOTA) en el benchmark PIE-Bench, logrando superioridad en reconstrucción y edición con menos pasos de muestreo y sin NFE adicionales.

4. Resultados Experimentales

Los experimentos se realizaron en el benchmark PIE-Bench (700 imágenes, 10 categorías de edición) utilizando el modelo Flux.1-dev como base.

• Reconstrucción e Inversión:
  • La integración de PMI con solucionadores base (Euler, Heun, RF-Solver, FireFlow) mejoró consistentemente las métricas de calidad de imagen: PSNR (relación señal-ruido) y SSIM (similitud estructural) aumentaron, mientras que MSE (error cuadrático medio) y LPIPS (similitud perceptual) disminuyeron.
  • Las mejoras fueron notables en configuraciones "sin condiciones" (unconditional), demostrando que el método estabiliza la trayectoria incluso sin prompts de texto.
• Edición de Imágenes:
  • El método mejoró significativamente la similitud con CLIP (fidelidad de la edición semántica) y la preservación del fondo (medida por PSNR/SSIM en áreas no editadas).
  • Se logró un equilibrio óptimo entre la capacidad de editar el contenido y mantener la estructura original.
• Eficiencia:
  • Los métodos propuestos permitieron alcanzar o superar el rendimiento de los métodos base con menos pasos de muestreo (NFE reducidos). Por ejemplo, FireFlow con la corrección logró resultados superiores con solo 18 pasos en comparación con variantes base que requerían más.
  • No se incurrió en costos computacionales adicionales de inferencia (no se requieren NFE extra).

5. Significado e Impacto

Este trabajo aborda una limitación fundamental en los modelos de flujo rectificado: la inestabilidad inherente en la inversión de alto rendimiento. Al proporcionar una solución teórica y práctica (PMI y Mimic-CFG) que es plug-and-play (fácil de integrar) y libre de entrenamiento, los autores demuestran que es posible estabilizar la inversión sin sacrificar la eficiencia computacional.

La importancia de este estudio radica en:

• Habilitar aplicaciones de edición y reconstrucción de imágenes más robustas y de mayor calidad.
• Reducir la brecha entre la teoría de los flujos rectificados y su aplicación práctica en tareas inversas.
• Ofrecer un marco teórico para entender y mitigar la inestabilidad en ODEs de alta dimensión, lo cual es crucial para el futuro desarrollo de modelos generativos eficientes.

En resumen, el artículo presenta un "almuerzo gratis" (free lunch) en el sentido de que obtiene mejoras sustanciales en estabilidad y calidad sin el costo de reentrenar modelos o aumentar la carga computacional de inferencia.