Scaling Online Optimal Transport Minibatches for Flow Matching

• Independent source-target pairing can give nearby intermediate states different velocity targets, increasing regression ambiguity and stochastic-gradient variance.Source: §1 · p. 1
• Online OT constructs more coherent source-target assignments that induce straighter flows, but small coupling minibatches expose OT to only a limited subset of each distribution.Source: §1 · p. 1
• Dense pairwise interactions restrict standard OT minibatches to hundreds of samples, while FlashSinkhorn enables online entropic coupling contexts to grow from hundreds to thousands.Source: §1 · p. 2
• Entropic regularization changes the coupling, with ε controlling deviation from standard OT and k controlling online coupling scale; strict effects on full-gradient variance and integration error were uncharacterized.Source: §1 · p. 2
• The paper proves that population quadratic OT strictly reduces full stochastic-gradient variance versus independent coupling, with parallel strict reductions in Euler integration error and guarantees that improve with minibatch size.Source: §1 · p. 2
• Prior scaling approaches either precompute couplings before training or require a separate potential-fitting stage, motivating affordable fresh-sample minibatch coupling at every iteration.Source: §1 · p. 2

• A continuous normalizing flow uses a time-dependent, parameterized vector field to assign velocities and transport source samples toward the target distribution.Source: §2 · pp. 2–3
• Flow matching paths are induced by endpoint couplings; each training tuple samples endpoints from the coupling and a time independently and uniformly from [0, 1].Source: §2 · p. 3
• Under squared loss, the optimal deterministic velocity prediction is the conditional mean, while possible velocity targets at an input can still vary.Source: §2 · p. 3
• The total per-sample stochastic-gradient variance is the trace of its covariance matrix and decomposes into target ambiguity at fixed input and conditional-mean variation across inputs.Source: §2 · p. 3

• EOT-CFM uses fresh source and target minibatches to compute an approximate entropic transport plan, then samples source-target pairs from that plan for regression.Source: §3 · p. 3
• Independent coupling pairs samples randomly, optimal OT minimizes quadratic transport cost, and minibatch OT solves the corresponding problem on source and target minibatches.Source: §3 · p. 3
• The coupling minibatch has size k, while the regression batch has size B ≤ k; the entropic plan uses uniform sample mass and a dimension-normalized squared-distance cost.Source: §3 · p. 3
• At each iteration, FlashSinkhorn approximates the target potential with L iterations, selects B source indices, samples target partners, and constructs interpolation points and velocity targets for training.Source: §3 · p. 4
• The theoretical EOT analyses assume exact entropic plans, whereas finite FlashSinkhorn iterations add numerical approximation error not included in the bounds.Source: §3 · p. 4

4.1.1 GRADIENT VARIANCE REDUCTION: STANDARD UNREGULARIZED OT

• Under the stated assumptions, population quadratic OT strictly reduces full stochastic-gradient variance relative to independent coupling.The section also notes that Jensen's inequality extends the result to full-batch SGD variance.Source: §4.1.1 · p. 5

4.1.2 GRADIENT-VARIANCE REDUCTION: ENTROPIC OT

• For sufficiently small ε, entropic OT strictly reduces full stochastic-gradient variance relative to independent coupling under the stated assumptions and Jacobian control.The proof first establishes a reduction in target uncertainty and then translates it into a gradient-variance comparison.Source: §4.1.2 · pp. 5–6

4.1.3 GRADIENT-VARIANCE REDUCTION: LARGER MINIBATCH

• For unregularized and entropic OT, increasing k can reduce full gradient variance; the entropic guarantee requires sufficiently small ε and sufficiently large k under Jacobian bounds.The larger-minibatch results use a finite baseline and establish the comparison for all sufficiently large k.Source: §4.1.2 · p. 5, §4.1.3 · pp. 6–7

4.2 ON INTEGRATION-ERROR REDUCTION

• For one-step distributional Euler error, larger k gives a non-increasing upper bound, while the error converges toward that of the population entropic coupling as k grows.The analysis relates the expected empirical entropic-OT value to integration error.Source: §4.2 · p. 7
• Larger contexts improve EB/PBMC prediction quality and reduce gradient variance; on the synthetic task, increasing k also lowers full-gradient variance and generation MSE.Table 1 reports EB and PBMC results; Figure 1 shows the synthetic-task trend.Source: Table 1 · p. 8, Figure 1 · p. 8

2D-couplings

• In the synthetic 2D problems, larger OT contexts produce more coherent source-target pairings and lower SGD variance, corresponding to tighter projected-gradient distributions.The comparison includes a 128 × 128 OT-CFM minibatch and EOT contexts of 512, 2,048, and 8,192.Source: §5 · p. 8

High-dimensional single-cell interpolation

• On the EB and PBMC interpolation tasks, larger minibatches improve empirical performance and training stability; PBMC results show a clear scaling effect at moderate-thousands contexts robust to entropic regularization intensity.EB uses five whitened principal components, while PBMC uses 100; the EB coupling minibatch is capped by its smallest time point of about 1,900 cells.Source: §5 · pp. 8–9

Image generation

• On CIFAR-100 and ImageNet-64×64, the image-flow evaluation uses FID, and Figure 3 compares generated samples across minibatch sizes and sampling methods.Larger minibatches yield brighter, more vivid samples, particularly at low NFEs.Source: §5 · p. 9, Figure 3 · p. 9
• Across single-cell and image-flow experiments, larger coupling minibatches generally reduce gradient variance and improve performance, with generation-quality gains especially apparent at low NFEs.Table 2 reports FIDs across minibatch sizes and NFEs on CIFAR and ImageNet; Figure 2 shows variance during training and generated samples.Source: Table 2 · p. 9, Figure 2 · p. 9, §5 · p. 10
• On ImageNet, FlashSinkhorn with a 10K coupling uses 41.08 GB peak memory versus 39.96 GB for standard OT-CFM, while full coupling and pair sampling take under 10% of optimization-step time.Measured coupling latency is approximately 0.03 seconds at 4K context and 0.08 seconds at 8K.Source: §5 · p. 10

6 CONCLUSION

• The conclusion presents strict full-gradient variance reduction from independent pairing to population OT, characterizes recovery of this advantage with larger minibatches, and extends the analysis to entropic OT and integration error.Across cell and image flow experiments, scaling online coupling improves training stability and generation quality.Source: §6 · p. 10

A FURTHER RELATED WORKS

• Prior work studies minibatch transport objectives, fixed-input gradient covariance, offline-scaled Sinkhorn couplings, and restricted integration settings, whereas this paper analyzes random online couplings in flow-matching training.Other approaches use dataset-level alignments or learned population maps/plans rather than recomputing couplings from fresh flow-training minibatches.Source: A FURTHER RELATED WORKS · p. 14

B MISSING PROOFS

• The appendix fixes jointly measurable versions of conditional means and covariances; alternative versions agree almost surely, so null-set values do not affect subsequent results.The stated measurability basis is that the relevant spaces are standard Borel.Source: B MISSING PROOFS · p. 15

B.1 PROPOSITION 4.3

• Under realizability, global optimality makes the model velocity equal the conditional mean almost surely, leaving conditional gradient covariance to yield the expected-trace expression.The proof notes that the second loss term is independent of the parameters and the first can attain zero.Source: §B.1 · p. 15

B.2 THEOREM 4.5

• The proof bounds covariance effects using essential lower and upper model sensitivities, with the operator norm giving a sharper upper bound than the Frobenius norm.The sensitivities mπ and Mπ are the essential lower and upper bounds across output directions under the distribution of (t, X_t).Source: §B.2 · pp. 15–16
• Population OT has zero gradient variance, while independent coupling has strictly positive variance under the theorem’s assumptions.The proof uses sensitivity bounds and a variance sandwich to establish the strict comparison.Source: §B.2 · pp. 15–17

B.3 PROPOSITION 4.9

• The interpolation induced by any coupling is admissible in the Benamou–Brenier formulation, so its integrated kinetic energy is at least the squared Wasserstein distance.For population OT, the transport cost equals that distance and the loss at optimality is zero.Source: §B.3 · p. 18

B.4 LEMMA 4.10

• As regularization ε increases, relative entropy to independent coupling is non-increasing while transport cost is non-decreasing.The proof also establishes a strict variance ordering for sufficiently small positive ε.Source: §B.4 · p. 18
• For finite ε, entropic coupling and independent coupling have the same null sets, so a velocity determined by (t, X_t) under one is determined under the other.The proof uses this equivalence together with the zero-loss characterization.Source: §B.4 · pp. 18–19

B.5 REMARK 4.11

• Every weak cluster point of entropic couplings as ε ↓ 0 is a quadratic-OT optimizer, and their transport costs converge to the quadratic-OT cost.The proof passes from weak convergence to convergence of the unbounded quadratic costs using uniform integrability.Source: §B.5 · pp. 19–20
• Quadratic costs are uniformly integrable over all couplings with the fixed marginals, supporting convergence of the entropic costs.This follows from the finite second moments of the two marginals and a de la Vallée–Poussin argument.Source: §B.5 · p. 19
• For sufficiently small positive ε, the entropic coupling’s excess transport cost is below the independent-coupling optimum loss.The proof constructs a competitor with finite relative entropy and uses it to obtain the small-ε bound.Source: §B.5 · pp. 20–21

B.6 THEOREM 4.12

• The proof identifies K_t(x, ·) as a version of the conditional law of X_1 given X_t = x.Under the stated conditions, this conditional law is equivalent to q_1 and has strictly positive covariance trace.Source: §B.6 · p. 22
• Population OT has zero gradient variance, whereas independent coupling has strictly positive variance under the stated density and non-Dirac conditions.The proof establishes positive conditional covariance under independent coupling and uses it to show positive gradient variance.Source: §B.6 · pp. 21–22
• For sufficiently small positive ε, entropic-coupling gradient variance is strictly below independent-coupling variance and remains above population-OT variance.The theorem’s uniform condition supplies a range of ε for which the upper bound falls below independent-coupling variance.Source: §B.6 · p. 22

B.7 PROPOSITION 4.13

• After averaging over endpoint batches, the selected source and target remain distributed according to their respective marginals, so the minibatch plan is a coupling.Conditional on the batches, the OT solution is a permutation and the selected index is uniform.Source: §B.7 · p. 23
• Conditional orthogonality equates the minibatch plan’s expected integrated kinetic energy with its expected empirical endpoint transport cost.The interpolation’s integrated kinetic energy is also bounded below by the population squared Wasserstein distance.Source: §B.7 · pp. 23–24
• The expected squared Wasserstein distance between the empirical endpoint batches converges to the population squared Wasserstein distance.The proof establishes convergence of the empirical measures in W2 and concludes that the cost gap Δ_k tends to zero.Source: §B.7 · p. 23

B.8 THEOREM 4.14

• For every finite k, the OT coupling is equivalent to independent coupling, preserving positive target uncertainty when q0 has an almost-everywhere positive density and q1 is non-Dirac.The conditional law of X1 given Xt is equivalent to q1, and the conditional covariance trace is strictly positive.Source: §B.8 · pp. 24–25
• Under the uniform condition and a positive fixed baseline variance, the optimal gradient variance converges to zero as k grows.The averaged matched-batch update also has variance converging to zero under the same vanishing-product condition.Source: §B.8 · pp. 25–26

B.9 PROPOSITION 4.15

• As the empirical entropic plans grow in size, their expected couplings converge in W2 to the population entropic coupling, and the associated risk satisfies the stated upper bound.The finite-batch entropic coupling remains equivalent to independent coupling, preserving the zero-target-uncertainty characterization.Source: §B.9 · pp. 29–30

B.10 THEOREM 4.16

• For sufficiently large k under the theorem’s joint uniform condition, entropic minibatch OT has lower gradient variance than the prescribed baseline, while its variance remains positive at every finite k.The theorem establishes inequalities for the gradient variance and target-uncertainty risk, then obtains the strict upper comparison beyond a finite threshold.Source: §B.10 · p. 32
• For finite k and ε>0, entropic minibatch couplings remain equivalent to independent coupling, so their target uncertainty and gradient variance are strictly positive.This positivity is used alongside the large-k comparison to conclude variance reduction.Source: §B.10 · p. 32

B.10.1 PROPOSITION 4.18

• Population OT has zero one-step error; population entropic OT has positive error but strictly less than independent coupling.For a finite minibatch, the error converges to zero as the coupling approaches the population OT plan, and is strictly reduced relative to the finite baseline.Source: §B.10.1 · pp. 33–34

B.11 THEOREM 4.19

• For ε>0, the entropic objective’s upper bound is nonincreasing in k and strictly lower for every k≥2 than at k=1.The strict decrease is established first at k=2 and extended by monotonicity.Source: §B.11 · p. 35
• The empirical entropic objective differs from its population counterpart by at most e_k, a bound that tends to zero as the empirical marginals converge.The bound is e_k = 2√Cind E[a_k] + E[a_k^2]/(2d), with a_k the sum of the two empirical-marginal W2 distances.Source: §B.11 · p. 35
• As k grows and ε_k tends to zero, the selected-pair error tends to zero with no rate condition relating k and ε_k.The selected couplings converge to the Brenier coupling in W2.Source: §B.11 · p. 36

Additional 2D examples

• Larger minibatches consistently produce more stable training and lower flow-generation error on the additional 2D ring task.Figure 4 compares independent random coupling, OT-CFM at minibatch size 128, and EOT-CFM contexts up to 8192.Source: C FURTHER EMPIRICAL RESULTS · p. 36, Figure 4 · p. 37

Ablation on ε

• Online couplings show a clear scaling effect with contexts of a few thousand while remaining robust to ε choice, unlike offline million-scale pipelines that require very low ε.The PBMC dataset size permits substantially larger minibatch contexts than EB; EB uses ε = 0.01 and scales to 1.5K.Source: C FURTHER EMPIRICAL RESULTS · p. 37
• Performance peaks around ε = 0.001 and varies little across the swept range on PBMC with an 8K entropic OT context.The ablation averages five seeds; image-flow settings at their largest contexts use ε = 0.002 for CIFAR and ε = 0.001 for ImageNet.Source: C FURTHER EMPIRICAL RESULTS · p. 37, Figure 5 · p. 37