One-Sided Quantile Coupling for Flow Matching

Technische Zusammenfassung: Einseitige Quantil-Kopplung für Flow Matching

Problemstellung
Flow Matching (FM) trainiert kontinuierliche Zeit-generative Modelle durch die Regression des Geschwindigkeitsfeldes eines Wahrscheinlichkeitspfades, der eine einfache Quellverteilung (typischerweise eine Standard-Normalverteilung) mit einer Ziel-Datenverteilung verbindet. Die Qualität der Kopplung – also die Paarung von Quell- und Zielproben – beeinflusst entscheidend die Geometrie der Trainings-Trajektorien. Standardmäßiges FM verwendet eine unabhängige Zufallskopplung, was oft zu langen, widersprüchlichen Pfaden führt, die die Varianz des Ziel-Geschwindigkeitsfeldes erhöhen und die Trainingseffizienz behindern.

Obwohl strukturierte Kopplungen, die von Optimal Transport (OT) inspiriert sind, diese Pfade begradigen und die Generierungsqualität verbessern können, beruhen bestehende Methoden (z. B. OT-CFM) auf dem Lösen eines Mini-Batch-Zuweisungsproblems. Diese Ansätze erfordern die Konstruktion paarweiser Kostenmatrizen und das Lösen von Zuweisungsproblemen (z. B. via Hungarian-Algorithmus oder Sinkhorn-Iterationen), was zu Rechenkosten führt, die mindestens quadratisch (oder kubisch bei exakten Lösungen) mit der Batch-Größe skalieren. Dies begrenzt die Skalierbarkeit und kann Instabilitäten verursachen. Das Paper adressiert die Frage, ob die geometrischen Vorteile einer strukturierten Kopplung erreicht werden können, ohne zwei vorab gesampte Batches koppeln zu müssen, wodurch der hohe Aufwand für den Batch-weiten Transport vermieden wird.

Methodik: Quantile Coupling Flow Matching (QC-FM)
Die Autoren schlagen Quantile Coupling Flow Matching (QC-FM) vor, ein leichtgewichtiges, einseitiges Kopplungsschema. Im Gegensatz zu traditionellen Methoden, die sowohl Rausch- als auch Daten-Batches sampeln und diese anschließend neu koppeln, sampelt QC-FM nur den Daten-Batch und konstruiert die gekoppelten Quell-Proben direkt.

1. Einseitige Konstruktion:

   • Gegeben sei ein Daten-Batch $\mathcal{B}_x$. Die Methode projiziert die Datenpunkte auf $k$ zufällige orthogonale Richtungen (Slices), die durch eine Matrix $U$ definiert sind.
   • Für jede Richtung werden die Datenpunkte innerhalb des Batches gerankt. Diese empirischen Ränge werden auf Gaußsche Quantile mittels der inversen Standard-Normal-CDF abgebildet.
   • Der Quell-Latent-Code $e^{(i)}$ wird so konstruiert, dass er diese projizierten Quantil-Constraints erfüllt. Konkret entspricht die Projektion des Quell-Codes auf die Slice-Richtungen den aus den Daten-Rängen abgeleiteten Gaußschen Quantilen.
   • Die verbleibenden Dimensionen (der orthogonale Komplement des Slice-Unterraums) werden durch Sampling aus einer bedingten Gaußschen Verteilung aufgefüllt. Dies stellt sicher, dass der Quell-Code in den unbeschränkten Richtungen stochastisch bleibt.

2. Hybrid-Erweiterungen (Gaussian-Remainder Hybrid):
   Da die vollständige QC-Kopplung (bei der alle Quell-Punkte via Rängen konstruiert werden) für $k \ge 2$ aufgrund der Abhängigkeiten zwischen den Slice-Codes von der Standard-i.i.d.-Gaußschen Prior-Verteilung abweicht, führen die Autoren zwei Hybrid-Strategien ein, um die Vorteile von QC zu nutzen und gleichzeitig die Prior-Verteilung zu bewahren:

   • QC-FM-Mixture: Wendet die QC-Konstruktion auf eine "Anker"-Teilmenge des Batches an (z. B. 20 %) und koppelt die verbleibenden "Remainder"-Quell-Proben mit den verbleibenden Daten-Samples mittels unabhängiger Zufallskopplung (dem Baseline-Ansatz). Dies erzeugt eine Mischung aus strukturierten und unstrukturierten Kopplungen.
   • QC-FM-Adjacency: Wendet QC auf die Anker-Teilmenge an, koppelt aber die verbleibenden Quell-Proben basierend auf der lokalen Geometrie, die durch die Anker induziert wird. Verbleibende Datenpunkte werden ihrem nächsten Anker (im strukturierten Unterraum) zugewiesen, und die verbleibenden Quell-Proben werden innerhalb dieser Gruppen gekoppelt. Dies propagiert den strukturierten Bias, ohne die Ziele zu duplizieren.

Wesentliche Beiträge

• QC-FM Algorithmus: Eine neuartige, batch-lokale strukturierte Kopplung, die eine monotone eindimensionale Translation entlang zufälliger Slices erzwingt. Sie eliminiert die Notwendigkeit von paarweisen Kostenmatrizen und Zuweisungs-Solvern, wodurch der dominante Rechenaufwand auf $O(kB \log B)$ (Sortieren) statt $O(B^2)$ oder $O(B^3)$ reduziert wird.
• Theoretische Analyse:
  • Varianzreduktion: Die Autoren beweisen, dass für einen festen Frame die QC-Kopplung die irreduzible Regressionsvarianz entlang jedes gewählten Slices eliminiert, wodurch der ideale Fluss in diesen Richtungen exakt gerade verläuft.
  • Erhalt der Prior-Verteilung: Sie quantifizieren, dass für $k=1$ die Randverteilung der Quelle exakt Gaußsch ist. Für $k \ge 2$ ist die Abweichung von der Gaußschen Prior-Verteilung auf die Copula der Slice-Codes beschränkt, und die Transportkosten dieser Abweichung sind begrenzt.
• Hybride Trainingsverfahren: Die Einführung der Mixture- und Adjacency-Varianten, die den geometrischen Bias und das explizite Gaußsche Signal der Baseline ausbalancieren.

Experimentelle Ergebnisse
Die Methode wurde auf CIFAR-10, CelebA-64, FFHQ-64 und ImageNet-64 unter angepassten Trainingsbudgets und Architekturen evaluiert.

• Performance: QC-FM-Mixture übertraf konsistent die unabhängige Kopplungs-Baseline und den exakten Hungarian OT-CFM über alle vier Datensätze hinweg. Die Verbesserungen gegenüber der Baseline reichten hinsichtlich des FID-50K von 6,3 % (CIFAR-10) bis zu 12,9 % (FFHQ).
• Effizienz: In Bezug auf die Rechenkosten ist QC-FM signifikant schneller als OT-basierte Methoden. Bei einer Batch-Größe von 2048 war QC-FM-Mixture über 800 Mal schneller als der Hungarian OT-Solver und 35 Mal schneller als der Sinkhorn-Solver.
• Toy-Diagnostik: Auf synthetischen Datensätzen (Checkerboard und 8-Gaussians) produzierten QC-FM-Varianten nahezu gerade gelernte Trajektorien (Path-Length Ratio $\approx 1$) und eine signifikant geringere Ziel-Geschwindigkeitsvarianz im Vergleich zur Baseline, was die geometrischen Vorteile der Kopplung bestätigt.

Bedeutung und Ansprüche
Das Paper behauptet, dass die Bewahrung der projizierten Rangstruktur ein einfacher, skalierbarer und effektiver Weg ist, um nützlichen geometrischen Bias in Flow-Matching-Kopplungen zu injizieren, ohne ein Mini-Batch-Transportproblem lösen zu müssen.

• Skalierbarkeit: QC-FM bietet einen attraktiven Mittelweg zwischen der unstrukturierten Baseline und den teuren batch-weiten Transportmethoden. Es ist rechentechnisch leichtgewichtig und über die Anzahl der Projektionen ($k$) steuerbar.
• Praktikabilität: Die Methode benötigt keine gelernten Komponenten für die Inferenz; die Generierung startet weiterhin von der Standard-Normalverteilung. Die Kopplung wird während des Trainings "on-the-fly" konstruiert.
• Limitierungen: Die Autoren merken bescheiden an, dass QC-FM ein batch-lokaler Surrogat für strukturierte Kopplung ist, nicht eine Approximation von globalem Optimal Transport. Die Quantile werden innerhalb der Mini-Batches geschätzt, und die Theorie kontrolliert die Varianz pro Slice, anstatt globale Trajektoriengradheit oder spezifische FID-Verbesserungen zu garantieren. Die Auswahl der Slice-Anzahl ($k$) und des Anker-Verhältnisses ($p$) basiert derzeit auf Heuristiken und datensatzspezifischem Tuning.

Zusammenfassend zeigt die Arbeit, dass eine einseitige Quantil-Kopplung die Stichprobenqualität und die Trainingseffizienz in Flow Matching verbessern kann, indem sie die Rausch-Daten-Paarung durch Rang-Erhaltung strukturiert, und bietet damit eine praktische Alternative zu rechentechnisch prohibitiven Optimal-Transport-Methoden.