Abstract
Minibatch optimal transport coupling straightens paths in unconditional flowmatching. This leads to computationally less demanding inference as fewerintegration steps and less complex numerical solvers can be employed whennumerically solving an ordinary differential equation at test time. However, inthe conditional setting, minibatch optimal transport falls short. This isbecause the default optimal transport mapping disregards conditions, resultingin a conditionally skewed prior distribution during training. In contrast, attest time, we have no access to the skewed prior, and instead sample from thefull, unbiased prior distribution. This gap between training and testing leadsto a subpar performance. To bridge this gap, we propose conditional optimaltransport C^2OT that adds a conditional weighting term in the cost matrix whencomputing the optimal transport assignment. Experiments demonstrate that thissimple fix works with both discrete and continuous conditions in8gaussians-to-moons, CIFAR-10, ImageNet-32x32, and ImageNet-256x256. Our methodperforms better overall compared to the existing baselines across differentfunction evaluation budgets. Code is available athttps://hkchengrex.github.io/C2OT