Scaling Online Optimal Transport Minibatches for Flow Matching

Scaling Online Optimal Transport Minibatches for Flow Matching. By Peter Chen and one other, from the University of California, Berkeley. Published on the thirteenth of October. No code repository is linked to it.

Scaling flow models runs into a stability problem, both in training and in integration. In flow matching, that stability depends critically on how source samples are paired with data. Pairing them independently at random introduces a lot of stochasticity. Minibatch optimal transport, following Tong and colleagues, builds the source-target couplings online to reduce that randomness. Two things remain unsettled. On larger, higher-dimensional datasets, small minibatches can limit the benefit, yet coupling much larger minibatches at every iteration is computationally prohibitive. And how optimal transport coupling affects the full stochastic gradient and the integration has not been settled in theory.

The work answers both. On the practical side, it develops a scalable online approximation to entropic optimal transport, which allows much larger source-target couplings while keeping the benefits of optimal transport. The point of the approximation is to break the trade described above, so that the coupling can grow with the dataset instead of being held small by its cost. On the theoretical side, the aim is to explain why the couplings help at all.

The theory closes the gap for both unregularized and entropic optimal transport, so it covers the exact coupling and the approximate one the method relies on. It proves strict reductions, relative to independent pairing, in the variance of the full stochastic gradient and in integration error. It further characterizes how those improvements strengthen as the online coupling minibatch grows, which is exactly the quantity the scalable approximation makes it possible to increase. Lower gradient variance speaks to stability in training. Lower integration error speaks to stability when samples are generated. Those are the two kinds of stability the problem started from.

Experiments on two kinds of training check these predictions: single-cell gene interpolation flows, and image flows. In both, larger online couplings consistently improve training stability and final generation performance, which is the direction the theory predicted.

Theory and method point the same way here. The proofs say bigger couplings should mean lower gradient variance and lower integration error, and the approximation is what makes bigger couplings affordable, so the experiments can test whether the predicted gains appear.