De Vloek van Condities: Analyse en Verbetering van Optimale Transport voor Voorwaardelijke Stroomgebaseerde Generatie

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

March 13, 2025

Auteurs: Ho Kei Cheng, Alexander Schwing

cs.AI

Samenvatting

Minibatch optimaal transportkoppeling vereenvoudigt paden in onvoorwaardelijke flow matching. Dit leidt tot minder rekenintensieve inferentie, omdat er minder integratiestappen en minder complexe numerieke oplossers nodig zijn bij het numeriek oplossen van een gewone differentiaalvergelijking tijdens testtijd. In de voorwaardelijke setting schiet minibatch optimaal transport echter tekort. Dit komt doordat de standaard optimale transporttoewijzing geen rekening houdt met voorwaarden, wat resulteert in een voorwaardelijk scheve priorverdeling tijdens de training. Tijdens testtijd hebben we echter geen toegang tot deze scheve prior, maar nemen we steekproeven uit de volledige, onbevooroordeelde priorverdeling. Deze kloof tussen training en testen leidt tot een ondermaatse prestaties. Om deze kloof te overbruggen, stellen we voorwaardelijk optimaal transport C^2OT voor, dat een voorwaardelijke wegingsterm toevoegt aan de kostenmatrix bij het berekenen van de optimale transporttoewijzing. Experimenten tonen aan dat deze eenvoudige aanpassing werkt met zowel discrete als continue voorwaarden in 8gaussians-to-moons, CIFAR-10, ImageNet-32x32 en ImageNet-256x256. Onze methode presteert over het algemeen beter in vergelijking met de bestaande baseline-methoden bij verschillende functie-evaluatiebudgetten. Code is beschikbaar op 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