What if you could generate high-quality images in one step instead of hundreds?

Stanford and ByteDance introduce W-Flow: a single-step generator that turns random noise into target data by following a Wasserstein gradient flow. It compresses a full diffusion-like evolution into a single neural network pass.

The result? A new state-of-the-art for one-step ImageNet 256x256 generation (1.29 FID), with better coverage of rare modes — and 100x faster sampling than multi-step diffusion models of similar quality.

One-Step Generative Modeling via Wasserstein Gradient Flows

Paper: https://arxiv.org/abs/2605.11755

Our report: https://mp.weixin.qq.com/s/3EFvWPBBLELPReQ3I0y0YQ…

#PapersAccepted by Jiqizhixin

One-Step Generative Modeling via Wasserstein Gradient Flows

Source: https://arxiv.org/html/2605.11755 Jiaqi Han111footnotemark:1&Puheng Li1&Qiushan Guo2Renyuan Xu1222Here we omit the second-batchq^′\widehat{q}^{\prime}in the subscript ofVVfor brevity.&Stefano Ermon1222Here we omit the second-batchq^′\widehat{q}^{\prime}in the subscript ofVVfor brevity.&Emmanuel J. Candès1 1Stanford University2ByteDance

Project Page:https://hanjq17.github.io/W-Flow/

Abstract

Diffusion models and flow-based methods have shown impressive generative capability, especially for images, but their sampling is expensive because it requires many iterative updates. We introduceW-Flow, a framework for training a generator that transforms samples from a simple reference distribution into samples from a target data distribution in a single step. This is achieved in two steps: we first define an evolution from the reference distribution to the target distribution through a Wasserstein gradient flow that minimizes an energy functional; second, we train a static neural generator to compress this evolution into one-step generation. We instantiate the energy functional with the Sinkhorn divergence, which yields an efficient optimal-transport-based update rule that captures global distributional discrepancy and improves coverage of the target distribution. We further prove that the finite-sample training dynamics converge to the continuous-time distributional dynamics under suitable assumptions.Empirically, W-Flow sets a new state of the art for one-step ImageNet 256×\times256 generation,achieving1.29 FID, with improved mode coverage and domain transfer. Compared to multi-step diffusion models with similar FID scores, our method yields approximately 100×\timesfaster sampling. These resultsshowthat Wasserstein gradient flows provide a principled and effective foundation for fast and high-fidelity generative modeling.

Figure 1:*(Left)1-NFE samplesfrom W-Flow-L/2 trained from scratch on ImageNet-256×\times256.(Right)*Sample quality (measured by FID) vs. effective sampling computeLu and Song (2024)(billion parameters×\timesnumber of function evaluations during sampling) evaluated on ImageNet 256×\times256.### 1Introduction

Generative modeling can be viewed as transporting a simple reference distribution (often referred to as a prior) to a complex target distribution (e.g., data distribution). In modern high-dimensional settings, this transport is rarely learnedfrom scratchas a single global map. Instead, the dominant paradigm constructs it through a sequence of simpler updates.Diffusion modelsSohl-Dicksteinet al.(2015); Songet al.(2020); Song and Ermon (2019); Hoet al.(2020)learn a sequence of ‘infinitesimal’ denoising steps and then unroll them overmanyinference steps. Flow modelsLipmanet al.(2022); Liuet al.(2022)learn a sequence of velocity fields and likewise unroll the dynamics over many steps.Despite high sample quality,multi-step inferenceleads to significant latency and computational cost.

It would be convenient to use an iterative procedureduring trainingto specify the evolution of the generated data distribution, and then compress this evolution into a static map. In this way, the resulting generator would directly transform an input from the reference distribution into an output from the target distribution in one step. This would combine the efficiency of one-step generation with the flexibility of a distributional evolution during training. The fundamental question would thus be:

What principles should govern the training dynamics of one-step generators?

Existing approaches fall short of providing a satisfactory answer. Generative Adversarial Networks (GANs)Goodfellowet al.(2014); Radfordet al.(2015); Arjovskyet al.(2017)evolve the model via discriminator-induced gradients, yet inherit the well-known instability of minimax optimization. More recently, Drifting ModelDenget al.(2026)achieves strong empirical performance by iteratively displacing particles111A particle is one candidate sample/image as it evolves through the algorithm.via handcrafted attractive and repulsive fields. Yet, the underlying dynamics remain heuristic and lack a clear and principled interpretation, leaving the model vulnerable to unpredictable convergence and mode collapse (see Fig.6).

In this paper, we propose W-Flow, a new framework for one-step generative modeling. W-Flow prescribes how theentire model distributionshould evolve toward the target distribution. This evolution is given by a Wasserstein gradient flow (WGF): at each training step, the current generated distribution moves along the steepest descent direction of a chosen energy functional to the target distribution. We instantiate this energy functional with the Sinkhorn divergence, an optimal-transport-based distance that can be estimated efficiently from mini-batches.The resulting Sinkhorn update combines two batch-level transport plans: a generated-to-real transport plan and a generated-to-generated self-transport plan. This makes each update depend on the batch-level transport structure, rather than on a single matched sample or nearest neighbor. We then train a generator to imitate these updates, thereby compressing the multi-step training evolution into a single one-step map at inference time. This fundamentally differs from few/one-step diffusion methodsSonget al.(2023); Zhouet al.(2025), which typically distill or follow a path originally designed for iterative reverse diffusion.

Notably, W-Flow achieves a new state of the art for one-step generation on ImageNet 256×\times256, obtaining 1.29 FID at XL scale and 1.35 at L scale, significantly outperformingDenget al.(2026)while also surpassing many multi-step diffusion and flow models (see Fig.1). Beyond class-conditional generation, experiments on FFHQ further demonstrate improved mode coverage and strong domain transfer capability, enabled by the globally coordinated OT dynamics. Overall, W-Flow establishes a new framework for designing the training dynamics of generators, showing that WGFs can serve as a powerful foundation for fast and high-fidelity generative modeling.

2Related work

(Few/one-step) diffusion and flow-based models.DiffusionSohl-Dicksteinet al.(2015); Songet al.(2020); Hoet al.(2020)and flow modelsLipmanet al.(2022); Liuet al.(2022)learn to transport asimple reference distributionto acomplex target distributionthrough a forward noising process and its time-reverse. However, they require many sequential steps for sampling, leading to substantial computational cost and latency. Many works distill multi-step teachers into few/one-step models via variational score distillationYinet al.(2024b,a); Wanget al.(2023); Salimanset al.(2024)or learning shortcuts along the diffusion ODE specified by the teacherSonget al.(2023); Song and Dhariwal (2023); Lu and Song (2024); Salimans and Ho (2022). Efforts have also been made to directly learn a few/one-step generator from scratch, typically by enforcing certain self-consistency conditions on the trajectoryGenget al.(2025a,b); Boffiet al.(2025); Songet al.(2023)or the intermediate marginalsZhouet al.(2025). These methods largely inherit their training signal from a predefined diffusion transport path, which is originally designed for iterative generation. We instead directly define the training-time evolution of a one-step generator via WGFs.

One-step generator mappings.An alternative is to directly learn a static map that takes an input from the reference distribution and outputs a sample from the target distribution.The most studied is the family of GANsGoodfellowet al.(2014); Arjovskyet al.(2017); Radfordet al.(2015); Genevayet al.(2018), which train the generator using discriminator-driven gradients in a minimax game. However, adversarial training is often unstable and requires considerable effort to stabilize. Recently, Drifting ModelDenget al.(2026)guides the generator via an empirically constructed drifting field using kernel interactions and achieves strong performance(Liet al.,2026; Gaoet al.,2026). Yet, the design of this technique remains largely heuristic, lacking a principled understanding of the induced dynamics. Subsequent works(Laiet al.,2026; Turan and Ovsjanikov,2026)connect Drifting Model to score-based formulations under Gaussian kernels, while(Heet al.,2026; Caoet al.,2026)explore links to Sinkhorn divergence and WGFs.In particular, the concurrent workHeet al.(2026)studies a closely related Sinkhorn-drifting field and interprets Drifting Model as a Sinkhorn-divergence WGF approximation. Our focus is different: we prove convergence of the induced particle dynamics, propose large-scale training framework, and scale the method to ImageNet 256×\times256.

Optimal transport in generative models.Existing applications of OT in generative models broadly fall into two categories. The first uses OT-inspired discrepancies as the training objectives. Wasserstein GANsArjovskyet al.(2017); Gulrajaniet al.(2017); Tolstikhinet al.(2017)optimize a dual form of the Wasserstein distance via an adversarial critic, while Sinkhorn GANsGenevayet al.(2018)leverage OT-based divergence computed by Sinkhorn iterations. The second employs OT to guide dynamic transport paths, including Schrödinger bridgesAlouadiet al.(2026); De Bortoliet al.(2021); Maet al.(2025)and OT-inspired flow models(Albergo and Vanden-Eijnden,2022; Lipmanet al.,2022; Liuet al.,2022; Linet al.,2025). Yet, static OT losses require adversarial approximation or differentiation through iterative solvers, while dynamic transport methods still rely on multi-step sampling.

Wasserstein gradient flows in generative models.WGFs provide a principled description of the evolution of probability measures. A standard discretization is the Jordan–Kinderlehrer–Otto (JKO) schemeJordanet al.(1998), which has inspired several recent approaches and theoretical frameworks(Mokrovet al.,2021; Choiet al.,2024; Fanet al.,2021; Xie and Cheng,2025). However, JKO-based methods require solving a complex inner optimization problem at each step, often involving adversarial training over dual potentials, which makes them difficult to scale in practice.

3Method

Figure 2:(a)The conceptual diagram of W-Flow.(b)Visualization of the training dynamics projected onto the Sinkhorn divergence landscape on 8 Gaussian mixtures, shown on a logarithmic scale.We introduce W-Flow, adeterministicone-stepgeneratorwith its training dynamics guided by a WGF. The main idea is to separate the design of training dynamics from the parameterization of the generator: we first prescribe how the model distribution should evolve toward the data distribution during training, and then learn a neural network to realize this evolution through particle updates.

3.1W-Flow: Learningone-step transport mapunder Wasserstein gradient flow

One-step transport maps.Given a reference distributionpref∈𝒫​(ℝm)p_{\mathrm{ref}}\in{\mathcal{P}}({\mathbb{R}}^{m})and a target distributionp∈𝒫​(ℝn)p\in{\mathcal{P}}({\mathbb{R}}^{n}), the goal is to learn a mappingf:ℝm→ℝnf:{\mathbb{R}}^{m}\to{\mathbb{R}}^{n}such that if an inputzzis sampled from the reference distribution, the outputf​(z)f(z)follows the target distributionpp. In distribution space, this is denoted by the pushforward operatorf#f_{\#}, meaning we aim to learn a mapffsuch thatf#​pref≈pf_{\#}p_{\mathrm{ref}}\approx p. In practice,ppis only observed through a finite dataset of samples.

We adopt an incremental construction offf, which breaks this complex transport into a sequence of simpler steps. Specifically, letf(0):ℝm→ℝnf^{(0)}:\mathbb{R}^{m}\rightarrow\mathbb{R}^{n}be an initial map and setq(0)=f#(0)​prefq^{(0)}=f^{(0)}_{\#}p_{\mathrm{ref}}. We then build the generator through a sequence of local transport maps{g(k)}k=0K−1:ℝn→ℝn\{g^{(k)}\}_{k=0}^{K-1}:\mathbb{R}^{n}\rightarrow\mathbb{R}^{n}satisfying

q(k+1)=g#(k)​q(k),f(k+1)=g(k)∘f(k)=g(k)∘⋯∘g(0)∘f(0),\displaystyle q^{(k+1)}=g^{(k)}_{\#}q^{(k)},\qquad f^{(k+1)}=g^{(k)}\circ f^{(k)}=g^{(k)}\circ\cdots\circ g^{(0)}\circ f^{(0)},(1)so that the final mapf(K)f^{(K)}pushesprefp_{\mathrm{ref}}toq(K)≈pq^{(K)}\approx p(see Fig.2(a)). Explicitly,f(K)f^{(K)}is astaticmap obtained by composing stepwise transport maps{g(k)}k=0K−1\{g^{(k)}\}_{k=0}^{K-1}and the initialf(0)f^{(0)}. The role of this sequence isnotto define a standalone multi-step generative model, but rather to provide an incremental procedure for constructing the target one-step mapf(K)f^{(K)}used at inference time.

Velocity field.We parameterize the mapsg(k)g^{(k)}to reflect incrementaldisplacementsbetween iterations:

g(k)​(x)≔x+η​V(k)​(x),\displaystyle g^{(k)}(x)\coloneqq x+\eta V^{(k)}(x),(2)whereV(k):ℝn→ℝnV^{(k)}:{\mathbb{R}}^{n}\to{\mathbb{R}}^{n}is a continuous velocity field, also referred to as a drifting fieldDenget al.(2026), andη>0\eta>0is the step size. Under this parameterization, each iteration movesa data samplex∼q(k)x\sim q^{(k)}only locally, while the composition of these maps yields a flexible global transport, and the fieldV(k)V^{(k)}determines how the intermediate distributions evolve.

In principle, the recursion in (1) would make the generator an increasingly long composition, so exact sampling fromf#(k)​preff^{(k)}_{\#}p_{\mathrm{ref}}becomes prohibitively costly whenkkis large. To keep the model one-step throughout training, we represent the generator by a neural network and learnfθ(k+1)f_{\theta^{(k+1)}}to match the updated mapf(k+1)f^{(k+1)}. Usingfθ(k)f_{\theta^{(k)}}as a surrogate forf(k)f^{(k)}, we optimize

ℒ(k)​(θ)≔𝔼z∼pref​‖fθ​(z)−g(k)∘f~(k)​(z)‖2,f~(k)​(z)≔sg​(fθ(k)​(z)),\displaystyle{\mathcal{L}}^{(k)}(\theta)\coloneqq\mathbb{E}_{z\sim p_{\mathrm{ref}}}\left\|f_{\theta}(z)-g^{(k)}\circ\tilde{f}^{(k)}(z)\right\|^{2},\qquad\tilde{f}^{(k)}(z)\coloneqq\mathrm{sg}\!\left(f_{\theta^{(k)}}(z)\right),(3)wheresg​(⋅)\mathrm{sg}(\cdot)denotes stop-gradient. The parameterθ(k+1)\theta^{(k+1)}is then obtained by taking a gradient step in (3). Atθ=θ(k)\theta=\theta^{(k)}, the loss scale is𝔼z∼pref​‖η​V(k)​(fθ(k)​(z))‖2\mathbb{E}_{z\sim p_{\mathrm{ref}}}\left\|\eta V^{(k)}(f_{\theta^{(k)}}(z))\right\|^{2}, henceℒ(k)=0{\mathcal{L}}^{(k)}=0impliesV(k)=0V^{(k)}=0.

A central challenge, however, is to design the vector fieldV(k)V^{(k)}to guide the iteration in (1) towards the target distribution. Existing approaches, such asDenget al.(2026), are largely heuristic and guarantee only that the target distribution is a stationary point of the dynamics:V(k)=0V^{(k)}=0ifq(k)=pq^{(k)}=p. However, the converse generally fails. In particular, during training, observing thatV(k)→0V^{(k)}\to 0does not imply thatq(k)→pq^{(k)}\to p. Moreover, the intermediate distributional dynamics{q(k)}\{q^{(k)}\}induced by such velocity fields remain largely uncharacterized, and convergence guarantees are generally unavailable.

Continuous-time Wasserstein velocity field.To obtain a principled and well-behaved velocity fieldV(k)V^{(k)}in Eq. (2), we take a step back and first specify the desired evolution of the model distribution in continuous time, and then derive the corresponding discrete updates. Specifically, we model the evolution of{q(k)}\{q^{(k)}\}via a WGF. Letℱ:𝒫​(ℝn)→ℝ{\mathcal{F}}:{\mathcal{P}}({\mathbb{R}}^{n})\to\mathbb{R}be an energy functional defined on probability distributions. The associated WGF is given by

∂tqt=∇⋅(qt​∇δ​ℱδ​q​(qt)),\displaystyle\partial_{t}q_{t}=\nabla\cdot\!\left(q_{t}\nabla\frac{\delta{\mathcal{F}}}{\delta q}(q_{t})\right),(4)whereδ​ℱδ​q​(qt)\frac{\delta\mathcal{F}}{\delta q}(q_{t})denotes the first variation ofℱ\mathcal{F}atqtq_{t}. Eq. (4) defines acontinuousevolution of distributionsqt∈𝒫​(ℝn)q_{t}\in{\mathcal{P}}({\mathbb{R}}^{n})that decreases the energyℱ{\mathcal{F}}. Equivalently, this evolution can be expressed through the continuity equation with a specific form of the velocity field

∂tqt+∇⋅(qt​Vt)=0,Vt​(x)=−∇δ​ℱδ​q​(qt)​(x).\displaystyle\partial_{t}q_{t}+\nabla\cdot(q_{t}V_{t})=0,\quad V_{t}(x)=-\nabla\frac{\delta{\mathcal{F}}}{\delta q}(q_{t})(x).(5)This characterization provides a clear variational interpretation: the distributionqtq_{t}evolves along thesteepest descentdirection ofℱ{\mathcal{F}}in Wasserstein space (see Fig.2(b) for an illustration). To drive the dynamics toward the target distributionpp, we choose the energy functionalℱ​(q)=𝒟​(q∥p){\mathcal{F}}(q)={\mathcal{D}}(q\,\|\,p)where𝒟{\mathcal{D}}is a suitable divergence such thatp=arg⁡minq⁡ℱ​(q)p=\arg\min_{q}{\mathcal{F}}(q). Under this choice, the induced velocity field in (5) transportsqtq_{t}to progressively reduce its discrepancy frompp.

Time discretization.Discretizing the continuous-time evolution provides a concrete prescription for the stepwise transport maps introduced earlier. Applying an explicit Euler scheme to the WGF over time steps{t(k):t(k)=k​η,k∈ℕ≥0}\{t^{(k)}:t^{(k)}=k\eta,\;k\in{\mathbb{N}}_{\geq 0}\}yields

g(k)​(x)≔x+η​V(k)​(x),V(k)​(x)=−∇δ​ℱδ​q​(q(k))​(x),\displaystyle g^{(k)}(x)\coloneqq x+\eta V^{(k)}(x),\qquad V^{(k)}(x)=-\nabla\frac{\delta{\mathcal{F}}}{\delta q}(q^{(k)})(x),(6)whereq(k)=qt(k)q^{(k)}=q_{t^{(k)}}. This recovers exactly the stepwise velocity-based update in Eq. (2), now with the velocity fieldV(k)V^{(k)}specified by the underlying gradient flow.Each mapg(k)g^{(k)}operates directly on individual data points inℝn\mathbb{R}^{n}. By moving these points, the resulting empirical distributionapproximates the continuous transport trajectory in distribution space, while composing these local updates yields a mapf(k)f^{(k)}that approximates the flow at timet(k)t^{(k)}. In sum, the gradient flow defines thetraining-time dynamicsin distribution space, whereas the learned model is ultimately astatic mapobtained by compressing these local transports into a single generator. The diagram is provided in Fig.2(a).

3.2Instantiation via Sinkhorn divergence

Choices of the energy functionalℱ\mathcal{F}.The choice ofℱ\mathcal{F}is critical, as it determines both the geometry of the WGF and the tractability of the induced velocity field. Ideally,ℱ\mathcal{F}should satisfy two properties: (i) its minimizer coincides with the target distributionpp, and (ii) the induced flow is well-behaved and can be implemented efficiently in high dimensions.

Table 1:Comparison of energy functionals and their induced continuous velocity fields under Wasserstein gradient flow. Detailed derivations are in AppendixA.1.Common choices include Maximum Mean Discrepancy (MMD)Liet al.(2015)and KL divergence (see Table1for the corresponding velocity fields). However, both face significant limitations in practice. MMD induces a kernel-based interaction between thesampled data points.Whenqtq_{t}is far frompp, the kernel values can become small, causing the gradients to vanish and the process to stall.The KL divergence involves the score function∇log⁡qt​(x)\nabla\log q_{t}(x), which is generally intractable and must be approximated, introducing additional bias and instability. Our ablations in Table2aconfirm that both choices lead to inferior performance.

These limitations motivate us to seek an alternative energy functional that is computationally tractable and geometrically well-behaved. To this end, we adopt the Sinkhorn divergence, an OT-based discrepancy that avoids explicit score estimation and admits a natural and direct implementation on data samples. While the convergence of WGFs to a global minimizer is not automatic, a recent result inHardion and Lacombe (2026)shows that, under Gaussian distributional assumptions, the WGF of the Sinkhorn divergence converges to its global minimizer.

Entropic optimal transport and Sinkhorn divergence.A natural candidate is the entropically regularized OT cost, which is computationally attractive due to Sinkhorn iterations. However, it is biased: in general,OTε​(p,p)≠0\mathrm{OT}_{\varepsilon}(p,p)\neq 0, soOTε​(⋅,p)\mathrm{OT}_{\varepsilon}(\cdot,p)does not define an energy whose minimizer is exactly the target distribution. This motivates the use of thedebiasedSinkhorn divergence.Concurrently,Heet al.(2026)also employs Sinkhorn divergence; our implementation incorporates the two-batch self-transport and velocity guidance, which are essential to achieving strong empirical performance in large-scale training (Sec.3.3and Sec.3.4).

We start by introducing some general notations. Let𝒫2​(ℝn)\mathcal{P}_{2}(\mathbb{R}^{n})denote the set of probability measures onℝn\mathbb{R}^{n}with finite second moment. Takeq,p∈𝒫2​(ℝn)q,p\in\mathcal{P}_{2}(\mathbb{R}^{n}). Define the set of couplingsΠ​(q,p):={π∈𝒫​(ℝn×ℝn):π​(⋅,ℝn)=q,π​(ℝn,⋅)=p},\Pi(q,p):=\big\{\pi\in\mathcal{P}(\mathbb{R}^{n}\times\mathbb{R}^{n}):\ \pi(\cdot,\mathbb{R}^{n})=q,\ \pi(\mathbb{R}^{n},\cdot)=p\big\},which is contained in𝒫2​(ℝn×ℝn)\mathcal{P}_{2}(\mathbb{R}^{n}\times\mathbb{R}^{n})as both marginals have finite second moments. We then define the entropic optimal transport (EOT) functional under quadratic cost

OTε​(q,p):=infπ∈Π​(q,p){∫12​‖x−y‖2​π​(d​x,d​y)+ε​DKL​(π∥q⊗p)},\mathrm{OT}_{\varepsilon}(q,p):=\inf_{\pi\in\Pi(q,p)}\left\{\int\tfrac{1}{2}\|x-y\|^{2}\,\pi(dx,dy)+\varepsilon\,D_{\mathrm{KL}}(\pi\,\|\,q\otimes p)\right\},(7)whereε>0\varepsilon>0is the entropic regularization parameter. The associated Sinkhorn divergence is given by

Sε​(q,p)=OTε​(q,p)−12​OTε​(q,q)−12​OTε​(p,p).S_{\varepsilon}(q,p)=\mathrm{OT}_{\varepsilon}(q,p)-\tfrac{1}{2}\mathrm{OT}_{\varepsilon}(q,q)-\tfrac{1}{2}\mathrm{OT}_{\varepsilon}(p,p).(8) Velocity field induced by Sinkhorn divergence.Letπq,pε∈Π​(q,p)\pi_{q,p}^{\varepsilon}\in\Pi(q,p)andπq,qε∈Π​(q,q)\pi_{q,q}^{\varepsilon}\in\Pi(q,q)be the optimal couplings forOTε​(q,p)\mathrm{OT}_{\varepsilon}(q,p)andOTε​(q,q)\mathrm{OT}_{\varepsilon}(q,q), respectively. Expressing these joint couplings in terms of their conditional distributions asπq,pε​(d​x,d​y)=q​(d​x)​πq,pε​(d​y∣x)\pi_{q,p}^{\varepsilon}(dx,dy)=q(dx)\,\pi_{q,p}^{\varepsilon}(dy\mid x)andπq,qε​(d​x,d​x′)=q​(d​x)​πq,qε​(d​x′∣x),\pi_{q,q}^{\varepsilon}(dx,dx^{\prime})=q(dx)\,\pi_{q,q}^{\varepsilon}(dx^{\prime}\mid x),we define the barycentric projections

Tq,pε​(x):=∫y​πq,pε​(d​y∣x),Tq,qε​(x):=∫x′​πq,qε​(d​x′∣x).T_{q,p}^{\varepsilon}(x):=\int y\,\pi_{q,p}^{\varepsilon}(dy\mid x),\qquad T_{q,q}^{\varepsilon}(x):=\int x^{\prime}\,\pi_{q,q}^{\varepsilon}(dx^{\prime}\mid x).(9)It is well known (see,e.g.,Genevayet al.(2018); Peyré and Cuturi (2019); Nutz and Wiesel (2022)and the proof in AppendixA.1) that the velocity field coincides with the Wasserstein gradient direction of the Sinkhorn divergence

Vq,pε​(x)=−∇δ​Sε​(⋅,p)δ​q​(q)​(x)=Tq,pε​(x)−Tq,qε​(x).V_{q,p}^{\varepsilon}(x)=-\nabla\frac{\delta S_{\varepsilon}(\cdot,p)}{\delta q}(q)(x)=T_{q,p}^{\varepsilon}(x)-T_{q,q}^{\varepsilon}(x).(10)Characterization of equilibrium.The velocity field vanishesVq,pε≡0V_{q,p}^{\varepsilon}\equiv 0if and only ifq=pq=p(seeFeydyet al.(2019)). In particular, the Sinkhorn divergence admits no spurious stationary points: the only equilibrium is the target distribution. This is a notable advantage over local drifting dynamics, where the interaction field can vanish at someq≠pq\neq p, causing the dynamics to stall before reaching the target.

Interpretation of the velocity field.The velocity field in Eq. (10) admits a natural interpretation as a balance between attraction and self-interaction. The first termTq,pε​(x)T_{q,p}^{\varepsilon}(x)transports the current distribution toward the target distribution, while the correction termTq,qε​(x)T_{q,q}^{\varepsilon}(x)enforces consistency with the current distribution through a global mass-conserving coupling. This structure ensures that the resulting dynamics are globally coordinated, in contrast to purely local interaction rules, and plays a key role in stabilizing the transport and preventing mode collapse (see Sec.4.3for experimental illustrations).

3.3Interactive particle dynamics

In practice, we operate on finite mini-batches. GivenNNi.i.d. samples{xi}i=1N∈ℝn\{x_{i}\}_{i=1}^{N}\in\mathbb{R}^{n}fromqtq_{t}andMMi.i.d. samples{yj}j=1M∈ℝn\{y_{j}\}_{j=1}^{M}\in\mathbb{R}^{n}frompp, let the weighted empirical measures beq^t=∑i=1Nai​δxi\widehat{q}_{t}=\sum_{i=1}^{N}a_{i}\,\delta_{x_{i}},p^=∑j=1Mbj​δyj,\widehat{p}=\sum_{j=1}^{M}b_{j}\,\delta_{y_{j}},withai,bj>0,a_{i},b_{j}>0,and∑iai=∑jbj=1\sum_{i}a_{i}=\sum_{j}b_{j}=1. In this discrete setting, the general definition in Eq. (7) reduces to the finite-dimensional problem over the coupling matrixπ∈Π​(q^t,p^)\pi\in\Pi(\widehat{q}_{t},\widehat{p}):

OTε​(q^t,p^):=minπ∈Π​(q^t,p^)​∑i,j12​‖xi−yj‖2​πi​j+ε​∑i,jπi​j​(log⁡πi​j−1),\mathrm{OT}_{\varepsilon}\big(\widehat{q}_{t},\widehat{p}\big):=\min_{\pi\in\Pi(\widehat{q}_{t},\widehat{p})}\ \sum_{i,j}\frac{1}{2}\|x_{i}-y_{j}\|^{2}\pi_{ij}+\varepsilon\sum_{i,j}\pi_{ij}(\log\pi_{ij}-1),(11) Figure 3:Comparison between one-batch and two-batch estimators on learning a 2D Gaussian.whereΠ​(q^t,p^)\Pi(\widehat{q}_{t},\widehat{p})is the set of matrices with prescribed marginals. Denote the optimal solutionπq^t,p^ε,∗\pi^{\varepsilon,*}_{\widehat{q}_{t},\widehat{p}}.

Two-batch estimate for self-transport.Naïvely estimating the self-entropic OT termOTε​(q^t,q^t){\rm OT}_{\varepsilon}(\widehat{q}_{t},\widehat{q}_{t})from a single empirical batch introduces a self-matching artifact: since each particle can be matched to itself at zero transport cost, the resulting coupling tends to overemphasize the diagonal entriesπi​i\pi_{ii}. To mitigate it, we adopt atwo-batchstrategy. We draw an independent second batch{xl′}l=1N\{x^{\prime}_{l}\}_{l=1}^{N}fromqtq_{t}and define the empirical measureq^t′=∑l=1Nal​δxl′\widehat{q}^{\prime}_{t}=\sum_{l=1}^{N}a_{l}\delta_{x^{\prime}_{l}}. We then compute the coupling betweenq^t\widehat{q}_{t}andq^t′\widehat{q}^{\prime}_{t}, yielding the optimal planπq^t,q^t′ε,∗\pi^{\varepsilon,*}_{\widehat{q}_{t},\widehat{q}^{\prime}_{t}}. This construction removes the need for the diagonal masking heuristic used inDenget al.(2026)and admits a cleaner theoretical interpretation. Fig.3compares the two-batch strategy with the naïve one-batch estimate on learning a 2D Gaussian. The one-batch strategy fails to capture the Gaussian tails, whereas the two-batch strategy recovers the target faithfully. Additional illustrations are in AppendixA.2.

The discrete counterparts of the barycentric maps in (9) are, therefore,

Tq^t,p^ε​(xi)=1ai​∑j(πq^t,p^ε,∗)i​j​yj,Tq^t,q^t′ε​(xi)=1ai​∑l(πq^t,q^t′ε,∗)i​l​xl′,T_{\widehat{q}_{t},\widehat{p}}^{\varepsilon}(x_{i})=\frac{1}{a_{i}}\sum_{j}(\pi_{\widehat{q}_{t},\widehat{p}}^{\varepsilon,*})_{ij}y_{j},\quad T_{\widehat{q}_{t},\widehat{q}^{\prime}_{t}}^{\varepsilon}(x_{i})=\frac{1}{a_{i}}\sum_{l}(\pi_{\widehat{q}_{t},\widehat{q}^{\prime}_{t}}^{\varepsilon,*})_{il}x^{\prime}_{l},(12)where in practice the optimal couplingsπq^t,p^ε,∗\pi_{\widehat{q}_{t},\widehat{p}}^{\varepsilon,*}andπq^t,q^t′ε,∗\pi_{\widehat{q}_{t},\widehat{q}^{\prime}_{t}}^{\varepsilon,*}can be efficiently approximated using the Sinkhorn-Knopp algorithmSinkhorn and Knopp (1967). We present the subroutine in Alg.2. With these, the particle-based estimate of the velocity field in Eq. (10) is

Vq^t,p^ε​(xi)=Tq^t,p^ε​(xi)−Tq^t,q^t′ε​(xi).\displaystyle V^{\varepsilon}_{\widehat{q}_{t},\widehat{p}}(x_{i})=T^{\varepsilon}_{\widehat{q}_{t},\widehat{p}}(x_{i})-T^{\varepsilon}_{\widehat{q}_{t},\widehat{q}^{\prime}_{t}}(x_{i}).(13) Convergence of the particle dynamics.Now thatVq,pεV^{\varepsilon}_{q,p}andVq^,p^εV^{\varepsilon}_{\widehat{q},\widehat{p}}are defined, we present a convergence guarantee from the empirical dynamics to the population dynamics. This theorem shows that the finite-particle, discrete-time transport mechanism used in our method is a consistent approximation of the intended population-level WGF. As a result, the training dynamics induced by our algorithm faithfully track the continuous distributional evolution toward the target, thereby giving a principled justification for learning the generator through these sequential local transports. See Fig.2(b) for an illustration. The convergence result holds for more general velocity fields, but we present below using the notation of Sinkhorn velocity. The full statement, assumptions, and proof are deferred to AppendixB.

Theorem 3.1(Informal).

Letp^M\widehat{p}^{M}be the empirical measure based onMMi.i.d. samples{yj}j=1M\{y_{j}\}_{j=1}^{M}from the target measurepp. The particles{xi(k)}i=1N\{x_{i}^{(k)}\}_{i=1}^{N}evolve asxi(k+1)=xi(k)+η​Vq^kN,p^Mε​(xi(k))x_{i}^{(k+1)}=x_{i}^{(k)}+\eta\,V_{\widehat{q}_{k}^{N},\widehat{p}^{M}}^{\varepsilon}(x_{i}^{(k)}), whereq^kN\widehat{q}_{k}^{N}denotes the empirical measure of{xi(k)}i=1N\{x_{i}^{(k)}\}_{i=1}^{N}. Define the continuous-time interpolation asxiη​(t):=xi(k)+(t−tk)​Vq^kN,p^Mε​(xi(k))x_{i}^{\eta}(t):=x_{i}^{(k)}+(t-t_{k})\,V^{\varepsilon}_{\widehat{q}_{k}^{N},\widehat{p}^{M}}(x_{i}^{(k)}),t∈[tk,tk+1)t\in[t_{k},t_{k+1}), with the empirical measureq^tN,M,η≔1N​∑i=1Nδxiη​(t)\widehat{q}_{t}^{N,M,\eta}\coloneqq\frac{1}{N}\sum_{i=1}^{N}\delta_{x_{i}^{\eta}(t)}. Under suitable regularity conditions, for everyT>0T>0, there exists a uniqueq∈C​([0,T];𝒫2​(ℝd))q\in C([0,T];\mathcal{P}_{2}(\mathbb{R}^{d}))that solves the continuity equation

∂tqt+∇⋅(qt​Vqt,pε)=0\partial_{t}q_{t}+\nabla\cdot\big(q_{t}\,V^{\varepsilon}_{q_{t},p}\big)=0in the weak sense. Moreover, asη→0\eta\to 0andN,M→∞N,M\to\infty, it holds that

supt∈[0,T]𝒲2​(q^tN,M,η,qt)→0.\sup_{t\in[0,T]}\mathcal{W}_{2}\big(\widehat{q}_{t}^{N,M,\eta},q_{t}\big)\to 0.

The main result is formulated for the W-Flow setting, capturing both the two-measure velocity fieldVq,pεV^{\varepsilon}_{q,p}and the empirical target approximation. Its proof uses stability and particle-approximation techniques for nonlocal continuity equations. These techniques are largely inspired by related arguments in the literature, though the assumptions and settings are different; seePiccoli and Rossi (2013).Relatedly, Proposition 3.8 ofHeet al.(2026)studies the statistical approximation of the Sinkhorn drifting field, whereas our result establishes the convergence of the particle dynamics under a more general setting (see AssumptionB.1).

Algorithm 1W-Flow TrainingInput:Generatorfθf_{\theta}, reference distributionprefp_{\mathrm{ref}}, target distributionpp, step sizeη\eta, regularizationε\varepsilon, iterationsLL

1:repeat

2:Sample input

{zi}i=1N∼pref,{zl′}l=1N∼pref\{z_{i}\}_{i=1}^{N}\sim p_{\mathrm{ref}},\quad\{z^{\prime}_{l}\}_{l=1}^{N}\sim p_{\mathrm{ref}} 3:Sample target

{yj}j=1M∼p\{y_{j}\}_{j=1}^{M}\sim p 4:

xi←fθ​(zi),xl′←sg​(fθ​(zl′))x_{i}\leftarrow f_{\theta}(z_{i}),\quad x^{\prime}_{l}\leftarrow\mathrm{sg}(f_{\theta}(z^{\prime}_{l})) 5:

{Tq^θ,p^ε​(xi)}←Sinkhorn​({xi},{yj},ε,L)\{T^{\varepsilon}_{\hat{q}_{\theta},\hat{p}}(x_{i})\}\leftarrow{\color[rgb]{0,0,1}\textsc{Sinkhorn}}(\{x_{i}\},\{y_{j}\},\varepsilon,L) 6:

{Tq^θ,q^θ′ε​(xi)}←Sinkhorn​({xi},{xl′},ε,L)\{T^{\varepsilon}_{\hat{q}_{\theta},\hat{q}^{\prime}_{\theta}}(x_{i})\}\leftarrow{\color[rgb]{0,0,1}\textsc{Sinkhorn}}(\{x_{i}\},\{x^{\prime}_{l}\},\varepsilon,L) 7:

Vε​(xi)←Tq^θ,p^ε​(xi)−Tq^θ,q^θ′ε​(xi)V^{\varepsilon}(x_{i})\leftarrow T^{\varepsilon}_{\hat{q}_{\theta},\hat{p}}(x_{i})-T^{\varepsilon}_{\hat{q}_{\theta},\hat{q}^{\prime}_{\theta}}(x_{i}) 8:

x~i←sg​(xi+η​Vε​(xi))\tilde{x}_{i}\leftarrow\mathrm{sg}(x_{i}+\eta V^{\varepsilon}(x_{i})) 9:Take a gradient descent step on

10:

∇θ1N​∑i=1N‖fθ​(zi)−x~i‖2\qquad\nabla_{\theta}\frac{1}{N}\sum_{i=1}^{N}\|f_{\theta}(z_{i})-\tilde{x}_{i}\|^{2} 11:untilconverged

Algorithm 2Sinkhorn Projection (Sinkhorn)Input:Source particles{xi}i=1N\{x_{i}\}_{i=1}^{N}, target particles{yj}j=1M\{y_{j}\}_{j=1}^{M}, regularizationε\varepsilon, Sinkhorn iterationsLL

1:

a←1N​𝟏N,b←1M​𝟏Ma\leftarrow\frac{1}{N}\mathbf{1}_{N},\quad b\leftarrow\frac{1}{M}\mathbf{1}_{M}⊳\trianglerightUniform weights

2:

Ci​j←12​‖xi−yj‖2C_{ij}\leftarrow\frac{1}{2}\|x_{i}-y_{j}\|^{2}⊳\trianglerightCost matrix

3:

Ki​j=exp⁡(−Ci​j/ε)K_{ij}=\exp(-C_{ij}/\varepsilon)⊳\trianglerightGibbs kernel

4:

v(0)←𝟏Mv^{(0)}\leftarrow\mathbf{1}_{M}⊳\trianglerightInitialization

5:for

ℓ=0,…,L−1\ell=0,\dotsc,L-1do⊳\trianglerightSinkhorn-Knopp

6:

u(ℓ+1)←a⊘(K​v(ℓ))u^{(\ell+1)}\leftarrow a\oslash(Kv^{(\ell)}) 7:

v(ℓ+1)←b⊘(K⊤​u(ℓ+1))v^{(\ell+1)}\leftarrow b\oslash(K^{\top}u^{(\ell+1)}) 8:endfor

9:

Πε←diag⁡(u(L))​K​diag⁡(v(L))\Pi^{\varepsilon}\leftarrow\operatorname{diag}(u^{(L)})\,K\,\operatorname{diag}(v^{(L)})⊳\trianglerightCoupling

10:

Tε​(xi)=1ai​∑j=1MΠi​jε​yjT^{\varepsilon}(x_{i})=\frac{1}{a_{i}}\sum_{j=1}^{M}\Pi^{\varepsilon}_{ij}y_{j}⊳\trianglerightProjection

11:return

{Tε​(xi)}i=1N\{T^{\varepsilon}(x_{i})\}_{i=1}^{N}

3.4Implementation

Training.Empirically, we construct the training loss by learning the generatorfθf_{\theta}to simulate the discrete particle dynamics driven by the velocity in Eq. (13). This can be viewed as a finite-sample counterpart of the population-level objective in (3), where expectations with respect toq(k)q^{(k)}andppare replaced by empirical measures. Specifically, at each training step, we optimize

ℒW-Flow​(θ)=1N​∑i=1N‖xi−sg​(xi+η​Vq^θ,p^ε​(xi))‖2,xi=fθ​(zi),\mathcal{L}_{\text{W-Flow}}(\theta)=\frac{1}{N}\sum_{i=1}^{N}\left\|x_{i}-\mathrm{sg}\bigl(x_{i}+\eta\,V^{\varepsilon}_{\widehat{q}_{\theta},\widehat{p}}(x_{i})\bigr)\right\|^{2},\quad x_{i}=f_{\theta}(z_{i}),(14)where{zi}i=1N\{z_{i}\}_{i=1}^{N}are i.i.d. samples fromprefp_{\mathrm{ref}}. Crucially, the stop-gradient operatorsg​(⋅)\mathrm{sg}(\cdot)precisely aligns with the discrete WGF dynamics by treating the regression target as fixed, which also avoids differentiation through the OT plan and stabilizes the optimization. We summarize the overall training procedure in Alg.1.

Classifier-free guidance.Our framework naturally supports baking classifier-free guidance (CFG)Ho and Salimans (2022)into the pushforward mapfθf_{\theta}attraining time. The idea is to use a small batch of unconditional samples as additional reference particles, thereby introducing an extra repulsive effect on the generated samples. A natural implementation inDenget al.(2026)is to modify the repulsive distribution as

q~θw​(x|c)≔1w+1​qθ​(x|c)+ww+1​p​(x|∅),\displaystyle\tilde{q}^{w}_{\theta}(x|c)\coloneqq\frac{1}{w+1}q_{\theta}(x|c)+\frac{w}{w+1}p(x|\varnothing),(15)wherew≥0w\geq 0is the guidance scale andp​(x|∅)p(x|\varnothing)is the unconditional (marginal) distribution. Operationally, this amounts to sampling additional particles fromp​(x|∅)p(x|\varnothing)with importance weightwww.r.t.qθq_{\theta}. However, this approach implicitlyaltersthe target being learned. One can show that it corresponds to learning alinearlyextrapolated target distributionp~w​(x|c)=p​(x|c)+w​(p​(x|c)−p​(x|∅))\tilde{p}^{w}(x|c)=p(x|c)+w(p(x|c)-p(x|\varnothing)), which deviates from the standard CFG formulation based on exponential tiltingHo and Salimans (2022); Dhariwal and Nichol (2021).

In light of this, we instead introduce guidance directly at thevelocity field level. Specifically, we define

V~qθ,pε,w​(xi)≔(Tqθ(⋅|c),p(⋅|c)ε​(xi)−Tqθ(⋅|c),qθ(⋅|c)ε​(xi))+w​(Tqθ(⋅|c),p(⋅|c)ε​(xi)−Tqθ(⋅|c),p(⋅|∅)ε​(xi))¯,\displaystyle\tilde{V}^{\varepsilon,w}_{q_{\theta},p}(x_{i})\!\coloneqq\!\left(T^{\varepsilon}_{q_{\theta}(\cdot|c),p(\cdot|c)}(x_{i})\!-\!T^{\varepsilon}_{q_{\theta}(\cdot|c),q_{\theta}(\cdot|c)}(x_{i})\right)\!+\!w\underline{\left(T^{\varepsilon}_{q_{\theta}(\cdot|c),p(\cdot|c)}(x_{i})\!-\!T^{\varepsilon}_{q_{\theta}(\cdot|c),p(\cdot|\varnothing)}(x_{i})\right)},(16)whereTqθ(⋅|c),p(⋅|∅)εT^{\varepsilon}_{q_{\theta}(\cdot|c),p(\cdot|\varnothing)}denotes the Sinkhorn barycentric projection (Alg.2) betweenqθ(⋅|c)q_{\theta}(\cdot|c)and the unconditional data distributionp(⋅|∅)p(\cdot|\varnothing). This construction injects guidance via the underlined term that computes the difference between the conditional and unconditional velocity fields (see AppendixA.3). Due to these conceptual differences, we refer to Eq. (15) asdistribution guidanceand to Eq. (16) asvelocity guidance. Although the effective target distribution induced by Eq. (16) is generally not available in closed form, it admits a natural geometric interpretation through a barycentric construction. Moreover, we show in AppendixA.3that, when the KL divergence is used asℱ{\mathcal{F}}, this formulation exactly recovers the standard exponentially tilted target. This provides a conceptual justification for adopting the same design under the Sinkhorn divergence, whose empirical benefits are demonstrated in Sec.4.1.

Sinkhorn divergence in feature space.Rather than evaluating the Sinkhorn divergence directly in the target spaceℝn{\mathbb{R}}^{n}, we may first lift the particles{xi},{xl′},{yj}\{x_{i}\},\{x^{\prime}_{l}\},\{y_{j}\}into a feature space through a feature mapϕ:ℝn→ℝd\phi:\mathbb{R}^{n}\rightarrow\mathbb{R}^{d}. The velocity field and training loss are then computed using the embedded particles{ϕ​(xi)},{ϕ​(xl′)},{ϕ​(yj)}\{\phi(x_{i})\},\{\phi(x^{\prime}_{l})\},\{\phi(y_{j})\}, while gradients are backpropagated throughϕ\phito update the generator. In practice, followingDenget al.(2026), we use multiple feature maps and sum the corresponding losses.

4Experiments

Experiment setup.We conduct our main experiments on ImageNetDenget al.(2009)generation with 256×\times256 resolution. We directly adopt the DiT-stylePeebles and Xie (2023)generator architecture inDenget al.(2026). We also use the officially released pretrained MAE feature extractors fromDenget al.(2026). More details are deferred to AppendixC.

4.1Ablation study

Ablation setting.To enable an informative comparison in a controlled setting, we follow the default ablation setup ofDenget al.(2026). Specifically, all ablation experiments are conducted in the SD-VAERombachet al.(2022)latent space, using a B/2 DiT generator together with the pretrained latent-MAE feature encoder fromDenget al.(2026). The generator is trained for 100 epochs, and FIDHeuselet al.(2017)is evaluated on 50K generated images. We setε=0.05\varepsilon=0.05by default with additional detailed hyperparameters provided in AppendixC.

Energy functionals in gradient flow.We implement and compare three different energy functionals: squared MMD, KL divergence, and Sinkhorn divergence (Eq. (8)). Results in Table2aconfirm that Sinkhorn divergence performs the best, whereas both MMD and KL divergence yield noticeably weaker results, highlighting the advantage of the OT-driven dynamics. The Sinkhorn-based gradient flow also significantly outperforms the heuristic design of Drifting Model when both use distribution guidance.

Classifier-free guidance.We further equip the Sinkhorn W-Flow with our proposed velocity-guidance CFG (denoted asvelo.). As reported in Table2b, this guidance mechanism is fully compatible with Sinkhorn W-Flow and further improves the FID to 7.08, outperforming distribution guidance (dist.). We adopt it as the default configuration in the subsequent studies.

Entropy regularization parameterε\varepsilon.We next study the effect of the parameterε\varepsilonin the EOT problems; the results are reported in Table2b. Overall, W-Flow is fairly robust to the choice ofε\varepsilon. Among the values we consider,ε=0.05\varepsilon=0.05yields the best performance under the ablation setting. See Table9for more details. For the main experiments, we setε=0.05\varepsilon=0.05by default, with detailed choices specified in Table5.

DivergenceFID, 1-NFEDriftingDenget al.(2026)8.46MMD10.40KL10.17Sinkhorn7.29(a)Choice of divergence. Sinkhorn divergence leads to the lowest FID, outperforming heuristic designs inDenget al.(2026). ε\varepsilonCFGFID, 1-NFE0.05dist.7.290.05velo.7.080.025velo.7.110.075velo.7.26(b)Ablation on CFG andε\varepsilon. Our proposedveloCFG withε=0.05\varepsilon=0.05for Sinkhorn divergence performs the best. CostMask2-batchFID, 1-NFE∥⋅∥\|\cdot\|-✓7.1612∥⋅∥2\frac{1}{2}\|\cdot\|^{2}-✓7.0812∥⋅∥2\frac{1}{2}\|\cdot\|^{2}--17.5712∥⋅∥2\frac{1}{2}\|\cdot\|^{2}✓-7.45(c)OT cost andVtV_{t}estimation.“Mask” and “2-batch” refer to thediagonal maskandtwo-batch estimatorinTq,qεT^{\varepsilon}_{q,q}(§3.3).

Table 2:Ablation study on ImageNet 256×\times256 generation.We use a B/2 backbone with 100-epoch training from scratch. The FID on 50K images is reported. The default settings are marked ingray.The cost function in OT.We compare our default choice of quadratic costc​(x,y)=12​‖x−y‖2c(x,y)=\frac{1}{2}\|x-y\|^{2}with theℓ2\ell_{2}distance costc​(x,y)=‖x−y‖c(x,y)=\|x-y\|. The results in Table2cdemonstrate that the quadratic cost yields superior performance, in line with our theoretical assumptions in Eq. (7).

Velocity field estimation.We validate the two-batch estimator in Sec.3.3for the self-transport map. As shown in Table2c, estimatingTq^θ,q^θεT^{\varepsilon}_{\widehat{q}_{\theta},\widehat{q}_{\theta}}from a single batch leads to a substantially worse FID of 17.57, due to the intrinsic self-matching bias of this estimator. Our two-batch strategy also outperforms the mask-diagonal heuristic used inDenget al.(2026), while admitting a principled theoretical interpretation.

4.2Main results on ImageNet

Setup.Putting together the key components validated in the ablation study, we scale up training to achieve state-of-the-art results on ImageNet 256×\times256 generation. We follow the setup ofDenget al.(2026)with the SD-VAE latent space, adopting the same generator architecture and the same pretrained feature encoder weights. We also largely retain their hyperparameter configuration (see AppendixC).

Table 3:Class-conditional generation on ImageNet 256×\times256.FID is computed on 50K images. All methods are reported with CFG if applicable. “#Params” denotes the number of parameters of “generator + decoder”. All generators are trained from scratch without distilling from a teacher model.#ParamsNFEFID↓\downarrowIS↑\uparrowMulti-step Diffusion/FlowsADM-GDhariwal and Nichol(2021)554M250×\times24.59186.7DiT-XL/2Peebles and Xie(2023)675M+49M250×\times22.27278.2SiT-XL/2Maet al.(2024)675M+49M250×\times22.06270.3SiT-XL/2+REPAYuet al.(2024)675M+49M250×\times21.42305.7LightningDiT-XL/2Yaoet al.(2025)675M+70M250×\times21.35295.3RAE+DiTDH{}^{\text{DH}}-XL/2Zhenget al.(2025)839M+415M50×\times21.13262.6Masking & AutoregressiveMaskGITChanget al.(2022)227M+31M86.18182.1VAR-d​30d30Tianet al.(2024)2B+61M10×\times21.92323.1MAR-HLiet al.(2024)943M+41M256×\times21.55303.7GANsBigGANBrocket al.(2018)112M16.95152.8GigaGANKanget al.(2023)569M13.45225.5StyleGAN-XLSaueret al.(2022)166M12.30265.1

#ParamsNFEFID↓\downarrowIS↑\uparrowSingle-step Diffusion/FlowsiCT-XL/2Song and Dhariwal(2023)675M+49M134.24–Shortcut-XL/2Franset al.(2024)675M+49M110.60–MeanFlow-XL/2Genget al.(2025a)676M+49M13.43–TiM-XL/2Wanget al.(2025)664M+49M13.26210.3α\alpha-Flow-XL/2+Zhanget al.(2025)676M+49M12.58–AdvFlow-XL/2Linet al.(2025)673M+49M12.38284.2iMeanFlow-XL/2Genget al.(2025b)610M+49M11.72282.0Drifting ModelsDrifting Model, B/2Denget al.(2026)133M+49M11.75263.2Drifting Model, L/2Denget al.(2026)463M+49M11.54258.9W-Flow, B/2133M+49M11.52271.8W-Flow, L/2463M+49M11.35272.5W-Flow, XL/2679M+49M11.29265.4

Results.We compare W-Flow with state-of-the-art one-step generators and multi-step diffusion and flow models. Results are reported in Table3. Notably, W-Flow establishes a new state of the art for one-step class-conditional generation on ImageNet 256×\times256, achieving an FID of 1.29 at XL scale and 1.35 at L scale, outperforming MeanFlow-based methods and Drifting Models by a clear margin. Remarkably, W-Flow B/2, with only 133M generator parameters, achieves an FID of 1.52, surpassing Drifting Model L/2 despite its substantially larger 463M-parameter generator. Furthermore, despite 1-NFE sampling, W-Flow outperforms most diffusion models requiring up to 250 steps, such as LightningDiT-XL/2. See AppendixDfor a detailed inference throughput comparison. These strong empirical results support our central claim that principled WGF dynamics can translate into exceptional generation performance. W-Flow also exhibits a favorable scaling behavior with model size. More image samples produced by W-Flow are provided in AppendixF.

Figure 4:Classifier-free guidance.*Left:*The FID and Inception Score curve when sweeping over CFG scales.*Right:*Image samples by W-Flow, L/2 with CFG increasing from 0.0 to 2.0.The effect of CFG.Fig.4shows the FID and Inception Score with different CFG scales at inference time; W-Flow clearly achieves a better tradeoff compared with DriftingDenget al.(2026). Samples in Fig.4demonstrate that increasing CFG leads to clearer class-specific details and higher perceptual quality, which aligns with the expected CFG behavior.

Table 4:FID vs. training epoch.Convergence speed.W-Flow induces training dynamics that follow the steepest descent of the energy functional, leading to faster convergence at training time. We empirically validate this effect in Table4, where W-Flow converges substantially faster than Drifting Model. In particular, both W-Flow B/2 and L/2 trained for only 384 epochs closely match the FID of Drifting Model trained for 1280 epochs.

4.3Domain transfer and mode coverage

Domain transfer.Unlike standard generative models that typically map Gaussian noise to data, W-Flow can instead define a WGF between arbitrary distributions. This makes it naturally suited for domain transfer tasks where the source and target share the same semantic space. By settingprefp_{\mathrm{ref}}to be the source data distribution, it learns a one-step domain-transfer map,e.g., from senior to young-adult faces. To preserve source semantics, we parameterize the map as a residual networkHeet al.(2016),fθ​(z)=z+𝒩θ​(z)f_{\theta}(z)=z+\mathcal{N}_{\theta}(z), and zero-initialize the final layer of𝒩θ\mathcal{N}_{\theta}, such thatfθ0​(z)=zf_{\theta_{0}}(z)=z.

In a 2D oval-to-circle toy example in Fig.5(a), Drifting Model shows erratic intermediate states and often moves particles far from their initial positions. In contrast, the OT-induced velocity in W-Flow produces coordinated, nearly straight, and much shorter trajectories toward the target.

We further evaluate one-step age translation on FFHQKarraset al.(2019), mapping senior faces (ages 55–100) to young-adult faces (ages 18–30) in the 512-dimensional latent space of a pretrained autoencoderPidhorskyiet al.(2020). We train a 4-layer MLP that maps source latent codes to the target domain. Fig.5(b) and5(c) show that W-Flow achieves shorter latentℓ2\ell_{2}transport distances and better preserves source identity.

Figure 5:(a) Oval-to-circle domain transfer.Source and target are constructed by sampling angles uniformly from[0,2​π)[0,2\pi)with parametric curves corrupted by Gaussian noise.(b) & (c) One-step facial age translation on FFHQ, mapping older faces to younger ones.(b)Histogram of the latentℓ2\ell_{2}distance between 2,000 source images and their generated targets.(c)Visual comparison.Figure 6:Evaluation of mode coverage under imbalanced target distributions.(a)Evaluation of mode coverage on a 2D Gaussian mixture dataset featuring six dominant modes and two distant minority modes.(b)PCA scatter plot of generated latent codes for an artificially imbalanced FFHQ target distribution (95% senior faces, 5% child faces). See AppendixFfor generated samples comparing mode coverage behavior.Mode coverage.We further investigate whether the globally structured OT plan in W-Flow can mitigate mode collapse, a known vulnerability of Drifting Model (see AppendixA.4). We construct a 2D Gaussian mixture dataset with six large, tightly concentrated modes and two small distant minority modes. As shown in Fig.6(a), Drifting Model greedily collapses on the dominant clusters, while the Sinkhorn barycentric projections in W-Flow instead enforce global marginal constraints and successfully preserve the coverage of all modes.

We further extend this observation to image generation on FFHQ under an imbalanced training distribution, where target batches contain 95% senior faces (ages 55–100) and 5% child faces (ages 0–12). The PCA visualization in Fig.6(b) shows that Drifting Model almost entirely drops the minority child mode, whereas W-Flow successfully captures it while maintaining high fidelity on the dominant senior mode. See AppendixFfor representative samples.

5Conclusion

We introduced W-Flow, a new paradigm for one-step generative modeling that guides the training dynamics of image generators by Wasserstein gradient flow. Our method derives the velocity field via the steepest descent of Sinkhorn divergence, yielding a globally coordinated OT-based update rule with consistent particle dynamics. Empirically, W-Flow sets a new state of the art for one-step ImageNet 256×\times256 generation with 1.29 FID. The results establish the WGF dynamics as a foundation for fast, high-fidelity generation. Future work includes exploring a broader class of energy functionals within the WGF framework and extending this paradigm beyond class-conditional image generation to more complex settings such as text-to-image generation.

Acknowledgments

We gratefully acknowledge the Stanford Marlowe clusterKapferet al.(2025)for supporting part of the GPU computing resources. RX was partially supported by an NSF CAREER Award DMS-2614933 and a gift fund from Point72. SE was supported by ARO (W911NF-21-1-0125), ONR (N00014-23-1-2159), and the CZ Biohub. EC was supported by ONR (N00014-24-1-2305) and the NIH (1R01AG08950901A1). The authors thank Yaniv Romano, Chieh-Hsin Lai, Gabriel Peyré, Mingyang Deng, and Tudor Manole for feedback on an early version of this paper.

References

• [1]M. S. Albergo and E. Vanden-Eijnden(2022)Building normalizing flows with stochastic interpolants.arXiv preprint arXiv:2209.15571.Cited by:§2.
• [2]A. Alouadi, P. Henry-Labordère, G. Loeper, O. Mazhar, H. Pham, and N. Touzi(2026)LightSBB-m: bridging Schrödinger and bass for generative diffusion modeling.arXiv preprint arXiv:2601.19312.Cited by:§2.
• [3]L. Ambrosio, N. Gigli, and G. Savaré(2008)Gradient flows: in metric spaces and in the space of probability measures.Birkhäuser.Cited by:Remark B.2.
• [4]M. Arjovsky, S. Chintala, and L. Bottou(2017)Wasserstein generative adversarial networks.InInternational conference on machine learning,pp. 214–223.Cited by:§1,§2,§2.
• [5]N. M. Boffi, M. S. Albergo, and E. Vanden-Eijnden(2025)How to build a consistency model: learning flow maps via self-distillation.arXiv preprint arXiv:2505.18825.Cited by:§2.
• [6]A. Brock, J. Donahue, and K. Simonyan(2018)Large scale GAN training for high fidelity natural image synthesis.arXiv preprint arXiv:1809.11096.Cited by:Table 3.
• [7]J. Cao, Z. Wei, and Y. Liu(2026)Gradient flow drifting: generative modeling via Wasserstein gradient flows of KDE-approximated divergences.arXiv preprint arXiv:2603.10592.Cited by:§2.
• [8]G. Carlier, L. Chizat, and M. Laborde(2024)Displacement smoothness of entropic optimal transport.ESAIM: Control, Optimisation and Calculus of Variations30,pp. 25.Cited by:§B.2,§B.2.
• [9]R. Carmona and F. Delarue(2018)Probabilistic theory of mean field games with applications i.Springer.Cited by:Remark B.2.
• [10]H. Chang, H. Zhang, L. Jiang, C. Liu, and W. T. Freeman(2022)Maskgit: masked generative image transformer.InProceedings of the IEEE/CVF conference on computer vision and pattern recognition,pp. 11315–11325.Cited by:Table 3.
• [11]J. Choi, J. Choi, and M. Kang(2024-21–27 Jul)Scalable Wasserstein gradient flow for generative modeling through unbalanced optimal transport.InProceedings of the 41st International Conference on Machine Learning,Proceedings of Machine Learning Research, Vol.235,pp. 8629–8650.Cited by:§A.4,§A.4,§2.
• [12]V. De Bortoli, J. Thornton, J. Heng, and A. Doucet(2021)Diffusion Schrödinger bridge with applications to score-based generative modeling.arXiv preprint arXiv:2106.01357.Cited by:§2.
• [13]J. Deng, W. Dong, R. Socher, L. Li, K. Li, and L. Fei-Fei(2009)ImageNet: a large-scale hierarchical image database.InCVPR,pp. 248–255.Cited by:§4.
• [14]M. Deng, H. Li, T. Li, Y. Du, and K. He(2026)Generative modeling via drifting.arXiv preprint arXiv:2602.04770.Cited by:Figure 9,§A.2,§A.4,Appendix C,Appendix E,§1,§1,§2,§3.1,§3.1,§3.3,§3.4,§3.4,§4.1,§4.1,§4.2,§4.2,2a,2a,Table 3,Table 3,Table 4,§4.
• [15]P. Dhariwal and A. Nichol(2021)Diffusion models beat GANs on image synthesis.NeurIPS34,pp. 8780–8794.Cited by:§3.4,Table 3.
• [16]J. Fan, Q. Zhang, A. Taghvaei, and Y. Chen(2021)Variational Wasserstein gradient flow.arXiv preprint arXiv:2112.02424.Cited by:§2.
• [17]J. Feydy, T. Séjourné, F. Vialard, S. Amari, A. Trouvé, and G. Peyré(2019)Interpolating between optimal transport and mmd using Sinkhorn divergences.InThe 22nd international conference on artificial intelligence and statistics,pp. 2681–2690.Cited by:§3.2.
• [18]K. Frans, D. Hafner, S. Levine, and P. Abbeel(2024)One step diffusion via shortcut models.arXiv preprint arXiv:2410.12557.Cited by:Table 3.
• [19]Y. Gao, Y. Shen, S. Zhang, W. Yu, Y. Duan, J. Wu, J. Deng, Y. Zhang,et al.(2026)Drift-based policy optimization: native one-step policy learning for online robot control.arXiv preprint arXiv:2604.03540.Cited by:§2.
• [20]A. Genevay, G. Peyré, and M. Cuturi(2018)Learning generative models with Sinkhorn divergences.InInternational Conference on Artificial Intelligence and Statistics,pp. 1608–1617.Cited by:§A.4,§2,§2,§3.2.
• [21]Z. Geng, M. Deng, X. Bai, J. Z. Kolter, and K. He(2025)Mean flows for one-step generative modeling.arXiv preprint arXiv:2505.13447.Cited by:§2,Table 3.
• [22]Z. Geng, Y. Lu, Z. Wu, E. Shechtman, J. Z. Kolter, and K. He(2025)Improved mean flows: on the challenges of fastforward generative models.arXiv preprint arXiv:2512.02012.Cited by:§2,Table 3.
• [23]I. J. Goodfellow, J. Pouget-Abadie, M. Mirza, B. Xu, D. Warde-Farley, S. Ozair, A. Courville, and Y. Bengio(2014)Generative adversarial nets.NeurIPS.Cited by:§1,§2.
• [24]I. Gulrajani, F. Ahmed, M. Arjovsky, V. Dumoulin, and A. C. Courville(2017)Improved training of Wasserstein GANs.Advances in neural information processing systems30.Cited by:§2.
• [25]M. Hardion and T. Lacombe(2026)The Wasserstein gradient flow of the Sinkhorn divergence between gaussian distributions.arXiv preprint arXiv:2602.10726.Cited by:§3.2.
• [26]K. He, X. Zhang, S. Ren, and J. Sun(2016)Deep residual learning for image recognition.InCVPR,pp. 770–778.Cited by:§4.3.
• [27]P. He, O. Khangaonkar, H. Pirsiavash, Y. Bai, and S. Kolouri(2026)Sinkhorn-drifting generative models.arXiv preprint arXiv:2603.12366.Cited by:§2,§2,§3.2,§3.3.
• [28]M. Heusel, H. Ramsauer, T. Unterthiner, B. Nessler, and S. Hochreiter(2017)GANs trained by a two time-scale update rule converge to a local Nash equilibrium.NeurIPS.Cited by:§4.1.
• [29]J. Ho, A. Jain, and P. Abbeel(2020)Denoising diffusion probabilistic models.NeurIPS33,pp. 6840–6851.Cited by:§1,§2.
• [30]J. Ho and T. Salimans(2022)Classifier-free diffusion guidance.arXiv preprint arXiv:2207.12598.Cited by:§3.4,§3.4.
• [31]R. Jordan, D. Kinderlehrer, and F. Otto(1998)The variational formulation of the Fokker–Planck equation.SIAM journal on mathematical analysis29(1),pp. 1–17.Cited by:§2.
• [32]M. Kang, J. Zhu, R. Zhang, J. Park, E. Shechtman, S. Paris, and T. Park(2023)Scaling up GANs for text-to-image synthesis.InCVPR,pp. 10124–10134.Cited by:Table 3.
• [33]C. Kapfer, K. Stine, B. Narasimhan, C. Mentzel, and E. Candes(2025)Marlowe: stanford’s gpu-based computational instrument.January.Cited by:Acknowledgments.
• [34]T. Karras, S. Laine, and T. Aila(2019)A style-based generator architecture for generative adversarial networks.InProceedings of the IEEE/CVF conference on computer vision and pattern recognition,pp. 4401–4410.Cited by:§4.3.
• [35]C. Lai, B. Nguyen, N. Murata, Y. Takida, T. Uesaka, Y. Mitsufuji, S. Ermon, and M. Tao(2026)A unified view of drifting and score-based models.arXiv preprint arXiv:2603.07514.Cited by:§2.
• [36]T. Li, Y. Tian, H. Li, M. Deng, and K. He(2024)Autoregressive image generation without vector quantization.NeurIPS37,pp. 56424–56445.Cited by:Table 3.
• [37]Y. Li, K. Swersky, and R. Zemel(2015)Generative moment matching networks.InICML,pp. 1718–1727.Cited by:§A.4,§3.2.
• [38]Z. Li, S. Mei, W. Wu, A. Maier, L. Gölz, and Y. Xia(2026)Generative drifting for conditional medical image generation.arXiv preprint arXiv:2604.19736.Cited by:§2.
• [39]S. Lin, C. Yang, Z. Lin, H. Chen, and H. Fan(2025)Adversarial flow models.arXiv preprint arXiv:2511.22475.Cited by:§2,Table 3.
• [40]Y. Lipman, R. T. Chen, H. Ben-Hamu, M. Nickel, and M. Le(2022)Flow matching for generative modeling.arXiv preprint arXiv:2210.02747.Cited by:§1,§2,§2.
• [41]X. Liu, C. Gong, and Q. Liu(2022)Flow straight and fast: learning to generate and transfer data with rectified flow.arXiv preprint arXiv:2209.03003.Cited by:§1,§2,§2.
• [42]C. Lu and Y. Song(2024)Simplifying, stabilizing and scaling continuous-time consistency models.arXiv preprint arXiv:2410.11081.Cited by:Figure 1,§2.
• [43]J. Ma, Y. Tan, and R. Xu(2025)Schrödinger bridge for generative ai: soft-constrained formulation and convergence analysis.arXiv preprint arXiv:2510.11829.Cited by:§2.
• [44]N. Ma, M. Goldstein, M. S. Albergo, N. M. Boffi, E. Vanden-Eijnden, and S. Xie(2024)SiT: exploring flow and diffusion-based generative models with scalable interpolant transformers.InECCV,pp. 23–40.Cited by:Table 6,Appendix D,Table 3.
• [45]P. Mokrov, A. Korotin, L. Li, A. Genevay, J. M. Solomon, and E. Burnaev(2021)Large-scale Wasserstein gradient flows.Advances in Neural Information Processing Systems34,pp. 15243–15256.Cited by:§2.
• [46]M. Nutz and J. Wiesel(2022)Entropic optimal transport: convergence of potentials.Probability Theory and Related Fields184(1),pp. 401–424.Cited by:§3.2.
• [47]W. Peebles and S. Xie(2023)Scalable diffusion models with transformers.InCVPR,pp. 4195–4205.Cited by:Table 3,§4.
• [48]G. Peyré and M. Cuturi(2019)Computational optimal transport: with applications to data science.Now Foundations and Trends.Cited by:§3.2.
• [49]B. Piccoli and F. Rossi(2013)Transport equation with nonlocal velocity in Wasserstein spaces: convergence of numerical schemes.Acta applicandae mathematicae124(1),pp. 73–105.Cited by:§3.3.
• [50]S. Pidhorskyi, D. A. Adjeroh, and G. Doretto(2020)Adversarial latent autoencoders.InProceedings of the IEEE/CVF conference on computer vision and pattern recognition,pp. 14104–14113.Cited by:§4.3.
• [51]A. Radford, L. Metz, and S. Chintala(2015)Unsupervised representation learning with deep convolutional generative adversarial networks.arXiv preprint arXiv:1511.06434.Cited by:§1,§2.
• [52]R. Rombach, A. Blattmann, D. Lorenz, P. Esser, and B. Ommer(2022)High-resolution image synthesis with latent diffusion models.InCVPR,pp. 10684–10695.Cited by:§4.1.
• [53]T. Salimans and J. Ho(2022)Progressive distillation for fast sampling of diffusion models.arXiv preprint arXiv:2202.00512.Cited by:§2.
• [54]T. Salimans, T. Mensink, J. Heek, and E. Hoogeboom(2024)Multistep distillation of diffusion models via moment matching.NeurIPS37,pp. 36046–36070.Cited by:§2.
• [55]A. Sauer, K. Schwarz, and A. Geiger(2022)StyleGAN-XL: scaling StyleGAN to large diverse datasets.InSIGGRAPH,pp. 1–10.Cited by:Table 3.
• [56]R. Sinkhorn and P. Knopp(1967)Concerning nonnegative matrices and doubly stochastic matrices.Pacific Journal of Mathematics21(2),pp. 343–348.Cited by:§3.3.
• [57]J. Sohl-Dickstein, E. Weiss, N. Maheswaranathan, and S. Ganguli(2015)Deep unsupervised learning using nonequilibrium thermodynamics.InICML,pp. 2256–2265.Cited by:§1,§2.
• [58]Y. Song, P. Dhariwal, M. Chen, and I. Sutskever(2023)Consistency models.Cited by:§1,§2.
• [59]Y. Song and P. Dhariwal(2023)Improved techniques for training consistency models.arXiv preprint arXiv:2310.14189.Cited by:§2,Table 3.
• [60]Y. Song and S. Ermon(2019)Generative modeling by estimating gradients of the data distribution.Advances in neural information processing systems32.Cited by:§1.
• [61]Y. Song, J. Sohl-Dickstein, D. P. Kingma, A. Kumar, S. Ermon, and B. Poole(2020)Score-based generative modeling through stochastic differential equations.arXiv preprint arXiv:2011.13456.Cited by:§1,§2.
• [62]K. Tian, Y. Jiang, Z. Yuan, B. Peng, and L. Wang(2024)Visual autoregressive modeling: scalable image generation via next-scale prediction.NeurIPS37,pp. 84839–84865.Cited by:Table 3.
• [63]I. Tolstikhin, O. Bousquet, S. Gelly, and B. Schoelkopf(2017)Wasserstein auto-encoders.arXiv preprint arXiv:1711.01558.Cited by:§2.
• [64]E. Turan and M. Ovsjanikov(2026)Generative drifting is secretly score matching: a spectral and variational perspective.arXiv preprint arXiv:2603.09936.Cited by:§2.
• [65]Z. Wang, C. Lu, Y. Wang, F. Bao, C. Li, H. Su, and J. Zhu(2023)Prolificdreamer: high-fidelity and diverse text-to-3d generation with variational score distillation.Advances in neural information processing systems36,pp. 8406–8441.Cited by:§2.
• [66]Z. Wang, Y. Zhang, X. Yue, X. Yue, Y. Li, W. Ouyang, and L. Bai(2025)Transition models: rethinking the generative learning objective.arXiv preprint arXiv:2509.04394.Cited by:Table 3.
• [67]Y. Xie and X. Cheng(2025)Flow-based generative models as iterative algorithms in probability space.arXiv preprint arXiv:2502.13394.Cited by:§2.
• [68]J. Yao, B. Yang, and X. Wang(2025)Reconstruction vs. generation: taming optimization dilemma in latent diffusion models.InCVPR,pp. 15703–15712.Cited by:Table 6,Appendix D,Table 3.
• [69]T. Yin, M. Gharbi, T. Park, R. Zhang, E. Shechtman, F. Durand, and W. T. Freeman(2024)Improved distribution matching distillation for fast image synthesis.Advances in neural information processing systems37,pp. 47455–47487.Cited by:§2.
• [70]T. Yin, M. Gharbi, R. Zhang, E. Shechtman, F. Durand, W. T. Freeman, and T. Park(2024)One-step diffusion with distribution matching distillation.InCVPR,pp. 6613–6623.Cited by:§2.
• [71]S. Yu, S. Kwak, H. Jang, J. Jeong, J. Huang, J. Shin, and S. Xie(2024)Representation alignment for generation: training diffusion transformers is easier than you think.arXiv preprint arXiv:2410.06940.Cited by:Table 3.
• [72]H. Zhang, A. Siarohin, W. Menapace, M. Vasilkovsky, S. Tulyakov, Q. Qu, and I. Skorokhodov(2025)AlphaFlow: understanding and improving MeanFlow models.arXiv preprint arXiv:2510.20771.Cited by:Table 3.
• [73]B. Zheng, N. Ma, S. Tong, and S. Xie(2025)Diffusion transformers with representation autoencoders.arXiv preprint arXiv:2510.11690.Cited by:Table 3.
• [74]L. Zhou, S. Ermon, and J. Song(2025)Inductive moment matching.arXiv preprint arXiv:2503.07565.Cited by:§A.4,§1,§2.

Appendix

Appendix AAdditional discussions

A.1Wasserstein gradient flows of energy functionals

In this section, we derive the WGFs induced by the functionals listed in Table1. For MMD and KL divergence, we also provide their particle-based estimators in the same spirit as the Sinkhorn estimator used in Sec.3.3. The estimators are then used in the ablation study in Table2a.

Throughout this section,ppdenotes the target data distribution. Given an energy functionalℱ\mathcal{F}, recall that the corresponding WGF follows

∂tqt=∇⋅(qt​∇δ​ℱδ​q​(qt)),\partial_{t}q_{t}=\nabla\cdot\left(q_{t}\nabla\frac{\delta\mathcal{F}}{\delta q}(q_{t})\right),or equivalently, through the continuity equation, we have

∂tqt+∇⋅(qt​Vt)=0,Vt​(x)=−∇δ​ℱδ​q​(qt)​(x).\partial_{t}q_{t}+\nabla\cdot(q_{t}V_{t})=0,\qquad V_{t}(x)=-\nabla\frac{\delta\mathcal{F}}{\delta q}(q_{t})(x).

Gradient flow of Sinkhorn divergence.

To compute the first variation ofSε​(q,p)S_{\varepsilon}(q,p)under quadratic cost with respect toqq, we first consider the dual formulation of the EOT problem:

OTε​(q,p)=maxu,v\displaystyle\mathrm{OT}_{\varepsilon}(q,p)=\max_{u,v}∫ℝdu​(x)​𝑑q​(x)+∫ℝdv​(y)​𝑑p​(y)\displaystyle\int_{\mathbb{R}^{d}}u(x)dq(x)+\int_{\mathbb{R}^{d}}v(y)dp(y)−ε​∬ℝd×ℝdexp⁡(u​(x)+v​(y)−12​‖x−y‖2ε)​𝑑q​(x)​𝑑p​(y)+ε,\displaystyle-\varepsilon\iint_{\mathbb{R}^{d}\times\mathbb{R}^{d}}\exp\left(\frac{u(x)+v(y)-\frac{1}{2}\|x-y\|^{2}}{\varepsilon}\right)dq(x)dp(y)+\varepsilon,(17)whereu​(x)u(x)andv​(y)v(y)are the dual potentials.

By the envelope theorem, the first variation of the EOT cost Eq. (A.1) with respect to the marginal distributionqqis exactly the optimal dual potentialuq,pε​(x)u_{q,p}^{\varepsilon}(x):

δ​OTε​(q,p)δ​q​(x)=uq,pε​(x).\frac{\delta\mathrm{OT}_{\varepsilon}(q,p)}{\delta q}(x)=u_{q,p}^{\varepsilon}(x). Consequently, the first variation ofSε​(⋅,p)S_{\varepsilon}(\cdot,p)evaluated atqqis:

δ​Sε​(⋅,p)δ​q​(q)​(x)=uq,pε​(x)−uq,qε​(x).\frac{\delta S_{\varepsilon}(\cdot,p)}{\delta q}(q)(x)=u_{q,p}^{\varepsilon}(x)-u_{q,q}^{\varepsilon}(x). From the optimality conditions of the dual problem,uq,pε​(x)u_{q,p}^{\varepsilon}(x)satisfies the following structural equation:

uq,pε​(x)=−ε​log​∫ℝdexp⁡(vq,pε​(y)−12​‖x−y‖2ε)​𝑑p​(y).u_{q,p}^{\varepsilon}(x)=-\varepsilon\log\int_{\mathbb{R}^{d}}\exp\left(\frac{v_{q,p}^{\varepsilon}(y)-\frac{1}{2}\|x-y\|^{2}}{\varepsilon}\right)dp(y).(18)Taking the gradient with respect toxxyields:

∇xuq,pε​(x)=\displaystyle\nabla_{x}u_{q,p}^{\varepsilon}(x)=∫ℝd∇x12​‖x−y‖2​exp⁡(vq,pε​(y)−12​‖x−y‖2ε)​d​p​(y)∫ℝdexp⁡(vq,pε​(y)−12​‖x−y‖2ε)​𝑑p​(y)\displaystyle\frac{\int_{\mathbb{R}^{d}}\nabla_{x}\frac{1}{2}\|x-y\|^{2}\exp\left(\frac{v_{q,p}^{\varepsilon}(y)-\frac{1}{2}\|x-y\|^{2}}{\varepsilon}\right)dp(y)}{\int_{\mathbb{R}^{d}}\exp\left(\frac{v_{q,p}^{\varepsilon}(y)-\frac{1}{2}\|x-y\|^{2}}{\varepsilon}\right)dp(y)}=\displaystyle=∫ℝd∇x12​‖x−y‖2​πq,pε​(d​y|x)\displaystyle\int_{\mathbb{R}^{d}}\nabla_{x}\frac{1}{2}\|x-y\|^{2}\pi_{q,p}^{\varepsilon}(dy|x)=\displaystyle=∫ℝd(x−y)​πq,pε​(d​y|x)\displaystyle\int_{\mathbb{R}^{d}}(x-y)\pi_{q,p}^{\varepsilon}(dy|x)=\displaystyle=x−Tq,pε​(x).\displaystyle x-T_{q,p}^{\varepsilon}(x).(19) Combining the previous results, we obtain

Vqt,pε​(x)\displaystyle V_{q_{t},p}^{\varepsilon}(x)=−∇xδ​Sε​(⋅,p)δ​q​(qt)​(x)\displaystyle=-\nabla_{x}\frac{\delta S_{\varepsilon}(\cdot,p)}{\delta q}(q_{t})(x)=−∇x(uqt,pε​(x)−uqt,qtε​(x))\displaystyle=-\nabla_{x}\left(u_{q_{t},p}^{\varepsilon}(x)-u_{q_{t},q_{t}}^{\varepsilon}(x)\right)=−(x−Tqt,pε​(x))+(x−Tqt,qtε​(x))\displaystyle=-\left(x-T_{q_{t},p}^{\varepsilon}(x)\right)+\left(x-T_{q_{t},q_{t}}^{\varepsilon}(x)\right)=Tqt,pε​(x)−Tqt,qtε​(x).\displaystyle=T_{q_{t},p}^{\varepsilon}(x)-T_{q_{t},q_{t}}^{\varepsilon}(x). This recovers the Sinkhorn divergence velocity field reported in Table1.

Gradient flow of MMD.

Letk:ℝd×ℝd→ℝk:\mathbb{R}^{d}\times\mathbb{R}^{d}\to\mathbb{R}be a differentiable positive definite kernel. For probability measuresp,q∈𝒫​(ℝd)p,q\in\mathcal{P}(\mathbb{R}^{d}), we consider the energy

ℱMMD​(q):=12​MMD2​(q,p),\mathcal{F}_{\mathrm{MMD}}(q):=\frac{1}{2}\mathrm{MMD}^{2}(q,p),(20)where

MMD2​(q,p)=∬k​(x,y)​q​(d​x)​q​(d​y)−2​∬k​(x,y)​q​(d​x)​p​(d​y)+∬k​(x,y)​p​(d​x)​p​(d​y).\mathrm{MMD}^{2}(q,p)=\iint k(x,y)q(dx)q(dy)-2\iint k(x,y)q(dx)p(dy)+\iint k(x,y)p(dx)p(dy).(21)Equivalently,

ℱMMD​(q)=12​∬k​(x,y)​(q−p)​(d​x)​(q−p)​(d​y).\mathcal{F}_{\mathrm{MMD}}(q)=\frac{1}{2}\iint k(x,y)(q-p)(dx)(q-p)(dy).To compute the first variation, consider a perturbationqs=q+s​χq_{s}=q+s\,\chiwith signed measureχ\chisatisfying∫χ​(d​x)=0\int\chi(dx)=0. Then

dd​s​ℱMMD​(qs)|s=0\displaystyle\frac{d}{ds}\mathcal{F}_{\mathrm{MMD}}(q_{s})\bigg|_{s=0}=∬k​(x,y)​χ​(d​x)​(q−p)​(d​y)\displaystyle=\iint k(x,y)\chi(dx)(q-p)(dy)=∫[∫k​(x,y)​(q−p)​(d​y)]​χ​(d​x).\displaystyle=\int\left[\int k(x,y)(q-p)(dy)\right]\chi(dx).Therefore, up to an additive constant,

δ​ℱMMDδ​q​(q)​(x)=∫k​(x,y)​q​(d​y)−∫k​(x,y)​p​(d​y).\frac{\delta\mathcal{F}_{\mathrm{MMD}}}{\delta q}(q)(x)=\int k(x,y)q(dy)-\int k(x,y)p(dy).Taking the negative spatial gradient gives the Wasserstein velocity field

VtMMD​(x)\displaystyle V_{t}^{\mathrm{MMD}}(x)=−∇xδ​ℱMMDδ​q​(qt)​(x)\displaystyle=-\nabla_{x}\frac{\delta\mathcal{F}_{\mathrm{MMD}}}{\delta q}(q_{t})(x)=∫∇xk​(x,y)​p​(d​y)−∫∇xk​(x,y)​qt​(d​y),\displaystyle=\int\nabla_{x}k(x,y)p(dy)-\int\nabla_{x}k(x,y)q_{t}(dy),withqtq_{t}the corresponding flow under the initial conditionq0q_{0}. This is the MMD velocity field reported in Table1.

Particle-based estimator for MMD.

Given generated particles{xi}i=1N∼qt\{x_{i}\}_{i=1}^{N}\sim q_{t}and target particles{yj}j=1M∼p\{y_{j}\}_{j=1}^{M}\sim p, define the empirical measures

q^t=∑i=1Nai​δxi,p^=∑j=1Mbj​δyj,∑iai=∑jbj=1.\widehat{q}_{t}=\sum_{i=1}^{N}a_{i}\delta_{x_{i}},\qquad\widehat{p}=\sum_{j=1}^{M}b_{j}\delta_{y_{j}},\qquad\sum_{i}a_{i}=\sum_{j}b_{j}=1.As in the Sinkhorn self-transport estimator, we draw an independent second batch{xℓ′}ℓ=1N′∼qt\{x^{\prime}_{\ell}\}_{\ell=1}^{N^{\prime}}\sim q_{t}and define

q^t′=∑ℓ=1N′aℓ′​δxℓ′.\widehat{q}^{\prime}_{t}=\sum_{\ell=1}^{N^{\prime}}a^{\prime}_{\ell}\delta_{x^{\prime}_{\ell}}.Then the particle estimator of the MMD-induced velocity atxix_{i}is

V^q^t,p^MMD​(xi)=∑j=1Mbj​∇xk​(xi,yj)−∑ℓ=1N′aℓ′​∇xk​(xi,xℓ′).\widehat{V}_{\widehat{q}_{t},\widehat{p}}^{\mathrm{MMD}}(x_{i})=\sum_{j=1}^{M}b_{j}\nabla_{x}k(x_{i},y_{j})-\sum_{\ell=1}^{N^{\prime}}a^{\prime}_{\ell}\nabla_{x}k(x_{i},x^{\prime}_{\ell}).For uniform weights, this becomes

V^q^t,p^MMD​(xi)=1M​∑j=1M∇xk​(xi,yj)−1N′​∑ℓ=1N′∇xk​(xi,xℓ′).\widehat{V}_{\widehat{q}_{t},\widehat{p}}^{\mathrm{MMD}}(x_{i})=\frac{1}{M}\sum_{j=1}^{M}\nabla_{x}k(x_{i},y_{j})-\frac{1}{N^{\prime}}\sum_{\ell=1}^{N^{\prime}}\nabla_{x}k(x_{i},x^{\prime}_{\ell}).Ifk​(⋅,⋅)k(\cdot,\cdot)is the Gaussian kernel

kσ​(x,y)=exp⁡(−‖x−y‖22​σ2),k_{\sigma}(x,y)=\exp\left(-\frac{\|x-y\|^{2}}{2\sigma^{2}}\right),then

∇xkσ​(x,y)=y−xσ2​kσ​(x,y),\nabla_{x}k_{\sigma}(x,y)=\frac{y-x}{\sigma^{2}}k_{\sigma}(x,y),and hence

V^q^t,p^MMD​(xi)=1M​σ2​∑j=1M(yj−xi)​kσ​(xi,yj)−1N′​σ2​∑ℓ=1N′(xℓ′−xi)​kσ​(xi,xℓ′).\widehat{V}_{\widehat{q}_{t},\widehat{p}}^{\mathrm{MMD}}(x_{i})=\frac{1}{M\sigma^{2}}\sum_{j=1}^{M}(y_{j}-x_{i})k_{\sigma}(x_{i},y_{j})-\frac{1}{N^{\prime}\sigma^{2}}\sum_{\ell=1}^{N^{\prime}}(x^{\prime}_{\ell}-x_{i})k_{\sigma}(x_{i},x^{\prime}_{\ell}).The corresponding particle update follows:

xi←xi+η​V^q^t,p^MMD​(xi).x_{i}\leftarrow x_{i}+\eta\widehat{V}_{\widehat{q}_{t},\widehat{p}}^{\mathrm{MMD}}(x_{i}).

Gradient flow of KL divergence.

Forp,q∈𝒫​(ℝd)p,q\in\mathcal{P}(\mathbb{R}^{d}), we next consider the KL energy

ℱKL​(q):=DKL​(q∥p)=∫q​(x)​log⁡q​(x)p​(x)​d​x,\mathcal{F}_{\mathrm{KL}}(q):=D_{\mathrm{KL}}(q\|p)=\int q(x)\log\frac{q(x)}{p(x)}dx,assuming thatqqandppadmit smooth positive densities with respect to the Lebesgue measure. For a perturbationqs=q+s​χq_{s}=q+s\,\chiwith∫χ​(x)​𝑑x=0\int\chi(x)dx=0, we have

dd​s​ℱKL​(qs)|s=0\displaystyle\frac{d}{ds}\mathcal{F}_{\mathrm{KL}}(q_{s})\bigg|_{s=0}=∫(log⁡q​(x)−log⁡p​(x)+1)​χ​(x)​𝑑x.\displaystyle=\int\left(\log q(x)-\log p(x)+1\right)\chi(x)dx.Therefore, up to an additive constant,

δ​ℱKLδ​q​(q)​(x)=log⁡q​(x)−log⁡p​(x).\frac{\delta\mathcal{F}_{\mathrm{KL}}}{\delta q}(q)(x)=\log q(x)-\log p(x).The induced Wasserstein velocity field is

VtKL​(x)\displaystyle V_{t}^{\mathrm{KL}}(x)=−∇xδ​ℱKLδ​q​(qt)​(x)\displaystyle=-\nabla_{x}\frac{\delta\mathcal{F}_{\mathrm{KL}}}{\delta q}(q_{t})(x)=∇xlog⁡p​(x)−∇xlog⁡qt​(x),\displaystyle=\nabla_{x}\log p(x)-\nabla_{x}\log q_{t}(x),withqtq_{t}the corresponding flow under the initial conditionq0q_{0}. This recovers the KL velocity field reported in Table1. The first term attracts particles toward high-density regions of the target distribution, while the second term is a repulsive score term induced by the current model distribution.

Particle-based estimator for KL.

Unlike MMD and Sinkhorn divergence, the KL velocity field requires access to the score functions∇log⁡p\nabla\log pand∇log⁡qt\nabla\log q_{t}. In the absence of analytic scores, a particle-based estimator can be constructed through kernel density estimation (KDE).

Letkσ​(⋅,⋅)k_{\sigma}(\cdot,\cdot)be a smooth kernel with bandwidthσ>0\sigma>0, for example,

kσ​(x,y)=exp⁡(−‖x−y‖22​σ2).k_{\sigma}(x,y)=\exp\left(-\frac{\|x-y\|^{2}}{2\sigma^{2}}\right).Using target particles{yj}j=1M∼p\{y_{j}\}_{j=1}^{M}\sim pand an independent generated batch{xℓ′}ℓ=1N′∼qt\{x^{\prime}_{\ell}\}_{\ell=1}^{N^{\prime}}\sim q_{t}, define

p^σ​(x)=∑j=1Mbj​kσ​(x,yj),q^t,σ​(x)=∑ℓ=1N′aℓ′​kσ​(x,xℓ′).\widehat{p}_{\sigma}(x)=\sum_{j=1}^{M}b_{j}k_{\sigma}(x,y_{j}),\qquad\widehat{q}_{t,\sigma}(x)=\sum_{\ell=1}^{N^{\prime}}a^{\prime}_{\ell}k_{\sigma}(x,x^{\prime}_{\ell}).The KDE-based score estimators are

s^p​(x):=∇xlog⁡p^σ​(x)=∑j=1Mbj​∇xkσ​(x,yj)∑j=1Mbj​kσ​(x,yj),\widehat{s}_{p}(x):=\nabla_{x}\log\widehat{p}_{\sigma}(x)=\frac{\sum_{j=1}^{M}b_{j}\nabla_{x}k_{\sigma}(x,y_{j})}{\sum_{j=1}^{M}b_{j}k_{\sigma}(x,y_{j})},and

s^qt​(x):=∇xlog⁡q^t,σ​(x)=∑ℓ=1N′aℓ′​∇xkσ​(x,xℓ′)∑ℓ=1N′aℓ′​kσ​(x,xℓ′).\widehat{s}_{q_{t}}(x):=\nabla_{x}\log\widehat{q}_{t,\sigma}(x)=\frac{\sum_{\ell=1}^{N^{\prime}}a^{\prime}_{\ell}\nabla_{x}k_{\sigma}(x,x^{\prime}_{\ell})}{\sum_{\ell=1}^{N^{\prime}}a^{\prime}_{\ell}k_{\sigma}(x,x^{\prime}_{\ell})}.The particle estimator of the KL-induced velocity is

V^q^t,p^KL​(xi)=s^p​(xi)−s^qt​(xi).\widehat{V}_{\widehat{q}_{t},\widehat{p}}^{\mathrm{KL}}(x_{i})=\widehat{s}_{p}(x_{i})-\widehat{s}_{q_{t}}(x_{i}).For the Gaussian kernel, we have

∇xkσ​(x,y)=y−xσ2​kσ​(x,y).\nabla_{x}k_{\sigma}(x,y)=\frac{y-x}{\sigma^{2}}k_{\sigma}(x,y).Therefore,

s^p​(xi)=1σ2​(∑j=1Mbj​kσ​(xi,yj)​yj∑j=1Mbj​kσ​(xi,yj)−xi),\widehat{s}_{p}(x_{i})=\frac{1}{\sigma^{2}}\left(\frac{\sum_{j=1}^{M}b_{j}k_{\sigma}(x_{i},y_{j})y_{j}}{\sum_{j=1}^{M}b_{j}k_{\sigma}(x_{i},y_{j})}-x_{i}\right),and similarly

s^qt​(xi)=1σ2​(∑ℓ=1N′aℓ′​kσ​(xi,xℓ′)​xℓ′∑ℓ=1N′aℓ′​kσ​(xi,xℓ′)−xi).\widehat{s}_{q_{t}}(x_{i})=\frac{1}{\sigma^{2}}\left(\frac{\sum_{\ell=1}^{N^{\prime}}a^{\prime}_{\ell}k_{\sigma}(x_{i},x^{\prime}_{\ell})x^{\prime}_{\ell}}{\sum_{\ell=1}^{N^{\prime}}a^{\prime}_{\ell}k_{\sigma}(x_{i},x^{\prime}_{\ell})}-x_{i}\right).The particle-based estimate of the velocity field simplifies to the difference between two kernel-weighted local means:

V^q^t,p^KL​(xi)=1σ2​[∑j=1Mbj​kσ​(xi,yj)​yj∑j=1Mbj​kσ​(xi,yj)−∑ℓ=1N′aℓ′​kσ​(xi,xℓ′)​xℓ′∑ℓ=1N′aℓ′​kσ​(xi,xℓ′)].\widehat{V}_{\widehat{q}_{t},\widehat{p}}^{\mathrm{KL}}(x_{i})=\frac{1}{\sigma^{2}}\left[\frac{\sum_{j=1}^{M}b_{j}k_{\sigma}(x_{i},y_{j})y_{j}}{\sum_{j=1}^{M}b_{j}k_{\sigma}(x_{i},y_{j})}-\frac{\sum_{\ell=1}^{N^{\prime}}a^{\prime}_{\ell}k_{\sigma}(x_{i},x^{\prime}_{\ell})x^{\prime}_{\ell}}{\sum_{\ell=1}^{N^{\prime}}a^{\prime}_{\ell}k_{\sigma}(x_{i},x^{\prime}_{\ell})}\right].The corresponding particle update then follows:

xi←xi+η​V^q^t,p^KL​(xi).x_{i}\leftarrow x_{i}+\eta\widehat{V}_{\widehat{q}_{t},\widehat{p}}^{\mathrm{KL}}(x_{i}).

A.2More discussion on the estimators for self-transport

Figure 7:Evaluation of self-transport estimators on a 2D Gaussian mixture dataset featuring six dominant modes and two distant minority modes. “One-batch” refers to the single batch estimator with diagonal masking. “Two-batch” refers to our two-batch estimator with resampling.As introduced in Sec.3.3, estimating the self-interaction velocity fieldTq,qε​(x)T_{q,q}^{\varepsilon}(x)requires computing the OT plan of the model distributionqqwith itself. A naive empirical approximation based on a single mini-batch of sizeBB, i.e.,{xi}i=1B\{x_{i}\}_{i=1}^{B}, inherently suffers from a self-matching bias, as shown in Fig.3. Following[14], which addresses a similar issue, one may instead apply a diagonal mask to the cost matrix (settingCi​i=∞C_{ii}=\infty), which explicitly prevents any particle from matching with itself.

However, this masking strategy introduces a new source of bias that is sensitive to the batch sizeBB. For any given particlexix_{i}, the diagonal mask effectively forces it to transport its mass to an empirical target distribution composed of the remaining particles:

q^−i=1B−1​∑j≠iδxj.\widehat{q}_{-i}=\frac{1}{B-1}\sum_{j\neq i}\delta_{x_{j}}.Becausexix_{i}is excluded from its own target space, the particle is artificially forced to move a strictly positive distance towards its neighbors, even if it is already situated perfectly at a high-density peak of the true distributionqq. The magnitude of this structural bias scales proportionally to𝒪​(1/B)\mathcal{O}(1/B), effectively meaning that a fraction of the mass corresponding to roughly(B−1)/B(B-1)/Bis mapped sub-optimally. AsBBbecomes smaller, this forced redistribution dominates the velocity field, pushing particles away from their true local density centers (as observed in the left column of Fig.7).

By drawing an independent second batch{xl′}l=1B∼qt\{x_{l}^{\prime}\}_{l=1}^{B}\sim q_{t}, our two-batch estimator effectively avoids this issue. The OT plan is computed between two independent empirical measures, allowing a particlexix_{i}to match freely with a particlexl′x_{l}^{\prime}that happens to be infinitesimally close to it, without any artificial∞\infty-cost penalty. This yields a consistent estimate of the continuous self-transport map that remains robust even at small batch sizes.

A.3More discussions on velocity guidance

In this section, we first show that our proposed velocity guidance recovers the standard exponentially tilted target distribution when the energy functional is chosen as the KL divergence. We then show how this understanding is leveraged to construct the proposed velocity guidance in Eq. (16), which also serves as a conceptual justification for Eq. (16).

Letp(⋅|c)p(\cdot|c)denote the conditional data distribution,p(⋅|∅)p(\cdot|\varnothing)denote the unconditional data distribution, andw≥0w\geq 0be the guidance weight. In analogy with Eq. (16), we construct the velocity-guidance field by adding a term scaled bywwthat points from the unconditional to the conditional distribution:

V~tKL,w​(x):=(∇xlog⁡p​(x|c)−∇xlog⁡qt​(x|c))⏟Vqt(⋅|c),p(⋅|c)KL+w​(∇xlog⁡p​(x|c)−∇xlog⁡p​(x|∅))⏟Vqt(⋅|c),p(⋅|c)KL−Vqt(⋅|c),p(⋅|∅)KL.\widetilde{V}_{t}^{\mathrm{KL},w}(x):=\underbrace{\Big(\nabla_{x}\log p(x|c)-\nabla_{x}\log q_{t}(x|c)\Big)}_{V^{\mathrm{KL}}_{q_{t}(\cdot|c),p(\cdot|c)}}+w\underbrace{\Big(\nabla_{x}\log p(x|c)-\nabla_{x}\log p(x|\varnothing)\Big)}_{V^{\mathrm{KL}}_{q_{t}(\cdot|c),p(\cdot|c)}-V^{\mathrm{KL}}_{q_{t}(\cdot|c),p(\cdot|\varnothing)}}.(22) We formalize the connection between this constructed velocity field and CFG in the following proposition.

Proposition A.1.

The velocity-guidance fieldV~tKL,w\widetilde{V}_{t}^{\mathrm{KL},w}in Eq.(22)is exactly the WGF velocity field of the combined energy functional

ℱKLw(qt(⋅|c))=(1+w)DKL(qt(⋅|c)∥p(⋅|c))−wDKL(qt(⋅|c)∥p(⋅|∅)).\mathcal{F}_{\mathrm{KL}}^{w}(q_{t}(\cdot|c))=(1+w)D_{\mathrm{KL}}\big(q_{t}(\cdot|c)\|p(\cdot|c)\big)-wD_{\mathrm{KL}}\big(q_{t}(\cdot|c)\|p(\cdot|\varnothing)\big).Furthermore, the unique global minimizer ofℱw\mathcal{F}_{w}is the exponentially tilted target distributionpw(⋅|c)p_{w}(\cdot|c), defined as

pw(x|c):=1Zw​(c)p(x|c)(p​(x|c)p​(x|∅))w,p_{w}(x|c):=\frac{1}{Z_{w}(c)}p(x|c)\left(\frac{p(x|c)}{p(x|\varnothing)}\right)^{w},whereZw​(c)Z_{w}(c)is the normalizing constant.

Proof.

We first expand the combined energy functionalℱKLw(qt(⋅|c))\mathcal{F}_{\mathrm{KL}}^{w}(q_{t}(\cdot|c)):

ℱKLw​(qt)\displaystyle\mathcal{F}_{\mathrm{KL}}^{w}(q_{t})=(1+w)​∫qt​log⁡qtp(⋅|c)−w​∫qt​log⁡qtp(⋅|∅)\displaystyle=(1+w)\int q_{t}\log\frac{q_{t}}{p(\cdot|c)}-w\int q_{t}\log\frac{q_{t}}{p(\cdot|\varnothing)}=∫qtlogqt−∫qtlog(p(⋅|c)1+wp(⋅|∅)−w)\displaystyle=\int q_{t}\log q_{t}-\int q_{t}\log\Big(p(\cdot|c)^{1+w}p(\cdot|\varnothing)^{-w}\Big)=∫qt​log⁡qtpw(⋅|c)−log⁡Zw​(c)\displaystyle=\int q_{t}\log\frac{q_{t}}{p_{w}(\cdot|c)}-\log Z_{w}(c)=DKL(qt(⋅|c)∥pw(⋅|c))−logZw(c).\displaystyle=D_{\mathrm{KL}}\big(q_{t}(\cdot|c)\|p_{w}(\cdot|c)\big)-\log Z_{w}(c).Sincelog⁡Zw​(c)\log Z_{w}(c)is a constant with respect to the distributionqtq_{t}, it is evident thatℱKLw(qt(⋅|c))\mathcal{F}_{\mathrm{KL}}^{w}(q_{t}(\cdot|c))achieves its unique global minimum whenqt(⋅|c)=pw(⋅|c)q_{t}(\cdot|c)=p_{w}(\cdot|c), which proves the second part of the proposition.

To derive the induced WGF, we compute the first variation ofℱKLw(qt(⋅|c))\mathcal{F}_{\mathrm{KL}}^{w}(q_{t}(\cdot|c)). Because the additive constant−log⁡Zw​(c)-\log Z_{w}(c)vanishes upon differentiation, the first variation is identical to that ofDKL(qt(⋅|c)∥pw(⋅|c))D_{\mathrm{KL}}\big(q_{t}(\cdot|c)\|p_{w}(\cdot|c)\big):

δ​ℱKLwδ​q(qt(⋅|c))(x)=logqt(x|c)−logpw(x|c)+const.\frac{\delta\mathcal{F}_{\mathrm{KL}}^{w}}{\delta q}(q_{t}(\cdot|c))(x)=\log q_{t}(x|c)-\log p_{w}(x|c)+\mathrm{const.}The corresponding Wasserstein velocity field is the negative spatial gradient of the first variation:

−∇xδ​ℱKLwδ​q(qt(⋅|c))(x)\displaystyle-\nabla_{x}\frac{\delta\mathcal{F}_{\mathrm{KL}}^{w}}{\delta q}(q_{t}(\cdot|c))(x)=∇xlog⁡pw​(x|c)−∇xlog⁡qt​(x|c)\displaystyle=\nabla_{x}\log p_{w}(x|c)-\nabla_{x}\log q_{t}(x|c)=∇x((1+w)​log⁡p​(x|c)−w​log⁡p​(x|∅)−log⁡Zw​(c))−∇xlog⁡qt​(x|c)\displaystyle=\nabla_{x}\Big((1+w)\log p(x|c)-w\log p(x|\varnothing)-\log Z_{w}(c)\Big)-\nabla_{x}\log q_{t}(x|c)=(1+w)​∇xlog⁡p​(x|c)−w​∇xlog⁡p​(x|∅)−∇xlog⁡qt​(x|c).\displaystyle=(1+w)\nabla_{x}\log p(x|c)-w\nabla_{x}\log p(x|\varnothing)-\nabla_{x}\log q_{t}(x|c).Rearranging the terms, we obtain:

−∇xδ​ℱKLwδ​q(qt(⋅|c))(x)=(∇xlogp(x|c)−∇xlogqt(x|c))+w(∇xlogp(x|c)−∇xlogp(x|∅)),-\nabla_{x}\frac{\delta\mathcal{F}_{\mathrm{KL}}^{w}}{\delta q}(q_{t}(\cdot|c))(x)=\Big(\nabla_{x}\log p(x|c)-\nabla_{x}\log q_{t}(x|c)\Big)+w\Big(\nabla_{x}\log p(x|c)-\nabla_{x}\log p(x|\varnothing)\Big),which is exactly the constructed velocity-guidance fieldV~tKL,w​(x)\widetilde{V}_{t}^{\mathrm{KL},w}(x)defined in Eq. (22). ∎

Therefore, under the KL energy, our velocity-guidance formulation is precisely equivalent to running the WGF toward the exponentially tilted distributionpw​(x|c)p_{w}(x|c). This corresponds exactly to the standard CFG target form, where the conditional distribution is amplified relative to the unconditional distribution by the density ratiop​(x|c)/p​(x|∅)p(x|c)/p(x|\varnothing). The results in AppendixEverify that velocity guidance yields enhanced performance over distribution guidance when KL divergence is leveraged as the energy functional.

Figure 8:Comparison of velocity guidance and distribution guidance for conditional generation on a 2D 3-mode Gaussian mixture dataset.##### Connection to the Sinkhorn velocity guidance.

For the Sinkhorn divergence, the same closed-form target distribution is generally unavailable, because the velocity field is expressed through barycentric projections rather than score functions. Nevertheless, the KL case shows that adding the velocity-level correction

Vqt(⋅|c),p(⋅|c)−Vqt(⋅|c),p(⋅|∅)V_{q_{t}(\cdot|c),p(\cdot|c)}-V_{q_{t}(\cdot|c),p(\cdot|\varnothing)}has the correct interpretation in the score-based limit: it changes the effective target fromp​(x|c)p(x|c)to the exponentially tilted distributionp​(x|c)​(p​(x|c)/p​(x|∅))wp(x|c)(p(x|c)/p(x|\varnothing))^{w}. This motivates applying the analogous guidance design to the Sinkhorn-induced velocity field. Specifically, we have

V~qθ,pε,w​(xi)=(Tqθ(⋅|c),p(⋅|c)ε​(xi)−Tqθ(⋅|c),qθ(⋅|c)ε​(xi))⏟Vqθ(⋅|c),p(⋅|c)ϵ+w​(Tqθ(⋅|c),p(⋅|c)ε​(xi)−Tqθ(⋅|c),p(⋅|∅)ε​(xi))⏟Vqθ(⋅|c),p(⋅|c)ϵ−Vqθ(⋅|c),p(⋅|∅)ϵ,\displaystyle\tilde{V}^{\varepsilon,w}_{q_{\theta},p}(x_{i})\!=\!\underbrace{\left(T^{\varepsilon}_{q_{\theta}(\cdot|c),p(\cdot|c)}(x_{i})\!-\!T^{\varepsilon}_{q_{\theta}(\cdot|c),q_{\theta}(\cdot|c)}(x_{i})\right)}_{V^{\epsilon}_{q_{\theta}(\cdot|c),p(\cdot|c)}}\!+\!w\underbrace{\left(T^{\varepsilon}_{q_{\theta}(\cdot|c),p(\cdot|c)}(x_{i})\!-\!T^{\varepsilon}_{q_{\theta}(\cdot|c),p(\cdot|\varnothing)}(x_{i})\right)}_{V^{\epsilon}_{q_{\theta}(\cdot|c),p(\cdot|c)}-V^{\epsilon}_{q_{\theta}(\cdot|c),p(\cdot|\varnothing)}},where

Vqθ(⋅|c),p(⋅|c)ϵ​(xi)\displaystyle V^{\epsilon}_{q_{\theta}(\cdot|c),p(\cdot|c)}(x_{i})=Tqθ(⋅|c),p(⋅|c)ε​(xi)−Tqθ(⋅|c),qθ(⋅|c)ε​(xi),\displaystyle=T^{\varepsilon}_{q_{\theta}(\cdot|c),p(\cdot|c)}(x_{i})\!-\!T^{\varepsilon}_{q_{\theta}(\cdot|c),q_{\theta}(\cdot|c)}(x_{i}),(23)Vqθ(⋅|c),p(⋅|∅)ϵ​(xi)\displaystyle V^{\epsilon}_{q_{\theta}(\cdot|c),p(\cdot|\varnothing)}(x_{i})=Tqθ(⋅|c),p(⋅|∅)ε​(xi)−Tqθ(⋅|c),qθ(⋅|c)ε​(xi),\displaystyle=T^{\varepsilon}_{q_{\theta}(\cdot|c),p(\cdot|\varnothing)}(x_{i})\!-\!T^{\varepsilon}_{q_{\theta}(\cdot|c),q_{\theta}(\cdot|c)}(x_{i}),(24)which recovers Eq. (16).

Toy illustration.To empirically compare our proposed velocity guidance with distribution guidance, we evaluate conditional generation on a 2D 3-mode Gaussian mixtures. Here, each mode is treated as a distinct class. We train a generatorx=fθ​(z,c,w)x=f_{\theta}(z,c,w)where the guidance scalewwis randomly sampled during training.

As illustrated in Fig.8(top row), velocity guidance reflects the behavior of an exponentially tilted target distributionpw​(x|c)p_{w}(x|c). Aswwincreases, the probability mass is amplified at the highest-density regions of the conditional mode. In contrast, the distribution guidance implicitly targets a linear mixture between the conditional and unconditional distributions. Fig.8(bottom row) demonstrates that increasingwwcauses particles to scatter away from the modes.

A.4More discussions on related work

Our proposed W-Flow framework is related to several prominent directions in generative modeling. We highlight the key distinctions below.

Figure 9:Illustrations of the difference in the velocity field computation between Drifting Model (Alg. 2 in[14]) and our approach. “⋅\sqrt{\cdot}” refers to the element-wise square root, and thegray blocksrefer to the diagonal mask adopted in[14]. Drifting Model adopts a heuristic construction with separate row-wise and column-wise Softmax normalizations. Our approach leverages Sinkhorn iteration to compute the barycentric projection.Connection to Drifting Model.Our formulation shares the forward, iterative pushforward structure with Drifting Model[14]. The main difference lies in the construction of the velocity field. Drifting Model employs heuristic attraction and repulsion terms (Vp+V_{p}^{+}andVq−V_{q}^{-}) based on localized kernel interactions, resulting in a mean-shift-type dynamics driven by local proximity. Such local updates can lead to unstable behavior and require careful tuning to avoid mode collapse.

In contrast, W-Flow derives the velocity field from a WGF and instantiates it via Sinkhorn divergence. The resulting dynamics are governed by global OT couplings, which enforce marginal constraints and coordinate particle movement at the distribution level. This yields a principled and globally consistent transport mechanism.

Distinction from JKO-based WGF models.Another line of work formulates generative processes as WGFs using the JKO scheme[11], which corresponds to an implicit time discretization. However, each JKO step requires solving a nontrivial optimization problem, which in practice often necessitates adversarial training, thereby introducing training instability and limiting scalability.

Distinct from[11], W-Flow adopts an explicit Euler discretization of the flow. By computing the Sinkhorn velocity fieldVεV^{\varepsilon}via efficient Sinkhorn scaling, the learning problem reduces to a simple regression objective, avoiding inner optimization loops and eliminating the need for adversarial training. This enables scalable learning of a one-step pushforward map.

Distinction from Sinkhorn GANs.The use of EOT in generative modeling was popularized by Sinkhorn GAN[20], which directly employs the Sinkhorn divergence as the target loss function. This approach requires auto-differentiating through the Sinkhorn iterations and backpropagating gradients through the transport plan, inducing numerical instability and memory overhead.

W-Flow instead interprets the Sinkhorn divergence as an energy functional that defines the velocity field of the dynamics. The transport plan is treated as an empirical oracle, without backpropagation through Sinkhorn iterations. This avoids instability and allows the method to scale effectively.

Distinction from moment matching.Generative Moment Matching Networks (GMMNs)[37]learn generators through particle updates driven by the squared MMD, which can be viewed as an instance of the MMD-based variant of W-Flow in Table1. However, MMD relies on kernel-based interactions between particles, which can suffer from the curse of dimensionality and yield vanishing gradients when the current pushforward distribution is far from the target, as also reflected in our ablation experiments. Inductive Moment Matching[74]performs kernel-based moment matching along a prescribed static interpolation between the prior and target distributions. As a result, it inherits the same kernel-related limitations and typically requires multiple pushforward sampling steps to simulate transitions along the interpolator. W-Flow instead leverages the OT-based dynamics and directly parametrizes a one-step map that does not require any intermediate simulations at inference time.

Appendix BProofs

In this section, we present the detailed statements of assumptions and theorems discussed in the main paper, along with their proofs.

In the main text, Theorem3.1is stated for the Sinkhorn velocityVq,pεV_{q,p}^{\varepsilon}defined in (10). For generality, here we prove the result under AssumptionB.1for a generic velocity field

V:𝒫2​(ℝd)×𝒫2​(ℝd)×ℝd→ℝd,(q,p,x)↦Vq,p​(x),V:\mathcal{P}_{2}(\mathbb{R}^{d})\times\mathcal{P}_{2}(\mathbb{R}^{d})\times\mathbb{R}^{d}\to\mathbb{R}^{d},\qquad(q,p,x)\mapsto V_{q,p}(x),whereqqis the evolving source distribution andppis the fixed target distribution. PropositionB.7then verifies the assumption for the Sinkhorn velocity whenq0q_{0}andpphave bounded support.

B.1Complete statement and the proof of Theorem3.1

ForN∈ℕN\in\mathbb{N}and step sizeη>0\eta>0, let particles{xi(k)}i=1N\{x_{i}^{(k)}\}_{i=1}^{N}evolve according to

xi(k+1)=xi(k)+η​Vq^kN,p^M​(xi(k)),q^kN:=1N​∑i=1Nδxi(k),tk:=k​η,x_{i}^{(k+1)}=x_{i}^{(k)}+\eta\,V_{\widehat{q}_{k}^{N},\widehat{p}^{M}}(x_{i}^{(k)}),\qquad\widehat{q}_{k}^{N}:=\frac{1}{N}\sum_{i=1}^{N}\delta_{x_{i}^{(k)}},\qquad t_{k}:=k\eta,(25)with the initial data pointsxi(0)x_{i}^{(0)}sampled i.i.d. fromq0q_{0}.

Define the piecewise-linear interpolationxiη:[0,∞)→ℝdx_{i}^{\eta}:[0,\infty)\to\mathbb{R}^{d}by

xiη​(t):=xi(k)+(t−tk)​Vq^kN,p^M​(xi(k)),t∈[tk,tk+1),i=1,2,⋯,N,x_{i}^{\eta}(t):=x_{i}^{(k)}+(t-t_{k})\,V_{\widehat{q}_{k}^{N},\widehat{p}^{M}}(x_{i}^{(k)}),\quad t\in[t_{k},t_{k+1}),\quad i=1,2,\cdots,N,(26)and the associated empirical measure

q^tN,M,η:=1N​∑i=1Nδxiη​(t).\widehat{q}_{t}^{N,M,\eta}:=\frac{1}{N}\sum_{i=1}^{N}\delta_{x_{i}^{\eta}(t)}. We study the limiting behavior of (26) under a joint asymptotic regime in whichN,M→∞N,M\to\inftyandη→0\eta\to 0.

Letm2​(μ)m_{2}(\mu)be the root second moment ofμ∈𝒫2​(ℝd)\mu\in\mathcal{P}_{2}(\mathbb{R}^{d}):

m2​(μ):=(∫ℝd‖x‖2​𝑑μ​(x))1/2.m_{2}(\mu):=\left(\int_{\mathbb{R}^{d}}\|x\|^{2}\,d\mu(x)\right)^{1/2}.

Assumption B.1.

Assume thatVq,pV_{q,p}is Borel measurable and belongs toL2​(q)L^{2}(q)for everyq,p∈𝒫2​(ℝd)q,p\in\mathcal{P}_{2}(\mathbb{R}^{d}). Assume that there exist constantsL0,L1,L2,L3<∞L_{0},L_{1},L_{2},L_{3}<\inftysuch that, for allq,q~,p,p~∈𝒫2​(ℝd)q,\widetilde{q},p,\widetilde{p}\in\mathcal{P}_{2}(\mathbb{R}^{d}),

‖Vq,p‖L2​(q)\displaystyle\|V_{q,p}\|_{L^{2}(q)}≤\displaystyle\leqL0​(m2​(q)+m2​(p)),\displaystyle L_{0}\bigl(m_{2}(q)+m_{2}(p)\bigr),(27)‖Vq,p​(x)−Vq,p​(x′)‖\displaystyle\|V_{q,p}(x)-V_{q,p}(x^{\prime})\|≤\displaystyle\leqL1​‖x−x′‖,x,x′∈ℝd,\displaystyle L_{1}\|x-x^{\prime}\|,\qquad x,x^{\prime}\in\mathbb{R}^{d},(28)‖Vq,p−Vq~,p‖L2​(q~)\displaystyle\|V_{q,p}-V_{\widetilde{q},p}\|_{L^{2}(\widetilde{q})}≤\displaystyle\leqL2​𝒲2​(q,q~),\displaystyle L_{2}\mathcal{W}_{2}(q,\widetilde{q}),(29)‖Vq,p−Vq,p~‖L2​(q)\displaystyle\|V_{q,p}-V_{q,\widetilde{p}}\|_{L^{2}(q)}≤\displaystyle\leqL3​𝒲2​(p,p~).\displaystyle L_{3}\mathcal{W}_{2}(p,\widetilde{p}).(30)

Under these conditions, one can establish a quantitative convergence result showing that the interpolated empirical measureq^tN,M,η\widehat{q}_{t}^{N,M,\eta}converges to the unique solution of a nonlinear PDE asN,M→∞N,M\to\inftyandη→0\eta\to 0.

Now we present theformal versionof Theorem3.1below, followed by the detailed proof.

Theorem B.4.

Assume AssumptionB.1holds andq0,p∈𝒫2​(ℝd)q_{0},p\in\mathcal{P}_{2}(\mathbb{R}^{d}). For everyT>0T>0, there exists a uniqueq∈C​([0,T];𝒫2​(ℝd))q\in C([0,T];\mathcal{P}_{2}(\mathbb{R}^{d}))solving the nonlinear continuity equation

∂tqt+∇⋅(qt​Vqt,p)=0\partial_{t}q_{t}+\nabla\cdot\big(q_{t}\,V_{q_{t},p}\big)=0(31)in the weak sense

∫ℝdφ​𝑑qt−∫ℝdφ​𝑑q0=∫0t∫ℝd∇φ​(x)⋅Vqs,p​(x)​𝑑qs​(x)​𝑑s,∀φ∈Cc1​(ℝd).\int_{\mathbb{R}^{d}}\varphi\,dq_{t}-\int_{\mathbb{R}^{d}}\varphi\,dq_{0}=\int_{0}^{t}\int_{\mathbb{R}^{d}}\nabla\varphi(x)\cdot V_{q_{s},p}(x)\,dq_{s}(x)\,ds,\qquad\forall\varphi\in C_{c}^{1}(\mathbb{R}^{d}).Moreover, for allN,MN,Mand allη>0\eta>0,

supt∈[0,T]𝒲2​(q^tN,M,η,qt)\displaystyle\sup_{t\in[0,T]}\mathcal{W}_{2}\big(\widehat{q}_{t}^{N,M,\eta},q_{t}\big)≤e(L1+L2)​T​𝒲2​(q^0N,q0)\displaystyle\leq e^{(L_{1}+L_{2})T}\mathcal{W}_{2}(\widehat{q}_{0}^{N},q_{0})+L3L1+L2​(e(L1+L2)​T−1)​𝒲2​(p^M,p)\displaystyle\quad+\frac{L_{3}}{L_{1}+L_{2}}\bigl(e^{(L_{1}+L_{2})T}-1\bigr)\mathcal{W}_{2}(\widehat{p}^{M},p)+CT​η​[m2​(q^0N)+m2​(p^M)],\displaystyle\quad+C_{T}\,\eta\bigl[m_{2}(\widehat{q}_{0}^{N})+m_{2}(\widehat{p}^{M})\bigr],whereCT<∞C_{T}<\inftydepends only onT,L0,L1,L2T,L_{0},L_{1},L_{2}. Consequently, asη→0\eta\rightarrow 0andN,M→∞N,M\to\infty,

supt∈[0,T]𝒲2​(q^tN,M,η,qt)→0,almost surely.\sup_{t\in[0,T]}\mathcal{W}_{2}\big(\widehat{q}_{t}^{N,M,\eta},q_{t}\big)\to 0,\qquad\text{almost surely.}(32)

Proof.

For generic𝒙=(x1,…,xN)∈(ℝd)N\boldsymbol{x}=(x_{1},\ldots,x_{N})\in(\mathbb{R}^{d})^{N}, define

q^N:=1N​∑i=1Nδxi,VN,M​(𝒙):=(Vq^N,p^M​(x1),…,Vq^N,p^M​(xN)).\displaystyle\widehat{q}^{N}:=\frac{1}{N}\sum_{i=1}^{N}\delta_{x_{i}},\quad V^{N,M}(\boldsymbol{x}):=\bigl(V_{\widehat{q}^{N},\widehat{p}^{M}}(x_{1}),\ldots,V_{\widehat{q}^{N},\widehat{p}^{M}}(x_{N})\bigr).Equip(ℝd)N(\mathbb{R}^{d})^{N}with the norm

‖𝒙‖2,N:=(1N​∑i=1N‖xi‖2)1/2.\|\boldsymbol{x}\|_{2,N}:=\left(\frac{1}{N}\sum_{i=1}^{N}\|x_{i}\|^{2}\right)^{1/2}.Let𝒙′=(x1′,…,xN′)\boldsymbol{x}^{\prime}=(x^{\prime}_{1},\ldots,x^{\prime}_{N}), and setq^N⁣′:=N−1​∑i=1Nδxi′\widehat{q}^{N\prime}:=N^{-1}\sum_{i=1}^{N}\delta_{x^{\prime}_{i}}. By (28) and (29),

‖VN,M​(𝒙)−VN,M​(𝒙′)‖2,N\displaystyle\|V^{N,M}(\boldsymbol{x})-V^{N,M}(\boldsymbol{x}^{\prime})\|_{2,N}≤(1N​∑i=1N‖Vq^N,p^M​(xi)−Vq^N,p^M​(xi′)‖2)1/2\displaystyle\leq\left(\frac{1}{N}\sum_{i=1}^{N}\bigl\|V_{\widehat{q}^{N},\widehat{p}^{M}}(x_{i})-V_{\widehat{q}^{N},\widehat{p}^{M}}(x^{\prime}_{i})\bigr\|^{2}\right)^{1/2}+(1N​∑i=1N‖Vq^N,p^M​(xi′)−Vq^N⁣′,p^M​(xi′)‖2)1/2\displaystyle\quad+\left(\frac{1}{N}\sum_{i=1}^{N}\bigl\|V_{\widehat{q}^{N},\widehat{p}^{M}}(x^{\prime}_{i})-V_{\widehat{q}^{N\prime},\widehat{p}^{M}}(x^{\prime}_{i})\big\|^{2}\right)^{1/2}≤L1​‖𝒙−𝒙′‖2,N+L2​𝒲2​(q^N,q^N⁣′)\displaystyle\leq L_{1}\|\boldsymbol{x}-\boldsymbol{x}^{\prime}\|_{2,N}+L_{2}\mathcal{W}_{2}(\widehat{q}^{N},\widehat{q}^{N\prime})≤(L1+L2)​‖𝒙−𝒙′‖2,N,\displaystyle\leq(L_{1}+L_{2})\|\boldsymbol{x}-\boldsymbol{x}^{\prime}\|_{2,N},where

𝒲2​(q^N,q^N⁣′)≤(1N​∑j=1N‖xj−xj′‖2)1/2=‖𝒙−𝒙′‖2,N\mathcal{W}_{2}(\widehat{q}^{N},\widehat{q}^{N\prime})\leq\left(\frac{1}{N}\sum_{j=1}^{N}\|x_{j}-x_{j}^{\prime}\|^{2}\right)^{1/2}=\|\boldsymbol{x}-\boldsymbol{x}^{\prime}\|_{2,N}by couplingxix_{i}withxi′x^{\prime}_{i}. Moreover, by (27),

‖VN,M​(𝒙)‖2,N=‖Vq^N,p^M‖L2​(q^N)≤L0​(m2​(q^N)+m2​(p^M)).\|V^{N,M}(\boldsymbol{x})\|_{2,N}=\|V_{\widehat{q}^{N},\widehat{p}^{M}}\|_{L^{2}(\widehat{q}^{N})}\leq L_{0}\bigl(m_{2}(\widehat{q}^{N})+m_{2}(\widehat{p}^{M})\bigr).Thus the finite-particle vector field is globally Lipschitz and has linear growth in∥⋅∥2,N\|\cdot\|_{2,N}.

Consider now the ODE

𝒙˙N,M​(t)=VN,M​(𝒙N,M​(t)),𝒙N,M​(0)=(x1(0),…,xN(0)),\dot{\boldsymbol{x}}^{N,M}(t)=V^{N,M}(\boldsymbol{x}^{N,M}(t)),\qquad\boldsymbol{x}^{N,M}(0)=(x^{(0)}_{1},\dots,x^{(0)}_{N}),(33)with the initializationxi(0)x^{(0)}_{i}defined below (25). SinceVN,MV^{N,M}is globally Lipschitz in∥⋅∥2,N\|\cdot\|_{2,N}, it holds that (33) has a unique global solution by the Picard–Lindelöf theorem.

Let

q^tN,M:=1N​∑i=1NδxiN,M​(t).\widehat{q}_{t}^{N,M}:=\frac{1}{N}\sum_{i=1}^{N}\delta_{x_{i}^{N,M}(t)}.For anyφ∈Cc1​(ℝd)\varphi\in C_{c}^{1}(\mathbb{R}^{d}),

dd​t​∫φ​𝑑q^tN,M\displaystyle\frac{d}{dt}\int\varphi\,d\widehat{q}_{t}^{N,M}=\displaystyle=dd​t​1N​∑i=1Nφ​(xiN,M​(t))=1N​∑i=1N∇φ​(xiN,M​(t))⋅x˙iN,M​(t)\displaystyle\frac{d}{dt}\frac{1}{N}\sum_{i=1}^{N}\varphi(x_{i}^{N,M}(t))=\frac{1}{N}\sum_{i=1}^{N}\nabla\varphi(x_{i}^{N,M}(t))\cdot\dot{x}_{i}^{N,M}(t)=\displaystyle=∫∇φ​(x)⋅Vq^tN,M,p^M​(x)​𝑑q^tN,M​(x).\displaystyle\int\nabla\varphi(x)\cdot V_{\widehat{q}_{t}^{N,M},\widehat{p}^{M}}(x)\,d\widehat{q}_{t}^{N,M}(x).Thusq^tN,M\widehat{q}_{t}^{N,M}is a weak solution of

∂tq^tN,M+∇⋅(q^tN,M​Vq^tN,M,p^M)=0.\partial_{t}\widehat{q}_{t}^{N,M}+\nabla\cdot\left(\widehat{q}_{t}^{N,M}V_{\widehat{q}_{t}^{N,M},\widehat{p}^{M}}\right)=0. We next derive the stability estimate. Letμ,ν∈C​([0,T];𝒫2​(ℝd))\mu,\nu\in C([0,T];\mathcal{P}_{2}(\mathbb{R}^{d}))be weak solutions of

∂tμt+∇⋅(μt​Vμt,p)=0,∂tνt+∇⋅(νt​Vνt,p~)=0,\partial_{t}\mu_{t}+\nabla\cdot(\mu_{t}V_{\mu_{t},p})=0,\qquad\partial_{t}\nu_{t}+\nabla\cdot(\nu_{t}V_{\nu_{t},\widetilde{p}})=0,wherep,p~∈𝒫2​(ℝd)p,\widetilde{p}\in\mathcal{P}_{2}(\mathbb{R}^{d})are fixed. First note that (27) and (28) imply the pointwise linear-growth bound

‖Vr,a​(x)‖≤L1​‖x‖+(L0+L1)​m2​(r)+L0​m2​(a),r,a∈𝒫2​(ℝd).\|V_{r,a}(x)\|\leq L_{1}\|x\|+(L_{0}+L_{1})m_{2}(r)+L_{0}m_{2}(a),\qquad r,a\in\mathcal{P}_{2}(\mathbb{R}^{d}).(34)Indeed,

‖Vr,a​(0)‖≤(∫‖Vr,a​(0)−Vr,a​(z)‖2​𝑑r​(z))1/2+‖Vr,a‖L2​(r)≤L1​m2​(r)+L0​(m2​(r)+m2​(a)).\|V_{r,a}(0)\|\leq\left(\int\|V_{r,a}(0)-V_{r,a}(z)\|^{2}\,dr(z)\right)^{1/2}+\|V_{r,a}\|_{L^{2}(r)}\leq L_{1}m_{2}(r)+L_{0}\bigl(m_{2}(r)+m_{2}(a)\bigr).Sinceμ,ν∈C​([0,T];𝒫2​(ℝd))\mu,\nu\in C([0,T];\mathcal{P}_{2}(\mathbb{R}^{d})), for these fixed solution trajectories we have

MT:=supt∈[0,T](m2​(μt)+m2​(νt))<∞.M_{T}:=\sup_{t\in[0,T]}\bigl(m_{2}(\mu_{t})+m_{2}(\nu_{t})\bigr)<\infty.Letγ0∈Π​(μ0,ν0)\gamma_{0}\in\Pi(\mu_{0},\nu_{0})be optimal for𝒲2​(μ0,ν0)\mathcal{W}_{2}(\mu_{0},\nu_{0}). For eachx0,y0∈ℝdx_{0},y_{0}\in\mathbb{R}^{d}, defineX​(⋅;x0)X(\cdot;x_{0})andY​(⋅;y0)Y(\cdot;y_{0})as the solutions on[0,T][0,T]of

{dd​t​X​(t;x0)=Vμt,p​(X​(t;x0)),X​(0;x0)=x0,{dd​t​Y​(t;y0)=Vνt,p~​(Y​(t;y0)),Y​(0;y0)=y0.\begin{cases}\dfrac{d}{dt}X(t;x_{0})=V_{\mu_{t},p}(X(t;x_{0})),\\ X(0;x_{0})=x_{0},\end{cases}\qquad\begin{cases}\dfrac{d}{dt}Y(t;y_{0})=V_{\nu_{t},\widetilde{p}}(Y(t;y_{0})),\\ Y(0;y_{0})=y_{0}.\end{cases}The global well-posedness of these ODEs on[0,T][0,T]follows from (34) and (28), sinceμ,ν∈C​([0,T];𝒫2​(ℝd))\mu,\nu\in C([0,T];\mathcal{P}_{2}(\mathbb{R}^{d})). Sinceμt\mu_{t}andνt\nu_{t}solve the corresponding continuity equations, their characteristic representations give

X​(t;⋅)#​μ0=μt,Y​(t;⋅)#​ν0=νt.X(t;\cdot)_{\#}\mu_{0}=\mu_{t},\qquad Y(t;\cdot)_{\#}\nu_{0}=\nu_{t}.Therefore

(X​(t;⋅),Y​(t;⋅))#​γ0∈Π​(μt,νt).\bigl(X(t;\cdot),Y(t;\cdot)\bigr)_{\#}\gamma_{0}\in\Pi(\mu_{t},\nu_{t}).Define

D​(t):=(∫ℝd×ℝd‖X​(t;x0)−Y​(t;y0)‖2​𝑑γ0​(x0,y0))1/2.D(t):=\left(\int_{\mathbb{R}^{d}\times\mathbb{R}^{d}}\|X(t;x_{0})-Y(t;y_{0})\|^{2}\,d\gamma_{0}(x_{0},y_{0})\right)^{1/2}.Since(X​(t;⋅),Y​(t;⋅))#​γ0∈Π​(μt,νt)\bigl(X(t;\cdot),Y(t;\cdot)\bigr)_{\#}\gamma_{0}\in\Pi(\mu_{t},\nu_{t}), we have

𝒲2​(μt,νt)≤D​(t).\mathcal{W}_{2}(\mu_{t},\nu_{t})\leq D(t). Forδ>0\delta>0, set

Dδ​(t):=(D​(t)2+δ)1/2.D_{\delta}(t):=(D(t)^{2}+\delta)^{1/2}.For a.e.tt,

dd​t​Dδ​(t)\displaystyle\frac{d}{dt}D_{\delta}(t)=1Dδ​(t)​∫(X​(t;x0)−Y​(t;y0))⋅[Vμt,p​(X​(t;x0))−Vνt,p~​(Y​(t;y0))]​𝑑γ0​(x0,y0)\displaystyle=\frac{1}{D_{\delta}(t)}\int\bigl(X(t;x_{0})-Y(t;y_{0})\bigr)\cdot\Bigl[V_{\mu_{t},p}(X(t;x_{0}))-V_{\nu_{t},\widetilde{p}}(Y(t;y_{0}))\Bigr]\,d\gamma_{0}(x_{0},y_{0})≤(∫‖Vμt,p​(X​(t;x0))−Vνt,p~​(Y​(t;y0))‖2​𝑑γ0​(x0,y0))1/2\displaystyle\leq\left(\int\Bigl\|V_{\mu_{t},p}(X(t;x_{0}))-V_{\nu_{t},\widetilde{p}}(Y(t;y_{0}))\Bigr\|^{2}\,d\gamma_{0}(x_{0},y_{0})\right)^{1/2}≤(∫‖Vμt,p​(X​(t;x0))−Vμt,p​(Y​(t;y0))‖2​𝑑γ0)1/2\displaystyle\leq\left(\int\Bigl\|V_{\mu_{t},p}(X(t;x_{0}))-V_{\mu_{t},p}(Y(t;y_{0}))\Bigl\|^{2}\,d\gamma_{0}\right)^{1/2}+(∫‖Vμt,p​(Y​(t;y0))−Vνt,p​(Y​(t;y0))‖2​𝑑γ0)1/2\displaystyle\quad+\left(\int\Bigl\|V_{\mu_{t},p}(Y(t;y_{0}))-V_{\nu_{t},p}(Y(t;y_{0}))\Bigr\|^{2}\,d\gamma_{0}\right)^{1/2}+(∫‖Vνt,p​(Y​(t;y0))−Vνt,p~​(Y​(t;y0))‖2​𝑑γ0)1/2\displaystyle\quad+\left(\int\Bigl\|V_{\nu_{t},p}(Y(t;y_{0}))-V_{\nu_{t},\widetilde{p}}(Y(t;y_{0}))\Bigr\|^{2}\,d\gamma_{0}\right)^{1/2}≤L1​D​(t)+‖Vμt,p−Vνt,p‖L2​(νt)+‖Vνt,p−Vνt,p~‖L2​(νt)\displaystyle\leq L_{1}D(t)+\|V_{\mu_{t},p}-V_{\nu_{t},p}\|_{L^{2}(\nu_{t})}+\|V_{\nu_{t},p}-V_{\nu_{t},\widetilde{p}}\|_{L^{2}(\nu_{t})}≤(L1+L2)​D​(t)+L3​𝒲2​(p,p~).\displaystyle\leq(L_{1}+L_{2})D(t)+L_{3}\mathcal{W}_{2}(p,\widetilde{p}).Here we used that the law ofY​(t;y0)Y(t;y_{0})underγ0\gamma_{0}isνt\nu_{t}, together with𝒲2​(μt,νt)≤D​(t)\mathcal{W}_{2}(\mu_{t},\nu_{t})\leq D(t). Integrating the preceding inequality and lettingδ↓0\delta\downarrow 0gives

D​(t)≤D​(0)+∫0t((L1+L2)​D​(s)+L3​𝒲2​(p,p~))​𝑑s.D(t)\leq D(0)+\int_{0}^{t}\Bigl((L_{1}+L_{2})D(s)+L_{3}\mathcal{W}_{2}(p,\widetilde{p})\Bigr)\,ds.SinceD​(0)=𝒲2​(μ0,ν0)D(0)=\mathcal{W}_{2}(\mu_{0},\nu_{0}), Gronwall’s inequality yields

𝒲2​(μt,νt)≤e(L1+L2)​t​𝒲2​(μ0,ν0)+L3L1+L2​(e(L1+L2)​t−1)​𝒲2​(p,p~).\mathcal{W}_{2}(\mu_{t},\nu_{t})\leq e^{(L_{1}+L_{2})t}\mathcal{W}_{2}(\mu_{0},\nu_{0})+\frac{L_{3}}{L_{1}+L_{2}}\bigl(e^{(L_{1}+L_{2})t}-1\bigr)\mathcal{W}_{2}(p,\widetilde{p}).(35)In particular, takingp~=p\widetilde{p}=pgives uniqueness.

We now prove existence. Let(q^0N)N(\widehat{q}_{0}^{N})_{N}and(p^M)M(\widehat{p}^{M})_{M}be the empirical measures so that

q^0N→q0,p^M→pin​𝒲2\widehat{q}_{0}^{N}\to q_{0},\qquad\widehat{p}^{M}\to p\qquad\text{in }\mathcal{W}_{2}almost surely. Applying (35) toqt^N,M\widehat{q_{t}}^{N,M}andqt^N′,M′\widehat{q_{t}}^{N^{\prime},M^{\prime}}yields

supt∈[0,T]𝒲2​(q^tN,M,q^tN′,M′)\displaystyle\sup_{t\in[0,T]}\mathcal{W}_{2}(\widehat{q}_{t}^{N,M},\widehat{q}_{t}^{N^{\prime},M^{\prime}})≤e(L1+L2)​T​𝒲2​(q^0N,q^0N′)+L3L1+L2​(e(L1+L2)​T−1)​𝒲2​(p^M,p^M′).\displaystyle\leq e^{(L_{1}+L_{2})T}\mathcal{W}_{2}(\widehat{q}_{0}^{N},\widehat{q}_{0}^{N^{\prime}})+\frac{L_{3}}{L_{1}+L_{2}}\bigl(e^{(L_{1}+L_{2})T}-1\bigr)\mathcal{W}_{2}(\widehat{p}^{M},\widehat{p}^{M^{\prime}}).Sinceq^0N→q0\widehat{q}_{0}^{N}\to q_{0}andp^M→p\widehat{p}^{M}\to pin𝒲2\mathcal{W}_{2}almost surely, the right-hand side tends to0asN,N′,M,M′→∞N,N^{\prime},M,M^{\prime}\to\infty. Hence(q^N,M)N,M(\widehat{q}^{N,M})_{N,M}is a Cauchy family inC​([0,T];𝒫2​(ℝd))C([0,T];\mathcal{P}_{2}(\mathbb{R}^{d})), where𝒫2​(ℝd)\mathcal{P}_{2}(\mathbb{R}^{d})is equipped with𝒲2\mathcal{W}_{2}. Since(𝒫2​(ℝd),𝒲2)(\mathcal{P}_{2}(\mathbb{R}^{d}),\mathcal{W}_{2})is complete, there exists

q¯∈C​([0,T];𝒫2​(ℝd))\overline{q}\in C([0,T];\mathcal{P}_{2}(\mathbb{R}^{d}))such that

supt∈[0,T]𝒲2​(qt^N,M,q¯t)→0almost surelyas​N,M→∞.\sup_{t\in[0,T]}\mathcal{W}_{2}(\widehat{q_{t}}^{N,M},\overline{q}_{t})\to 0\qquad\text{almost surely}\qquad\text{as }N,M\to\infty.(36)In particular,q¯0=q0\overline{q}_{0}=q_{0}.

We also need a uniform second-moment bound. Set

m2N,M​(t):=(∫‖x‖2​𝑑q^tN,M​(x))1/2.m_{2}^{N,M}(t):=\left(\int\|x\|^{2}\,d\widehat{q}_{t}^{N,M}(x)\right)^{1/2}.For a.e.tt, using the ODE and AssumptionB.1,

dd​t​m2N,M​(t)≤‖Vq^tN,M,p^M‖L2​(q^tN,M)≤L0​(m2N,M​(t)+m2​(p^M)).\frac{d}{dt}m_{2}^{N,M}(t)\leq\|V_{\widehat{q}_{t}^{N,M},\widehat{p}^{M}}\|_{L^{2}(\widehat{q}_{t}^{N,M})}\leq L_{0}\bigl(m_{2}^{N,M}(t)+m_{2}(\widehat{p}^{M})\bigr).Therefore, by Gronwall’s inequality,

m2N,M​(t)≤eL0​t​m2​(q^0N)+(eL0​t−1)​m2​(p^M),0≤t≤T.m_{2}^{N,M}(t)\leq e^{L_{0}t}m_{2}(\widehat{q}_{0}^{N})+\bigl(e^{L_{0}t}-1\bigr)m_{2}(\widehat{p}^{M}),\qquad 0\leq t\leq T.(37)Sinceq^0N→q0\widehat{q}_{0}^{N}\to q_{0}andp^M→p\widehat{p}^{M}\to pin𝒲2\mathcal{W}_{2}almost surely, the sequencesm2​(q^0N)m_{2}(\widehat{q}_{0}^{N})andm2​(p^M)m_{2}(\widehat{p}^{M})are bounded. Hence there existsCT>0C_{T}>0such that

supN,Msupt∈[0,T]∫‖x‖2​𝑑q^tN,M​(x)≤CT.\sup_{N,M}\sup_{t\in[0,T]}\int\|x\|^{2}\,d\widehat{q}_{t}^{N,M}(x)\leq C_{T}.From (37) and the convergence in (36), passing to the limit also gives

supt∈[0,T]∫‖x‖2​𝑑q¯t​(x)≤CT.\sup_{t\in[0,T]}\int\|x\|^{2}\,d\overline{q}_{t}(x)\leq C_{T}. It remains to pass to the limit in the weak formulation. Fixφ∈Cc1​(ℝd)\varphi\in C_{c}^{1}(\mathbb{R}^{d}). For eachs∈[0,T]s\in[0,T], letγsN,M\gamma_{s}^{N,M}be an optimal coupling ofq^sN,M\widehat{q}_{s}^{N,M}andq¯s\overline{q}_{s}. Write

IN,M​(s)\displaystyle I_{N,M}(s):=∫∇φ​(x)⋅Vq^sN,M,p^M​(x)​𝑑q^sN,M​(x),\displaystyle:=\int\nabla\varphi(x)\cdot V_{\widehat{q}_{s}^{N,M},\widehat{p}^{M}}(x)\,d\widehat{q}_{s}^{N,M}(x),I​(s)\displaystyle I(s):=∫∇φ​(y)⋅Vq¯s,p​(y)​𝑑q¯s​(y).\displaystyle:=\int\nabla\varphi(y)\cdot V_{\overline{q}_{s},p}(y)\,d\overline{q}_{s}(y).Then

|IN,M​(s)−I​(s)|\displaystyle|I_{N,M}(s)-I(s)|≤|∫(∇φ​(x)−∇φ​(y))⋅Vq^sN,M,p^M​(x)​𝑑γsN,M​(x,y)|\displaystyle\leq\left|\int\bigl(\nabla\varphi(x)-\nabla\varphi(y)\bigr)\cdot V_{\widehat{q}_{s}^{N,M},\widehat{p}^{M}}(x)\,d\gamma_{s}^{N,M}(x,y)\right|+‖∇φ‖∞​(∫‖Vq^sN,M,p^M​(x)−Vq¯s,p​(y)‖2​𝑑γsN,M​(x,y))1/2.\displaystyle\quad+\|\nabla\varphi\|_{\infty}\left(\int\bigl\|V_{\widehat{q}_{s}^{N,M},\widehat{p}^{M}}(x)-V_{\overline{q}_{s},p}(y)\bigr\|^{2}\,d\gamma_{s}^{N,M}(x,y)\right)^{1/2}.The second term is bounded by

‖∇φ‖∞​[(L1+L2)​𝒲2​(q^sN,M,q¯s)+L3​𝒲2​(p^M,p)],\|\nabla\varphi\|_{\infty}\Bigl[(L_{1}+L_{2})\mathcal{W}_{2}(\widehat{q}_{s}^{N,M},\overline{q}_{s})+L_{3}\mathcal{W}_{2}(\widehat{p}^{M},p)\Bigr],and converges to zero uniformly inss. For the first term, letωφ\omega_{\varphi}be a modulus of continuity of∇φ\nabla\varphi. For everyr>0r>0, Cauchy–Schwarz and the optimality ofγsN,M\gamma_{s}^{N,M}give

|∫(∇φ​(x)−∇φ​(y))⋅Vq^sN,M,p^M​(x)​𝑑γsN,M​(x,y)|\displaystyle\left|\int\bigl(\nabla\varphi(x)-\nabla\varphi(y)\bigr)\cdot V_{\widehat{q}_{s}^{N,M},\widehat{p}^{M}}(x)\,d\gamma_{s}^{N,M}(x,y)\right|≤ωφ​(r)​‖Vq^sN,M,p^M‖L1​(q^sN,M)+2​‖∇φ‖∞​‖Vq^sN,M,p^M‖L2​(q^sN,M)​𝒲2​(q^sN,M,q¯s)r.\displaystyle\qquad\leq\omega_{\varphi}(r)\,\|V_{\widehat{q}_{s}^{N,M},\widehat{p}^{M}}\|_{L^{1}(\widehat{q}_{s}^{N,M})}+2\|\nabla\varphi\|_{\infty}\|V_{\widehat{q}_{s}^{N,M},\widehat{p}^{M}}\|_{L^{2}(\widehat{q}_{s}^{N,M})}\frac{\mathcal{W}_{2}(\widehat{q}_{s}^{N,M},\overline{q}_{s})}{r}.The velocity norms are uniformly bounded on[0,T][0,T]by (27) and (37). TakingN,M→∞N,M\to\inftyand thenr↓0r\downarrow 0proves thatIN,M→II_{N,M}\to Iuniformly on[0,T][0,T]. Passing to the limit in the weak formulation gives

∫φ​𝑑q¯t−∫φ​𝑑q0=∫0t∫∇φ​(x)⋅Vq¯s,p​(x)​𝑑q¯s​(x)​𝑑s.\int\varphi\,d\overline{q}_{t}-\int\varphi\,dq_{0}=\int_{0}^{t}\int\nabla\varphi(x)\cdot V_{\overline{q}_{s},p}(x)\,d\overline{q}_{s}(x)\,ds.Thusq¯\overline{q}is a weak solution of (31). By uniqueness,q¯\overline{q}is the unique weak solution; we denote it byqq.

Applying the stability estimate in (35) with

μt=qt,νt=q^tN,M,p~=p^M\mu_{t}=q_{t},\qquad\nu_{t}=\widehat{q}_{t}^{N,M},\qquad\widetilde{p}=\widehat{p}^{M}gives, for allt∈[0,T]t\in[0,T],

𝒲2​(q^tN,M,qt)≤e(L1+L2)​t​𝒲2​(q^0N,q0)+L3L1+L2​(e(L1+L2)​t−1)​𝒲2​(p^M,p).\mathcal{W}_{2}(\widehat{q}_{t}^{N,M},q_{t})\leq e^{(L_{1}+L_{2})t}\mathcal{W}_{2}(\widehat{q}_{0}^{N},q_{0})+\frac{L_{3}}{L_{1}+L_{2}}(e^{(L_{1}+L_{2})t}-1)\mathcal{W}_{2}(\widehat{p}^{M},p).(38)Taking the supremum overt∈[0,T]t\in[0,T]gives the corresponding bound for the continuous-time particle system.

On the other hand, let𝒙N,M,η​(t)\boldsymbol{x}^{N,M,\eta}(t)be the piecewise linear Euler interpolation of (33), namely

𝒙N,M,η​(tk+1)=𝒙N,M,η​(tk)+η​VN,M​(𝒙N,M,η​(tk)),\boldsymbol{x}^{N,M,\eta}(t_{k+1})=\boldsymbol{x}^{N,M,\eta}(t_{k})+\eta\,V^{N,M}(\boldsymbol{x}^{N,M,\eta}(t_{k})),and

𝒙N,M,η​(t)=𝒙N,M,η​(tk)+(t−tk)​VN,M​(𝒙N,M,η​(tk)),t∈[tk,tk+1).\boldsymbol{x}^{N,M,\eta}(t)=\boldsymbol{x}^{N,M,\eta}(t_{k})+(t-t_{k})V^{N,M}(\boldsymbol{x}^{N,M,\eta}(t_{k})),\qquad t\in[t_{k},t_{k+1}).Let

q^tN,M,η:=1N​∑i=1NδxiN,M,η​(t).\widehat{q}_{t}^{N,M,\eta}:=\frac{1}{N}\sum_{i=1}^{N}\delta_{x_{i}^{N,M,\eta}(t)}. Set

e​(t):=‖𝒙N,M,η​(t)−𝒙N,M​(t)‖2,N.e(t):=\|\boldsymbol{x}^{N,M,\eta}(t)-\boldsymbol{x}^{N,M}(t)\|_{2,N}. SinceVN,MV^{N,M}is globally Lipschitz with constantL1+L2L_{1}+L_{2}, the standard Euler estimate gives

supt∈[0,T]e​(t)≤12​(e(L1+L2)​T−1)​η​sups∈[0,T]‖VN,M​(𝒙N,M​(s))‖2,N.\sup_{t\in[0,T]}e(t)\leq\frac{1}{2}\bigl(e^{(L_{1}+L_{2})T}-1\bigr)\eta\sup_{s\in[0,T]}\|V^{N,M}(\boldsymbol{x}^{N,M}(s))\|_{2,N}.Moreover, by (27) and (37), for everys∈[0,T]s\in[0,T],

‖VN,M​(𝒙N,M​(s))‖2,N\displaystyle\|V^{N,M}(\boldsymbol{x}^{N,M}(s))\|_{2,N}=‖Vq^sN,M,p^Mε‖L2​(q^sN,M)\displaystyle=\|V_{\widehat{q}_{s}^{N,M},\widehat{p}^{M}}^{\varepsilon}\|_{L^{2}(\widehat{q}_{s}^{N,M})}≤L0​(m2​(q^sN,M)+m2​(p^M))\displaystyle\leq L_{0}\bigl(m_{2}(\widehat{q}_{s}^{N,M})+m_{2}(\widehat{p}^{M})\bigr)≤L0​eL0​T​[m2​(q^0N)+m2​(p^M)].\displaystyle\leq L_{0}e^{L_{0}T}\bigl[m_{2}(\widehat{q}_{0}^{N})+m_{2}(\widehat{p}^{M})\bigr].Therefore,

supt∈[0,T]e​(t)≤L0​eL0​T2​(e(L1+L2)​T−1)​η​[m2​(q^0N)+m2​(p^M)].\sup_{t\in[0,T]}e(t)\leq\frac{L_{0}e^{L_{0}T}}{2}\bigl(e^{(L_{1}+L_{2})T}-1\bigr)\eta\bigl[m_{2}(\widehat{q}_{0}^{N})+m_{2}(\widehat{p}^{M})\bigr].(39) Since the coupling that matchesxiN,M,η​(t)x_{i}^{N,M,\eta}(t)withxiN,M​(t)x_{i}^{N,M}(t)is admissible, we have

𝒲2​(q^tN,M,η,q^tN,M)≤‖𝒙N,M,η​(t)−𝒙N,M​(t)‖2,N=e​(t).\mathcal{W}_{2}(\widehat{q}_{t}^{N,M,\eta},\widehat{q}_{t}^{N,M})\leq\|\boldsymbol{x}^{N,M,\eta}(t)-\boldsymbol{x}^{N,M}(t)\|_{2,N}=e(t).Combining this estimate with (38) and (39), and taking the supremum overt∈[0,T]t\in[0,T], proves the quantitative bound in the theorem.

Finally, sinceq0,p∈𝒫2​(ℝd)q_{0},p\in\mathcal{P}_{2}(\mathbb{R}^{d}), the empirical measures satisfy

𝒲2​(q^0N,q0)→0,𝒲2​(p^M,p)→0,almost surely.\mathcal{W}_{2}(\widehat{q}_{0}^{N},q_{0})\to 0,\qquad\mathcal{W}_{2}(\widehat{p}^{M},p)\to 0,\qquad\text{almost surely}.Moreover,

∫ℝd‖x‖2​𝑑q^0N​(x)→∫ℝd‖x‖2​𝑑q0​(x),∫ℝd‖x‖2​𝑑p^M​(x)→∫ℝd‖x‖2​𝑑p​(x),almost surely.\displaystyle\int_{\mathbb{R}^{d}}\|x\|^{2}\,d\widehat{q}_{0}^{N}(x)\to\int_{\mathbb{R}^{d}}\|x\|^{2}\,dq_{0}(x),\quad\int_{\mathbb{R}^{d}}\|x\|^{2}\,d\widehat{p}^{M}(x)\to\int_{\mathbb{R}^{d}}\|x\|^{2}\,dp(x),\qquad\text{almost surely}.LettingN,M→∞N,M\to\inftyandη↓0\eta\downarrow 0, we obtain

supt∈[0,T]𝒲2​(q^tN,M,η,qt)→0,almost surely.\sup_{t\in[0,T]}\mathcal{W}_{2}(\widehat{q}_{t}^{N,M,\eta},q_{t})\to 0,\qquad\text{almost surely}.∎

B.2Assumption verification for Sinkhorn

Recall that the Sinkhorn velocity in (10):

Vq,pε​(x)=Tq,pε​(x)−Tq,qε​(x),V_{q,p}^{\varepsilon}(x)=T_{q,p}^{\varepsilon}(x)-T_{q,q}^{\varepsilon}(x),whereTq,pεT_{q,p}^{\varepsilon}denotes the barycentric projection of the entropic optimal coupling betweenqqandppfor the quadratic cost.

Before verifying AssumptionB.1, we first localize the dynamics. Ifsupp⁡(q0)⊂BR0\operatorname{supp}(q_{0})\subset B_{R_{0}}andsupp⁡(p)⊂BR\operatorname{supp}(p)\subset B_{R}, then the WGF under Sinkhorn divergence remains supported in a bounded ball on every finite time interval; see PropositionB.5. On this bounded domain, we then use compact-domain stability and regularity estimates for Sinkhorn barycentric projections, collected in LemmaB.6. Combining these two ingredients allows us to verify AssumptionB.1for the Sinkhorn velocity along the localized flow and its finite-particle approximations; see PropositionB.7.

Proposition B.5(Bounded support propagation for the Sinkhorn flow).

Assume that

supp⁡(q0)⊂BR0,supp⁡(p)⊂BR,\operatorname{supp}(q_{0})\subset B_{R_{0}},\qquad\operatorname{supp}(p)\subset B_{R},whereBR:={x∈ℝd:‖x‖≤R}B_{R}:=\{x\in\mathbb{R}^{d}:\|x\|\leq R\}. Letqtq_{t}solve

∂tqt+∇⋅(qt​Vqt,pε)=0.\partial_{t}q_{t}+\nabla\cdot(q_{t}V_{q_{t},p}^{\varepsilon})=0.Then, on every finite time interval[0,T][0,T],

supp⁡(qt)⊂Bet​R0+(et−1)​R.\operatorname{supp}(q_{t})\subset B_{e^{t}R_{0}+(e^{t}-1)R}.

The same support bound holds for the discrete finite-particle dynamics in (25) and for the piecewise-linear interpolation in (26), provided the initial particles are supported inBR0B_{R_{0}}and the target particles are supported inBRB_{R}.

Proof.

Let us define the radius

ρ(t):=sup{∥x∥:x∈supp(qt)}.\rho(t):=\sup\{\|x\|:x\in\operatorname{supp}(q_{t})\}.SinceTqt,pε​(x)T_{q_{t},p}^{\varepsilon}(x)is the barycenter of a probability measure supported insupp⁡(p)⊂BR\operatorname{supp}(p)\subset B_{R}, andTqt,qtε​(x)T_{q_{t},q_{t}}^{\varepsilon}(x)is the barycenter of a probability measure supported insupp⁡(qt)\operatorname{supp}(q_{t}), we have

‖Tqt,pε​(x)‖≤R,‖Tqt,qtε​(x)‖≤ρ​(t).\|T_{q_{t},p}^{\varepsilon}(x)\|\leq R,\qquad\|T_{q_{t},q_{t}}^{\varepsilon}(x)\|\leq\rho(t).Therefore,

‖Vqt,pε​(x)‖≤‖Tqt,pε​(x)‖+‖Tqt,qtε​(x)‖≤R+ρ​(t),x∈supp⁡(qt).\|V_{q_{t},p}^{\varepsilon}(x)\|\leq\|T_{q_{t},p}^{\varepsilon}(x)\|+\|T_{q_{t},q_{t}}^{\varepsilon}(x)\|\leq R+\rho(t),\qquad x\in\operatorname{supp}(q_{t}).Equivalently, for any particle trajectoryx​(t)x(t)in the support of the flow,

dd​t​‖x​(t)‖≤‖Vqt,pε​(x​(t))‖≤R+ρ​(t).\frac{d}{dt}\|x(t)\|\leq\|V_{q_{t},p}^{\varepsilon}(x(t))\|\leq R+\rho(t).Taking the supremum over all such trajectories gives

ρ​(t)≤R0+∫0t(R+ρ​(s))​𝑑s.\rho(t)\leq R_{0}+\int_{0}^{t}(R+\rho(s))\,ds.By Gronwall’s inequality,

ρ​(t)≤et​R0+(et−1)​R.\rho(t)\leq e^{t}R_{0}+(e^{t}-1)R.This proves

supp⁡(qt)⊂Bet​R0+(et−1)​R.\operatorname{supp}(q_{t})\subset B_{e^{t}R_{0}+(e^{t}-1)R}. For the continuous-time finite-particle dynamics, the same argument applies withqtq_{t}replaced by the empirical measure. For the Euler scheme, ifq^kN\widehat{q}_{k}^{N}is supported inBρkB_{\rho_{k}}, then

‖xi(k+1)‖≤‖xi(k)‖+η​‖Vq^kN,p^Mε​(xi(k))‖≤ρk+η​(R+ρk).\left\|x_{i}^{(k+1)}\right\|\leq\left\|x_{i}^{(k)}\right\|+\eta\left\|V_{\widehat{q}_{k}^{N},\widehat{p}^{M}}^{\varepsilon}(x_{i}^{(k)})\right\|\leq\rho_{k}+\eta(R+\rho_{k}).Hence

ρk+1≤(1+η)​ρk+η​R.\rho_{k+1}\leq(1+\eta)\rho_{k}+\eta R.Iterating this recursion gives

ρk≤(1+η)k​R0+((1+η)k−1)​R≤etk​R0+(etk−1)​R.\rho_{k}\leq(1+\eta)^{k}R_{0}+\bigl((1+\eta)^{k}-1\bigr)R\leq e^{t_{k}}R_{0}+(e^{t_{k}}-1)R.The same estimate controls the piecewise-linear interpolation betweentkt_{k}andtk+1t_{k+1}. ∎

Lemma B.6(Compact-domain stability of Sinkhorn barycentric projections).

LetK⊂ℝdK\subset\mathbb{R}^{d}be compact, letDK:=diam⁡(K)D_{K}:=\operatorname{diam}(K), and fixε>0\varepsilon>0. There existsCK,ε<∞C_{K,\varepsilon}<\inftysuch that, for allq,q~,p,p~∈𝒫​(K)q,\widetilde{q},p,\widetilde{p}\in\mathcal{P}(K),

supx∈K‖Tq,pε​(x)−Tq~,p~ε​(x)‖≤CK,ε​(𝒲2​(q,q~)+𝒲2​(p,p~)),\sup_{x\in K}\bigl\|T_{q,p}^{\varepsilon}(x)-T_{\widetilde{q},\widetilde{p}}^{\varepsilon}(x)\bigr\|\leq C_{K,\varepsilon}\bigl(\mathcal{W}_{2}(q,\widetilde{q})+\mathcal{W}_{2}(p,\widetilde{p})\bigr),(40)and, for allx,x′∈Kx,x^{\prime}\in K,

‖Tq,pε​(x)−Tq,pε​(x′)‖≤DK24​ε​‖x−x′‖.\bigl\|T_{q,p}^{\varepsilon}(x)-T_{q,p}^{\varepsilon}(x^{\prime})\bigr\|\leq\frac{D_{K}^{2}}{4\varepsilon}\|x-x^{\prime}\|.(41)

Proof.

We first prove (40). By the compact-domain stability result for entropic OT potentials[8, Corollary 2.4], applied to the rescaled quadratic cost(x,y)↦‖x−y‖2/(2​ε)(x,y)\mapsto\|x-y\|^{2}/(2\varepsilon)onK×KK\times K, there existsCK,ε<∞C_{K,\varepsilon}<\inftysuch that, for allq,q~,p,p~∈𝒫​(K)q,\widetilde{q},p,\widetilde{p}\in\mathcal{P}(K),

supx∈K‖∇uq,pε​(x)−∇uq~,p~ε​(x)‖≤CK,ε​(𝒲2​(q,q~)+𝒲2​(p,p~)),\sup_{x\in K}\bigl\|\nabla u_{q,p}^{\varepsilon}(x)-\nabla u_{\widetilde{q},\widetilde{p}}^{\varepsilon}(x)\bigr\|\leq C_{K,\varepsilon}\bigl(\mathcal{W}_{2}(q,\widetilde{q})+\mathcal{W}_{2}(p,\widetilde{p})\bigr),whereuq,pεu^{\varepsilon}_{q,p}denotes the Schrödinger potential introduced in (18). SinceKKis compact and the quadratic cost is smooth onK×KK\times K, the assumptions required for[8, Corollary 2.4]to hold are satisfied. Using the identity in (19), we immediately obtain (40).

It remains to relate the Schrödinger potential to the barycentric projection. Recall from (19) that∇uq,pε​(x)=x−Tq,pε​(x)\nabla u_{q,p}^{\varepsilon}(x)=x-T_{q,p}^{\varepsilon}(x). Therefore,

supx∈K‖Tq,pε​(x)−Tq~,p~ε​(x)‖\displaystyle\sup_{x\in K}\bigl\|T_{q,p}^{\varepsilon}(x)-T_{\widetilde{q},\widetilde{p}}^{\varepsilon}(x)\bigr\|=supx∈K‖∇uq,pε​(x)−∇uq~,p~ε​(x)‖\displaystyle=\sup_{x\in K}\bigl\|\nabla u_{q,p}^{\varepsilon}(x)-\nabla u_{\widetilde{q},\widetilde{p}}^{\varepsilon}(x)\bigr\|≤CK,ε​(𝒲2​(q,q~)+𝒲2​(p,p~)).\displaystyle\leq C_{K,\varepsilon}\bigl(\mathcal{W}_{2}(q,\widetilde{q})+\mathcal{W}_{2}(p,\widetilde{p})\bigr).This proves (40).

It remains to prove the spatial estimate. Since

Tq,pε​(x)=∫Ky​πq,pε​(d​y∣x),T_{q,p}^{\varepsilon}(x)=\int_{K}y\,\pi_{q,p}^{\varepsilon}(dy\mid x),differentiating the conditional mean with respect toxxgives

∇xTq,pε​(x)=1ε​Covπq,pε​(d​y∣x)⁡(Y).\nabla_{x}T_{q,p}^{\varepsilon}(x)=\frac{1}{\varepsilon}\operatorname{Cov}_{\pi_{q,p}^{\varepsilon}(dy\mid x)}(Y).For any probability measure supported onKK, the operator norm of its covariance matrix is at mostDK2/4D_{K}^{2}/4. Therefore

‖∇xTq,pε​(x)‖op≤DK24​ε,x∈K.\|\nabla_{x}T_{q,p}^{\varepsilon}(x)\|_{\mathrm{op}}\leq\frac{D_{K}^{2}}{4\varepsilon},\qquad x\in K.This proves (41). ∎

Proposition B.7(Verification of AssumptionB.1for bounded-support Sinkhorn flow).

Assume

supp⁡(q0)⊂BR0,supp⁡(p)⊂BR,\operatorname{supp}(q_{0})\subset B_{R_{0}},\qquad\operatorname{supp}(p)\subset B_{R},and fixT>0T>0. Then all measures appearing in the Sinkhorn flow and in the finite-particle approximations on[0,T][0,T]are supported inBR¯TB_{\bar{R}_{T}}with

R¯T:=max⁡{R,eT​R0+(eT−1)​R}.\bar{R}_{T}:=\max\{R,\ e^{T}R_{0}+(e^{T}-1)R\}.Moreover, for the Sinkhorn velocityVq,pε​(x)V_{q,p}^{\varepsilon}(x)defined in(10), AssumptionB.1holds on𝒫​(BR¯T)\mathcal{P}(B_{\bar{R}_{T}}), with constants depending only on(T,R0,R,ε)(T,R_{0},R,\varepsilon). More explicitly, there existL0,L1,L2,L3<∞L_{0},L_{1},L_{2},L_{3}<\infty, depending only on(R¯T,ε)(\bar{R}_{T},\varepsilon), such that(27)–(30)hold for allq,q~,p,p~∈𝒫​(BR¯T)q,\widetilde{q},p,\widetilde{p}\in\mathcal{P}(B_{\bar{R}_{T}}), with the spatial estimate required only forx,x′∈BR¯Tx,x^{\prime}\in B_{\bar{R}_{T}}.

Proof.

The support statement follows from PropositionB.5. It remains to verify AssumptionB.1onBR¯TB_{\bar{R}_{T}}.

TheL2L^{2}-growth estimate follows from Jensen’s inequality. Specifically, if(X,Y)∼πq,pε(X,Y)\sim\pi_{q,p}^{\varepsilon}, then

∫∥Tq,pε(x)∥2dq(x)=∫∥𝔼[Y∣X=x]∥2dq(x)≤∫∥y∥2dp(y).\int\|T_{q,p}^{\varepsilon}(x)\|^{2}\,dq(x)=\int\|\mathbb{E}[Y\mid X=x]\|^{2}\,dq(x)\leq\int\|y\|^{2}\,dp(y).Similarly,

∫‖Tq,qε​(x)‖2​𝑑q​(x)≤∫‖y‖2​𝑑q​(y).\int\|T_{q,q}^{\varepsilon}(x)\|^{2}\,dq(x)\leq\int\|y\|^{2}\,dq(y).Hence

‖Vq,pε‖L2​(q)≤‖Tq,pε‖L2​(q)+‖Tq,qε‖L2​(q)≤m2​(p)+m2​(q),\|V_{q,p}^{\varepsilon}\|_{L^{2}(q)}\leq\|T_{q,p}^{\varepsilon}\|_{L^{2}(q)}+\|T_{q,q}^{\varepsilon}\|_{L^{2}(q)}\leq m_{2}(p)+m_{2}(q),so (27) holds withL0=1L_{0}=1.

By LemmaB.6, forx,x′∈BR¯Tx,x^{\prime}\in B_{\bar{R}_{T}},

‖Vq,pε​(x)−Vq,pε​(x′)‖\displaystyle\|V_{q,p}^{\varepsilon}(x)-V_{q,p}^{\varepsilon}(x^{\prime})\|≤‖Tq,pε​(x)−Tq,pε​(x′)‖+‖Tq,qε​(x)−Tq,qε​(x′)‖\displaystyle\leq\|T_{q,p}^{\varepsilon}(x)-T_{q,p}^{\varepsilon}(x^{\prime})\|+\|T_{q,q}^{\varepsilon}(x)-T_{q,q}^{\varepsilon}(x^{\prime})\|≤diam(BR¯T)22​ε​‖x−x′‖.\displaystyle\leq\frac{\operatorname{diam}(B_{\bar{R}_{T}})^{2}}{2\varepsilon}\|x-x^{\prime}\|.Thus (28) holds onBR¯TB_{\bar{R}_{T}}.

For theqq-stability estimate, LemmaB.6gives

supx∈BR¯T‖Vq,pε​(x)−Vq~,pε​(x)‖\displaystyle\sup_{x\in B_{\bar{R}_{T}}}\|V_{q,p}^{\varepsilon}(x)-V_{\widetilde{q},p}^{\varepsilon}(x)\|≤supx∈BR¯T‖Tq,pε​(x)−Tq~,pε​(x)‖\displaystyle\leq\sup_{x\in B_{\bar{R}_{T}}}\|T_{q,p}^{\varepsilon}(x)-T_{\widetilde{q},p}^{\varepsilon}(x)\|+supx∈BR¯T‖Tq,qε​(x)−Tq~,q~ε​(x)‖\displaystyle\quad+\sup_{x\in B_{\bar{R}_{T}}}\|T_{q,q}^{\varepsilon}(x)-T_{\widetilde{q},\widetilde{q}}^{\varepsilon}(x)\|≤3​CBR¯T,ε​𝒲2​(q,q~).\displaystyle\leq 3C_{B_{\bar{R}_{T}},\varepsilon}\mathcal{W}_{2}(q,\widetilde{q}).Integrating with respect toq~\widetilde{q}gives (29).

For thepp-stability estimate, the self-interaction term cancels:

Vq,pε​(x)−Vq,p~ε​(x)=Tq,pε​(x)−Tq,p~ε​(x).V_{q,p}^{\varepsilon}(x)-V_{q,\widetilde{p}}^{\varepsilon}(x)=T_{q,p}^{\varepsilon}(x)-T_{q,\widetilde{p}}^{\varepsilon}(x).Therefore, by LemmaB.6,

supx∈BR¯T‖Vq,pε​(x)−Vq,p~ε​(x)‖≤CBR¯T,ε​𝒲2​(p,p~).\sup_{x\in B_{\bar{R}_{T}}}\|V_{q,p}^{\varepsilon}(x)-V_{q,\widetilde{p}}^{\varepsilon}(x)\|\leq C_{B_{\bar{R}_{T}},\varepsilon}\mathcal{W}_{2}(p,\widetilde{p}).Integrating with respect toqqgives (30). ∎

Appendix CAdditional implementation details

ImageNet experiments.We provide detailed hyperparameters and configurations in Table5. Note that we largely follow the configurations and setups adopted in[14], without extensive hyperparameter tuning. We conduct experiments on 8 Nvidia H100 GPU nodes, where each GPU has 80 GB of memory. The training on L/2 takes around 3 days.

Table 5:Configurations and hyperparameters for ImageNet experiments.FFHQ experiments.For the FFHQ experiments, we operate in the 512-dimensional latent space of a pretrained Adversarial Latent Autoencoder (ALAE), utilizing a training split of 60,000 images. The generator is parameterized as a 4-layer MLP with hidden dimensions of 1024, utilizing LayerNorm and SiLU activations in the hidden blocks. Models are trained using the AdamW optimizer with a constant learning rate of2×10−42\times 10^{-4},β1=0.9\beta_{1}=0.9,β2=0.95\beta_{2}=0.95, and zero weight decay. We apply gradient clipping at 5.0 and maintain an exponential moving average (EMA) of the weights with a decay rate of 0.999. During training, we use a batch size ofN=64N=64for the generated particles andM=64M=64for the target particles, with a fixed step size ofη=1\eta=1. For the Sinkhorn divergence, we use the quadratic cost with an entropy regularization parameterε=0.02\varepsilon=0.02, with 10 Sinkhorn iterations. Consistent with Sec.3.3, we explicitly use the two-batch debiased self-transport estimator, avoiding any heuristic diagonal masking.

We evaluate this setup on two distinct tasks: domain transfer and mode coverage, as shown in Sec.4.3. For the domain transfer task, we map a source distribution of senior faces (ages 55–100) to a target distribution of young adult faces (ages 18–30). The initial particles are drawn directly from the source data distribution, and we apply the residual identity initialization by zero-initializing the final layer of the MLP to learn the residual update. For the mode coverage experiment, we perform unconditional generation starting from a standard Gaussian prior. The target batches are constructed to be highly imbalanced, comprising 95% senior faces and 5% child faces (ages 0–12). In this setting, the standard network initialization is used without the residual connection.

Settings for plotting Fig.2(b).We train W-Flow with the Sinkhorn divergence on toy datasets for 20,000 iterations, saving a checkpoint every 400 iterations. This yields a sequence of parameter vectors{θk}k=0K\{\theta_{k}\}_{k=0}^{K}, whereK=50K=50. We visualize the Sinkhorn-divergence energy landscape over a two-dimensional affine slice of the parameter space. Each point(x,y)(x,y)on thexx-yyplane is mapped to a parameter vectorθx,y\theta_{x,y}, and the corresponding height is given bySε​(qθx,y,p),S_{\varepsilon}(q_{\theta_{x,y}},p),whereqθx,yq_{\theta_{x,y}}denotes the model distribution induced byθx,y\theta_{x,y}.

We construct thexx-yygrid as follows. First, we place the training trajectory along thexx-axis by assigning checkpointθk\theta_{k}to coordinate(k,0)(k,0)fork=0,…,Kk=0,\ldots,K. Let

dk=θk+1−θk,k=0,…,K−1,d_{k}=\theta_{k+1}-\theta_{k},\qquad k=0,\ldots,K-1,denote the displacement directions along the checkpoint trajectory. We then sample a random directionvvin parameter space and project it onto the orthogonal complement ofspan​{d0,…,dK−1}\mathrm{span}\{d_{0},\ldots,d_{K-1}\}, followed by normalization. This directionvvis used as a common transverse direction for theyy-axis.

For each intervalx∈[k,k+1]x\in[k,k+1], we interpolate along the local trajectory directiondkd_{k}and extend transversely alongvv:

θx,y=θk+(x−k)​dk+y​v,k≤x≤k+1,−Y≤y≤Y.\theta_{x,y}=\theta_{k}+(x-k)d_{k}+yv,\qquad k\leq x\leq k+1,\quad-Y\leq y\leq Y.This defines the landscape grid over0≤x≤K0\leq x\leq Kand−Y≤y≤Y-Y\leq y\leq Y, while ensuring that the saved checkpoints lie exactly on thexx-axis.

To visualize the landscape beyond the training trajectory, we further extend thexx-axis to−X≤x≤K+X-X\leq x\leq K+X. For the left extension−X≤x≤0-X\leq x\leq 0, we use the initial trajectory directiond0=θ1−θ0d_{0}=\theta_{1}-\theta_{0}:

θx,y=θ0+x​d0+y​v.\theta_{x,y}=\theta_{0}+xd_{0}+yv.For the right extensionK≤x≤K+XK\leq x\leq K+X, we use the final trajectory directiondK−1=θK−θK−1d_{K-1}=\theta_{K}-\theta_{K-1}:

θx,y=θK+(x−K)​dK−1+y​v.\theta_{x,y}=\theta_{K}+(x-K)d_{K-1}+yv.Together, this construction yields a continuous two-dimensional slice of the parameter space on which we evaluate and plot the Sinkhorn-divergence energy.

Appendix DSampling throughput comparison

We compare our one-step W-Flow against multi-step diffusion model SiT[44]and LightningDiT[68]in terms of actual sampling throughput measured by wall-clock time. We benchmark on a single Nvidia H100 GPU with a batch size of 4 for all models. For SiT-XL/2 and LightningDiT-XL/2, we adopt 250 sampling steps with CFG enabled and strictly follow the configuration in the original papers[44,68]that leads to a reported FID of 2.06 and 1.35, respectively. Note that we also include the latent preparation and VAE decoding time when recording the actual runtime and calculating the throughput. The results are depicted in Table6, where W-Flow clearly yields significant speedup while achieving competitive generation quality measured by FID on ImageNet 256×\times256.

Table 6:Generation quality and throughput comparison.Throughput is measured inimages per second. The time is recorded as theactual wallclock time. The speedup is computed*w.r.t.*SiT-XL/2.

Appendix EAdditional experimental results

Sinkhorn iterations.We provide more ablation results on number of Sinkhorn iterationsLLused when computing the Sinkhorn barycentric projection (Alg.2). We follow the same ablation setup in Sec.4.1with distribution guidance. The results in Table9show that performance improves as the number of Sinkhorn iterations increases, plateauing at around 10 iterations.

Table 7:Ablation on Sinkhorn iterationsLL. Table 8:Velocity guidance vs. distribution guidance applied to KL divergence. Table 9:Ablation onε\varepsilonwithℓ2\ell_{2}OT cost.

Velocity guidance.We further show in Table9that the proposed velocity guidance not only enhances the performance of Sinkhorn divergence-driven WGF, but also that of KL divergence. This observation aligns with our theoretical analysis in AppendixA.3, where we draw a connection to the exponentially tilted target distribution when using velocity guidance with KL divergence. Note that KL divergence with velocity guidance also outperforms the design in[14]with 8.46 FID, further highlighting the significance of our principled WGF framework.

Entropic regularization parameterε\varepsilon.We provide more results under the ablation study setting for the entropic regularization parameterε\varepsilonwith theℓ2\ell_{2}distance cost, under velocity guidance. As shown in Table9, theℓ2\ell_{2}distance OT cost yields consistently higher FIDs than the quadratic OT cost, while both achieve the best performance underε=0.05\varepsilon=0.05.

Appendix FAdditional generated samples

In this section, we provide additional generated image samples. We provide uncurated samples on ImageNet 256×\times256 in Fig.10and Fig.11, and additional samples in Fig.12with a larger guidance scale. We provide uncurated samples on FFHQ for the mode coverage experiment in Fig.13and Fig.14for Drifting Model and our W-Flow, respectively.

Figure 10:Uncurated samples generated by W-Flow, L/2 with CFGw=0.15w=0.15.Figure 11:Uncurated samples generated by W-Flow, XL/2*(Top)and B/2(Bottom)*with CFGw=0.15w=0.15.Figure 12:Uncurated samples generated by W-Flow, XL/2 with CFGw=2.0w=2.0.Figure 13:Uncurated samples generated by Drifting Model in the mode coverage experiment (Sec.4.3). None of the 100 randomly generated images is a child’s face, showing failure in covering the minority mode.Figure 14:Uncurated samples generated by W-Flow in the mode coverage experiment (Sec.4.3). 4 of the 100 randomly generated images are child faces, which is close to the 5% child face ratio in the down-sampled population.

Appendix GLimitations & Broader impacts

Limitations.

While W-Flow achieves strong one-step generation performance, several limitations remain. First, our empirical evaluation is focused on ImageNet 256×\times256 and FFHQ; extending the framework to higher-resolution, text-conditioned, video, or other multimodal generation settings remains future work. Second, training still requires pretrained feature encoders or autoencoders in our large-scale experiments, and the choice of the feature encoders remains heuristic. Third, our theoretical convergence result analyzes the particle dynamics under regularity and asymptotic assumptions, and does not fully characterize finite-network optimization or finite-compute training in high-dimensional neural generators.

Broader impacts.

This work advances fast, high-fidelity generative modeling by reducing sampling from many iterative steps to a single generator evaluation, which can lower latency and inference cost and make generative models more accessible for creative tools, simulation, data augmentation, and domain-transfer applications. At the same time, improved image generation can also amplify risks associated with synthetic media, including impersonation, disinformation, biased representations, and privacy concerns, especially in face-related applications such as age translation. Responsible deployment should therefore include safeguards such as provenance tracking or watermarking, clear usage restrictions, dataset and model documentation, consent-aware use of face data, and monitoring or filtering mechanisms when models are released beyond research settings.