Introduction
Modern generative AI relies heavily on Flow Matching models to transform simple noise into complex data. However, the efficiency of these models is often limited by the "tangled" paths they learn during training. This paper demonstrates that Flow Matching performance improves significantly when models are trained on massive Optimal Transport (OT) couplings involving millions of points. By scaling the Sinkhorn algorithm to unprecedented sizes and introducing a scale-invariant "renormalized entropy" metric, the authors achieve straighter trajectories that lead to higher image quality and much faster inference.
The Problem: Curvy Paths and the Curse of Dimensionality
Generative flow models learn a "velocity field" to map noise to data. Ideally, these models should follow straight lines—the shortest path—to simplify simulation during inference.
In practice, training often uses "Independent Flow Matching" (I-FM), pairing random noise with random images. This results in curvy, inefficient paths that require many steps of an ODE solver to navigate. While researchers have attempted to use Optimal Transport (OT) on small mini-batches (around 256 points) to "untangle" these paths, these small batches suffer from the curse of dimensionality. They provide a poor approximation of the global transport map, leaving the model with a noisy and inefficient signal.
Scaling Batch-OT to Millions of Points
The authors propose scaling "Batch-OT" by three to four orders of magnitude, moving from hundreds of points to millions. This is achieved through several technical innovations:
- Sharded Sinkhorn: Utilizing the OTT-JAX library to distribute the Sinkhorn algorithm across multiple GPUs and nodes, enabling couplings for up to $n \approx 2 \times 10^6$ points.
- Renormalized Entropy: A new metric, $\mathcal{E}(\mathbf{P})$, is introduced to map the "blurriness" of a Sinkhorn coupling to a scale of 0 (perfectly sharp) to 1 (completely independent). This allows for robust tuning of the regularization parameter $\epsilon$ across different dimensions and batch sizes.
- Numerical Stability and Efficiency: The use of negative dot-products instead of squared Euclidean distances improves stability. Additionally, techniques like warm-starting and PCA dimensionality reduction speed up the OT solver by up to 10x without sacrificing quality.
Results: Straighter Paths and Faster Inference
The core finding of this research is that the benefits of OT in flow matching do not saturate at small batch sizes. Increasing the coupling size $n$ consistently reduces the curvature of learned trajectories, leading to "free speed" during inference.
Quantitatively, the improvements are stark. On ImageNet-64, using large-batch OT significantly reduced the FID score at low step counts. For instance, at 4 integration steps (NFE=4), the FID dropped from 80.1 with Independent FM to 46.9 with Large Batch OT ($n=131,072$).
Across all benchmarks, including CIFAR-10 and ImageNet, larger batch sizes and lower entropy led to significantly lower curvature. The authors successfully ran Sinkhorn on over 2 million points, a scale previously considered prohibitive for training flow models. They argue that while computing these couplings is computationally intensive, it is a one-time "data prepping" cost that is negligible compared to the overall training of the neural network.
Conclusion and Future Work
By scaling Optimal Transport to millions of points, this work provides a path toward more efficient generative models that require fewer steps to produce high-quality results. Future directions include exploring how these large-scale couplings interact with "Reflow" procedures and whether they can eliminate the need for multiple training rounds entirely. While some limitations exist regarding dataset size for smaller benchmarks like CIFAR-10, the scalability of this approach offers a promising new standard for training high-performance flow models.