條件之困：分析並改進基於條件流的最優傳輸生成方法

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

March 13, 2025

作者: Ho Kei Cheng, Alexander Schwing

cs.AI

摘要

在小批量最優傳輸耦合中，無條件流匹配的路徑被拉直。這導致了計算上更為簡便的推理，因為在測試時數值求解常微分方程時，可以使用更少的積分步驟和更簡單的數值求解器。然而，在條件設置下，小批量最優傳輸則顯得不足。這是因為默認的最優傳輸映射忽略了條件，導致在訓練過程中產生條件性偏斜的先驗分佈。相反，在測試時，我們無法訪問這個偏斜的先驗，而是從完整、無偏的先驗分佈中進行採樣。這種訓練與測試之間的差距導致了性能不佳。為彌補這一差距，我們提出了條件最優傳輸C^2OT，它在計算最優傳輸分配時，在成本矩陣中添加了一個條件加權項。實驗表明，這一簡單的修復方法在8gaussians-to-moons、CIFAR-10、ImageNet-32x32和ImageNet-256x256等數據集上，無論是離散還是連續條件下均能有效工作。與現有的基線方法相比，我們的方法在不同的函數評估預算下總體表現更佳。代碼可在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