Flow Matching

El ajuste de flujo es un método para entrenar modelos generativos que convierten ruido aleatorio en datos siguiendo un campo de velocidades aprendido. Una red neuronal aprende, para cualquier punto y cualquier instante entre 0 y 1, la dirección y la velocidad a las que debe desplazarse ese punto; generar una muestra consiste en partir de ruido gaussiano e integrar esa velocidad con un solucionador de ecuaciones diferenciales ordinarias (ODE) hasta el instante 1. Yaron Lipman y sus colaboradores de Meta AI y el Instituto Weizmann de Ciencias lo presentaron en Flow Matching for Generative Modeling en ICLR 2023, y dos artículos publicados al mismo tiempo llegaron a un objetivo casi idéntico con los nombres de flujo rectificado e interpolantes estocásticos. El ajuste de flujo entrena flujos normalizantes continuos sin simular la ODE durante el entrenamiento, abarca los modelos de difusión como caso particular y es el objetivo de entrenamiento de Stable Diffusion 3, FLUX.1, Movie Gen de Meta y muchos modelos recientes de voz y robótica.

Cómo funciona el ajuste de flujo#

El entrenamiento no genera nada. En cada paso se extrae una muestra de datos x1, una muestra de ruido x0 de una distribución gaussiana estándar y un instante t entre 0 y 1, y se sitúa un punto en la línea recta que los une: xt = (1 - t) * x0 + t * x1. Un punto de esa línea se mueve a velocidad constante x1 - x0, así que la red v(xt, t) se entrena con un error cuadrático medio simple para predecir x1 - x0. Cada par de ruido y datos define una sola trayectoria, pero muchas trayectorias pasan por un mismo punto y la velocidad del flujo completo de ruido a datos en ese punto es la media de todas ellas. El artículo demuestra que ajustar la regresión a los objetivos simples de cada par produce los mismos gradientes que ajustarla a esa media intratable; por tanto, el entrenamiento utiliza una función de pérdida de regresión simple sin simular ninguna ODE.

La generación integra dx/dt = v(x, t) desde t = 0 hasta t = 1, normalmente con pasos de Euler o de un solucionador de orden superior. Este código de PyTorch entrena una pequeña red de velocidades con dos grupos de puntos 2D y, después, genera muestras con 1, 2 y 8 pasos de Euler:

import torch
from torch import nn

torch.manual_seed(0)
# Velocity network v(x, t) for 2D points; the data are two clusters at (2, 2) and (-2, -2)
net = nn.Sequential(nn.Linear(3, 128), nn.SiLU(), nn.Linear(128, 128), nn.SiLU(), nn.Linear(128, 2))
optimizer = torch.optim.Adam(net.parameters(), lr=1e-3)
centers = torch.tensor([[2.0, 2.0], [-2.0, -2.0]])

for step in range(2000):
    x1 = centers[torch.randint(0, 2, (256,))] + 0.2 * torch.randn(256, 2)  # data samples
    x0 = torch.randn(256, 2)  # noise samples
    t = torch.rand(256, 1)
    xt = (1 - t) * x0 + t * x1  # point on the straight path from noise to data
    loss = ((net(torch.cat([xt, t], dim=1)) - (x1 - x0)) ** 2).mean()  # regress the path's velocity
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()

# Generate: start from noise and follow the learned velocity field with Euler steps
for steps in (1, 2, 8):
    x = torch.randn(1000, 2)
    with torch.no_grad():
        for i in range(steps):
            x = x + net(torch.cat([x, torch.full((1000, 1), i / steps)], dim=1)) / steps
    near = (torch.cdist(x, centers).min(dim=1).values < 0.6).float().mean()
    print(f"{steps}-step sampling: {near:.0%} of samples land near a cluster")

Imprime 1-step sampling: 0% of samples land near a cluster, después un 47 % con 2 pasos y un 94 % con 8 pasos. En t = 0, todos los puntos de ruido tienen el mismo destino esperado, la media de los datos, por lo que un solo paso envía todas las muestras al espacio vacío entre los dos grupos. Cada trayectoria de entrenamiento es recta, pero las trayectorias que empiezan cerca unas de otras terminan en grupos distintos, y el flujo aprendido que las promedia sigue una curva. Por eso, los generadores de pocos pasos añaden técnicas como el flujo repetido, que vuelve a entrenar con sus propios pares de ruido y muestras para enderezar el flujo, o la destilación en modelos de consistencia.

Ajuste de flujo frente a métodos relacionados#

• Modelos de difusión: El entrenamiento por difusión corrompe los datos siguiendo un programa de ruido y enseña a una red a predecir el ruido añadido; después, genera muestras invirtiendo ese proceso. El artículo sobre ajuste de flujo demostró que su objetivo admite cualquier trayectoria de probabilidad gaussiana, incluidas las que utilizan los modelos de difusión, y determinó que era una forma más estable de entrenarlos. La trayectoria recta es una opción distinta y más sencilla. Ambas familias comparten arquitecturas como el transformador de difusión (DiT), la guía sin clasificador y los muestreadores ODE.
• Flujo rectificado: Utiliza la misma interpolación en línea recta y la misma pérdida de velocidad, por lo que los dos nombres son prácticamente sinónimos para la trayectoria lineal. El artículo sobre flujo rectificado añade el procedimiento de flujo repetido para enderezar aún más las trayectorias.
• Flujos normalizantes: Los flujos normalizantes discretos apilan capas invertibles y se entrenan mediante verosimilitud exacta. Los flujos normalizantes continuos, una forma de ODE neuronal, definen la transformación mediante una ODE, pero se entrenaban por máxima verosimilitud, lo que requiere resolver la ODE en cada paso de entrenamiento. El ajuste de flujo entrena el mismo tipo de modelo ODE mediante regresión, sin resolverla.
• Flujo óptico: No guarda relación con el ajuste de flujo, pese al nombre. El flujo óptico es el movimiento aparente de los píxeles entre fotogramas de vídeo, mientras que el ajuste de flujo desplaza muestras por el espacio de datos posibles.

Usos del ajuste de flujo#

• Generación de imágenes: Stability AI entrenó Stable Diffusion 3 con un transformador de flujo rectificado y extrajo los instantes de entrenamiento de una distribución logit-normal que favorece los niveles de ruido intermedios; FLUX.1 de Black Forest Labs es un transformador de flujo rectificado de 12.000 millones de parámetros. Las imágenes de estos modelos de texto a imagen sirven como datos sintéticos que añaden objetos poco frecuentes, condiciones meteorológicas y puntos de vista a los conjuntos de datos para entrenar modelos de detección YOLO26.
• Generación de vídeo: Movie Gen de Meta es un transformador de 30.000 millones de parámetros entrenado con ajuste de flujo que genera vídeos de 1080p de hasta 16 segundos a 16 fotogramas por segundo.
• Voz: Voicebox de Meta entrenó un modelo no autorregresivo de ajuste de flujo con más de 50.000 horas de voz y, en la conversión de texto a voz de cero disparos, alcanzó una tasa de error por palabra del 1,9 %, frente al 5,9 % de VALL-E, y funcionó hasta 20 veces más rápido. Más tarde, F5-TTS combinó el ajuste de flujo con un transformador de difusión.
• Control robótico: π0 de Physical Intelligence añade un «experto en acciones» de ajuste de flujo con 300 millones de parámetros a un modelo de visión y lenguaje de 3.000 millones de parámetros. Así genera fragmentos de 50 acciones futuras con 10 pasos de integración y controla robots a una frecuencia de hasta 50 Hz.