Technisches Resümee: Verbesserung von Rectified Flow durch Randbedingungen

Problemstellung

Rectified Flow (RF), auch bekannt als Flow Matching, hat sich als ein leistungsstarkes Framework für die generative Modellierung etabliert, das ein Geschwindigkeitsfeld lernt, um eine Rauschverteilung mittels einer gewöhnlichen Differentialgleichung (ODE) in eine Ziel-Datenverteilung zu transformieren. Trotz seines Erfolgs in Aufgaben wie der Text-zu-Bild- oder Videoerzeugung identifizieren die Autoren eine kritische Einschränkung in Standard-RF-Implementierungen: Das gelernte Geschwindigkeitsfeld erfüllt oft nicht die theoretisch abgeleiteten Randbedingungen.

Speziell unter linearer Interpolation zwischen Rauschen ($t=0$) und Daten ($t=1$) sollte das optimale Geschwindigkeitsfeld folgende Bedingungen erfüllen:

1. Linke Randbedingung ($t=0$): $\mathbf{v}(\mathbf{x}, 0) = \mathbb{E}[\mathbf{X}_1] - \mathbf{x}$
2. Rechte Randbedingung ($t=1$): $\mathbf{v}(\mathbf{x}, 1) = \mathbf{x}$

Standardmäßige neuronale Netzwerkparameterisierungen des Geschwindigkeitsfeldes erzwingen diese Bedingungen nicht inhärent. Wie in der Arbeit visualisiert, weicht die vorhergesagte Geschwindigkeit bei $t=1$ oft von der erwarteten Identitätsabbildung ab. Diese Verletzung ist besonders schädlich für das stochastische Sampling (z. B. mittels Langevin-Dynamik). Da die Schätzung der Score-Funktion (via Tweedies Formel) eine Division durch $(1-t)$ beinhaltet, führt eine ungenaue Geschwindigkeit nahe $t=1$ dazu, dass die Score-Schätzung divergiert. Dies führt zu instabilem Sampling, was in überglatten Bildern oder „cartoonartigen“ Artefakten resultiert, insbesondere bei Verwendung größerer Schrittweiten oder Rauschskalen.

Methodik

Um dies zu adressieren, schlagen die Autoren das Boundary-enforced Rectified Flow Model (Boundary RF Model) vor. Dieser Ansatz integriert Randbedingungen explizit in die Parameterisierung des Geschwindigkeitsfeldes durch zwei unterschiedliche Varianten:

1. Masken-basiertes Boundary RF Model

Diese Variante erzwingt beide Randbedingungen ($t=0$ und $t=1$), indem sie das Geschwindigkeitsfeld als Linearkombination von Randtermen und einem lernbaren neuronalen Netz-Output strukturiert:
$$\mathbf{v}(\mathbf{x}, t) = g(t) \cdot (C - \mathbf{x}) + f(t) \cdot \mathbf{x} + h(t) \cdot m_\theta(\mathbf{x}, t)$$
Wobei:

• $m_\theta(\mathbf{x}, t)$ ein Standard-neuronales Netz (z. B. U-Net oder DiT) ist.
• $C = \mathbb{E}[\mathbf{X}_1]$.
• $f(t), g(t), h(t)$ Skalarfunktionen sind, die an spezifischen Grenzen verschwinden oder aktiviert werden (z. B. $g(1)=f(0)=h(0)=h(1)=0$ und $g(0)=f(1)=1$).
• Die Autoren schlagen eine „Standard Cosine“-Wahl vor: $g(t) = \cos(\frac{\pi}{2}t)$, $f(t) = \sin(\frac{\pi}{2}t)$ und $h(t) = \sin(\pi t)$.

Diese Konstruktion garantiert, dass die Randbedingungen per Design erfüllt sind, unabhängig vom Output des neuronalen Netzes.

2. Subtraktions-basiertes Boundary RF Model

Motiviert durch die Beobachtung, dass die alleinige Durchsetzung der $t=1$-Bedingung signifikante empirische Vorteile bringt, schlagen die Autoren eine einfachere Variante mit einer einzelnen Randbedingung vor:
$$\mathbf{v}(\mathbf{x}, t) = \mathbf{x} + m_\theta(\mathbf{x}, t) - m_\theta(\mathbf{x}, 1)$$
Diese Formulierung erfüllt inhärent $\mathbf{v}(\mathbf{x}, 1) = \mathbf{x}$, da der Term $m_\theta(\mathbf{x}, 1)$ bei $t=1$ herausfällt. Während dieser Ansatz eine einfache Gestaltung und Flexibilität bietet, verursacht er einen Rechenaufwand, da für jede Evaluierung ein zusätzlicher Vorwärtspass zur Berechnung von $m_\theta(\mathbf{x}, 1)$ erforderlich ist.

Auswirkungen auf das stochastische Sampling

Durch die Erzwingung von $\mathbf{v}(\mathbf{x}, 1) = \mathbf{x}$ stellt das Boundary RF Model sicher, dass der Zähler in der Tweedie-Formel ($t\mathbf{v}(\mathbf{x}, t) - \mathbf{x}$) gegen Null geht, wenn $t \to 1$. Dies verhindert, dass die Score-Funktion $\nabla \log \rho_t$ divergiert, wodurch stochastische Sampler (SDE-Sampler) stabilisiert werden und eine robuste Generierung selbst mit größeren Rauschskalen oder Schrittweiten ermöglicht wird.

Kernbeiträge

1. Identifizierung einer Einschränkung: Die Arbeit hebt hervor, dass unbeschränkte Geschwindigkeitsfelder in Rectified Flow die theoretischen Randbedingungen verletzen, was zu ungenauen ODE-Trajektorien und instabilem stochastischem Sampling führt.
2. Neuartige Parameterisierung: Die Einführung des Boundary RF Modells, das zwei praktische Varianten (maskenbasiert und subtraktionsbasiert) bietet, um diese Bedingungen mit minimalem Code-Aufwand zu erzwingen.
3. Stabilisierung von Score-Funktionen: Der Nachweis, dass die Durchsetzung von Randbedingungen die Score-Funktion nahe $t=1$ direkt stabilisiert und somit die Divergenzprobleme mildert, die das klassische (vanilla) RF bei stochastischem Sampling plagen.
4. Empirische Validierung: Umfassende Experimente, die konsistente Verbesserungen der Generierungsqualität und der Sampling-Stabilität über verschiedene Datensätze und Modellskalen hinweg zeigen.

Experimentelle Ergebnisse

Die Autoren evaluierten ihre Modelle auf CIFAR-10 und ImageNet (256$\times$256 und 512$\times$512) unter Verwendung von Diffusion Transformer (DiT) Architekturen.

• Quantitative Verbesserungen:

  • Auf ImageNet 256$\times$256 erreichte das Subtraktions-basierte Boundary RF Model einen FID von 6,32 (ODE-Sampling) und 6,38 (SDE-Sampling), verglichen mit dem Vanilla-RF-Baseline-Wert von 6,87 bzw. 6,79.
  • Das masken-basierte Modell erreichte einen FID von 6,63 (ODE) und 6,18 (SDE).
  • Verbesserungen wurden auch beim Inception Score (IS), Precision und Recall über beide Datensätze hinweg beobachtet.
  • Die Subtraktions-basierte Variante erzielte konsistent die höchsten IS-Werte.

• Sampling-Robustheit:

  • In Experimenten mit dem Overshooting-Sampler und variierenden Rauschskalen ($c$) zeigten die Boundary RF Modelle eine signifikant stabilere Performance und weniger Qualitätsverlust im Vergleich zu Vanilla RF, wenn die Stochastizität zunahm.
  • Qualitative Ergebnisse zeigten, dass Boundary RF Modelle schärfere Bilder mit weniger Artefakten und besserer Objekttreue produzieren, insbesondere in stochastischen Sampling-Szenarien, in denen Vanilla RF dazu neigt, zu stark zu glätten.

• Skalierbarkeit:

  • Die Vorteile des Boundary RF Modells waren konsistent beim Skalieren auf größere Modelle (DiT-L, DiT-XL) und höhere Auflösungen (512$\times$512) feststellbar.

Bedeutung und Behauptungen

Die Arbeit behauptet, dass das Boundary RF Model eine einfache, leicht implementierbare Lösung für eine fundamentale Einschränkung von Rectified Flow darstellt. Durch die explizite Durchsetzung der Randbedingungen bewirkt die Methode:

1. Verbesserung der Generierungsqualität: Es übertrifft Vanilla RF konsistent in Bezug auf FID und andere Metriken.
2. Ermöglichung von stabilem stochastischem Sampling: Es löst die Divergenz der Score-Funktion nahe der Endzeit, was robustere Sampling-Strategien erlaubt.
3. Einfache Anpassbarkeit: Es kann mit minimalen architektonischen Änderungen in bestehende RF-Modelle (wie jene mit U-Net oder DiT) integriert werden.

Die Autoren positionieren diese Arbeit als notwendige Verfeinerung für den praktischen Einsatz von Rectified Flow, insbesondere in Szenarien, die ein hochwertiges stochastisches Sampling erfordern. Sie deuten an, dass das Kernkonzept der Durchsetzung von Randbedingungen potenziell auch auf Diffusionsmodelle mit geringfügigen Anpassungen übertragen werden könnte. Zukünftige Arbeiten sollen die adaptive Randbedingungs-Durchsetzung und die Anwendung auf die Text-zu-Video-Generierung untersuchen.