Scaling Online Optimal Transport Minibatches for Flow Matching

Preprint. Under review.

S CALING O NLINE O PTIMAL T RANSPORT M INIBATCHES FOR F LOW M ATCHING Peter Chen

Sergey Levine

University of California, Berkeley

arXiv:2610.18838v1 [cs.LG] 13 Oct 2026

A BSTRACT 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.

1

I NTRODUCTION

Much of the recent progress in generative AI has been driven by scale. Autoregressive language models benefit from a standardized learning objective: next-token prediction under a softmax distribution. Once this objective is fixed, much of the challenge of scaling shifts to systems and optimization. Flow models introduce an additional dimension: their supervision is constructed along a generative path from noise to data, and the resulting regression targets depend on how source noise samples are coupled with target data samples. Thus, scaling flow models requires consideration not only of how the model is optimized, but also of how its training targets are constructed. The choice of such source-target coupling is therefore part of the learning problem itself. In flow matching, each training example pairs a source sample from the noise distribution with a target data sample and regresses the velocity along the resulting interpolation path. Under the standard independent coupling, these pairings are random. Consequently, nearby intermediate states can arise from different endpoint pairs and receive substantially different velocity targets, increasing ambiguity in the regression problem and the variance of stochastic gradients. A natural way to reduce this randomness is to determine the source-target correspondence during training. Optimal transport (OT) then provides a natural solution by constructing a coupling between source and target samples within each minibatch during training. Such online coupling minimizes the transport cost between the two empirical distributions, yielding more coherent source-target assignments that induce straighter flows and reduce the ambiguity of independent pairing (Pooladian et al., 2023; Tong et al., 2024). However, as flow matching is applied to larger and more heterogeneous datasets, such as high-dimensional single-cell trajectories and large-scale image datasets, a fixed small coupling minibatch can become increasingly restrictive, where each minibatch only exposes OT to a limited subset of the source and target distributions, constraining the candidate correspondences available at each iteration. As the underlying data distribution grows in diversity and complexity, it therefore becomes natural to scale the OT minibatch accordingly. 1

Preprint. Under review.

Still, scaling the minibatch size is computationally challenging, where dense pairwise interactions in standard OT restrict practical minibatches to hundreds of samples. Nevertheless, recent advances in scalable entropic OT make substantially larger contexts accessible. In particular, FlashSinkhorn (Ye et al., 2026) implements Sinkhorn iterations with FlashAttention-style HBM kernels without materializing the full coupling matrix in high-bandwidth memory. To this end, we propose a simple, tractable approximation by replacing online OT coupling with scalable entropic OT, enabling coupling contexts to grow from hundreds to thousands of samples during training. The price of this tractability is that entropic regularization changes the coupling itself. The regularization strength ε controls this deviation from standard OT, while the minibatch size k controls the scale of the online coupling. Understanding whether the stability benefits of OT persist therefore requires characterizing the joint effects of regularization and coupling scale. Moreover, this question remains unresolved even for standard unregularized minibatch OT. Existing analyses motivate OT through straighter paths and reduced conditional uncertainty (Pooladian et al., 2023; Boïté et al., 2026; Zhang et al., 2026; Chambolle & Hertrich, 2026), but do not provide a general strict characterization of its effect on the full stochastic-gradient variance or the induced numerical integration error. By making larger online couplings feasible, we determine whether their stability benefits survive the approximation required for scaling. Our contributions can be summarized below: Contributions. We provide a unified and principled characterization of how online OT coupling affects optimization and numerical integration. On the optimization side, we characterize the full stochastic-gradient variance by (i) proving that population quadratic OT strictly reduces this variance relative to independent coupling, (ii) showing that unregularized minibatch-OT variance reduces as the minibatch size k grows, and (iii) extending same guarantees to entropic-regularized couplings. On the integration side, we establish parallel strict reductions in Euler integration error and characterize how these guarantees improve with scaled minibatch size. These results provide a better theoretical understanding for using and scaling entropic OT as a replacement to the standard online OT couplings. Related works. OT-based conditional flow matching (OT-CFM) improves both training and generation through better source-target pairings (Tong et al., 2024; Pooladian et al., 2023). Given the difficulty to scale OT-CFM, recent works scale the minibatch context by separating pairing from training: Zhang et al. (2026) precomputed million-scale couplings offline, which requires buffering enough pairs to cover the entire training run before optimization begins; semidiscrete approximation methods offer better efficiency by reusing fitted dataset-level potentials, but need a separate potentialoptimization stage (Mousavi-Hosseini et al., 2026) or face non-convergence when scaled to large unconditional datasets (Kong et al., 2026). Therefore, we return to the original challenge, by finding an affordable and stable method for scaling the minibatch that flow matching recomputes from fresh samples at every training iteration. The existing theory has so far addressed separate pieces. Tong et al. (2024) connect population OT to dynamic OT, while Pooladian et al. (2023) study conditional gradient covariance and asymptotic BatchOT straightness. Finite-batch theory remains limited: Zhang et al. (2026) characterize uncertainty in the resulting OT coupling, while Boïté et al. (2026) analyze its integration benefit only in a restricted two-atom model; more recent work also studies full parameter-gradient variance in restricted Gaussian models (Reu et al., 2025) and integration error through the geometry of a prescribed velocity field (Tao & Choi, 2026). Thus, no prior result establishes that replacing independent pairing with minibatch OT strictly reduces full stochastic-gradient variance while inducing a general integration advantage; we defer further related work to Appendix A.

2

P RELIMINARIES

We introduce flow matching (FM) and coupling definitions used throughout the later analyses. Continuous normalizing flows. Let q0 , q1 ∈ P2 (Rd ) denote the source (noise) and target (imagedata) distributions, respectively. A time-dependent vector field vθ : [0, 1] × Rd → Rd , parameterized by θ ∈ Rp , assigns an instantaneous velocity to each time-state pair (t, x), thereby defining a flow from the source to the target. Specifically, a continuous normalizing flow (CNF) (Lipman et al., 2023) transports source samples through the flow map d ϕt (x) = vθ (t, ϕt (x)), ϕ0 (x) = x, dt 2

Preprint. Under review.

with the goal that (ϕ1 )# q0 ≈ q1 . Flow matching trains this vector field from paths induced by a coupling of the source and target endpoints. Coupling-based FM. Let π be a coupling of q0 and q1 . Each FM training tuple is generated first by sampling (X0 , X1 ) ∼ π and then independently sampling t ∼ U ([0, 1]). Unless stated otherwise, we take all expectations Eπ [·] and covariances Covπ (·) over the randomness induced by the sampling procedure associated with coupling π. For simplicity, we denote the complete stochastic FM training tuple by ξ π = (t, X0 , X1 ) and its associated variables, including intermediate interpolation point Xtπ and pathwise velocity target U π , by Xtπ = (1 − t)X0 + tX1 ,

U π = X1 − X0 .

At a fixed model input (t, Xtπ ), different endpoint pairs may produce different velocity targets. Given one sampled FM training tuple ξ π , its loss is ℓπ (θ; ξ π ) = ∥vθ (t, Xtπ ) − U π ∥22 . Under this squared loss, the optimal deterministic velocity prediction at input (t, x) is the conditional mean, defined by uπ (t, Xtπ ) := Eπ [U π | t, Xtπ ]. The conditional covariance, Covπ (U π | t, Xtπ ), then measures how much the possible velocity targets still vary at that input; the population FM objective further averages ℓπ (θ; ξ π ) over all possible training tuples, Lπ (θ) := Eπ [ℓπ (θ; ξ π )]. Let Jθ (t, x) := ∇θ vθ (t, x) ∈ Rd×p , the gradient over ℓπ (θ; ξ π ) can be expressed as
Gπ (θ) := ∇θ ℓπ (θ; ξ π ) = 2Jθ (t, Xtπ )⊤ vθ (t, Xtπ ) − U π . Whenever Eπ [∥Gπ (θ)∥22 ] < ∞, we can derive the total variance of the per-sample stochastic gradient induced by sampling FM training tuples under coupling π:   2 σgrad,π (θ) := Tr{Covπ (Gπ (θ))} = Eπ ∥Gπ (θ) − Eπ Gπ (θ)∥22 . Here, Covπ (Gπ (θ)) ∈ Rp×p is the covariance matrix of the per-sample stochastic gradient, whereas 2 σgrad,π (θ) is its trace, measuring the total sample-to-sample gradient variance across FM training tuples under coupling π. The law of total covariance further gives us the following decomposition, 2 σgrad,π (θ) = Eπ [Tr{Covπ (Gπ (θ) | t, Xtπ )}] + Tr{Covπ (Eπ [Gπ (θ) | t, Xtπ ])} .

(1)

The first term is the gradient noise from target ambiguity at a fixed model input, and the second term is the variation of the conditional mean gradient across model inputs.

3

S CALABLE OT FOR F LOW M ATCHING

We refer to our method as conditional flow matching with entropic OT (EOT-CFM). Our approach applies a larger context of source and target samples to construct better couplings, with FlashSinkhorn (Ye et al., 2026) enabling efficient computation at scale. Specifically, at each training iteration, we draw fresh source and target minibatches, compute an approximate entropic transport plan, and sample source-target pairs from this plan to form the regression batch. We consider three coupling constructions as background: (i) independent coupling πind , which pairs samples from q0 and q1 randomly; (ii) optimal OT coupling π ⋆ , which minimizes Eπ [∥X0 − X1 ∥22 ] over all noise and data points1 ; (iii) in practice, minibatch OT solves the corresponding problem on source and target minibatches of size k, where we denote this empirical plan as π bk . To formalize our method, we denote k as the number of samples to be paired in the coupling minibatch and B ≤ k as the regression batch used per model update. Given independent samples xi ∼ q0 and yj ∼ q1 , we write Cij = ∥xi − yj ∥22 /(2d), with dimension normalization that preserves the unregularized OT matching. Assigning uniform mass 1/k to each sample, we aim to solve the following entropic transport plan: 
π bk,ε = arg min ⟨C, P ⟩ + ε KL P 11⊤ /k 2 . (2) P ≥0, P 1=1/k P ⊤ 1=1/k

1 If q0 ≪ Ld , Brenier’s theorem (Brenier, 1991) states that this coupling is unique, with the deterministic transport map: π ⋆ = (id, T )# q0 , where T = ∇φ is q0 -almost everywhere for some convex function φ.

3

Preprint. Under review.

Algorithm 1 Continuous normalizing flow with online entropic couplings (EOT-CFM) Require: Minibatch size k, batch size B, regularization ε > 0, Sinkhorn iterations L, training steps T . 1: for t = 1, . . . , T do 2: Sample noise and data from source xi ∼ q0 and target yj ∼ q1 with minibatch size k 3: Approximate the target dual potential ge of Eq. (2) with L FlashSinkhorn iterations 4: Select B distinct source indices I1 , . . . , IB uniformly 5: for b = 1, . . . , B do    g ej −CI j k b 6: Compute (e rIb ,j )kj=1 = softmax ε j=1
7: Sample Jb ∼ Cat (e rIb ,j )kj=1 and tb ∼ U([0, 1]) 8: end for PB 2 1 b 9: Compute L(θ) = Bd b=1 ∥vθ (tb , Xtb ,b ) − Ub ∥2 b 10: Update θ with the optimizer using ∇θ L(θ) 11: end for

The regularization ε > 0 balances between the cost and entropy. Specifically, the exact row conditional is P(J = j | I = i) = k · π bk,ε,ij , which normalizes probabilities from the approximate target potential ge. Algorithm 1 then constructs the interpolation points Xtb ,b and velocity targets Ub defined in §2. Without materializing the full transport plan or invoking a separate JAX process, FlashSinkhorn solves the minibatch coupling directly within the PyTorch training loop using tiled cost evaluations and dual potentials, with iteration count L that controls the numerical approximation. We also note that our later EOT analyses consider exact entropic transport plans. For computational efficiency, FlashSinkhorn uses a finite number of Sinkhorn iterations, introducing additional numerical approximation error that is not included in the bounds below.

4

M AIN R ESULTS

4.1

O N STOCHASTIC GRADIENT VARIANCE

We first study the benefit of scaling OT minibatch size through its effect on stochastic-gradient variance. The closest existing result we are aware of is Pooladian et al. (2023, Lemma 3.2), which analyzes the conditional gradient variance of minibatch OT at model input (t, x), corresponding to the fixed-input component in Eq. (1). However, it remains a gap to prove that OT coupling strictly reduces the stochastic-gradient variance relative to independent coupling. Therefore, in §4.1.1, we begin by establishing a strict gradient-variance reduction by filling the gap between optimal coupling π ⋆ and independent coupling πind . In §4.1.2, we extend this result to entropic-OT couplings. Finally, in §4.1.3, we establish that scaling to larger minibatch size could further reduce gradient variance for both unregularized and entropic OT couplings. 4.1.1

G RADIENT VARIANCE REDUCTION : STANDARD UNREGULARIZED OT

We first state two standard assumptions used in theoretical analyses of flow matching and neural differential equations: realizability of the population target field and global optimality of the population objective (Islam et al., 2026; Khan et al., 2026; Baheri, 2026). Since the population regression target depends on the coupling π, we note that these assumptions are imposed coupling-wise; all the detailed math derivations and proofs for the later theoretical results are deferred to Appendix B. Assumption 4.1 (Realizability). For every coupling π under comparison, the model class V = {vθ : θ ∈ Rp } contains the coupling-induced conditional mean velocity, i.e., there exists θπ ∈ Rp such that vθπ (t, Xtπ ) = uπ (t, Xtπ ) almost surely. Assumption 4.2 (Optimality). For every coupling π under comparison, θπ⋆ is a global minimizer of Lπ , the map θ 7→ vθ (t, Xtπ ) is differentiable at θπ⋆ almost surely, and Eπ [∥Gπ (θπ⋆ )∥22 ] < ∞. From an optimization perspective, gradient noise at an optimum provides a reference quantity for analyzing SGD convergence (Gower et al., 2019). In flow matching setting, evaluating gradients at population optima also removes the variance contribution from imperfect fitting of the conditional-mean velocity, providing a common basis for comparing gradient noise across couplings 4

Preprint. Under review.

and motivating later flow matching analyses for using optimality conditions to study stochastic gradients (Vaitl et al., 2022b;a; Kim et al., 2025; Zhang et al., 2025; Vaitl & Klein, 2025). Proposition 4.3. Under Assumptions 4.1 and 4.2, along with a trainable additive output bias b ∈ Rd under OT-CFM architecture (see Fact B.3 for detail), we have    2 σgrad,π (θπ⋆ ) = 4 Eπ Tr Jθπ⋆ (t, Xtπ )⊤ Covπ (U π | t, Xtπ )Jθπ⋆ (t, Xtπ ) . (3) Thus, at the population optimum, conditional target covariance is mapped into parameter-gradient space by the model Jacobian. Its average magnitude also satisfies     Eπ Tr{Covπ (U π | t, Xtπ )} = Eπ ∥U π − uπ (t, Xtπ )∥22 . To obtain an explicit positive lower bound for independent coupling, we directly follow from the Assumption (A2’) in Pooladian et al. (2023), using the following technical sufficient condition: Assumption 4.4. There exist a compact convex set Ω ⊂ Rd with nonempty interior, α ∈ (0, 1], γ > 0, and α-Hölder functions fi : Ω → [γ, ∞) such that, for i ∈ {0, 1} and every Borel A ⊆ Rd , Z qi (A) = fi (x) dx, supp(qi ) = Ω. A∩Ω

Finally, combining the preceding variance identity with this common-support condition, we establish a strict gradient-variance gap between standard unregularized OT and independent coupling. Theorem 4.5. Let π ⋆ = (id, T )# q0 be the unregularized OT coupling and πind be the independent coupling under the quadratic cost. Suppose Assumptions 4.1, 4.2, and 4.4 hold for both couplings and define ωd = Ld (B(0, 1)), then, for every x̄ ∈ Ω◦ and r > 0 such that B(x̄, 2r) ⊂ Ω, we have 4dγ 2 ωd2 2d+2 r > 0. (4) d+2 Remark 4.6. Given Theorem 4.5, one can bound full-batch SGD variance by single-sample variance via Jensen inequality (see Eq. (30)), which extends similar later results also to batch-level. 2 2 ⋆ σgrad,π (θπ⋆ ind ) − σgrad,π ⋆ (θπ ⋆ ) ≥ ind

4.1.2

G RADIENT- VARIANCE REDUCTION : ENTROPIC OT

We next study the influence of entropic regularization on gradient variance and introduce following setup inherited from previous minibatch flow matching analysis. Definition 4.7 (Uniformly bounded parameter Jacobians). Following the maximum-Jacobian setup in the flow-matching gradient-variance bound of Pooladian et al. (2023, Lemma 3.2), a coupling family Π has uniformly bounded parameter Jacobians if there exists M < ∞ such that max

t∈[0,1], x∈Rd

∥Jθπ⋆ (t, x)∥2op ≤ M,

∀ π ∈ Π.

(5)

Remark 4.8. We apply Definition 4.7 to obtain a common Jacobian bound M for all couplings considered in later theorems, allowing us to translate target-uncertainty bounds into full stochasticgradient variance comparisons. We also note that such Jacobian bound control is widely used in neural-network optimization (Oymak & Soltanolkotabi, 2019, Assumption 2) and recent flowmatching theory analysis (Zeng & Yan, 2025, Proposition 3). We use πε to denote the population entropic-OT coupling between q0 and q1 :   1 πε ∈ arg min Eπ [∥X0 − X1 ∥22 ] + ε KL(π∥πind ) . π∈Π(q0 ,q1 ) 2d

(6)

To compare its variance to that under independent coupling, we first relate its irreducible flow matching regression error to its quadratic transport cost, where W2 denotes the 2-Wasserstein distance: Proposition 4.9. Under Assumptions 4.1 and 4.2, every coupling π satisfies 0 ≤ Lπ (θπ⋆ ) = Eπ [Tr{Covπ (U π | t, Xtπ )}] ≤ Eπ [∥X1 − X0 ∥22 ] − W22 (q0 , q1 ).

(7)

Thus, the population loss at optimality is exactly the average target uncertainty remaining at a fixed model input. As an intermediate step, we establish Lemma 4.10 to show that sufficiently small entropic regularization strictly reduces this uncertainty comparing to independent coupling. 5

Preprint. Under review.

Lemma 4.10. Suppose Assumptions 4.1 and 4.2 hold for πind , π ⋆ , and every πε under comparison, and assume Lπind (θπ⋆ ind ) > 0, which holds whenever q0 has a strictly positive Lebesgue density and q1 is not a Dirac mass, including the standard choice q0 = N (0, Id ). Then, for any ε0 > 0 that satisfies Eπε0 [∥X1 − X0 ∥22 ] − W22 (q0 , q1 ) < Lπind (θπ⋆ ind ),

(8)

every 0 < ε ≤ ε0 gives us Eq. (9), and such ε0 is guaranteed to exist by Remark 4.11, Lπ⋆ (θπ⋆ ⋆ ) < Lπε (θπ⋆ ε ) < Lπind (θπ⋆ ind ).

(9)

Remark 4.11. We note that such ε0 always exists. For any π e ∈ Π(q0 , q1 ) with 0 < KL(e π ∥πind ) < ∞ and Eπe [∥X1 − X0 ∥22 ] − W22 (q0 , q1 ) < Lπind (θπ⋆ ind ), Eq. (9) holds whenever 0<ε<

Lπind (θπ⋆ ind ) + W22 (q0 , q1 ) − Eπe [∥X1 − X0 ∥22 ] . 2d KL(e π ∥πind )

(10)

A coupling π e satisfying these conditions always exists; we defer its proof to Appendix B.5. Finally, we transform Lemma 4.10 into the stochastic-gradient variance reduction. Theorem 4.12. Suppose the conditions of Lemma 4.10 hold. Then, every 0 < ε ≤ ε0 satisfies that 2 0 < σgrad,π (θπ⋆ ε ) ≤ 4{Eπε [∥X1 − X0 ∥22 ] − W22 (q0 , q1 )} ess sup ∥Jθπ⋆ε (t, Xtπε )∥2op . ε

(11)

Moreover, under Definition 4.7 over all sufficiently small ε, there exists ε1 ∈ (0, ε0 ] such that every 0 < ε ≤ ε1 satisfies 2 ⋆ 2 ⋆ 2 ⋆ σgrad,π (12) ⋆ (θπ ⋆ ) < σgrad,π (θπ ) < σgrad,π (θπ ). ε ε ind ind 4.1.3

G RADIENT- VARIANCE REDUCTION : LARGER MINIBATCH

We now move to characterize how the gradient variance changes when minibatch size is scaled. Similar to the previous proof structure, we first derive the result upon standard unregularized OT and then extend to the results to entropic OT. Our math idea follows from by connecting the cost residue in Eq. (13) to the velocity uncertainty and then bridging this uncertainty to the gradient variance. Unregularized OT. For each minibatch, we independently draw k source endpoints from q0 and k target endpoints from q1 and denote the distribution of this selected pair by πk , with π1 = πind . Since πk is a coupling between q0 and q1 , its quadratic transport cost cannot be smaller than the cost of optimal coupling π ∗ , i.e., W22 (q0 , q1 ). Therefore, we can represent this gap as the excess cost residue ∆k : ∆k := Eπk [∥X1 − X0 ∥22 ] − W22 (q0 , q1 ). (13) Clearly, ∆k ≥ 0, and we apply this cost residual in an interesting way by directly connecting it to the velocity-target uncertainty that drives stochastic gradient variance: Proposition 4.13. Given q0 , q1 ∈ P2 (Rd ), then without any assumptions, for every k ≥ 1, we have k→∞

Eπk [Tr{Covπk (U πk | t, Xtπk )}] ≤ ∆k −−−−→ 0.

(14)

Indeed, Proposition 4.13 plays an important role by bridging the gap between the velocity-target uncertainty and OT minibatch, which allows us to further connect it to stochastic gradient variance: Theorem 4.14. Suppose that q0 , q1 ∈ P2 (Rd ) and Assumptions 4.1 and 4.2 hold for every πk . Then, 2 0 ≤ σgrad,π (θπ⋆ k ) ≤ 4∆k ess sup ∥Jθπ⋆ (t, Xtπk )∥2op . k k

(15)

Given the setup in Definition 4.7, the right-hand side of Eq. (15) converges to zero, so for any fixed 2 k0 with σgrad,π (θπ⋆ k ) > 0, Eq. (16) strictly holds for all sufficiently large k: k 0

0

2 2 σgrad,π (θπ⋆ k ) < σgrad,π (θπ⋆ k ). k k 0

6

0

(16)

Preprint. Under review.

Entropic OT. For each k ≥ 1 and ε > 0, we denote πk,ε as the distribution of the endpoint pair generated by the entropic minibatch procedure in §3 and π bk,ε as the random empirical plan within minibatches. Similarly, π1,ε = πind . For simplicity, we abbreviate the average target uncertainty in Proposition 4.13 by R(π) = Eπ [Tr{Covπ (U π | t, Xtπ )}]. Shown in Proposition 4.15, we first characterize the effect of increased k using the similar cost residue as in Eq. (13): Proposition 4.15. Suppose q0 , q1 ∈ P2 (Rd ). For every fixed ε > 0, πk,ε → πε in W2 , and moreover, lim sup R(πk,ε ) ≤ R(πε ) ≤ Eπε [∥X1 − X0 ∥22 ] − W22 (q0 , q1 ).

(17)

k→∞

Note that πk,ε and πind are mutually absolutely continuous for every finite k, so R(πk,ε ) = 0 if and only if R(πind ) = 0. Increasing k at a fixed ε indeed approaches the population entropic coupling, but this convergence alone cannot give the complete reduction result, which also depends on the coupling-specific model Jacobian. So, we next establish the gradient variance reduction under larger minibatch size: Theorem 4.16. Fix a finite baseline k0 ≥ 1 and consider that q0 , q1 ∈ P2 (Rd ), R(πind ) > 0, and Assumptions 4.1 and 4.2 hold for every πk,ε under comparison. Then, for every k ≥ 1 and ε > 0, we have π 2 4R(πk,ε ) ≤ σgrad,π (θπ⋆ k,ε ) ≤ 4R(πk,ε ) ess sup ∥Jθπ⋆ (t, Xt k,ε )∥2op . (18) k,ε k,ε

Under Definition 4.7 over all sufficiently large k and sufficiently small ε, there exists ε̄ > 0 such that, for every fixed 0 < ε ≤ ε̄, there is a finite batch threshold K > k0 for which every k ≥ K satisfies 2 2 σgrad,π (θπ⋆ k,ε ) < σgrad,π (θπ⋆ k ,ε ). k,ε k ,ε 0

4.2

(19)

0

O N I NTEGRATION -E RROR R EDUCTION

Beyond optimization, OT couplings induce straighter probability paths, but how this geometry translates into numerical benefits remains unclear. We study integration error at the distribution level by comparing the output of Euler discretization with the target distribution. Definition 4.17 (One-step distributional Euler error). For a coupling π ∈ Π(q0 , q1 ), one Euler step from t = 0 produces X0 + Eπ [X1 − X0 | X0 ] = Eπ [X1 | X0 ]. Its distributional integration error is err(π) := W22 ((Eπ [X1 | X0 = ·])# q0 , q1 ) ,

(20)

i.e., the Wasserstein distance between the one-step output distribution and the target distribution q1 . Shown in Proposition 4.18, we first establish the relationship of distributional integration error across optimal, entropic, and independent couplings: Proposition 4.18. Suppose q0 , q1 ∈ P2 (Rd ), q0 ≪ Ld , and q1 is not a Dirac mass. Let π ⋆ = (id, T )# q0 be the population quadratic-OT coupling, let πk be the unregularized OT plan at batch size k, with π1 = πind , and for every fixed ε > 0, let πε be the population entropic coupling. Then, err(π ∗ ) < err(πε ) < err(πind ). Finally, based on Proposition 4.18, we establish the error reduction over larger k by connecting the expected empirical EOT value in Eq. (21) to the integration error. 
 Sk,ε := E ⟨C, π bk,ε ⟩ + ε KL π bk,ε 11⊤ /k 2 . (21) Theorem 4.19. Under the conditions of Proposition 4.18, for every fixed ε > 0 and k ≥ 2, we have 2

err(πk,ε ) ≤ E[∥X1 − E[X1 ]∥22 ] −

d2 (S1,ε − Sk,ε ) < err(πind ). E[∥X0 − E[X0 ]∥22 ]

(22)

k→∞

The upper bound in Eq. (22) is also non-increasing in k, and err(πk,ε ) −−−−→ err(πε ). Let εk > 0 be any sequence with εk → 0, for every fixed finite baseline k0 , there exists k1 > k0 such that all k > k1 satisfies err(πk,εk ) < err(πk0 ,εk ). (23) Figure 3 illustrates the empirical benefit of Eq. (23) through a qualitative comparison of generated samples at the same NFE; we defer further explanations to §5. 7

Preprint. Under review.

(a) Sampled couplings Independent

(b) Projected training gradients OT-128

Independent

(c) Gradient variance and flow MSE

OT-128 Gradient variance Flow MSE

EOT-8192

EOT-512

EOT-8192

10 2

Flow MSE

CFM loss + backprop

EOT-512

2 grad, ( )

10

1

10 3 0.1 10 4

Source

Target

Optimal coupling

Suboptimal coupling

0

Projected gradient

2048

Minibatch size k

8192

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. Table 1: Results on EB and PBMC (mean ± std over 5 seeds), with performance metrics (W1 , Energy) and gradient variance. Larger contexts improve prediction quality and reduce gradient variance. EB

PBMC

Method

Exact W1 (↓)

2 σgrad,π (θ) (↓)

OT-CFM-128

0.756 ± 0.027

0.247 ± 0.047

OT-CFM-128

2.612 ± 0.021

0.693 ± 0.102

1.183 ± 0.343

EOT-CFM-512 EOT-CFM-1024 EOT-CFM-1536

0.746 ± 0.021 0.719 ± 0.017 0.714 ± 0.016

0.101 ± 0.017 0.080 ± 0.012 0.071 ± 0.011

EOT-CFM-1024 EOT-CFM-4096 EOT-CFM-8192

2.531 ± 0.017 2.505 ± 0.028 2.432 ± 0.015

0.435 ± 0.115 0.385 ± 0.161 0.329 ± 0.046

0.697 ± 0.096 0.515 ± 0.035 0.384 ± 0.037

5

Method

Exact W1 (↓)

CellFlow Energy (↓)

2 σgrad,π (θ) (↓)

E XPERIMENTS

We study three classes of experiments to show the empirical benefits of scaling minibatch size: (i) first, we consider classical 2D synthetic datasets, where the effects of minibatch scaling can be clearly visualized and interpreted; (ii) we then work on single-cell interpolation flows; apart from the traditional medium-dimensional dataset used in OT-CFM, we include a large-scale, high-dimensional dataset to demonstrate the benefits of scaled minibatches; (iii) finally, we evaluate image flow training on CIFAR-100 (Krizhevsky, 2009) and ImageNet-64×64 (Deng et al., 2009). 2D-couplings. We consider two synthetic problems: a symmetric problem with 64 Gaussian modes arranged on an 8 × 8 grid, and a double-ring problem with 8 modes on each ring. In both problems, the source and target distributions are separated by a horizontal translation, and the ground-truth coupling matches points with the same mode label. We compare independent coupling, OT-CFM with a 128 × 128 minibatch, and EOT with context sizes n ∈ {512, 2048, 8192}, using squared Euclidean cost and training a three-layer MLP for 6,000 updates with batch size 128 with AdamW. Results are shown in Figure 1. Specifically, larger OT contexts produce increasingly coherent sourcetarget pairings, as shown by the reduction in red suboptimal crossing connections and the increase in green optimal coupling connections. This improvement directly translates into lower SGD variance, with tighter projected-gradient distributions that shows scaling the OT context provides cleaner pairings and substantially better training stability. More 2D examples are provided in Appendix C. High-dimensional single-cell interpolation. We evaluate on both the embryoid body (EB) benchmark (Moon et al., 2019) used in OT-CFM (Tong et al., 2024) and a larger, higher-dimensional PBMC cytokine-response dataset (Oesinghaus et al., 2025), processed by CellFlow (Klein et al., 2025). For EB, it represents 16,819 cells from five time points (the smallest point has 1900 cells) using five whitened principal components and independently split each time point into 80%/10%/10% training, validation, and test sets. We leave out each intermediate time point t ∈ {1, 2, 3}, train a three-layer MLP of width 64 on the remaining transitions for at most 1,000 epochs. For PBMC, we use 90,452 PBS control cells and 20,471 IFN-β-stimulated cells from Donor 1, represent each cell using 100 whitened principal components, and apply the same data split. We train 8

Preprint. Under review.

Table 2: FIDs across different minibatch size and NFEs on CIFAR and ImageNet dataset. CIFAR-100 Minibatch context

NFE 25

NFE 50

NFE 100

ImageNet-64×64

NFE 1000

Adaptive NFE

FID

Minibatch context

NFE 25

NFE 50

NFE 100

NFE 1000

Adaptive 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 EOT-CFM-4096 EOT-CFM-8192

12.77 12.05 11.85

10.28 9.80 9.63

8.77 8.59 8.44

7.79 7.46 7.34

147.16 147.80 146.86

7.71 7.36 7.25

EOT-CFM-6144 EOT-CFM-8192 EOT-CFM-10240

16.82 16.44 16.36

14.69 14.57 14.47

13.80 13.73 13.59

13.24 13.12 13.05

188.17 187.32 187.19

13.20 13.01 12.94

0.016

CIFAR-100 Ind-coupling OT-coupling 2048 OT-coupling 4096 OT-coupling 8192

Full SGD gradient variance

0.014 0.012 0.010

0.0020

0.008

0.0010

ImageNet

0.006

Ind-coupling OT-coupling 6144 OT-coupling 8192 OT-coupling 10240

0.005 0.0010

0.004

0.0008

0.0015

0.0006

0.003

300

0.006

350

EOT-CFM 6K

0.0004

400

400

450

500

0.002

0.004

EOT-CFM 10K

0.001

0.002 0.000

Qualitative samples OT-CFM 128

0

50

100

150

200

250

Training steps (k)

300

350

400

0.000

0

100

200

300

400

Training steps (k)

500

Figure 2: Full SGD gradient variance and qualitative samples. Variance during training on CIFAR100 (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). NFE=4

NFE=8

NFE=16

Adaptive

NFE=4

NFE=8

NFE=16

Adaptive

NFE=4

NFE=8

NFE=16

Adaptive

2K

4K

8K

Figure 3: Generated images from initial noise. Rows correspond to minibatch sizes of 2K, 4K, and 8K, while columns compare Euler sampling at NFE 4, 8, and 16 with adaptive Dopri5. Smaller coupling minibatches tend to produce darker, less vivid images, whereas larger minibatches yield brighter samples with more vibrant colors, particularly at low NFEs. a three-layer MLP of width 256 for 100,000 steps. Both experiments use AdamW with learning rate 10−3 , weight decay 10−5 , and a regression batch size of 128, and generate cells using 100 Euler steps. For EB, we report exact W1 averaged over three held-out time points and five seeds. For PBMC, we compare 2,048 generated and held-out target cells and report exact assignment W1 , CellFlow energy, and gradient variance, averaged over five seeds. Results are presented in Table 1. Since the smallest time point from EB contains around 1900 cells, the maximum minibatch cannot be scaled as PBMC. For both dataset, larger minibatch leads to both better empirical performance and more stable training, where the improvement consistently gets better with larger minibatch size. With the result deferred to Appendix C, we also use the high dimensional PBMC dataset to study the ablation over ε, moderate thousands-level context already gives a clear effect of scaling and remains robust to the entropic regularization intensity. Image generation. We evaluate on CIFAR-100 and ImageNet-64×64. For CIFAR, all methods follow the same U-Net architecture and optimization hyperparameters from OT-CFM’s CIFAR training script (Tong et al., 2024). We evaluate EMA checkpoints using fixed-step and adaptive Dopri5 integration, and FIDs are computed with CleanFID with 50,000 generated samples. For ImageNet training, we follow the unconditional ADM-style U-Net configuration of Pooladian et al. (2023). We train for 500,000 updates with a global batch size of 1,024 across 8×NVIDIA H100 GPUs and a constant learning rate of 10−4 . Similar to CIFAR-100, all coupling methods share the same architecture and training configuration, and we generate 50,000 samples from the EMA checkpoint and report FID against the canonical ADM ImageNet-64 reference statistics. Evaluation results are shown in Table 2, and we further present in-training gradient variance and qualitative image generation results in Figure 2. Experiments on both single-cell and image flow matching 9

Preprint. Under review.

show that larger coupling minibatches generally reduce gradient variance and improve performance, consistent with our entropic analysis in §4.1.3. Figure 3 further shows that the gains in generation quality are particularly better at low NFEs. Together, these findings suggest that scaling the entropic coupling context benefits both optimization and sampling, combining lower gradient noise during training with better generation quality under limited integration budgets. Thanks to its bandwidth-efficient design and compatibility with PyTorch/Triton over HBM GPUs, FlashSinkhorn integrates naturally into efficient FM training. Increasing the minibatch size incurs only a modest increase in peak GPU memory. On ImageNet, the standard OT-CFM baseline reaches a peak memory usage of 39.96 GB, while FlashSinkhorn with a coupling size of 10K uses only 41.08 GB. For runtime, even at the largest context 10K, full OT coupling and pair sampling wall time account for less than 10% of the total time per gradient optimization step. The measured coupling latency is approximately 0.03 seconds at the 4K context and around 0.08 seconds at the 8K context. In addition to the benefit brought by entropic regularization, FlashAttention-style engineering also offers efficient GPU kernels, the native PyTorch integration removes the need for PyTorch-JAX tensor interchange and coordination between separate execution and memory-management runtimes, keeping coupling computation and pair sampling within the flow-matching optimization loop.

6

C ONCLUSION

In this work, we study online coupling as a fundamental scaling axis for flow matching and fill the gap by establishing a strict reduction in full stochastic-gradient variance from independent pairing to population OT, characterize how larger online minibatch couplings recover this advantage, and extend the analysis to entropically regularized OT together with parallel guarantees for integration error. These results motivate a scalable approximation based on large online entropic couplings. Across both cell flow and image flow experiments, scaling the online coupling consistently improves training stability and generation quality. In the end, our results provide both a principled theoretical framework for rigorously analyzing the concrete effect brought by the more informative couplings and simple yet practical foundation for scaling flow matching.

10

Preprint. Under review.

R EFERENCES Jason Altschuler, Jonathan Niles-Weed, and Philippe Rigollet. Near-linear time approximation algorithms for optimal transport via sinkhorn iteration. In NeurIPS, 2017. Brandon Amos, Lei Xu, and J Zico Kolter. Input convex neural networks. In ICML, 2017. Arip Asadulaev, Aleksandr Korotin, Vage Egiazarian, Petr Mokrov, and Evgeny Burnaev. Neural optimal transport with general cost functionals. In ICLR, 2024. Ali Baheri. Geometry-grounded flow matching on compact manifolds. In ICLR 2026 Workshop on Geometry-grounded Representation Learning and Generative Modeling, 2026. Mathias Beiglböck, Marcel Nutz, and Nizar Touzi. Complete duality for martingale optimal transport on the line. 2017. Jean-David Benamou and Yann Brenier. A computational fluid mechanics solution to the mongekantorovich mass transfer problem. Numerische Mathematik, 84(3):375–393, 2000. Samuel Boïté, Julie Delon, and Kimia Nadjahi. Expected batch optimal transport plans and consequences for flow matching. arXiv preprint arXiv:2605.12174, 2026. Yann Brenier. Polar factorization and monotone rearrangement of vector-valued functions. In Communications on Pure and Applied Mathematics, 1991. Sergio Calvo-Ordonez, Matthieu Meunier, Alvaro Cartea, Christoph Reisinger, Yarin Gal, and Jose Miguel Hernandez-Lobato. Weighted conditional flow matching. arXiv preprint arXiv:2507.22270, 2025. Antonin Chambolle and Johannes Hertrich. Limit points of reflow with minibatch optimal transport. arXiv preprint arXiv:2608.07042, 2026. Peter Chen, Yue Xie, and Qingpeng Zhang. SICNN: Sparsity-induced input convex neural network for optimal transport. In Optimization for Machine Learning, 2024. Ricky TQ Chen, Yulia Rubanova, Jesse Bettencourt, and David K Duvenaud. Neural ordinary differential equations. In NeurIPS, 2018. Marco Cuturi. Sinkhorn distances: Lightspeed computation of optimal transport. In NeurIPS, 2013. Jia Deng, Wei Dong, Richard Socher, Li-Jia Li, Kai Li, and Li Fei-Fei. Imagenet: A large-scale hierarchical image database. In CVPR, 2009. Stephan Eckstein and Marcel Nutz. Quantitative stability of regularized optimal transport and convergence of sinkhorn’s algorithm. SIAM Journal on Mathematical Analysis, 54(6):5922–5948, 2022. Jiaojiao Fan, Amirhossein Taghvaei, and Yongxin Chen. Scalable computations of wasserstein barycenter via input convex neural networks. In ICML, 2021. Kilian Fatras, Younes Zine, Rémi Flamary, Rémi Gribonval, and Nicolas Courty. Learning with minibatch wasserstein: asymptotic and gradient properties. In AISTATS, 2020. Rémi Flamary, Nicolas Courty, Alexandre Gramfort, Mokhtar Z Alaya, Aurélie Boisbunon, Stanislas Chambon, Laetitia Chapel, Adrien Corenflos, Kilian Fatras, Nemo Fournier, et al. POT: Python optimal transport. In JMLR, 2021. Alberto González-Sanz, Marcel Nutz, and Austin J Stromme. Finite-sample bounds for regularized optimal transport. arXiv preprint arXiv:2606.25947, 2026. Robert Mansel Gower, Nicolas Loizou, Xun Qian, Alibek Sailanbayev, Egor Shulgin, and Peter Richtárik. Sgd: General analysis and improved rates. In ICML, 2019. Mohammad Mohaiminul Islam, Thijs P Kuipers, Sharvaree Vadgama, Coen de Vente, Afsana Khan, Clara I Sánchez, and Erik J Bekkers. Longitudinal flow matching for trajectory modeling. In AISTATS, 2026. 11

Preprint. Under review.

Fairoz Nower Khan, Nabuat Zaman Nahim, Ruiquan Huang, Haibo Yang, and Peizhong Ju. Flow matching for offline reinforcement learning with discrete actions. arXiv preprint arXiv:2602.06138, 2026. Myunsoo Kim, Donghyeon Ki, Seong-Woong Shim, and Byung-Jun Lee. Adaptive non-uniform timestep sampling for accelerating diffusion model training. In CVPR, 2025. Dominik Klein, Jonas Simon Fleck, Daniil Bobrovskiy, Lea Zimmermann, Sören Becker, Alessandro Palma, Leander Dony, Alejandro Tejada-Lapuerta, Guillaume Huguet, Hsiu-Chuan Lin, et al. Cellflow enables generative single-cell phenotype modeling with flow matching. bioRxiv, 2025. Lingkai Kong, Molei Tao, Yang Liu, Bryan Wang, Jinmiao Fu, Chien-Chih Wang, and Huidong Liu. Alignflow: Improving flow-based generative models with semi-discrete optimal transport. In ICLR, 2026. Nikita Kornilov, Petr Mokrov, Alexander Gasnikov, and Alexander Korotin. Optimal flow matching: Learning straight trajectories in just one step. In NeurIPS, 2024. Alexander Korotin, Vage Egiazarian, Arip Asadulaev, Alexander Safin, and Evgeny Burnaev. Wasserstein-2 generative networks. In ICLR, 2021a. Alexander Korotin, Lingxiao Li, Justin Solomon, and Evgeny Burnaev. Continuous wasserstein-2 barycenter estimation without minimax optimization. In ICLR, 2021b. Alexander Korotin, Daniil Selikhanovych, and Evgeny Burnaev. Neural optimal transport. In ICLR, 2023. Alex Krizhevsky. Learning multiple layers of features from tiny images. Tech Report, 2009. Tianyi Lin, Nhat Ho, and Michael Jordan. On efficient optimal transport: An analysis of greedy and accelerated mirror descent algorithms. In ICML, 2019. Tianyi Lin, Nhat Ho, Marco Cuturi, and Michael I Jordan. On the complexity of approximating multimarginal optimal transport. In JMLR, 2022. Yaron Lipman, Ricky TQ Chen, Heli Ben-Hamu, Maximilian Nickel, and Matt Le. Flow matching for generative modeling. In ICLR, 2023. Elon Litman. Scaled-dot-product attention as one-sided entropic optimal transport. arXiv preprint arXiv:2508.08369, 2025. Elon Litman and Gabe Guo. You need better attention priors. In ICML, 2026. Xingchao Liu, Chengyue Gong, and Qiang Liu. Flow straight and fast: Learning to generate and transfer data with rectified flow. In ICLR, 2023. Ashok Makkuva, Amirhossein Taghvaei, Sewoong Oh, and Jason Lee. Optimal transport mapping via input convex neural networks. In ICML, 2020. Hugo Malamut and Maxime Sylvestre. Convergence rates of the regularized optimal transport: Disentangling suboptimality and entropy. In SIAM Journal on Mathematical Analysis, 2025. Kevin R Moon, David Van Dijk, Zheng Wang, Scott Gigante, Daniel B Burkhardt, William S Chen, Kristina Yim, Antonia van den Elzen, Matthew J Hirn, Ronald R Coifman, et al. Visualizing structure and transitions in high-dimensional biological data. Nature biotechnology, 2019. Gilles Mordant. The entropic optimal (self-) transport problem: Limit distributions for decreasing regularization with application to score function estimation. arXiv preprint arXiv:2412.12007, 2024. Alireza Mousavi-Hosseini, Stephen Zhang, Michal Klein, and Marco Cuturi. Flow matching with semidiscrete couplings. In ICLR, 2026. Marcel Nutz and Johannes Wiesel. Entropic optimal transport: Convergence of potentials. Probability Theory and Related Fields, 2022. 12

Preprint. Under review.

Lukas Oesinghaus, Sören Becker, Larsen Vornholz, Efthymia Papalexi, Joey Pangallo, Amir Ali Moinfar, Jenni Liu, Alyssa La Fleur, Maiia Shulman, Simone Marrujo, et al. A single-cell cytokine dictionary of human peripheral blood. BioRXiv, 2025. Samet Oymak and Mahdi Soltanolkotabi. Overparameterized nonlinear learning: Gradient descent takes the shortest path? In ICML, 2019. Aram-Alexandre Pooladian, Heli Ben-Hamu, Carles Domingo-Enrich, Brandon Amos, Yaron Lipman, and Ricky TQ Chen. Multisample flow matching: Straightening flows with minibatch couplings. In ICML, 2023. Teodora Reu, Sixtine Dromigny, Michael Bronstein, and Francisco Vargas. Gradient variance reveals failure modes in flow-based generative models. In NeurIPS, 2025. Danilo Rezende and Shakir Mohamed. Variational inference with normalizing flows. In ICML, 2015. Philippe Rigollet and Austin J Stromme. On the sample complexity of entropic optimal transport. The Annals of Statistics, 2025. Chenxi Tao and Seung-Kyum Choi. On the role of strain and vorticity in numerical integration error for flow matching. arXiv preprint arXiv:2605.06680, 2026. Alexander Tong, Kilian Fatras, Nikolay Malkin, Guillaume Huguet, Yanlei Zhang, Jarrid RectorBrooks, Guy Wolf, and Yoshua Bengio. Improving and generalizing flow-based generative models with minibatch optimal transport. In TMLR, 2024. Lorenz Vaitl and Leon Klein. Path gradients after flow matching. In NeurIPS, 2025. Lorenz Vaitl, Kim A Nicoli, Shinichi Nakajima, and Pan Kessel. Gradients should stay on path: better estimators of the reverse-and forward kl divergence for normalizing flows. Machine Learning: Science and Technology, 2022a. Lorenz Vaitl, Kim Andrea Nicoli, Shinichi Nakajima, and Pan Kessel. Path-gradient estimators for continuous normalizing flows. In ICML, 2022b. Jonathan Weed. An explicit analysis of the entropic penalty in linear programming. In COLT, 2018. Felix X-F Ye, Xingjie Li, An Yu, Ming-Ching Chang, Linsong Chu, and Davis Wertheimer. Flashsinkhorn: Io-aware entropic optimal transport on gpu. In ICML, 2026. Weili Zeng and Yichao Yan. Flow matching in the low-noise regime: Pathologies and a contrastive remedy. arXiv preprint arXiv:2509.20952, 2025. Juyan Zhang, Rhys Newbury, Xinyang Zhang, Tin Tran, Dana Kulic, and Michael Burke. Why heuristic weighting works: A theoretical analysis of denoising score matching. arXiv preprint arXiv:2508.01597, 2025. Stephen Y. Zhang, Alireza Mousavi-Hosseini, Michal Klein, and Marco Cuturi. On fitting flow models with large sinkhorn couplings. In TMLR, 2026.

13

Preprint. Under review.

A

F URTHER R ELATED W ORKS

We walk through more related works and discuss some related theoretical results from both the FM and OT perspectives: Flow matching and minibatch couplings. Normalizing flows and neural ODEs provide the modeling background for flow matching (Rezende & Mohamed, 2015; Chen et al., 2018), while flow matching and rectified flow introduce regression-based objectives for learning continuous transport dynamics (Lipman et al., 2023; Liu et al., 2023). More broadly, minibatch Wasserstein objectives have been studied as stochastic surrogates for population OT, including their asymptotic and gradient properties (Fatras et al., 2020); these results concern the minibatch transport objective itself rather than using its random plan as flow-matching supervision. Tong et al. (2024) introduce OT-CFM and connect the population OT coupling to dynamic OT, whereas Pooladian et al. (2023) analyze gradient covariance conditional on a fixed model input and establish asymptotic BatchOT straightness and transport-cost convergence. Very recent work characterizes limit points of Reflow combined with fixed-size minibatch OT through cyclical monotonicity and straightness properties (Chambolle & Hertrich, 2026); we note that this concerns repeated Reflow rather than the optimization or integration behavior of a single online minibatch-CFM objective. Scaling OT couplings for flow matching. Zhang et al. (2026) scales Sinkhorn couplings to millions of samples through an offline construction, where paired samples are precomputed independently of model training and served through an augmented data loader. We note that their finite-batch lower bound motivates increasing coupling size, but it does not establish a strict OT-versus-independent comparison or a full stochastic-gradient result. Boïté et al. (2026) characterize the expected exact-OT batch plan and its transport-cost convergence, while their integration result is restricted to an exact velocity field in a one-dimensional two-atom model, and they also show that when the population w OT plan is unique, π bk converges weakly to π ⋆ , i.e., π bk −−−−→ π ⋆ . However, we note that this k→∞

coupling-level convergence does not imply any level of convergence for stochastic gradient variance. Semidiscrete methods instead construct dataset-level noise-data alignments, avoiding a fresh discrete OT solve at every iteration (Mousavi-Hosseini et al., 2026; Kong et al., 2026). Calvo-Ordonez et al. (2025) replace minibatch OT with Gibbs-weighted independent pairs, avoiding an explicit OT solve but potentially altering the marginals, so larger minibatches provide more weighted pairs rather than scaling a larger candidate pool for optimizing source-target assignments. Our setting instead retains online couplings recomputed from fresh minibatches, with clear scaling benefits merely to thousand-level of minibatch size. Scalable optimal transport. Brenier’s characterization motivates representing quadratic-cost OT maps as gradients of convex potentials (Brenier, 1991), for which input-convex neural networks provide a convenient parameterization (Amos et al., 2017). This approach has been used to learn OT maps, plans, and Wasserstein barycenters (Makkuva et al., 2020; Korotin et al., 2021a; Fan et al., 2021; Korotin et al., 2021b; 2023; Chen et al., 2024; Asadulaev et al., 2024). Within flow matching, optimal flow matching constrains the velocity parameterization to recover the population quadratic-OT displacement directly (Kornilov et al., 2024). These methods learn a population-level map or plan through model structure or a separate neural optimization problem, rather than analyzing the random coupling recomputed from each flow-training minibatch. Optimization and integration theory. Reu et al. (2025) analyze the variance of the full parameter gradient in restricted Gaussian transport models under specific analytic velocity parameterizations, and show that low gradient variance can also coexist with undesirable deterministic pairings. We note that these results do not associate independent pairing with random online minibatch OT in a general model class. For numerical integration, Tao & Choi (2026) decompose the velocity Jacobian into strain and vorticity and establish favorable Euler properties for exact OT displacement fields. This analysis concerns the geometry of a given velocity field rather than an integration-error ordering induced by finite minibatch couplings. Thus, neither line characterizes how full-gradient variance and integration error jointly change with coupling size and entropic regularization. Computational entropic OT. Entropic regularization makes discrete OT amenable to Sinkhorn matrix scaling (Cuturi, 2013). Subsequent work establishes approximation-complexity guarantees for 14

Preprint. Under review.

Sinkhorn, Greenkhorn, accelerated variants, and multimarginal extensions (Altschuler et al., 2017; Lin et al., 2019; 2022), while POT provides widely used implementations (Flamary et al., 2021). FlashSinkhorn targets GPU efficiency by fusing Sinkhorn updates and avoiding materialization of the full coupling matrix in high-bandwidth memory (Ye et al., 2026), along with prior results with scaled dotproduct attention for onesided EOT problem (Litman, 2025) and later extension with improved learning prior (Litman & Guo, 2026). Statistical and limiting theory of entropic OT. For finite linear programs, Weed (2018) give explicit non-asymptotic guarantees for convergence of entropically penalized solutions toward unregularized optima. At the population level, EOT potentials, plans, and costs converge toward their unregularized counterparts as ε → 0 (Nutz & Wiesel, 2022; Malamut & Sylvestre, 2025). Statistical results establish parametric sample-complexity bounds at fixed regularization (Rigollet & Stromme, 2025), limit distributions when regularization decreases with sample size (Mordant, 2024), and finite-sample bias and variance bounds with explicit dependence on regularization and intrinsic dimension (González-Sanz et al., 2026).

B

M ISSING P ROOFS

We first clarify a measurability point used throughout the proofs. Since [0, 1]×Rd and Rd are standard Borel spaces, a regular conditional law of U π given (t, Xtπ ) exists. We fix jointly measurable versions of its conditional mean and covariance. Any two such versions agree almost surely, so their values on null sets do not affect the results below. B.1

P ROPOSITION 4.3

Proof. For any θ with Lπ (θ) < ∞, conditioning on (t, Xtπ ) and expanding the squared loss gives     Lπ (θ) = Eπ ∥vθ (t, Xtπ ) − uπ (t, Xtπ )∥22 + Eπ ∥U π − uπ (t, Xtπ )∥22 . The second term is independent of θ. By realizability, the first term can attain zero. Therefore, global optimality implies vθπ⋆ (t, Xtπ ) = uπ (t, Xtπ ) almost surely, and substituting this into the per-sample gradient yields Gπ (θπ⋆ ) = 2Jθπ⋆ (t, Xtπ )⊤ (uπ (t, Xtπ ) − U π ). So, we have Eπ [Gπ (θπ⋆ ) | t, Xtπ ] = 0. Thus, the conditional-mean term in Eq. (1) vanishes. Conditional on (t, Xtπ ), the Jacobian is fixed, so the conditional gradient covariance is 4Jθ⊤π⋆ Covπ (U π | t, Xtπ )Jθπ⋆ , with every Jacobian evaluated at (t, Xtπ ). Taking its expected trace would give us Eq. (3). B.2

T HEOREM 4.5

Before proving Theorem 4.5, we derive two supporting facts. First, we define the essential lower and upper model sensitivities at the coupling-specific population optimum by mπ and Mπ , respectively: mπ := ess inf λmin (Jθπ⋆ (t, Xtπ )Jθπ⋆ (t, Xtπ )⊤ ),

Mπ := ess sup ∥Jθπ⋆ (t, Xtπ )∥2op .

Here, the essential extrema are taken under the distribution of (t, Xtπ ), and together, they can yield the following variance sandwich bound: Proposition B.1 (Variance sandwich). Under Assumptions 4.1 and 4.2, every coupling π under comparison satisfies   2 4mπ Eπ ∥U π − uπ (t, Xtπ )∥22 ≤ σgrad,π (θπ⋆ ), (24) and, if Mπ < ∞, we have   2 σgrad,π (θπ⋆ ) ≤ 4Mπ Eπ ∥U π − uπ (t, Xtπ )∥22 . (25) Proof. For notation simplicity, we denote C = Covπ (U π | t, Xtπ ) ⪰ 0 and A = Jθπ⋆ Jθ⊤π⋆ ⪰ 0 (with the Jacobian evaluated at (t, Xtπ )). Cyclicity of the trace gives Tr(Jθ⊤π⋆ CJθπ⋆ ) = Tr(CA), and λmin (A) Tr(C) ≤ Tr(CA) ≤ ∥A∥op Tr(C) = ∥Jθπ⋆ ∥2op Tr(C). 15

Preprint. Under review.

By the definition of the essential infimum, λmin (A) ≥ mπ almost surely. Applying the lower trace inequality pointwise and taking expectations in Eq. (3), we have 2 σgrad,π (θπ⋆ ) = 4Eπ [Tr(CA)] ≥ 4Eπ [mπ Tr(C)] = 4mπ Eπ [Tr(C)].

Similarly, when Mπ < ∞, the definition of the essential supremum implies that ∥A∥op = ∥Jθπ⋆ (t, Xtπ )∥2op ≤ Mπ

almost surely.

Applying the upper trace inequality pointwise and taking expectations yields that 2 σgrad,π (θπ⋆ ) = 4Eπ [Tr(CA)] ≤ 4Eπ [Mπ Tr(C)] = 4Mπ Eπ [Tr(C)].

Finally, since uπ (t, Xtπ ) = Eπ [U π | t, Xtπ ], the definition of conditional covariance and the tower property gives us:    Eπ [Tr(C)] = Eπ Eπ ∥U π − uπ (t, Xtπ )∥22 t, Xtπ   = Eπ ∥U π − uπ (t, Xtπ )∥22 . Substituting this identity into the preceding lower and upper bounds proves Eqs. (24) and (25), respectively. We also provide some further discussions over Jacobian norm and essential extrema in Remark B.2. Remark B.2. At each (t, x), λmin (Jθπ⋆ (t, x)Jθπ⋆ (t, x)⊤ ) and ∥Jθπ⋆ (t, x)∥2op measure the weakest and strongest model sensitivities across output directions. Accordingly, mπ and Mπ are their essential lower and upper bounds under the distribution of (t, Xtπ ) (also note that the essential extrema are suitable here, because Assumption 4.2 requires differentiability only almost everywhere and are therefore unaffected by arbitrary values on null sets). We also note that the operator norm in Mπ gives a sharper bound than the Frobenius norm, while playing the same role as the maximum Jacobian factor in Pooladian et al. (2023, Lemma 3.2). Second, we show the nondegeneracy supplied by an affine output layer with trainable bias. Fact B.3. Over standard OT-CFM’s MLP structure of width-d output layer trainable bias, let θ = (η, b), where b ∈ Rd , one always have vη,b (t, x) = v̄η (t, x) + b. We note that is explicitly shown in Tong et al. (2024, Appendix E.2) and its code implementation2 . With Fact B.3, we then establish Lemma B.4: Lemma B.4. Under Assumptions 4.1 4.2, we have Jθπ⋆ (t, x)Jθπ⋆ (t, x)⊤ ⪰ Id almost everywhere,  and 2 ⋆ π so mπ ≥ 1 and σgrad,π (θπ ) ≥ 4Eπ ∥U − uπ (t, Xtπ )∥22 . Proof. Since ∇b vη,b (t, x) = Id , the full parameter Jacobian at the coupling-specific optimum is   Jθπ⋆ (t, x) = ∇η v̄ηπ⋆ (t, x) Id , Jθπ⋆ (t, x)Jθπ⋆ (t, x)⊤ = ∇η v̄ηπ⋆ (t, x)∇η v̄ηπ⋆ (t, x)⊤ + Id ⪰ Id . The conclusion follows from Eq. (24). We now prove the theorem: Proof of Theorem 4.5. We first consider population OT. Choose a Borel representative of the Brenier map and a Borel set D of full q0 -measure on which T = ∇φ. Monotonicity of gradients of convex functions gives us ⟨T (x) − T (y), x − y⟩ ≥ 0, x, y ∈ D. For t < 1, we write Ft (x) = (1 − t)x + tT (x). If x ̸= y, then we have ⟨Ft (x) − Ft (y), x − y⟩ ≥ (1 − t)∥x − y∥22 > 0, so Ft is injective on D. It’s clear that Borel map Φ(t, x) = (t, Ft (x)) from [0, 1) × D into [0, 1) × Rd is injective. By the Lusin-Souslin theorem, its image is Borel and its inverse is Borel as well. Extending the inverse by a fixed point off its image yields a Borel map Ψ such that, almost surely, we have ⋆

X0 = pr2 [Ψ(t, Xtπ )]. 2

See code at commit cacd4dc, where it uses torch.nn.Linear with bias=True by default

16

Preprint. Under review.

⋆

⋆

Since t = 1 has probability zero, this recovers X0 from (t, Xtπ ) almost surely, so U π = T (X0 )−X0 ⋆ ⋆ ⋆ is measurable with respect to σ(t, Xtπ ) (and also Covπ⋆ (U π | t, Xtπ ) = 0 almost surely). ⋆

⋆

⋆

⋆

Assumptions 4.1 and 4.2 then give vθπ⋆ ⋆ (t, Xtπ ) = uπ (t, Xtπ ) = U π almost surely. Therefore, 2 ⋆ every per-sample gradient vanishes almost surely and σgrad,π ⋆ (θπ ⋆ ) = 0. We next consider independent coupling and first extend f0 , f1 by zero outside Ω. For fixed t, the map (x0 , x1 ) 7→ (x, u) = ((1 − t)x0 + tx1 , x1 − x0 ) has inverse x0 = x − tu, x1 = x + (1 − t)u, and absolute Jacobian determinant one. Therefore, (Xt , U ) has joint density ht (x, u) = f0 (x − tu)f1 (x + (1 − t)u). For notation clarity, we define Z Z ρt (x) = ht (x, u) du, mt (x) = ρt (x)−1 u ht (x, u) du, and, when ρt (x) > 0, we can derive that Covπind (U

πind

| t, Xtπind = x) = ρt (x)−1

Z

(u − mt (x))(u − mt (x))⊤ ht (x, u) du,

which sets mt (x) = 0 when ρt (x) = 0. Thus, ρt is the density of Xt , while note that mt and above are versions of its conditional mean velocity and covariance, respectively. Next, we fix x̄ ∈ Ω◦ and r > 0 with B(x̄, 2r) ⊂ Ω. For x ∈ B(x̄, r), u ∈ B(0, r), and t ∈ [0, 1], it’s clear that both x − tu and x + (1 − t)u lie in Ω. Therefore, we have ht (x, u) ≥ γ 2 , ρt (x) ≥ γ 2 ωd rd > 0. For every a ∈ Rd and such (t, x), we can further derive Z πind ⊤ πind 2 ρt (x)a Covπind (U | t, Xt = x)a ≥ γ (a⊤ u − a⊤ mt (x))2 du B(0,r)

γ 2 ωd rd+2 ∥a∥22 , ≥ d+2 where the second step follows because the cross term integrates to zero by symmetry and R uu⊤ du = ωd rd+2 Id /(d + 2). Therefore, we can finally derive B(0,r) ρt (x) Covπind (U πind | t, Xtπind = x) ⪰

γ 2 ωd rd+2 Id . d+2

The joint law of (t, Xt ) has density ρt (x) dt dx. Taking the trace and integrating over t ∈ [0, 1] and x ∈ B(x̄, r) eventually yields   dγ 2 ωd2 2d+2 Eπind ∥U πind − uπind (t, Xtπind )∥22 ≥ r . d+2 By Lemma B.4, we have 4dγ 2 ωd2 2d+2 r > 0. d+2 Subtracting the zero population-OT variance, we can conclude Eq. (4). 2 σgrad,π (θπ⋆ ind ) ≥ ind

B.3

P ROPOSITION 4.9

Proof. By Assumptions 4.1 and 4.2, vθπ⋆ (t, Xtπ ) = uπ (t, Xtπ ) almost surely. Therefore, the loss at optimality equals the conditional target variance, and conditional orthogonality gives us that Eπ [∥U π ∥22 ] = Eπ [∥uπ (t, Xtπ )∥22 ] + Lπ (θπ⋆ ).

(26)

For every smooth compactly supported test function ϕ, differentiation along the linear interpolation and conditioning on (t, Xtπ ) at each fixed t, we have d E[ϕ(Xtπ )] = E[∇ϕ(Xtπ )⊤ U π ] = E[∇ϕ(Xtπ )⊤ uπ (t, Xtπ )]. dt 17

Preprint. Under review.

Thus, together with uπ , the marginal path induced by π satisfies the continuity equation and is admissible in the Benamou–Brenier formulation (Benamou & Brenier, 2000). So, its integrated kinetic energy is at least W22 (q0 , q1 ), i.e., Eπ [∥uπ (t, Xtπ )∥22 ] ≥ W22 (q0 , q1 ). Since Eπ [∥U π ∥22 ] = Eπ [∥X1 − X0 ∥22 ], combining this inequality with Eq. (26) proves Eq. (7). For π ⋆ , recall its transport cost is exactly W22 (q0 , q1 ), so the same bound also gives Lπ⋆ (θπ⋆ ⋆ ) = 0. B.4

L EMMA 4.10

Proof. Let 0 < ε1 < ε2 . Since q0 , q1 ∈ P2 (Rd ), independent coupling is a finite-objective competitor, so both quantities are finite. We use πε2 as a competitor for the ε1 -problem and πε1 as a competitor for the ε2 -problem. Their optimality gives us that Eπε1 [∥X1 − X0 ∥22 ] Eπε2 [∥X1 − X0 ∥22 ] + ε1 KL(πε1 | πind ) ≤ + ε1 KL(πε2 | πind ), 2d 2d Eπε1 [∥X1 − X0 ∥22 ] Eπε2 [∥X1 − X0 ∥22 ] + ε2 KL(πε2 | πind ) ≤ + ε2 KL(πε1 | πind ). 2d 2d Adding these inequalities and canceling the transport costs yields (ε2 − ε1 ) (KL(πε2 | πind ) − KL(πε1 | πind )) ≤ 0. Because ε2 > ε1 , we have KL(πε1 | πind ) ≥ KL(πε2 | πind ). Substituting this into the first optimality inequality further gives Eπε1 [∥X1 − X0 ∥22 ] − Eπε2 [∥X1 − X0 ∥22 ] ≤ 2dε1 (KL(πε2 | πind ) − KL(πε1 | πind )) ≤ 0. Therefore, the relative entropy is non-increasing in ε, while the transport cost is non-decreasing: KL(πε1 | πind ) ≥ KL(πε2 | πind ), Eπε1 [∥X1 − X0 ∥22 ] ≤ Eπε2 [∥X1 − X0 ∥22 ].

(27)

For every 0 < ε ≤ ε0 , we can now apply Eqs. (7), (8) and (27) in order to obtain Lπε (θπ⋆ ε ) ≤ Eπε [∥X1 − X0 ∥22 ] − W22 (q0 , q1 ) ≤ Eπε0 [∥X1 − X0 ∥22 ] − W22 (q0 , q1 ) < Lπind (θπ⋆ ind ), which proves the strict upper inequality in Eq. (9).   φε (x0 ) + ψε (x1 ) − ∥x1 − x0 ∥22 /(2d) >0 exp ε for πind -almost every (x0 , x1 ) (Nutz & Wiesel, 2022). Thus, πε ∼ πind , meaning that the two couplings have exactly the same null sets. Therefore, we have Eπ [U π | t, Xtπ ] = f (t, Xtπ ) = U π

almost surely,

so the loss is zero.
t, x − tu, x + (1 − t)u . We write Pπ = Φ# (λ ⊗ π) for the induced law of (t, Xtπ , U π ). For any Borel set A, the definition of pushforward gives Pπε (A) = 0 ⇐⇒ (λ ⊗ πε )(Φ−1 (A)) = 0 ⇐⇒ ⇐⇒

(λ ⊗ πind )(Φ−1 (A)) = 0 Pπind (A) = 0.

Therefore, Pπε ∼ Pπind . For any Borel measurable f : [0, 1] × Rd → Rd , its graph Af := {(t, x, u) : u = f (t, x)} 18

Preprint. Under review.

is Borel. Applying the preceding equivalence to Acf gives Pπε (Af ) = 1

⇐⇒

Pπind (Af ) = 1.

Thus, a measurable function determines U from (t, Xt ) almost surely under one law exactly when it does under the other. Combining this with the zero-loss characterization above, we obtain Lπε (θπ⋆ ε ) = 0

⇐⇒

Lπind (θπ⋆ ind ) = 0.

Therefore, as Lπε (θπ⋆ ε ) > 0, with Proposition 4.9 and optimality of π ⋆ , we end up with 0 ≤ Lπ⋆ (θπ⋆ ⋆ ) ≤ Eπ⋆ [∥X1 − X0 ∥22 ] − W22 (q0 , q1 ) = 0. Consequently, Lπ⋆ (θπ⋆ ⋆ ) = 0 < Lπε (θπ⋆ ε ), which proves the strict lower inequality. B.5

R EMARK 4.11

Proof. Because q0 , q1 ∈ P2 (Rd ), the normalized quadratic cost satisfies Z  Z Z 1 ∥x1 − x0 ∥22 2 2 ∥x0 ∥2 q0 (dx0 ) + ∥x1 ∥2 q1 (dx1 ) < ∞. (q0 ⊗ q1 )(dx0 , dx1 ) ≤ 2d d The cost is also continuous, so vanishing-regularization convergence of entropic OT implies that every weak cluster point of πε as ε ↓ 0 is a quadratic-OT optimizer (Nutz & Wiesel, 2022). To pass from weak convergence to convergence of the unbounded quadratic costs, we make the uniform-integrability argument explicit. This is the same de la Vallée–Poussin argument used below in Lemma B.9; for similar one usage, also see the proof of Lemma 6.4 in Beiglböck et al. (2017). Since the squared norms under both marginals are integrable, the de la Vallée–Poussin criterion supplies a common nondecreasing, convex function Φ : [0, ∞) → [0, ∞) such that Z Z Φ(s) lim = ∞, Φ(∥x0 ∥22 ) q0 (dx0 ) + Φ(∥x1 ∥22 ) q1 (dx1 ) < ∞. s→∞ s Indeed, we can apply the criterion to the equally weighted mixture of the two squared-norm distributions, which has finite first moment. For every pair (x0 , x1 ), we have ∥x1 − x0 ∥22 ∥x0 ∥22 + ∥x1 ∥22 ≤ . 4 2 Therefore, monotonicity and convexity of Φ give us the following:   1 ∥x1 − x0 ∥22 1 Φ ≤ Φ(∥x0 ∥22 ) + Φ(∥x1 ∥22 ). 4 2 2 Integrating under any coupling π ∈ Π(q0 , q1 ), its fixed marginals yield the uniform bound  Z  Z Z ∥x1 − x0 ∥22 1 1 sup Φ dπ ≤ K = Φ(∥x0 ∥22 ) q0 (dx0 )+ Φ(∥x1 ∥22 ) q1 (dx1 ) < ∞. 4 2 2 π∈Π(q0 ,q1 ) For any a > 0, superlinearity gives a threshold Ma < ∞ such that s ≤ aΦ( 4s ) whenever s > Ma . Consequently, for every M ≥ Ma , we have  Z Z  ∥x1 − x0 ∥22 sup ∥x1 − x0 ∥22 1{∥x1 −x0 ∥22 >M } dπ ≤ a sup Φ dπ ≤ aK. 4 π∈Π(q0 ,q1 ) π∈Π(q0 ,q1 ) Since a > 0 is arbitrary, we obtain Z τ (M ) := sup ∥x1 − x0 ∥22 1{∥x1 −x0 ∥22 >M } dπ −→ 0

as M → ∞.

π∈Π(q0 ,q1 )

Thus, the quadratic costs are uniformly integrable over all couplings with these fixed marginals. We now use this bound to establish the cost limit. First, the family {πε }ε>0 is tight. To see this, for any δ > 0, we choose compact sets K0 , K1 ⊂ Rd such that qi (Kic ) < δ/2. The marginal constraints give us:
sup πε (K0 × K1 )c ≤ q0 (K0c ) + q1 (K1c ) < δ. ε>0

19

Preprint. Under review.

Therefore, from any sequence εn ↓ 0, we can extract a subsequence such that πεnj ⇒ π̄. The limit π̄ still has marginals q0 , q1 , and the vanishing-regularization result above gives us Eπ̄ [∥X1 − X0 ∥22 ] = W22 (q0 , q1 ). For fixed M > 0, we define cM := min{∥x1 − x0 ∥22 , M }. As cM is bounded and continuous, weak convergence further gives us Z Z cM dπεnj −→ cM dπ̄. Moreover, since 0 ≤ ∥x1 − x0 ∥22 − cM ≤ ∥x1 − x0 ∥22 1{∥x1 −x0 ∥22 >M } , so we have Eπεn [∥X1 −X0 ∥22 ]−Eπ̄ [∥X1 −X0 ∥22 ] j

Z

Z cM dπεnj −

≤

cM dπ̄

Z

Z + (∥x1 − x0 ∥22 −cM ) dπεnj + (∥x1 − x0 ∥22 −cM ) dπ̄ Z Z cM dπεnj − cM dπ̄ + 2τ (M ). ≤ Taking j → ∞ at fixed M , and then M → ∞, we have Eπεn [∥X1 − X0 ∥22 ] −→ Eπ̄ [∥X1 − X0 ∥22 ] = W22 (q0 , q1 ). j

This subsequence argument also gives convergence of the full family. Otherwise, since Eπε [∥X1 − X0 ∥22 ] ≥ W22 (q0 , q1 ), there would exist δ > 0 and a sequence εn ↓ 0 such that Eπεn [∥X1 − X0 ∥22 ] − W22 (q0 , q1 ) ≥ δ for every n. The preceding claim would produce a subsequence whose costs converge to W22 (q0 , q1 ), which is a contradiction, so we get Eπε [∥X1 − X0 ∥22 ] −→ W22 (q0 , q1 )

as ε ↓ 0.

(28)

Since Lπind (θπ⋆ ind ) > 0, we can now choose ε0 > 0 sufficiently small that 0 ≤ Eπε [∥X1 − X0 ∥22 ] − W22 (q0 , q1 ) <

Lπind (θπ⋆ ind ) < Lπind (θπ⋆ ind ), 2

0 < ε ≤ ε0 .

In particular, Eq. (8) holds at ε0 . For the explicit cutoff, fix a competitor π e ∈ Π(q0 , q1 ) satisfying 0 < KL(e π ∥πind ) < ∞,

Eπe [∥X1 − X0 ∥22 ] − W22 (q0 , q1 ) < Lπind (θπ⋆ ind ).

Multiplying the objective in Eq. (6) by 2d, optimality of πε against π e gives Eπε [∥X1 − X0 ∥22 ] + 2dε KL(πε ∥πind ) ≤ Eπe [∥X1 − X0 ∥22 ] + 2dε KL(e π ∥πind ). Since KL(πε ∥πind ) ≥ 0, dropping its contribution and subtracting W22 (q0 , q1 ) yields Eπε [∥X1 − X0 ∥22 ] − W22 (q0 , q1 ) ≤ Eπe [∥X1 − X0 ∥22 ] − W22 (q0 , q1 ) + 2dε KL(e π ∥πind ). The competitor condition ensures that ∆ := Lπind (θπ⋆ ind ) + W22 (q0 , q1 ) − Eπe [∥X1 − X0 ∥22 ] > 0. 20

Preprint. Under review.

∆ Thus, whenever 0 < ε < 2d KL(e π ∥πind ) , which is exactly Eq. (10), we have

2dε KL(e π ∥πind ) < ∆. Substituting this into the preceding cost bound gives Eπε [∥X1 − X0 ∥22 ] − W22 (q0 , q1 ) ≤ Eπe [∥X1 − X0 ∥22 ] − W22 (q0 , q1 ) + 2dε KL(e π ∥πind ) < Eπe [∥X1 − X0 ∥22 ] − W22 (q0 , q1 ) + ∆ = Lπind (θπ⋆ ind ). So, Eq. (8) holds with this ε as the reference ε0 , and Lemma 4.10 gives the claimed strict ordering in Eq. (9). Finally, we verify that a competitor satisfying both conditions always exists. By Eq. (28), we choose a sufficiently small η > 0 such that Eπη [∥X1 − X0 ∥22 ] − W22 (q0 , q1 ) < Lπind (θπ⋆ ind ), and set π e = πη . To establish finiteness of its relative entropy, we use πind as a competitor in the η-problem. Obviously, KL(πind ∥πind ) = 0, and optimality gives us Eπη [∥X1 − X0 ∥22 ] + 2dη KL(πη ∥πind ) ≤ Eπind [∥X1 − X0 ∥22 ] < ∞. Consequently, we have 0 ≤ KL(πη ∥πind ) ≤

Eπind [∥X1 − X0 ∥22 ] − Eπη [∥X1 − X0 ∥22 ] Eπind [∥X1 − X0 ∥22 ] ≤ < ∞. 2dη 2dη

We still need to show that KL(πη ∥πind ) > 0. Again, we can consider the contradiction, i.e., suppose instead that KL(πη ∥πind ) = 0. Since relative entropy vanishes only when the two probability measures agree, this would imply πη = πind . Applying Eq. (7) to independent coupling then gives Lπind (θπ⋆ ind ) ≤ Eπind [∥X1 − X0 ∥22 ] − W22 (q0 , q1 ) = Eπη [∥X1 − X0 ∥22 ] − W22 (q0 , q1 ) < Lπind (θπ⋆ ind ), which is a contradiction. Therefore, we can now conclude 0 < KL(πη ∥πind ) < ∞, and π e = πη satisfies all the required conditions, so the guarantee we claimed holds true. B.6

T HEOREM 4.12

Before establish the proof, we first justify the sufficient condition for Lπind (θπ⋆ ind ) > 0 stated in Lemma 4.10. Suppose q0 has a Lebesgue density f0 > 0 almost everywhere and q1 is not a Dirac mass. Under πind , the joint law of (X0 , X1 ) is f0 (x0 ) dx0 q1 (dy). Fix t ∈ (0, 1) and we write Xt = (1 − t)X0 + tX1 , then for any Borel sets A, B ⊆ Rd , independence gives us that Z Z
P(Xt ∈ A, X1 ∈ B) = 1A (1 − t)x0 + ty f0 (x0 ) dx0 q1 (dy). B

Rd

For each fixed y, we can do the change of variables x = (1 − t)x0 + ty, so that x0 = (x − ty)/(1 − t) and dx0 = (1 − t)−d dx. Substituting into the preceding integral and applying Tonelli’s theorem yields the following   Z Z x − ty −d dx q1 (dy) P(Xt ∈ A, X1 ∈ B) = (1 − t) f0 1−t B A   Z Z x − ty q1 (dy) dx. = (1 − t)−d f0 1−t A B 21

Preprint. Under review.

By taking B = Rd , we can show that Xt has Lebesgue density   Z x − tz −d ρt (x) = (1 − t) f0 q1 (dz). 1−t Rd For every x such that 0 < ρt (x) < ∞, we abbreviate the following for clarity:   (1 − t)−d f0 x−ty 1−t Kt (x, dy) = q1 (dy). ρt (x) R The preceding joint-probability identity then becomes P(Xt ∈ A, X1 ∈ B) = A Kt (x, B)ρt (x) dx, so Kt (x, ·) is a version of the conditional law of X1 given Xt = x. Consequently, for almost every x under the law of Xt , we have   x−ty f 0 1−t d L(X1 | Xt = x)   (y) = Z . x − tz dq1 f0 q1 (dz) 1−t Rd The denominator depends only on t and x, so the Radon-Nikodym derivative is proportional to x − ty
. y 7→ f0 1−t We note that the derivative is strictly positive almost everywhere, so the conditional law is equivalent to the non-Dirac law q1 and therefore has strictly positive covariance trace. Since X1 − Xtπind , 1−t its conditional covariance also has strictly positive trace. By integrating over (t, Xtπind ), we can prove Lπind (θπ⋆ ind ) > 0, and the condition holds for q0 = N (0, Id ). U πind =

We now prove the theorem: Proof. Following from Fact B.3, we write the affine output layer as vη,b (t, x) = v̄η (t, x) + b. Since ∇b vη,b (t, x) = Id , its bias-gradient block at the population optimum is
∇b ℓπ (θπ⋆ ; ξ π ) = 2 uπ (t, Xtπ ) − U π . Its conditional mean at (t, Xtπ ) is zero, so the law of total covariance therefore gives us that Tr{Covπ (∇b ℓπ (θπ⋆ ; ξ π ))} = 4Lπ (θπ⋆ ).

(29)

For the complete gradient, Eq. (3) and the positive-semidefinite trace bound gives us that 2 σgrad,π (θπ⋆ ε ) ≤ 4Lπε (θπ⋆ ε ) ess sup ∥Jθπ⋆ε (t, Xtπε )∥2op . ε

Combining this inequality with Eq. (7) proves the upper bound in Eq. (11). For the strict lower bound, Eq. (29) and Lemma 4.10 show that the bias-gradient block has strictly positive variance for every finite ε. Since the total gradient-variance trace is at least the trace of this principal block, we can show that 2 σgrad,π (θπ⋆ ε ) > 0. At the same time, population OT has zero gradient variance by Eqs. (3) and (7). ε Therefore, whenever the right-hand side of Eq. (11) is smaller than the independent-coupling variance, the strict chain in Eq. (12) follows. Finally, the independent-coupling variance is strictly positive by Eq. (29). The theorem’s uniform condition means that there exist εJ > 0 and M < ∞ such that {πε : 0 < ε ≤ εJ } satisfies Definition 4.7 with bound M . Thus, for every 0 < ε ≤ εJ , we have ess sup ∥Jθπ⋆ε (t, Xtπε )∥2op ≤

max

t∈[0,1], x∈Rd

∥Jθπ⋆ε (t, x)∥2op ≤ M.

We denote Dε := Eπε [∥X1 − X0 ∥22 ] − W22 (q0 , q1 ), Eq. (28) then gives us that 4M Dε → 0 as ε ↓ 0. Therefore, there exists εg > 0 such that 2 4M Dε < σgrad,π (θπ⋆ ind ) ind

for every 0 < ε ≤ εg .

By setting ε1 = min{ε0 , εJ , εg } > 0, then for every 0 < ε ≤ ε1 , the right-hand side of Eq. (11) is smaller than the independent-coupling variance, which concludes Eq. (12). 22

Preprint. Under review.

B.7

P ROPOSITION 4.13

Proof. We take qb0,k and qb1,k to denote the empirical distributions of the two endpoint batches used to construct πk . Note that conditional on these batches, the OT solution is a permutation and the selected index is uniform. Therefore, the selected source and target each remain uniformly distributed over their respective batches; after averaging over the batches, we have πk ∈ Π(q0 , q1 ). Moreover, one can derive that Eπk [∥X1 − X0 ∥22 ] = E[W22 (b q0,k , qb1,k )]. To establish this, we first show that, for any q ∈ P2 (Rd ) along with its empirical distribution qbk , E[W22 (b qk , q)] will converge to zero. For any R > 0, let PR be the metric projection onto the closed ball B(0, R), and let q R = (PR )# q and qbkR = (PR )# qbk , coupling each sample with its projection gives us that E[W22 (b qk , qbkR )] ≤ aR := Eq [∥X − PR (X)∥22 ],

W22 (q, q R ) ≤ aR .

For fixed R > 0, the empirical measure qbkR converges to q R almost surely in W2 . Both measures are supported on B(0, R), whose diameter is 2R, so W22 (b qkR , q R ) ≤ 4R2 . The almost-sure convergence R R of qbk to q in W2 further gives W22 (b qkR , q R ) −→ 0

almost surely.

Since both measures are supported on the radius-R ball, we also have 0 ≤ W22 (b qkR , q R ) ≤ 4R2 . This deterministic bound is integrable and independent of k, so dominated convergence gives   2 R R 2 R R lim E[W2 (b qk , q )] = E lim W2 (b qk , q ) = 0. k→∞

k→∞

Next, the triangle inequality gives W2 (b qk , q) ≤ W2 (b qk , qbkR ) + W2 (b qkR , q R ) + W2 (q R , q). First using (a + b + c)2 ≤ 3(a2 + b2 + c2 ) and then simply taking expectations, we have E[W22 (b qk , q)] ≤ 3E[W22 (b qk , qbkR )] + 3E[W22 (b qkR , q R )] + 3W22 (q R , q) ≤ 3aR + 3E[W22 (b qkR , q R )] + 3aR = 6aR + 3E[W22 (b qkR , q R )], where the second inequality uses the two projection bounds established above. Taking the limit superior as k → ∞, while keeping R fixed, therefore gives lim sup E[W22 (b qk , q)] ≤ 6aR + 3 lim E[W22 (b qkR , q R )] = 6aR . k→∞

k→∞

Since q ∈ P2 (Rd ), aR → 0 as R → ∞, which proves our claim. Applying it to both marginals yields |W2 (b q0,k , qb1,k ) − W2 (q0 , q1 )| ≤ W2 (b q0,k , q0 ) + W2 (b q1,k , q1 ), The right-hand side converges to 0 in L2 , so W2 (b q0,k , qb1,k ) → W2 (q0 , q1 ) in L2 . Its square therefore 1 converges in L , which proves ∆k → 0 in Eq. (13). Since πk is a coupling of q0 and q1 , its cost is at least W22 (q0 , q1 ), so ∆k ≥ 0. For any coupling π, conditional orthogonality gives us that Eπ [∥U π ∥22 ] = Eπ [∥uπ (t, Xtπ )∥22 ] + Eπ [Tr{Covπ (U π | t, Xtπ )}]. As earlier shown in the proof of Proposition 4.9, the interpolation induced by π, together with uπ , is admissible in the Benamou–Brenier formulation. Therefore, Eπ [∥uπ (t, Xtπ )∥22 ] ≥ W22 (q0 , q1 ). Applying the preceding kinetic-energy bound to πk , we have   Eπk ∥uπk (t, Xtπk )∥22 ≥ W22 (q0 , q1 ). 23

Preprint. Under review.

We can then rearrange the conditional-orthogonality identity and use U πk = X1 − X0 to get 0 ≤ Eπk [Tr{Covπk (U πk | t, Xtπk )}]   = Eπk ∥U πk − uπk (t, Xtπk )∥22     = Eπk ∥U πk ∥22 − Eπk ∥uπk (t, Xtπk )∥22   ≤ Eπk ∥U πk ∥22 − W22 (q0 , q1 )   = Eπk ∥X1 − X0 ∥22 − W22 (q0 , q1 ) = ∆k .

(∵ conditional-variance identity) (∵ conditional orthogonality) (∵ mean-velocity energy lower bound)

This proves the upper bound in Eq. (14). Since we have already established that ∆k → 0, the squeeze theorem further gives Eπk [Tr{Covπk (U πk | t, Xtπk )}] −→ 0

as k → ∞.

So, we prove ∆k → 0. B.8

T HEOREM 4.14

Before moving into the proof of Theorem 4.14. We first provide the following supporting lemma, which shows that every finite-batch coupling has zero target uncertainty exactly when independent coupling does. Lemma B.5. For every finite k ≥ 1, πk and πind are mutually absolutely continuous, so it has Eπk [Tr{Covπk (U πk | t, Xtπk )}] = 0

⇐⇒

Eπind [Tr{Covπind (U πind | t, Xtπind )}] = 0.

If q0 = f0 Ld with f0 > 0 Lebesgue-almost everywhere and q1 is not a Dirac mass, both quantities are strictly positive. Proof. We write the two endpoint batches as (X1 , . . . , Xk ) and (Y1 , . . . , Yk ), and let I be uniform on {1, . . . , k}, independently of both batches, along with τ be the permutation returned by a fixed measurable tie-breaking rule for empirical OT. Then, (XI , Yτ (I) ) ∼ πk . For any Borel set E ⊂ Rd × Rd , if (q0 ⊗ q1 )(E) = 0, all finitely many candidate pairs (Xi , Yj ) avoid E almost surely, so πk (E) = 0. Conversely, we denote Ex = {y : (x, y) ∈ E}, and the event that every candidate target paired with XI lies in EXI forces the selected pair to lie in E. Therefore, we get Z πk (E) ≥ P((XI , Yj ) ∈ E for every j) = q1 (Ex )k q0 (dx). If (q0 ⊗ q1 )(E) > 0, then q1 (Ex ) > 0 on a set of positive q0 -measure, so the final integral is strictly positive. Thus, one have πk ∼ πind . Next, as we discussed before, the map (t, x0 , x1 ) 7→ (t, (1 − t)x0 + tx1 , x1 − x0 ) is a Borel bijection with inverse x0 = x − tu, x1 = x + (1 − t)u. Therefore, equivalent endpoint couplings induce equivalent laws of (t, Xt , U ). As the covariance trace is nonnegative, its expectation is zero exactly when Cov(U | t, Xt ) = 0 almost surely, or equivalently when U is almost surely a measurable function of (t, Xt ). This graph property is preserved under equivalent laws, proving the stated equivalence. Finally, using the stated density conditions and fix t ∈ (0, 1). Under independent coupling, for almost every x under the law of Xt , the conditional law of X1 given Xt = x has Radon–Nikodym derivative with respect to q1 proportional to   x − ty y 7−→ f0 . 1−t We claim that this derivative is strictly positive for q1 -almost every y. To prove it, we set t ∈ (0, 1) and let N := {z ∈ Rd : f0 (z) = 0}. The assumption f0 > 0 Lebesgue-almost everywhere gives Ld (N ) = 0. For every fixed y ∈ Rd , the change of variables x = (1 − t)z + ty gives   Z Z x − ty d dx = (1 − t) 1N (z) dz = (1 − t)d Ld (N ) = 0. 1N 1−t Rd Rd 24

Preprint. Under review.

Thus, for each fixed y, the numerator f0 ((x − ty)/(1 − t)) is positive for Lebesgue-almost every x. We now interchange the order of integration. Since the indicator is nonnegative, Tonelli’s theorem gives us       Z Z Z Z x − ty x − ty 1N q1 (dy) dx = dx q1 (dy) = 0. 1N 1−t 1−t Rd Rd Rd Rd The inner integral on the left is nonnegative, so it must vanish for Lebesgue-almost every x. Equivalently, we have     x − ty q1 y : f0 =0 =0 for Lebesgue-almost every x. 1−t Therefore, for Lebesgue-almost every x, we have   x − ty >0 for q1 -almost every y. f0 1−t We still remain for considering the normalizing factor in the conditional density. Under independent coupling, Xt has Lebesgue density   Z x − ty ρt (x) = (1 − t)−d q1 (dy). f0 1−t Rd The preceding positivity gives ρt (x) > 0 for Lebesgue-almost every x. Moreover, we can show that     Z Z  Z Z Z x − ty ρt (x) dx = (1 − t)−d f0 dx q1 (dy) = f0 (z) dz q1 (dy) = 1, 1−t Rd Rd Rd Rd Rd so ρt (x) < ∞ for Lebesgue-almost every x. Since the law of Xt is absolutely continuous with respect to Lebesgue measure, all these properties also hold for almost every x under the law of Xt . For such x, the conditional Radon–Nikodym derivative therefore satisfies   x−ty −d (1 − t) f 0 1−t dL(X1 | Xt = x) >0 for q1 -almost every y. (y) = dq1 ρt (x) Thus, for almost every x under the law of Xt , the conditional law of X1 given Xt = x is equivalent to q1 , which concludes our claim. Furthermore, the conditional law is equivalent to the non-Dirac distribution q1 and has strictly positive covariance trace. Since U = (X1 − Xt )/(1 − t), the conditional covariance of U also has strictly positive trace. Integrating first over Xt and then over t ∈ (0, 1) proves positivity for independent coupling and therefore, by the preceding equivalence, for every k < ∞. We now formally prove Theorem 4.14: Proof of Theorem 4.14. At the coupling-specific population optimum, Eq. (3) and the positivesemidefinite trace bound used in Proposition B.1 give 2 0 ≤ σgrad,π (θπ⋆ k ) ≤ 4 ess sup ∥Jθπ⋆ (t, Xtπk )∥2op Eπk [Tr{Covπk (U πk | t, Xtπk )}]. k k

Applying Eq. (14) proves the upper bound in Eq. (15). The uniform condition means that there exist KJ ∈ N and M < ∞ such that, for every k ≥ KJ , the following essential optima control exists: ess sup ∥Jθπ⋆ (t, Xtπk )∥2op ≤ k

max

t∈[0,1], x∈Rd

∥Jθπ⋆ (t, x)∥2op ≤ M. k

Since ∆k → 0, the right-hand side of Eq. (15) is at most 4M ∆k → 0 for k ≥ KJ . The squeeze 2 theorem therefore gives σgrad,π (θπ⋆ k ) → 0, so Eq. (16) holds for all sufficiently large k whenever the k fixed baseline variance is positive. For completeness, we also connect the per-sample statement to an update averaged over all matched pairs. We denote {Gk,1 , . . . , Gk,k } as the gradients generated at θπ⋆ k by the k matched pairs in one empirical OT batch, with an independent uniform time sampled for each pair, and let I be an independent 25

Preprint. Under review.

uniform index. Then, Gk,I has the same law as Gπk (θπ⋆ k ). The proof of Proposition 4.3 gives E[Gk,I ] = 0, so the averaged update also has mean zero. Jensen’s inequality further gives us the following bound:   n  o 2 −1 Pk −1 Pk 2 Tr Cov k =E k ≤ E[∥Gk,I ∥22 ] = σgrad,π (θπ⋆ k ). (30) i=1 Gk,i i=1 Gk,i k 2

Thus, the averaged matched-batch update also has variance converging to zero under the same vanishing-product condition. B.9

P ROPOSITION 4.15

To ensure rigor, we first establish several lemmas required to pass from convergence of entropic couplings to convergence of their target uncertainty. Lemma B.6. If πn , π ∈ Π(q0 , q1 ), with q0 , q1 ∈ P2 (Rd ), and πn → π in W2 , then we have lim sup R(πn ) ≤ R(π).

(31)

n→∞

Proof. We write Pn and P as the laws of ((t, Xt ), U ) induced by πn and π, respectively. Then, we can couple πn and π optimally and use the same independent t for both couplings. Since Xt = (1 − t)X0 + tX1 and U = X1 − X0 are linear in the endpoints, Pn → P in W2 . The conditional-mean characterization of squared-error regression gives R(π) = inf EP [∥U − f (t, Xt )∥22 ], f

where the infimum may be restricted to bounded continuous f , which are dense in the relevant L2 space. For every such fixed f , the integrand is continuous with quadratic growth, so its expectation converges under W2 convergence. Therefore, we end up with lim sup R(πn ) ≤ EP [∥U − f (t, Xt )∥22 ]. n→∞

Taking the infimum over f proves Eq. (31). For the convergence of empirical entropic plans, we first try to conclude the following local value bound. Before that, for notation simplicity, given probability measure µ and ν, we define   1 Eρ [∥X1 − X0 ∥22 ] + ε KL(ρ∥µ ⊗ ν) . Sε (µ, ν) := inf ρ∈Π(µ,ν) 2d Lemma B.7. Let qb0,k and qb1,k be the two empirical endpoint distributions, set ak = W2 (b q0,k , q0 ) + W2 (b q1,k , q1 ), and let Cind = Eπind [∥X1 − X0 ∥22 ]. Then, for every ε > 0, we have p 1 Sε (b q0,k , qb1,k ) ≤ Sε (q0 , q1 ) + 2d (2 Cind ak + a2k ). (32) Proof. Let ρε minimize Sε (q0 , q1 ), and transport each of its marginals to the corresponding empirical marginal through optimal W2 kernels. Applying the product of these kernels to ρε produces a coupling ρ′ε ∈ Π(b q0,k , qb1,k ). We claim that KL(ρ′ε ∥b q0,k ⊗ qb1,k ) ≤ KL(ρε ∥q0 ⊗ q1 ). To prove it, we fix the sampled endpoint batches and let Ki (xi , db xi ) be R the transport kernel from qi to qbi,k , for i ∈ {0, 1}. Thus, for every Borel set Ai ⊆ Rd , we have Ki (xi , Ai ) qi (dxi ) = qbi,k (Ai ). We denote x = (x0 , x1 ) and x b = (b x0 , x b1 ), and the product kernel is K0 (x0 , db x0 )K1 (x1 , db x1 ). For any probability measure µ on the space, the law obtained by applying this kernel to µ
R endpoint R assigns to each Borel set A with K0 (x0 , db x0 )K1 (x1 , db x1 ) µ(dx). So, we then have A  Z Z ′ ρε (A) = K0 (x0 , db x0 )K1 (x1 , db x1 ) ρε (dx). A

26

Preprint. Under review.

Applying the same kernel to the independent reference measure gives, for every Borel rectangle A0 × A1 , we can do the following derivation: Z Z

 Z K0 (x0 , db x0 )K1 (x1 , db x1 ) (q0 ⊗ q1 )(dx) = K0 (x0 , A0 )K1 (x1 , A1 ) q0 (dx0 )q1 (dx1 ) A0 ×A1 Z  Z  = K0 (x0 , A0 ) q0 (dx0 ) K1 (x1 , A1 ) q1 (dx1 ) = qb0,k (A0 )b q1,k (A1 ) = (b q0,k ⊗ qb1,k )(A0 × A1 ).

Since these rectangles determine the probability measure, for every Borel set A, we have Z Z



A

K0 (x0 , db x0 )K1 (x1 , db x1 ) (q0 ⊗ q1 )(dx) = (b q0,k ⊗ qb1,k )(A).

We next derive the entropy inequality directly. Since independent coupling is a finite-objective competitor and ε > 0, optimality gives KL(ρε | q0 ⊗ q1 ) < ∞. Indeed, since ρε ≪ q0 ⊗ q1 , the ε density ratio d(qdρ (x) is then well defined. 0 ⊗q1 ) Now, we consider the joint reference law (q0 ⊗ q1 )(dx)K0 (x0 , db x0 )K1 (x1 , db x1 ). Under this law, the b has law qb0,k ⊗ qb1,k . original endpoint pair X has law q0 ⊗ q1 , while the transformed endpoint pair X We note that all the expectations below are taken under this joint reference law. For any Borel set A in the transformed endpoint space, the definition of ρε and the conditionalexpectation identity give  K0 (x0 , db x0 )K1 (x1 , db x1 ) ρε (dx) A Z  Z dρε (x) K0 (x0 , db x0 )K1 (x1 , db x1 ) (q0 ⊗ q1 )(dx) = d(q0 ⊗ q1 ) A   dρε =E (X)1{X∈A} b d(q0 ⊗ q1 )     dρε b =E E (X) X 1{X∈A} b d(q0 ⊗ q1 )  Z  dρε b = E (X) X = x b (b q0,k ⊗ qb1,k )(db x). d(q0 ⊗ q1 ) A

ρ′ε (A) =

Z Z

Therefore, the transformed density ratio is   dρε dρ′ε b =x (b x) = E (X) X b d(b q0,k ⊗ qb1,k ) d(q0 ⊗ q1 ) for (b q0,k ⊗ qb1,k )-almost every x b. Finally, the function s 7→ s log s (with 0 log 0 = 0) is convex on [0, ∞). Conditional Jensen’s inequality therefore gives  E

    dρε dρε b b (X) X log E (X) X d(q0 ⊗ q1 ) d(q0 ⊗ q1 )     dρε dρε b . ≤E (X) log (X) X d(q0 ⊗ q1 ) d(q0 ⊗ q1 ) 27

Preprint. Under review.

Taking expectations and using the two marginal laws of the joint reference law, we get   Z
dρ′ε dρ′ε KL ρ′ε | qb0,k ⊗ qb1,k = (b x) log (b x) (b q0,k ⊗ qb1,k )(db x) d(b q0,k ⊗ qb1,k ) d(b q0,k ⊗ qb1,k )       dρε dρε b log E b (X) X (X) X =E E d(q0 ⊗ q1 ) d(q0 ⊗ q1 )      dρε dρε b ≤E E (X) log (X) X d(q0 ⊗ q1 ) d(q0 ⊗ q1 )    dρε dρε =E (X) log (X) d(q0 ⊗ q1 ) d(q0 ⊗ q1 )   Z dρε dρε = (x) log (x) (q0 ⊗ q1 )(dx) d(q0 ⊗ q1 ) d(q0 ⊗ q1 ) = KL(ρε | q0 ⊗ q1 ). So our claim is proved, and by using Minkowski’s inequality, we have q q Eρ′ε [∥X1 − X0 ∥22 ] ≤ Eρε [∥X1 − X0 ∥22 ] + ak . Since independent coupling is feasible and has zero relative entropy, the raw transport cost of ρε is at most Cind . Substituting these bounds into the empirical objective, we can then conclude Eq. (32). Lemma B.8. Considering two sampled endpoint batches (X1 ,. . ., Xk ) and (Y1 ,. . ., Yk ), then for any uniform-marginal matrix P = (pij )ki,j=1 ∈ Rk×k satisfying P 1 = P ⊤ 1 = 1/k, we write + Pk ΓP = i,j=1 pij δ(Xi ,Yj ) ∈ Π(b q0,k , qb1,k ) as its pushforward to endpoint space. Then, the index-level problem in Eq. (2) and the measure-level problem defining Sε (b q0,k , qb1,k ) have the same minimum. Proof. Let {x1 , . . . , xA } and {y 1 , . . . , y B } be the distinct values in the two batches. We write Ia := {i : Xi = xa },

Jb := {j : Yj = y b },

ma := |Ia |,

nb := |Jb |.

For a feasible matrix P , we define its location-level block masses by XX P ab := pij . i∈Ia j∈Jb B Then, P is a coupling between the location masses (ma /k)A a=1 and (nb /k)b=1 , and

ΓP =

B A X X

P ab δ(xa ,yb ) .

a=1 b=1 ab Whenever P ab > 0, we set rij := pij /P ab on Ia × Jb ; when P ab = 0, one can choose rab arbitrarily. The relative-entropy chain rule gives us that 

 X
KL P 11⊤ /k 2 = KL P ma nb /k 2 a,b + P ab KL rab Unif(Ia × Jb ) . (33) a,b

So, clearly, the first term in Eq. (33) is exactly KL(ΓP ∥b q0,k ⊗ qb1,k ), while the second term is nonnegative. Moreover, the quadratic cost is constant on every block Ia × Jb , so we have Z 1 ⟨C, P ⟩ = ∥x − y∥22 ΓP (dx, dy). 2d Conversely, every location-level coupling P has the feasible block-uniform lift pij :=

P ab , ma nb

X

pij =

Indeed, for i ∈ Ia , we have j

i ∈ Ia , j ∈ J b .

1 X 1 P ab = , ma k b

28

Preprint. Under review.

and the column constraints follow identically. For this lift every conditional entropy term in Eq. (33) vanishes. Thus, minimizing the index-level objective first over lifts of a fixed endpoint coupling, and then over endpoint couplings, gives exactly the measure-level entropic problem. The block-uniform lift is the unique entropy-minimizing lift of each endpoint coupling. Since relative entropy is strictly convex on the finite location-level coupling polytope, the measurelevel optimizer is unique, and its block-uniform lift is the unique index-level optimizer. With these lemmas, we can finally now prove Proposition 4.15: Proof of Proposition 4.15. We denote qb0,k and qb1,k as the empirical marginals and set γ bk,ε = Γπbk,ε . By Lemma B.8, conditional on the endpoint batches, γ bk,ε is exactly the unique minimizer of Sε (b q0,k , qb1,k ). The empirical marginals converge almost surely to q0 and q1 in W2 . Fix ε > 0, we can work on an event of probability on which both empirical W2 convergences hold. We claim that the sequence {b γk,ε }k is tight as its two marginals are tight. Let γ be the weak limit of an arbitrary convergent subsequence {b γkj ,ε }j≥1 , its marginals are q0 and q1 . Joint lower semicontinuity of the quadratic cost and relative entropy along this subsequence gives us 1 Eγ ∥X1 − X0 ∥22 + ε KL(γ∥q0 ⊗ q1 ) ≤ lim inf Sε (b q0,kj , qb1,kj ). j→∞ 2d For the entropy term, the joint lower semicontinuity follows directly from its variational representation Z  Z KL(µ∥ν) = sup f dµ − log ef dν , f ∈Cb

which is because qb0,k ⊗ qb1,k ⇒ q0 ⊗ q1 . On the other hand, Lemma B.7 and ak → 0 (almost surely) give lim sup Sε (b q0,k , qb1,k ) ≤ Sε (q0 , q1 ). (34) k→∞

The full-sequence bound in Eq. (34) also holds along {kj }. Therefore, we have 1 Eγ ∥X1 −X0 ∥22 +ε KL(γ∥q0 ⊗q1 ) ≤ lim inf Sε (b q0,kj , qb1,kj ) ≤ lim sup Sε (b q0,kj , qb1,kj ) ≤ Sε (q0 , q1 ), j→∞ 2d j→∞ so γ minimizes the population objective in Eq. (6). We note the objective has a unique minimizer, i.e., independent coupling has finite objective, and relative entropy is strictly convex on the finite-objective domain, so γ = πε . Since every weak cluster point is the same, we end up with γ bk,ε ⇒ πε

almost surely.

Furthermore, we have Z Z Z (∥x∥22 + ∥y∥22 ) γ bk,ε (dx, dy) = ∥x∥22 qb0,k (dx) + ∥y∥22 qb1,k (dy), which converges almost surely to the corresponding second moment under q0 ⊗ q1 . Thus, the almost-sure convergence is in W2 . Conditional on the endpoint batches, γ bk,ε is precisely the law of the selected endpoint pair. Therefore, for every bounded continuous g, we have Z  Z Z g dπk,ε = E g db γk,ε −→ g dπε by bounded convergence, so πk,ε ⇒ πε . Every expected coupling πk,ε has marginals q0 , q1 , and therefore has the same product-space second moment as πε . It follows that πk,ε → πε in W2 . For any coupling π, the conditional orthogonality and Benamou–Brenier argument in the proof of Proposition 4.9 gives 0 ≤ R(π) ≤ Eπ [∥X1 − X0 ∥22 ] − W22 (q0 , q1 ). (35) Applying Lemma B.6 to the established W2 convergence yields lim supk→∞ R(πk,ε ) ≤ R(πε ), and applying Eq. (35) to πε proves Eq. (17). 29

Preprint. Under review.

We still need to prove the finite-batch equivalence. Let E ⊆ Rd × Rd be Borel, we write Ex = {y : (x, y) ∈ E}. If πind (E) = 0, every one of the finitely many candidate source-target pairs avoids E almost surely, so πk,ε (E) = 0. Conversely, if πind (E) > 0, then Z P(Yj ∈ EXI for every j) = q1 (Ex )k q0 (dx) > 0, and this event forces the selected pair to belong to E, so πk,ε ∼ πind . The Borel bijection
(t, x0 , x1 ) 7→ t, (1 − t)x0 + tx1 , x1 − x0 preserves equivalence of the induced laws. Finally, R(π) = 0 exactly when U π is almost surely a measurable function of (t, Xtπ ), and this graph property is invariant under equivalent laws. We now prove the proposition. B.10

T HEOREM 4.16

We first establish Lemma B.9 before prove Theorem 4.16. Lemma B.9. Suppose q0 , q1 ∈ P2 (Rd ), and we choose any k ∈ N. Conditional on two sampled endpoint batches, let Ok be the face of uniform-marginal matrices minimizing the empirical quadratic cost in Eq. (2), and we define
ME π bk,0 := arg min KL P 11⊤ /k 2 . (36) P ∈Ok

ME Let πk,0 be the averaged selected-pair coupling induced by this plan. Then, it gives that ME πk,ε − πk,0 −→ 0, TV

ME R(πk,ε ) −→ R(πk,0 )

(37)

as ε ↓ 0. Moreover, we have ME πk,0 ∼ πind ,

ME R(πk,0 )=0

⇐⇒

R(πind ) = 0.

(38)

Proof. Let  Tk :=



1 1 P ∈ Rk×k : P 1 = 1, P ⊤ 1 = 1 +

k k be the empirical transport poly-tope. Conditional on the sampled batches, we write
C⋆ := min C(P ). C(P ) := ⟨C, P ⟩, H(P ) := KL P 11⊤ /k 2 , P ∈Tk

We note that Tk is both compact and convex, and the optimal face Ok = {P ∈ Tk : C(P ) = C⋆ } is nonempty, compact, and convex. Because the reference matrix has strictly positive entries, H is continuous and strictly convex on Tk . Therefore, Eq. (36) has a unique solution. ME We denote π bk,ε as the unique empirical entropic optimizer, and its optimality against π bk,0 gives ME C(b πk,ε ) + εH(b πk,ε ) ≤ C⋆ + εH(b πk,0 ).

Since C(b πk,ε ) ≥ C⋆ and H ≥ 0, we can show that ME 0 ≤ C(b πk,ε ) − C⋆ ≤ εH(b πk,0 ),

ME H(b πk,ε ) ≤ H(b πk,0 ).

ME Therefore, every cluster point of π bk,ε as ε ↓ 0 lies in Ok and has entropy no larger than that of π bk,0 . ME The uniqueness in Eq. (36) forces every cluster point to equal π bk,0 . Compactness then gives convergence of the full sequence in the finite-dimensional polytope, equivalently in entrywise ℓ1 and total variation.

For every positive ε, the unique optimizer is a Borel function of the sampled cost matrix by the measurable maximum theorem. Taking, for example, εn = 1/n, the pointwise convergence just ME proved shows that π bk,0 is the pointwise limit of measurable optimizers and is therefore measurable as well. 30

Preprint. Under review.

Pk For a matrix P ∈ Tk , we denote ΓP = i,j=1 Pij δ(Xi ,Yj ) as its selected-pair endpoint law conditional on the batches. Pushforward contracts total variation, so we have 1 ME ME ∥TV ≤ ∥Γπbk,ε − Γπbk,0 ∥b πk,ε − π bk,0 ∥1 −→ 0 2 for every realized pair of batches. Convexity of total variation and dominated convergence now yield h i ME ME ∥TV ∥πk,ε − πk,0 ∥TV ≤ E ∥Γπbk,ε − Γπbk,0 −→ 0. This proves the first convergence in Eq. (37). We next prove the continuity of R needed for the second convergence. More generally, let ρn , ρ ∈ Π(q0 , q1 ) satisfy ∥ρn − ρ∥TV → 0, and let Pn , P be the corresponding laws of ((t, Xt ), U ), where U = X1 − X0 . Appending the same independent t ∼ U([0, 1]) preserves total-variation convergence, while measurable pushforward contracts it. Thus, we have ∥Pn − P ∥TV → 0.   Let TL : Rd → B(0, L) denote metric projection, we set RL (ρ′ ) = inf f Eρ′ ∥TL (U ) − f (t, Xt )∥22 . The infimum may be restricted to predictors with values in B(0, L), because projection onto this ball cannot increase the distance to TL (U ), so the resulting losses are bounded by 4L2 . Given ∥µ − ν∥TV = supA |µ(A) − ν(A)|, comparison with an arbitrarily accurate minimizer under the other law gives us that |RL (ρn ) − RL (ρ)| ≤ 8L2 ∥Pn − P ∥TV −→ 0. (39) The final missing point is that we need to remove the truncation uniformly. We write A = ∥X0 ∥22 and B = ∥X1 ∥22 ; because their two fixed marginal laws are integrable, the de la Vallée–Poussin criterion supplies a nondecreasing, convex, superlinear function Φ such that Eq0 [Φ(A)] + Eq1 [Φ(B)] < ∞. ′ For every ρ ∈ Π(q0 , q1 ), we have ∥X1 − X0 ∥22 A+B ≤ , 4 2 and monotonicity and convexity give the following bound    ∥X1 − X0 ∥22 1 1 sup Eρ′ Φ ≤ Eq0 [Φ(A)] + Eq1 [Φ(B)] < ∞. 4 2 2 ρ′ ∈Π(q0 ,q1 ) Therefore, {∥X1 − X0 ∥22 : ρ′ ∈ Π(q0 , q1 )} is uniformly integrable, and we end up with   eL = sup Eρ′ ∥U − TL (U )∥22 −→ 0.

(40)

ρ′ ∈Π(q0 ,q1 )

For any coupling ρ′ , the square root of its Bayes risk is the ℓ2 -distance from U to the closed subspace of (t, Xt )-measurable vector-valued functions. Since the distance to a closed set is 1-Lipschitz, we have p p
1/2 1/2 R(ρ′ ) − RL (ρ′ ) ≤ Eρ′ ∥U − TL (U )∥22 ≤ eL . (41) Combining Eqs. (39) to (41), first taking n → ∞ at fixed L and then taking L → ∞, we get R(ρn ) → R(ρ). Applying this fact to the endpoint couplings above proves the second convergence in Eq. (37). Finally, consider E ⊆ Rd × Rd be Borel with Ex = {y : (x, y) ∈ E}. If (q0 ⊗ q1 )(E) = 0, ME all k 2 sampled candidate pairs avoid E almost surely, which gives πk,0 (E) = 0. Conversely, if (q0 ⊗ q1 )(E) > 0, we just consider I be any uniformly selected source index, which gives us Z P(Yj ∈ EXI for every j) = q1 (Ex )k q0 (dx) > 0. On this event, whichever target index is selected from row I, the selected pair still lies in E. Therefore, ME ME ∼ πind . Again, recall that the map πk,0 (E) > 0, proving πk,0
(t, x0 , x1 ) 7−→ t, (1 − t)x0 + tx1 , x1 − x0 is a Borel bijection, with inverse (t, x, u) 7→ (t, x − tu, x + (1 − t)u). Therefore, equivalent endpoint laws induce equivalent laws of (t, Xt , U ). Since R(ρ) = 0 ⇐⇒ U is almost surely a measurable function of (t, Xt ), so this zero-risk graph property is invariant under equivalence, and we can then conclude Eq. (38). 31

Preprint. Under review.

Before establishing the proof for Theorem 4.16, we define Dε = Eπε [∥X1 − X0 ∥22 ] − W22 (q0 , q1 ). The theorem’s joint uniform condition means that there exist KJ ∈ N, εJ > 0, and M < ∞ such that ess sup ∥Jθπ⋆

π

k,ε

(t, Xt k,ε )∥2op ≤

max

t∈[0,1], x∈Rd

∥Jθπ⋆

k,ε

(t, x)∥2op ≤ M, k ≥ KJ , 0 < ε ≤ εJ . (42)

Moreover, Eq. (28) gives Dε → 0 as ε ↓ 0. Now, we can start the proof: Proof of Theorem 4.16. At a realizable coupling-specific population optimum, Eq. (3) gives the exact gradient-variance identity. The trainable additive output bias implies Jθπ⋆ Jθ⊤π⋆ ⪰ Id by Fact B.3, k,ε

π

k,ε

with both Jacobians evaluated at (t, Xt k,ε ). Applying the positive-semidefinite trace bounds used in Proposition B.1 proves both inequalities in Eq. (18). We next obtain a positive lower bound for the prescribed finite baseline as ε ↓ 0 and set r0 = R(πkME ). 0 ,0 By Lemma B.9 and the assumption R(πind ) > 0, we have R(πk0 ,ε ) −→ r0

r0 > 0,

as ε ↓ 0.

So, there exists εb > 0 such that r0 2 By the lower inequality in Eq. (18), this gives

for every 0 < ε ≤ εb .

R(πk0 ,ε ) >

2 σgrad,π (θπ⋆ k ,ε ) > 2r0 k ,ε 0

0

for every 0 < ε ≤ εb .

(43)

Now fix 0 < ε ≤ εJ . Proposition 4.15 and Eq. (42) imply lim sup R(πk,ε ) ≤ Dε , k→∞

lim sup ess sup ∥Jθπ⋆ k→∞

π

k,ε

(t, Xt k,ε )∥2op ≤ M.

Note that both sequences are nonnegative and the two right-hand sides are finite. Therefore, applying Eq. (18) and the elementary limsup product inequality gives 2 lim sup σgrad,π (θπ⋆ k,ε ) ≤ 4M Dε . k,ε

(44)

k→∞

Since Dε → 0, there exists εs > 0 such that 4M Dε < r0

for every 0 < ε ≤ εs .

(45)

We set ε̄ = min{εJ , εb , εs } > 0. For every fixed 0 < ε ≤ ε̄, Eqs. (43) to (45) give 2 2 lim sup σgrad,π (θπ⋆ k,ε ) < r0 < 2r0 < σgrad,π (θπ⋆ k ,ε ). k,ε k ,ε 0

k→∞

0

By the limit of superior, we have a finite threshold after which the strict upper comparison holds. Increasing that threshold if necessary, choose a finite K > max{k0 , KJ }, which may depend on k0 and ε; then, Eq. (19) holds for every k ≥ K. Now, it remains for us only to verify strict positivity. For every finite k and ε > 0, Proposition 4.15 gives πk,ε ∼ πind . As in Lemma B.9, equivalence preserves the graph property characterizing zero 2 target uncertainty. Therefore, R(πk,ε ) > 0, and the lower inequality in Eq. (18) gives σgrad,π > 0. k,ε This completes the proof, and we can conclude the variance reduction. B.10.1

P ROPOSITION 4.18

Before establishing the proof, we need to prove two supporting lemmas. We first prove the following one-step coupling bound: Lemma B.10. Let q0 , q1 ∈ P2 (Rd ), let π ∈ Π(q0 , q1 ), and we denote mπ (x) = Eπ [X1 | X0 = x]. Then, we have 0 ≤ err(π) ≤ Eπ [∥X1 − mπ (X0 )∥22 ] = E[∥X1 − E[X1 ]∥22 ] − E[∥mπ (X0 ) − E[X1 ]∥22 ].

(46)

Moreover, for every measurable g ∈ L2 (q0 ; Rd ), the following reduction holds true: err(π) ≤ Eπ [∥X1 − mπ (X0 )∥22 ] ≤ Eπ [∥X1 − g(X0 )∥22 ]. If π ∼ q0 ⊗ q1 and q1 is not a Dirac mass, then err(π) > 0. 32

(47)

Preprint. Under review.

Proof. The proof is quite straightforward. Specifically, the pair (mπ (X0 ), X1 ) couples (mπ )# q0 with q1 , which gives the first upper bound. Conditional orthogonality gives the equality in Eq. (46) and the predictor inequality in Eq. (47). Now suppose π ∼ q0 ⊗ q1 . For q0 -almost every x, the conditional law π(dy | X0 = x) is equivalent to the non-Dirac law q1 , so its conditional variance is strictly positive. If err(π) = 0, then (mπ )# q0 = q1 , while conditional orthogonality gives us that E[∥X1 ∥22 ] = E[∥mπ (X0 )∥22 ] + Eπ [∥X1 − mπ (X0 )∥22 ] > E[∥X1 ∥22 ], which is a contradiction, so this proves the lemma. We next prove the convergence to Monge coupling: Lemma B.11. Let T ∈ L2 (q0 ; Rd ) and let γn ∈ Π(q0 , q1 ) satisfy γn → (id, T )# q0 in W2 . Then, Eγn [∥X1 − T (X0 )∥22 ] −→ 0,

err(γn ) −→ 0.

(48)

Proof. Fix g ∈ Cb (Rd ; Rd ). Since (x, y) 7→ ∥y − g(x)∥22 is continuous with quadratic growth, W2 convergence and the common first marginal give lim sup Eγn [∥X1 − T (X0 )∥22 ] ≤ 4∥g − T ∥2L2 (q0 ) . n→∞

Bounded continuous functions are dense in L2 (q0 ), so the first limit follows. The second follows from Eq. (47) with g = T . Given these two supporting lemmas, we can now prove Proposition 4.18: Proof of Proposition 4.18. For population OT, X1 = T (X0 ) almost surely, so the one-step output has law q1 and err(π ⋆ ) = 0. Under independent coupling, E[X1 | X0 ] = E[X1 ], so we can show that err(πind ) = E[∥X1 − E[X1 ]∥22 ] > 0. We next prove the two population-entropic claims. First, we write objective as 1 Eπ [∥X1 − X0 ∥22 ] + ε KL(π∥πind ). 2d Since πind has finite objective, πε is absolutely continuous with respect to πind , and there exists a Radon–Nikodym density f such that dπε = f dπind . Jε (π) =

We claim that f > 0 almost everywhere, which can be proved via contradiction, i.e., otherwise if A = {f = 0} had positive πind -mass, the feasible mixture πs = (1 − s)πε + sπind would satisfy, with h(r) = r log r, that Eπind [∥X1 − X0 ∥22 ] − Eπε [∥X1 − X0 ∥22 ] Jε (πs ) − Jε (πε ) ≤ + ε{πind (A) log s − KL(πε ∥πind )} , s 2d which tends to −∞ as s ↓ 0 and contradicts the optimality, so πε ∼ πind , and Lemma B.10 gives err(πε ) > 0. We next compare πε strictly with independence. We choose w such that ⟨w, X1 − E[X1 ]⟩ is not a constant and set b(X1 ) = tanh⟨w, X1 − E[X1 ]⟩ − E[tanh⟨w, X1 − E[X1 ]⟩], and it’s clear that v := E[(X1 − E[X1 ])b(X1 )] ̸= 0. Since q0 ≪ Ld , the variable ⟨v, X0 − E[X0 ]⟩ is not a constant. We define a(X0 ) by applying and centering tanh in the same way, then rescaling it so that both ∥a∥∞ , ∥b∥∞ ≤ 1. For |s| < 1, the density 1 + s, a(x)b(y) defines a coupling ρs ∈ Π(q0 , q1 ), and its entropy derivative at s = 0 vanishes, whereas d 1 1 Eρ [∥X1 − X0 ∥22 ] = − ⟨E[(X0 − E[X0 ])a(X0 )], v⟩ < 0. ds 2d s d s=0 Therefore, Jε (πε ) ≤ Jε (ρs ) < Jε (πind ) for all sufficiently small s > 0, and non-negativity of relative entropy gives us that Eπε [∥X1 − X0 ∥22 ] < Eπind [∥X1 − X0 ∥22 ]. 33

Preprint. Under review.

Expanding the two costs shows that Eπε [X1 | X0 ] is not almost surely equal to E[X1 ]. The strict upper comparison err(πε ) < err(πind ) now follows from Eq. (46). Finally, choosing any minibatch batch size k, we denote (Z1 , . . . , Zk ) and (Y1 , . . . , Yk ) as the endpoint batches, let I be any uniform source index, and let τ be the empirical matching. If (q0 ⊗ q1 )(E) = 0, every candidate pair (Zi , Yj ) avoids E almost surely, so πk (E) = 0. Conversely, if (q0 ⊗ q1 )(E) > 0 and Ez = {y : (z, y) ∈ E}, then Z
πk (E) ≥ P Yj ∈ EZI for every j = q1 (Ez )k q0 (dz) > 0. Therefore, πk ∼ q0 ⊗ q1 , and Lemma B.10 gives err(πk ) > 0. Now, recall that Proposition 4.13 already proves the convergence of πk to W22 (q0 , q1 ). Also, the tightness, lower semicontinuity, and uniqueness of the Brenier coupling imply πk ⇒ π ⋆ . Since all couplings have the same product-space second moment, this convergence holds in W2 . Lemma B.11 therefore gives err(πk ) ≤ Eπk [∥X1 − T (X0 )∥22 ] −→ 0. Since every fixed finite baseline has positive error, we can conclude the strict error reduction. B.11

T HEOREM 4.19

Proof. For simplicity, we abbreviate the following notations: µi = E[Xi ], Vi = E[∥Xi − µi ∥22 ], and C(π) = Eπ [∥X1 − X0 ∥22 ] with Cind = C(πind ). Since q0 ≪ Ld and q0 is non-Dirac, we have V0 > 0. For any coupling π, expanding the two quadratic costs gives Cind − C(π) = 2Eπ [⟨X0 − µ0 , Eπ [X1 | X0 ] − µ1 ⟩]. Together with Cauchy–Schwarz and Lemma B.10, we have err(π) ≤ V1 −

{Cind − C(π)}2 . 4V0

(49)

To apply this inequality, we consider B as two sampled endpoint batches. Under the actual sampling law Pot , the conditional distribution of the selected indices (I, J) given B is π bk,ε . We can define a reference law Pref using the same batches but conditionally independent uniform indices, so (I, J) | B ∼ 11⊤ /k 2 . The map (B, I, J) 7→ (XI , YJ ) pushes Pot forward to πk,ε and Pref forward to πind , so the conditional KL chain rule, along with the same data processing proof structure from Lemma B.7, give us that 
 KL(πk,ε ∥πind ) ≤ KL(Pot ∥Pref ) = E KL π bk,ε 11⊤ /k 2 . (50) Moreover, selecting (I, J) according to π bk,ε gives us that C(πk,ε ) = E[⟨C, π bk,ε ⟩] . 2d Adding Eq. (50) to Eq. (51) and recalling the definition of Sk,ε , we have

(51)

C(πk,ε ) + ε KL(πk,ε ∥πind ) ≤ Sk,ε . 2d Since S1,ε = Cind /(2d), we obtain Cind − C(πk,ε ) ≥ 2d(S1,ε − Sk,ε ). Substitution into Eq. (49) proves the middle inequality in Eq. (22). Proposition 4.15 gives πk,ε ∼ πind at every finite k, so strict positivity follows from Lemma B.10, and we finish proving the first statement in the theorem. We now prove the second statement using the idea of induction, which the upper bound is nonincreasing and strict for k ≥ 2. First, we set n = k + 1 and fix two realized n-point batches. For each source index a and target index b, remove those two points, let P ab solve the resulting k × k problem, embed it into an n × n matrix 34

Preprint. Under review.

Qab by inserting a zero row and column, and set P = n−2 uniform n-point marginals. If M (a, b, i, j) = n−2 Qab ij ,

ab a,b Q . We note that this matrix has

P

R(a, b, i, j) =

1{i̸=a, j̸=b} , n2 k 2

then their (i, j)-marginals are respectively P and the uniform n × n matrix. Data processing gives 1 X KL(P ∥11⊤ /n2 ) ≤ 2 KL(P ab ∥11⊤ /k 2 ). n a,b

The same averaging identity holds for the cost, so taking expectations proves Sk+1,ε ≤ Sk,ε . We now prove the strictness part, and it’s enough for us to start with k = 2.
1 −1 Let U2 = 11⊤ /4 and D = −1 1 . The feasible perturbation U2 + sD has zero entropy derivative at s = 0, whereas its cost derivative is 1 − ⟨Z1 − Z2 , Y1 − Y2 ⟩. d Since q1 is non-Dirac, Y1 − Y2 ̸= 0 with positive probability; conditional on any such value, absolute continuity of q0 makes the displayed inner product nonzero almost surely. Therefore, the empirical optimum is strictly better than U2 with positive probability. Also, the nonnegative gap is integrable, so S2,ε < S1,ε , and monotonicity extends this strict inequality to every k ≥ 2. This proves both the strict final inequality and monotonicity of the upper bound in Eq. (22). We now proceed to prove the final Eq. (23). Similar to before, we write the population value as   C(π) Sε := inf + ε KL(π∥πind ) . 2d π∈Π(q0 ,q1 ) The index-to-measure reduction and uniform shadow bound in Lemmas B.8 and B.7 give the following, respectively, √ 2 Cind E[ak ] + E[a2k ] , ak = W2 (b q0,k , q0 ) + W2 (b q1,k , q1 ). (52) 0 ≤ Sk,ε − Sε ≤ ek := 2d For every q ∈ P2 (Rd ), E[W22 (b qk , q)] → 0 (i.e., project onto a fixed ball, use bounded-support empirical W2 convergence and dominated convergence, and then send the radius to infinity), so ek → 0. We set C⋆ = W22 (q0 , q1 ) and try to first show Sε → C⋆ /(2d). First, we note that the optimal coupling π ⋆ may have infinite relative entropy with respect to πind , so we do the following math technique by approximating it through finite-entropy couplings. Specifically, we first partition both endpoint spaces into increasingly fine finite partitions and set pAB = π ⋆ (A × B), and we set ρn by assigningPmass pAB to each block P A × B according to the product measure q0 (dx | A)q1 (dy | B). Since B pAB = q0 (A) and A pAB = q1 (B), ρn has marginals q0 and q1 . Moreover, we have X pAB KL(ρn ∥πind ) = pAB log < ∞. q0 (A)q1 (B) A,B

As the partitions become finer, ρn converges weakly to π ⋆ . Since all ρn have the same marginals as π ⋆ , their product-space second moments also agree, so ρn → π ⋆ in W2 . Therefore, C(ρn ) → C⋆ , and C⋆ C(ρn ) ≤ Sε ≤ + ε KL(ρn ∥πind ). 2d 2d To prove the claim, we simply need to take ε ↓ 0 and n → ∞. Now, let εk → 0. The raw selected-pair cost, i.e., Eq. (52), and the preceding population limit give C⋆ C(πk,εk ) C⋆ ≤ ≤ Sk,εk ≤ Sεk + ek −→ . 2d 2d 2d 35

Preprint. Under review.

Tightness and uniqueness of the Brenier coupling imply πk,εk ⇒ π ⋆ , and the common product-space second moment upgrades this to W2 convergence. Lemma B.11 therefore gives err(πk,εk ) → 0, with no rate condition relating k and εk . To derive Eq. (23), we still need to use the following continuity fact: let ρn , ρ ∈ Π(q0 , q1 ) satisfy ∥ρn − ρ∥TV → 0, and we can apply the technique by disintegrating them as q0 (dx)Kn (x, dy) and q0 (dx)K(x, dy). Their common first marginal gives Z ∥ρn − ρ∥TV = ∥Kn (x, ·) − K(x, ·)∥TV q0 (dx). For metric projection PL onto the radius-L ball, the corresponding truncated conditional means satisfy Z Z Z 2 PL (y)Kn (x, dy) − PL (y)K(x, dy) q0 (dx) ≤ 4L2 ∥ρn − ρ∥TV . 2

R Conditional Jensen bounds each truncation tail by ∥y − PL (y)∥22 q1 (dy), and we can take n → ∞ and then L → ∞, which proves convergence of the conditional means in L2 (q0 ). Coupling their pushforward laws through X0 ∼ q0 gives err(ρn ) → err(ρ), and we now derive the final ingredient. Fix ε > 0, and recall the population objective C(ρ) + ε KL(ρ∥πind ), Sε = Jε (πε ). 2d Since πk,ε ∈ Π(q0 , q1 ), population optimality and the preceding bounds on the averaged coupling and empirical value give Sε ≤ Jε (πk,ε ) ≤ Sk,ε ≤ Sε + ek . Therefore, the population objective gap satisfies Jε (ρ) :=

0 ≤ Jε (πk,ε ) − Jε (πε ) ≤ ek −→ 0. To translate this objective convergence into convergence of the couplings, we apply the entropy Pythagorean inequality (see Lemma 4.4 in Eckstein & Nutz (2022)) to the rescaled objective Jε /ε. This gives ek Jε (πk,ε ) − Jε (πε ) ≤ . KL(πk,ε ∥πε ) ≤ ε ε Using the convention ∥µ − ν∥TV = supA |µ(A) − ν(A)|, Pinsker’s inequality then yields r r 1 ek ∥πk,ε − πε ∥TV ≤ KL(πk,ε ∥πε ) ≤ −→ 0, 2 2ε where the limit follows because ε > 0 is fixed and ek → 0. Both couplings have the same marginals q0 , q1 , so the total-variation continuity established above applies and gives err(πk,ε ) −→ err(πε )

as k → ∞.

This proves the fixed-regularization convergence claimed in the theorem. For any fixed finite k0 , Lemma B.9 gives πk0 ,ε → πkME in total variation and πkME ∼ πind . With the 0 ,0 0 ,0 continuity, along with Lemma B.10, we have err(πk0 ,ε ) −→ err(πkME ) > 0, 0 ,0 i.e., the candidate error tends to zero, so Eq. (23) holds for all sufficiently large k, and we can finally conclude the theorem.

C

F URTHER E MPIRICAL R ESULTS

Additional 2D examples. We provide additional experimental results and details. First, we present additional 2D synthetic experiment on rings in Figure 4, using the same training settings as the multimodal block synthetic experiment in Figure 1. The results show a similar conclusion, where larger minibatch sizes consistently lead to more stable training and lower flow-generation error. 36

Preprint. Under review.

(a) Sampled couplings Independent

(b) Projected training gradients OT-128

Independent

(c) Gradient variance and flow MSE

OT-128 Gradient variance Flow MSE

EOT-8192

EOT-512

EOT-8192

10 2

Flow MSE

CFM loss + backprop

EOT-512

2 grad, ( )

10

1

10 3

0.1 10 4

Source

Target

Optimal coupling

Suboptimal coupling

0

Projected gradient

2048

Minibatch size k

8192

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. 0.60

Gradient variance

Exact W1

2.60

2.50

2.40

2.30

0.0025

0.001

0.0005

0.00025

0.50

0.40

0.30

0.20

0.0001

ε

0.0025

0.001

0.0005

0.00025

0.0001

ε

Figure 5: Ablation over ε on the PBMC dataset with an 8K entropic OT context, with the exact W1 error of generated samples (left) and training gradient variance (right) shown. We sweep ε from {0.0025, 0.0015, 0.001, 0.0009, 0.00075, 0.0006, 0.0005, 0.0004, 0.00025, 0.0001}. Ablation on ε. We also ablate ε on the PBMC dataset, whose size lets us reach substantially larger minibatch contexts than the EB dataset. Offline pipelines that scale to million-level contexts must drive ε to very low levels to keep the coupling sharp, which makes its tuning delicate. Online couplings instead pair each freshly sampled source point within the current minibatch, so that a moderate context of a few thousand points already yields a clear scaling effect while remaining robust to the choice of ε. As shown in Figure 5, we sweep ε from 0.0001 to 0.0025 at 8K minibatch entropic OT context. Following the exactly same training configuration detailed in §5, each ε trial is running over 5 seeds, with the mean performance and its standard deviation shown. Performance peaks around ε = 0.001 and varies little across the swept range. For EB training, we use a constant ε = 0.01 across all context sizes, since the largest context can only be scaled to 1.5K. For image flow training, where the context can be scaled similarly to PBMC, the optimal ε at the largest minibatch context is 0.002 for CIFAR and 0.001 for ImageNet, consistent with the 8K-context ablation.

37

Related documents

• Scaling Online Optimal Transport Minibatches for Flow Matching#1344328
• Beyond the Manifold Hypothesis: Hybrid Spectral Parameterizations for Flow Matching#1116763
• Prediction--Loss Alignment for Sampler--Robust Flow Matching Training#998105
• Universality and Convergence of Generative Flows#1201503
• Theoretical guarantees in KL for Diffusion Flow Matching#1203887