Research map / Probabilistic and scientific modeling / Variational and generative sampling
On Fitting Flow Models with Large Sinkhorn Couplings
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Scales mini-batch entropic OT couplings for flow matching to batch sizes up to ~2M with multi-GPU Sinkhorn, introduces a scale-free renormalized coupling entropy, and shows large-batch, low-epsilon Batch-OT flows beat independent flow matching on synthetic and ImageNet/CIFAR generation.
Problem
Independent-coupling flow matching trains slow and yields curved trajectories that are costly to integrate, while batch-OT pairing has so far been limited to small mini-batches (~256) and a single regularization value.
Approach
Treat Hungarian (epsilon to 0) and independent coupling (epsilon to infinity) as endpoints of entropic OT, scale mini-batch OT to millions of points via sharded Sinkhorn, and quantify coupling sharpness with a scale-free renormalized entropy in [0,1].
Key results
- On ImageNet-64 with Euler 4 NFE, FID drops from 80.1 for I-FM to 46.9 with OT batch size 131072 (epsilon=0.1).
- On ImageNet-32 adaptive Dopri5, FID improves from 5.55 (I-FM) to 4.85 with batch size 524288 at epsilon=0.1.
- Sinkhorn with dot-product cost stays convergent at small epsilon where squared-Euclidean cost fails within 50,000 iterations; PCA reduces Sinkhorn time ~10x with no significant FID change.
Surprising findings
- At the Dopri5 adaptive solver, the gains from large batch OT are much smaller since adaptive solvers rely less on flow straightness.
- Using dot-product cost instead of squared-Euclidean avoids numerical failure of Sinkhorn at low epsilon in single precision.
Limitations
- Hyperparameters such as learning rates were not re-tuned across batch sizes and epsilon values, reusing I-FM settings.
- Large-scale Sinkhorn computation is expensive; gains for very large problems are less clear.
- CIFAR-10 is small, so large batches resample duplicate data, limiting conclusions there.
Experimental setup
Datasets: CIFAR-10 ImageNet-32 ImageNet-64 piecewise affine OT map (synthetic) Korotin benchmark (Gaussian mixture with ICNN map)
Benchmarks: CIFAR-10 ImageNet-32 ImageNet-64 piecewise affine OT map benchmark Korotin benchmark
Metrics: FID curvature reconstruction loss bits per dimension renormalized coupling entropy
Baselines: Independent-FM Batch OT-FM (Hungarian) Batch OT-FM (entropic) Korotin benchmark ICNN solvers
Builds on: flow matching Batch OT-FM Sinkhorn algorithm stochastic interpolants OTT-JAX
Closest papers by meaning
- Faster Inference of Flow-Based Generative Models via Improved Data-Noise Coupling (ICLR 2025 · similarity 0.87)
- Adaptive Flow Matching for Resolving Small-Scale Physics (ICML 2025 · similarity 0.87)
- Simple ReFlow: Improved Techniques for Fast Flow Models (ICLR 2025 · similarity 0.86)
- Shortcutting Pre-trained Flow Matching Diffusion Models is Almost Free Lunch (NeurIPS 2025 · similarity 0.86)
- Beckmann Transport Models: From Autonomous Flows to One-Step Maps (NeurIPS 2026 · similarity 0.86)
- Bespoke Solvers for Generative Flow Models (ICLR 2024 · similarity 0.86)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.