Rectified Flow

El flujo rectificado es un método de modelado generativo que aprende un campo de velocidades que transporta muestras de una distribución sencilla, normalmente ruido gaussiano, a la distribución de datos por trayectorias lo más rectas posible. El entrenamiento consiste en una regresión sencilla: se empareja una muestra de ruido con una muestra de datos, se elige un punto de la línea recta que las une y se enseña a una red neuronal a predecir la dirección de la línea. La generación integra la velocidad aprendida con un solucionador de ecuaciones diferenciales ordinarias (ODE), el mismo mecanismo que se utiliza en las ODE neuronales, y las trayectorias rectas se pueden recorrer en unos pocos pasos grandes. Xingchao Liu, Chengyue Gong y Qiang Liu, de la Universidad de Texas en Austin, presentaron el método en septiembre de 2022 en Flow Straight and Fast: Learning to Generate and Transfer Data with Rectified Flow, junto con el método estrechamente relacionado de emparejamiento de flujos, de Yaron Lipman y sus colaboradores de Meta.

Cómo se entrena el flujo rectificado#

Para el ruido x0 y los datos x1, el punto x_t = (1 - t) * x0 + t * x1 pasa del ruido en t = 0 a los datos en t = 1 con la velocidad constante x1 - x0. La red v(x_t, t) se entrena para predecir esa velocidad mediante una función de pérdida de error cuadrático, con pares y tiempos aleatorios. Por un mismo punto pasan muchas líneas, así que la red aprende su dirección media; las trayectorias de ODE que define no se cruzan entre sí, pero se curvan donde confluyen líneas que se dirigen a distintos puntos de datos. Un flujo rectificado unidimensional en PyTorch aprende a convertir ruido gaussiano en dos grupos, en -2 y +2:

import torch
from torch import nn

torch.manual_seed(0)
velocity = nn.Sequential(nn.Linear(2, 64), nn.SiLU(), nn.Linear(64, 64), nn.SiLU(), nn.Linear(64, 1))
optimizer = torch.optim.Adam(velocity.parameters(), lr=1e-3)

for step in range(3000):
    x0 = torch.randn(256, 1)  # noise
    x1 = 2.0 * torch.sign(torch.randn(256, 1)) + 0.1 * torch.randn(256, 1)  # data: two modes at -2 and +2
    t = torch.rand(256, 1)
    xt = (1 - t) * x0 + t * x1  # a point on the straight line from noise to data
    loss = ((velocity(torch.cat([xt, t], dim=1)) - (x1 - x0)) ** 2).mean()  # regress the line's direction
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()

# Generate by Euler integration of dx/dt = v(x, t), with different step counts
for steps in [1, 2, 4, 32]:
    x = torch.randn(10000, 1)
    with torch.no_grad():
        for i in range(steps):
            x = x + velocity(torch.cat([x, torch.full_like(x, i / steps)], dim=1)) / steps
    near_mode = ((x.abs() - 2).abs() < 0.3).float().mean()
    print(f"{steps:>2} Euler steps: {near_mode:.0%} of samples within 0.3 of a mode")

Con un paso de Euler, el 0 % de las muestras cae cerca de un modo: en t = 0, la dirección media desde cada punto de ruido conduce a la media de los datos, 0, a medio camino entre los grupos. Con dos pasos, el 39 % de las muestras queda a menos de 0,3 de un modo; con cuatro pasos, el 81 %, y con 32 pasos, el 96 %, porque las trayectorias aprendidas se curvan hacia los grupos y solo los pasos pequeños siguen esa curva.

Reflow: cómo enderezar las trayectorias#

Reflow, llamado rectificación recursiva en el artículo, elimina esas curvas. El flujo entrenado genera un conjunto de pares (ruido, muestra), y un nuevo flujo se entrena con esos pares en lugar de con pares aleatorios. Cada par ya está unido por una trayectoria del primer flujo, así que las líneas rectas entre pares rara vez entran en conflicto, y el artículo demuestra que el nuevo acoplamiento nunca aumenta el coste de transporte. Las iteraciones de reflow producen trayectorias casi rectas, y los autores obtuvieron imágenes y traducciones de imágenes de alta calidad con un único paso de Euler:

El mismo equipo aplicó reflow y destilación de conocimiento a Stable Diffusion 1.5 en InstaFlow (2023). Su modelo de un solo paso obtuvo un FID de 23,3 en MS COCO 2017-5k, frente a 37,2 de la destilación progresiva, generó una imagen en 0,09 segundos y necesitó 199 días de GPU A100 para entrenarse.

Flujo rectificado en modelos de texto a imagen#

El artículo de Stable Diffusion 3 de Stability AI, «Escalado de transformadores de flujo rectificado para la síntesis de imágenes de alta resolución» (marzo de 2024), comparó distintas formulaciones de entrenamiento y optó por el flujo rectificado con tiempos de entrenamiento extraídos de una distribución logit-normal, que concentra el entrenamiento en la parte central de la trayectoria, donde es más difícil predecir la velocidad. El modelo combina ese objetivo con un transformador de difusión multimodal (DiT) que opera en un espacio latente. Black Forest Labs, fundada por antiguos investigadores de Stability AI, describe su modelo FLUX.1 dev como un transformador de flujo rectificado de 12 000 millones de parámetros, y su variante destilada FLUX.1 [schnell] genera en uno a cuatro pasos. Estos modelos de texto a imagen son una fuente habitual de datos sintéticos para entrenar sistemas de visión artificial. Ultralytics Platform integra la generación de texto a imagen en la curación de conjuntos de datos con Generar imágenes similares: a partir de una imagen del conjunto de datos, redacta nuevas descripciones de escenas, genera hasta 16 variaciones y añade al conjunto las que conserves, sin etiquetas y listas para la anotación.

Flujo rectificado frente a términos relacionados#

• Emparejamiento de flujos: Marco general para entrenar generadores de ODE mediante regresión de velocidades a lo largo de trayectorias de probabilidad elegidas. La línea recta entre ruido y datos independientes del flujo rectificado es la trayectoria más utilizada, y ambos términos suelen usarse indistintamente.
• Modelos de difusión: Añaden ruido según un programa cuya trayectoria de flujo de probabilidad se curva, y sus muestreadores deterministas suelen requerir de 20 a 50 pasos.
• Flujos normalizadores: Transforman el ruido en datos mediante una secuencia fija de capas invertibles entrenadas con máxima verosimilitud exacta.
• Modelos de consistencia: Aprenden a saltar desde cualquier punto de una trayectoria curva directamente hasta el final, de modo que una sola llamada a la red puede generar una muestra.