条件付けの呪い：条件付きフローベース生成のための最適輸送の分析と改善

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