Rectified Flow

O fluxo retificado é um método de modelagem generativa que aprende um campo de velocidades que transporta amostras de uma distribuição simples, geralmente ruído gaussiano, para a distribuição dos dados por trajetórias tão retas quanto possível. O treinamento é uma regressão simples: associa-se uma amostra de ruído a uma amostra de dados, escolhe-se um ponto na linha reta entre elas e ensina-se uma rede neural a prever a direção da linha. A geração integra a velocidade aprendida com um solucionador de equações diferenciais ordinárias (ODE), a mesma ferramenta usada nas ODEs neurais, e as trajetórias retas podem ser percorridas em poucos passos grandes. Xingchao Liu, Chengyue Gong e Qiang Liu, da Universidade do Texas em Austin, apresentaram o método em setembro de 2022 no artigo Flow Straight and Fast: Learning to Generate and Transfer Data with Rectified Flow, juntamente com o método estreitamente relacionado de correspondência de fluxo, de Yaron Lipman e colegas da Meta.

Como o fluxo retificado é treinado#

Para o ruído x0 e os dados x1, o ponto x_t = (1 - t) * x0 + t * x1 vai do ruído em t = 0 aos dados em t = 1 com a velocidade constante x1 - x0. A rede v(x_t, t) é treinada para prever essa velocidade usando uma função de perda de erro quadrático sobre pares e tempos aleatórios. Muitas linhas passam por qualquer ponto, por isso a rede aprende a direção média delas; as trajetórias ODE que ela define nunca se cruzam, mas se curvam onde se encontram linhas que levam a diferentes pontos de dados. Um fluxo retificado unidimensional em PyTorch aprende a transformar ruído gaussiano em dois agrupamentos em -2 e +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")

Com um passo de Euler, 0% das amostras ficam perto de uma moda: em t = 0, a direção média a partir de qualquer ponto de ruído leva à média dos dados, que é 0, a meio caminho entre os agrupamentos. Dois passos colocam 39% das amostras a menos de 0,3 de uma moda; quatro passos, 81%; e 32 passos, 96%, porque as trajetórias aprendidas se curvam em direção aos agrupamentos e só os passos pequenos acompanham essa curva.

Reflow: endireitando as trajetórias#

Reflow, chamado de retificação recursiva no artigo, elimina essas curvas. O fluxo treinado gera um conjunto de pares (ruído, amostra), e um novo fluxo é treinado com esses pares em vez de pares aleatórios. Cada par já está ligado por uma trajetória do primeiro fluxo, por isso as linhas retas entre os pares raramente entram em conflito; o artigo prova que o novo acoplamento nunca aumenta o custo de transporte. Reflows repetidos produzem trajetórias quase retas, e os autores relataram imagens de alta qualidade e traduções de imagens com um único passo de Euler:

A mesma equipe aplicou reflow e destilação de conhecimento ao Stable Diffusion 1.5 no InstaFlow (2023). O modelo de um passo obteve FID de 23,3 no MS COCO 2017-5k, contra 37,2 da destilação progressiva, gerou uma imagem em 0,09 segundos e custou 199 dias de GPU A100 para treinar.

Fluxo retificado em modelos de texto para imagem#

O artigo sobre Stable Diffusion 3 da Stability AI, Scaling Rectified Flow Transformers for High-Resolution Image Synthesis (março de 2024), comparou formulações de treinamento e escolheu fluxo retificado com tempos de treinamento extraídos de uma distribuição logit-normal, que concentra o treinamento no meio da trajetória, onde é mais difícil prever a velocidade. O modelo combina esse objetivo com um transformer de difusão multimodal (DiT) que opera em um espaço latente. A Black Forest Labs, fundada por antigos pesquisadores da Stability AI, descreve seu modelo FLUX.1 dev como um transformer de fluxo retificado com 12 bilhões de parâmetros, e sua variante FLUX.1 [schnell] destilada gera imagens em um a quatro passos. Esses modelos de texto para imagem são uma fonte comum de dados sintéticos para o treinamento de visão. A Ultralytics Platform integra a geração de texto para imagem à curadoria de conjuntos de dados com o recurso Gerar imagens semelhantes: a partir de uma imagem do conjunto de dados, escreve novas descrições de cenas, gera até 16 variações e adiciona ao conjunto de dados as que você escolher, sem rótulos e prontas para anotação.

Fluxo retificado vs. termos relacionados#

• Correspondência de fluxo: estrutura geral para treinar geradores ODE por regressão das velocidades ao longo de trajetórias de probabilidade escolhidas. A linha reta entre ruído e dados independentes do fluxo retificado é a trajetória mais utilizada, e os dois nomes são muitas vezes usados como sinônimos.
• Modelos de difusão: adicionam ruído ao longo de um cronograma cujas trajetórias de fluxo de probabilidade são curvas; seus amostradores determinísticos normalmente usam de 20 a 50 passos.
• Fluxos normalizantes: mapeiam o ruído para os dados por meio de uma sequência fixa de camadas invertíveis treinadas com máxima verossimilhança exata.
• Modelos de consistência: aprendem a saltar de qualquer ponto de uma trajetória curva diretamente para o seu fim, permitindo que uma única chamada à rede produza uma amostra.