Rectified Flow

Rectified Flow ist eine Methode des generativen Modellierens. Dabei wird ein Geschwindigkeitsfeld gelernt, das Stichproben entlang möglichst gerader Pfade von einer einfachen Verteilung – üblicherweise Gaußschem Rauschen – zur Datenverteilung überführt. Das Training ist eine einfache Regression: Man kombiniert eine Rauschstichprobe mit einer Datenstichprobe, wählt einen Punkt auf der geraden Linie zwischen beiden und bringt einem neuronalen Netz bei, die Richtung der Linie vorherzusagen. Bei der Generierung wird die erlernte Geschwindigkeit mit einem Löser für gewöhnliche Differentialgleichungen (ODE) integriert – derselben Methode wie bei neuronalen ODEs. Gerade Pfade lassen sich mit wenigen großen Schritten durchlaufen. Xingchao Liu, Chengyue Gong und Qiang Liu von der University of Texas at Austin stellten die Methode im September 2022 in Flow Straight and Fast: Learning to Generate and Transfer Data with Rectified Flow vor, gemeinsam mit dem eng verwandten Flow Matching von Yaron Lipman und Kollegen bei Meta.

So wird Rectified Flow trainiert#

Für das Rauschen x0 und die Daten x1 bewegt sich der Punkt x_t = (1 - t) * x0 + t * x1 mit der konstanten Geschwindigkeit x1 - x0 bei t = 0 vom Rauschen zu den Daten bei t = 1. Das Netz v(x_t, t) wird darauf trainiert, diese Geschwindigkeit anhand zufälliger Paare und Zeitpunkte mit einer quadratischen Verlustfunktion vorherzusagen. Durch jeden Punkt verlaufen viele Linien, daher lernt das Netz deren mittlere Richtung. Die von ihm definierten ODE-Pfade kreuzen sich nie, sondern krümmen sich dort, wo Linien zu unterschiedlichen Datenpunkten zusammentreffen. Ein eindimensionaler Rectified Flow in PyTorch lernt, Gaußsches Rauschen in zwei Cluster bei -2 und +2 umzuwandeln:

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

Mit einem Euler-Schritt landet keine Stichprobe in der Nähe eines Modus: Bei t = 0 führt die durchschnittliche Richtung jedes Rauschpunkts zum Datenmittelwert 0, der mittig zwischen den Clustern liegt. Mit zwei Schritten landen 39 % der Stichproben höchstens 0,3 Einheiten von einem Modus entfernt, mit vier Schritten 81 % und mit 32 Schritten 96 %. Das liegt daran, dass sich die erlernten Pfade zu den Clustern hin krümmen und nur kleine Schritte dieser Krümmung folgen.

Reflow: Die Pfade begradigen#

Reflow, in der Arbeit rekursive Begradigung genannt, beseitigt diese Krümmungen. Der trainierte Flow erzeugt eine Menge von (Rauschen, Stichprobe)-Paaren, und ein neuer Flow wird statt mit zufälligen Paaren mit diesen Paaren trainiert. Jedes Paar ist bereits durch einen Pfad des ersten Flows verbunden. Daher stehen die geraden Linien zwischen den Paaren selten im Konflikt, und die Arbeit beweist, dass die neue Kopplung die Transportkosten nie erhöht. Wiederholte Reflow-Schritte ergeben nahezu gerade Pfade. Die Autorinnen und Autoren berichteten von hochwertigen Bildern und Bildübersetzungen mit nur einem Euler-Schritt:

Dasselbe Team wandte Reflow und Wissensdestillation in InstaFlow (2023) auf Stable Diffusion 1.5 an. Das Ein-Schritt-Modell erreichte auf MS COCO 2017-5k einen FID von 23,3, gegenüber 37,2 bei der progressiven Destillation. Es erzeugte ein Bild in 0,09 Sekunden und benötigte für das Training 199 A100-GPU-Tage.

Rectified Flow in Text-zu-Bild-Modellen#

In der Arbeit zu Stable Diffusion 3 verglich Stability AI verschiedene Trainingsformulierungen und entschied sich für Rectified Flow mit Trainingszeitpunkten aus einer Logit-Normalverteilung. Diese konzentriert das Training auf die Mitte des Pfads, wo die Vorhersage der Geschwindigkeit am schwierigsten ist. Das Modell kombiniert dieses Ziel mit einem multimodalen Diffusion-Transformer (DiT), der in einem latenten Raum arbeitet. Black Forest Labs, gegründet von ehemaligen Forschenden von Stability AI, beschreibt sein FLUX.1-dev-Modell als Rectified-Flow-Transformer mit 12 Milliarden Parametern. Die destillierte Variante FLUX.1 [schnell] erzeugt Bilder in einem bis vier Schritten. Diese Text-zu-Bild-Modelle sind eine häufige Quelle für synthetische Daten beim Training von Bildverarbeitungsmodellen. Die Ultralytics Platform integriert die Text-zu-Bild-Generierung mit Ähnliche Bilder generieren in die Zusammenstellung von Datensätzen: Ausgehend von einem einzelnen Datensatzbild verfasst sie neue Szenenbeschreibungen, erzeugt bis zu 16 Varianten und fügt die von dir ausgewählten ohne Labels zum Datensatz hinzu, bereit für die Annotation.

Rectified Flow im Vergleich zu verwandten Begriffen#

• Flow Matching: Der allgemeine Rahmen zum Trainieren von ODE-Generatoren durch Regression von Geschwindigkeiten entlang ausgewählter Wahrscheinlichkeitsverläufe. Die gerade Linie zwischen unabhängigem Rauschen und Daten ist der am häufigsten verwendete Verlauf von Rectified Flow. Die beiden Begriffe werden oft synonym verwendet.
• Diffusionsmodelle: Sie fügen nach einem festgelegten Zeitplan Rauschen hinzu, dessen Wahrscheinlichkeitsfluss-Pfade gekrümmt sind. Ihre deterministischen Sampler benötigen typischerweise 20 bis 50 Schritte.
• Normalisierende Flows: Sie bilden Rauschen durch eine feste Abfolge invertierbarer Schichten auf Daten ab und werden mit exakter Maximum-Likelihood-Schätzung trainiert.
• Konsistenzmodelle: Sie lernen, von jedem Punkt einer gekrümmten Trajektorie direkt zu deren Ende zu springen, sodass ein einziger Modellaufruf eine Stichprobe erzeugen kann.