Der Fluch der Bedingungen: Analyse und Verbesserung des Optimalen Transports für bedingte flussbasierte Generierung

The Curse of Conditions: Analyzing and Improving Optimal Transport for Conditional Flow-Based Generation

March 13, 2025

Autoren: Ho Kei Cheng, Alexander Schwing

cs.AI

Zusammenfassung

Minibatch Optimal Transport-Kopplung begradigt Pfade im unkonditionellen Flow-Matching. Dies führt zu einem rechnerisch weniger aufwändigen Inferenzprozess, da weniger Integrationsschritte und weniger komplexe numerische Löser eingesetzt werden können, wenn zur Testzeit eine gewöhnliche Differentialgleichung numerisch gelöst wird. Im konditionellen Setting hingegen reicht Minibatch Optimal Transport nicht aus. Dies liegt daran, dass die standardmäßige Optimal-Transport-Abbildung die Bedingungen ignoriert, was zu einer konditionell verzerrten Prior-Verteilung während des Trainings führt. Zur Testzeit haben wir jedoch keinen Zugriff auf diese verzerrte Prior-Verteilung, sondern stichproben aus der vollständigen, unverzerrten Prior-Verteilung. Diese Diskrepanz zwischen Training und Test führt zu einer suboptimalen Leistung. Um diese Lücke zu schließen, schlagen wir Conditional Optimal Transport (C²OT) vor, das einen konditionellen Gewichtungsterm in die Kostenmatrix einfügt, wenn die Optimal-Transport-Zuordnung berechnet wird. Experimente zeigen, dass diese einfache Anpassung sowohl mit diskreten als auch kontinuierlichen Bedingungen in 8gaussians-to-moons, CIFAR-10, ImageNet-32x32 und ImageNet-256x256 funktioniert. Unsere Methode schneidet insgesamt besser ab als die bestehenden Baselines über verschiedene Funktionsauswertungsbudgets hinweg. Der Code ist verfügbar unter https://hkchengrex.github.io/C2OT.

English

Minibatch optimal transport coupling straightens paths in unconditional flow matching. This leads to computationally less demanding inference as fewer integration steps and less complex numerical solvers can be employed when numerically solving an ordinary differential equation at test time. However, in the conditional setting, minibatch optimal transport falls short. This is because the default optimal transport mapping disregards conditions, resulting in a conditionally skewed prior distribution during training. In contrast, at test time, we have no access to the skewed prior, and instead sample from the full, unbiased prior distribution. This gap between training and testing leads to a subpar performance. To bridge this gap, we propose conditional optimal transport C^2OT that adds a conditional weighting term in the cost matrix when computing the optimal transport assignment. Experiments demonstrate that this simple fix works with both discrete and continuous conditions in 8gaussians-to-moons, CIFAR-10, ImageNet-32x32, and ImageNet-256x256. Our method performs better overall compared to the existing baselines across different function evaluation budgets. Code is available at https://hkchengrex.github.io/C2OT