La Malédiction des Conditions : Analyse et Amélioration du Transport Optimal pour la Génération Conditionnelle par Flots

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

Résumé

Le couplage par transport optimal sur minibatch redresse les trajectoires dans l'appariement de flux inconditionnel. Cela conduit à une inférence moins exigeante en calcul, car moins d'étapes d'intégration et des solveurs numériques moins complexes peuvent être utilisés lors de la résolution numérique d'une équation différentielle ordinaire au moment du test. Cependant, dans le cadre conditionnel, le transport optimal sur minibatch montre ses limites. En effet, la cartographie par transport optimal par défaut ignore les conditions, ce qui entraîne une distribution a priori biaisée conditionnellement pendant l'entraînement. En revanche, au moment du test, nous n'avons pas accès à cette distribution a priori biaisée et échantillonnons plutôt à partir de la distribution a priori complète et non biaisée. Cet écart entre l'entraînement et le test conduit à une performance médiocre. Pour combler cet écart, nous proposons le transport optimal conditionnel C^2OT qui ajoute un terme de pondération conditionnelle dans la matrice de coût lors du calcul de l'affectation par transport optimal. Les expériences démontrent que cette correction simple fonctionne à la fois avec des conditions discrètes et continues dans les tâches 8gaussians-to-moons, CIFAR-10, ImageNet-32x32 et ImageNet-256x256. Notre méthode surpasse globalement les approches de référence existantes pour différents budgets d'évaluation de fonctions. Le code est disponible à l'adresse 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