Sinkhorn-Drifting Generative Models

Ping He; Om Khangaonkar; Hamed Pirsiavash; Yikun Bai; Soheil Kolouri

Source-linked AI summary

Sinkhorn-Drifting Generative Models

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

arXiv:2603.12366v2Mar 12, 2026cs.LG

TL;DR

Prior drifting formulations lack a general guarantee that zero drift identifies the target distribution. This paper connects drifting to Sinkhorn-divergence gradient flows through shared cross-minus-self couplings, showing that two-sided scaling resolves the identifiability gap and improves temperature stability at extra training cost.

Problem
Prior drifting formulations do not generally establish that zero drift implies equality between generated and target distributions.
Method
The paper interprets drifting as an approximate Sinkhorn-divergence gradient flow, replacing one-sided normalized Gibbs kernels with entropic OT couplings from two-sided Sinkhorn scaling.
Results
Zero Sinkhorn drift implies model–target equality through Sinkhorn divergence's definiteness and unique stationary point.
Takeaways & Limitations
Sinkhorn normalization provides temperature-stable, theoretically grounded drift fields without changing inference time, while increasing training overhead.
Takeaways & Limitations
Some empirical-measure identifiability cases remain open for n ≥ 3, and prior drifting identifiability depends on a non-degeneracy condition that can fail.

Abstract

from arXiv · show

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/

1 Introduction

The paper links drifting generative dynamics to Sinkhorn-divergence gradient flows through a shared cross-minus-self structure, while addressing drifting's identifiability gap. Two-sided Sinkhorn scaling provides a theoretically grounded alternative to drifting's one-sided normalization and improves temperature stability without changing inference.

Motivation: Drifting uses an attractive target term and a repulsive self-correction term in a stop-gradient training signal for one-step generation.The construction pulls samples toward the data distribution and away from the current model; inference remains a single forward pass.
Open problem: The prior drift formulation leaves open whether zero drift guarantees equality between generated and target distributions.Its identifiability argument depends on a non-degeneracy assumption that may fail for some kernels and distributions.
Connection: Both drifting and Sinkhorn-divergence flows use a cross-minus-self barycentric structure with row-stochastic coupling matrices.Their difference lies in how the couplings are constructed rather than in the overall algebraic form.
Connection: Drifting approximately row-normalizes a Gibbs kernel, whereas Sinkhorn divergence uses entropic optimal-transport plans from full two-sided Sinkhorn scaling.This places drifting as an approximate Sinkhorn-divergence gradient flow.
Implications: Two-sided normalization moves couplings closer to the doubly stochastic Sinkhorn solution and yields a rigorous zero-drift identifiability guarantee.The Sinkhorn divergence is definite, so zero drift implies matching model and target measures; training incurs extra normalization cost while inference time stays unchanged.

2 Background and Notations

The background defines drifting as a one-step particle-based velocity field built from normalized kernel barycenters, then introduces entropic OT, Sinkhorn scaling, and Wasserstein gradient-flow notation. These components provide the formal machinery for comparing one-sided drift updates with Sinkhorn-divergence dynamics.

Drift Method: The drift velocity field is defined for a fixed target distribution p and current model distribution q.Its optimization uses gradient descent with a stop-gradient operator, and the parameter update induces the drift flow.
Drift Method: Each drift component is a normalized kernel-weighted average of displacement vectors.The kernel and its normalizers determine the corresponding barycentric updates.
Drift Method: The target component attracts samples toward p, while the model component subtracts an analogous attraction toward q as self-correction.Together these terms form the attractive–repulsive drift structure.
Drift Method: The model is one-step: a generator maps prior noise to samples, while its drift flow is approximated by forward Euler discretization during training.The training loss is defined from this discretized drift process.
Entropic OT: The Sinkhorn divergence is built from entropic OT couplings constrained by prescribed marginals and KL divergence relative to the outer-product measure.For discrete measures, the couplings satisfy row and column marginal constraints.
Sinkhorn Scaling: The Sinkhorn–Knopp algorithm alternates row and column rescalings of a Gibbs kernel to enforce the coupling constraints.Its convergence typically depends polylogarithmically on problem size and inversely on regularization strength.
Wasserstein Gradient Flows: Wasserstein gradient flows describe steepest-descent evolution of probability measures and induce particle dynamics through first variations and velocity vectors.Forward Euler discretization converts the continuous dynamics into particle updates.

3 Wasserstein Gradient Flow of Sinkhorn-Divergence

The paper formulates Sinkhorn drifting as a particle approximation to the Wasserstein gradient flow of the Sinkhorn divergence, using Sinkhorn-scaled couplings in a cross-minus-self field. Under stated regularity or finite-sample conditions, the analysis establishes several forms of zero-drift identity, while retaining open cases and assumptions.

Sinkhorn gradient flow: Sinkhorn couplings are obtained by iterating two-sided scaling with prescribed marginals, then row-normalized to form barycentric weights for the drift.The implementation constructs positive and negative Gibbs kernels, computes separate Sinkhorn couplings, and applies row normalization before the cross-minus-self update.
Sinkhorn gradient flow: The resulting field is a cross-minus-self update: a positively weighted average toward target samples minus a negatively weighted average toward self or model samples.The algorithm returns V = P_pos @ Y_pos - P_neg @ Y_neg; self-distance masking is optional for Sinkhorn drifting.
Connection to drifting: When l = 1, the Euler discretization coincides with the Drift dynamic proposed by Deng et al.Thus the construction connects one-step Drift dynamics to successive Sinkhorn iterations, with T = 1 recovering the prior method in the implementation.
Identity results: For smooth, strictly positive densities on a connected compact domain, zero Sinkhorn drift implies equality of the target and model measures.The argument uses spatial constancy of the first variation, feasible zero-mass perturbations, and strict convexity and definiteness of the Sinkhorn divergence.
Identity results: With sufficiently large batches and sufficiently small empirical drift under the stated connected-domain setting, the statistical result gives high confidence that the model converges to the target distribution.The identity statement relies on sampling from the continuous distributions rather than only their empirical approximations.
Identity results: For empirical measures, identifiability is proved at τ = 0 under stationarity and non-degeneracy, and for 0 < τ < ∞ when n = 2.For n ≥ 3, zero drift is shown to imply stationarity on the empirical-measure submanifold, while the implication to equality remains open.

4 Computational tradeoff

Sinkhorn drifting increases training cost through iterative scaling but offers more stable temperature behavior and stronger theoretical structure without changing inference cost.

Computational tradeoff: Sinkhorn scaling makes each training step more expensive than one-sided or partial two-sided normalization because it requires iterative updates.The additional computation is incurred during training rather than generation.
Computational tradeoff: Inference retains the same computational cost as drifting approaches because generation remains a single learned-network evaluation.The paper therefore confines the computational tradeoff to optimization.
Computational tradeoff: The extra training cost is exchanged for improved temperature stability and principled two-sided balancing, reducing reliance on ad hoc drifting heuristics.The cited heuristics include self-distance masking and averaging drift across temperatures.

5 Numerical Studies

The experiments compare one-sided, two-sided, and full Sinkhorn normalization across temperatures on synthetic distributions, MNIST, and FFHQ. Sinkhorn improves stability, mode coverage, and generative quality, especially at low temperatures, while requiring additional training computation.

Toy Experiments: At τ = 0.1 on 8-Gaussians, one-sided normalization collapses to a single mode with W 2 2 ≈7–8, while Sinkhorn covers all 8 modes with W 2 2 < 1.Two-sided normalization partially improves over one-sided but remains unstable.
Toy Experiments: At small τ, one-sided and two-sided normalization increasingly suffer mode collapse, whereas Sinkhorn remains comparatively stable across temperatures.The synthetic experiments use one-sided, geometric-mean two-sided, and full Sinkhorn normalization under Gaussian kernels.
Toy Experiments: On Checkerboard, all three methods achieve similar W 2 2, indicating Sinkhorn’s advantage is most pronounced for multimodal targets with isolated modes.The synthetic study reports final generated samples and W 2 2 convergence curves.
MNIST Experiments: On MNIST, Sinkhorn keeps EMD between 6.88–8.57 and class accuracy ≥99.97% across τ ∈[0.005, 0.1], while the baseline collapses for every τ ≤0.05.At small temperatures, baseline accuracy falls to ≈10%; at τ = 0.1 it recovers to EMD = 5.63 and Acc = 95.0%.
Image Generation Experiments: At τ=0.1 on FFHQ-ALAE, mean FID drops from 187.7 to 37.1 and mean EMD drops from 453.3 to 144.4 with Sinkhorn drifting.Sinkhorn consistently outperforms the baseline across all temperatures and six demographic classes in both metrics.

Conclusion

The paper connects drifting dynamics to Sinkhorn-divergence Wasserstein gradient flows through a shared cross-minus-self barycentric structure. This formulation strengthens identifiability and yields improved low-temperature stability and mode coverage without changing one-step inference.

Conclusion: Drifting and Sinkhorn gradient flows share a cross-minus-self barycentric structure, differing in one-sided versus doubly-stochastic coupling construction.Drifting uses one-sided kernel normalization, while Sinkhorn uses entropic optimal-transport couplings enforcing both marginals.
Conclusion: Sinkhorn drifting can be interpreted as a single-iteration approximation of the Sinkhorn gradient flow and explains the benefits of partial two-sided normalization.The connection places drifting within a variational framework.
Conclusion: Vanishing Sinkhorn drift implies equality of the model and target distributions, resolving the identifiability gap identified for prior drifting constructions.This follows from the definiteness of the Sinkhorn divergence.
Conclusion: Experiments demonstrate improved stability and mode coverage, particularly at low temperatures, while preserving drifting’s one-step inference procedure.The empirical evaluations cover synthetic distributions, MNIST, and FFHQ.
Conclusion: The experiments are smaller in scale than those of Deng et al., largely because ImageNet-1K generation requires heavy computation.The authors nevertheless report consistent improvements in smaller-scale experiments.

A.2 Identifiability of Drift Generative Model

This section formalizes how drifting losses induce particle dynamics and derives their connection to Sinkhorn-divergence Wasserstein gradient flows. It also distinguishes rigorous conclusions from formal derivations and identifies assumptions needed for stationarity arguments.

Identifiability: The prior drift construction is non-identifiable: zero drift on the support of p need not imply p = q.The section provides an explicit counterexample with empirical measures and a Gaussian kernel.
Gradient Descent in Drift Field: A continuous-time stop-gradient drift loss yields ∇x_i L_drift = −V_drift(x_i), inducing the particle-space ODE ẋ_i = V_drift(x_i).The result assumes particles are optimized directly without parametrization restrictions.
Gradient Descent in Drift Field: Under generator parametrization, output velocity depends on the network Jacobian, so parameter updates do not generally realize an explicit Euler step along V_drift.The deviation arises unless the parametrization and update rule can realize the drift in output space.
Sinkhorn Gradient Flow: Sinkhorn couplings are obtained by alternating marginal normalizations, with odd and even iterations enforcing the two marginal constraints.Under mild integrability conditions, the iterates converge to the unique entropic optimal-transport minimizer.
Sinkhorn Gradient Flow: The Sinkhorn-divergence gradient flow has a cross-minus-self form derived from converged entropic optimal-transport couplings for the cross and self terms.The derivation uses first variations and envelope or shape-derivative arguments.

E.2 Proof of Proposition 3.5

The proof establishes cancellation when empirical model and target measures coincide and studies when zero Sinkhorn drift yields stationarity on the empirical-measure manifold. The strongest identity result relies on non-degeneracy and convexity assumptions.

Zero Drift at Equality: When empirical measures coincide up to relabeling, Sinkhorn scaling is permutation-equivariant and the cross and self barycentric projections cancel exactly.Thus the truncated Sinkhorn drift is zero at every particle.
Zero Drift at Equality: The general measure result likewise shows that p = q makes the cross and self integrals identical, yielding zero level-l Sinkhorn drift almost everywhere.The argument applies to finite Sinkhorn truncation levels under the stated coupling construction.
Stationarity and Identifiability: Sinkhorn divergence is nonnegative, definite, strictly convex in q for fixed p, and has q = p as its unique minimizer.Differentiability converts stationarity into equality of the measures.
Stationarity and Identifiability: Zero particle-coordinate gradient implies stationarity of Sinkhorn divergence restricted to the empirical manifold, but restricted convexity is left for future work.This limits the available converse identifiability conclusion in the empirical setting.
Empirical Measure Manifold: For empirical measures with distinct particles, the configuration-to-measure differential is an isomorphism onto the tangent space.Every tangent direction can therefore be represented by a unique particle displacement.

F.5 Identity for n = 2.

For n = 2, the entropic Sinkhorn construction forces the two empirical measures to coincide whenever the limiting drift vanishes, subject to distinct-particle assumptions. The proof reduces the scaling factors and analyzes aligned and anti-aligned configurations, with a degenerate repeated-particle case excluded.

Sinkhorn construction: Symmetric two-point costs force equal scaling factors within each scaling vector.For symmetric costs with equal diagonal and off-diagonal entries, a1 = a2 and b1 = b2.
Sinkhorn construction: The entropic transport plan is uniquely represented by diagonal scaling of the kernel, with both marginals enforced.The construction uses π = diag(a) K diag(b), where the transport costs determine K and the scaling vectors are positive.
Case analysis: The drift equations imply that the two support directions are parallel or anti-parallel.The nonzero drift relation is parallel to the first support direction, yielding b̂ = ±â.
Case analysis: In the aligned case, strict monotonicity forces equal radii and pointwise equality between the two configurations.When b̂ = â, f(s) = f(r) implies s = r, so y1 = x1 and y2 = x2.
Case analysis: Combining both directional cases shows that X is a permutation of Y.The anti-aligned case is handled separately, and both cases yield equality up to permutation.
Scope: The distinctness assumption is essential because coincident source particles make the self-plan non-unique and permit counterexamples.If x1 = x2, one can have pX ≠ pY while satisfying the zero-drift condition.

G.2 One-sample complexity

The one-sample analysis bounds the empirical entropic transport map and drift under regularity assumptions using variational reductions, concentration, covering arguments, and entropy bounds. The resulting error decreases with sample size but worsens as the regularization temperature becomes small.

One-sample bound: Under Assumptions (A1)–(A4), the entropic map admits a one-sample bound for sufficiently small τ.The proposition applies when τ ≤ τ0 for a sufficiently small constant τ0.
Proof strategy: The proof decomposes the estimation error into terms controlled by entropic duality, concentration, empirical-process bounds, and covering-number arguments.The analysis proceeds through variational reduction, fixed-x empirical processes, and Dudley chaining.
Rate: The key empirical-process contribution scales as (τ^-1 + τ^-d/2) log(n) n^-1/2.This bound is obtained after combining the component estimates.
Proof strategy: The analysis controls conditional entropic couplings through strong concavity and pointwise bounds on their densities.For sufficiently small τ, the auxiliary function Hx is strongly concave, enabling normalization and concentration estimates.
Statistical implication: At sufficiently large batch size and sufficiently small empirical drift, the population drift is small with high confidence.The stated conclusion is statistical and depends on the supported compact-domain and regularity setting.
Gradient-flow interpretation: Stop-gradient regression with q-weighted gradient descent realizes a forward Euler step for the particle gradient-flow dynamics.The same argument extends to parametric generators, up to Jacobian-dependent preconditioning.

I Additional Drift Trajectory Results for Varying τ

Additional trajectory and low-dimensional experiments show that Sinkhorn normalization is more stable as kernel temperature decreases and generally improves convergence, coverage, or class preservation over classical normalizations.

Drift trajectories: Sinkhorn trajectories remain most stable as τ decreases, while one-sided and two-sided normalization become increasingly temperature-sensitive.Across the temperature sweep, Sinkhorn convergence is at least as good as and often better than classical Drift methods.
2D targets: Across 2-Moons, Spiral, 8-Gaussians, and Checkerboard, Sinkhorn improves mode coverage and convergence, especially for τ ∈ {0.01, 0.05, 0.1}.The experiments use Gaussian and Laplacian kernels across multiple normalizations.
Temperature sweep: With the Laplacian kernel, the baseline completely collapses for τ ≤ 0.02, whereas Sinkhorn remains stable throughout the tested range.The baseline recovers from τ = 0.03 onward but remains unstable at lower temperatures.
2D targets: Sinkhorn gives better coverage and lower W 2 over 5,000 iterations in the reported Laplacian-kernel comparisons.The figure reports final samples alongside W 2 trajectories.
MNIST: At τ = 0.01 on MNIST, the baseline collapses while Sinkhorn generates all classes correctly.At τ = 0.05, Sinkhorn remains stable while the baseline shows partial recovery with visible class confusion.

L Additional ALAE Qualitative Results

The additional FFHQ qualitative panel compares baseline and Sinkhorn class-conditional generations at τ = 0.1, with one demographic class represented per row.

Qualitative comparison: At τ = 0.1, each row corresponds to one class, with Baseline on the left and Sinkhorn on the right.The panel provides a low-temperature qualitative comparison for class-conditional FFHQ generation.