Flow Matching

O matching de fluxo é um método para treinar modelos generativos que transformam ruído aleatório em dados seguindo um campo de velocidades aprendido. Uma rede neural aprende, para qualquer ponto e qualquer instante entre 0 e 1, a direção e a velocidade com que esse ponto deve mover-se, e gerar uma amostra significa começar com ruído gaussiano e integrar essa velocidade com um solucionador de equações diferenciais ordinárias (ODE) até ao instante 1. Yaron Lipman e colegas da Meta AI e do Weizmann Institute of Science apresentaram o método em Flow Matching for Generative Modeling, na ICLR 2023, e dois artigos publicados em simultâneo chegaram a um objetivo quase igual, chamado fluxo retificado e interpolantes estocásticos. O matching de fluxo treina fluxos normalizantes contínuos sem simular a ODE durante o treinamento, inclui os modelos de difusão como caso especial e é o objetivo de treinamento por trás do Stable Diffusion 3, do FLUX.1, do Movie Gen da Meta e de muitos modelos recentes de fala e robótica.

Como funciona o matching de fluxo#

O treinamento nunca gera nada. Em cada passo, escolhe-se uma amostra de dados x1, uma amostra de ruído x0 de uma distribuição gaussiana padrão e um instante t entre 0 e 1, e posiciona-se um ponto na linha reta entre eles: xt = (1 - t) * x0 + t * x1. Um ponto nessa linha move-se com a velocidade constante x1 - x0, por isso a rede v(xt, t) é treinada com um simples erro quadrático médio para prever x1 - x0. Cada par de ruído e dados define apenas um caminho, mas muitos caminhos passam por qualquer ponto, e a velocidade do fluxo completo de ruído para dados nesse ponto é a média das velocidades desses caminhos. O artigo prova que regredir para os alvos simples por par produz os mesmos gradientes que regredir para essa média intratável; assim, o treinamento usa uma função de perda de regressão simples, sem simulação de ODE.

A geração integra dx/dt = v(x, t) de t = 0 a t = 1, normalmente com passos do solucionador de Euler ou de ordem superior. Este código em PyTorch treina uma pequena rede de velocidades em dois agrupamentos de pontos 2D e depois gera amostras com 1, 2 e 8 passos 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")

Ele imprime 1-step sampling: 0% of samples land near a cluster e, em seguida, mostra 47% para 2 passos e 94% para 8 passos. Em t = 0, o destino esperado de cada ponto de ruído é a média dos dados; por isso, um único passo envia todas as amostras para o espaço vazio entre os dois agrupamentos. Cada trajetória de treinamento é reta, mas as trajetórias que começam próximas umas das outras terminam em agrupamentos diferentes, e o fluxo aprendido que calcula a média delas ainda é curvo. Por isso, os geradores de poucos passos incorporam técnicas como o reflow, que retreina o modelo usando seus próprios pares de ruído e amostra para endireitar o fluxo, ou a destilação em modelos de consistência.

Matching de fluxo e métodos relacionados#

• Modelos de difusão: O treinamento por difusão corrompe os dados seguindo um esquema de ruído e ensina uma rede a prever o ruído adicionado; depois, gera amostras revertendo esse processo. O artigo sobre matching de fluxo mostrou que o seu objetivo aceita qualquer caminho de probabilidade gaussiana, incluindo os usados pelos modelos de difusão, e considerou-o uma forma mais estável de os treinar. O caminho reto é uma escolha diferente e mais simples. Ambas as famílias partilham arquiteturas como o transformador de difusão (DiT), orientação sem classificador e amostradores ODE.
• Fluxo retificado: Usa a mesma interpolação em linha reta e a mesma perda de velocidade, pelo que, para o caminho linear, os dois nomes são praticamente sinónimos. O artigo sobre fluxo retificado acrescenta o procedimento de reflow para endireitar ainda mais as trajetórias.
• Fluxos normalizantes: Os fluxos normalizantes discretos empilham camadas invertíveis e treinam com verosimilhança exata. Os fluxos normalizantes contínuos, uma forma de ODE neural, definem a transformação com uma ODE, mas eram treinados por máxima verosimilhança, o que resolve a ODE em cada passo de treinamento. O matching de fluxo treina o mesmo tipo de modelo ODE por regressão, sem essa resolução.
• Fluxo ótico: Não tem relação, apesar do nome. O fluxo ótico é o movimento aparente dos píxeis entre fotogramas de vídeo, enquanto o matching de fluxo move amostras pelo espaço dos dados possíveis.

Onde o matching de fluxo é usado#

• Geração de imagens: A Stability AI treinou o Stable Diffusion 3 com um transformador de fluxo retificado e escolheu os instantes de treinamento por amostragem de uma distribuição logit-normal que favorece níveis de ruído intermédios; o FLUX.1 da Black Forest Labs é um transformador de fluxo retificado com 12 mil milhões de parâmetros. As imagens destes modelos de texto para imagem servem de dados sintéticos que acrescentam objetos raros, condições meteorológicas e pontos de vista aos conjuntos de dados usados para treinar modelos de deteção YOLO26.
• Geração de vídeo: O Movie Gen da Meta é um transformador com 30 mil milhões de parâmetros, treinado com matching de fluxo, que gera vídeos de 1080p com até 16 segundos de duração a 16 fotogramas por segundo.
• Fala: O Voicebox da Meta treinou um modelo de matching de fluxo não autorregressivo com mais de 50.000 horas de fala e, na conversão de texto em fala zero-shot, atingiu uma taxa de erro por palavra de 1,9%, contra 5,9% do VALL-E, funcionando até 20 vezes mais depressa. Mais tarde, o F5-TTS combinou matching de fluxo com um transformador de difusão.
• Controle de robôs: O π0 da Physical Intelligence acrescenta um «especialista em ações» de matching de fluxo com 300 milhões de parâmetros a um modelo de visão e linguagem com 3 mil milhões de parâmetros, produzindo blocos de 50 ações futuras com 10 passos de integração e controlando robôs a até 50 Hz.