cs.LGMar 12, 2026

Sinkhorn-Drifting Generative Models

Authors: Ping He, Om Khangaonkar, Hamed Pirsiavash, Yikun Bai, Soheil Kolouri

Organizations: Department of Computer Science, Vanderbilt University · Department of Computer Science, University of California, Davis · Department of Computer Science, Purdue University

Read PDF
Ask AI
Cite

Copy citation

DOINot available for this paper

Abstract

We establish a theoretical link between the recently proposed "drifting" generative dynamics and gradient flows induced by the Sinkhorn divergence. In a particle discretization, the drift field admits a cross-minus-self decomposition: an attractive term toward the target distribution and a repulsive/self-correction term toward the current model, both expressed via one-sided normalized Gibbs kernels. We show that Sinkhorn divergence yields an analogous cross-minus-self structure, but with each term defined by entropic optimal-transport couplings obtained through two-sided Sinkhorn scaling (i.e., enforcing both marginals). This provides a precise sense in which drifting acts as a surrogate for a Sinkhorn-divergence gradient flow, interpolating between one-sided normalization and full two-sided Sinkhorn scaling. Crucially, this connection resolves an identifiability gap in prior drifting formulations: leveraging the definiteness of the Sinkhorn divergence, we show that zero drift (equilibrium of the dynamics) implies that the model and target measures match. Experiments show that Sinkhorn drifting reduces sensitivity to kernel temperature and improves one-step generative quality, trading off additional training time for a more stable optimization, without altering the inference procedure used by drift methods. These theoretical gains translate to strong low-temperature improvements in practice: on FFHQ-ALAE at the lowest temperature setting we evaluate, Sinkhorn drifting reduces mean FID from 187.7 to 37.1 and mean latent EMD from 453.3 to 144.4, while on MNIST it preserves full class coverage across the temperature sweep. Project page: https://mint-vu.github.io/SinkhornDrifting/

Figures & tables

Figure 1: Drift trajectories across different values of τ\tau under one-sided, two-sided, and Sinkhorn normalization, with and without self-distance masking. The Sinkhorn trajectories are generated with a fixed number of iterations, 3030 .
Figure 2: Generative model training on 2D distributions across τ∈{0.01,0.05,0.1}\tau\in\{0.01,0.05,0.1\} and three normalization schemes (one-sided, two-sided, Sinkhorn) with Gaussian kernel. Left six columns : final generated samples (orange) vs. target (blue). Right two columns : W22W_{2}^{2} convergence curves over 5,0005{,}000 iterations. Sinkhorn consistently achieves lower W22W_{2}^{2} and better mode coverage, especially at small τ\tau .
 EMD ↓\downarrowAccuracy ↑\uparrow EMD ↓\downarrowAccuracy ↑\uparrow
τ\tauBaselineSinkhornBaselineSinkhornτ\tauBaselineSinkhornBaselineSinkhorn
0.00573.218.579.97%†9.97\%^{\dagger}100.00%0.03073.226.969.95%†9.95\%^{\dagger}100.00%
0.01073.217.719.99%†9.99\%^{\dagger}100.00%0.04077.196.9110.00%†10.00\%^{\dagger}100.00%
0.02073.217.129.98%†9.98\%^{\dagger}100.00%0.05079.736.889.96%†9.96\%^{\dagger}100.00%
0.02573.217.019.99%†9.99\%^{\dagger}100.00%0.1005.636.9195.00%100.00%
Table 1: MNIST τ\tau -sweep with Gaussian kernel. † Acc ≈ ⁣10%\approx\!10\% indicates mode collapse.
Figure 3: Generated MNIST samples (Gaussian kernel). Each panel shows 10 classes ×\times 8 samples. (a) Baseline at τ=0.01\tau=0.01 collapses to a single degenerate mode; class accuracy is ≈10%{\approx}10\% (random chance). (b) Sinkhorn at τ=0.01\tau=0.01 correctly generates all ten classes with 100% class accuracy. (c,d) At τ=0.1\tau=0.1 both methods produce recognizable digits, with Sinkhorn remaining sharper and more consistent.
τ\tauMaleFemale
ClassEMD ↓\downarrowFID ↓\downarrowClassEMD ↓\downarrowFID ↓\downarrow
BaselineSinkhornBaselineSinkhornBaselineSinkhornBaselineSinkhorn
0.1Adult441.1145.0198.733.8Adult414.0146.5143.739.7
Children489.8157.6228.540.8Children465.8149.1185.638.6
Old508.7138.5219.131.4Old400.6130.0150.538.3
Average453.3144.4187.737.1     
Table 2: Class-conditional face generation on FFHQ with ALAE across temperatures τ\tau . We report EMD ↓\downarrow and FID ↓\downarrow for Baseline and Sinkhorn. Average is the mean value of all 6 classes for each τ\tau .
Figure 4: Qualitative comparison of class-conditional FFHQ generation at τ=1.0\tau{=}1.0 (top) and τ=10.0\tau{=}10.0 (bottom). In each panel, each row corresponds to one class; Baseline is on the left and Sinkhorn is on the right . The corresponding low-temperature qualitative panel ( τ=0.1\tau{=}0.1 ) is shown in Figure 10 of Appendix L .
Appendix figures & tables7 assets

Supplementary material from the paper’s appendix.

Appendix

Figure 5: Full drift-trajectory grid for the temperature sweep in Section 5.1 . Columns correspond to one-sided, two-sided, and Sinkhorn normalization, with and without self-distance masking where applicable; rows correspond to different values of τ\tau . Sinkhorn exhibits the most stable trajectories as τ\tau decreases, whereas one-sided and two-sided normalization become increasingly sensitive in the low-temperature regime. Empirically, for any choice of τ\tau , the Sinkhorn method admits convergence that is at least as good as (and often better than) the classical one-/two-sided Drift methods.
Figure 6: 2D generative training on 2-Moons and Spiral with a Gaussian kernel across τ∈{0.01,0.05,0.1}\tau\in\{0.01,0.05,0.1\} and three normalizations (one-sided, two-sided, Sinkhorn). Left: final samples (orange) vs. targets (blue); right: W22W_{2}^{2} over 5,0005{,}000 iterations. Sinkhorn gives better coverage and lower W22W_{2}^{2} , especially at small τ\tau .
Figure 7: 2D generative training on 8-Gaussians and Checkerboard with a Laplacian kernel across τ∈{0.01,0.05,0.1}\tau\in\{0.01,0.05,0.1\} and three normalizations (one-sided, two-sided, Sinkhorn). Left : final samples (orange) vs. targets (blue); Right : W22W_{2}^{2} over 5,0005{,}000 iterations. Sinkhorn gives better coverage and lower W22W_{2}^{2} , especially at small τ\tau .
Figure 8: 2D generative training on 2-Moons and Spiral with a Laplacian kernel across τ∈{0.01,0.05,0.1}\tau\in\{0.01,0.05,0.1\} and three normalizations (one-sided, two-sided, Sinkhorn). Left : final samples (orange) vs. targets (blue); Right : W22W_{2}^{2} over 5,0005{,}000 iterations. Sinkhorn gives better coverage and lower W22W_{2}^{2} , especially at small τ\tau .
 EMD ↓\downarrowAccuracy ↑\uparrow EMD ↓\downarrowAccuracy ↑\uparrow
τ\tauBaselineSinkhornBaselineSinkhornτ\tauBaselineSinkhornBaselineSinkhorn
0.00573.215.409.97%†9.97\%^{\dagger}99.97%0.0305.706.4894.86%99.98%
0.01073.217.149.99%†9.99\%^{\dagger}100.00%0.0404.716.3896.07%99.97%
0.02077.106.6710.00%†10.00\%^{\dagger}100.00%0.0504.466.3496.17%99.97%
0.02512.256.5584.21%‡84.21\%^{\ddagger}100.00%0.1004.186.6196.56%100.00%
Appendix
Table 3: Additional Results for MNIST Experiments. † Accuracy ≈10%\approx 10\% indicates mode collapse. ‡ Partial collapse.
Figure 9: Generated MNIST samples (Laplacian kernel). (a) Baseline collapses at τ=0.01\tau=0.01 . (b) Sinkhorn generates all classes correctly at τ=0.01\tau=0.01 . (c) Baseline at τ=0.05\tau=0.05 shows partial recovery but with visible class confusion. (d) Sinkhorn at τ=0.05\tau=0.05 remains stable.
Figure 10: Additional qualitative comparison for class-conditional FFHQ generation at τ=0.1\tau{=}0.1 . Each row corresponds to one class; Baseline is on the left and Sinkhorn is on the right .
More in this topic

Sinkhorn Divergence

See all 9 papers

Explore similar work

CardsList
Mar 10, 2026

Generative Drifting is Secretly Score Matching: a Spectral and Variational Perspective

Generative Modeling via Drifting~\citep{deng2026drifting} has recently achieved state-of-the-art one-step image generation through a kernel-based drift operator, yet its success is largely empirical and its theoretical foundations remain poorly understood. We observe that under a Gaussian kernel, the drift operator is exactly a score difference on smoothed distributions. This answers three questions left open in the original work: (1) whether a vanishing drift guarantees equality of distributions (Vp,q=0⇒p=qV_{p,q}=0\Rightarrow p=q), (2) how to choose between kernels, and (3) why the stop-gradient operator is indispensable for stable training. Our observations position drifting within the score-matching family. By linearizing the McKean-Vlasov dynamics and probing them in Fourier space, we reveal frequency-dependent convergence timescales comparable to Landau damping in plasma kinetic theory: the Gaussian kernel suffers an exponential high-frequency bottleneck, potentially explaining the empirical preference for the Laplacian kernel. This suggests a fix: an exponential bandwidth annealing schedule σ(t)=σ0e−rtσ(t)=σ_0 e^{-rt} that reduces convergence time from exp⁡(O(Kmax⁡2))\exp(O(K_{\max}^2)) to O(log⁡Kmax⁡)O(\log K_{\max}). Finally, by formalizing drifting as a Wasserstein gradient flow of the smoothed KL divergence, we prove that the stop-gradient operator is not a heuristic but is derived from the frozen-field discretization mandated by the Jordan-Kinderlehrer-Otto (JKO) scheme, and removing it severs training from any gradient-flow guarantee. This variational perspective further provides a general template for constructing novel drift operators, which we demonstrate with a Sinkhorn divergence drift. We validate our analysis on toy datasets and scale it up to ImageNet.

Erkan Turan, Nicolas Dufour, Maks Ovsjanikov

LIX, Ecole Polytechnique, IP Paris · Kyutai

Drifting ModelsScore Matching
AskPDF
May 12, 2026

DriftXpress: Faster Drifting Models via Projected RKHS Fields

Drifting Models have emerged as a new paradigm for one-step generative modeling, achieving strong image quality without iterative inference. The premise is to replace the iterative denoising process in diffusion models with a single evaluation of a generator. However, this creates a different trade-off: drifting reduces inference cost by moving much of the computation into training. We introduce DriftXpress, an accelerated formulation of drifting models based on projected RKHS fields. DriftXpress approximates the drifting kernel in a low-rank feature space. This preserves the attraction-repulsion structure of the original drifting field while reducing the cost of field evaluation. Across image-generation benchmarks, DriftXpress achieves comparable FID to standard drifting while reducing wall-clock training cost. These results show that the training-inference trade-off of drifting models can be pushed further without giving up their one-step inference advantage.

Ali Falahati, Elliot Creager, Gautam Kamath +1

Cheriton School of Computer Science, University of Waterloo and Vector Institute. · Department of Electrical and Computer Engineering, University of Waterloo and Vector Institute. · A*STAR Centre for Frontier AI Research, Singapore.

Image GenerationOne-Step Generative ModelingKernel MethodsDrifting Models
AskPDF
May 11, 2026

Kernel-Gradient Drifting Models

We propose kernel-gradient drifting, a one-step generative modeling framework that replaces the fixed Euclidean displacement direction in drifting models with directions induced by the kernel itself. Standard drifting is attractive because it enables fast, high-quality generation without distilling a large pretrained diffusion model, but its theory is currently understood mainly for Gaussian kernels, where the drift coincides with smoothed score matching and is identifiable. Our gradient-based reformulation exposes this score-based structure for general kernels: the resulting drift is the score difference between kernel-smoothed data and model distributions, yielding identifiability for characteristic kernels and a smoothed-KL descent interpretation of the drifting dynamics. Since kernel gradients are intrinsic tangent vectors, the same construction extends naturally to Riemannian manifolds and to discrete data via the Fisher-Rao geometry of the probability simplex. Across spherical geospatial data, promoter DNA and molecule generation, kernel-gradient drifting enables state-of-the-art one-step generation beyond the Euclidean setting without distillation.

Maria Esteban-Casadevall, Jorge Carrasco-Pollo, Max Welling +3

AMLab · University of Amsterdam · CuspAI +1

Synthetic Data GenerationOne-Step Generative ModelingGenerative ModelingDrifting Models
AskPDF