Flow Matching

Il flow matching è un metodo per addestrare modelli generativi che trasformano rumore casuale in dati seguendo un campo di velocità appreso. Una rete neurale apprende, per qualsiasi punto e in qualsiasi istante tra 0 e 1, la direzione e la velocità con cui quel punto dovrebbe muoversi; generare un campione significa partire da rumore gaussiano e integrare tale velocità con un risolutore di equazioni differenziali ordinarie (ODE) fino al tempo 1. Yaron Lipman e colleghi di Meta AI e del Weizmann Institute of Science l'hanno introdotto in Flow Matching for Generative Modeling a ICLR 2023; due articoli pubblicati nello stesso periodo sono giunti quasi al medesimo obiettivo con i nomi rectified flow e interpolanti stocastici. Il flow matching addestra flussi continui normalizzanti senza simulare l'ODE durante l'addestramento, include i modelli di diffusione come caso particolare ed è l'obiettivo di addestramento alla base di Stable Diffusion 3, FLUX.1, Movie Gen di Meta e molti recenti modelli per il parlato e la robotica.

Come funziona il flow matching#

L'addestramento non genera mai alcun contenuto. A ogni passaggio estrae un campione di dati x1, un campione di rumore x0 da una distribuzione gaussiana standard e un istante t tra 0 e 1, collocando un punto sulla retta che li unisce: xt = (1 - t) * x0 + t * x1. Un punto su quella retta si muove con velocità costante x1 - x0; la rete v(xt, t) viene quindi addestrata con un semplice errore quadratico medio per predire x1 - x0. Ogni coppia rumore-dati definisce un solo percorso, ma molti percorsi attraversano un punto qualsiasi, e in quel punto la velocità del flusso completo da rumore a dati ne è la media. L'articolo dimostra che eseguire la regressione sui semplici obiettivi delle singole coppie produce gli stessi gradienti della regressione su tale media intrattabile: l'addestramento consiste quindi in una semplice funzione di perdita di regressione, senza simulazione dell'ODE.

La generazione integra dx/dt = v(x, t) da t = 0 a t = 1, di solito usando passaggi del risolutore di Eulero o di ordine superiore. Questo codice PyTorch addestra una piccola rete di velocità su due gruppi di punti 2D e poi genera campioni con 1, 2 e 8 passaggi di Eulero:

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")

Stampa 1-step sampling: 0% of samples land near a cluster, poi il 47% con 2 passaggi e il 94% con 8 passaggi. In t = 0, ogni punto di rumore ha la stessa destinazione attesa, ossia la media dei dati, quindi un singolo passaggio invia tutti i campioni nello spazio vuoto tra i due cluster. Ogni traiettoria di addestramento è rettilinea, ma le traiettorie che iniziano vicine terminano in cluster diversi, e il flusso appreso che ne calcola la media rimane comunque curvo. I generatori a pochi passaggi integrano quindi tecniche come il reflow, che riaddestra il modello sulle coppie rumore-campione generate dal modello stesso per raddrizzare il flusso, oppure la distillazione in modelli di consistenza.

Flow matching e metodi correlati#

• Modelli di diffusione: l'addestramento per diffusione corrompe i dati seguendo un programma del rumore e insegna a una rete a predire il rumore aggiunto; i campioni vengono poi generati invertendo il processo. L'articolo sul flow matching ha dimostrato che il suo obiettivo accetta qualsiasi percorso di probabilità gaussiano, compresi quelli usati dai modelli di diffusione, e lo ha ritenuto un metodo di addestramento più stabile per questi modelli. Il percorso rettilineo è una scelta distinta e più semplice. Entrambe le famiglie condividono architetture come il transformer di diffusione (DiT), la guida senza classificatore e i risolutori ODE.
• Rectified flow: usa la stessa interpolazione in linea retta e la stessa perdita di velocità, perciò, nel caso del percorso lineare, i due nomi sono in gran parte sinonimi. L'articolo sul rectified flow aggiunge la procedura di reflow per raddrizzare ulteriormente le traiettorie.
• Flussi normalizzanti: i flussi normalizzanti discreti concatenano livelli invertibili e si addestrano tramite verosimiglianza esatta. I flussi normalizzanti continui, una forma di ODE neurale, definiscono la trasformazione con un'ODE ma venivano addestrati tramite massima verosimiglianza, che risolve l'ODE a ogni passaggio di addestramento. Il flow matching addestra lo stesso tipo di modello ODE tramite regressione, senza quella risoluzione.
• Flusso ottico: non è correlato, nonostante il nome. Il flusso ottico è il movimento apparente dei pixel tra i fotogrammi di un video, mentre il flow matching sposta i campioni nello spazio dei dati possibili.

Dove viene usato il flow matching#

• Generazione di immagini: Stability AI ha addestrato Stable Diffusion 3 con un transformer a flusso rettificato e ha estratto gli istanti di addestramento da una distribuzione logit-normale che favorisce livelli di rumore intermedi; FLUX.1 di Black Forest Labs è un transformer a flusso rettificato con 12 miliardi di parametri. Le immagini generate da questi modelli da testo a immagine fungono da dati sintetici, aggiungendo oggetti rari, condizioni meteorologiche e punti di vista ai dataset per addestrare i modelli di rilevamento YOLO26.
• Generazione video: Movie Gen di Meta è un transformer con 30 miliardi di parametri addestrato con il flow matching, che genera video a 1080p lunghi fino a 16 secondi a 16 fotogrammi al secondo.
• Parlato: Voicebox di Meta è stato addestrato su oltre 50.000 ore di parlato con un modello di flow matching non autoregressivo; nella sintesi vocale zero-shot ha raggiunto un tasso di errore sulle parole dell'1,9%, contro il 5,9% di VALL-E, funzionando fino a 20 volte più velocemente. In seguito, F5-TTS ha combinato il flow matching con un transformer di diffusione.
• Controllo robotico: π0 di Physical Intelligence aggiunge un «esperto di azioni» di flow matching con 300 milioni di parametri a un modello visione-linguaggio con 3 miliardi di parametri, generando blocchi di 50 azioni future con 10 passaggi d'integrazione e controllando i robot fino a 50 Hz.