Flow Matching

L’appariement de flux est une méthode d’entraînement des modèles génératifs qui transforment du bruit aléatoire en données en suivant un champ de vitesse appris. Un réseau neuronal apprend, pour tout point et tout instant entre 0 et 1, la direction et la vitesse auxquelles le point doit se déplacer. Pour générer un échantillon, il faut partir d’un bruit gaussien et intégrer cette vitesse à l’aide d’un solveur d’équation différentielle ordinaire (EDO) jusqu’à l’instant 1. Yaron Lipman et ses collègues de Meta AI et de l’Institut Weizmann des sciences ont présenté cette méthode dans Flow Matching for Generative Modeling, à ICLR 2023. Deux articles publiés simultanément sont parvenus à un objectif presque identique sous les noms de flux rectifié et d’interpolants stochastiques. L’appariement de flux entraîne des flux normalisants continus sans simuler l’EDO pendant l’entraînement, inclut les modèles de diffusion comme cas particulier et constitue l’objectif d’entraînement de Stable Diffusion 3, FLUX.1, Movie Gen de Meta et de nombreux modèles récents de parole et de robotique.

Fonctionnement de l’appariement de flux#

L’entraînement ne génère jamais rien. À chaque étape, on tire un échantillon de données x1, un échantillon de bruit x0 d’une loi gaussienne standard et un instant t entre 0 et 1, puis on place un point sur la ligne droite qui les relie : xt = (1 - t) * x0 + t * x1. Un point sur cette ligne se déplace à la vitesse constante x1 - x0. Le réseau v(xt, t) est donc entraîné avec une simple erreur quadratique moyenne pour prédire x1 - x0. Chaque paire bruit-données définit un seul chemin, mais plusieurs chemins passent par un point donné ; la vitesse du flux complet du bruit vers les données à ce point correspond à leur moyenne. L’article prouve que la régression sur les cibles simples de chaque paire produit les mêmes gradients qu’une régression sur cette moyenne impossible à calculer. L’entraînement consiste donc en une simple fonction de perte de régression, sans simulation d’EDO.

La génération intègre dx/dt = v(x, t) de t = 0 à t = 1, généralement par des étapes du solveur d’Euler ou d’ordre supérieur. Ce code PyTorch entraîne un petit réseau de vitesse sur deux groupes de points 2D, puis génère des échantillons avec 1, 2 et 8 étapes d’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")

Le code affiche 1-step sampling: 0% of samples land near a cluster, puis 47 % pour 2 étapes et 94 % pour 8 étapes. À t = 0, chaque point de bruit a la même destination attendue, la moyenne des données ; une seule étape envoie donc tous les échantillons dans le vide entre les deux groupes. Chaque chemin d’entraînement est rectiligne, mais des chemins qui commencent à proximité se terminent dans des groupes différents, et le flux appris qui en fait la moyenne reste incurvé. Les générateurs à peu d’étapes ajoutent donc des techniques comme le reflow, qui réentraîne le modèle sur ses propres paires bruit-échantillon pour redresser le flux, ou la distillation en modèles de cohérence.

L’appariement de flux par rapport aux méthodes associées#

• Modèles de diffusion : L’entraînement par diffusion corrompt les données selon un calendrier de bruit et apprend à un réseau à prédire le bruit ajouté ; l’échantillonnage inverse ensuite ce processus. L’article sur l’appariement de flux a montré que son objectif accepte n’importe quel chemin de probabilité gaussien, y compris ceux utilisés par les modèles de diffusion, et qu’il constitue une méthode plus stable pour les entraîner. Le chemin rectiligne est un choix distinct et plus simple. Les deux familles partagent des architectures comme le transformeur de diffusion (DiT), le guidage sans classifieur et les solveurs d’EDO.
• Flux rectifié : Utilise la même interpolation en ligne droite et la même perte de vitesse ; les deux noms sont donc largement synonymes pour le chemin linéaire. L’article sur le flux rectifié ajoute la procédure de reflow pour redresser davantage les trajectoires.
• Flux normalisants : Les flux normalisants discrets empilent des couches inversibles et s’entraînent par vraisemblance exacte. Les flux normalisants continus, une forme d’EDO neuronale, définissent la transformation par une EDO, mais étaient entraînés par maximum de vraisemblance, ce qui résout l’EDO à chaque étape d’entraînement. L’appariement de flux entraîne le même type de modèle d’EDO par régression, sans cette résolution.
• Flux optique : Sans rapport malgré son nom. Le flux optique désigne le mouvement apparent des pixels entre les images vidéo, tandis que l’appariement de flux déplace des échantillons dans l’espace des données possibles.

Utilisations de l’appariement de flux#

• Génération d’images : Stability AI a entraîné Stable Diffusion 3 avec un transformeur à flux rectifié et échantillonné les pas de temps d’entraînement selon une distribution logit-normale qui privilégie les niveaux de bruit intermédiaires. FLUX.1 de Black Forest Labs est un transformeur à flux rectifié de 12 milliards de paramètres. Les images produites par ces modèles texte-image servent de données synthétiques pour ajouter des objets rares, des conditions météo et des points de vue aux jeux de données d’entraînement des modèles de détection YOLO26.
• Génération vidéo : Movie Gen de Meta est un transformeur de 30 milliards de paramètres entraîné par appariement de flux, qui génère des vidéos 1080p d’une durée maximale de 16 secondes à 16 images par seconde.
• Parole : Voicebox de Meta a entraîné un modèle non autorégressif d’appariement de flux sur plus de 50 000 heures de parole, et, pour la synthèse vocale en mode zéro-shot, il a atteint un taux d’erreur sur les mots de 1,9 %, contre 5,9 % pour VALL-E, tout en fonctionnant jusqu’à 20 fois plus vite. F5-TTS a ensuite associé l’appariement de flux à un Transformer de diffusion.
• Commande robotique : π0 de Physical Intelligence ajoute un « expert en actions » fondé sur l’appariement de flux et comptant 300 millions de paramètres à un modèle vision-langage de 3 milliards de paramètres. Il produit des séquences de 50 actions futures en 10 étapes d’intégration et commande des robots jusqu’à 50 Hz.