Scaling Online Optimal Transport Minibatches for Flow Matching

Peter Chen; Sergey Levine

cs.LGOct 13, 2026

Scaling Online Optimal Transport Minibatches for Flow Matching

Authors: Peter Chen, Sergey Levine

Organizations: University of California, Berkeley

Abstract

A central challenge in scaling flow models is maintaining stability in training and integration. In flow matching, this stability depends critically on how source samples are paired with data. Independent random pairing introduces substantial stochasticity, motivating minibatch optimal transport (OT) to construct online source-target couplings that reduce this randomness (Tong et al., 2024). However, as applications extend to larger and higher-dimensional datasets, small minibatches can limit these benefits, while scaling the coupling to substantially larger minibatches at every iteration is computationally prohibitive. Moreover, how OT coupling affects the full stochastic gradient and integration remains theoretically unresolved. In this work, we develop a scalable online entropic OT approximation that enables substantially larger source-target couplings while preserving the benefits of OT. We close the theoretical gap for both unregularized and entropic OT by proving strict reductions in full stochastic-gradient variance and integration error relative to independent pairing, and further characterize how these improvements strengthen as the online coupling minibatch grows. Experiments on single-cell gene interpolation-flow and image-flow training validate these predictions, showing that larger online couplings consistently improve training stability and final generation performance.

Figures & tables

Figure 1: Results on a synthetic multimodal task over independent random coupling, OT-CFM using the minibatch size of $128$ , and EOT-CFM with contexts up to $8192$ . (a) Source-target pairings, with green and red lines indicating optimal and suboptimal couplings, respectively. Larger minibatches produce more optimal assignments for subsequent flow training. (b) Stochastic gradients evaluated at a trained flow model with regression batch size $128$ and projected onto common two-dimensional axes; ellipses summarize their spread. (c) Full gradient variance (left axis) and final flow-generation MSE (right axis) averaged over three training seeds, scaling along the OT minibatch size, where $k=0$ denotes independent coupling. Scaling the minibatch size consistently reduces both gradient variance and MSE error, yielding more stable training and better flow generation.

EB			PBMC
Method	Exact $W_{1}$ $(\downarrow)$	$\sigma_{\text{grad},\pi}^{2}(\theta)$ $(\downarrow)$	Method	Exact $W_{1}$ $(\downarrow)$	CellFlow Energy $(\downarrow)$	$\sigma_{\text{grad},\pi}^{2}(\theta)$ $(\downarrow)$
OT-CFM-128	$0.756\pm 0.027$	$0.247\pm 0.047$	OT-CFM-128	$2.612\pm 0.021$	$0.693\pm 0.102$	$1.183\pm 0.343$
EOT-CFM-512	$0.746\pm 0.021$	$0.101\pm 0.017$	EOT-CFM-1024	$2.531\pm 0.017$	$0.435\pm 0.115$	$0.697\pm 0.096$
EOT-CFM-1024	$0.719\pm 0.017$	$0.080\pm 0.012$	EOT-CFM-4096	$2.505\pm 0.028$	$0.385\pm 0.161$	$0.515\pm 0.035$
EOT-CFM-1536	$0.714\pm 0.016$	$0.071\pm 0.011$	EOT-CFM-8192	$2.432\pm 0.015$	$0.329\pm 0.046$	$0.384\pm 0.037$

Table 1: Results on EB and PBMC (mean

\pm

std over 5 seeds), with performance metrics (

W_{1}

, Energy) and gradient variance. Larger contexts improve prediction quality and reduce gradient variance.

	CIFAR- $100$						ImageNet- $64\!\times\!64$
Minibatch context	NFE 25	NFE 50	NFE 100	NFE 1000	Adaptive		Minibatch context	NFE 25	NFE 50	NFE 100	NFE 1000	Adaptive
					NFE	FID						NFE	FID
OT-CFM- $128$	12.85	10.49	9.20	8.00	142.40	7.89	OT-CFM- $128$	16.95	14.88	13.98	13.33	183.91	13.26
EOT-CFM- $2048$	12.77	10.28	8.77	7.79	147.16	7.71	EOT-CFM- $6144$	16.82	14.69	13.80	13.24	188.17	13.20
EOT-CFM- $4096$	12.05	9.80	8.59	7.46	147.80	7.36	EOT-CFM- $8192$	16.44	14.57	13.73	13.12	187.32	13.01
EOT-CFM- $8192$	11.85	9.63	8.44	7.34	146.86	7.25	EOT-CFM- $10240$	16.36	14.47	13.59	13.05	187.19	12.94

Table 2: FIDs across different minibatch size and NFEs on CIFAR and ImageNet dataset.

Panel of Figure 2 — Figure 2: Full SGD gradient variance and qualitative samples. Variance during training on CIFAR-100 ( left ) and ImageNet ( middle ) under different minibatch size, with insets highlighting late-stage behavior. Samples from OT-CFM and EOT-CFM from 128 to 10K, generated at NFE 100. ( right ).

Appendix figures & tables2 assets

Supplementary material from the paper’s appendix.

Appendix

Figure 4: Results on a synthetic ring task over independent random coupling, OT-CFM using the minibatch size of $128$ , and EOT-CFM with contexts up to $8192$ . From left to right, each panel corresponds the flow couplings, the projected training gradients, and the final gradient variance and generation quality under different context size. All the training setup follows from Figure 1 .

Figure 5: Ablation over

\varepsilon

on the PBMC dataset with an 8K entropic OT context, with the exact

W_{1}

error of generated samples (left) and training gradient variance (right) shown. We sweep

\varepsilon

from

\{0.0025,0.0015,0.001,0.0009,0.00075,0.0006,0.0005,0.0004,0.00025,0.0001\}

Flow Matching

Flow matching learns continuous-time vector fields that transport simple noise distributions into data for efficient generative modeling.

See all 656 papers

Expected Batch Optimal Transport Plans and Consequences for Flow Matching
May 12, 2026Samuel Boïté, Julie Delon, Kimia NadjahiFlow Matching Optimal Transport
1DMA, ENS Paris · Université Paris Cité · DI, ENS Paris, CNRS
Optimal Transport Flow Matching by Design
Jun 2, 2026Shimon Malnick, Matan Rusanovsky, Ohad Fried +1Flow Matching One-Step Generative Modeling
1Tel Aviv University · 2Reichman University
One-Sided Quantile Coupling for Flow Matching
Aug 2, 2026Jin-Young Kim, So-Yoon Cho, Hyun-Gyoon KimFlow Matching Image Generation
Independent Researcher · Division of Finance & AI, Hankuk University of Foreign Studies · Dept. of Financial Engineering, Ajou University

May 12, 2026J/K move · Enter open · S save

Expected Batch Optimal Transport Plans and Consequences for Flow Matching

Samuel Boïté, Julie Delon, Kimia Nadjahi

1DMA, ENS Paris · Université Paris Cité · DI, ENS Paris, CNRS

Flow Matching Optimal Transport

Open paper AskPDF arXiv

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.

Abstract

Flow Matching

Explore similar work

Expected Batch Optimal Transport Plans and Consequences for Flow Matching

Optimal Transport Flow Matching by Design

One-Sided Quantile Coupling for Flow Matching