Expected Batch Optimal Transport Plans and Consequences for Flow Matching

arXiv:2605.12174 · cs.LG, math.PR · Submitted 2026-05-12 · Read on arXiv

Listen

Transcript

Introduction to the show: ident: AI Radio. Generated commentary on the latest Artificial Intelligence papers.

Tom: Today's paper: "Expected Batch Optimal Transport Plans and Consequences for Flow Matching".

Jane: Detailed Research Summary:

Tom: First, who's behind it and why it matters.

Title and authors: Tom: Let's talk about the title and who wrote this. It’s "Expected Batch Optimal Transport Plans and Consequences for Flow Matching". The authors are Samuel Boïté, Julie Delon, Kimia Nadjahi, all from ENS Paris.

Jane: So what does this actually mean for us? It sounds like they're giving us a mathematical way to handle the complexity of optimal transport in these huge generative models.

Lu: They are formalizing this coupling as the expected batch OT plan pi k, which is obtained by averaging empirical OT plans over independent minibatches of size k. This helps us understand how the population-level coupling induced by repeated minibatch OT works out.

Meng: I’m still stuck on what this implies for running things at scale. Does this mean we can just use a bigger batch and be done? Or is there something else that breaks at scale?

Lalam: The paper shows that they establish large-batch consistency for pi k, meaning as the batch size k gets bigger, this approximation gets closer to the true optimal plan pi.

The paper's summary: Tom: The summary of this paper is pretty direct: they study pi k, and in the semidiscrete setting relevant to flow matching, they derive rates for both the transport-cost bias and the convergence of pi k to the OT plan.

Jane: Okay, so they quantify how much error we can expect in two places: how far off our cost calculation is from perfect, and how fast pi k moves toward that perfect plan.

Lu: They show a specific convergence rate for the Wasserstein-two distance between these plans, W two(pi k, pi) = O(k-one/two) <ref:2605.12174#pg2>. And they also quantify the transport-cost bias as it relates to this convergence, decaying at a rate of O(k-one) <ref:2605.12174#pg2>.

Meng: O(k-one) for the bias is better than I thought <ref:2605.12174#pg2>. So if we increase our batch size k, we get a better cost estimate faster. That’s actually helpful for training if it means less waste on inaccurate costs.

Lalam: And that convergence rate, the O(k-one/two), shows how quickly pi k gets to the true plan pi when the OT plan is unique, which is a key theoretical result here <ref:2605.12174#pg2,when the OT plan is unique>.

The paper's improvements: Tom: The paper suggests a few major improvements. First, they show that this expected batch OT plan pi k induces a well-posed transport map when used in flow matching because the velocity field derived from this coupling is locally Lipschitz.

Jane: That’s huge for flow matching, right? It means we can actually define a unique flow from the source distribution to the discrete target distribution using this averaged plan, even if individual minibatch plans aren't always rectifiable.

Lu: Exactly. They prove that the velocity field derived from pi k is locally Lipschitz, which guarantees the existence of a unique flow on the interval

zero one), and it admits a terminal map capable of transporting the Gaussian source distribution to the discrete target distribution <ref:2605.12174#pg3>. [Meng: So we can use this instead of dealing with potentially ill-defined couplings from single batches? That’s a practical win for model stability.

Lalam: They also provide a concrete bound on the transport cost of the resulting flow map phi pi k, showing that zero R d phi pi k(x) - x squared d mu(x) - W two two(mu, nu) = O(k-one/two).

Tom: That cost bound ties the batch size directly to the quality of the transport map itself. It gives us a way to measure how much error we can expect in that specific mapping process.

Conclusion: Jane: So, wrapping up this whole discussion on "Expected Batch Optimal Transport Plans and Consequences for Flow Matching," what’s the main point for us right now? It seems they're telling us that using these expected batch plans is a mathematically sound way to approximate the true optimal transport plan.

Lu: They show that increasing the OT batch size k can yield meaningful improvements in generation quality, especially when we are constrained by a low inference budget, which is measured by the number of function evaluations, or NFE.

Meng: I see that connection between batch size and inference budget mentioned. It suggests we have to think about how much computation we spend training versus how much it costs to sample later.

Lalam: They also provide a way to balance this trade-off by showing that increasing the OT batch size k improves low-noise generation, as shown by the observation that increasing k improves performance mainly in the low-NFE regime.

Tom: So, while increasing k helps when NFE is limited, they make it clear that this doesn't get us out of needing accurate inference-time integration. The gains from larger batches have to be weighed against the extra computation during training.

Jane: Right, so the paper suggests that using expected batch OT plans is a robust surrogate, but we still need to be mindful of how our training cost interacts with our sampling budget when we use this specific approach.

Samuel Boíté, Julie Delon, Kimia Nadjahi

ENS Paris · Université Paris Cité · CNRS

cs.LG, math.PR

Submitted: 2026-05-12

Updated: 2026-10-08

Comments: NeurIPS 2026

Code: https://github.com/samuelbx/expected-batch-ot

License: http://arxiv.org/licenses/nonexclusive-distrib/1.0/

Importance score: 88/100

The gist: Detailed Research Summary: Expected Batch Optimal Transport Plans and Consequences for Flow Matching This research paper investigates a crucial surrogate method for solving optimal transport (OT)

Key concepts

Expected Batch OT Plan ($\pi_k$): This is an approximation of the true optimal transport plan derived by averaging empirical plans calculated from many small random minibatches. It serves as a computationally feasible surrogate for the exact, intractable optimal transport plan required in large-scale generative models.
Plan Convergence Rate ($O(k^{-1/2})$): This quantifies how quickly the approximation $\pi_k$ gets closer to the true optimal plan $\pi^\star$ as the batch size $k$ increases. A rate of $O(k^{-1/2})$ means that doubling the batch size reduces the error by a factor of approximately four, indicating good convergence.
Flow Well-Posedness: This concept ensures that when using $\pi_k$ in flow matching, the resulting transport map is mathematically sound. The paper proves that this plan induces a locally Lipschitz velocity field, which guarantees a unique flow exists and can successfully move the source distribution to the target distribution.

Terminology

Flow Matching — 2× in this paper · explained in 10 episodes
Optimal Transport — 2× in this paper · explained in 5 episodes
Consistency — 1× in this paper · explained in 3 episodes
Flow Matching (FM) — 1× in this paper · explained in 2 episodes
Generative Modeling — 1× in this paper · explained in 3 episodes
Optimal Transport (OT) — 1× in this paper · explained in 5 episodes

Summary

Detailed Research Summary: Expected Batch Optimal Transport Plans and Consequences for Flow Matching

This research paper investigates a crucial surrogate method for solving optimal transport (OT) problems on large-scale data, specifically within the context of flow matching (FM) models used in generative modeling. The core contribution is the formalization and rigorous analysis of the expected batch OT plan (pi k), which serves as an approximation to the true optimal transport plan (pi).

Core Methodology and Formalization

The authors address a common practical challenge: solving exact OT on massive datasets is computationally prohibitive. They propose using a surrogate approach by averaging empirical OT plans obtained from independent minibatches of size k. This averaged plan, pi k, is formally defined as the expected batch OT plan, obtained by averaging the empirical OT plans (pi bk) over these random minibatches. The paper establishes the large-batch consistency of this approximation and derives convergence rates for both the transport-cost bias and the convergence of pi k to pi in a semidiscrete setting pertinent to generative modeling.

Key Theoretical Contributions

The paper yields several significant theoretical results that ground minibatch OT on firmer mathematical footing:

Plan Convergence Rate: In Section 3, the authors quantify how quickly the expected batch plan pi k converges to the true optimal plan pi when the OT plan is unique. They establish a convergence rate for the Wasserstein-2 distance between these plans:

W 2(pi k, pi) = O(k-1/2)

Furthermore, they quantify the transport-cost bias as it relates to this convergence: E[W 2 2(mu bk, nu bk)] - W 2 2(mu, nu) decays at a rate of O(k-1).

Flow Well-Posedness: A critical application of this result is its impact on flow matching. The expected batch OT plan pi k induces a well-posed transport map when used in FM. Specifically, they prove that the associated velocity field derived from this coupling is locally Lipschitz, which guarantees the existence of a unique flow on the interval [0, 1). Moreover, this flow admits a terminal map capable of transporting the Gaussian source distribution (mu) to the discrete target distribution (nu).
Transport Cost Rate: For the Gaussian-to-discrete setting, they provide a concrete bound on the transport cost of the resulting flow map phi pi k, showing that:

0 R d phi pi k(x) - x squared d mu(x) - W 2 2(mu, nu) = O(k-1/2)

Practical Implications and Trade-offs

Beyond the theoretical guarantees, the paper provides crucial insights into the practical trade-offs between computational resources:

Batch Size vs. Inference Budget (NFE): The authors rigorously quantify how increasing the OT batch size (k) interacts with inference time, measured by the number of function evaluations (NFE). In a tractable two-atom model, they demonstrate an asymmetry: increasing NFE reduces integration error much faster than increasing k. This suggests that the benefit of larger batches is most pronounced in regimes where the NFE budget is limited.
Asymptotic Comparison: A key theoretical result comparing batch size and inference budget shows that E n,1 about-sqrt n 2/3, which informs how these resources scale together.
Experimental Validation: Empirical results from image experiments (Figure 7) corroborate the theory, showing that increasing k primarily improves metrics like FID in the low-NFE regime.

Conclusion and Limitations

The paper concludes that while increasing the OT batch size (k) can meaningfully improve generation quality in scenarios constrained by a low inference budget (low NFE), it does not eliminate the fundamental need for accurate inference-time integration. The gains achieved through larger batches must be carefully weighed against the additional computational cost incurred during training due to the increased OT computation.

Limitations: The current quantitative results are strictly derived within the semidiscrete setting relevant to generative modeling. A significant open problem identified is understanding how these theoretical bounds transfer when applied to learned velocity fields (as used in practical FM implementations) rather than true, analytical velocity fields. Furthermore, deriving general bounds for E n,k as a function of both n and k remains an area for future research.

**In essence, this paper provides the necessary mathematical foundation to justify using expected batch OT plans (pi k) as a robust and theoretically sound surrogate for exact OT in large-scale flow matching, detailing precisely when and how increasing batch size yields computational benefits.

Improvements for AI systems

The AI system can utilize expected batch optimal transport plans to obtain population-level couplings that are regular enough to define a unique flow from the source to the discrete target, allowing for more stable and well-posed flow matching dynamics.
The system can implement a training strategy where increasing the OT batch size k improves low-noise generation, as shown by the observation that increasing k improves performance mainly in the low-NFE regime, thereby optimizing the trade-off between training computation and inference time for generative models.
The system can leverage the convergence rate of expected batch OT cost, which decays as O(k−1) in non-degenerate settings, to effectively balance training complexity against desired plan accuracy when using minibatch OT surrogates.
The system can perform more accurate sampling during inference by dynamically adjusting the number of function evaluations (NFE) based on the relationship derived from Proposition 4.4, which shows that increasing NFE reduces the integration error much faster than increasing k in a tractable model where the benefit vanishes at larger NFE.
The system can use the terminal flow map to achieve high-fidelity data generation by exploiting its properties: the terminal flow map ϕk partitions R d into regions Γk, j = (ϕk−1)−1(vj), j ∈ 1... M, which converges to the exact OT Laguerre partition as k increases.