Resumen Técnico: Acoplamientos Supresores de Divergencia para Flujo Rectificado

Enunciado del Problema

El Flujo Rectificado (RectFlow) busca acelerar el modelado generativo aprendiendo campos de velocidad con trayectorias rectas, permitiendo la generación de muestras de alta calidad con pocos o incluso un solo paso de integración numérica (NFE). El método se basa en un procedimiento iterativo de "rectificación": un modelo base entrenado con un acoplamiento independiente genera nuevos pares fuente-objetivo, que luego se utilizan para reentrenar un modelo con un campo de velocidad más recto.

Sin embargo, el artículo identifica un cuello de botella fundamental: la distorsión del acoplamiento. En la práctica, las trayectorias generadas por el modelo base a menudo se doblan y entrelazan. Esta distorsión se asocia frecuentemente con regiones de divergencia no nula en el campo de velocidad aprendido. En estas regiones, la expansión o contracción local hace que las trayectorias se separen o compriman, desviando las partículas de sus puntos finales ideales. En consecuencia, el modelo reentrenado hereda estos errores, y rondas adicionales de rectificación producen rendimientos decrecientes porque el modelo está efectivamente "rectificando" un transporte que ya está doblado. La calidad del integrador hacia adelante utilizado para generar los pares de acoplamiento, y no la función de pérdida de entrenamiento ni la capacidad del modelo, limita la eficacia de RectFlow.

Metodología: DS-RectFlow

Los autores proponen DS-RectFlow (Flujo Rectificado Supresor de Divergencia), un método de corrección offline diseñado para generar pares de acoplamiento más limpios durante el proceso de rectificación sin alterar el objetivo de entrenamiento, la arquitectura del modelo ni el procedimiento de inferencia.

Mecanismo Central

El método se basa en la observación de que el doblamiento de trayectorias (distorsión del acoplamiento) es impulsado por el componente compresible del campo de velocidad. Utilizando la descomposición de Helmholtz, el campo de velocidad $v$ se divide en una parte de transporte sin divergencia $u$ y una parte dipolar $\nabla \phi$ que lleva toda la compresibilidad del campo. El artículo argumenta que suprimir la magnitud de la divergencia $|\nabla \cdot v|$ mantiene las trayectorias rectas.

El Integrador Supresor de Divergencia

Durante la generación offline de pares de acoplamiento (el paso de "rectificación"), DS-RectFlow reemplaza el integrador de Euler estándar con un integrador de Euler supresor de divergencia:

1. Corrección de Estado: En cada paso de integración, en lugar de avanzar directamente desde el estado actual $x_t$, el algoritmo busca un estado cercano $x^*_t$ donde la divergencia estimada $|\hat{\nabla} \cdot v_\theta(x^*_t, t)|$ esté minimizada.
2. Estimación de Hutchinson: Para evitar el alto costo del cálculo exacto de la divergencia en altas dimensiones, el método utiliza el estimador de traza de Hutchinson para aproximar $\nabla \cdot v$. Genera $m$ perturbaciones gaussianas alrededor del estado actual y selecciona el candidato que minimiza la magnitud de la divergencia estimada.
3. Integración: El paso de Euler se calcula luego desde este estado corregido $x^*_t$ utilizando el campo de velocidad original, sin modificar: $x_{t+\Delta t} = x^*_t + v_\theta(x^*_t, t) \cdot \Delta t$.

Detalles de Implementación

• Solo Offline: La corrección se aplica únicamente durante la generación del conjunto de datos de acoplamiento para el reentrenamiento.
• Inferencia: El modelo final entrenado se despliega utilizando integración de Euler estándar sin sobrecarga computacional adicional. La lógica supresora de divergencia no está presente en el momento de la inferencia.
• Control del Presupuesto: La corrección se aplica solo hasta un tiempo $t_{stop}$ (típicamente 0.5) y está gobernada por un parámetro de radio de búsqueda $\delta$.

Contribuciones Clave

1. Diagnóstico de las Limitaciones de RectFlow: El artículo identifica que la calidad del acoplamiento generado por el modelo base, específicamente la presencia de divergencia no nula que causa compresión/expansión de trayectorias, es el principal cuello de botella que impide que RectFlow alcance su potencial teórico de generación en un solo paso.
2. Estrategia de Corrección Offline: Introduce una corrección offline novedosa que atenúa el componente divergente del campo de velocidad durante la generación de acoplamientos. A diferencia de los refinamientos en tiempo de inferencia (por ejemplo, solucionadores de orden superior o guía), este método traslada el costo de la corrección completamente a la fase de generación de datos de entrenamiento.
3. Cero Sobrecarga en Inferencia: El método produce un modelo que se ejecuta con un costo de reloj idéntico al de RectFlow estándar durante el despliegue, ya que la ruta de inferencia permanece como un simple paso de Euler.
4. Restauración de Efectos Compuestos: Al proporcionar pares de acoplamiento más limpios, DS-RectFlow restaura el beneficio "compuesto" de la rectificación, donde cada ronda de rectificación mejora significativamente la rectitud del campo, un fenómeno que a menudo se pierde en RectFlow estándar debido a la distorsión acumulada.

Resultados Experimentales

Los autores evalúan DS-RectFlow en benchmarks sintéticos 2D y tareas de generación de imágenes (CIFAR-10, CelebA-64).

• Benchmarks 2D (Tablero de Ajedrez y GMM):

  • En el benchmark de Tablero de Ajedrez, DS-RectFlow reduce significativamente la Distancia de Wasserstein Recortada (SWD) en comparación con RectFlow estándar. Con $NFE=1$, alcanza una SWD de 0.102 (frente a 0.165 para la versión estándar), igualando efectivamente el rendimiento de una línea base de Flujo de Emparejamiento de 20 pasos pero en un solo paso.
  • En el benchmark de cruce de GMM, DS-RectFlow muestra inicialmente un error mayor con NFE bajos (debido a capturar con mayor fidelidad decisiones complejas de asignación de modos), pero converge rápidamente a un rendimiento superior. Para la segunda ronda de rectificación ($k=2$), logra una mejora de 3x en SWD con $NFE=1$ en comparación con RectFlow estándar.
  • Las visualizaciones confirman que DS-RectFlow produce modos más nítidos y separados con menos muestras "filtradas" hacia regiones prohibidas.

• Generación de Imágenes (CIFAR-10 y CelebA-64):

  • CIFAR-10: DS-RectFlow logra consistentemente una menor Distancia Fréchet Inception (FID) en todos los presupuestos de NFE. Con un solucionador offline Euler-20, el método reduce el FID de 17.41 a 12.03 con $NFE=1$ (una mejora relativa del 31%). Notablemente, el modelo DS-RectFlow con $NFE=1$ supera al modelo estándar en todos los pasos de NFE en la exploración.
  • CelebA-64: Se observan tendencias similares. Con $NFE=2$, DS-RectFlow alcanza un FID de 4.90, superando el mejor rendimiento del modelo estándar en todos los NFE (que es 5.64 con $NFE=20$).
  • Los resultados demuestran que la corrección offline se transfiere directamente a un modelo capaz de generación de alta calidad en un solo paso.

Significado y Afirmaciones

El artículo afirma que la calidad del acoplamiento, y no la capacidad de entrenamiento ni el presupuesto de inferencia, es el principal cuello de botella de RectFlow. Al abordar la distorsión inducida por la divergencia de forma offline, DS-RectFlow permite las mejoras compuestas que promete la rectificación pero que rara vez se materializan en la práctica.

Los autores enfatizan que este enfoque ofrece una ventaja práctica: proporciona los beneficios de un integrador corregido por divergencia (trayectorias más rectas, mejor alineación de modos) mientras mantiene el perfil de inferencia de latencia cero de RectFlow estándar. La corrección es un costo único pagado durante la generación de datos, amortizado sobre toda la ejecución de entrenamiento. El artículo sugiere que este mecanismo puede ser aplicable a otros métodos basados en flujos, como interpolantes estocásticos y emparejamiento de flujos de transporte óptimo, siempre que componentes de deriva divergentes degraden la calidad de las muestras.