On Fitting Flow Models with Large Sinkhorn Couplings

• URL: http://arxiv.org/abs/2506.05526v2
• Date: Mon, 09 Jun 2025 12:15:36 GMT
• Title: On Fitting Flow Models with Large Sinkhorn Couplings
• Authors: Michal Klein, Alireza Mousavi-Hosseini, Stephen Zhang, Marco Cuturi,
• Abstract summary: Flow models transform data gradually from one modality (e.g. noise) onto another (e.g. images)<n>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.<n>We show that flow models greatly benefit when fitted with large Sinkhorn couplings, with a low entropic regularization.
• Score: 21.33171004803672
• License: http://creativecommons.org/licenses/by/4.0/
• Abstract: 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\approx 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$.

Related papers

• VAE-DNN: Energy-Efficient Trainable-by-Parts Surrogate Model For Parametric Partial Differential Equations [49.1574468325115]
  We propose a trainable-by-parts surrogate model for solving forward and inverse parameterized nonlinear partial differential equations.<n>The proposed approach employs an encoder to reduce the high-dimensional input $y(bmx)$ to a lower-dimensional latent space, $bmmu_bmphi_y$.<n>A fully connected neural network is used to map $bmmu_bmphi_y$ to the latent space, $bmmu_bmphi_h$, of the P
  arXiv  Detail & Related papers  (2025-08-05T18:37:32Z)
• Self-Ensembling Gaussian Splatting for Few-Shot Novel View Synthesis [55.561961365113554]
  3D Gaussian Splatting (3DGS) has demonstrated remarkable effectiveness in novel view synthesis (NVS)<n>In this paper, we introduce Self-Ensembling Gaussian Splatting (SE-GS)<n>We achieve self-ensembling by incorporating an uncertainty-aware perturbation strategy during training.<n> Experimental results on the LLFF, Mip-NeRF360, DTU, and MVImgNet datasets demonstrate that our approach enhances NVS quality under few-shot training conditions.
  arXiv  Detail & Related papers  (2024-10-31T18:43:48Z)
• Truncated Consistency Models [57.50243901368328]
  Training consistency models requires learning to map all intermediate points along PF ODE trajectories to their corresponding endpoints.<n>We empirically find that this training paradigm limits the one-step generation performance of consistency models.<n>We propose a new parameterization of the consistency function and a two-stage training procedure that prevents the truncated-time training from collapsing to a trivial solution.
  arXiv  Detail & Related papers  (2024-10-18T22:38:08Z)
• Winner-Take-All Column Row Sampling for Memory Efficient Adaptation of Language Model [89.8764435351222]
  We propose a new family of unbiased estimators called WTA-CRS, for matrix production with reduced variance. Our work provides both theoretical and experimental evidence that, in the context of tuning transformers, our proposed estimators exhibit lower variance compared to existing ones.
  arXiv  Detail & Related papers  (2023-05-24T15:52:08Z)
• Smoothly Giving up: Robustness for Simple Models [30.56684535186692]
  Examples of algorithms to train such models include logistic regression and boosting. We use $Served-Served joint convex loss functions, which tunes between canonical convex loss functions, to robustly train such models. We also provide results for boosting a COVID-19 dataset for logistic regression, highlighting the efficacy approach across multiple relevant domains.
  arXiv  Detail & Related papers  (2023-02-17T19:48:11Z)
• FeDXL: Provable Federated Learning for Deep X-Risk Optimization [105.17383135458897]
  We tackle a novel federated learning (FL) problem for optimizing a family of X-risks, to which no existing algorithms are applicable. The challenges for designing an FL algorithm for X-risks lie in the non-decomability of the objective over multiple machines and the interdependency between different machines.
  arXiv  Detail & Related papers  (2022-10-26T00:23:36Z)
• An Improved Analysis of Gradient Tracking for Decentralized Machine Learning [34.144764431505486]
  We consider decentralized machine learning over a network where the training data is distributed across $n$ agents. The agent's common goal is to find a model that minimizes the average of all local loss functions. We improve the dependency on $p$ from $mathcalO(p-1)$ to $mathcalO(p-1)$ in the noiseless case.
  arXiv  Detail & Related papers  (2022-02-08T12:58:14Z)
• STEM: A Stochastic Two-Sided Momentum Algorithm Achieving Near-Optimal Sample and Communication Complexities for Federated Learning [58.6792963686231]
  Federated Learning (FL) refers to the paradigm where multiple worker nodes (WNs) build a joint model by using local data. It is not clear how to choose the WNs' minimum update directions, the first minibatch sizes, and the local update frequency. We show that there is a trade-off curve between local update frequencies and local mini sizes, on which the above complexities can be maintained.
  arXiv  Detail & Related papers  (2021-06-19T06:13:45Z)

This list is automatically generated from the titles and abstracts of the papers in this site.

This site does not guarantee the quality of this site (including all information) and is not responsible for any consequences.