On Fitting Flow Models with Large Sinkhorn Couplings

Alireza Mousavi-Hosseini, Stephen Y. Zhang, Michal Klein, marco cuturi

Keywords: Optimal Transport, Flow Matching, Sinkhorn, Generative Modeling

TL;DR: We scale the Sinkhorn algorithm to large batch sizes of ~ 1M samples, and show that with a proper choice of entropic regularization, the resulting coupling can significantly improve flow matching.

Abstract:

Flow models transform data gradually from one modality (e.g. noise) onto another (e.g. images) following 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) coupling, since this would lead to an optimal flow in the sense of Benamou and Brenier.
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\approx 256$, and considered OT solvers that return couplings that are either sharp (using e.g. the Hungarian algorithm) or blurred (using the Sinkhorn algorithm).
We follow in the footsteps of these works by exploring the benefits of increasing this mini-batch size $n$ by three to four orders of magnitude, and look more carefully on the effect of the entropic regularization $\varepsilon$ used in Sinkhorn. Our analysis is facilitated by new scale invariant quantities to report the sharpness of a coupling, while our sharded computations across multiple GPUs and 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$.

Submission Number: 73