Rectified Flow

Le flux rectifié est une méthode de modélisation générative qui apprend un champ de vitesse transportant des échantillons d’une distribution simple, généralement du bruit gaussien, vers la distribution des données en suivant des trajectoires aussi droites que possible. L’entraînement est une simple régression : on associe un échantillon de bruit à un échantillon de données, on choisit un point sur la ligne droite qui les relie et on apprend à un réseau de neurones à prédire la direction de cette ligne. La génération intègre la vitesse apprise à l’aide d’un solveur d’équations différentielles ordinaires (ODE), le même mécanisme que celui des ODE neuronales ; les trajectoires droites peuvent être parcourues en quelques grandes étapes. Xingchao Liu, Chengyue Gong et Qiang Liu, de l’Université du Texas à Austin, ont présenté cette méthode en septembre 2022 dans Flow Straight and Fast: Learning to Generate and Transfer Data with Rectified Flow, en parallèle du flow matching, étroitement apparenté, proposé par Yaron Lipman et ses collègues de Meta.

Entraînement du flux rectifié#

Pour le bruit x0 et les données x1, le point x_t = (1 - t) * x0 + t * x1 va du bruit à t = 0 aux données à t = 1, avec la vitesse constante x1 - x0. Le réseau v(x_t, t) est entraîné à prédire cette vitesse avec une fonction de perte d’erreur quadratique, calculée sur des paires et des instants aléatoires. De nombreuses lignes passent par un point donné ; le réseau apprend donc leur direction moyenne, et les trajectoires ODE qu’il définit ne se croisent jamais, mais se courbent là où se rencontrent les lignes menant à différents points de données. Un flux rectifié unidimensionnel en PyTorch apprend à transformer du bruit gaussien en deux amas situés en -2 et +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")

Avec une seule étape d’Euler, 0 % des échantillons se trouvent près d’un mode : à t = 0, la direction moyenne depuis chaque point de bruit mène à la moyenne des données, 0, à mi-chemin entre les amas. Deux étapes placent 39 % des échantillons à moins de 0,3 d’un mode, quatre étapes en placent 81 % et 32 étapes, 96 %, car les trajectoires apprises se courbent vers les amas et seules de petites étapes suivent cette courbure.

Reflow : redresser les trajectoires#

Le reflow, appelé rectification récursive dans l’article, supprime ces courbures. Le flux entraîné génère un ensemble de paires (bruit, échantillon), puis un nouveau flux est entraîné sur ces paires plutôt que sur des paires aléatoires. Chaque paire est déjà reliée par une trajectoire du premier flux ; les lignes droites entre les paires entrent donc rarement en conflit, et l’article démontre que le nouveau couplage n’augmente jamais le coût de transport. Des reflows répétés produisent des trajectoires presque droites, et les auteurs ont rapporté des images et des traductions d’images de grande qualité avec une seule étape d’Euler :

La même équipe a appliqué le reflow et la distillation des connaissances à Stable Diffusion 1.5 dans InstaFlow (2023). Son modèle à une étape a obtenu un FID de 23,3 sur MS COCO 2017-5k, contre 37,2 pour la distillation progressive ; il a généré une image en 0,09 seconde et son entraînement a coûté 199 jours-GPU A100.

Flux rectifié dans les modèles texte-image#

Dans son article sur Stable Diffusion 3, Scaling Rectified Flow Transformers for High-Resolution Image Synthesis (mars 2024), Stability AI a comparé plusieurs formulations d’entraînement et choisi le flux rectifié avec des temps d’entraînement tirés d’une distribution logit-normale, qui concentre l’entraînement au milieu de la trajectoire, là où la prédiction de la vitesse est la plus difficile. Le modèle associe cet objectif à un transformer de diffusion multimodal (DiT) opérant dans un espace latent. Black Forest Labs, fondée par d’anciens chercheurs de Stability AI, décrit son modèle FLUX.1 dev comme un transformer à flux rectifié de 12 milliards de paramètres ; sa variante FLUX.1 [schnell] distillée génère des images en une à quatre étapes. Ces modèles texte-image sont une source courante de données synthétiques pour l’entraînement en vision. Ultralytics Platform intègre la génération texte-image à la curation des jeux de données avec Générer des images similaires : à partir d’une image du jeu de données, l’outil rédige de nouvelles descriptions de scène, génère jusqu’à 16 variantes et ajoute celles que tu conserves au jeu de données, sans étiquettes, prêtes pour l’annotation.

Flux rectifié et termes associés#

• Flow matching : cadre général pour entraîner des générateurs ODE par régression des vitesses le long de trajectoires de probabilité choisies. La ligne droite entre bruit et données indépendants est la trajectoire la plus utilisée du flux rectifié ; les deux noms sont souvent employés indifféremment.
• Modèles de diffusion : ajoutent du bruit selon un calendrier dont les trajectoires de flux de probabilité sont courbes ; leurs échantillonneurs déterministes nécessitent généralement de 20 à 50 étapes.
• Flux normalisants : transforment le bruit en données à travers une pile fixe de couches inversibles entraînées par vraisemblance maximale exacte.
• Modèles de cohérence : apprennent à aller directement de n’importe quel point d’une trajectoire courbe à son terme, de sorte qu’un seul appel réseau peut produire un échantillon.