La Maledizione delle Condizioni: Analisi e Miglioramento del Trasporto Ottimale per la Generazione Condizionata Basata su Flussi

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

March 13, 2025

Autori: Ho Kei Cheng, Alexander Schwing

cs.AI

Abstract

L'accoppiamento del trasporto ottimale su minibatch raddrizza i percorsi nel flusso di corrispondenza incondizionato. Ciò porta a un'inferenza computazionalmente meno impegnativa, poiché è possibile utilizzare meno passi di integrazione e risolutori numerici meno complessi quando si risolve numericamente un'equazione differenziale ordinaria al momento del test. Tuttavia, nel contesto condizionato, il trasporto ottimale su minibatch non è sufficiente. Questo perché la mappatura predefinita del trasporto ottimale ignora le condizioni, risultando in una distribuzione a priori condizionata distorta durante l'addestramento. Al contrario, al momento del test, non abbiamo accesso alla distribuzione a priori distorta, ma campioniamo dalla distribuzione a priori completa e imparziale. Questo divario tra addestramento e test porta a una performance inferiore. Per colmare questo divario, proponiamo il trasporto ottimale condizionato C^2OT, che aggiunge un termine di ponderazione condizionata nella matrice dei costi quando si calcola l'assegnazione del trasporto ottimale. Gli esperimenti dimostrano che questa semplice correzione funziona sia con condizioni discrete che continue in 8gaussians-to-moons, CIFAR-10, ImageNet-32x32 e ImageNet-256x256. Il nostro metodo performa complessivamente meglio rispetto alle baseline esistenti attraverso diversi budget di valutazione delle funzioni. Il codice è disponibile all'indirizzo 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