# On Fitting Flow Models with Large Sinkhorn Couplings

arXiv: https://arxiv.org/abs/2506.05526 · HTML: https://mlsanity.com/paper/2506.05526/

## Key facts

- Cited by 11 papers in the map (data through 14 Oct 2026).
- Builds on 28 of its 49 references that are in the map.
- First posted on arXiv 5 Jun 2025.

## Builds on (most cited first, 28 of 28)

- [Flow Matching for Generative Modeling](https://mlsanity.com/paper/2210.02747/) (2022)
- [Neural Ordinary Differential Equations](https://mlsanity.com/paper/1806.07366/) (2018)
- [Flow Straight and Fast: Learning to Generate and Transfer Data with Rectified Flow](https://mlsanity.com/paper/2209.03003/) (2022)
- [Sinkhorn Distances: Lightspeed Computation of Optimal Transportation Distances](https://mlsanity.com/paper/1306.0895/) (2013)
- [Variational Inference with Normalizing Flows](https://mlsanity.com/paper/1505.05770/) (2015)
- [Building Normalizing Flows with Stochastic Interpolants](https://mlsanity.com/paper/2209.15571/) (2022)
- [Improving and generalizing flow-based generative models with minibatch optimal transport](https://mlsanity.com/paper/2302.00482/) (2023)
- [Computational Optimal Transport](https://mlsanity.com/paper/1803.00567/) (2018)
- [FFJORD: Free-form Continuous Dynamics for Scalable Reversible Generative Models](https://mlsanity.com/paper/1810.01367/) (2018)
- [Input Convex Neural Networks](https://mlsanity.com/paper/1609.07152/) (2016)
- [A Downsampled Variant of ImageNet as an Alternative to the CIFAR datasets](https://mlsanity.com/paper/1707.08819/) (2017)
- [Flow Matching Guide and Code](https://mlsanity.com/paper/2412.06264/) (2024)
- [Rectified Flow: A Marginal Preserving Approach to Optimal Transport](https://mlsanity.com/paper/2209.14577/) (2022)
- [Multisample Flow Matching: Straightening Flows with Minibatch Couplings](https://mlsanity.com/paper/2304.14772/) (2023)
- [On Neural Differential Equations](https://mlsanity.com/paper/2202.02435/) (2022)
- [Equivariant flow matching](https://mlsanity.com/paper/2306.15030/) (2023)
- [Optimal Transport Tools (OTT): A JAX Toolbox for all things Wasserstein](https://mlsanity.com/paper/2201.12324/) (2022)
- [Action Matching: Learning Stochastic Dynamics from Samples](https://mlsanity.com/paper/2210.06662/) (2022)
- [Do Neural Optimal Transport Solvers Work? A Continuous Wasserstein-2 Benchmark](https://mlsanity.com/paper/2106.01954/) (2021)
- [Minimizing Trajectory Curvature of ODE-based Generative Models](https://mlsanity.com/paper/2301.12003/) (2023)
- [Learning with minibatch Wasserstein : asymptotic and gradient properties](https://mlsanity.com/paper/1910.04091/) (2019)
- [Statistical bounds for entropic optimal transport: sample complexity and the central limit theorem](https://mlsanity.com/paper/1905.11882/) (2019)
- [Non-Denoising Forward-Time Diffusions](https://mlsanity.com/paper/2312.14589/) (2023)
- [Statistical optimal transport](https://mlsanity.com/paper/2407.18163/) (2024)
- [Simple ReFlow: Improved Techniques for Fast Flow Models](https://mlsanity.com/paper/2410.07815/) (2024)
- [A note on overrelaxation in the Sinkhorn algorithm](https://mlsanity.com/paper/2012.12562/) (2020)
- [Modeling Stochastic Conditional Dynamics from Sparse Observations via Kernel-Stabilized Flow Matching](https://mlsanity.com/paper/2411.08314/) (2024)
- [EquiFlow: Equivariant Conditional Flow Matching with Optimal Transport for 3D Molecular Conformation Prediction](https://mlsanity.com/paper/2412.11082/) (2024)

## Built on by (most cited first, 11 of 11)

- [Flow Matching with Semidiscrete Couplings](https://mlsanity.com/paper/2509.25519/) (2025)
- [Diffusion Models in Simulation-Based Inference: A Tutorial Review](https://mlsanity.com/paper/2512.20685/) (2025)
- [Counterfactual Identifiability via Dynamic Optimal Transport](https://mlsanity.com/paper/2510.08294/) (2025)
- [Weighted Conditional Flow Matching](https://mlsanity.com/paper/2507.22270/) (2025)
- [The Coupling Within: Flow Matching via Distilled Normalizing Flows](https://mlsanity.com/paper/2603.09014/) (2026)
- [Optimal Transport Flow Matching by Design](https://mlsanity.com/paper/2606.04092/) (2026)
- [Federated Flow Matching](https://mlsanity.com/paper/2509.21250/) (2025)
- [Quantifying the noise sensitivity of the Wasserstein metric for images](https://mlsanity.com/paper/2510.01015/) (2025)
- [DrainSinkhorn: Safe Elimination for Batched Entropic Optimal Transport](https://mlsanity.com/paper/2607.24741/) (2026)
- [Beyond Straightness: Non-Crossing Flow Matching via Quantile AlignTree Coupling](https://mlsanity.com/paper/2609.03443/) (2026)
- [Scaling Online Optimal Transport Minibatches for Flow Matching](https://mlsanity.com/paper/2610.18838/) (2026)

## Abstract (arXiv)

Flow models transform data gradually from one modality (e.g. noise) onto another (e.g. images). Such models are parameterized by a time-dependent velocity field, trained to fit segments connecting pairs of source and target points. When the pairing between source and target points is given, training flow models boils down to a supervised regression problem. When no such pairing exists, as is the case when generating data from noise, training flows is much harder. A popular approach lies in picking source and target points independently. This can, however, lead to velocity fields that are slow to train, but also costly to integrate at inference time. In theory, one would greatly benefit from training flow models by sampling pairs from an optimal transport (OT) measure coupling source and target, since this would lead to a highly efficient flow solving the Benamou and Brenier dynamical OT problem. In practice, recent works have proposed to sample mini-batches of n source and n target points and reorder them using an OT solver to form better pairs. These works have advocated using batches of size n≈ 256, and considered OT solvers that return couplings that are either sharp (using e.g. the Hungarian algorithm) or blurred (using e.g. entropic regularization, a.k.a. Sinkhorn). We follow in the footsteps of these works by exploring the benefits of increasing n by three to four orders of magnitude, and look more carefully on the effect of the entropic regularization varepsilon used in the Sinkhorn algorithm. Our analysis is facilitated by new scale invariant quantities to report the sharpness of a coupling, while our sharded computations across multiple GPU or GPU nodes allow scaling up n. We show that in both synthetic and image generation tasks, flow models greatly benefit when fitted with large Sinkhorn couplings, with a low entropic regularization varepsilon.
