PaperlayerView exact revision on arXiv

Source-linked AI summary

On the Wasserstein Gradient Flow Interpretation of Drifting Models

Arthur Gretton, Li Kevin Wenliang, Alexandre Galashov, James Thornton, Valentin De Bortoli, Arnaud Doucet

arXiv:2605.05118v2cs.LGcs.AIstat.ML

TL;DR

GMD has practical strengths, but its theoretical mechanics remain relatively unexplored. This note analyzes its drift algorithms as fixed-point procedures related to Wasserstein Gradient Flows, finding that the implemented algorithm only partly resembles a Sinkhorn-divergence flow and that drifting methods have limits when modes are far apart.

  • Problem

    Despite GMD's practical strengths, its underlying theoretical mechanics remain relatively unexplored.

    Source: §1 · p. 1

  • Method

    The note interprets drifting algorithms through Wasserstein Gradient Flows, including fixed-point analyses for KL and Sinkhorn-based procedures and extensions to other flows.

    Source: §1 · p. 2, §2.3.2 · p. 6, §3 · p. 6

  • Results

    The implemented algorithm resembles a Sinkhorn-divergence flow fixed point but its velocity generally is not a Wasserstein gradient, and neither it nor KL-KDE shifts mass between modes far apart relative to kernel width τ.

    Source: §1 · pp. 1–2, §3.3 · p. 10

  • Takeaways & Limitations

    The fixed-point framing also applies to MMD, Sliced-Wasserstein, and GAN critic functions derived from f-divergences.

    Source: §1 · p. 2

    Neither KL-KDE nor the Sinkhorn Proxy shifts mass between modes far apart relative to kernel width τ.

    Source: §3.3 · p. 10

Abstract

from arXiv · show

Recently, Deng et al. (2026) proposed Generative Modeling via Drifting (GMD), a novel framework for generative tasks. This note presents an analysis of GMD through the lens of Wasserstein Gradient Flows (WGF), i.e., the path of steepest descent for a functional in the space of probability measures, equipped with the geometry of optimal transport. Unlike previous WGF-based contributions, GMD can be thought of as directly targeting a fixed point of a specific WGF flow. We demonstrate three main results: first, that one algorithm proposed by Deng et al. (2026) corresponds to finding the limiting point of a WGF on the KL divergence, with Parzen smoothing on the densities. Second, that the algorithm actually implemented by Deng et al. (2026) corresponds to a different procedure, which bears some resemblance to the fixed point of a WGF on the Sinkhorn divergence, but lacks certain desirable properties of the latter. Third, the same same idea can be extended to the limiting point of other WGFs, including the Maximum Mean Discrepancy (MMD), the sliced Wasserstein distance, and GAN critic functions.

1 Introduction

This note examines Generative Modeling via Drifting (GMD) through Wasserstein Gradient Flows, interpreting drift algorithms as seeking flow fixed points rather than simulating the full flow. It analyzes a KL-based variant, the algorithm implemented by Deng et al. (2026), and extensions to other divergences and critic functions.

1 Introduction

  • GMD moves generated samples toward the target distribution during training with a non-parametric, mean-shift-style drift operator, without adversarial training, distillation, or multi-step sampling.The note presents GMD as an empirically successful approach while identifying its theoretical mechanics as relatively unexplored.Source: §1 · p. 1
  • A simpler GMD instantiation approximates a fixed point of a Wasserstein gradient flow on KL divergence using Parzen-smoothed densities.The note distinguishes this interpretation from the algorithm Deng et al. (2026) implement in practice.Source: §1 · pp. 1–2
  • The implemented algorithm resembles drifting toward a fixed point of a Wasserstein gradient flow on the Sinkhorn divergence, but its velocity generally is not a Wasserstein gradient and its convergence can fail for distant, non-overlapping modes.The velocity is zero if and only if p and q agree under stated assumptions; the mode-shifting limitation is relative to kernel width τ.Source: §1 · pp. 1–2
  • The fixed-point perspective extends to MMD, Sliced-Wasserstein, and GAN critic functions derived from f-divergences.The note describes MMD drifting as close in practice to Deng et al. (2026)'s algorithm and as representing the MMD-flow fixed point.Source: §1 · p. 2
  • The note addresses an uncovered interpretation of Deng et al. (2026)'s implemented Algorithm 2, beyond prior work on other drifting/Wasserstein-flow connections.It situates this gap alongside prior work on KL and MMD drifting and on flows using the fully converged Sinkhorn divergence.Source: §1 · p. 2

2 A Wasserstein Gradient Flow Perspective

The section frames drifting models as methods for approaching fixed points of Wasserstein gradient flows, focusing on KL and Parzen-smoothed score estimates. It distinguishes this fixed-point interpretation from the true flow on smoothed densities.

2.1 Generative Modeling via Drifting models

  • A generator maps noise to model samples, while a drift field displaces those samples toward the target during training.The stop-gradient update targets the prescribed transport step without back-propagating through the drift estimator.Source: §2.1 · pp. 2–3

2.2 Wasserstein Gradient Flows

  • Under regularity assumptions, a Wasserstein gradient flow decreases its energy over time and can be expressed as a continuity equation.The density-level KL derivation assumes positive, sufficiently regular densities; singular model distributions require formal interpretation or smoothing.Source: §2.2 · p. 3

2.3 Drifting Models as Fixed Point of a Wasserstein Gradient Flow on the KL

  • The section positions the practical algorithm as a different procedure from the simpler score-difference drift analyzed as a KL-flow fixed-point approximation.The practical algorithm is treated separately from the simpler KL interpretation.Source: §2.3 · p. 4, §2.3.2 · p. 5

2.3.1 A score-difference drift field

  • The score-difference drift is a mean-shift field with attraction from target samples and repulsion from model samples.For the Gaussian kernel, prior work cited here proves that the field vanishes exactly when p and q agree; non-Gaussian kernels may require a different normalization.Source: §2.3.1 · p. 5

2.3.2 Wasserstein Gradient Flow for Kullback–Leibler (KL) divergence

  • The simpler drifting model approximates a KL-flow fixed point using Parzen-smoothed estimates of the target and model scores.It directly targets the limiting point rather than simulating the flow trajectory.Source: §2.3.2 · p. 6
  • The paper distinguishes the Parzen-smoothed score approximation from the actual Wasserstein velocity for KL(qτ||pτ).It states that the field attributed to that smoothed-density flow is incorrect, so the cited dissipation result is not established for particles following it.Source: §2.3.2 · p. 6

3 A Sinkhorn Proxy for Drifting Models

The Sinkhorn Proxy connects drifting to a debiased Sinkhorn gradient flow through a one-shot approximation. It is consistent under stated assumptions, but its velocity is generally not a Wasserstein gradient flow and it shares a distant-mode limitation with KL-KDE.

3 A Sinkhorn Proxy for Drifting Models

  • The section connects drifting to the Wasserstein gradient flow on the debiased Sinkhorn divergence and develops the Sinkhorn Proxy as a related drift field.The fully converged Sinkhorn velocity depends on transport plans for both data-to-model and model-to-model terms.Source: §3 · p. 7

3.1 The empirical Sinkhorn proxy

  • The Sinkhorn Proxy replaces converged Sinkhorn computation with a one-shot approximation, normalizing positive and negative affinities separately rather than jointly as in Deng et al.’s Algorithm 2.The separate normalization more closely matches the marginal structures of the two transport problems.Source: §3.1 · pp. 7–8

3.2 Analysis of the Sinkhorn Proxy

  • The Proxy resembles the KL score-difference drift modulated by a spatial preconditioner, and its Gaussian-kernel velocity is zero everywhere if and only if p = q under finite-mass assumptions.The field is generally not a valid Wasserstein gradient flow; that requires the preconditioner and score-difference gradients to be linearly dependent everywhere.Source: §3.2 · pp. 9–10

3.3 Illustration: Sinkhorn Proxy flow vs KL flow and Wasserstein-2 flow

  • When modes are far apart relative to kernel width τ, neither KL-KDE nor the Sinkhorn Proxy shifts mass between modes, unlike the expected behavior of a mass-transport gradient flow.For the Wasserstein-2 flow in the illustration, velocity increases with mode separation D.Source: §3.3 · pp. 10–11

4 Extensions

The extensions describe a generic energy-based drifting procedure and connect it to several Wasserstein gradient flows. Proof-of-concept experiments on synthetic data show that each tested flow can reach a similar best performance at an appropriate hyperparameter value.

4 Extensions

  • Earlier WGF generative methods implement time-discretizations of the flow ODE, whereas drifting offers an approach described as readily implementable for any Wasserstein gradient flow.Source: §4 · p. 11

4.1 Practical algorithm

  • The general method develops a one-stage generator to sample from p using an energy function Fp(q) with Fp(p) = 0; its detached drifted-target gradient step at the current iterate equals an energy-functional step up to 2η.The paper sets η = 1 in Algorithm 2.Source: §4.1 · pp. 11–12

4.2 Applications

  • The MMD drifting approach of Deng et al. is recovered as a Wasserstein gradient flow on MMD², and a characteristic kernel makes MMD zero if and only if p = q.The paper also notes that adaptive kernel choice can improve MMD-flow convergence behavior.Source: §4.2 · p. 13
  • The extensions include a Sliced-Wasserstein flow and GAN-critic flows based on dual forms of f-divergences, including a regularized dual-KL estimator.For the sliced flow, projected distributions and one-dimensional transport maps contribute to the velocity approximation.Source: §4.2 · pp. 13–14

4.3 Illustration

  • Across five synthetic 2D datasets, all tested flows attain similarly best performance at a suitable hyperparameter value, while the Sinkhorn Proxy tolerates smaller τ better than DA2.The study trained six flow objectives and evaluated sample quality using MMD² with a median-heuristic Gaussian kernel.Source: §4.3 · p. 14

A Wasserstein Gradient flow of KL on Parzen Window density estimates

  • The appendix derives the first variation and Wasserstein velocity for forward KL between Parzen-smoothed densities qτ and pτ.The smoothed densities are formed by convolution with a symmetric kernel.Source: A Wasserstein Gradient flow of KL on… · p. 18

B Details of the Sinkhorn Proxy

B.1 Drifting Models as Fixed Point of a Wasserstein Gradient Flow on an Approximate Sinkhorn Divergence

The section relates the implemented drifting model to a proxy for the Wasserstein gradient flow on the debiased Sinkhorn divergence, distinguishing it from the exact flow.

B.1 Drifting Models as Fixed Point of a Wasserstein Gradient Flow on an Approximate Sinkhorn Divergence

  • The implemented drifting algorithm is a fixed point of a probability flow approximating the Sinkhorn optimal-transport distance, but its velocity field is not generally the gradient of a scalar distributional loss.The field is non-conservative for generic pairs (p, q) in dimension at least two.Source: §B.1 · p. 19
  • The debiased Sinkhorn divergence is built from entropy-regularized optimal transport, with τ > 0 controlling the entropic regularization.The section considers the squared Euclidean transport cost C(x, y) = 1/2∥x − y∥².Source: §B.1 · p. 19
  • The exact Sinkhorn-flow particle velocity uses row-stochastic conditional optimal-transport plans for the cross and self terms.Its derivation uses the difference between optimal dual potentials for those terms.Source: §B.1 · p. 19

B.2 The One-Shot Sinkhorn Proxy

The one-shot Sinkhorn Proxy approximates the converged Sinkhorn solution used by the exact flow, then uses cross-weighting to form a drift; its population form includes a spatially varying preconditioner.

B.2 The One-Shot Sinkhorn Proxy

  • The one-shot proxy replaces the fully converged Sinkhorn solution with an approximation to closely recover Deng et al.’s Algorithm 2.It approximates the solution using the geometric mean of row-wise and column-wise softmax normalizations; its unscaled weights need not have unit row or column sums, so the envelope-theorem gradient derivation no longer applies.Source: §B.2 · pp. 20–21
  • Cross-weighting yields a drift vector corresponding to the probability flow of the approximate Sinkhorn solution.The derivation uses a pseudo-plan rather than a fully converged Sinkhorn solution.Source: §B.2 · pp. 21–22

B.3 Alternative form, Sinkhorn proxy velocity field

  • The population-limit proposition expresses the proxy velocity for Gaussian-kernel modified measures and identifies Z(x) as a spatially varying preconditioner.The stated modified measures use q_τ(x) in their density form.Source: §B.3 · p. 22

B.4 Consistency of the Sinkhorn Proxy

With a Gaussian kernel and finite-mass weighted measures, the Sinkhorn Proxy velocity vanishes everywhere exactly when the model and target distributions agree.

B.4 Consistency of the Sinkhorn Proxy

  • For finite-mass weighted measures and a Gaussian kernel, V = 0 everywhere if and only if p = q.The proof uses equality of kernel-smoothed modified measures and injectivity of convolution with the characteristic Gaussian kernel.Source: §B.4 · pp. 23–24
  • When p = q, the modified measures and their kernel convolutions coincide, making the proxy velocity identically zero.The preconditioner is strictly positive for the Gaussian kernel.Source: §B.4 · p. 23
  • If the velocity vanishes everywhere, equality of the smoothed score gradients implies the smoothed measures are proportional; normalization of the probability densities then supports equality of the underlying distributions.The argument relies on the modified measures having finite mass so the Gaussian-kernel embedding is injective.Source: §B.4 · pp. 23–24

B.5 The Sinkhorn Proxy Velocity does not generally correspond to a Wasserstein Gradient Flow

The Sinkhorn Proxy is a valid Wasserstein gradient flow only under a restrictive alignment condition between its spatial preconditioner and score difference.

B.5 The Sinkhorn Proxy Velocity does not generally correspond to a Wasserstein Gradient Flow

  • For a Gaussian kernel, the proxy is a valid Wasserstein gradient flow only if the preconditioner and score-difference gradients are linearly dependent everywhere.Equivalently, their wedge product must vanish: ∇Z ∧ ∇Φ = 0.Source: §B.5 · p. 24
  • A Wasserstein gradient-flow velocity must be a pure gradient field, so the proxy’s velocity must admit a scalar potential.The condition follows from requiring a symmetric Jacobian for the continuously differentiable field on a simply connected domain.Source: §B.5 · p. 24
  • The required gradient alignment generally fails because the preconditioner and score difference have distinct definitions.The paper states that the alignment condition will not hold in general.Source: §B.5 · p. 24

C Geometric mean as approximation to Sinkhorn solution

The practical update is related to Sinkhorn normalization: it combines row-wise and column-wise softmax operations as a geometric mean rather than alternating them. This approximation can suppress affinity to dense target hubs, but does not guarantee the marginal constraints of a fully converged Sinkhorn solution.

C Geometric mean as approximation to Sinkhorn solution

  • Sinkhorn seeks scaling vectors u and v so P = diag(u)Kdiag(v) satisfies the target marginals, with K given by the strictly positive Gibbs kernel.The proxy assumes uniform marginals to recover its Sinkhorn structure.Source: C Geometric mean as approximation to… · p. 25
  • The update in Deng et al. combines row-wise and column-wise softmax normalizations by taking their geometric mean, rather than iterating between them.The row and column softmaxes correspond to normalizing the Gibbs kernel by its row and column sums.Source: C Geometric mean as approximation to… · p. 25
  • When N+ = N−, repeated application of the procedure converges to the entropy-regularized transport map, as does alternating Sinkhorn normalization.This convergence statement concerns repeated iteration, not the single geometric-mean update.Source: C Geometric mean as approximation to… · p. 26
  • The geometric-mean correction does not guarantee Sinkhorn’s marginal constraints, but its column scaling down-weights affinity to high-density target hubs.A large column sum c_j makes the factor 1/√c_j smaller.Source: C Geometric mean as approximation to… · pp. 25–26

D Failure mode on a synthetic example

In a two-mode setting with β < α, both the KL-KDE flow and Sinkhorn Proxy have drift that becomes negligible as the kernel narrows relative to mode separation. The true Wasserstein-2 flow instead transports mass between the modes, highlighting that Sinkhorn Proxy does not reproduce Earth Mover’s mass-transport dynamics in this setting.

D Failure mode on a synthetic example

  • As τ decreases with fixed mode separation D and β < α, the KL-KDE drift points in the correct direction but is scaled by ϵ = exp(−4D^2/τ).When ϵ falls below machine precision, the drift may be effectively zero.Source: D Failure mode on a synthetic example · pp. 26–27
  • The Sinkhorn Proxy has the same narrow-kernel failure mode, and in this setting behaves more like a KL flow than an Earth Mover’s distance flow.Its resemblance to Sinkhorn does not give it the mass-transport dynamics of the Wasserstein-2 flow.Source: D Failure mode on a synthetic example · pp. 27, 29
  • For β < α, the Wasserstein-2 optimal plan sends α − β mass from the overcrowded +D mode to the starved −D mode.The plan leaves the available β mass at −D in place and keeps 1 − α mass at +D in place.Source: D Failure mode on a synthetic example · p. 29

E Proof of Proposition 8

The proposition connects a generator’s stop-gradient displacement loss to the Wasserstein-gradient-flow objective when the prescribed velocity field is exact. Its proof expresses parameter-induced particle motion as an effective spatial velocity and uses the continuity equation to recover the objective’s parameter gradient.

E Proof of Proposition 8

  • When multiple latent samples map to the same spatial point, the proof defines the effective Eulerian velocity by conditional expectation.This links parameter-induced particle displacements to motion at locations in data space.Source: E Proof of Proposition 8 · p. 31
  • The proof substitutes the continuity equation for the density’s parameter evolution into the first variation, assuming boundary terms vanish.It then rewrites the spatial integral as an expectation over latent noise.Source: E Proof of Proposition 8 · p. 31
  • With the exact velocity field, the stop-gradient displacement loss yields ∇θL(θ) = 2η∇θFp(qθ).The result is stated for qθ represented as the pushforward of latent noise through the generator.Source: E Proof of Proposition 8 · pp. 30–31

F Additional results

The additional-results section presents experiments on toy datasets, with captions identifying results for the 8 Gaussian, Circles, Pinwheel, and Swiss roll datasets.

F Additional results

  • Figure 4 presents results for the 8 Gaussian dataset; empty panels indicate that samples diverged.The passage identifies the dataset and divergence convention, but does not state a comparative outcome.Source: Figure 4 · p. 32
  • Figure 5 presents results for the Circles dataset; empty panels indicate that samples diverged.The passage does not state which method or setting performs best.Source: Figure 5 · p. 33
  • Figures 6 and 7 present results for the Pinwheel and Swiss roll datasets, respectively; empty panels indicate divergence.The passages identify the datasets but provide no comparative outcomes.Source: Figure 6 · p. 34, Figure 7 · p. 35

Related papers

Loading 2605.05118v2…