arXiv CSopen access

On the Wasserstein Gradient Flow Interpretation of Drifting Models

2026 · arxiv_cs

arXiv CS · Papers · License: Open Access · 2026

artificial intelligence, reasoning, knowledge representation

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.05118v1 [cs.LG] 6 May 2026

Google DeepMind

Abstract 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 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.

Introduction

Generative Modeling via Drifting (GMD), introduced by Deng et al. (2026), represents a novel and empirically successful approach to generative tasks. The method achieves strong performance without adversarial training, distillation, or multi-step sampling, relying instead on a non-parametric, mean-shift style drift operator to iteratively move generated samples towards the target data distribution during training. Despite these practical strengths, the underlying theoretical mechanics of GMD remain relatively unexplored. In this note, we discuss the relation between drifting models and Wasserstein Gradient Flows (WGF). We begin in Section 2.1 with a short review of the Generative Modeling via Drifting (GMD) framework proposed by Deng et al., and then cover the basics of Wasserstein Gradient Flows (WGF) (Ambrosio et al., 2008) in Section 2.2. Next, in Section 2.3, we illustrate a re-interpretation of one of the simpler GMD instantiations as a technique to find an approximation to a fixed point of a WGF for the Kullback–Leibler divergence. The algorithm implemented in practice by Deng et al. (Algorithm 2 in their paper) is more complex, however. In Section 3, we analyse this algorithm (with minor variations), demonstrating a number of results. First, we show that this model bears a resemblance to a drifting model on the Wasserstein gradient flow on the Sinkhorn divergence, and shares a specific property of this flow; thus, we refer to it as the Sinkhorn Proxy. Second, we prove that the velocity field of the Sinkhorn Proxy is zero if and only if p and q agree when the Gaussian kernel is used. In addition to these positive results, we also present two negative results for the Gaussian kernel case: first, that the velocity field does not in fact correspond to the Wasserstein gradient of any distributional loss (unless a highly implausible condition holds); second, that convergence of the Sinkhorn Proxy 1

does not inherit the good convergence behaviour one would expect of gradient flow on a true optimal transport distance, in the event that source and target distributions have distant modes with non-overlapping support. We next explore a broader perspective on drifting in Section 4, beginning with a generalpurpose formulation of a drifting algorithm. We then consider instantiations for specific measures of divergence. Deng et al. (2026, Appendix C.2) discuss drifting on the maximum mean discrepancy (MMD), which is very close in practice to their drifting algorithm, and represents the fixed point of the Wasserstein gradient flow on the MMD (Arbel et al., 2019). We further introduce drifting on the Sliced-Wasserstein flow (Cozzi and Santambrogio, 2025), and on the class of GAN critic functions derived from f-divergences (Nowozin et al., 2016), providing among other results an alternative to the KL-drifting method of Section 2.3. Related work: the interpretation of the drifting algorithm as the stationary point of a Wasserstein gradient flow was established by Cao et al. (2026); Turan and Ovsjanikov (2026); He et al. (2026). Cao et al. (2026) describes a connection of drifting to the fixed point of a Wasserstein gradient flow on the KL between Parzen estimates of the densities, although there is a subtle difference with our presentation - we provide more detail in Section 2.3. They further note the interpretation of MMD drifting as the stationary point of the MMD gradient flow. Turan and Ovsjanikov (2026) and Lai et al. (2026) show that for the Gaussian kernel, the velocity field of the initial drifting instantiation is expressed as a score difference, which is zero only when the distributions agree. He et al. (2026) and Turan and Ovsjanikov (2026, Section 6.3, Appendix I.4) each describe a drifting algorithm built on the Wasserstein gradient flow of the fully converged Sinkhorn divergence. Alternative perspectives on drifting models, complementary to the one presented here, are proposed by Lai et al. (2026) and Li and Zhu (2026). It is notable that the WGF interpretation of the drifting model actually implemented by Deng et al. (2026, Algorithm 2) is not covered in these prior works.

A Wasserstein Gradient Flow Perspective

2.1

Generative Modeling via Drifting models

Let p be the target data distribution on Rd . To approximate p, we use a neural network fθ : Rc → Rd whose input is some noise ϵ ∼ µ (e.g., ϵ ∼ N (0, Ic )) and whose output is x = fθ (ϵ). The implicit distribution of x is denoted qθ . We assume that fθ (ϵ) is differentiable w.r.t. θ for all ϵ ∈ Rc . The drifting-model paradigm evolves θ during training for qθ to match the target p. This method relies on a drift field Vp,qθ : Rd → Rd which we use to displace samples from our current model. We will discuss several instances of drifting fields in subsequent sections. We update parameters θ by minimizing the regression loss h i L(θ) = Eϵ∼µ ∥fθ (ϵ) − sg fθ (ϵ) + Vbp,qθ (fθ (ϵ)) ∥2 ; (1) i.e. the average mean square error between a sample and its displaced version; see (Deng et al., 2026, Equation (6)). Here Vbp,qθ is an approximation to Vp,qθ . The stop-gradient sg(·) operator appearing in (1) is essential: it ensures that optimization moves the generator toward the prescribed transport step, rather than changing the drift estimator itself (Deng et al., 2026). Crucially, it also does not require back-propagating through Vbp,qθ . Additionally, it is expected that this approach is more robust than minimizing Eϵ [∥Vbp,qθ (fθ (ϵ))∥2 ] where the generator can reduce the loss by collapsing mass so attraction and repulsion cancel, spreading mass to make kernel weights flat, or shaping qθ so that Vbp,qθ becomes small without matching p. In Section 4.1, Proposition 8, we show that the gradient of this loss turns out to be correct under the WGF perspective presented next, and the stop gradient operator is crucial. 2

2.2

Wasserstein Gradient Flows

We will sketch here the basics of WGFs. We refer the reader to (Ambrosio et al., 2008; Santambrogio, 2017) for a comprehensive treatment and to Appendix A in (Arbel et al., 2019) for a gentle introduction. Let P2 (Rd ) be the space of distributions with finite second-order moments on Rd . Let Fp : P2 (Rd ) → R be an energy functional minimized at q = p; e.g. the Kullback–Leibler divergence Fp (q) = KL(q∥p) or the Maximum Mean Discrepancy Fp (q) = MMD(p, q). Under standard regularity assumptions, the Wasserstein-2 (W2 ) gradient flow of Fp is characterized by the partial differential equation (PDE) (Ambrosio et al., 2008; Santambrogio, 2017) δFp ∂t qt (x) = ∇ · qt (x)∇ (x) , (2) δqt δF

where δqp : Rd → R is the first variation of the functional Fp with respect to the density q and ∇· is the divergence operator. If it exists, this first variation is defined as the unique function (up to an additive constant) that satisfies the following linear expansion for any perturbation χ in the R R tangent space (i.e. χ(x)dx = 0, ∥x∥2 |χ(x)|dx < ∞), where h ∈ R Z δFp d Fp (q + hχ) = (x) χ(x)dx. dh h=0 δq The PDE (2) implies that Fp (qt ) is non-increasing over time as, under regularity conditions, we have Z δFp d Fp (qt ) = (x) ∂t qt (x) dx dt δqt Z δFp δFp (x) ∇ · qt (x)∇ (x) dx = δqt δqt Z 2 δFp (x) dx ≤ 0, = − qt (x) ∇ δqt where we have used first the chain rule, then (2) and finally the divergence theorem. The PDE (2) can be rewritten as a continuity equation ∂t qt (x) + ∇ · (qt (x)Vp,qt (x)) = 0,

Vp,qt (x) = −∇

δFp (x). δqt

(3)

However, note that this is not a standard continuity equation as the drift not only depends on x but also on the distribution qt of xt at time t. Formally, this continuity equation corresponds to the following McKean–Vlasov process for individual samples, dxt = Vp,qt (xt )dt,

x0 ∼ q0 ;

(4)

i.e., this ODE is such that xt ∼ qt . In order to implement such an ODE, we need to discretize it and approximate its drift. For example, using an Euler method leads to xt+η = xt + η Vbp,qt (xt ),

(5)

where η > 0 is a stepsize, and we use Vbp,qt estimated from finite samples, which depends on the energy functional being minimized. 3

2.3

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

We provide two alternative views of Algorithm 1 in Deng et al. (2026) through the score function of distributions, and then present the Sinkhorn Proxy method that stems from the WGF of the Sinkhorn divergence, which better resembles Deng’s Algorithm 2 (main method). We then link the Sinkhorn Proxy to the score objective and show common failure modes. 2.3.1

A score-difference drift field

For probability densities p, q on Rd , we first consider the following drift field defined by Deng et al. (2026, eq. (8)): Vp,q (x) = Vp (x) − Vq (x) (6) with

Ex′ ∼q kτ (x, x′ ) x′ −x Vq (x) = Ex′ ∼q kτ (x, x′ )

Ey∼p kτ (x, y) y − x, Vp (x) = Ey∼p kτ (x, y)

(7)

where k : Rd × Rd → (0, ∞) is a kernel.1 This defines a mean-shift style drifting field (6) built from an “attraction” term (from positives y ∼ p) and a “repulsion” term (from negatives x′ ∼ q). In particular, for the Gaussian kernel in (Deng et al., 2026, Section 3.3) defined as τ > 0, (8) kτ (x, y) = (πτ )−d/2 exp − τ −1 ∥x − y∥2 , we have for pτ (x) ≜ Ey∼p [kτ (x, y)] and qτ (x) ≜ Ex′ ∼q [kτ (x, x′ )] that ∇pτ (x) 2 2 Ey∼p (y − x)kτ (x, y) ∇ log pτ (x) = = Vp (x) = pτ (x) τ τ Ey∼p kτ (x, y)

(9)

and similarly ∇ log qτ (x) = τ2 Vq (x). Note that (9) is an instance of Tweedie’s formula. Therefore we have for (8) that τ ∇ log pτ (x) − ∇ log qτ (x) ; (10) Vp,q (x) = 2 that is, up to a constant factor, the drifting field is the difference of the scores of noised versions of p and q. This result was shown by Lai et al. (2026, Theorem 1) and Turan and Ovsjanikov (2026, Theorem 4.1), where the latter also prove that for the Gaussian kernel, Vp,q = 0 iff p = q (their Theorem 5.1). Franz et al. (2026) show that when a non-Gaussian kernel is used, the velocity field does not correspond to a Wasserstein gradient, and propose an alternative normalization which restores conservatism. Practically, we are interested in Vp,qθ (x) = Vp (x) − Vqθ (x). Here Vp and Vqθ are approximated +

i.i.d.

− i.i.d.

N by sampling mini-batches {yj }N j=1 ∼ p and {xi }i=1 ∼ qθ . We define normalized weights:

kτ (x, yj ) αj+ (x) = PN + , j=1 kτ (x, yj )

kτ (x, xi ) αi− (x) = PN − . i=1 kτ (x, xi )

Then, from (7), we have the drift estimate −

Vbp,qθ (x) = Vbp (x) − Vbqθ (x) with Vbp (x) =

N X

αj+ (x)(yj − x),

j=1

Vbqθ (x) =

N X

αi− (x)(xi − x). (11)

i=1

We emphasize that this is not the algorithm Deng et al. use in practice, which has a different interpretation, and will be discussed in Section 3.

2.3.2

Wasserstein Gradient Flow for Kullback–Leibler (KL) divergence

Consider the WGF associated to the Kullback–Leibler discrepancy Fp (q) = KL(q||p). In this case, it can be easily checked that δFp (x) = log q(x) − log p(x) + 1 δq

(12)

so the ODE (4) becomes dxt = (∇ log p(xt ) − ∇ log qt (xt ))dt,

x0 ∼ q0

(13)

i.e. the drift of (13) is given by the difference of the (intractable) scores of both the data distribution p and of qt .2 An Euler discretization (i.e. (5)) of (13) with step size η = τ /2 is given by τ xt+ τ2 = xt + (∇ log p(xt ) − ∇ log qt (xt )), 2

x0 ∼ q0 .

(15)

Now, in particular, this shows that for xt ∼ qt = p, the particle displacement over the time interval of size τ2 is given by τ2 (∇ log p(xt ) − ∇ log qt (xt )), which evaluates exactly to zero. The key point of GMD is that it is not trying to simulate the flow (xt )t≥0 . It just directly attempts to find its limiting point by setting the following fixed-point equation. Indeed, for x∞ ∼ q = p, the displacement is zero, and we have from (15) that τ x∞ = x∞ + (∇ log p(x∞ ) − ∇ log q(x∞ )), 2

(16)

where both ∇ log p and ∇ log q are intractable. By making the approximations ∇ log p(x) ≈ ∇ log pτ (x) and ∇ log q(x) ≈ ∇ log qτ (x), we thus get the approximation of fixed-point equation (16) x∞ = x∞ + Vp,q (x∞ ), (17) where Vp,q (x) is defined in (10). Going back to the parametric framework, this suggests identifying θ as solving the fixed-point equation (17) using fθ (ϵ) = fθ (ϵ) + Vp,qθ (fθ (ϵ)), ϵ ∼ µ; see (Deng et al., 2026, Equation (4)). This is achieved by performing a Monte Carlo approximation Vbp,qθ in (11) of Vp,qθ and minimizing (1). So, to summarize, drifting models for the drift field defined by (6) and (7) can be re-interpreted as a method for approximating the fixed point of the WGF minimizing KL(q||p); that is, the data distribution p. We now review the related finding of Cao et al. (2026, Corollary 4.10), which states that the drifting model’s velocity field is the Wasserstein-2 gradient flow of the KL divergence between (τ ) smoothed densities, Fp (q) := KL(qτ ||pτ ). Cao et al. (2026, eq. 7) claim that (10) represents this velocity field. The correct expression for the velocity field of this KL divergence is not (10), however, but Z qτ (x) (τ ) dx. Vp,q (y) = −∇y kτ (x, y) log pτ (x) 2

This ODE is the probability flow version of the following Langevin diffusion √ dxt = ∇ log p(xt )dt + 2dBt , x0 ∼ q0

i.e., the processes defined in (13) and (14) have the same marginal distributions.

(14)

The proof is in Appendix A. More precisely, if the kernel kτ (x, y) = k̄τ (x − y) for some function k̄, then we have the following result using a change of variables Z Z τ (τ ) Vp,q (y) = kτ (x, y)∇ log pτ (x)dx − kτ (x, y)∇ log qτ (x)dx ̸= (∇ log pτ (y) − ∇ log qτ (y)) . 2 The same issue applies for the additional flows on Parzen window f -divergences defined in that work. As a downstream consequence, the dissipation result of Cao et al. (2026, Corollary 4.8 (ii)) (τ ) is not established, since particles following (10) are not following the velocity field for Fp (q).

A Sinkhorn Proxy for Drifting Models

The score-difference drift of Section 2.3 approximates scores via kernel density estimates. The algorithm used in practice by Deng et al. (2026), however, follows a different procedure (Algorithm 2 in their paper). We develop a closely related drift field—which we term the Sinkhorn Proxy—by connecting it to the Wasserstein gradient flow on the debiased Sinkhorn divergence (Ramdas et al., 2017; Feydy et al., 2019): Sτ (q, p) = OTτ (q, p) − 12 OTτ (q, q) − 12 OTτ (p, p), where OTτ denotes entropy-regularized optimal transport as defined by Cuturi (2013), with regularization τ > 0. This section is a summary of a more in-depth exposition in Appendices B-D. PN − P + Proposition 1 (Exact Sinkhorn particle velocity). Let p = N i=1 δxi be empirical j=1 δyj and q = −

−

×N ×N measures and let C(x, y) = 12 ∥x − y∥2 . Let W + ∈ RN and W − ∈ RN be the row+ + stochastic conditional optimal transport plans for OTτ (q, p) and OTτ (q, q), respectively, defined via the Sinkhorn dual potentials. Then the Wasserstein gradient flow velocity of the debiased Sinkhorn divergence Sτ (q, p) at location xi is −

S Vbp,q (xi ) =

N X

Wij+ yj −

j=1

N X

− xk . Wik

(18)

k=1

For reference and to make the paper self-contained, a proof is given in Appendix B.1, using the result from (Feydy et al., 2019). He et al. (2026) and Turan and Ovsjanikov (2026, Section 6.3, Appendix I.4) propose a drifting algorithm built on the fully converged Sinkhorn solution.

3.1

The empirical Sinkhorn proxy

Computing (18) requires running the Sinkhorn algorithm to convergence at each training step. We + replace it with a one-shot approximation. Let zim = −∥xi − ym ∥2 /τ be the scaled affinities between − 2 model samples and data, and zim = −∥xi − xm ∥ /τ be the scaled affinities within model samples. We define the pseudo-plan s s + + − − exp(zim ) exp(zim ) exp(zim ) exp(zim ) + − Aim = PN + · · , A = , P P P − − − im N N N + + − − u=1 exp(ziu ) v=1 exp(zvm ) u=1 exp(ziu ) v=1 exp(zvm ) i.e., the geometric mean of the row-wise and column-wise softmax normalizations. This corresponds to composing a simultaneous Sinkhorn half-iteration from each side (Appendix C) - repeated iteration of this procedure converges to the exact transport map (Knight et al., 2014, Theorem 3.1). 6

Cross-weighting. Unlike the fully converged solution, both A− and A+ no longer have row P unit + + or column sums. To restore unit row sums, we apply a cross-weighting step. Let si = j Aij and P − s− k Aik be the row sums, and set i = f + = A+ · s− , W ij ij i

f − = A− · s+ . W i ik ik

− Both rows now share the same total mass Zi = s+ i si , which will cancel out the explicit dependence + + + − − − on xi . Writing π̃ij = Aij /si and π̃ik = Aik /si , the empirical Sinkhorn proxy drift at xi becomes   N+ N− X X + − SP Vbp,q (xi ) = Zi  π̃ij yj − π̃ik xk  . (19) j=1

k=1

which shares with the true Sinkhorn flow (18) the property that the x cancels across both terms. This drift for a parametric qθ is then used to optimize the θ, similar to the method in section 2.3.2 for the KL drift in section 2.3.1. The complete procedure for computing the drift is summarized in Algorithm 1.3 As indicated in the algorithm, the Sinkhorn Proxy differs from the practical algorithm of Deng et al. (2026) in the row-normalization scope: we normalize A+ and A− over N + and N − columns respectively, whereas Deng et al. first concatenate positive and negative logits and normalize jointly over N + + N − . Our formulation yields a closer approximation to the true Sinkhorn solution (Appendix B.2). Algorithm 1 Computing the Sinkhorn Proxy drift field Ve −

N Require: Model samples {xi }N i=1 , target samples {yj }j=1 , temperature τ . + − ▷ compute affinities 1: zij = −∥xi − yj ∥2 /τ , zik = −∥xi − xk ∥2 /τ 2: Iterations are executed for j = 1, . . . , N + , k = 1, . . . , N − and m = 1, . . . , N + + N − ranges. Sinkhorn Proxy (ours) Deng et al., Algorithm 2 (DA2) − + − 2: z = [z + , z − ] ∈ RN ×(N +N )

2: for (⋄, ℓ) ∈ {(+, j), (−, k)} do 3:

row,⋄ = αiℓ

col,⋄ αiℓ =

⋄) exp(ziℓ PN ⋄ ⋄ u=1 exp(ziu ) ⋄) exp(ziℓ PN − ⋄ v=1 exp(zvℓ )

exp(z

im row = 3: αim PN + +N −

exp(ziu ) u=1 exp(zim ) col 4: αim = PN − v=1 exp(zvm )

q row ·αcol ; αim im 6: [A+, A− ] = Aim

q col,⋄ row,⋄ · αiℓ A⋄iℓ = αiℓ 6: end for PN + + PN − − + 7: si = s− j=1 Aij , k=1 Aik i = + + − − f =A ·s , W f = A− · s+ 8: W ij ij i i ik ik PN + f + PN − f − e 9: return V (xi ) = j=1 Wij yj − k=1 Wik xk

5: Aim =

3.2

▷ row concat

)

▷ row split ▷ cross-weighting

Analysis of the Sinkhorn Proxy

We now introduce the population counterpart to the velocity field induced by the one-shot approximation. We show two main results: first, that the velocity is zero if and only if p = q under reasonable assumptions. Second, that the velocity field of the Sinkhorn proxy does not in general correspond to a Wasserstein gradient flow. 3

We do not address here the ignore self step of Deng et al. (2026, Algorithm 2), where self-interactions are removed, which is important in the implementation.

Let k(x, y) = exp(−∥x − y∥2 /τ ) be the unnormalized Gibbs kernel.4 In the continuous population limit, recall5 pτ (y) = (k ∗ p)(y), qτ (x) = (k ∗ q)(x). The geometric mean proxy affinities can be written symmetrically as: k(x, y) , A+ (x, y) = p qτ (y)pτ (x)

A− (x, x′ ) = p

The row sums of the independent proxies are Z + s (x) = A+ (x, y)dp(y),

−

s (x) =

k(x, x′ ) qτ (x)qτ (x′ )

A− (x, x′ )dq(x′ ).

The cross-weighting step explicitly defines the final update weights as W + (x, y) = s− (x)A+ (x, y),

W − (x, x′ ) = s+ (x)A− (x, x′ )

The population drift vector for the Sinkhorn Proxy is Z Z SP + Vp,q (x) = W (x, y)y dp(y) − W − (x, x′ )x′ dq(x′ ).

(20)

The velocity field (20) may be written in the following alternative form, as proved in Appendix B.3. p Proposition 2 (Population form of the Sinkhorn Proxy). Let p̃(y) = p(y)/ qτ (y), and q̃(x) = p q(x)/ qτ (x). Then the Sinkhorn Proxy velocity field (19) with Gaussian kernel has the populationlimit form τ SP Vp,q (x) = Z(x) ∇ log(k ∗ p̃)(x) − ∇ log(k ∗ q̃)(x) , (21) 2 where Z(x) = s+ (x) s− (x) is a spatially varying pre-conditioner. This reveals a structural resemblance to the KL score-difference drift of Section 2.3, modulated by Z(x). The velocity field is consistent, as shown in Appendix B.4. Theorem 3 (Consistency). V = 0 if and only if p = q for the Gaussian kernel k. SP does not correspond to a valid Wasserstein gradient As we show in Appendix B.5, however, Vp,q flow for the Gaussian kernel, similar to the finding of Franz et al. (2026) for the KL-KDE drift with non-Gaussian kernels (Section 2.3.1). SP in (21), with Gaussian kernel, Proposition 4 (Non-gradient structure). The velocity field Vp,q corresponds to a valid Wasserstein gradient flow only when the gradients of the pre-conditioner and the score difference are linearly dependent everywhere:

∂i Z(x)∂j Φ(x) = ∂j Z(x)∂i Φ(x)

for all i, j,

where Φ(x) = τ2 (log(k ∗ p̃)(x) − log(k ∗ q̃)(x)). Equivalently, the wedge product must vanish: ∇Z ∧ ∇Φ = 0. This condition will not hold in general. 4 Note that this is now a positive definite kernel, not a Parzen window, so we do not employ the normalizer in (8). We henceforth write pτ and qτ to denote convolutions with the unnormalized kernel, with some abuse of notation. 5 InR the original algorithm, the logits for the positive and negative targets are concatenated, which would give a term k(x, x′ )d(p + q)(x′ ) = (k ∗ (p + q))(x) in place of pτ (x). A closer approximation to the Sinkhorn is attained in the equations we use below.

3.3

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

To illustrate a key limitation of both KL-KDE and the Sinkhorn Proxy, we demonstrate that neither is able to shift mass from one mode to another when the modes are far apart relative to kernel width τ . This is less surprising for score-based drifts like KL-KDE, but more so for the Sinkhorn Proxy, and indicates that the latter has not recovered the behaviour one would expect for gradient flow on a mass transport metric. We consider a synthetic example where the target distribution p and the model distribution q are mixtures of two identical point masses at +D and −D, differing only in their mixing proportions α and β, with α, β ∈ (0, 1) and α ̸= β: p(y) = αδ−D (y) + (1 − α)δ+D (y), q(x) = βδ−D (x) + (1 − β)δ+D (x). We assume β < α, meaning the model has insufficient mass at the left mode (−D) and excess mass at the right mode (+D). We evaluate the continuous vector fields for a model particle currently residing within the overcrowded cluster at x = +D. We use the unnormalized Gaussian kernel k(x, y) = exp(−∥x − y∥2 /τ ). We require the kernel bandwidth τ to be sufficiently narrow compared to the mode separation 2D, such that 4D2 ϵ = k(D, −D) = exp − ≪ 1. (22) τ For the algorithm to successfully correct the mixing proportions, it must yield a strict negative drift at x = +D, routing excess mass across the manifold to the starved mode at −D. The proofs of all the results in this Section are given in Appendix D. We first describe the failure of KL-KDE. Proposition 5. When D and τ jointly satisfy (22) and a Gaussian kernel is used, the velocity field VKL of the KL-KDE flow at x = +D is given by α β KL Vp,q (D) = −2Dϵ − + o(ϵ). 1−α 1−β Unsurprisingly, the velocity vanishes for small ϵ. We now show that the Sinkhorn proxy method has the same failure mode. Proposition 6. When D and τ jointly satisfy (22) and a Gaussian kernel is used, the velocity field of the Sinkhorn Proxy flow at x = D is given by s " # 1 − β α β SP Vp,q (D) = −2Dϵ Z(D) − + o(ϵ), β 1−α 1−β | {z } (*)

where the term (*) represents the difference with the KL-KDE flow. This confirms that in the present setting, the Sinkhorn Proxy does not share the mass transport dynamics of an optimal transport Distance, as demonstrated below. Proposition 7. The velocity field VW2 of the Wasserstein gradient flow on the Wasserstein-2 distance at x = D is given by α−β W2 Vp,q (D) = −2D . 1−β 9

Algorithm 2 Generic drifted-target method to minimize an energy function Fp (q) Require: data sampler y ∼ p, noise sampler ϵ ∼ µ, generator fθ , initial parameter θ 1: for ℓ = 0, 1, 2, . . . do + i.i.d. 2: sample {yj }N ▷ Real batch j=1 ∼ p − i.i.d.

− sample {ϵi }N i=1 ∼ µ, set xi = fθ (ϵi ) for all i ∈ {1, . . . , N } + − N 4: compute drifts Vbp,qθ (xi ) = Vbp,qθ (xi ; {yj }N j=1 , {xk }k=1 ) PN − 5: update θ by minimizing i=1 ∥fθ (ϵi ) − sg(xi + Vb (xi ))∥2 6: end for

▷ Model batch

Extensions

WGF techniques have been previously proposed to perform generative modeling for various functionals; e.g., Maximum Mean Discrepancy (Arbel et al., 2019), Sobolev witness (Mroueh et al., 2019) and Slice-Wasserstein distance (Liutkus et al., 2019; Cozzi and Santambrogio, 2025). WGF have also been used to solve inverse problems (Crucinio et al., 2024). Contrary to GMD, all these approaches implement time-discretization of the ODE (4). As noted by Cao et al. (2026); Turan and Ovsjanikov (2026), the methodology of Deng et al. (2026) offers an alternative that can be readily implemented for any Wassertein gradient flow. This might prove a useful addition to existing techniques.

4.1

Practical algorithm

Given an energy function Fp (q) such that Fp (p) = 0, the general method to develop a one-stage generator fθ to sample from p is as follows: • identify its first variation δFp /δq(x); • set Vp,q (x) = −∇δFp /δq(x); • approximate Vp,qθ (x) via Vbp,qθ (x) (e.g. expectations are approximated by mini-batches, api.i.d.

proximation of Tweedie’s formula etc.). When empirical samples {yj } ∼ p and xk = fθ (ϵk ), + − i.i.d. {ϵk } ∼ µ are available, we write Vbp,q (x; {yj }N , {xk }N ). j=1

k=1

Finally instead of minimizing Eϵ∼µ [∥Vbp,qθ (fθ (ϵ))∥2 ] w.r.t. θ, minimize h i L(θ) = Eϵ∼µ ∥fθ (ϵ) − sg fθ (ϵ) + η Vbp,qθ (fθ (ϵ)) ∥2 where η > 0 is a hyperparameter. The generic algorithm is detailed in Algorithm 2. The following Proposition, proven in Appendix E, shows that a gradient step on L is equivalent to a gradient step on the energy functional Fp (qθ ), and η simply scales the magnitude of the gradient, so we set η = 1. Proposition 8. Let qθ = (fθ )# µ be the pushforward of µ through fθ , and suppose we have the δF exact velocity field Vp,qθ (x) = −∇x δqθp (x) for some functional Fp . Then ∇θ L(θ) = 2η ∇θ Fp (qθ ).

4.2

Applications

The approximation of Vp,qθ is very dependent on the energy functional being minimized. We give a few potential applications using different types of flows, in addition to the KL flow in Section 2.3.2. Maximum mean discrepancy (MMD) flow Let k : Rd × Rd → R be a differentiable reproducing kernel (e.g. RBF), and let H be the corresponding reproducing kernel Hilbert space. The Maximum Mean Discrepancy (MMD) (Gretton et al., 2012) is defined as !2 MMD2k (p, q) =

Ey∼p [g(y)] − Ex∼q [g(x)]

sup g∈H,∥g∥≤1

= Ex,x′ ∼q [k(x, x′ )] − 2Ex∼q,y∼p [k(x, y)] + Ey,y′ ∼p [k(y, y ′ )]. The function g ∗ attaining the supremum is the MMD witness function. For a characteristic kernel, MMD2k (p, q) = 0 is equivalent to p = q. The WGF applied to FpMMD (q) = 21 MMD2k (p, q) was first studied by Arbel et al. (2019) who showed that δFpMMD (x) = g ∗ (x) := Ex′ ∼q [k(x, x′ )] − Ey∼p [k(x, y)], δq so MMD Vp,q (x) = −Ex′ ∼q [∇x k(x, x′ )] + Ey∼p [∇x k(x, y)].

This can be approximated by the mini-batch estimator N+

N−

i=1

j=1

1 X 1 X MMD Vbp,q (x) = + ∇x k(x, yi ) − − ∇x k(x, xj ). N N which recovers the MMD drifting approach proposed by Deng et al. (2026, Appendix C.2), as demonstrated by Cao et al. (2026, Section 4.5). The meaning of “kernel” is fundamentally different to the KL case: here, it is a reproducing kernel (a positive definite function, which need not be a Parzen window), informing the smoothness of the witness function (first variation). When a characteristic kernel (which includes Gauss and Laplace kernels) is used, the MMD is zero if and only if p = q (Sriperumbudur et al., 2010). As a practical matter, Wasserstein gradient flow on the MMD has much better convergence behaviour when the kernel is chosen adaptively, according to the current location of the q particles relative to those from p (Galashov et al., 2025; Chen et al., 2025). As an alternative to adapting the kernel, a sum of kernels over different bandwidths can be used, as was done by Deng et al. (see Cortes et al., 2009, Section 5, for the multiple kernel learning perspective). A related perspective is the use of the MMD as a distributional loss (critic) in GAN training Li et al. (2015); Dziugaite et al. (2015), and the related diffusion model of Zhou et al. (2025). The present discussion makes explicit the connection with Wasserstein gradient flows. Sliced-Wasserstein (SW) flow. Let Sd−1 be the unit hypersphere in Rd , and for any direction ϑ ∈ Sd−1 , let πϑ (x) = x · ϑ denote the 1D projection of a point x onto ϑ. We define the squared Sliced-Wasserstein distance as Z 2 SW2 (q, p) = W22 (πϑ )# q, (πϑ )# p dϑ, Sd−1

where (πϑ )#q denotes the push-forward (i.e., the 1D marginal) of the distribution q along the direction ϑ, and dϑ represents the uniform probability measure on the sphere. We set the energy function to FpSW (q) ≜ 12 SW22 (q, p). Cozzi and Santambrogio (2025) showed that δFpSW (x) = δq

Z φϑ (ϑ · x) dϑ, Sd−1

SW Vp,q (x) = −∇x

δFpSW (x) = δq

Tϑ (ϑ · x) − ϑ · x ϑ dϑ,

Sd−1

where φϑ is the 1D Kantorovich potential and Tϑ : R → R is the 1D optimal transport map pushing (πϑ )# q to (πϑ )# p. This can be approximated using a mini-batch of real data samples and model SW (x) then approximated by uniformly sampled directions (ϑi )L d−1 . samples, and Vp,q i=1 on S Dual Forms of f -Divergences. The f -divergences (notably Jensen-Shannon) are the standard divergence measures used in GANs. As shown by Nowozin et al. (2016), critics based on f -divergences can be formulated using the Fenchel dual of their corresponding divergences. This dual formulation introduces an optimal witness function g ∗ whose gradient is the Wasserstein gradient of the f -divergence. Building upon this, Glaser et al. (2021) proposed the KL Approximate Lower-bound Estimator (KALE) as an alternative to the KL gradient flow discussed in Section 2.3. KALE operates by regularizing the dual KL divergence over a designated function class H. For a regularization parameter λ > 0, the estimator of Glaser et al. (2021, Proposition 2) is: λ 2 KALE g(x) (23) Fp (q) = (1 + λ) max Ey∼q [g(y)] − Ex∼p [e ] + 1 − ∥g∥H . g∈H 2 This is computed via a quadratic program when H is a reproducing kernel Hilbert space Glaser et al. (2021), and by SGD on the weights ϕ of gϕ , when H is a neural network (Nowozin et al., 2016; Galashov et al., 2025). δF As for the MMD, ∇ δqp (x) = (1 + λ)∇g ∗ , where g ∗ ∈ H achieves the supremum in (23). A neural network implementation of the KL gradient flow, which parameterizes the critic function with a deep architecture, is given by (Galashov et al., 2025, Section 6 and Appendix D). A potential advantage of the Fenchel dual instantiation of the Wasserstein Gradient on the KL, over the flow described in Section 2.3, is that the former does not require kernel density estimates in its definition. As an alternative to the KL, a gradient flow on the dual approximation to the χ2 divergence is given by Chen et al. (2025), which has closed-form updates, unlike the KALE.

4.3

Illustration

As a proof of concept, we ran experiments to train generators using 6 flow objectives above on 5 synthetic 2D datasets. We use τ to denote hyperparameters in different flow, and vary its value over a wide range to test the sensitivity and effectiveness of the different approaches. For the KL flow, τ represents the additive Gaussian noise standard deviation; for the MMD flow, τ represents the bandwidth of the Gaussian kernel k(x, y) = exp(−∥x − y∥22 /τ ); for Sinkhorn-based flows, including Deng et al. (2026)’s Algorithm 2 (DA2) and the approximate Sinkhorn algorithm we derived, τ denotes the entropic regularization weight. The plain Sinkhorn flows had 100 Sinkhorn-Knopp iterations. All Sinkhorn-based algorithms operate in log-probability space for numerical stability. for the SW flow, we swept the number of random slices L, but did not find any effect on sample quality. The generator network architecture is a feedforward ResNet, with parameters optimized by the Adam algorithm with learning rate 10−4 , trained for 106 steps. Each data-drift type combination is repeated 20 times with different random seeds for the dataset and parameter initializations.

Moons

Circles

8 Gaussians

MMD²

10−1 10−3 10−5

10−6

Pinwheel

MMD²

10−2

100

102

Swiss roll

10−1

Drift type

10−3 10−5 10−6

10−4

10−2

100

102

10−6

10−4

10−2

100

102

KL MMD Sinkhorn Deng et al. (2026, Alg 2) Sinkhorn Proxy (ours)

Figure 1: MMD between true and generated samples trained by different drift types. We used the MMD2 under a median-heuristic Gaussian kernel as the metric for evaluating sample quality at the end of training. The key results are shown in Figure 1, and we show the true and the generated sample histograms for the Moons dataset in Figure 2. Histograms of other datasets are left in Appendix F. All flows can attain similarly best performances at some sweet spot of τ . For KL, a too small noise level leads to the mode-collapse issue common for score-based methods (Wenliang and Kanagawa, 2020), and a too large noise level causes the learned distribution to ignore local noise structure (Figures 2, 4 and 3). The Sinkhorn drift can tolerate a small τ and only gets worse for large τ . Interestingly, our Sinkhorn Proxy has better tolerance to smaller values of τ compared to DA2.

Acknowledgments The authors are grateful to Guillaume Couairon and Christopher Marouani for their helpful comments.

References Ambrosio, L., Gigli, N., and Savaré, G. (2008). Gradient Flows in Metric Spaces and in the Space of Probability Measures. Birkhäuser. Arbel, M., Korba, A., Salim, A., and Gretton, A. (2019). Maximum mean discrepancy gradient flow. In Advances in Neural Information Processing Systems. Cao, J., Wei, Z., and Liu, Y. (2026). Gradient flow drifting: Generative modeling via Wasserstein gradient flows of KDE-approximated divergences. Chen, Z., Mustafi, A., Glaser, P., Korba, A., Gretton, A., and Sriperumbudur, B. K. (2025). (De)regularized maximum mean discrepancy gradient flow. Journal of Machine Learning Research, 26(235):1–77.

Figure 2: True and generated samples for different types of drift and hyperparameters. Empty panel means the samples have diverged.

Cortes, C., Mohri, M., and Rostamizadeh, A. (2009). L2 regularization for learning kernels. In Proceedings of the 25th Conference on Uncertainty in Artificial Intelligence (UAI 2009), pages 109–116. Cozzi, G. and Santambrogio, F. (2025). Long-time asymptotics of the sliced-Wasserstein flow. SIAM Journal on Imaging Sciences, 18(1):1–19. Crucinio, F. R., De Bortoli, V., Doucet, A., and Johansen, A. M. (2024). Solving Fredholm integral equations of the first kind via Wasserstein gradient flows. Stochastic Processes and Their Applications, 173. Cuturi, M. (2013). Sinkhorn distances: Lightspeed computation of optimal transport. In Advances in Neural Information Processing Systems. Deng, M., Li, H., Li, T., Du, Y., and He, K. (2026). Generative modeling via drifting. arXiv preprint arXiv:2602.04770. Dziugaite, G. K., Roy, D. M., and Ghahramani, Z. (2015). Training generative neural networks via maximum mean discrepancy optimization. In Uncertainty in Artificial Intelligence. Feydy, J., Séjourné, T., Vialard, F.-X., Amari, S.-i., Trouvé, A., and Peyré, G. (2019). Interpolating between optimal transport and MMD using Sinkhorn divergences. In International Conference on Artificial Intelligence and Statistics. Franz, L., Hoffmann, S., and Martius, G. (2026). Drifting fields are not conservative. arXiv preprint arXiv:2604.06333. Galashov, A., De Bortoli, V., and Gretton, A. (2025). Deep MMD gradient flow without adversarial training. In International Conference on Learning Representations. Glaser, P., Arbel, M., and Gretton, A. (2021). KALE flow: A relaxed KL gradient flow for probabilities with disjoint support. In Advances in Neural Information Processing Systems. Gretton, A., Borgwardt, K. M., Rasch, M. J., Schölkopf, B., and Smola, A. J. (2012). A kernel two-sample test. Journal of Machine Learning Research, 13. He, P., Khangaonkar, O., Pirsiavash, H., Bai, Y., and Kolouri, S. (2026). Sinkhorn-drifting generative models. arXiv preprint arXiv:2603.12366. Knight, P. A., Ruiz, D., and Uçar, B. (2014). A symmetry preserving algorithm for matrix scaling. SIAM Journal on Matrix Analysis and Applications, 35(3):931–955. hal-00569250. Lai, C.-H., Nguyen, B., Murata, N., Takida, Y., Uesaka, T., Mitsufuji, Y., Ermon, S., and Tao, M. (2026). A unified view of drifting and score-based models. arXiv preprint arXiv:2603.07514. Li, Y., Swersky, K., and Zemel, R. (2015). Generative moment matching networks. In International Conference on Machine Learning. Li, Z. and Zhu, B. (2026). A long-short flow-map perspective for drifting models. arXiv preprint arXiv:2602.20463. Liutkus, A., Simsekli, U., Majewski, S., Durmus, A., and Stöter, F.-R. (2019). Sliced-Wasserstein flows: Nonparametric generative modeling via optimal transport and diffusions. In International Conference on Machine Learning. 15

Mroueh, Y., Sercu, T., and Raj, A. (2019). Sobolev descent. In International Conference on Artificial Intelligence and Statistics. Nowozin, S., Cseke, B., and Tomioka, R. (2016). f-GAN: training generative neural samplers using variational divergence minimization. In Advances in Neural Information Processing Systems. Ramdas, A., Trillos, N., and Cuturi, M. (2017). On Wasserstein two-sample testing and related families of nonparametric tests. Entropy, 19(2). Santambrogio, F. (2017). {Euclidean, metric, and Wasserstein} gradient flows: an overview. Bulletin of Mathematical Sciences, 7(1):87–154. Sriperumbudur, B., Fukumizu, K., and Lanckriet, G. (2011). Universality, characteristic kernels and RKHS embedding of measures. Journal of Machine Learning Research, 12:2389–2410. Sriperumbudur, B., Gretton, A., Fukumizu, K., Lanckriet, G., and Schölkopf, B. (2010). Hilbert space embeddings and metrics on probability measures. Journal of Machine Learning Research, 11:1517–1561. Turan, E. and Ovsjanikov, M. (2026). Generative drifting is secretly score matching: a spectral and variational perspective. arXiv preprint arXiv:2603.09936. Wenliang, L. K. and Kanagawa, H. (2020). Blindness of score-based methods to isolated components and mixing proportions. arXiv preprint arXiv:2008.10087. Zhou, L., Ermon, S., and Song, J. (2025). Inductive moment matching. In Proceedings of the 42nd International Conference on Machine Learning, volume 267. PMLR.

Wasserstein Gradient flow of KL on Parzen Window density estimates

We derive the first variation and the Wasserstein gradient flow velocity for the KDE-smoothed (τ ) forward KL divergence Fp (q) := KL(qτ ∥pτ ), where qτ = kτ ∗ q and pτ = kτ ∗ p are Parzen window density estimates with a symmetric kernel kτ . (τ ) The functional Fp is a composition: q first maps to qτ via the linear convolution q 7→ kτ ∗ q, and then the standard KL divergence is evaluated. By (12), the first variation of KL(·∥pτ ) with respect to its first argument is δ KL(·∥pτ ) (x) = log qτ (x) − log pτ (x) + 1. δqτ (τ )

To obtain the first variation of Fp R with respect to q, we apply the chain rule through the convolution. For any perturbation χ with χ(y) dy = 0, linearity gives kτ ∗ (q + hχ) = kτ ∗ q + h (kτ ∗ χ) = qτ + h (kτ ∗ χ), so d d Fp(τ ) (q + hχ) = KL qτ + h (kτ ∗ χ) pτ dh h=0 dh Z h=0 δ KL(·∥pτ ) = (x) (kτ ∗ χ)(x) dx δqτ Z Z = log qτ (x) − log pτ (x) + 1 kτ (x, y) χ(y) dy dx,

where the second line is the definition of first variation applied to KL(·∥pτ ) with perturbation kτ ∗ χ. Exchanging the order of integration (Fubini) and using the symmetry kτ (x, y) = kτ (y, x), this equals Z Z kτ (y, x) log qτ (x) − log pτ (x) + 1 dx χ(y) dy. | {z } (τ ) δFp (y) = δq R Because kτ (y, x) dx is constant in y for translation-invariant kernels, the +1 term contributes only a constant and vanishes under ∇y . The WGF velocity of the KDE-smoothed KL divergence is therefore Z (τ ) δFp qτ (x) (τ ) Vp,q (y) = −∇y (y) = −∇y kτ (y, x) log dx. δq pτ (x)

Details of the Sinkhorn Proxy

This appendix contains the full derivation and analysis of the Sinkhorn Proxy drift field summarized in Section 3. We will switch from the generic notation Vp,q to method-specific notation, such as KL , V SP , in the following analysis to avoid ambiguity. Vp,q p,q

B.1

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

As noted in their paper, the drifting model as implemented does not use (Deng et al., 2026, eq. (8)) in practice, but the approach described in (Deng et al., 2026, Algorithm 2, Appendix A). We show that this represents a fixed point of a probability flow on an approximation to the Sinkhorn Optimal Transport distance, where the velocity field is zero if and only if p and q coincide. Despite this resemblance to a Sinkhorn distance, however, the resulting velocity field is not a valid Wasserstein gradient flow on any distributional loss, and the transport dynamics are in certain important aspects closer to a gradient flow on the KL. We begin with a brief review of the Sinkhorn distance and its associated Wasserstein gradient flow on a set of particles. Let q, p ∈ P2 (Rd ) be probability measures with bounded second moment and let C(x, y) = 12 ∥x − y∥2 . The debiased Sinkhorn divergence (Ramdas et al., 2017) is defined as: 1 1 Sτ (q, p) = OTτ (q, p) − OTτ (q, q) − OTτ (p, p), 2 2 where OTτ (q, p) is the entropy-regularized OT distance (Cuturi, 2013), and τ > 0 is the entropic regularization parameter. P + PN − Proposition 1 (Exact Sinkhorn particle velocity). Let p = N j=1 δyj and q = i=1 δxi be empirical −

−

×N ×N measures and let C(x, y) = 21 ∥x − y∥2 . Let W + ∈ RN and W − ∈ RN be the row+ + stochastic conditional optimal transport plans for OTτ (q, p) and OTτ (q, q), respectively, defined via the Sinkhorn dual potentials. Then the Wasserstein gradient flow velocity of the debiased Sinkhorn divergence Sτ (q, p) at location xi is −

S Vbp,q (xi ) =

N X

Wij+ yj −

j=1

N X k=1

− Wik xk .

(18)

Proof. By the envelope theorem (Feydy et al., 2019, Proposition 1), the first variation of Sτ with respect to its first argument is the difference of optimal dual potentials, δSτ (q, p) (x) = gq,p (x) − gq,q (x), δq where gq,p and gq,q are the optimal potentials on the q-side of OTτ (q, p) and OTτ (q, q), respectively. The WGF velocity is therefore δSτ (q, p) S (x) = −∇x gq,p (x) + ∇x gq,q (x). Vbp,q (x) = −∇x δq It remains to compute these gradients. The optimal transport plan W + and the potentials g, h (on the q and p sides of OTτ (q, p)) are related by ! g(xi ) + h(yj ) − 21 ∥xi − yj ∥2 + Wij = exp . τ The potentials satisfy the coupled fixed-point equations (solved by the Sinkhorn algorithm, initializing h = 0): ! ! N+ N− X X h(yj ) − 21 ∥xi − yj ∥2 g(xi ) − 12 ∥xi − yj ∥2 g(xi ) = −τ log exp , h(yj ) = −τ log exp . τ τ j=1

i=1

Taking ∇xi of the first equation and using the fixed-point relation exp(g(xi )/τ ) = −1 1 2 )/τ ∥x − y ∥ : i k 2 PN + ∇xi g(xi ) = −τ

j=1 exp

PN +

h(y )− 1 ∥x −y ∥2

k=1 exp

−(xi −yj ) τ

h(y )− 1 ∥x −y ∥2 k

N X j=1

Wij+ (xi − yj ) +

N X

Wij+ (xi − yj ).

− k Wik (xi − xk ). Substituting back:

−

− Wik (xi − xk ) =

N X

Wij+ yj −

j=1

k=1

(h(yk )−

j=1

−

S Vbp,q (xi ) = −

k exp

An identical calculation for the self-term gives ∇xi gq,q (xi ) =

N X

− Wik xk ,

(24)

k=1

P P − since W + and W − both have unit row sums ( j Wij+ = k Wik = 1).

B.2

The One-Shot Sinkhorn Proxy

Equation (18) requires computing the true conditional optimal transport plans (i.e., running Sinkhorn to convergence) at every timestep. A drifting model using the fully converged Sinkhorn method was proposed by He et al. (2026) and Turan and Ovsjanikov (2026, Section 6.3, Appendix I.4). To recover (Deng et al., 2026, Algorithm 2, Appendix A) to a close approximation, we replace the exact solution with a one-shot approximation. + + − − Let zim = −Cim /τ and zim = −Cim /τ be the scaled affinities, where i is an index of a q sample, + and m may be either a p-sample index or a q-sample index, depending on whether we refer to Cij or 18

− + Cij , respectively. We will focus on the case of the squared Euclidean distance, Cim = ∥xi − ym ∥2 .6 We approximate the Sinkhorn solution via the geometric mean of the row-wise and column-wise softmax normalizations,7 v v u u + − + − u exp(zim u exp(z ) ) exp(z exp(zim ) ) + − im im t , A , · = · Aim = t PN + PN − PN − PN − im + + − − p=1 exp(zip ) q=1 exp(zqm ) p=1 exp(zip ) q=1 exp(zqm ) − where we again use A+ im when m is a p-sample index, and Aim for a q-sample index m. This serves as a pseudo-plan, (A ≈ A, where A is the fully converged Sinkhorn solution). The relation of this approximation to the full Sinkhorn algorithm is described in Appendix C. We emphasize that unlike the full solution, the rows and columns of A are no longer guaranteed to sum to 1, and the use of the envelope theorem (Feydy et al., 2019) in obtaining the spatial gradient is no longer valid: treating the affinities as constants with respect to xi thus acts as a stop-gradient approximation. Because Aim does not perfectly enforce the normalization conditions (i.e., row and column sums are not exactly equal to 1), we may not directly use A+ and A− to compute the velocity field. This is addressed by a cross-weighting trick. Define X X − s+ A+ A− i = ij and si = ik j

Writing A+ ij

f + = A+ · s− = W ij ij i

− × s+ i · si | {z }

s+ i | {z }

+ conditional plan (π̃ij )

scaling(Zi )

− and defining Zi = s+ i si , we observe that we recover a conditional plan, scaled by a particle-specific Zi . We emphasize that due to the cross-weighting, the cancellation of x in (24) still holds, giving us the final update equation for x = x′i for every i ∈ {1, . . . , N − },   X X + − (t)  Ve (x) = Zi  π̃ij yj − π̃ik xk j

Thus, the drift vector Vei corresponds to the probability flow of the approximate Sinkhorn solution.

B.3

Alternative form, Sinkhorn proxy velocity field

An alternative form for the velocity field is given in the following proposition. 6

Note that the exact Sinkhorn flow derived in Proposition 1 uses the standard optimal transport ground cost C(x, y) = 21 ∥x − y∥2 . The empirical algorithm’s use of the unscaled distance ∥x − y∥2 mathematically introduces an overall constant factor of 2 into the theoretical spatial gradient, which in practice is simply absorbed by the step-size hyperparameter η during optimization. 7 In the original paper, the expression was v ! ! u u exp(zim ) exp(zim ) + Aim = t PN+ +N− . PN− exp(zip ) exp(zqm ) p q We replace the sum over N+ + N− samples with a sum over only one set of samples, for reasons that will become clear.

p Proposition 2 (Population form of the Sinkhorn Proxy). Let p̃(y) = p(y)/ qτ (y), and q̃(x) = p q(x)/ qτ (x). Then the Sinkhorn Proxy velocity field (19) with Gaussian kernel has the populationlimit form τ SP Vp,q (x) = Z(x) ∇ log(k ∗ p̃)(x) − ∇ log(k ∗ q̃)(x) , (21) 2 where Z(x) = s+ (x) s− (x) is a spatially varying pre-conditioner. Proof. Recall the original expression of the drift vector from (20), Z Z V (x) = s− (x) A+ (x, y)y dp(y) − s+ (x) A− (x, x′ )x′ dq(x′ ), where

k(x, x′ ) A− (x, x′ ) = p . qτ (x)qτ (x′ )

k(x, y) , A+ (x, y) = p pτ (x)qτ (y)

and s+ (x) =

(25)

s− (x) =

A+ (x, y) dp(y),

A− (x, x′ ) dq(x′ ).

p p Using the modified density definitions p̃(y) = p(y)/ qτ (y) and q̃(x′ ) = q(x′ )/ qτ (x′ ), we can express the row sums as kernel convolutions over these modified measures: Z 1 p(y) 1 + s (x) = p k(x, y) p dy = p (k ∗ p̃)(x), (26) pτ (x) qτ (y) pτ (x) Z 1 q(x′ ) 1 − s (x) = p k(x, x′ ) p dx′ = p (k ∗ q̃)(x). (27) ′ qτ (x) qτ (x ) qτ (x) Next, we replace the proxy affinities A+ and A− inside the expected destination integrals in (25), yielding: Z Z s− (x) s+ (x) V (x) = p yk(x, y)p̃(y) dy − p x′ k(x, x′ )q̃(x′ ) dx′ . (28) pτ (x) qτ (x) To evaluate these spatial integrals, we use Tweedie’s formula for the unnormalized Gaussian kernel k(x, z) = exp(−∥x − z∥2 /τ ). For any generic measure µ(z), Tweedie’s formula provides the identity: R zk(x, z)µ(z) dz τ R = x + ∇x log(k ∗ µ)(x). 2 k(x, z)µ(z) dz R Multiplying both sides by the denominator k(x, z)µ(z) dz = (k ∗ µ)(x) yields: Z h i τ zk(x, z)µ(z) dz = (k ∗ µ)(x) x + ∇x log(k ∗ µ)(x) . 2 We apply this result to (28) for both modified measures p̃ and q̃. From our row sum p p derivations + in (26) and (27), we recognize that (k ∗ p̃)(x) = pτ (x)s (x) and (k ∗ q̃)(x) = qτ (x)s− (x). Substituting these relations into the right-hand side of Tweedie’s formula gives: Z h i p τ yk(x, y)p̃(y) dy = pτ (x)s+ (x) x + ∇x log(k ∗ p̃)(x) , 2 Z h i p τ x′ k(x, x′ )q̃(x′ ) dx′ = qτ (x)s− (x) x + ∇x log(k ∗ q̃)(x) . 2 20

Finally, we substitute these evaluated integrals back into the drift vector expression (28): ! h i τ s− (x) p pτ (x)s+ (x) x + ∇x log(k ∗ p̃)(x) V (x) = p 2 pτ (x) ! h i s+ (x) p τ −p qτ (x)s− (x) x + ∇x log(k ∗ q̃)(x) . 2 qτ (x) The

p p pτ (x) and qτ (x) terms cancel, giving h i h i τ τ V (x) = s− (x)s+ (x) x + ∇x log(k ∗ p̃)(x) − s+ (x)s− (x) x + ∇x log(k ∗ q̃)(x) . 2 2

Factoring out the common pre-conditioner Z(x) = s+ (x)s− (x) and expanding the brackets, the coordinates x cancel: i h τ τ V (x) = Z(x) x + ∇x log(k ∗ p̃)(x) − x − ∇x log(k ∗ q̃)(x) 2 2 τ = Z(x) ∇x log(k ∗ p̃)(x) − ∇x log(k ∗ q̃)(x) . 2

B.4

Consistency of the Sinkhorn Proxy

Theorem 3 (Consistency). V = 0 if and only if p = q for the Gaussian kernel k. SP (x) vanishes evProof. We analyze the condition under which the cross-weighted vector field Vp,q erywhere.

( ⇐= ) Assume p = q. It immediately follows that the modified measures are identical: p(y) q(y) p̃(y) = p =p = q̃(y) qτ (y) qτ (y)

∀y.

Consequently, their convolutions with the kernel k are identical, yielding ∇x log(k∗p̃)(x) = ∇x log(k∗ q̃)(x) for all x, and the velocity field is identically zero everywhere. SP (x) = 0 for all x. Because the pre-conditioner Z(x) is strictly positive every( =⇒ ) Assume Vp,q where (as k, p, and q are strictly positive), we must have:

∇x log(k ∗ p̃)(x) = ∇x log(k ∗ q̃)(x)

∀x.

Integrating this spatial gradient implies that the kernel-smoothed measures are proportional up to a global constant λ > 0: (k ∗ p̃)(x) = λ(k ∗ q̃)(x) ∀x. When k is characteristic (Sriperumbudur et al., 2010) and translation invariant on Rd , then the mean embedding map µ 7→ k ∗ µ is injective for finite signed measures (Sriperumbudur et al., 2011, Proposition 2 and Section 3.2). This property holds for the Gaussian kernel considered in this study. Assuming the mild regularity condition that p and q have sufficiently fast-decaying tails (e.g., they are compactly supported) such that the modified measures p̃ and q̃ have finite mass, the

injectivity applies. Therefore, if the smoothed measures are proportional, the underlying modified measures must themselves be strictly proportional: p̃(y) = λq̃(y)

∀y.

Substituting the definitions of p̃ and q̃ yields: p(y) q(y) p = λp . qτ (y) qτ (y) Because the kernel-smoothed model density is strictly positive (qτ (y) > 0), we can safely multiply p both sides by qτ (y) to completely eliminate the dependence on the algorithmic smoothing: p(y) = λq(y)

∀y.

Since p and q are both valid probability density functions representing the target data distribution and the model distribution respectively, they must integrate to 1. Integrating both sides over the domain yields: Z Z p(y) dy = λ

q(y) dy =⇒ λ = 1.

This concludes the proof.

B.5

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

SP in (21), with Gaussian kernel, Proposition 4 (Non-gradient structure). The velocity field Vp,q corresponds to a valid Wasserstein gradient flow only when the gradients of the pre-conditioner and the score difference are linearly dependent everywhere:

∂i Z(x)∂j Φ(x) = ∂j Z(x)∂i Φ(x)

for all i, j,

where Φ(x) = τ2 (log(k ∗ p̃)(x) − log(k ∗ q̃)(x)). Equivalently, the wedge product must vanish: ∇Z ∧ ∇Φ = 0. SP (x) driving Proof. To be a valid Wasserstein gradient flow, the associated velocity vector field Vp,q the particles must take the form of a pure gradient field (see (3)). Thus, there must exist a scalar δF function (the first variation δqp (x)) whose gradient produces this field. We start with the form in proposition 2. Define the term inside the parenthesis as τ Φ(x) = log(k ∗ p̃)(x) − log(k ∗ q̃)(x) . 2

We can then rewrite the Sinkhorn Proxy drift concisely as: SP Vp,q (x) = Z(x)∇x Φ(x).

By Poincaré’s Lemma, a continuously differentiable vector field on a simply connected domain is a pure gradient field if and only if its Jacobian matrix is symmetric. Let Vi denote the i-th SP . The symmetry condition ∂ V = ∂ V implies: component of Vp,q j i i j ∂j Z(x)∂i Φ(x) = ∂i Z(x)∂j Φ(x) ∂j Z(x)∂i Φ(x) + Z(x)∂j ∂i Φ(x) = ∂i Z(x)∂j Φ(x) + Z(x)∂i ∂j Φ(x). 22

By Schwarz’s theorem, the mixed partial derivatives of Φ are equal (∂j ∂i Φ = ∂i ∂j Φ), so ∂j Z(x)∂i Φ(x) = ∂i Z(x)∂j Φ(x). In other words, the gradient of the scalar pre-conditioner ∇Z(x) must be perfectly parallel to the gradient of the score difference ∇Φ(x) at all points in space. Given their distinct definitions, this geometric condition will not hold in general.

Geometric mean as approximation to Sinkhorn solution

We demonstrate the relation between A+ and the optimal transport plan obtained via Sinkhorn iteration. The analogous result applies for A− . In entropy-regularized optimal transport (Cuturi, 2013), the interaction between points is governed by a strictly positive Gibbs kernel matrix K. The algorithm defines the independent pairwise + + logits as Zij = −Cij /τ . Exponentiating these logits yields the Gibbs kernel: + Kij = exp(Zij ) = exp

−

+ Cij

The objective of the Sinkhorn algorithm is to find the strictly positive vectors u and v such that the resulting transport plan P = diag(u)Kdiag(v) satisfies target marginal constraints. To recover the Sinkhorn proxy algorithm, these marginals are assumed to be uniform, (P 1)⊤ = 1⊤ P = 1. The resulting optimal transport plan P ∗ takes the form Pij∗ = ui Kij vj The standard iterative Sinkhorn algorithm finds u and v by repeatedly alternating between dividing by row sums and column sums until the marginals match the target distributions. Assume first that Sinkhorn initializes the column scaling vector as a vector of ones: v (0) = 1. The first update for the source scaling vector u is computed by dividing the target marginals by the row sums of the current scaled kernel, (1)

1 1 1 = = (K1)i ri (Kv (0) )i

(1)

Applying this scaling alone yields ui Kij = Kij /ri , which corresponds exactly to the softmax row operation. Consider now the case that we begin the Sinkhorn iterations with (u(0) = 1). We would then make the update 1 1 1 (1) vj = = = . T (0) T (K 1)j cj (K u )j Turning now to the approach of Deng et al. (2026), we examine the relation between the matrix + A+ and the Sinkhorn structure. First, the row-wise P softmax of the logits Z is equivalent to the Gibbs kernel divided by its row sums. Let ri = m Kim be the i-th row sum: softmaxrow (Z + )ij = P

+ exp(Zij )

= + m exp(Zim )

Kij ri

Similarly, the column-wise softmax is the Gibbs kernel divided by its column sums. Let cj = be the j-th column sum:

n Knj

+ exp(Zij ) Kij softmaxcol (Z )ij = P + = c j n exp(Znj ) +

The update of Deng et al. corresponds to a geometric mean of these normalization operations, Kij A+ ij = √ ri · cj rather than iterating between them. Repeated application of the procedure converges to the exact transport map (Knight et al., 2014, Theorem 3.1), as with the better known alternating approach. While the above correction does not guarantee the marginal constraints of a fully converged Sinkhorn approach (which would require infinite sequential iterations), it does capture the property of “hub suppression:” If a target yj is a hub (a region of high density surrounded by many source √ points), its column sum cj will be large. Consequently, the column scaling factor 1/ cj becomes small, which down-scales the affinity of particles to that hub.

Failure mode on a synthetic example

We first describe the velocity field of the KL-KDE Proposition 5. When D and τ jointly satisfy (22) and a Gaussian kernel is used, the velocity field VKL of the KL-KDE flow at x = +D is given by α β KL Vp,q (D) = −2Dϵ − + o(ϵ). 1−α 1−β Proof. The continuous KL-KDE drift evaluates the difference between the scores of the kernelsmoothed distributions: τ KL Vp,q (x) = ∇x log pτ (x) − ∇x log qτ (x) . 2 We evaluate the score of the target distribution pτ at x = +D. The smoothed density is pτ (D) = (1 − α)k(D, D) + αk(D, −D) = (1 − α) + αϵ. For the unnormalized Gaussian kernel k(x, z) = exp(−∥x−z∥2 /τ ), the spatial gradient is ∇x k(x, z) = − 2(x−z) k(x, z). Evaluating this at x = +D gives ∇x k(D, D) = 0 and ∇x k(D, −D) = − 4D τ τ ϵ. The spatial gradient of the smoothed target distribution is 4D ∇x pτ (D) = (1 − α)∇x k(D, D) + α∇x k(D, −D) = α − ϵ . τ The score function is the ratio of these two quantities: ∇x log pτ (D) =

− 4D ∇x pτ (D) τ αϵ = . pτ (D) (1 − α) + αϵ

Assuming the mixing proportion α ̸= 1), 1 1 = (1 − α) + αϵ 1−α

1 α ϵ 1 + 1−α

! .

Expanding (1 + z)−1 = 1 − z + O(z 2 ) for small z, 4D α 1 2 ∇x log pτ (D) = − 1− αϵ ϵ + O(ϵ ) τ 1−α 1−α 2 ϵ α 4D +O . ϵ =− τ 1−α τ 2

To verify that the remainder term O( ϵτ ) is indeed o(ϵ), we must show that limτ →0 ϵ ϵ/τ = 0. 2 Using the fact that ϵ = exp(−4D2 /τ ), we can express τ as τ = − 4D ln ϵ . Then: ϵ2 /τ ϵ −ϵ ln ϵ =0 = lim = lim ϵ→0 ϵ ϵ→0 τ ϵ→0 4D 2

(29)

lim

Since the limit is zero, the remainder term decays strictly faster than ϵ. Therefore, O( ϵτ ) ⊂ o(ϵ), and thus, 2 τ τ 4D α ϵ ∇x log pτ (D) = − ϵ +O 2 2 τ 1−α τ α = −2Dϵ + o(ϵ). 1−α Applying the same reasoning to the model distribution qτ , with mixing proportion β, yields: τ β ∇x log qτ (D) = −2Dϵ + o(ϵ). 2 1−β Subtracting the two scaled scores gives the net KL-KDE drift: τ τ KL Vp,q (D) = ∇x log pτ (D) − ∇x log qτ (D) 2 2 α β = −2Dϵ + o(ϵ) − −2Dϵ + o(ϵ) 1−α 1−β α β − + o(ϵ). = −2Dϵ 1−α 1−β

The drift points in the correct direction (negative, since we assume β < α), but its magnitude is bounded by ϵ. Because τ is narrow relative to the mode separation 2D, ϵ may drop below machine precision (ϵ ≈ 0). We now show that the Sinkhorn proxy method has the exact same failure mode. Proposition 6. When D and τ jointly satisfy (22) and a Gaussian kernel is used, the velocity field of the Sinkhorn Proxy flow at x = D is given by s " # 1−β α β SP Vp,q (D) = −2Dϵ Z(D) − + o(ϵ), β 1−α 1−β | {z } (*)

where the term (*) represents the difference with the KL-KDE flow. 25

Proof. We begin with the population-limit velocity field established in proposition 2, τ Z(x) ∇x log(k ∗ p̃)(x) − ∇x log(k ∗ q̃)(x) , 2 p p SP (x) at where p̃(y) = p(y)/ qτ (y) and q̃(x) = q(x)/ qτ (x). We evaluate the algorithmic drift Vp,q the right-hand overcrowded mode x = +D. First, we evaluate the kernel-smoothed model density qτ exactly at both modes, expressing the result as 1 + O(ϵ) perturbations: SP Vp,q (x) =

qτ (D) = βk(D, −D) + (1 − β)k(D, D) = (1 − β) + βϵ = (1 − β)(1 + O(ϵ)) qτ (−D) = βk(−D, −D) + (1 − β)k(−D, D) = β + (1 − β)ϵ = β(1 + O(ϵ)) Using the Taylor expansion (1 + O(ϵ))−1/2 = 1 + O(ϵ), we substitute these into the modified density functions p̃ and q̃: α 1−α α 1−α p̃(y) = p δ−D (y) + p δ+D (y) = √ + O(ϵ) δ−D (y) + √ + O(ϵ) δ+D (y) β 1−β qτ (−D) qτ (D) p p 1−β β δ−D (x) + p δ+D (x) = β + O(ϵ) δ−D (x) + 1 − β + O(ϵ) δ+D (x) q̃(x) = p qτ (−D) qτ (D) Next, we evaluate the kernel-smoothed modified measures at x = +D. Since k(D, −D) = ϵ and k(D, D) = 1: α 1−α 1−α (k ∗ p̃)(D) = √ + O(ϵ) ϵ + √ + O(ϵ) (1) = √ + O(ϵ) β 1−β 1−β p p p (k ∗ q̃)(D) = β + O(ϵ) ϵ + 1 − β + O(ϵ) (1) = 1 − β + O(ϵ) To compute the score functions, we evaluate the spatial gradients at x = +D. For the unnormalized Gaussian kernel k(x, z) = exp(−∥x−z∥2 /τ ), the spatial gradients are ∇x k(x, −D)|x=D = − 4D τ ϵ and ∇x k(x, D)|x=D = 0. This yields: 2 α 4D 1−α 4D α ϵ ∇x (k ∗ p̃)(D) = √ + O(ϵ) − ϵ + √ + O(ϵ) (0) = − ϵ√ + O τ τ τ β 1−β β 2 p 4D p p 4D ϵ ∇x (k ∗ q̃)(D) = β + O(ϵ) − ϵ + 1 − β + O(ϵ) (0) = − ϵ β+O τ τ τ The score functions at x = +D are then the ratio of the to the densities. Utilizing gradients Aϵ+O(ϵ2 /τ ) A ϵ2 the formal quotient Taylor expansion B+O(ϵ) = B ϵ + O τ and our previous finding in (29) that O(ϵ2 /τ ) = o(ϵ), we have: √α − 4D τ ϵ β +O

2 ϵ τ

s α 1−β ∇x log(k ∗ p̃)(D) = + o(ϵ) 1−α √ 1−α β + O(ϵ) 1−β 2 s √ ϵ − 4D β + O ϵ τ τ 4D β ∇x log(k ∗ q̃)(D) = √ =− ϵ + o(ϵ) τ 1−β 1 − β + O(ϵ)

4D =− ϵ τ

Notice that

β 1−β =

β 1−β

1−β β . Substituting these into the algorithmic drift equation gives:

" !# s s 4D α 4D β τ τ 1 − β 1 − β SP ϵ − − ϵ + Z(D)o(ϵ) Vp,q (D) = Z(D) − 2 τ 1−α β τ 1−β β 2 s " # α 1−β β = −2Dϵ Z(D) + o(ϵ) − β 1−α 1−β The o(ϵ) bound strictly holds because the spatial pre-conditioner operates smoothly as Z(D) = O(1) with respect to ϵ. Proof. We begin with the population-limit velocity field established in proposition 2, τ Z(x) ∇x log(k ∗ p̃)(x) − ∇x log(k ∗ q̃)(x) , 2 p p where p̃(y) = p(y)/ qτ (y) and q̃(x) = q(x)/ qτ (x). We wish to evaluate the algorithmic drift SP (x) at the right-hand overcrowded mode x = +D. Vp,q First, we evaluate the kernel-smoothed model density qτ at both modes. Substituting q(x) into the convolution and dropping terms where the exponentially small ϵ is added to a quantity of O(1) yields: SP Vp,q (x) =

qτ (+D) = βk(D, −D) + (1 − β)k(D, D) = βϵ + (1 − β) ≈ 1 − β qτ (−D) = βk(−D, −D) + (1 − β)k(−D, D) = β + (1 − β)ϵ ≈ β We now substitute these marginal approximations and the original point masses into the modified density functions p̃ and q̃: α 1−α p̃(y) ≈ √ δ−D (y) + √ δ+D (y) β 1−β p p β 1−β q̃(x) ≈ √ δ−D (x) + √ δ+D (x) = βδ−D (x) + 1 − βδ+D (x) β 1−β Next, we evaluate the kernel-smoothed modified measures at x = +D, applying the same ϵ ≈ 0 approximation for additive terms: α 1−α α 1−α 1−α (k ∗ p̃)(D) = √ k(D, −D) + √ k(D, D) = √ ϵ + √ ≈√ β 1−β β 1−β 1−β p p p p p (k ∗ q̃)(D) = βk(D, −D) + 1 − βk(D, D) = βϵ + 1 − β ≈ 1 − β To compute the score functions, we need the spatial gradients evaluated at x = +D. For the unnormalized Gaussian kernel k(x, z) = exp(−∥x − z∥2 /τ ), the spatial gradients are ∇x k(x, −D)|x=D = − 4D τ ϵ and ∇x k(x, D)|x=D = 0. Applying this gives: α 1−α 4D α ∇x (k ∗ p̃)(D) = √ ∇x k(D, −D) + √ ∇x k(D, D) = − ϵ√ τ β 1−β β p p 4D p ∇x (k ∗ q̃)(D) = β∇x k(D, −D) + 1 − β∇x k(D, D) = − ϵ β τ

The score functions at x = +D are then: √α − 4D ∇x (k ∗ p̃)(D) 4D α τ ϵ β ∇x log(k ∗ p̃)(D) = = − ≈ ϵ √1−α (k ∗ p̃)(D) τ 1−α 1−β s √ ϵ β − 4D 4D β ∇x (k ∗ q̃)(D) τ =− ∇x log(k ∗ q̃)(D) = ≈ √ ϵ (k ∗ q̃)(D) τ 1−β 1−β

1−β β

Thus !# " s s τ 1 − β 1 − β 4D α 4D β SP Vp,q (D) ≈ Z(D) − ϵ − − ϵ 2 τ 1−α β τ 1−β β s " # 1−β α β = −2Dϵ Z(D) − β 1−α 1−β

This confirms that in the present setting, despite its superficial resemblance to a Sinkhorn flow, the Sinkhorn Proxy algorithm does not share the mass transport dynamics of an Earth Mover’s Distance, but behaves more like a flow on KL. We suspect that Deng’s Algorithm 2 also has this problem. Proposition 7. The velocity field VW2 of the Wasserstein gradient flow on the Wasserstein-2 distance at x = D is given by α−β W2 Vp,q (D) = −2D . 1−β Proof. We compute the velocity field at x = +D for a true Wasserstein-2 (W2 ) gradient flow. As defined in the main text, the target data distribution p and the model distribution q are mixtures of two isolated point masses: p(y) = αδ−D (y) + (1 − α)δ+D (y), q(x) = βδ−D (x) + (1 − β)δ+D (x). We assume β < α. This means the model q is “starved” of mass at the left mode (−D) and “overcrowded” at the right mode (+D). To find the Wasserstein-2 velocity, we must first find the exact optimal transport plan π(x, y) that minimizes the quadratic cost C(x, y) = 21 ∥x − y∥2 to move mass from q to p. The joint distribution π(x, y) must satisfy the following marginal constraints X X π(x, y) = q(x), π(x, y) = p(y). y

To minimize the cost, mass should remain in place wherever possible (cost = 0). • All mass β at x = −D in q will stay at y = −D in p. Thus, π(−D, −D) = β. • Because the target p requires a total mass of α at −D, the remaining deficit of (α − β) must be transported from x = +D. Thus, π(+D, −D) = α − β. • The remaining mass at x = +D stays at y = +D. Thus, π(+D, +D) = (1−β)−(α−β) = 1−α. 28

• Finally, to obtain that the probability sum at one, we have π(−D, +D) = 0. For probability mass at x = +D, the Wasserstein gradient flow velocity is the expected displacement under the optimal transport plan, where Z W2 Vp,q (x) = − ∇x C(x, y) dπ(y|x) over the optimal conditional transport plan dπ(y|x) (Ambrosio et al., 2008). We compute the conditional probability π(y | x) = π(x,y) q(x) for the overcrowded mode x = +D: The total available mass at x = +D is q(+D) = 1 − β. • Fraction of mass moving to −D: π(−D | +D) = α−β 1−β • Fraction of mass staying at +D: π(+D | +D) = 1−α 1−β W2 for the mass at x is the difference between its expected target destination The velocity vector Vp,q and its current position: W2 Vp,q (x) = EY ∼π(·|x) [Y ] − x

Evaluating this at x = +D: W2 Vp,q (+D) = [π(−D | +D)(−D) + π(+D | +D)(+D)] − D.

Substituting the conditional probabilities, α−β 1−α W2 Vp,q (+D) = (−D) + (+D) − D 1−β 1−β −(α − β) + (1 − α) − (1 − β) =D 1−β α−β = −2D 1−β which concludes the proof.

Proof of Proposition 8

Proposition 8. Let qθ = (fθ )# µ be the pushforward of µ through fθ , and suppose we have the δF exact velocity field Vp,qθ (x) = −∇x δqθp (x) for some functional Fp . Then ∇θ L(θ) = 2η ∇θ Fp (qθ ). Proof. Write x = fθ (ϵ) and let t = sg(x + ηVp,qθ (x)) denote the stop-gradient target, which is treated as a constant with respect to θ. Then L(θ) = Eϵ∼µ [∥fθ (ϵ) − t∥2 ], and applying the chain rule yields: h i ⊤ ∇θ L(θ) = Eϵ∼µ 2 fθ (ϵ) − t ∇θ fθ (ϵ) h i ⊤ = Eϵ∼µ 2 fθ (ϵ) − fθ (ϵ) − ηVp,qθ (fθ (ϵ)) ∇θ fθ (ϵ) h i ⊤ = −2η Eϵ∼µ Vp,qθ fθ (ϵ) ∇θ fθ (ϵ) . 29

A perturbation δθ changes the functional through the marginal distribution qθ . By the definition of the first variation: Z δFp ∇θ Fp (qθ ) = (x) ∇θ qθ (x) dx. δqθ Since qθ = (fθ )# µ, the parameter perturbation δθ induces a displacement field on the particles in data space. Because multiple noise instances ϵ could map to the same spatial point x (if fθ is noninvertible), we define the effective Eulerian velocity field at spatial location x via the conditional expectation: vθ (x) = Eϵ∼µ ∇θ fθ (ϵ) | fθ (ϵ) = x . The evolution of the density qθ with respect to θ is mathematically governed by the continuity equation ∇θ qθ (x) = −∇x · (qθ (x)vθ (x)). Substituting this into the first variation and applying the divergence theorem (assuming boundary terms vanish) gives: Z δFp ∇θ Fp (qθ ) = − (x) ∇x · qθ (x) vθ (x) dx δqθ Z δFp (x)⊤ vθ (x) qθ (x) dx. = ∇x δqθ Using the law of total expectation (the tower property), we can rewrite this integral over the spatial domain Rd as an expectation directly over the latent noise distribution µ: h δF i p ∇θ Fp (qθ ) = Ex∼qθ ∇x (x)⊤ Eϵ∼µ ∇θ fθ (ϵ) | fθ (ϵ) = x δqθ h δF i ⊤ p = Eϵ∼µ ∇x fθ (ϵ) ∇θ fθ (ϵ) . δqθ

(30) δF

Finally, substituting the definition of the Wasserstein gradient Vp,qθ (x) = −∇x δqθp (x) into (30) directly recovers the relation ∇θ L(θ) = 2η ∇θ Fp (qθ ).

Additional results

In Figure 3 to Figure 6, we present experimental results for various toy examples.

Figure 3: Results for the 8 Gaussian dataset. Empty panel means the samples have diverged.

Figure 4: Results for the Circles dataset. Empty panel means the samples have diverged.

Figure 5: Results for the Pinwheel dataset. Empty panel means the samples have diverged.

Figure 6: Results for the Swiss roll dataset. Empty panel means the samples have diverged.

Record · ID 158545 · SHA-256 04edf0406fef8b39
Retrieved via Conceptio — every document is proof-bundled with source, license, and retrieval metadata.