Flow Matching

Flow Matching ist ein Verfahren zum Trainieren generativer Modelle, die zufälliges Rauschen in Daten umwandeln, indem sie einem erlernten Geschwindigkeitsfeld folgen. Ein neuronales Netz lernt für jeden Punkt und jeden Zeitpunkt zwischen 0 und 1, in welche Richtung und mit welcher Geschwindigkeit sich der Punkt bewegen soll. Um eine Stichprobe zu erzeugen, startet man mit gaußschem Rauschen und integriert diese Geschwindigkeit mithilfe eines gewöhnlichen Differentialgleichungslösers (ODE) bis zum Zeitpunkt 1. Yaron Lipman und seine Kolleginnen und Kollegen bei Meta AI und dem Weizmann-Institut für Wissenschaften stellten das Verfahren auf der ICLR 2023 in Flow Matching für generative Modellierung vor. Zwei gleichzeitig erschienene Arbeiten gelangten unter den Bezeichnungen rektifizierter Fluss und stochastische Interpolanten zu nahezu demselben Optimierungsziel. Flow Matching trainiert kontinuierliche normalisierende Flüsse, ohne während des Trainings die ODE zu simulieren. Es umfasst Diffusionsmodelle als Spezialfall und liegt dem Trainingsziel von Stable Diffusion 3, FLUX.1, Metas Movie Gen und vielen aktuellen Sprach- und Robotikmodellen zugrunde.

So funktioniert Flow Matching#

Beim Training wird nie etwas generiert. Bei jedem Schritt werden eine Datenstichprobe x1, eine Rauschstichprobe x0 aus einer Standardnormalverteilung und ein Zeitpunkt t zwischen 0 und 1 gezogen. Daraus wird ein Punkt auf der Geraden zwischen den beiden Stichproben gebildet: xt = (1 - t) * x0 + t * x1. Ein Punkt auf dieser Geraden bewegt sich mit der konstanten Geschwindigkeit x1 - x0. Deshalb wird das Netz v(xt, t) mit einem einfachen mittleren quadratischen Fehler darauf trainiert, x1 - x0 vorherzusagen. Jedes Paar aus Rauschen und Daten definiert nur einen Pfad, doch durch jeden Punkt verlaufen viele Pfade. Die Geschwindigkeit des gesamten Übergangs von Rauschen zu Daten an diesem Punkt ist ihr Mittelwert. Die Arbeit beweist, dass die Regression auf die einfachen Ziele einzelner Paare dieselben Gradienten ergibt wie die Regression auf diesen unlösbaren Mittelwert. Das Training verwendet daher eine einfache Regressions-Verlustfunktion, ohne ODE-Simulation.

Bei der Generierung wird dx/dt = v(x, t) von t = 0 bis t = 1 integriert, üblicherweise mit Euler-Schritten oder Schritten eines Lösers höherer Ordnung. Dieser PyTorch-Code trainiert ein kleines Geschwindigkeitsnetz mit zwei Clustern aus 2D-Punkten und zieht anschließend Stichproben mit 1, 2 und 8 Euler-Schritten:

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

Die Ausgabe zeigt zuerst 1-step sampling: 0% of samples land near a cluster, dann 47 % bei 2 Schritten und 94 % bei 8 Schritten. Bei t = 0 hat jeder Rauschpunkt denselben erwarteten Zielpunkt – den Mittelwert der Daten. Daher schickt ein einzelner Schritt alle Stichproben in den leeren Raum zwischen den beiden Clustern. Jeder Trainingspfad ist zwar gerade, doch nahe beieinander startende Pfade enden in unterschiedlichen Clustern, und der gelernte Fluss, der sie mittelt, ist dennoch gekrümmt. Generatoren mit wenigen Schritten verwenden daher Techniken wie Reflow, bei dem mit den eigenen Rausch-Stichproben-Paaren des Modells erneut trainiert wird, um den Fluss zu begradigen, oder die Destillation in Konsistenzmodelle.

Flow Matching und verwandte Verfahren#

• Diffusionsmodelle: Beim Training von Diffusionsmodellen werden Daten gemäß einem Rauschplan verrauscht und einem Netz wird beigebracht, das hinzugefügte Rauschen vorherzusagen. Anschließend werden Stichproben erzeugt, indem dieser Prozess umgekehrt wird. Die Arbeit zu Flow Matching zeigte, dass das Optimierungsziel jeden gaußschen Wahrscheinlichkeitsverlauf akzeptiert, einschließlich der von Diffusionsmodellen verwendeten Verläufe, und dass sich diese damit stabiler trainieren lassen. Der gerade Pfad ist eine eigenständige, einfachere Wahl. Beide Modellfamilien verwenden gemeinsame Architekturen wie den Diffusion Transformer (DiT), die klassifikatorfreie Steuerung und ODE-Abtastverfahren.
• Rektifizierter Fluss: Verwendet dieselbe lineare Interpolation und denselben Geschwindigkeitsverlust. Für den linearen Pfad sind die beiden Bezeichnungen daher weitgehend synonym. Die Arbeit zum rektifizierten Fluss ergänzt das Reflow-Verfahren, um die Trajektorien weiter zu begradigen.
• Normalisierende Flüsse: Diskrete normalisierende Flüsse stapeln umkehrbare Schichten und werden mit exakter Likelihood trainiert. Kontinuierliche normalisierende Flüsse, eine Form der neuronalen ODE, definieren die Transformation mithilfe einer ODE, wurden aber mit Maximum-Likelihood trainiert. Dazu muss die ODE bei jedem Trainingsschritt gelöst werden. Flow Matching trainiert denselben ODE-Modelltyp durch Regression und ohne diese Berechnung.
• Optischer Fluss: Trotz des ähnlichen Namens besteht kein Zusammenhang. Der optische Fluss bezeichnet die scheinbare Bewegung von Pixeln zwischen Videobildern. Flow Matching bewegt dagegen Stichproben durch den Raum möglicher Daten.

Einsatzbereiche von Flow Matching#

• Bildgenerierung: Stability AI trainierte Stable Diffusion 3 mit einem Transformer für rektifizierten Fluss und zog Trainingszeitpunkte aus einer Logit-Normalverteilung, die mittlere Rauschpegel bevorzugt. FLUX.1 von Black Forest Labs ist ein Transformer für rektifizierten Fluss mit 12 Milliarden Parametern. Bilder aus diesen Text-zu-Bild-Modellen dienen als synthetische Daten, die Datensätze zum Trainieren von YOLO26-Erkennungsmodellen um seltene Objekte, Wetterbedingungen und Blickwinkel ergänzen.
• Videogenerierung: Metas Movie Gen ist ein Transformer mit 30 Milliarden Parametern, der mit Flow Matching trainiert wurde. Er erzeugt bis zu 16 Sekunden lange Videos in 1080p mit 16 Bildern pro Sekunde.
• Sprache: Metas Voicebox wurde als nicht-autoregressives Flow-Matching-Modell mit mehr als 50.000 Stunden Sprachmaterial trainiert. Bei der Sprachsynthese aus Text ohne Beispiele erreichte es eine Wortfehlerrate von 1,9 % gegenüber 5,9 % bei VALL-E und lief bis zu 20-mal schneller. F5-TTS kombinierte Flow Matching später mit einem Diffusion Transformer.
• Robotersteuerung: π0 von Physical Intelligence ergänzt ein Bild-Sprach-Modell mit 3 Milliarden Parametern um einen 300 Millionen Parameter großen „Aktionsspezialisten“ für Flow Matching. Dieser erzeugt Folgen von 50 zukünftigen Aktionen mit 10 Integrationsschritten und steuert Roboter mit bis zu 50 Hz.