MixFlow: Mixed Source Distributions Improve Rectified Flows

1. Problemstellung

Generative Modelle, insbesondere Diffusionsmodelle und deren Variante „Rectified Flows" (geradlinige Flüsse), sind in der Lage, hochwertige und vielfältige Bilder zu erzeugen. Ein Hauptnachteil dieser Modelle ist jedoch die langsame iterative Probennahme (Sampling), die durch stark gekrümmte generative Pfade verursacht wird.

• Ursache der Krümmung: Die Krümmung der Pfade hängt stark von der Forward Coupling (der gemeinsamen Verteilung von Quelle und Ziel) ab. Herkömmliche Ansätze gehen von einer unabhängigen Kopplung aus, bei der die Quellverteilung eine Standard-Gauß-Verteilung $p_0(x) = \mathcal{N}(0, I)$ ist und die Zielverteilung $p_1(x)$ die Datenverteilung darstellt. Diese Unabhängigkeit führt zu stark sich kreuzenden Trajektorien, was den Vektorfeld-Modellen erschwert, gerade Pfade zu lernen.
• Bisherige Lösungen: Methoden wie Distillation oder Consistency Models versuchen, die Abbildung zu linearisieren, erfordern aber oft mehrfaches Nachtrainieren oder verschlechtern die Qualität bei höheren Schrittzahlen. Andere Ansätze optimieren die Forward Coupling (z. B. Fast-ODE), stoßen jedoch an Grenzen, wenn sie versuchen, die Quellverteilung zu stark an die Daten anzupassen, was zu Problemen wie dem „Prior Hole" (Lücken im Prior) führt oder eine starke Regularisierung erfordert, die den Vorteil wieder zunichtemacht.

2. Methodik: MixFlow

Die Autoren schlagen MixFlow vor, eine Trainingsstrategie, die die Krümmung der generativen Pfade reduziert, indem sie die Quellverteilung optimiert, ohne die Nachteile starker Regularisierung in Kauf nehmen zu müssen.

2.1 $\kappa$-Forward Coupling ($\kappa$-FC)

Als theoretische Grundlage wird $\kappa$-FC eingeführt. Dies ist eine allgemeine Formulierung für lernbare Forward Couplings, die auf einem beliebigen Signal $\kappa$ basieren.

• $\kappa$ kann ein informatives Signal sein (z. B. ein Datenbeispiel $x_1$, ein Klassenlabel) oder ein unabhängiges Signal (z. B. Rauschen).
• Die Idee ist, die Quellverteilung $q(x_0|\kappa)$ so zu lernen, dass sie besser mit der Zielverteilung $p_1(x)$ abgestimmt ist.
• Problem: Eine naive Optimierung von $q(x_0|\kappa)$ (z. B. als Gauß-Verteilung mit lernbarem Mittelwert und Varianz) erfordert eine Regularisierung (KL-Divergenz zum Standard-Gauß), um sicherzustellen, dass das Modell auch ohne $\kappa$ (während des Inference) funktioniert. Ein zu starker Regularisierungsfaktor $\beta$ führt jedoch dazu, dass die Verteilung wieder unabhängig von $\kappa$ wird; ein zu schwacher Faktor führt zum „Prior Hole"-Problem.

2.2 Die MixFlow-Lösung: Verteilungsmischung

Um diesen Zielkonflikt zu lösen, schlägt MixFlow vor, die Quellverteilung als lineare Mischung aus zwei Verteilungen zu trainieren:

1. Eine bedingte Verteilung $q_\phi(x_0|\kappa)$, die durch das Signal $\kappa$ gesteuert wird und die Datenstruktur lernt.
2. Eine unbedingte Standard-Gauß-Verteilung $\mathcal{N}(0, I)$.

Die gemischte Quellverteilung wird definiert als:
$$q_\phi(x_0 | \kappa, w) = \mathcal{N}(w\mu_\phi(\kappa), w\Sigma_\phi(\kappa) + (1-w)I)$$
Dabei ist $w \sim U(0, 1)$ ein zufälliger Gewichtungsfaktor, der während des Trainings für jedes Beispiel neu gezogen wird.

Vorteile dieses Ansatzes:

• Robustheit: Da das Modell während des Trainings auch direkt von der Standard-Gauß-Verteilung ($w=0$) lernt, bleibt der gesamte Gauß-Raum abgedeckt. Dies eliminiert das „Prior Hole"-Problem.
• Geringe Regularisierung: Durch die Mischung kann der Regularisierungsfaktor $\beta$ extrem klein gewählt werden (z. B. $10^{-5}$), ohne dass die Stabilität leidet. Dies erlaubt der bedingten Verteilung, sich stark von der Standard-Gauß-Verteilung zu unterscheiden und komplexe Datenmodi besser zu erfassen.
• Transfer der Struktur: Die effiziente Kopplung zwischen bedingter Quelle und Ziel wird während des Trainings auf die unbedingte Quelle übertragen. Das Modell lernt also, wie man von einem reinen Rauschen (Standard-Gauß) startet, aber die Struktur der Daten nutzt, um gerade Pfade zu erzeugen.

2.3 Training und Sampling

• Training: Das Vektorfeld $v_\theta$ wird trainiert, um von dieser gemischten Verteilung zu fließen. Der Verlust beinhaltet den Standard-Rectified-Flow-Verlust plus einen sehr kleinen KL-Regularisierungsterm.
• Sampling: Während des Inference kann das Modell entweder mit einem gegebenen $\kappa$ (falls verfügbar) oder ohne (Standard-Gauß) starten. Interessanterweise kann der Parameter $w$ auch während des Sampling angepasst werden, um einen Trade-off zwischen Geschwindigkeit und Qualität zu steuern (z. B. mehr Gewichtung auf die bedingte Verteilung für wenige Schritte).

3. Wichtige Beiträge

1. $\kappa$-FC Formulierung: Eine allgemeine Methode für lernbare Forward Couplings, die beliebige Signale nutzen kann, um Quellverteilungen zu optimieren.
2. MixFlow-Strategie: Ein einfacher, aber effektiver Ansatz, der lineare Mischungen von Verteilungen nutzt, um die Regularisierung zu umgehen und gleichzeitig die Trajektorienkrümmung zu minimieren.
3. Verbesserte Effizienz: Die Methode ermöglicht eine signifikante Beschleunigung des Trainings (konvergiert schneller) und des Sampling-Prozesses (weniger Schritte für gleiche Qualität).

4. Ergebnisse

Die Autoren evaluieren MixFlow auf den Datensätzen CIFAR10, FFHQ (64x64) und AFHQv2 (64x64) und vergleichen sie mit State-of-the-Art-Methoden wie Rectified Flow, Fast-ODE und QAC.

• Qualitätsverbesserung (FID):
  • Im Vergleich zum Standard-Rectified-Flow verbessert MixFlow den FID (Fréchet Inception Distance) um durchschnittlich 12 %.
  • Im Vergleich zu den besten vorherigen Baselines (wie Fast-ODE) wird eine Verbesserung von 7 % erzielt.
• Reduzierte Krümmung: Die Trajektorienkrümmung wird um ca. 22 % im Vergleich zu Rectified Flow und 5 % im Vergleich zu Fast-ODE reduziert (gemessen durch den Krümmungswert $I(q)$).
• Sampling-Budget:
  • Bei niedrigem Sampling-Budget (z. B. 5–9 Funktionsevaluationen) übertrifft MixFlow die Konkurrenz deutlich (bis zu 20 % FID-Verbesserung gegenüber Fast-ODE bei 5 Schritten).
  • Auch bei hohem Budget (viele Schritte) bleibt die Qualität besser.
• Trainingseffizienz: MixFlow erreicht die gleiche Leistung wie Fast-ODE mit nur 60 % des Trainingsbudgets (Iterationen).
• Flexibilität: Das Modell funktioniert auch ohne $\kappa$ während des Inference (da es von der Mischung lernt) und ermöglicht durch den Parameter $w$ eine Steuerung des Geschwindigkeits-Qualitäts-Trade-offs zur Laufzeit.

5. Bedeutung und Fazit

MixFlow adressiert das fundamentale Problem der langsamen Probennahme in Flow-basierten Modellen nicht durch komplexere Solver oder mehrfaches Nachtrainieren, sondern durch eine optimierte Trainingsstrategie der Quellverteilung.

Die zentrale Erkenntnis ist, dass die strikte Trennung zwischen bedingter und unbedingter Quellverteilung (wie bei reinen Conditional-Approaches) vermieden werden muss. Durch die Einführung einer gemischten Quellverteilung während des Trainings gelingt es:

1. Die Regularisierung (KL-Divergenz) extrem zu schwächen, was zu flacheren, geradlinigeren Trajektorien führt.
2. Die Stabilität des Modells zu gewährleisten, sodass es auch ohne Konditionierung (unconditional generation) robust funktioniert.
3. Die Sampling-Geschwindigkeit signifikant zu erhöhen, ohne Einbußen bei der Bildqualität hinnehmen zu müssen.

Dieser Ansatz stellt einen wichtigen Schritt dar, um die „Generative Learning Trilemma" (Geschwindigkeit, Vielfalt, Qualität) zugunsten von Geschwindigkeit und Qualität zu verschieben, und bietet eine neue Perspektive auf die Optimierung von Forward Couplings in generativen Modellen.