Flow Matching : Guide complet — Génération par Correspondance de Flux

Résumé — Le Flow Matching, introduit par Lipman et al. en 2022, est une méthode de génération qui apprend un champ de vecteurs neural définissant une trajectoire ODE continue transformant une distribution de bruit simple gaussienne en la distribution des données. Contrairement aux modèles de diffusion basés sur des SDEs stochastiques, le Flow Matching est déterministe. Rectified Flow (Liu et al., 2022) améliore cette approche avec un ré-entraînement itératif produisant des trajectoires plus droites et donc une génération plus rapide avec seulement 10-20 steps d’intégration.

Principe mathématique

1. Formulation ODE

Le Flow Matching définit une équation différentielle ordinaire qui décrit comment transformer du bruit en données :

dx/dt = v_theta(x, t)

On intègre cette ODE de t=0 où x_0 suit une gaussienne N(0,I) à t=1 où x_1 suit la distribution des données p_data. La génération d’échantillons se fait simplement en intégrant l’ODE en avant à partir d’un bruit initial aléatoire.

2. Conditional Flow Matching

Au lieu d’apprendre le champ marginal directement, on définit des chemins conditionnels x_t reliant chaque paire (x_0, x_1) :

x_t = t * x_1 + (1-t) * x_0

C’est une interpolation linéaire entre le bruit et la donnée. La vitesse cible est simplement x_1 – x_0. La loss est une régression quadratique :

L_FM = E[||v_theta(x_t,t) - (x_1 - x_0)||^2]

3. Rectified Flow

Liu et al. montrent qu’en ré-entraînant le modèle avec les paires générées, les trajectoires deviennent plus droites. Après 2-3 itérations, 10-20 steps d’Euler suffisent pour une bonne qualité.

4. Comparaison avec la diffusion

Les modèles de diffusion utilisent des SDEs stochastiques avec bruitage progressif puis débruitage appris et nécessitent 50-1000 steps. Le Flow Matching est déterministe ODE, plus rapide, reproductible et ne nécessite pas de planification de bruit (noise schedule).

Intuition

Imaginez du brouillard d’un côté d’une vallée et un jardin fleuri de l’autre côté. La diffusion ajoute du brouillard progressivement puis essaie de le retirer à l’aveugle, comme quelqu’un qui marche dans le brouillard avec une boussole. Le Flow Matching apprend directement les courants de vent qui transportent chaque goutte de brouillard vers la bonne fleur. Le chemin est continu, réversible et direct.

C’est comme la différence entre une route goudronnée et un sentier forestier sinueux : plus direct, plus prévisible, moins d’étapes nécessaires pour arriver à destination.

Implémentation Python

import torch
import torch.nn as nn
import torch.nn.functional as F


class VelocityField(nn.Module):
    def __init__(self, data_dim, hidden_dim=256):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(data_dim + 1, hidden_dim), nn.SiLU(),
            nn.Linear(hidden_dim, hidden_dim), nn.SiLU(),
            nn.Linear(hidden_dim, hidden_dim), nn.SiLU(),
            nn.Linear(hidden_dim, data_dim)
        )

    def forward(self, x, t):
        t_exp = t.view(-1, 1).expand(x.size(0), 1)
        return self.net(torch.cat([x, t_exp], dim=1))


class FlowMatcher:
    def __init__(self, model, lr=1e-3):
        self.model = model
        self.optimizer = torch.optim.AdamW(model.parameters(), lr=lr)

    def train_step(self, x_0, x_1):
        batch_size = x_0.size(0)
        t = torch.rand(batch_size, 1)
        x_t = t * x_1 + (1 - t) * x_0
        v_target = x_1 - x_0
        self.optimizer.zero_grad()
        v_pred = self.model(x_t, t.squeeze())
        loss = F.mse_loss(v_pred, v_target)
        loss.backward()
        self.optimizer.step()
        return loss.item()

    @torch.no_grad()
    def sample(self, n, data_dim, steps=50):
        x = torch.randn(n, data_dim)
        dt = 1.0 / steps
        for i in range(steps):
            t = torch.full((n,), i * dt)
            v = self.model(x, t)
            x = x + v * dt
        return x


# Entraînement sur des cercles 2D
def make_circle(n=5000):
    r = torch.randn(n, 1) * 0.3 + 1.0
    theta = torch.rand(n, 1) * 2 * 3.14159
    return torch.cat([r * torch.cos(theta), r * torch.sin(theta)], dim=1)

data = make_circle()
model = VelocityField(data_dim=2, hidden_dim=256)
fm = FlowMatcher(model, lr=1e-3)

for epoch in range(200):
    idx = torch.randint(0, data.size(0), (128,))
    x_1 = data[idx]
    x_0 = torch.randn_like(x_1)
    loss = fm.train_step(x_0, x_1)
    if epoch % 50 == 0:
        samples = fm.sample(50, 2, steps=50)
        print(f'Epoch {epoch} | Loss: {loss:.4f}')

Hyperparamètres

Avantages

1. Génération rapide : 10-20 steps avec Rectified Flow contre 50-1000+ pour la diffusion, soit 5-100x plus rapide.
2. Déterministe : Processus reproductible (même bruit initial = même résultat).
3. Réversible : L’ODE peut être intégrée en sens inverse pour encoder des données en bruit latent.
4. Simple : Loss MSE plus simple à implémenter que le planning de bruit de la diffusion.

Limites

1. Qualité visuelle : Les modèles de diffusion produisent généralement des images de meilleure qualité.
2. Immaturité : Moins de recherche, benchmarks et outils que pour les modèles de diffusion.

4 cas d’usage concrets

1. Génération d’images 2D

Transformation de bruit gaussien en formes complexes (cercles, spirales, lettres), idéal pour visualiser le processus de génération continue.

2. Synthèse audio et parole

Plus rapide que WaveGrad pour générer des waveforms audio de haute qualité avec un nombre réduit d’étapes.

3. Génération de molécules

Trajectoires continues dans l’espace des coordonnées atomiques ou des SMILES pour créer des molécules valides avec propriétés ciblées.

4. Transfert de style continu

Interpolation fluide entre deux distributions d’images (ex: visages âgés vers jeunes) en interpolant les champs de vecteurs appris.

Rectified Flow iteratif — la cle de l’efficacite

Liu et al. (2022) ont montre qu’en iterant le processus d’entrainement, on transforme les courbes sinueuses en trajets quasi-lineaires. Concretement :

1. On entraine un premier modele de Flow Matching
2. On genere des paires avec ce modele
3. On re-entraine sur ces nouvelles paires
4. On repete 2-3 fois

Apres chaque iteration, les trajectoires ODE deviennent plus droites. Le resultat : 10-20 steps d’Euler suffisent au lieu de 50-100, soit un gain de vitesse massif. C’est cette idee de Rectified Flow qui a popularise le Flow Matching aupres de la communaute generation.

Analyse comparative detaillee avec les modeles generatifs

Flow Matching vs GANs

Les GANs souffrent de l’effondrement de mode ou le generateur ne couvre qu’une fraction de la distribution. Le probleme vient de l’entrainement min-max instable. Le Flow Matching n’a aucun de ces problemes : sa loss MSE est convexe et l’entrainement est stable. De plus, le Flow Matching est reversible, ce que les GANs ne permettent pas.

Flow Matching vs VAEs

Les VAEs ont un trade-off fondamental entre fidelite et diversite controle par le weight du KL. Le Flow Matching n’a pas ce compromis car la trajectoire est entierement deterministe : chaque bruit initial correspond a exactement un echantillon genere.

Flow Matching vs Diffusion models

C’est la comparaison la plus interessante. Les deux approches sont conceptuellement similaires : transformer du bruit en donnees via un processus continu. Mais la diffusion utilise des SDEs stochastiques avec un planning de bruit a definir soigneusement, tandis que Flow Matching utilise des ODEs deterministes avec une interpolation lineaire simple. Le Flow Matching est plus rapide, plus facile a implementer, mais la diffusion a plus de recul et de benchmarks.

Rectified Flow iteratif

Liu et al (2022) ont montre qu’en iterant le processus d’entrainement avec les paires generees par le modele lui-meme, les trajectoires deviennent progressivement plus droites. Apres 2-3 iterations, seulement 10-20 steps d’integration Euler suffisent au lieu des 50-100 necessaires initialement. Ce gain de vitesse est massif pour les applications temps reel comme la generation d’images ou la synthese audio.

Comparaison avec les autres methodes

Les GANs souffrent de l’effondrement de mode et d’un entrainement min-max instable. Le Flow Matching a une loss MSE convexe et stable. Les VAEs ont un trade-off fidelite-diversite via le KL que le Flow Matching evite car le processus est deterministe. La diffusion stochastique est plus lente et necessite un planning de bruit tandis que le Flow Matching est deterministe avec une interpolation lineaire simple.

4 cas d’usage supplementaires

En plus de la generation d’images et la synthese audio, le Flow Matching s’applique a la generation de molecules pour la chimie computationnelle, au transfert de style continu entre distributions d’images, a l’augmentation de donnees par interpolation dans l’espace latent, et a l’interpolation semantique entre concepts dans les modeles de langage multimodal.

Voir aussi

• Découverte de la Fonction McCarthy 91 en Python : Comprenez et Implémentez cette Curiosité Algorithmique
• Multiplication Matricielle Optimisée : Implémenter l’Algorithme de Strassen en Python