Generative Models via Transportation

The preceding gradient-flow calculus is variational. Modern machine-learning models often use the same transportation language more broadly: one may prescribe an interpolation and regress its velocity, fit a one-step generator to a descent field, or view network depth as a continuous transport of token measures. The examples below separate what is genuinely a Wasserstein gradient flow from what is a transportation dynamics with a useful geometric interpretation.

from pathlib import Path
import sys

from IPython.display import Image as DisplayImage
from IPython.display import display

here = Path.cwd()
myst_dir = None
for candidate in [here, here.parent, here / "myst", here.parent / "myst", here.parent.parent / "myst"]:
    if (candidate / "ot4ml_web.py").exists():
        myst_dir = candidate.resolve()
        sys.path.insert(0, str(myst_dir))
        break

if myst_dir is None:
    raise RuntimeError("Could not locate myst/ot4ml_web.py")

repo_root = myst_dir.parent
thumbnails = repo_root / "notebooks-figures" / "thumbnails"

def show_book_figure(name, width=760):
    display(DisplayImage(filename=str(thumbnails / f"{name}.png"), width=width))

Generative Models via Flow Matching¶

Flow matching constructs a generative map by learning the velocity field of an interpolation. The key computational insight is that a constrained continuity-equation problem can be trained by an unconstrained regression.

Generative models aim to build a transportation map $T$ between a reference distribution $\alpha$ (typically an isotropic Gaussian) and the target data distribution $\beta$. Since such reference measures are non-atomic, a measurable map with $T_\sharp\alpha=\beta$ exists on standard Borel spaces, for instance by identifying both probability spaces with the unit interval and using a quantile-type rearrangement. This abstract existence statement is much weaker than having an explicit and numerically stable construction of $T$. Optimal transport is one approach to achieving this, but it is computationally expensive and raises questions about how to estimate it from samples. A different route is to prescribe an interpolation between noise and data, learn its velocity, and obtain $T$ by integrating a time-dependent vector field $v_t$. This point of view sits at the meeting point of two literatures, surveyed from a transport perspective in Peyré, 2025. The diffusion branch builds on score matching Hyvärinen, 2005, denoising score matching Vincent, 2011, nonequilibrium noising chains Sohl-Dickstein et al., 2015, denoising diffusion probabilistic models Ho et al., 2020, score-based generative modeling Song & Ermon, 2019, and the continuous-time score-SDE/probability-flow formulation Song et al., 2021. The deterministic regression branch was introduced, essentially in parallel, under three closely related names: flow matching Lipman et al., 2023, rectified flow Liu et al., 2023, and stochastic interpolants Albergo et al., 2025. In all three cases, the computational object is a velocity field whose regression loss avoids simulating the learned ODE during training. This vector field $v_t$ is obtained by constructing an interpolation $\alpha_t$ and then finding $v_t$ using the least-squares formula of the dynamic chapter. As we will explain, for a specific class of interpolation (obtained by a parametric push-forward), this $v_t$ can be obtained by avoiding explicitly inverting a Laplacian and instead computing a simple conditional expectation. This conditional expectation can itself be estimated by solving another least-squares problem, but this time unconstrained, making the estimation feasible from finite samples of $\alpha$ and $\beta$.

Stochastic interpolant.¶

The word “stochastic” can hide two different levels of randomness. We first use the simpler one: after drawing a latent variable $U\sim\pi$, the path $t\mapsto P_t(U)$ is deterministic and differentiable. The randomness only comes from the initial draw of $U$; after taking the push-forward law, $\alpha_t=(P_t)_\sharp\pi$ is a deterministic curve of measures and obeys an ordinary continuity equation. This is the setting behind the stochastic-interpolant construction recalled in Remark Remark: Static-noise stochastic interpolants, and behind the flow-matching and rectified-flow regressions below. Genuine temporal noise, where the path itself has Brownian fluctuations, is different and is discussed in Remark Remark: Brownian realizations of interpolant marginals.

We assume first that $\alpha_t$ is obtained by pushing a latent distribution $\pi \in \Pp(\RR^{d'})$ through a time-dependent map $P_t : \RR^{d'} \to \RR^d$; the latent dimension $d'$ may be larger than the data dimension $d$:

The basic two-endpoint construction already covers most flow-matching paths used in practice.

If $\pi = \alpha \otimes \beta$ and $\alpha = \frac{1}{n} \sum_i \delta_{x_i}$, $\beta = \frac{1}{m} \sum_j \delta_{y_j}$, then $\alpha_t$ consists of $n \times m$ Dirac masses

If $\pi = (\Id, T)_\sharp \alpha$ is a Brenier-type coupling, then $\alpha_t = ((1-t)\Id + tT)_\sharp \alpha$ is the so-called McCann OT interpolation.

Figure Div separates the effect of the endpoint coupling from the effect of the path joining each paired endpoint.

show_book_figure("generative-flow-matching-interpolants", width=760)

Flow matching interpolants between the same empirical source and target measures. A product-style random pairing produces crossing paths, an OT pairing gives direct displacement rays, and a curved bridge changes the path geometry while keeping the same endpoints. Gray arrows mark representative midpoint velocities $\partial_tP_t$.

Flow matching formula.¶

This interpolation is not directly useful for sampling from $\beta$, but it can be used to define a flow field $v_t$ so that the continuity equation, in Eulerian form, holds. This flow field is computed by solving an unconstrained least-squares problem, or equivalently, it is a conditional expectation.

The conditional expectation in (5) has a simple measure-theoretic meaning. Let $\alpha_t=(P_t)_\sharp\pi$ and define the vector-valued flux measure $\omega_t$ on $\RR^d$ by

for every bounded continuous vector field $\psi$. Since $\alpha_t(A)=0$ implies $\pi(P_t^{-1}(A))=0$, one has $\omega_t\ll\alpha_t$. The Radon--Nikodym decomposition of $\omega_t$ with respect to $\alpha_t$ is therefore

In the language of Lebesgue decomposition, $\omega_t$ has only an absolutely continuous part with respect to $\alpha_t$ and no singular part; the conditional expectation is precisely its density. This agrees with the flux notation used in the dynamic formulation. Equivalently, disintegrating $\pi$ with respect to the map $P_t$ gives $\pi(\d u)=\pi_{t,z}(\d u)\alpha_t(\d z)$, where $\pi_{t,z}$ is supported on the fiber $\{u\,:\,P_t(u)=z\}$, and

Thus the solution of (4) is the conditional expectation of the velocities $\partial_t P_t$: intuitively, $v_t(z)$ is the average velocity of all trajectories passing through $z$. Numerically, $(x,t) \to v_t(x)$ can be parameterized by a neural network (e.g., a U-Net for vision tasks) and estimated using stochastic gradient descent on the objective in (4). For the exact field $v_t$, integrating the ODE $\dot{x}=v_t(x)$ defines a transport map $T_t$. If $v_t$ is regular enough, or more generally if the continuity equation has a unique solution for this velocity, then $(T_t)_\sharp\alpha_0=\alpha_t$. Thus the same interpolation as (1) is represented by a deterministic flow rather than by the original coupling. The sampling procedure consists in first drawing $X_0 \sim \alpha$, and then integrating the ODE $\dot{X}_t = v_t(X_t)$ starting with $X_{t=0} = X_0$. In the ideal exact-field limit, the resulting $X_{t=1}$ is distributed according to $\alpha_1 = \beta$.

Connection with diffusion models.¶

In the special case where $P_t(x,y)=(1-t)x+ty$ is a linear interpolation and $\pi = \alpha \otimes \beta$, the curve $\alpha_t$ is a convolution of rescaled versions of $\alpha_0$ and $\alpha_1$. The flow-matching problem (4) becomes

When one endpoint is an isotropic Gaussian, this construction is closely related to the probability-flow formulation of diffusion models, up to the usual change of time parametrization Song et al., 2021. This is why flow matching can be viewed both as a deterministic alternative to diffusion training and as a common language for diffusion paths, OT-inspired paths, and rectified paths Lipman et al., 2023Liu et al., 2023Albergo et al., 2025. The next two propositions are written in the noising direction, from a data law $\alpha$ to a Gaussian; reversing time gives the corresponding sampling flow. They also give an explicit closed form for $v_t$ and show that it is a gradient field. In this setting, $v_t$ is also the solution of the constrained least-squares problem from the dynamic chapter. The regression (4) is computationally simpler because the continuity equation has already been enforced by the chosen interpolant. To prove this, we rely on Tweedie’s formula Efron, 2011, which expresses the optimal Gaussian denoiser through the score, i.e. the gradient of the log-density.

Figure Div shows both directions of this construction: the prescribed interpolation noising the data and the same probability-flow ODE integrated backward for sampling.

show_book_figure("generative-diffusion-1d-forward-backward", width=760)

One-dimensional diffusion bridge for a Gaussian-mixture data law. The forward path $Z_t=(1-t)X+tY$ smooths the red data density toward a blue Gaussian endpoint. Reversing the probability-flow ODE transports a denser set of blue noise samples back toward the data modes, making the splitting of trajectories across mixture components visible.

The same probability-flow intuition is visible in two dimensions. For a discrete data law, or more generally for a Gaussian mixture, the noising density is a Gaussian mixture whose score can be evaluated explicitly. This makes it possible to draw backward trajectories without training a neural network. In the plots below, the Gaussian endpoint has covariance $\sigma^2\Id$ to keep the geometry visible at the scale of the three atoms. For a scalar noising schedule $Z_t=a_tX+b_tY$, the intermediate law has component centers $a_t c_j$ and covariance $(b_t\sigma)^2\Id$. For the linear bridge, $p_t(z)=\sum_j w_j\Gaussian((1-t)c_j,(t\sigma)^2\Id)$, with $s_t=\nabla\log p_t$, and the scaled Gaussian-endpoint field gives $v_t(z)=-(z+t\sigma^2s_t(z))/(1-t)$.

In Figure Div, the Gaussian endpoint has covariance $\sigma^2\Id$ to keep the geometry visible at the scale of the three atoms.

show_book_figure("generative-diffusion-2d-forward-backward", width=760)

Two-dimensional noising paths from three Dirac masses to a single Gaussian. The linear interpolation $Z_t=(1-t)X+tY$ moves component centers linearly toward the origin and grows covariance like $(t\sigma)^2\Id$. The variance-preserving Ornstein--Uhlenbeck bridge has the same endpoints but a different speed of contraction and noising.

When is the induced map optimal?¶

Integrating the learned velocity gives a deterministic map from $\alpha_0$ to $\alpha_1$, but this map is not automatically the Brenier optimal map. It is optimal only in special cases where the accumulated flow remains the gradient of a convex potential. The Gaussian product-coupling case already shows the precise obstruction: the interpolated covariances are simple, the velocity is affine, but the terminal map can contain a hidden rotational part. This phenomenon, and its extensions to rectified flows and mixtures, is analyzed in depth in Hertrich et al., 2025.

The Gaussian optimality proposition explains why the statement “flow matching gives an optimal map” is fragile. The same terminal map (41) is obtained for any scalar schedule $Z_t=a_tX_0+b_tX_1$ with the same endpoints, because after whitening the covariance path remains $a_t^2\Id+b_t^2C$. Thus changing the speed of a scalar Gaussian bridge, for instance by using an OU schedule, cannot repair the non-optimality created by non-commuting covariances. Commuting covariances reduce the terminal map to independent one-dimensional scalings, whereas non-commuting covariances create a non-symmetric affine map, hence a transport with a rotational or shearing component. More generally, mixture-like paths can create the same obstruction even when every instantaneous velocity looks natural. This distinction is closely related to counterexamples showing that flow maps associated with Fokker--Planck or diffusion-type evolutions do not in general provide optimal transport maps Lavenant & Santambrogio, 2022. In particular, starting from an isotropic Gaussian does not by itself guarantee optimality once the target distribution is non-Gaussian; additional structural assumptions on the path or on the coupling are needed.

Variations on the interpolant.¶

The geometry of the generated trajectories depends on the chosen interpolant, not only on the two endpoint laws. There is first a harmless ambiguity: a monotone reparametrization $Z_t=(1-\lambda(t))X+\lambda(t)Y$ of the linear bridge only changes the speed of the flow,

It therefore leaves the spatial integral curves unchanged. Diffusion models use a genuinely different family of noising paths. If

then both the mixture centers and the component variances are changed. Writing $p_t$ for the density of $Z_t$ and $s_t=\nabla\log p_t$, Tweedie’s formula gives, away from times where $a_t=0$,

For the linear bridge, $a_t=1-t$ and $b_t=t$, this recovers the formula above. For the variance-preserving Ornstein--Uhlenbeck noising used in diffusion models,

one obtains the forward probability-flow velocity $v_\tau(z)=-z-\sigma^2\nabla\log p_\tau(z)$. Sampling follows the reverse field $z+\sigma^2\nabla\log p_\tau(z)$ as $\tau$ decreases. This is the noising law used in the diffusion trajectory panel below; the trajectories are more curved than for the linear bridge because the centers and variances evolve according to the OU/Fokker--Planck scaling rather than by affine interpolation. Numerically, the integration is stopped at a small positive time before the Dirac endpoint, where the score becomes singular.

The finite-time coefficients $a_t=\cos(\pi t/2)$ and $b_t=\sin(\pi t/2)$ are not a new spatial interpolant: they are exactly the OU coefficients after the time change $\tau=-\log\cos(\pi t/2)$. The schedule comparison below therefore places the OU bridge next to a genuinely different scalar bridge,

whose data coefficient changes sign before vanishing. This overshooting bridge is mainly a diagnostic example: it keeps the same endpoints, but its intermediate mixture reflects through the origin and produces visibly different reverse trajectories.

This is the noising law used in the left panel of Figure Div; the trajectories are more curved than for the linear bridge because the centers and variances evolve according to the OU/Fokker--Planck scaling rather than by affine interpolation.

show_book_figure("generative-diffusion-versus-ot-2d", width=760)

Diffusion-style sampling trajectories compared with OT rays in the three-Dirac setting. Red particles are sampled from the centered Gaussian endpoint and transported toward the three blue atoms. The diffusion panel integrates a reverse probability-flow field based on a Gaussian-mixture score, while the OT panel uses straight displacement rays selected by a quadratic matching.

Figure Div therefore compares OU with a genuinely different scalar bridge,

show_book_figure("generative-diffusion-schedule-comparison", width=760)

Effect of the interpolant on the exact reverse flow for the same three-Dirac target and the same Gaussian endpoint. The linear bridge $a_t=1-t$, $b_t=t$ produces almost radial curves. The variance-preserving OU bridge $a_\tau=e^{-\tau}$, $b_\tau=\sqrt{1-e^{-2\tau}}$ changes the relative speed of contraction and noising. The overshooting bridge $a_t=(1-t)(1-2t)$, $b_t=t$ is not a time reparameterization of either one and produces a more pronounced bending of the reverse trajectories.

One-Step Generative Models¶

One-step generative models try to keep the geometric training principle of flows while removing the expensive multi-step integration at sampling time. The idea is to evolve the model distribution during training, but to store the final evolution in a single generator evaluation.

One-Step Models via Parameter-Domain Discrepancy Flows¶

The most direct one-step strategy is to train the generator parameters themselves by descending a discrepancy between generated and data distributions. Let $\zeta$ be a simple latent distribution, let $f_\theta:\mathcal Z\to\X$ be a neural generator, and let $\beta$ be the data distribution. The objective is to find $\theta$ such that

or, in practice, such that the generated law is close to $\beta$. Given a discrepancy $\mathcal D$ on probability measures, define

This viewpoint includes Wasserstein GANs, MMD-GANs and Sinkhorn generative models Arjovsky et al., 2017Dziugaite et al., 2015Genevay et al., 2018; when $\mathcal D$ is written through a dual potential or discriminator, it connects directly with the adversarial formulation in GANs via Duality. The resulting training dynamics is the ordinary Euclidean gradient flow in parameter space,

This parameter flow induces a flow of generated measures, but this induced flow is not intrinsic in general. Set

If $t\mapsto\theta_t$ is smooth and $f_\theta$ is differentiable with respect to $\theta$, then

For $\alpha_t$-a.e. $x$, let $\eta_{t,x}$ be the disintegration of the latent law $\zeta$ with respect to the map $f_{\theta_t}$. Thus $\eta_{t,x}$ is supported on the fiber $\{z:f_{\theta_t}(z)=x\}$, and for every bounded measurable $\psi$,

The Eulerian velocity of the generated measure is therefore the conditional average

Indeed, for every smooth test function $\varphi$,

which is the weak form of the continuity equation

The velocity (60) depends on the parameter value $\theta_t$ and on the latent disintegration, not only on the measure $\alpha_t$. If the parametrization is non-identifiable, two different parameters can generate the same measure while inducing different velocities. Thus (56) is a Euclidean gradient flow on parameter space, and only after push-forward a model-dependent flow on measures. It coincides with an intrinsic Wasserstein gradient flow only in special cases where $v_t$ agrees with the Wasserstein velocity

More generally, a projected Wasserstein flow on the model manifold would require projecting this intrinsic velocity onto the infinitesimal velocities attainable by variations of $\theta$, using the $L^2(\alpha_t)$ metric. The ordinary Euclidean parameter gradient flow (56) does not perform this projection in general.

In machine learning, however, $\beta$ is accessed through minibatches. If

then the implemented update often replaces $H_\beta$ by the random objective $H_{\hat\beta_n}$. For nonlinear discrepancies this is not, in general, an unbiased gradient estimator:

The expectation is over the data minibatch. This bias is a finite-sample effect of inserting the empirical measure inside a nonlinear transport or adversarial optimization before differentiating; it is distinct from the optimization noise of stochastic gradient descent. MMD losses admit unbiased $U$-statistic variants, but OT and entropic OT objectives generally exhibit the bias--variance phenomena discussed in Bias and Variance of OT.

One-Step Model Using Wasserstein Flow of Discrepancy¶

This construction keeps the discrepancy-minimization viewpoint of the previous paragraph but changes the dynamics. Instead of following the Euclidean gradient of the discrepancy in parameter space, one prescribes the Wasserstein gradient flow of the generated law and then learns a generator update realizing this motion. This distribution-space, natural-gradient-type dynamics is less tied to a particular parametrization and can have better convergence behavior, in particular by reducing parameter-space traps when the discrepancy has a favorable geometry.

Let $\zeta$ be a simple latent distribution and let $\alpha_\theta=(G_\theta)_\sharp\zeta$ be the model distribution. Assume that the target data distribution is $\beta$. The Wasserstein-flow construction chooses a discrepancy

for instance a smoothed $\KL(\alpha|\beta)$, an MMD/IPM loss, or the debiased Sinkhorn divergence $\bar\MK_c^\epsilon(\alpha,\beta)$ introduced in the Sinkhorn divergence section. The associated formal descent is

Instead of integrating (67) at inference time, one fits, at each training time $t$, a parametric residual field $U_{\eta_t}$ along the current model distribution:

In a particle or generator implementation, the learned residual is then used to update the current generator by

This is an ideal function-space update. A genuine one-step implementation fits the updated outputs back into a fixed generator architecture, or distills the accumulated transport into one network. After many training updates, that fixed-architecture or distilled generator is evaluated once at test time. This is the organizing principle behind recent one-step methods based on Wasserstein gradient flows: W-Flow uses such a construction with the Sinkhorn divergence as a tractable global discrepancy Han et al., 2026, while drifting methods evolve the generated distribution during training through a fitted vector field and also admit one-step inference Deng et al., 2026. The gradient-flow interpretation of drifting models, and its relation to KL, MMD, sliced-Wasserstein and Sinkhorn-type discrepancies, is analyzed in Gretton et al., 2026He et al., 2026. These ideas are also connected to the Sinkhorn-type normalization dynamics used to model attention in Sinkformers Sander et al., 2022.

Sliced-Wasserstein Flow¶

A particularly transparent instance uses the sliced objective introduced for imaging and barycenters by Rabin, Peyré, Delon and Bernot Rabin et al., 2011. Take

with $\SW_2$ defined in Sliced Wasserstein Distances. With the continuity-equation convention $\partial_t\alpha_t+\operatorname{div}(\alpha_t w_t)=0$, the descent velocity points from the current projected law toward the target projected law. More precisely, assume that $\alpha$ has a density and sufficient moments so that the projected monotone maps are uniquely defined for almost every direction. If $T_\theta$ denotes the one-dimensional monotone transport from $(P_\theta)_\sharp\alpha$ to $(P_\theta)_\sharp\beta$, then the formal $\Wass_2$-gradient-flow velocity is

The sign follows from the one-dimensional identity: the first variation of $\frac12\Wass_2^2(\rho,\nu)$ has spatial derivative $s-T(s)$, so the descent velocity is $T(s)-s$, and composing with $P_\theta$ lifts it in the direction $\theta$.

Thus empirical implementations only require one-dimensional sorting along sampled directions, followed by averaging the lifted projected displacements. This is the sliced-Wasserstein flow studied by Cozzi and Santambrogio Cozzi & Santambrogio, 2024, who prove long-time convergence under their hypotheses when the target is Gaussian and also show that the limiting characteristic map is not, in general, the optimal transport map. When both the evolving law and the target are Gaussian, the averaged velocity is affine and the flow closes on means and covariances; this finite-dimensional closure is revisited in Flows over the Gaussian Manifold, where the sliced objective appears in the Gaussian closure catalogue of Proposition: Gaussian closure catalogue.

The difference between the full $\Wass_2$ descent and its sliced analogue is already visible on a simple shape example. Div compares the exact empirical flow of $\frac12\Wass_2^2(\cdot,\beta)$, where one fixed optimal assignment gives straight relaxation curves, with the $\Wass_2$-gradient flow of $\frac12\SW_2^2(\cdot,\beta)$, where the velocity is recomputed from projected monotone rearrangements.

show_book_figure("generative-w2-vs-sliced-flow-shapes", width=850)

Full Wasserstein and sliced-Wasserstein flows from a cat-shaped empirical law to a heart-shaped target. The left panel in each row shows representative particle trajectories colored from red to blue. The five panels on the right render kernel-density estimates of all particles at common normalized times. The $\Wass_2$ flow follows the fixed optimal assignment and therefore straight relaxation rays, whereas the sliced flow averages one-dimensional sorted rearrangements over projection directions, producing curved trajectories and a different transient density evolution.

Stein Variational Gradient Descent¶

Stein variational gradient descent (SVGD) is another deterministic particle flow that fits naturally in this one-step viewpoint Liu & Wang, 2016. Its original motivation is Bayesian sampling: given a target probability $\beta=\rho_\beta\,\d x$ known through its score $\nabla\log\rho_\beta=-\nabla V$, but not necessarily through its normalizing constant, drive a particle cloud toward $\beta$ without estimating the score of the current empirical law. Geometrically, this replaces the Wasserstein gradient flow of $\KL(\alpha|\beta)$, whose tangent norm is $L^2(\alpha)$, by the kernelized Benamou--Brenier geometry of Kernelized Benamou--Brenier Distances, whose velocities live in a vector-valued RKHS.

For the formal density-level calculation, assume $\alpha=\rho_\alpha\,\d x$ and $\beta=\rho_\beta\,\d x$ have smooth positive densities, and let $v$ be a smooth compactly supported vector field. For the perturbation $\alpha_\epsilon=(\Id+\epsilon v)_\sharp\alpha$, integration by parts gives

The bracket is the Langevin--Stein operator applied to $v$ and averaged under $\alpha$; because it only evaluates $v$, $\operatorname{div}v$, and the target score at sample locations, it remains meaningful when $\alpha$ is empirical. Optimizing this linear functional over the unit ball of $\RKHS_k^d$ and using the reproducing property gives the RKHS steepest-descent direction,

The associated mean-field equation is

For particles $\alpha^\ell_n=n^{-1}\sum_i\delta_{X_i^\ell}$, this becomes

The first term attracts particles toward increasing target log-density; the second term is a kernel repulsion that prevents immediate collapse. The gradient-flow interpretation and its many-particle limits are studied in Liu, 2017Duncan et al., 2023Nüsken & Renger, 2023. In generative modeling terms, SVGD transports a simple empirical latent law by repeated smooth particle updates. It is therefore close in spirit to the drifting fields above, but its velocity is not learned by regression: it is the closed-form RKHS steepest-descent direction of the KL functional.

The figure below contrasts this RKHS flow with a particle closure of the Wasserstein gradient flow of $\KL(\alpha\mid\beta)$. The latter evaluates the current score $\nabla\log\rho_\alpha$ by a Gaussian KDE, while SVGD avoids this density estimate and uses the target score together with a kernel repulsion.

Figure Div contrasts this RKHS flow with a particle closure of the Wasserstein gradient flow of $\KL(\alpha|\beta)$.

show_book_figure("generative-w2-vs-svgd-entropy-flow", width=760)

Particle trajectories for two deterministic descents of relative entropy toward the same three-Gaussian target, whose density is shown by gray contours. The left panel approximates the $\Wass_2$ gradient flow of $\KL(\alpha\mid\beta)$ by the KDE velocity $\nabla\log\rho_\beta-\nabla\log\rho_{\alpha,h}$. The right panel uses the RBF-kernel SVGD velocity (73), where target-score attraction is coupled with RKHS repulsion. Both panels use the same initial particles and color trajectory time from red to blue.

Self-corrected drifting fields.¶

Drifting methods need not start from an exact Wasserstein gradient. They often prescribe an attraction-minus-repulsion field and then regress this field in $L^2(\alpha_t)$. A simple continuous version uses a strictly positive kernel $K_\epsilon(x,y)$ for which the following integrals are finite, and defines, for any probability measure $\nu$,

For the Gaussian kernel $K_\epsilon(x,y)=\exp(-\norm{x-y}^2/(2\epsilon))$, this normalized field is a score of a smoothed density:

The drifting velocity is then

The first term pulls samples toward data, while the second term corrects self-attraction and prevents all particles from collapsing onto the same high-density region. For a fixed reference measure, $B_\epsilon[\nu]$ is precisely the Gaussian mean-shift displacement in (112): it moves $x$ toward the local kernel barycenter of $\nu$. Hence self-corrected drifting can be read as the difference between a target mean-shift field and the current model’s own mean-shift field. Sinkhorn drifting replaces these one-sided kernel normalizations by two-sided entropic OT couplings, so that the cross and self terms are normalized by Sinkhorn scaling rather than by a single denominator He et al., 2026.

Figure Div illustrates why the self term matters: it corrects the collapse created by target attraction alone and improves coverage of separated modes.

show_book_figure("generative-drifting-model-trajectories", width=760)

Drifting trajectories for a small particle generator. The raw kernel drift has weak long-range attraction and can leave particles away from the data modes. The self-corrected field uses the difference $B_\epsilon[\beta]-B_\epsilon[\alpha_t]$, so a longer integration brings particles to the blue modes while repelling them from their own current concentration.

Proof

Since $\al_t$ and $\phi_t$ are fixed in the variation with respect to $\alpha$, the first variation is

$$\delta_\alpha \mathcal R_t(\alpha|\al_t)(x)=-\phi_t(x).$$

(81)

By Proposition Proposition: Formal Wasserstein Gradient,

$$\Wgrad \mathcal R_t(\alpha|\al_t) = \nabla\delta_\alpha \mathcal R_t(\alpha|\al_t) = -\nabla\phi_t = -u_t.$$