Expected Batch Optimal Transport Plans and Consequences for Flow Matching

Abstract: Solving optimal transport (OT) on random minibatches is a common surrogate for exact OT in large-scale learning. In flow matching (FM), this surrogate is used to obtain OT-like couplings that can straighten probability paths and reduce numerical integration cost. Yet, the population-level coupling induced by repeated minibatch OT remains only partially understood. We formalize this coupling as the expected batch OT plan $\overlineπ{k}$, obtained by averaging empirical OT plans over independent minibatches of size $k$. We then establish its large-batch consistency and, in the semidiscrete case relevant to generative modeling, derive rates for both the transport-cost bias and the convergence of $\overlineπ{k}$ to the OT plan. For FM, this yields a population coupling whose induced velocity field is regular enough to define a unique flow from the source to the discrete target. We finally quantify how OT batch size interacts with numerical integration in a tractable two-atom model and in synthetic and image experiments.

• The paper demonstrates that the expected batch OT plan converges to the population-level OT plan with precise rates (e.g., O(k⁻¹/²), O(k⁻¹), and O(k⁻¹/⁴)).
• It shows that the induced velocity field in flow matching is locally Lipschitz, ensuring a unique ODE flow and rectifiability of the generated trajectories.
• The study uncovers tradeoffs between OT batch size and numerical integration budget, revealing diminishing returns in sample quality improvements at high NFE.

Formalization and Analysis of Expected Batch Optimal Transport for Flow Matching

Introduction

The paper "Expected Batch Optimal Transport Plans and Consequences for Flow Matching" (2605.12174) presents a rigorous framework for understanding the population-level behavior of minibatch optimal transport (OT) in large-scale learning, specifically its impact on flow matching (FM) generative models. The research addresses key theoretical gaps about the stochastic bias and plan convergence of minibatch OT surrogates and offers new quantitative results in the semidiscrete setting, with direct implications for the tractability and rectifiability of the induced velocity fields in FM.

Expected Batch OT: Definition, Consistency, and Convergence Rates

A central object in the paper is the expected batch OT plan $T_k$, defined as the averaged empirical OT solution across independently sampled minibatches of size $k$. This formalization resolves ambiguities from previous heuristics and establishes that, under uniqueness of the OT plan, $T_k$ converges to the population-level OT plan $T^*$ as $k\rightarrow\infty$.

The authors derive precise convergence rates in the semidiscrete case (absolutely continuous source, finitely supported target), showing:

• The transport cost bias decays as $O(k^{-1/2})$ under finite moment assumptions and as $O(k^{-1})$ under compact support.
• For Gaussian-to-discrete distributions, the Wasserstein distance between $T_k$ and $T^*$ satisfies $W_2(T_k, T^*) = O(k^{-1/4})$.

These results demonstrate that minibatch OT provides significant improvement over independent couplings and entropic regularizations for small to moderate batch sizes, with empirical evidence indicating that the $k$0 cost-bias rate holds even for Gaussian sources (unbounded support), suggesting possible generalization beyond current theoretical bounds.

Flow Matching with Minibatch OT Couplings: Regularity and Well-Posedness

A major theoretical contribution is the demonstration that the expected batch OT plan induces a velocity field in FM which is locally Lipschitz and defines a unique ODE flow, thus ensuring rectifiability. The induced terminal map partitions $k$1 into regions converging to the Laguerre cells of the semidiscrete OT map as $k$2 increases.

The posterior assignment probabilities concentrate more rapidly over time as $k$3 or ambient dimension $k$4 grows, and the assignment is fully controlled by minibatch statistics. This posterior concentration behavior clarifies how minibatch OT increasingly mimics exact OT trajectories as batch size increases, supporting straighter probability paths.

Computational Tradeoffs: OT Batch Size Versus Numerical Integration Budget

The paper quantitatively distinguishes the effect of OT batch size ($k$5) versus inference budget (number of function evaluations, NFE), both theoretically and empirically:

• In tractable models (e.g., two-atom Gaussian-to-discrete), increasing NFE reduces integration error much faster (stretched exponential decay) than increasing $k$6 (polynomial decay).
• Empirically, increasing $k$7 improves generated sample quality (FID) primarily in the low-NFE regime; at high NFE, further increases in $k$8 yield diminishing returns or can even degrade performance.

This reveals a fundamental asymmetry: while higher $k$9 can compensate for coarse numerical integration, frequent function evaluations during inference dominate in minimizing trajectory error.

Empirical Validation: Image Generation Experiments

On CIFAR-10 and SVHN datasets, the expected batch OT cost decreases steadily for large batch sizes (up to $T_k$0). Generation experiments confirm improved performance for low NFE with larger $T_k$1, but added OT batch size incurs practical overhead in training time. The results provide a comprehensive grid over batch sizes and NFE, enabling granular analysis of training-inference tradeoffs.

Implications and Future Directions

The formalism and convergence results lay theoretical foundations for minibatch OT and its application in FM, sharpening interpretability and guiding practical tuning. The findings imply:

• Minibatch OT can effectively straighten probability flows and reduce integration cost, but with diminishing returns beyond moderate batch sizes, especially as the inference budget increases.
• Rectifiability and well-posedness properties can transfer to alternative coupling surrogates (e.g., entropic, sliced OT), suggesting avenues for future extension.
• There is no universal modulus for cost-to-plan convergence: sufficient additional hypotheses on distributions are required.

The precise analysis motivates further research into broader source-target configurations, behavior with neural velocity field approximations, and computational bounds for other numerical schemes.

Conclusion

This paper rigorously characterizes the expected batch OT plan as an underpinning of minibatch OT in flow matching, providing consistency, convergence rates, and demonstrating practical tradeoffs between training and inference. The theoretical advancements support principled deployment of OT surrogates in generative modeling and illuminate their limitations in large-scale settings, with implications for developing efficient, well-posed generative flows.