Resumen Técnico: Mejora de Rectified Flow con Condiciones de Frontera

Planteamiento del Problema

El Flujo Rectificado (Rectified Flow, RF), también conocido como Flow Matching, ha emergido como un marco de modelado generativo poderoso que aprende un campo de velocidad para transformar una distribución de ruido en una distribución de datos objetivo mediante una Ecuación Diferencial Ordinaria (ODE). A pesar de su éxito en tareas como la generación de texto a imagen y de video, los autores identifican una limitación crítica en las implementaciones estándar de RF: el campo de velocidad aprendido a menudo no logra satisfacer las condiciones de frontera derivadas teóricamente.

Específicamente, bajo una interpolación lineal entre el ruido ($t=0$) y los datos ($t=1$), el campo de velocidad óptimo debe satisfacer:

1. Frontera Izquierda ($t=0$): $\mathbf{v}(\mathbf{x}, 0) = \mathbb{E}[\mathbf{X}_1] - \mathbf{x}$
2. Frontera Derecha ($t=1$): $\mathbf{v}(\mathbf{x}, 1) = \mathbf{x}$

Las parametrizaciones de redes neuronales estándar para el campo de velocidad no imponen inherentemente estas restricciones. Como se visualiza en el artículo, la velocidad predicha en $t=1$ a menudo se desvía del mapa identidad esperado. Esta violación es particularmente perjudicial durante el muestreo estocástico (por ejemplo, usando dinámica de Langevin). Dado que la estimación de la función de puntuación (score function) mediante la fórmula de Tweedie implica una división por $(1-t)$, una velocidad inexacta cerca de $t=1$ provoca que la estimación de la puntuación diverja. Esto conduce a un muestreo inestable, resultando en imágenes sobre-suavizadas o artefactos de tipo "caricatura", especialmente cuando se utilizan tamaños de paso o escalas de ruido mayores.

Metodología

Para abordar esto, los autores proponen el Modelo de Flujo Rectificado con Fronteras Forzadas (Boundary-enforced Rectified Flow Model). Este enfoque incorpora explícitamente las restricciones de frontera en la parametrización del campo de velocidad a través de dos variantes distintas:

1. Modelo de RF con Fronteras Basado en Máscaras (Mask-based)

Esta variante impone ambas condiciones de frontera ($t=0$ y $t=1$) estructurando el campo de velocidad como una combinación lineal de términos de frontera y la salida de una red neuronal aprendible:
$$\mathbf{v}(\mathbf{x}, t) = g(t) \cdot (C - \mathbf{x}) + f(t) \cdot \mathbf{x} + h(t) \cdot m_\theta(\mathbf{x}, t)$$
Donde:

• $m_\theta(\mathbf{x}, t)$ es una red neuronal estándar (por ejemplo, U-Net o DiT).
• $C = \mathbb{E}[\mathbf{X}_1]$.
• $f(t), g(t), h(t)$ son funciones escalares diseñadas para anularse o activarse en fronteras específicas (por ejemplo, $g(1)=f(0)=h(0)=h(1)=0$ y $g(0)=f(1)=1$).
• Los autores sugieren una elección de "Coseno Estándar": $g(t) = \cos(\frac{\pi}{2}t)$, $f(t) = \sin(\frac{\pi}{2}t)$, y $h(t) = \sin(\pi t)$.

Esta construcción garantiza que las condiciones de frontera se satisfagan por diseño, independientemente de la salida de la red neuronal.

2. Modelo de RF con Fronteras Basado en Sustracción (Subtraction-based)

Motivados por la observación de que imponer solo la condición de $t=1$ ofrece beneficios empíricos significativos, los autores proponen una variante más simple de frontera única:
$$\mathbf{v}(\mathbf{x}, t) = \mathbf{x} + m_\theta(\mathbf{x}, t) - m_\theta(\mathbf{x}, 1)$$
Esta formulación satisface inherentemente $\mathbf{v}(\mathbf{x}, 1) = \mathbf{x}$ porque el término $m_\theta(\mathbf{x}, 1)$ se cancela en $t=1$. Aunque este enfoque ofrece simplicidad de diseño y flexibilidad, incurre en un costo computacional, requiriendo una pasada hacia adelante adicional para calcular $m_\theta(\mathbf{x}, 1)$ para cada evaluación.

Impacto en el Muestreo Estocástico

Al forzar $\mathbf{v}(\mathbf{x}, 1) = \mathbf{x}$, el Modelo de RF con Fronteras asegura que el numerador en la fórmula de Tweedie ($t\mathbf{v}(\mathbf{x}, t) - \mathbf{x}$) se aproxime a cero cuando $t \to 1$. Esto evita que la función de puntuación $\nabla \log \rho_t$ diverja, estabilizando así los muestreadores estocásticos (muestreadores SDE) y permitiendo una generación robusta incluso con escalas de ruido o tamaños de paso más grandes.

Contribuciones Clave

1. Identificación de una Limitación: El artículo destaca que los campos de velocidad sin restricciones violan las condiciones de frontera teóricas en el Flujo Rectificado, lo que conduce a trayectorias de ODE inexactas y un muestreo estocástico inestable.
2. Nueva Parametrización: La introducción del Modelo de RF con Fronteras, que ofrece dos variantes prácticas (basada en máscaras y basada en sustracción) para imponer estas condiciones con una modificación mínima de código.
3. Estabilización de Funciones de Puntuación: Demostrar que el forzado de fronteras estabiliza directamente la función de puntuación cerca de $t=1$, mitigando los problemas de divergencia que afectan al muestreo estocástico en el RF vanilla.
4. Validación Empírica: Experimentos exhaustivos que muestran mejoras consistentes en la calidad de la generación y la estabilidad del muestreo a través de diferentes conjuntos de datos y escalas de modelos.

Resultados Experimentales

Los autores evaluaron sus modelos en CIFAR-10 e ImageNet (256$\times$256 y 512$\times$512) utilizando arquitecturas de Transformer de Difusión (DiT).

• Mejoras Cuantitativas:

  • En ImageNet 256$\times$256, el Modelo de RF con Fronteras basado en Sustracción logró un FID de 6.32 (muestreo ODE) y 6.38 (muestreo SDE), comparado con el 6.87 y 6.79 del baseline de RF vanilla, respectivamente.
  • El modelo basado en Máscaras logró un FID de 6.63 (ODE) y 6.18 (SDE).
  • También se observaron mejoras en el Inception Score (IS), Precisión y Recall en ambos conjuntos de datos.
  • La variante basada en Sustracción logró consistentemente los puntajes de IS más altos.

• Robustez de Muestreo:

  • En experimentos utilizando el Muestreador de Sobreimpulso (Overshooting Sampler) con variaciones en las escalas de ruido ($c$), los modelos de RF con Fronteras demostraron un rendimiento significativamente más estable y una menor degradación en la calidad de la muestra en comparación con el RF vanilla a medida que aumentaba la estocasticidad.
  • Los resultados cualitativos mostraron que los modelos de RF con Fronteras producen imágenes más nítidas con menos artefactos y mejor fidelidad de objetos, particularmente en configuraciones de muestreo estocástico donde el RF vanilla tiende a sobre-suavizar.

• Escalabilidad:

  • Los beneficios del Modelo de RF con Fronteras fueron consistentes al escalar a modelos más grandes (DiT-L, DiT-XL) y resoluciones más altas (512$\times$512).

Significado y Reivindicaciones

El artículo afirma que el Modelo de RF con Fronteras ofrece una solución simple y fácilmente implementable para una limitación fundamental en el Flujo Rectificado. Al forzar explícitamente las condiciones de frontera, el método:

1. Mejora la Calidad de Generación: Supera consistentemente al RF vanilla en FID y otras métricas.
2. Permite un Muestreo Estocástico Estable: Resuelve la divergencia de la función de puntuación cerca del tiempo terminal, permitiendo estrategias de muestreo más robustas.
3. Es Fácil de Adaptar: Puede integrarse en los modelos de RF existentes (como los que usan U-Net o DiT) con cambios arquitectónicos mínimos.

Los autores posicionan este trabajo como un refinamiento necesario para el despliegue práctico de Rectified Flow, particularmente en escenarios que requieren un muestreo estocástico de alta calidad. Sugieren que, si bien su enfoque es RF, el concepto central de forzar condiciones de frontera podría extenderse potencialmente a los modelos de difusión con adaptaciones menores. Se propone como trabajo futuro explorar el forzado de fronteras adaptativo y la aplicación a la generación de texto a video.