arXiv CSopen access

Generative Path-Finding Method for Wasserstein Gradient Flow

Unknown · 2026 · arxiv_cs

arXiv CS · Papers · License: Open Access · 2026

machine learning, deep learning, neural networks

Generative Path-Finding Method for Wasserstein Gradient Flow Chengyu Liua , Xiang Zhoub,∗ a Department of Data Science, City University of Hong Kong, Kowloon, Hong Kong SAR

arXiv:2604.11519v1 [cs.LG] 13 Apr 2026

b Department of Mathematics, City University of Hong Kong, Kowloon, Hong Kong SAR

ARTICLE INFO

ABSTRACT

Keywords: Wasserstein gradient ﬂow Minimum action method Large deviation Geometric action Normalizing ﬂow

Wasserstein gradient ﬂows (WGFs) describe the evolution of probability distributions in Wasserstein space as steepest-descent dynamics that minimize a free-energy functional. To compute the entire path from an arbitrary initial distribution to an equilibrium distribution requires a long physical time, posing signiﬁcant challenges for both forward and implicitEuler (JKO) timemarching schemes. Grid-based Eulerian approaches for spatial variables are severely limited by the curse of dimensionality, whether solving dynamic or stationary equations for the density function. On the other hand, existing Lagrangian approaches, which leverage particle movements or generative maps, struggle to adaptively improve eﬃciency through timestep tuning. To address these limitations, we propose a generative path-ﬁnding framework for the Wasserstein gradient path (GenWGP). This approach constructs a generative ﬂow model that geometrically transports mass from an initial density to the unknown equilibrium distribution, guided by a path loss function that encodes the full trajectory, including the terminal endpoint for the equilibrium distribution. The path loss is derived from the geometric action functional based on DawsonGärtner large-deviation theory, which characterizes the most probable evolution of empirical distributions in interacting diﬀusion systems. We ﬁrst construct the path loss over an arbitrary ﬁnite time horizon using physical time parametrization and then derive the reparameterizationinvariant geometric action functional, based on the Wasserstein arc-length parametrization. GenWGP employs normalizing ﬂows to compute the geometric curve converging to the equilibrium distribution. A key feature of GenWGP is its enforcement of intrinsic constant-speed movement between adjacent layers of the generative neural network, encouraging approximately that discretized distributions remain equidistant with respect to the Wasserstein metric along the entire path, even under complex free energy landscapes. Consequently, our approach circumvents the need for intricate time-stepping schemes that are typically constrained by stepsize limitations. Instead, it facilitates stable training that is robust and independent of the speciﬁc temporal or geometric discretization of the underlying continuous descent path. Furthermore, the learned generative map serves as a reusable sampler, enabling eﬃcient computation of statistical quantities along the gradient ﬂow. We evaluate GenWGP on a variety of benchmark problems, including FokkerPlanck equations with both convex and nonconvex potentials, as well as interacting particle systems governing aggregation, aggregationdrift, and aggregationdiﬀusion dynamics. Numerical results demonstrate that GenWGP matches or exceeds the accuracy of highﬁdelity reference solutions, using as few as a dozen of discretization points for the entire gradient ﬂow path, while eﬀectively capturing complex dynamical behaviors.

1. Introduction Free energy on probability measure spaces provides a uniﬁed framework for analyzing both equilibrium states and non-equilibrium dynamics of complex systems in physics and applied mathematics [31, 43]. It is a fundamental concept to modeling phenomena ranging from diﬀusion and chemotaxis to pattern formation. In this work, we consider the following free energy  ∶ 2 (Ω) → ℝ ∪ {+∞}  (𝜌) ∶=

[ −1 ] 1 𝛽 𝑈𝑚 (𝜌(𝑥)) + 𝑉 (𝑥)𝜌(𝑥) d𝑥 + 𝑊 (𝑥 − 𝑦)𝜌(𝑥)𝜌(𝑦) d𝑥 d𝑦, ∫Ω 2 ∫Ω×Ω

(1)

where 𝜌 denotes a probability density in 2 (Ω), the space of probability measures that are absolutely continuous with respect to the Lebesgue measure and possess ﬁnite second moments. Throughout this paper, we identify the probability measure with its Radon-Nikodym derivative (density) 𝜌 with respect to the Lebesgue measure. ⋆

This work was funded by the Hong Kong General Research Funds (11318522,11308323,11304525).

∗ Corresponding author

[email protected] (C. Liu); [email protected] (X. Zhou) ORCID (s): 0009-0002-4513-1752 (C. Liu); 0000-0002-3835-3894 (X. Zhou)

C. Liu and X. Zhou: Preprint submitted to Elsevier

Page 1 of 39

Generative Wasserstein Gradient Path Method

Equation (1) comprises three contributions. The ﬁrst term, ∫Ω 𝛽 −1 𝑈𝑚 (𝜌(𝑥)) d𝑥, corresponds to the internal energy, where 𝛽 = (𝑘𝐵 𝑇 )−1 denotes the inverse temperature, 𝑘𝐵 is the Boltzmann constant, and 𝑇 is the absolute temperature. The choice of the function 𝑈𝑚 corresponds to distinct diﬀusive behavior: for instance, 𝑈1 (𝜌) = 𝜌 log 𝜌 yields the classical Boltzmann entropy [4] and recovers linear diﬀusion [52]. Nonlinear diﬀusion is captured by the power-law 𝜌𝑚 form 𝑈𝑚 (𝜌) = 𝑚−1 (for 𝑚 > 0, 𝑚 ≠ 1), encompassing the porous medium regime (𝑚 > 1) and the fast diﬀusion regime (0 < 𝑚 < 1). Key properties of the solutionssuch as mass conservation, ﬁnite-time extinction, ﬁnite propagation, and self-similaritydepend critically on the interplay between the power 𝑚 and the dimension 𝑑 [38, 10, 50]. The second term involves an external potential 𝑉 (𝑥), representing spatial conﬁnement or environmental drift. The third term accounts for pairwise interactions, modeling attraction or repulsion between ( ) particles. A notable example of (1) is the relative entropy (KullbackLeibler divergence) 𝐷KL (𝜌‖𝑝∗ ) = ∫ 𝜌 log 𝜌∕𝑝∗ d𝑥, by setting 𝛽 = 1, 𝑚 = 1, 𝑊 ≡ 0, and 𝑉 = − log 𝑝∗ . A central goal in the applications is to identify the minimizers of the free energy  in (1), which correspond to stable equilibrium distributions of the underlying particle system. The relaxation dynamics from an initial state 𝜌0 toward an equilibrium is well described by the Wasserstein Gradient Flow (WGF) - the steepest descent of  in the Wasserstein space equipped with the 2-Wasserstein metric d : { 𝜕𝑡 𝑝𝑡 = −∇d  (𝑝𝑡 ), (2) 𝑝𝑡 |𝑡=0 = 𝜌0 . Using the Benamou-Brenier diﬀerential structure of the Wasserstein space [2], the abstract gradient ﬂow (2) takes the form of a continuity equation: ) ( 𝛿 (𝑝𝑡 ) 𝜕 𝑡 𝑝𝑡 = ∇ ⋅ 𝑝𝑡 ∇ , (3) 𝛿𝑝𝑡 where the ﬁrst variation (Fréchet derivative) of (1) is given by 𝛿 (𝜌) (𝑥) = 𝛽 −1 𝑈𝑚′ (𝜌(𝑥)) + 𝑉 (𝑥) + 𝑊 (𝑥 − 𝑦)𝜌(𝑦) d𝑦. ∫Ω 𝛿𝜌 Under suitable conditions, 𝑝𝑡 converges as 𝑡 → ∞ to a steady state (equilibrium distribution or invariant measure) 𝑝∞ hat corresponds to a minimizer of  [2, 40]. The gradient ﬂow (3) encompasses a broad class of fundamental evolution equations: selecting 𝑈 (𝜌) = 𝑈1 (𝜌) = 𝜌 log 𝜌 and 𝑊 ≡ 0 recovers the Fokker-Planck equation; incorporating the nonlocal term 𝑊 ∗ 𝜌 leads to the McKean-Vlasov equation; and the setting 𝑉 = 𝑊 = 0 and 𝑈𝑚 with 𝑚 > 1 yields the porous medium equation. Solving (3) numerically poses signiﬁcant challenges. Existing methods generally fall into two categories. Eulerian schemes, such as Finite Diﬀerence or Finite Volume methods, solve the PDE (3) by discretizing the spatial domain on a ﬁxed grid or mesh [8]. While accurate in low dimensions, they suﬀer from the curse of dimensionality as grid complexity scales exponentially with 𝑑. Recent deep learning-based solvers addresses this issue by parameterizing 𝑝𝑡 with neural networks [41, 56] and deﬁning the loss function as the mean squared error of the residual for (3). While these methods oﬀer ﬂexibility and generalization for high-dimensional problems, they often fail to strictly preserve essential physical invariants, such as probability mass conservation and non-negativity, and rely on sophisticated adaptive training techniques. In contrast, Lagrangian approaches employ particle dynamics to track individual particles or trajectories, avoiding the need for ﬁxed spatial grids. Numerically, these methods approximate the underlying particle ﬂow rather than directly solving for the probability density function. It preserves essential probability properties, adapts naturally to sparse or localized solutions, and avoids the curse of dimensionality when equipped with the modern generative methods. Prominent examples include score-based methods, which typically employ forward Euler time discretization [3, 35, 27], and the JordanKinderlehrerOtto (JKO) scheme, based on implicit Euler discretization [29], as well as its neural extensions [32, 55, 25]. However, regardless of the speciﬁc spatial discretization approaches, the above existing methods fundamentally operate as time-marching schemes. Thesse methods rely on sequential updates of the density or particle movements over short time steps, requiring the solution of a subproblem at each step.This creates a critical dependency on the timestep size: smaller step sizes - while necessary for accuracy and stability - signiﬁcantly increase computational C. Liu and X. Zhou: Preprint submitted to Elsevier

Page 2 of 39

Generative Wasserstein Gradient Path Method

costs, making longtime simulations particularly prohibitive. In principle, adaptive tuning of the step size, as done in traditional Euclidean spaces, is feasible. However, in practice, it is constrained by the intrinsic complexity of the Wasserstein metric. Moreover, the gradient ﬂow requires inﬁnite time to reach the exact equilibrium. Truncating the ﬂow to a ﬁnite time interval can result in bias, causing the ﬁnal state to deviate from the true equilibrium. There is no deﬁnitive way to select the truncation time interval a prior, as the accuracy of the truncation varies on a case-by-case basis. To address these limitations, we propose a Generative Wasserstein Gradient Path (GenWGP) method, which re-frames the problem from sequential time-stepping to global path optimization, emphasizing the geometric path viewpoint of the underlying Wasserstein gradient ﬂow. The Wasserstein Gradient Path (WGP) refers to a curve representing the Wasserstein-2 gradient ﬂow from an arbitrary initial density toward a nearby equilibrium density, parameterized either by the original physical time or by an alternative geometric parameter, such as the Wasserstein arc-length. To achieve such a path optimization, we connect the gradient ﬂow with the large deviation principle [13, 14, 19, 1], and construct the loss functional in path space based on the action functional of the DawsonGärtner large deviation principle for interacting particles [14]. The gradient ﬂow corresponds to the zero-action path [20, 18, 57, 1]. We shall see later that, our path loss function is equivalent to the mean-squared residual of the time-dependent PDE (3) measured in the 𝐻𝜌−1 metric associated with the Wasserstein structure. Thus our variational principle is to learn the entire trajectory up to a prescribed terminal time 𝑇 , rather than updating the solution one short time interval after another. With a suﬃciently large 𝑇 , our method identiﬁes the learned 𝜌𝑇 as an approximation to the equilibrium distribution. If the extra penalty of  (𝜌𝑇 ) is added to the path-loss functional, the optimization method can simultaneously determines the ﬁnal terminal state an equilibrium and resolve the whole gradient path quite well. Compared with the explicit time-marching schemes, our resulting time-discretized path loss avoids the stability restrictions and permits to take large time step size. But there still remains two accuracy issues: the choice of the truncation horizon 𝑇 and the non-uniform accuracy induced by a non-optimal temporal grid. In long-time relaxation problem, we face the typical situation that the early stage of the evolution may contain rapid transients, while the later stage approaches equilibrium only in a very slow pace. A physical-time parametrization is therefore not the most eﬃcient one for describing the whole trajectory. The classical action in the Dawson-Gärtner large deviation principle is expressed as an integral over physical time. Motivated by the Maupertuis principle, as in the geometric minimum action method for the Freidlin-Wentzell action functional [22], we reformulate the original Dawson-Gärtner action functional into a geometric form that is invariant under time reparameterization and free of the explicit dependence on the horizon truncation 𝑇 . In this formulation, the free energy at the terminal state is always added to the path-loss functional to guide the convergence to the equilibrium state. The optimization of the curve allows us to represent the same gradient path by its geometry alone, independently of how fast or slow the path is traversed in physical time. In particular, it enables more suitable parametrizations for long-time relaxation, such as the Wasserstein-2 arc-length parametrization. Our GenWGP ﬁrst computes the desired curve geometrically, and then recovers the corresponding physical-time parameter through a simple post-processing step based on the derived nonlinear relation between the geometric and physical parametrizations. In this way, the geometric formulation retains the essential dynamic information while avoiding the ineﬃciency of resolving the slow tail directly on a physical-time grid. Numerically, to parameterize a Wasserstein path with 𝐾 discrete points, our GenWGP method employs a single Normalizing Flow (NF) neural network [30, 39] composed of 𝐾 stacked layers, each of which represents a diﬀeomorphic map transporting mass between two neighboring distributions along the path. This yields a meshfree, generative solver that computes the whole Wasserstein gradient path globally from the initial distribution to the ﬁnal equilibrium distribution. Enforcing a meaningful arc-length parametrization is a central challenge in path-based methods. Existing approaches typically rely either on global path reparametrization, which is mainly designed for ﬁnite-dimensional state spaces [17], or on penalizing the Riemannian norm of the tangent vector to enforce constant speed [21]. In contrast, GenWGP exploits the structure of Normalizing Flows together with the Monge representation of the Wasserstein distance. For two neighboring path images 𝑝𝑘−1 and 𝑝𝑘 , the Wasserstein distance can be characterized by minimizing the 𝐿2 mean-square distance between admissible transport maps sharing the same reference measure 𝜌0 . This makes it natural to lift the geometric action from the density level to the map level and to train the NF directly through the discrete transport increments between neighboring layers. When the corresponding Monge minimizer is realized, the segment cost coincides exactly with the Wasserstein distance. Therefore, imposing a constant-speed constraint on these lifted segment lengths provides a practical and geometrically consistent mechanism for distributing the 𝐾 discretized path images approximately evenly along the Wasserstein curve. C. Liu and X. Zhou: Preprint submitted to Elsevier

Page 3 of 39

Generative Wasserstein Gradient Path Method

Main contributions. The main contributions of this work can be summarized as follows. 1. A global path-optimization formulation. Instead of computing the WGF (3) by sequential time marching, we reformulate the problem as the optimization of an entire probability path. This shifts the numerical viewpoint from local time stepping to a global least-action principle in path space, and provides a variational framework for learning the full relaxation trajectory from an initial distribution toward equilibrium. 2. A generative Lagrangian parameterization. We represent the discrete probability path by one normalizing ﬂow whose layerwise composition plays the role of transport along the path. This yields a mesh-free Lagrangian solver that avoids Eulerian spatial grids and directly parameterizes the evolving distributions and particle trajectories in a uniﬁed manner. 3. Physical-time and geometric action formulations. Starting from the Dawson-Gärtner action functional [14], we derive both a physical-time path loss on a prescribed horizon and a geometric, reparameterization-invariant formulation for long-time relaxation problems. The geometric formulation removes the explicit dependence on the physical-time horizon and is therefore better suited to computing paths that approach equilibrium. 4. A practical discrete optimization framework. We construct trainable discrete objectives for both formulations using Crank-Nicolson-type temporal discretization, Monte Carlo particle approximation, and a lifted map-level geometric loss. In the geometric setting, we further introduce an arc-length regularization and a terminal freeenergy penalty to stabilize training and to guide the endpoint toward a low-energy equilibrium state. 5. Supporting analytical results and numerical validation. Under the stated regularity assumptions, we establish an a priori KL-divergence estimate for the physical-time formulation, a trajectory-error decomposition for the discrete physical-time scheme, and a consistency result for the discrete geometric objective. Numerical experiments on Fokker-Planck and interacting-particle systems (aggregation, aggregation-drift, and aggregationdiﬀusion dynamics) demonstrate that the proposed framework can accurately approximate both relaxation paths and terminal states. Related work We review related works in the recent literature to highlight the distinct features of our approach. Equilibrium Solvers. Several recent works focus solely on identifying the equilibrium 𝜌∞ by training NFs via strong/weak/variational formulations of the static PDE 𝜕𝑡 𝑝𝑡 = 0 in (3) [48, 6, 54, 7]. While these methods eﬃciently ﬁnd the steady state using adaptive sampling, they do not capture the path of evolution. Score-Based and Flow-Based Time-Marching Methods. This new class of methods on discrete physical time grid learns the time-dependent velocity ﬁeld [𝑝𝑡 ](𝑥) or Stein score (∇ log 𝑝𝑡 ) [3, 44, 33, 27]. These approaches typically use forward Euler integration to advance the solution. Consequently, they inherit the standard limitations of ODE solvers, including the need for small time steps and lack of global error control. Furthermore, all score-based methods are essentially restricted to linear diﬀusion (entropy-based) models. JKO-based Variational Schemes. The JKO scheme has inspired methods like JKO-iFlow [55], EVNN [25], and Deep JKO [32], which use NFs to solve the variational subproblem at each time step. While these methods correctly leverage the Lagrangian framework, they solve a sequence of optimization problems at 𝑡1 , 𝑡2 , … , which is computationally expensive and prone to error drift. In contrast,our GenWGP optimizes the global action in a single training loop. Parametric WGF. The Parameterized WGF scheme [34, 36] projects the gradient ﬂow onto a neural network’s parameter space. This transforms the PDE into an ODE on parameters but requires to compute and invert the Fisher Information Matrix. This inversion is computationally prohibitive for a large number of network parameters, a bottleneck that GenWGP avoids entirely. Minimum Action Methods and Dawson-Gärtner Large Deviation Principle. Our approach is based on the equivalence of the zero-action path and the gradient ﬂow. The background of large deviation on Freidlin-Wentzell and Dawson-Gärtner theories can be found in [20, 14] as well as [1]. An nonexhaustive list for the development of minimum action methods, developed either in the ﬁnite dimensional conﬁguration space ℝ𝑑 or in the function space 𝐿2 (ℝ𝑑 ) include [18, 57, 53, 23, 49, 47]. However, none of these approaches have been extended to the Wasserstein manifold. More recent pathﬁnding frameworks that combine variational principles with deep neural networks are presented in [21, 45]. The remainder of this paper is organized as follows. Section 2 introduces the mathematical background on Wasserstein gradient ﬂows, the Dawson-Gärtner large deviation theory, and normalizing ﬂows. Section 3 develops the physical-time GenWGP formulation for approximating Wasserstein gradient ﬂows on a ﬁnite time horizon. C. Liu and X. Zhou: Preprint submitted to Elsevier

Page 4 of 39

Generative Wasserstein Gradient Path Method

Section 4 then extends this framework to a geometric, reparameterization-invariant formulation for paths converging to equilibrium. Section 5 presents numerical experiments to demonstrate the accuracy and eﬀectiveness of the proposed framework. Finally, Section 6 concludes the paper.

2. Background and Preliminaries This section lays out the foundational concepts underpinning our method. We begin by reviewing the Riemannian geometry of the Wasserstein space and its associated gradient ﬂow structure. We then present two complementary variational characterizations of the dynamicsone local and one global. Finally, we introduce normalizing ﬂows as a ﬂexible framework for parameterizing transport maps.

2.1. Wasserstein Geometry and Gradient Flows

Let 2 (Ω) denote the space of absolutely continuous probability measures on Ω ⊂ ℝ𝑑 with ﬁnite second moments. We identify each measure with its density function 𝜌. The Wasserstein-2 distance between 𝜌0 , 𝜌1 ∈ 2 (Ω) is deﬁned by the optimal transport problem [52]: d2 (𝜌0 , 𝜌1 ) =

inf

𝑇# 𝜌0 =𝜌1 ∫Ω

|𝑥 − 𝑇 (𝑥)|2 𝜌0 (𝑥) d𝑥,

(4)

where 𝑇 ∶ Ω → Ω is a transport map and 𝑇# 𝜌0 denotes the pushforward measure. Tangent Space and Metric. The space 2 (Ω) possesses a formal Riemannian structure [38]. Given 𝜌 ∈ 2 (Ω), the tangent space at 𝜌 can then be identiﬁed with density perturbations 𝑓 = 𝜕𝑡 𝜌𝑡 |𝑡=0 , namely { } | ( ) 𝜌 2 (Ω) = 𝑓 ∈ 𝐶 ∞ (Ω) || 𝑓 (𝑥)d𝑥 = 0, ∃𝜓 ∈ 𝐶 ∞ (Ω) such that 𝑓 = −∇ ⋅ 𝜌∇𝜓 , (5) | ∫Ω where 𝜓 is uniquely determined up to an additive constant under suitable boundary conditions. The Riemannian metric at 𝜌 is given by the weighted negative Sobolev norm ‖ ⋅ ‖−1,𝜌 . For two tangent vectors 𝑓𝑖 = −∇ ⋅ (𝜌∇𝜓𝑖 ) (𝑖 = 1, 2), the corresponding inner product is deﬁned by ⟨𝑓1 , 𝑓2 ⟩−1,𝜌 = ∫Ω ∇𝜓1 (𝑥) ⋅ ∇𝜓2 (𝑥)𝜌(𝑥) d𝑥. Equivalently, the induced norm admits the dynamic characterization ‖𝑓 ‖2−1,𝜌 =

inf

𝐮∶𝑓 =−∇⋅(𝜌𝐮) ∫Ω

|𝐮(𝑥)|2 𝜌(𝑥) d𝑥,

(6)

where the inﬁmum is attained at 𝐮 = ∇𝜓 whenever 𝑓 = −∇ ⋅ (𝜌∇𝜓). Wasserstein Gradient Flow. The Wasserstein gradient of a functional  (𝜌) (1) is given by ∇d  (𝜌) = −∇ ⋅ ( ) 𝜌∇ 𝛿𝛿𝜌(𝜌) . The gradient ﬂow equation 𝜕𝑡 𝑝𝑡 = −∇d  (𝑝𝑡 ), takes the form of a continuity equation: ( ) 𝜕𝑡 𝑝𝑡 = −∇ ⋅ 𝑝𝑡 [𝑝𝑡 ] , where [𝑝𝑡 ] is the speciﬁc Lagrangian velocity ﬁeld: ( ) ) ( 𝛿 (𝑝𝑡 ) [𝑝𝑡 ](𝑥) = −∇ (𝑥) = −∇ 𝛽 −1 𝑈𝑚′ (𝑝𝑡 (𝑥)) + 𝑉 (𝑥) + (𝑊 ∗ 𝑝𝑡 )(𝑥) . 𝛿𝜌

(7)

(8)

Here (𝑊 ∗ 𝑝𝑡 )(𝑥) = ∫Ω 𝑊 (𝑥 − 𝑦)𝑝𝑡 (𝑦)d𝑦 and 𝑈𝑚′ (⋅) denotes the derivative of the internal-energy density.

2.2. Variational and Dynamic Formulations We interpret the WGF (7) as two complementary variational characterizations. The local-in-time formulation motivates time-marching algorithms, while the global-in-time formulation motivates our path-ﬁnding approach. Local: Onsager’s Principle (Least Dissipation). At any given time 𝑡, the system selects an instantaneous velocity ﬁeld 𝐯 that minimizes the sum of dissipation potential and energy change rates [37]: { } { } 1 𝑑 1 𝛿 2 2 𝐯𝑡 = argmin |𝐯(𝑥)| 𝑝𝑡 (𝑥)d𝑥 +  (𝑝𝑡 ) = argmin |𝐯(𝑥)| 𝑝𝑡 (𝑥)d𝑥 + 𝜕𝑝 (𝑝 )d𝑥 , (9) ∫ 𝑡 𝑡 𝛿𝜌 𝑡 2 ∫Ω 𝑑𝑡 2 ∫Ω 𝐯 𝐯 C. Liu and X. Zhou: Preprint submitted to Elsevier

Page 5 of 39

Generative Wasserstein Gradient Path Method

subject to the continuity constraint 𝑡 𝑝𝑡 + ∇ ⋅ (𝑝𝑡 𝐯) = 0. By integral by part, the unique minimizer of (9) is exactly the ( 𝜕)

𝑑 velocity ﬁeld (8), i.e., 𝐯∗𝑡 = −∇ 𝛿 = [𝑝𝑡 ]. At this optimal velocity, 𝑑𝑡  (𝑝𝑡 ) = − ∫Ω ‖[𝑝𝑡 ]‖2 𝑝𝑡 d𝑥 ≤ 0, ensuring 𝛿𝜌 energy dissipation. Global: Least Action Principle. Integrating the gradient ﬂow (2) over the time interval [0, 𝑇 ] formally yields the Dawson-Gärtner action functional [13]. For a path 𝑝 = (𝑝𝑡 )𝑡∈[0,𝑇 ] , the action is deﬁned as:

⎧1 𝑇 ‖ ‖2 𝜕𝑡 𝑝𝑡 + ∇d  (𝑝𝑡 )‖ d𝑡, if 𝑝𝑡 is absolutely continuous and the integral converges, ⎪2 ∫ ‖ ‖ ‖−1,𝑝𝑡 𝑆𝑇 [𝑝] ∶= ⎨ 0 ⎪ otherwise. ⎩+∞,

(10)

where the norm ‖⋅‖2−1,𝑝 is deﬁned in (6). For entropy-driven diﬀusion, namely the case 𝑚 = 1, this 𝑆𝑇 is known as 𝑡 the rate function in the large deviation theory of interacting diﬀusion particle systems [12, 13, 14], and it quantiﬁes the likelihood of observing a prescribed trajectory of the empirical measure (𝑝𝑁 𝑡 )0≤𝑡≤𝑇 of 𝑁 diﬀusion particles following the McKean-Vlasov system. More precisely, consider the system of 𝑁 interacting Itô processes 𝑁

√ 1 ∑ d𝑋𝑖 (𝑡) = −∇𝑉 (𝑋𝑖 (𝑡)) d𝑡 − ∇𝑊 (𝑋𝑖 (𝑡) − 𝑋𝑗 (𝑡)) d𝑡 + 2𝛽 −1 d𝐵𝑡𝑖 , 𝑁 𝑗=1

𝑖 = 1, … , 𝑁,

(11)

1 ∑𝑁 where {𝐵𝑡𝑖 }𝑁 are independent standard Brownian motions. The associated empirical measure 𝑝𝑁 𝑡 ∶= 𝑁 𝑖=1 𝛿𝑋𝑖 (𝑡) 𝑖=1 converges, under standard assumptions on 𝑉 and 𝑊 , to a deterministic measure 𝑝(𝑡) solving the McKean-Vlasov SDE

d𝑋𝑡 = −∇𝑉 (𝑋𝑡 ) d𝑡 − (∇𝑊 ∗ 𝑝(𝑡))(𝑋𝑡 ) d𝑡 +

√ 2𝛽 −1 d𝐵𝑡 ,

𝑝(𝑡) ∶= Law(𝑋𝑡 ),

(12)

whose time-marginal density satisﬁes the WGF (7). The ﬂuctuations of the empirical-measure path {𝑝𝑁 𝑡 }𝑡∈[0,𝑇 ] around the macroscopic limit are, for entropy-driven diﬀusion (namely the case 𝑚 = 1), formally described by a path-space large deviation principle of Dawson-Gärtner type [13]. For a ﬁxed 𝑇 , this suggests that the path law of 𝑝𝑁 = (𝑝𝑁 𝑡 )𝑡∈[0,𝑇 ] admits a path-space large-deviation description with speed 𝑁 and rate functional 𝑆𝑇 . More precisely, for a prescribed absolutely continuous path 𝑞 = (𝑞𝑡 )𝑡∈[0,𝑇 ] , one formally writes ( ) ( ) Pr 𝑝𝑁 ≈ 𝑞 ≍ exp −𝑁𝑆𝑇 [𝑞] . Here the large deviation principle is understood in the usual setwise ( sense; )heuristically, for suﬃciently small Wasserstein neighborhoods of 𝑞, one expects probabilities of order exp −𝑁𝑆𝑇 [𝑞] . In particular, for admissible paths 𝑞 for which the action is well-deﬁned, 𝑆𝑇 [𝑞] = 0 if and only if 𝑞𝑡 satisﬁes the WGF equation (7) almost everywhere in 𝑡 ∈ [0, 𝑇 ]. Hence the deterministic gradient ﬂow trajectory is a zero-action path, reﬂecting the fact that it arises as the law-of-large-numbers limit of the underlying interacting particle system. Motivated by this large-deviation structure, for more general free energies of Wasserstein gradient-ﬂow type, we still use the same action functional as the natural global variational principle for path computation. For a prescribed terminal distribution 𝜌1 , it is therefore natural to consider the associated minimum action problem [18, 49, 57]: inf 𝑝

𝑆𝑇 (𝑝),

subject to 𝑝0 (𝑥) = 𝜌0 (𝑥), 𝑝𝑇 (𝑥) = 𝜌1 (𝑥).

(13)

This variational problem determines the least-action path connecting 𝜌0 to 𝜌1 . To ﬁnd the zero-action path, we simply drop the terminal state and minimize 𝑆𝑇 (𝑝) only with the initial 𝑝0 = 𝜌0 , since the solution of the WGF (7) indeed gives zero action under this initial constraint.

2.3. Normalizing Flows To numerically approximate the probability path, we employ Normalizing Flows (NFs)[42, 15, 11, 16, 30, 39]. An NF is a deep generative model that represents a complex probability distribution as the pushforward of a C. Liu and X. Zhou: Preprint submitted to Elsevier

Page 6 of 39

Generative Wasserstein Gradient Path Method

simple reference distribution 𝜌ref (e.g., the standard Gaussian) through an invertible map Φ. This map is typically parameterized as a composition of 𝐾 invertible neural network blocks Ψ𝑘 (commonly referred to as layers): Φ = Ψ𝐾 ◦ ⋯ ◦Ψ2 ◦Ψ1 . For a sample 𝑧0 ∼ 𝜌ref , the target density at 𝑧𝐾 = Φ(𝑧0 ) is computable via the change-of-variables formula: log 𝑝(𝑧𝐾 ) = log 𝜌ref (𝑧0 ) −

𝐾 ∑

| 𝜕Ψ𝑘 || log ||det . 𝜕𝑧𝑘−1 || | 𝑘=1

(14)

Various architectures facilitate eﬃcient computation. For instance, RealNVP coupling ﬂows [15] use triangular Jacobians to ensure linear-cost determinant evaluation (𝑂(𝑑)). More expressive variants include spline-based ﬂows [16] which implement ﬂexible monotone transforms, and diﬀeomorphic non-uniform B-spline ﬂows [24] which oﬀer 𝐶 2 smooth, bi-Lipschitz maps with controlled regularity. In our framework, we interpret the layer-wise composition of the ﬂow as a temporal discretization of the Lagrangian trajectory. By associating each layer Ψ𝑘 with the transport over a time step Δ𝑡, the intermediate activations 𝑧𝑘 = (Ψ𝑘 ◦ ⋯ ◦Ψ1 )(𝑧0 ) represent the positions of particles at time 𝑡𝑘 . Consequently, the full network Φ parameterizes the . This design diﬀers from Neural ODEs [11] as it retains the exact tractability of the entire discrete trajectory (𝑝𝑡𝑘 )𝐾 𝑘=0 density via Eq. (14) at every layer,circumventing the need for ODE solvers, numerical integration errors, and posthoc density estimation. As a result, the action functional can be evaluated directly and eﬃciently, making the method particularly wellsuited for pathbased variational problems in the Wasserstein space.

3. The Generative Wasserstein Gradient Path (GenWGP) Method for Wasserstein Gradient Flow This section develops our Lagrangian, generative framework for computing Wasserstein gradient path. We ﬁrst present a physical-time parameterized formulation, which learns the WGF dynamics on a ﬁxed time horizon [0, 𝑇 ].

3.1. Physical Time Parameterized Lagrangian Representation of the Action Functional Building on the variational principles established in Section 2, we derive a computational framework to minimize the action functional 𝑆𝑇 [𝑝] (10) by employing a Lagrangian particle approximation of the density 𝑝𝑡 . The continuum weighted 𝐻𝑝−1 norm appearing in the rate functional (10) is then replaced by a discrete 𝐿2 norm over a ﬁnite ensemble 𝑡 of particle trajectories, thereby avoiding the need to solve an elliptic PDE at every time step. We ﬁrst take the continuous time perspective and characterize the trajectory of the probability density using a time-dependent velocity ﬁeld 𝐟𝑡 ∶ Ω → ℝ𝑑 . This ﬁeld determines the motion of particles via the characteristic ODE: 𝜕𝑡 Φ(𝑡, 𝑧) = 𝐟𝑡 (Φ(𝑡, 𝑧)),

Φ(0, 𝑧) = 𝑧,

(15)

where 𝑧 ∼ 𝜌0 represents the initial Lagrangian coordinate. The time-dependent density 𝑝𝑡 is deﬁned as the pushforward 𝑝𝑡 = Φ(𝑡, ⋅)# 𝜌0 . By the transport theorem, 𝑝𝑡 satisﬁes the continuity equation driven by 𝐟𝑡 : 𝜕𝑡 𝑝𝑡 + ∇ ⋅ (𝑝𝑡 𝐟𝑡 ) = 0.

(16)

This path (𝑝𝑡 ) is used to match the target Wasserstein gradient ﬂow (7) governed by the thermodynamic driving force [𝑝𝑡 ] in (8), via the least action principle of minimizing the action functional (10). To do this, we substitute 𝜕𝑡 𝑝𝑡 in 𝑆𝑇 by the continuity equation (16) and see the “residual” term becomes the divergence of the velocity mismatch: ( ) ( ) 𝜕𝑡 𝑝𝑡 − (−∇d  (𝑝𝑡 )) = −∇ ⋅ 𝑝𝑡 (𝐟𝑡 − [𝑝𝑡 ]) = −∇ ⋅ 𝑝𝑡 (𝜕𝑡 Φ(𝑡, 𝑧) − [𝑝𝑡 ]) . Applying the deﬁnition of 𝐻𝜌−1 norm (6), we arrive at the following minimization problem for the loss function: 𝑝

𝑇

1 ‖𝐟 (𝑥) − [𝑝𝑡 ](𝑥)‖2 𝑝𝑡 (𝑥) d𝑥 d𝑡 ‖ 𝑝 𝐟 ∶𝜕𝑡 𝑝𝑡 +∇⋅(𝑝𝑡 𝐟𝑡 )=0 2 ∫0 ∫Ω ‖ 𝑡 [ ] 𝑇 1 2 = inf inf 𝔼𝑧∼𝜌0 ‖ 𝜕𝑡 Φ(𝑡, 𝑧) − [𝑝𝑡 ](Φ(𝑡, 𝑧))‖ d𝑡 =∶ inf 𝐽 [Φ]. ‖ ‖ 𝑝 Φ∶(Φ ) 𝜌 =𝑝 2 ∫ Φ

inf 𝑆𝑇 [𝑝] = inf

inf

𝑡 # 0

𝑡

(17)

C. Liu and X. Zhou: Preprint submitted to Elsevier

Page 7 of 39

Generative Wasserstein Gradient Path Method

This formulation seeks a ﬂow map Φ whose instantaneous kinematic velocity matches the thermodynamic driving force [𝑝𝑡 ] given by (8). This formulation can be interpreted as a Physics-Informed Neural Network (PINN) in the Wasserstein space, but unlike standard Eulerian version of the PINNs that minimizes PDE residuals for 𝑝𝑡 on spatial grids, our method minimizes the “residual” of the ﬂow map Φ as governed by the neural ODE (15). This Lagrangian perspective for the PINN in the Wasserstein space naturally builds the connections of many generative models and the evolution PDE for the density. We rigorously justify our path loss function in (17) below by providing an a priori bound on the error measured in Kullback-Leibler divergence. Theorem 1 (KL Divergence Bound). Let 𝑝𝑡 be the density induced by the ﬂow Φ via (15)(16) and 𝑝̂𝑡 be the exact solution to the WGF (7). Under the regularity assumptions 1, there exist constants 𝛼, 𝛾 > 0 such that sup 𝐷KL (𝑝𝑡 ‖̂ 𝑝𝑡 ) ≤ exp(𝛾𝑇 )𝛼𝐽 [Φ].

(18)

𝑡∈[0,𝑇 ]

PROOF. See Appendix A.1.

3.2. Discretization via Normalizing Flows. In the numerical implementation, we parameterize the ﬂow map Φ(𝑡, ⋅) at discrete time points using a Normalizing Flow. We partition the time horizon [0, 𝑇 ] into 𝐾 intervals - for instance, with a uniform step size Δ𝑡 = 𝑇 ∕𝐾. The particle positions at time 𝑡𝑘 = 𝑘Δ𝑡 are modelled by the generative map Φ𝑘 (𝑧) = (Ψ𝜃𝑘 ◦ … ◦Ψ𝜃1 )(𝑧), and the velocity ﬁeld 𝐟𝑡 is approximated by the ﬁnite diﬀerence scheme: 𝐟𝑡𝑘 (Φ𝑘 (𝑧)) ≈

Φ𝑘 (𝑧) − Φ𝑘−1 (𝑧) . Δ𝑡

(19)

To achieve secondorder temporal accuracy for the path loss (17), we employ the CrankNicolson scheme, which yields the following discrete empirical loss: 𝐾 [Φ] = 𝐽𝑁

𝐾

𝑁

‖ Δ𝑡 ∑ ∑ ‖ ‖ Φ𝑘 (𝑧𝑖 ) − Φ𝑘−1 (𝑧𝑖 ) − 𝑁 [𝑝𝑘 ](Φ𝑘 (𝑧𝑖 )) + 𝑁 [𝑝𝑘−1 ](Φ𝑘−1 (𝑧𝑖 )) ‖ , ‖ ‖ 𝑁 𝑘=1 𝑖=1 ‖ Δ𝑡 2 ‖ 2

(20)

where {𝑧𝑖 }𝑁 are i.i.d. samples drawn from 𝜌0 . Here, 𝑁 [𝑝𝑘 ] denotes the empirical approximation of the velocity ﬁeld 𝑖=1 (8) computed using the particle batch at step 𝑘: ) ( 𝑁 1 ∑ (𝑗) −1 ′ ∶= Φ𝑘 (𝑧𝑗 ). (21) 𝑁 [𝑝𝑘 ](𝑥) ∶= −∇ 𝛽 𝑈𝑚 (𝑝𝑘 (𝑥)) + 𝑉 (𝑥) + 𝑊 (𝑥 − 𝑥𝑘 ) , 𝑥(𝑗) 𝑘 𝑁 𝑗=1 The density 𝑝𝑘 (𝑥) is computed via the formula (14). The training procedure is summarized in Algorithm 1. Algorithm 1 GenWGP (Physical-Time) Require: NF parameters {𝜃𝑘 }𝐾 , initial distribution 𝜌0 , horizon 𝑇 , steps 𝐾. 𝑘=1 1: for each training iteration do 2: Sample a batch of 𝑁 particles {𝑧𝑖 }𝑁 ∼ 𝜌0 . 𝑖=1 3:

Compute particle trajectories 𝑥(𝑖) ∶= Φ𝑘 (𝑧𝑖 ) for 𝑘 = 0, … , 𝐾. 𝑘

For each step 𝑘, compute densities 𝑝𝑘 (𝑥(𝑖) ) via (14) and empirical velocities 𝑁 [𝑝𝑘 ](𝑥(𝑖) ) via (21). 𝑘 𝑘 𝐾 5: Evaluate the discrete loss 𝐽𝑁 [Φ] using (20). 6: Update parameters {𝜃𝑘 } via gradient descent (e.g., Adam). 7: end for 8: return Optimized NF parameters {𝜃𝑘 }𝐾 . 𝑘=1 4:

As Δ𝑡 → 0 and 𝑁 → ∞, the discrete loss (20) consistently approximates the continuous action (17). The following result gives a trajectory-error estimate in terms of a consistency residual measured against the exact Crank-Nicolson driving force. C. Liu and X. Zhou: Preprint submitted to Elsevier

Page 8 of 39

Generative Wasserstein Gradient Path Method

Theorem 2 (Residual-based trajectory-error bound). Let 𝑋 ∗ (𝑡, 𝑧) be the exact characteristic trajectory associated with the velocity ﬁeld (8) starting from 𝑧 ∼ 𝜌0 , and let 𝑋𝑘𝑁 (𝑧) ∶= Φ𝑘 (𝑧) denote the discrete numerical trajectory generated by Algorithm 1. We deﬁne the consistency residual 𝜀 as the maximum mismatch between the kinematic velocity and the Crank-Nicolson driving force: ‖ Φ (𝑧) − Φ (𝑧) 𝑁 [̂ 𝑝𝑡𝑘 ](Φ𝑘 (𝑧)) + 𝑁 [̂ 𝑝𝑡𝑘−1 ](Φ𝑘−1 (𝑧)) ‖ ‖ ‖ 𝑘−1 𝜀 ∶= sup ‖ 𝑘 − ‖. ‖ Δ𝑡 2 𝑘,𝑧 ‖ ‖ ‖

(22)

Under the assumptions in Appendix 1, the expected trajectory error at any step 𝑡𝑘 satisﬁes: ⎤ ⎡ ⎢ ( ⎥ ) ‖ ‖ ⎥ 𝑡𝑘 . 𝔼𝑧∼𝜌0 ‖𝑋 ∗ (𝑡𝑘 , 𝑧) − 𝑋𝑘𝑁 (𝑧)‖ ≤ ⎢  𝑁 −1∕2 + (𝜀) + (Δ𝑡2 ) ‖ ‖ ⎢ ⏟⏞⏞⏞⏟⏞⏞⏞⏟ ⎥ ⏟⏟⏟ ⏟⏟⏟ ⎥ ⎢ Consistency Residual Discretization Error Sampling Error ⎦ ⎣

(23)

PROOF. See Appendix A.2. In particular, Eq. (23) shows that the trajectory error is controlled by three contributions: the sampling error 𝑂(𝑁 −1∕2 ), the consistency residual 𝑂(𝜀), and the Crank-Nicolson discretization error 𝑂(Δ𝑡2 ). Remark 1. Our path formulation allows ﬂexible time discretization with no essential implementation barriers. Here we employ the Crank-Nicolson scheme, which uses the average of the velocity 𝑁 evaluated at 𝑡𝑘 and 𝑡𝑘+1 . Under suitable regularity, this gives a second-order accurate discretization in time. By incorporating information from both ends of each time interval, the scheme typically provides a more faithful approximation of the continuous trajectory than the ﬁrstorder explicit or implicit methods. Remark 2. While the above results extend straightforwardly to nonuniform time grids, the adaptive timemeshing strategy employed in the adaptive minimum action method [57, 47] is not an easy task in the inﬁnitedimensional space 2 (Ω). The reason lies in the architecture of the normalizing ﬂow: each layer corresponds exactly to a speciﬁed time point. Consequently, adjusting to a new time mesh requires an expensive reﬁtting of the entire network [55], which is fundamentally diﬀerent from the simple component-wise interpolation in ﬁnitedimensional settings. Remark 3. Since we are interested in the zero-action path for the WGP, we can introduce a weight function 𝜔𝑡 for the path loss (17) or the discrete 𝜔𝑘 for the discrete loss (20) as in the score-training approach [46]. For example, to enhance the path accuracy near the initial state, we may use a larger weight 𝜔1 than 𝜔2 , ⋯ , 𝜔𝐾 . This weighted training eﬀectively changes the 𝐿2 norm in (20) and may be interpreted as a new large deviation rate function for (12) associated with a time-dependent 𝛽𝑡 for the noise amplitude. For simplicity, we use the constant weight in our algorithms and examples.

3.3. Connections to Classical Results Our Lagrangian actionminimization framework (17) provides a valuable len for reinterpreting existing numerical topics, such as numerical timemarching schemes, geometric optimal transport, and the probabilistic theory of interacting particle systems. 1. Time-discretization schemes: The path loss functional 𝐽 (Φ) in (17), if restricted in a single time interval from 𝑡𝑘 to 𝑡𝑘+1 sequentially, is closely related to several well-known time-discretization schemes for WGFs: Φ

(𝑧)−Φ (𝑧)

• Forward Euler (explicit): choosing 𝑘+1 Δ𝑡 𝑘 − 𝑁 [𝑝𝑘 ](Φ𝑘 (𝑧)) in (20) results in an explicit marching scheme. This is related to scorebased transport modeling methods [3, 35, 27], where a score function, instead of the entire velocity ﬁeld  independently trained by a neural network based on data. • Backward Euler (implicit) and JKO: Replacing the midpoint velocity in (20) by the backward endpoint velocity produces a backward-Euler-type residual. The stationary condition is formally consistent with the Euler-Lagrange equation associated with the JKO minimization [29] { } 1 2 𝑝𝑘+1 = argmin  [𝜌] + d (𝜌, 𝑝𝑘 ) . 2Δ𝑡 𝜌∈2 (ℝ𝑑 ) This establishes the connection between our action-based formulation and implicit variational timestepping methods such as Deep JKO [32], JKO-iFlow [55] and EVNN [25]. C. Liu and X. Zhou: Preprint submitted to Elsevier

Page 9 of 39

Generative Wasserstein Gradient Path Method

2. Relation to dynamic optimal transport: In the special case  ≡ 0, the driving force vanishes and (17) reduces to a kinetic-energy minimization over transport maps. In this sense, the proposed formulation is consistent with the Benamou-Brenier dynamic characterization of the Wasserstein-2 distance [51]. 3. Entropy-driven ﬂows, Fokker-Planck equation, and score-based diﬀusion models: For 𝑈 = 𝑈1 (𝜌) = 𝜌 log 𝜌, the WGF becomes the Fokker-Planck equation. Minimizing the action (17) is equivalent to: [ ] 𝑇 ( )‖2 1 ‖ 𝔼𝑥∼𝑝𝑡 ‖𝐟𝑡 (𝑥) − −∇ log 𝑝𝑡 (𝑥) − ∇𝑉 (𝑥) ‖ d𝑡. ‖ ‖ 2 ∫0 This is a continuous-time analogue of score matching [28, 46]. We emphasize that while most of these results are limited to gradient ﬂows corresponding to the diﬀusion or entropy case of the free energy (3) with 𝑚 = 1 only, our actionminimization framework applies more broadly to gradient ﬂows in the Wasserstein space. This formulation can even characterize transition paths between two distinct local minima [12], although this application is not further pursued in the present work.

4. The GenWGP Method for Wasserstein Gradient Flow Converging to Equilibrium Distribution The physical-time path formulation in Section 3 provides a Lagrangian path loss to learn the gradient-ﬂow dynamics over a speciﬁed ﬁnite horizon [0, 𝑇 ]. However, when the purpose is to capture the full relaxation from an initial state 𝜌0 to an equilibrium 𝜌∞ - namely, a stationary point of the free energy - we suﬀer from the ﬁnite time interval truncation. Without any prior knowledge about the truncation error between 𝜌𝑇 and the true 𝜌∞ due to a ﬁnite 𝑇 , a safe play is to use a very large 𝑇 . Even though this practically works for classical adaptive minimum action method [57], our Remark 2 pointed out the fundamental diﬃculty of adaptive time-stepping strategy of repeated global redistribution of the temporal mesh and model reﬁtting of normalizing ﬂow. In fact, parameterizing the entire Wasserstein gradient ﬂow by physical time is suboptimal for describing its convergence to an equilibrium. The main purpose is to characterize the path by its geometry rather than by the speed at which it is traversed. By reparameterizing the trajectory with an intrinsic variable such as arclength one removes the explicit dependence on the physical time horizon and transforms the longhorizon relaxation problem into a ﬁnitelength path optimization problem on the Wasserstein manifold. In this section, we develop the geometric reformulation in the Wasserstein space, adapting the core principle of the geometric Minimum Action Method [22] originally formulated in Euclidean spaceto this inﬁnitedimensional Wasserstein space.

4.1. Reparameterization-Invariant Path Formulation of Geometric Action The basic idea is in the spirit of Maupertuis’s principle, which allows for possible variation in the ﬁnal time 𝑇 while keeping the beginning and end points ﬁxed, in contrast to Hamilton mechanics’s principle in Section 3 with ﬁxed initial state and ﬁnal time. We start with the establishment of a geometric action functional below. The optimal relation between the arc-length and the physical time as the result of the variation of the time interval 𝑇 is used in Section 4.4 to recover the physical time. Theorem 3 (Geometric reformulation of the Dawson-Gärtner action function). Assume 𝜌𝑎 , 𝜌𝑏 ∈ 2 (Ω)∩Dom( ). Let 𝐴𝐶𝜌𝑎 ,𝜌𝑏 ,𝑇 denote the set of absolutely continuous paths (with respect to Wasserstein metric) connecting two distributions from 𝜌𝑎 to 𝜌𝑏 over [0, 𝑇 ]. Under the regularity assumptions 2 speciﬁed in Appendix A.3, the variational problem associated with the time-dependent Dawson-Gärtner action 𝑆𝑇 [𝑝] in (10) admits the following geometric variational form: inf

inf

𝑇 >0 𝑝∈𝐴𝐶𝜌𝑎 ,𝜌𝑏 ,𝑇

𝑆𝑇 [𝑝] =

inf

𝑝∈𝐴𝐶𝜌𝑎 ,𝜌𝑏 ,1

̂ 𝑆[𝑝],

̂ is deﬁned on the interval 𝜏 ∈ [0, 1] by where the Geometric Action 𝑆[𝑝] ( ) ⟨ ⟩ ⎧ 1 ‖𝜕𝜏 𝑝𝜏 ‖−1,𝑝𝜏 ‖∇d  (𝑝𝜏 )‖−1,𝑝𝜏 + ∇d  (𝑝𝜏 ), 𝜕𝜏 𝑝𝜏 d𝜏, ⎪∫ −1,𝑝𝜏 0 ⎪ ̂ ∶= ⎨ 𝑆[𝑝] ⎪ ⎪ ⎩+∞, C. Liu and X. Zhou: Preprint submitted to Elsevier

(24)

if 𝑝𝜏 is absolutely continuous and the integral converges,

(25)

otherwise. Page 10 of 39

Generative Wasserstein Gradient Path Method

PROOF. See Appendix A.3. Here we slightly abuse the notation: (𝑝𝜏 )0≤𝜏≤1 and (𝑝𝑡 )0≤𝑡≤𝑇 (𝑇 could be inﬁnity) denote the same curve 𝑝, parametized by arclength 𝜏 and by physical time 𝑡, respectively. We highlight here that 𝑆̂ is invariant under any reparametrization. So 𝜏 here may refer to any curve parameter, not restricted to the arc-length parameter. The integrand in (25) is equivalent to ‖𝜕𝜏 𝑝𝜏 ‖−1,𝑝𝜏 ‖∇d  (𝑝𝜏 )‖−1,𝑝𝜏 (1 − cos 𝛼) where 𝛼 is the angle between the ̂ tangent and the negative Wasserstein gradient. When the action 𝑆[𝑝] is zero, this angle is exactly zero everywhere, indicating the path (𝑝𝜏 )0≤𝜏≤1 is indeed the Wasserstein gradient ﬂow. The second term in (25) is the endpoint contribution of the free energy along the path. When the terminal state is prescribed, namely 𝑝0 = 𝜌𝑎 and 𝑝1 = 𝜌𝑏 , the chain rule in Wasserstein space gives 1

∫0

⟨∇d  (𝑝𝜏 ), 𝜕𝜏 𝑝𝜏 ⟩−1,𝑝𝜏 d𝜏 =

d  (𝑝𝜏 ) d𝜏 =  (𝜌𝑏 ) −  (𝜌𝑎 ). ∫0 d𝜏

(26)

̂ is equivalent to minimizing only its ﬁrst term, namely Hence, for ﬁxed 𝑝0 and 𝑝1 , minimizing the geometric action 𝑆[𝑝] the Eulerian Geometric Action: Euler [𝑝] ∶=

‖ ‖ ‖𝜕 𝑝 ‖ d𝜏, ‖∇  (𝑝𝜏 )‖ ‖−1,𝑝𝜏 ‖ 𝜏 𝜏 ‖−1,𝑝𝜏 ∫0 ‖ d

(27)

which is certainly invariant under any reparameterization of the curve {𝑝𝜏 }. To implement this geometric principle by particle-based generative models, we utilize the same isometry discussed in Section 3.1 to minimize the equivalent Lagrangian Geometric Action: Lagrangian [Φ] ∶=

∫0

‖2 𝔼𝑧∼𝜌0 ‖ ‖[𝑝𝜏 ](Φ(𝜏, 𝑧))‖

)1∕2 (

‖2 𝔼𝑧∼𝜌0 ‖ ‖𝜕𝜏 Φ(𝜏, 𝑧)‖

)1∕2 d𝜏,

(28)

which expresses the same geometric cost as Euler [𝑝], but at the level of particles’ path and the velocity ﬁelds evaluated along them. In the present work of searching the geometric path connecting 𝑝0 = 𝜌𝑎 to an unknown equilibrium 𝜌𝑏 , the terminal state 𝑝𝜏=1 is optimized jointly with the path losses  by keeping the additional terminal free-energy penalty  (𝑝𝜏=1 ).

4.2. Discrete Geometric Optimization We approximate the continuous-integral action functional by discretizing 𝜏 ∈ [0, 1] into 𝐾 equal intervals. Let 𝑝 = {𝑝𝑘 }𝐾 be the sequence of densities (which are referred to as discrete “images” in the minimum action 𝑘=0 be the corresponding sequence of transport maps, with Φ0 = Id and 𝑝𝑘 = (Φ𝑘 )# 𝜌0 . A method[18]), and Φ = {Φ𝑘 }𝐾 𝑘=0 ‖ natural discretization of the Eulerian geometric action (27) approximates ‖ ‖𝜕𝜏 𝑝𝜏 ‖−1,𝑝𝜏 d𝜏 with the Wasserstein distance between two neighbors and leads to the following sum 𝐾 min Euler [𝑝] ∶= 𝑝

‖ ‖ ‖ ‖ + ‖∇  (𝑝𝑘 )‖ ‖∇d  (𝑝𝑘−1 )‖ ‖ ‖−1,𝑝𝑘−1 ‖ d ‖−1,𝑝𝑘 d (𝑝𝑘 , 𝑝𝑘−1 ) ⋅ , 2 𝑘=1 𝐾 ∑

s.t.

𝑝0 = 𝜌0 , 𝑝𝐾 = 𝜌1 . (29)

Likewise, we obtain the following equivalent discrete Lagrangian problem for Φ: 𝐾 min Lagrangian [Φ] ∶= Φ

𝐾 ∑ 𝑘=1

‖[𝑝𝑘−1 ]◦Φ𝑘−1 ‖ 2 ‖ ‖ ‖𝐿 (𝜌0 ) + ‖[𝑝𝑘 ]◦Φ𝑘 ‖𝐿2 (𝜌0 ) ‖Φ𝑘 − Φ𝑘−1 ‖ 2 ⋅ ‖ , ‖ ‖𝐿 (𝜌0 ) 2

s.t.

Φ0 = Id, (Φ𝐾 )# 𝜌0 = 𝜌1 . (30)

where 𝑝𝑘 = (Φ𝑘 )# 𝜌0 . The consistency between these two formulations is formally established below. Theorem 4 (Equivalence of Discrete Formulations). Let 𝜌0 and 𝜌1 be two given probability distributions. (a) If Φ∗ is a minimizer of the Lagrangian problem (30), then the induced density sequence 𝑝∗𝑘 = (Φ∗𝑘 )# 𝜌0 is a minimizer of the Eulerian problem (29). C. Liu and X. Zhou: Preprint submitted to Elsevier

Page 11 of 39

Generative Wasserstein Gradient Path Method ∗ (b) Conversely, if 𝑝∗ is a minimizer of (29), and Ψ𝑘 is the optimal transport map from 𝑝𝑘−1 to 𝑝∗𝑘 , then the composite map Φ∗𝑘 = Ψ𝑘 ◦ ⋯ ◦Ψ1 is a minimizer of (30).

PROOF. See Appendix A.4. We parameterize Φ𝑘 using Normalizing Flows and approximate the 𝐿2 (𝜌0 ) expectations in Eq. (30) via Monte Carlo sampling. We ﬁnally have the trainable Empirical Discrete Geometric Loss for a batch of data {𝑧𝑖 }𝑁 ∼ 𝜌0 as 𝑖=1 follows: 𝐾 𝐽̂𝑁 [Φ] ∶=

𝐾 ∑ 𝑘=1

𝑑𝑘 (Φ) ⋅

𝑣𝑘−1 (Φ) + 𝑣𝑘 (Φ) , 2

(31)

where 𝑑𝑘 (Φ) and 𝑣𝑘 (Φ) are the batch estimators: ( 𝑑𝑘 (Φ) ∶=

𝑁

1 ∑‖ 2 Φ𝑘 (𝑧𝑖 ) − Φ𝑘−1 (𝑧𝑖 )‖ ‖ ‖ 𝑁 𝑖=1

)1∕2

( ,

𝑣𝑘 (Φ) ∶=

𝑁

1 ∑‖ 2 𝑁 [𝑝𝑘 ](Φ𝑘 (𝑧𝑖 ))‖ ‖ ‖ 𝑁 𝑖=1

)1∕2 .

(32)

Theorem 5 (Consistency of the discrete geomtric objective). Let 𝐽̂[Φ] be the continuous geometric action (28), 𝐾 [Φ] be the discrete empirical objective (31). Under the regularity and moment assumptions stated in and let 𝐽̂𝑁 Appendix A.5, one has | 𝐾 | 𝔼|𝐽̂𝑁 (Φ) − 𝐽̂(Φ)| = (𝐾 −2 ) + (𝑁 −1∕2 ). | | In particular, the geometric discretization is second-order accurate in the number of path segments 𝐾 (equivalent to the number of layers in the neural networks), up to the Monte Carlo sampling error. PROOF. See Appendix A.5.

4.3. Arc-length Parametrization as Regularization in Training Algorithm However, direct optimization of the geometric objective can lead to a numerical artifact: degenerate parametrization, where many of the 𝐾 discrete images cluster in a small portion of the path. This occurs because the action 𝑆̂ is invariant under any parametrization, including numerically pathological ones. To obtain a stable and informative discretization, we impose the Wasserstein arc-length parametrization, which is associated with a constant-speed constraint, ‖𝜕𝜏 𝑝𝜏 ‖−1 ≡ const. In the discrete setting, this is enforced by a variance penalty on the segment lengths {𝑑𝑘 } computed via (32): arc [Φ] =

Var(𝑑1 , 𝑑2 , … , 𝑑𝐾 ) , Mean(𝑑1 , 𝑑2 , … , 𝑑𝐾 )

(33)

which is one of standard strategies in adaptive minimum action method [57, 21] to ensure the even distance {𝑑𝑘 } along the path . Because the equilibrium 𝜌𝑏 is unknown, we employ a penalized terminal cost for 𝑝𝐾 , as if our objective were solely to locate this equilibrium rather than the entire Wasserstein gradient ﬂow path. Speciﬁcally, the terminal density 𝑝𝐾 is treated as an optimization variable, and the penalty  (𝑝𝐾 ) is introduced [ ] 𝑈 (𝑝𝐾 (𝑥)) 1  (𝑝𝐾 ) = 𝔼𝑥∼𝑝𝐾 𝛽 −1 + 𝑉 (𝑥) + 𝔼𝑦∼𝑝𝐾 [𝑊 (𝑥 − 𝑦)] 𝑝𝐾 (𝑥) 2 (34) [ ] 1 −1 𝑈 (𝑝𝐾 (Φ𝐾 (𝑧))) ′ = 𝔼𝑧∼𝜌0 𝛽 + 𝑉 (Φ𝐾 (𝑧)) + 𝔼𝑧′ ∼𝜌0 [𝑊 (Φ𝐾 (𝑧) − Φ𝐾 (𝑧 ))] 𝑝𝐾 (Φ𝐾 (𝑧)) 2 to drive this endpoint toward a low-energy terminal state. The ﬁnal training objective then consists of the following three contributions : 𝐾 total = 𝐽̂𝑁 [Φ] + 𝛼term  (𝑝𝐾 ) + 𝛼arc arc [Φ],

C. Liu and X. Zhou: Preprint submitted to Elsevier

(35) Page 12 of 39

Generative Wasserstein Gradient Path Method

where 𝛼term and 𝛼arc are penalty parameters. The complete training procedure is summarized in Algorithm 2. Algorithm 2 GenWGP: Geometric Path Require: NF parameters {𝜃𝑘 }𝐾 , initial distribution 𝜌0 , steps 𝐾, weights 𝛼term , 𝛼arc . 𝑘=1 1: for each training iteration do 2: Sample a batch of 𝑁 particles {𝑧𝑖 }𝑁 ∼ 𝜌0 . 𝑖=1

Compute paths 𝑥(𝑖) = Φ𝑘 (𝑧𝑖 ) and densities 𝑝𝑘 (𝑥(𝑖) ) via Eq.(14). 𝑘 𝑘 Compute segment lengths 𝑑𝑘 and force magnitudes 𝑣𝑘 via Eq. (32). 𝐾 [Φ] via Eq. (31). 5: Evaluate Geometric Loss 𝐽̂𝑁 6: Evaluate Regularizers: terminal energy  (𝑝𝐾 ) via Eq. (34) and arc-length penalty arc [Φ] via Eq. (33). 7: Update parameters {𝜃𝑘 } via gradient descent on total (35). 8: end for 9: return Optimized NF parameters. 3: 4:

Remark 4 (Control of Lipschitz Regularity by transport cost). [26] indicates that minimizing the transport cost eﬀectively controls the Lipschitz constant of the learned ﬂow, which beneﬁts robustness and generalization. Invoking the triangle inequality, the divergence between two particles 𝑥, 𝑦 at the ﬁnal map Φ𝐾 is bounded by the accumulated transport cost: ‖ ‖ 𝐾 𝐾 ∑ ∑ ‖ ‖ ‖ ‖Φ𝐾 (𝑥) − Φ𝐾 (𝑦)‖ = ‖ (𝑥 − 𝑦) + (Φ (𝑥) − Φ (𝑥)) − (Φ (𝑦) − Φ (𝑦)) 𝑘 𝑘−1 𝑘 𝑘−1 ‖ ‖ ‖ ‖ 𝑘=1 𝑘=1 ‖ ‖ 𝐾 𝐾 ∑ ∑ ≤ ‖𝑥 − 𝑦‖ + ‖Φ𝑘 (𝑥) − Φ𝑘−1 (𝑥)‖ + ‖Φ𝑘 (𝑦) − Φ𝑘−1 (𝑦)‖, 𝑘=1

(36)

𝑘=1

where the summation terms represent the discrete path lengths. Our geometric action (30) is a weighted total arc-length length, where the weights encode the contributions of movements against the gradient ﬂow. So, the geometric action 𝐾 could be understood as a regularization, similar to the right-hand side of (36), to enhance the regularity of the 𝐽̂𝑁 trained diﬀeomorphic map.

4.4. Recovering Physical Time from the Geometric Path The geometric formulation in Section 4.1 produces a WGP 𝑝𝜏 parameterized by an intrinsic geometric variable 𝜏 ∈ [0, 1], implicitly performing the optimal adaptive discretization of the physical time march while hiding the temporal evolution. But from any initial 𝑝0 (which is not a stationary point of  ), we can indeed recover the original physical-time dynamics 𝑝𝑡 from our geometric path 𝑝𝜏 , at least up to the last second one, 𝑝𝐾−1 .

Time-Rescaling Relation. The zero-action trajectory follows the Wasserstein gradient ﬂow 𝜕𝑡 𝑝𝑡 = −∇d  (𝑝𝑡 ) in physical time 𝑡. Let 𝑡(𝜏) be the strictly increasing map from arc-length parameter 𝜏 to physical time 𝑡, then by the chain rule, we obtain the the geometric “velocity” as 𝜕 𝜏 𝑝𝜏 =

d𝑡 𝑑𝑡 𝜕 𝑝 = − ∇d  (𝑝𝜏 ). d𝜏 𝑡 𝑡 d𝜏

(37)

Taking the Wasserstein tangent norm ‖ ⋅ ‖−1,𝑝𝜏 on both sides, we obtain the scalar diﬀerential equation governing the time mapping: ‖𝜕𝜏 𝑝𝜏 ‖−1,𝑝𝜏 =

d𝑡 ‖∇  (𝑝𝜏 )‖−1,𝑝𝜏 . d𝜏 d

(38)

A key feature of the geometric training (Algorithm 2) is the regularization of the arc-length speed, so the path 𝑝𝜏 satisﬁes the condition of the arc-length parametrization ‖𝜕𝜏 𝑝𝜏 ‖−1,𝑝𝜏 ≈ 𝑐 for some constant (i.e., total length) 𝑐 > 0. Substituting this into (38) allows us to solve for the time scaling factor: d𝑡 𝑐 𝑐 = = . d𝜏 ‖∇d  (𝑝𝜏 )‖−1,𝑝𝜏 ‖[𝑝𝜏 ]‖𝑝𝜏 C. Liu and X. Zhou: Preprint submitted to Elsevier

(39) Page 13 of 39

Generative Wasserstein Gradient Path Method

Equation (39) oﬀers a clear physical interpretation: the physical time lapse d𝑡 required to traverse a ﬁxed geometric distance d𝜏 is inversely proportional to the magnitude of the driving force. Crucially, near equilibrium or metastable states where ‖‖ ≪ 1, the derivative d𝑡∕ d𝜏 naturally becomes large. This allows the method to capture the “slow tail” of the relaxation process accurately without the computational burden of inﬁnitesimal time-stepping required by Eulerian solvers.

Determination of the Time Constant. The constant 𝑐 represents the total path length in the Wasserstein metric and ﬁxes the global time scale. It is determined from the free-energy dissipation identity along the geometrically parameterized path. By (37), the rate of free energy dissipation along the geometric path is: d𝑡 𝑑 = ⟨∇d  [𝑝𝜏 ], 𝜕𝜏 𝑝𝜏 ⟩−1,𝑝𝜏 = − ‖∇d  [𝑝𝜏 ]‖2−1,𝑝 = −𝑐‖[𝑝𝜏 ]‖𝑝𝜏 . 𝜏 d𝜏 d𝜏 Integrating both sides over 𝜏 ∈ [0, 1] yields the formula for 𝑐: 𝑐=

 (𝑝𝜏=0 ) −  (𝑝𝜏=1 ) 1

∫0 ‖[𝑝𝜏 ]‖𝑝𝜏 d𝜏

(40)

(41)

This ensures that the reconstructed time evolution exactly matches the total free energy diﬀerence speciﬁed by the boundary conditions.

Numerical Reconstruction. Given the discrete sequence of transport maps {Φ𝑘 }𝐾 produced by the Normalizing 𝑘=0

Flow, we estimate the velocity magnitudes 𝑣𝑘 ≈ ‖[𝑝𝑘 ]‖ using the batch estimator deﬁned in Eq. (32). We approximate the integrals using the trapezoidal rule. First, the constant 𝑐 is estimated by (41)  (𝑝 ) −  (𝑝 ) 𝑐 ≈ ∑𝐾 0 𝑣 +𝑣 𝐾 , 𝑘−1 𝑘 Δ𝜏 𝑘=1 2

(42)

where Δ𝜏 = 1∕𝐾. Subsequently, the physical time increments Δ𝑡𝑘 = 𝑡𝑘 − 𝑡𝑘−1 are recovered by the mid-point scheme: 𝑘Δ𝜏

𝑐Δ𝜏 𝑐 𝑑𝜏 ≈ Δ𝑡𝑘 ≈ ∫(𝑘−1)Δ𝜏 ‖[𝑝𝜏 ]‖𝑝 2 𝜏

(

1 + 𝑣𝑘−1 𝑣𝑘

) .

(43)

This procedure is summarized in Algorithm 3. This reconstruction of the physical time is quite accurate up the last second distribution 𝑝𝐾−1 on the path; the terminal distribution is the equilibrium state, taking inﬁnitely long time to reach in theory. Algorithm 3 Recover Physical Time from Geometric Path Require: Trained NF parameters {𝜃𝑘 }𝐾 , initial distribution 𝜌0 . 𝑘=0 1: Compute boundary energies 0 =  ((Φ0 )# 𝜌0 ) and 𝐾 =  ((Φ𝐾 )# 𝜌0 ). 2: Estimate velocity norms 𝑣𝑘 for 𝑘 = 0, … , 𝐾 using batch samples via Eq. (32). 3: Compute path length constant 𝑐 via discrete approximation of Eq. (42). 4: Initialize 𝑡0 ← 0. 5: for 𝑘 = 1 to 𝐾 do 6: Compute time step Δ𝑡𝑘 ← 𝑐Δ𝜏 (𝑣−1 + 𝑣−1 ). 𝑘 𝑘−1 2 7: Update physical time 𝑡𝑘 ← 𝑡𝑘−1 + Δ𝑡𝑘 . 8: end for 9: return Physical timestamps {𝑡𝑘 }𝐾 . 𝑘=0 Remark 5. Our numerical recover of the physical time using Algorithm 3 also also yields an approximate terminal time 𝑡𝐾 , even though the theoretical time required to reach equilibrium is inﬁnite. Nevertheless, this numerical 𝑡𝐾 is practically meaningful: it indicates that the time interval [0, 𝑡𝐾 ] is suﬃciently long for the gradient ﬂow to approach equilibrium and for the free energy to converge close to its minimal value. Consequently, the setup of the terminal 𝑇 ≈ 𝑡𝐾 (or between 𝑡𝐾−1 and 𝑡𝐾 ) - together with the entire recovered time mesh (𝑡𝑘 )0≤𝑘≤𝐾−1 ) - can be directly used in the physical-time path optimization Algorithm 1 as an optimal adaptive time mesh. This allows reﬁnement of the path within an practically optimal interval without requiring any change to the network architecture. C. Liu and X. Zhou: Preprint submitted to Elsevier

Page 14 of 39

Generative Wasserstein Gradient Path Method

5. Numerical examples Our numerical examples focus on the validation and application of the geometric GenWGP approach (Algorithm 2), together with its time-recovery postprocessing (Algorithm 3). Unlike time-marching methods that operate on a prescribed ﬁnite horizon 𝑇 , our goal is to approximate the full Wasserstein gradient ﬂow toward equilibrium and to assess the learned path not only at the terminal state but also along the evolution in physical time. Section 5 is organized to validate one central numerical claim: the geometric GenWGP formulation provides a more eﬀective representation of long-time relaxation than uniform physical-time discretization, while retaining accurate recovered dynamics on the transient regime. We begin with analytically tractable FokkerPlanck examples, where exact solutions allow direct veriﬁcation of both recovered trajectories and terminal states. We then perform matched comparisons with the physical-time formulation under identical architectures and training setups, so that the eﬀect of geometric parametrization and time recovery can be isolated cleanly. Finally, for non-convex and interacting-particle systems where full transient references are unavailable or only partially reliable, we use partial-reference comparisons and structure-preserving diagnostics to test whether the learned path remains dynamically meaningful. Unless otherwise speciﬁed, training samples are drawn from the standard Gaussian base distribution  (𝑥; 0, 𝐼𝑑 ) with 𝑁 particles. Models are trained using Adam with exponential learning-rate decay. Our Python implementation is available at GitHub.

5.1. Diﬀusion Process: The Fokker-Planck Equation We begin with entropy-driven dynamics associated with the free energy  (𝜌) = ∫ 𝜌(𝑥) log 𝜌(𝑥) d𝑥+∫ 𝑉 (𝑥)𝜌(𝑥) d𝑥, which corresponds to the Fokker-Planck equation at 𝛽 = 1. In this subsection, the availability of exact reference solutions allows for rigorous direct quantitative validation of both the recovered physical-time dynamics and the terminal equilibrium state.

5.1.1. Quadratic Potentials We consider convex quadratic potentials 𝑉 (𝑥) = 12 (𝑥 − 𝜇)⊤ Σ−1 (𝑥 − 𝜇). The corresponding WGF is the Ornstein√ Uhlenbeck dynamics d𝑋𝑡 = −Σ−1 (𝑋𝑡 − 𝜇)d𝑡 + 2d𝐵𝑡 , whose law is  (𝜇(𝑡), Σ(𝑡)) given by ( −1 −1 −1 −1 ) 𝜇(𝑡) = 𝜇 − (𝜇 − 𝜇(0))𝑒−Σ 𝑡 , Σ(𝑡) = 𝑒−Σ 𝑡 Σ(0)𝑒−Σ 𝑡 + Σ 𝐼 − 𝑒−2Σ 𝑡 . and 𝜇, Σ are the equilibrium mean and covariance, respectively. These examples provide the cleanest setting for quantitative validation, since both the trajectory and the equilibrium state are explicitly known. We test the method on three cases: (i) a 2D isotropic potential, (ii) a 2D anisotropic potential, and (iii) a 10D block-structured potential. We use 𝑁 = 5000 particles. The transport map is parameterized by a RealNVP normalizing ﬂow with 𝐾 = 9 aﬃne coupling “layers”. The coupling sub-network for each layer is a four-layer MLP of width 128, with LeakyReLU activations and a ﬁnal tanh scale head. Training uses 1000 epochs, an initial learning rate 8 × 10−4 , and exponential decay factor 𝛾 = 0.9999.

2D isotropic diﬀusion. With the target 𝜌∞ =  (𝑥; [3, 3]⊤ , 0.25𝐼2 ) and base 𝜌0 =  (0, 𝐼2 ), the exact solution

remains Gaussian with 𝜇(𝑡) = 𝜇(1 − 𝑒−4𝑡 ) and Σ(𝑡) = 0.25𝐼2 + 0.75𝑒−8𝑡 𝐼2 . We ﬁrst show the results of GenWGP from Algorithm 2. Fig. 1 illustrates the transport map learned by Algorithm 2 with 𝐾 = 9 stacked layers, which produces a smooth contraction ﬂow toward equilibrium. Fig. 2 validates the learned curve is indeed nearly an arc-length parametrized Wasserstein gradient ﬂow: the panel (a) shows that the Wasserstein distance between each neighboring layers (segment length) is well maintained close to constant, indicating a good quality of arc-length parametrization; the panel (b) conﬁrms that the cosine alignment between 𝜕𝜏 𝑝𝜏 and −∇d  (𝑝𝜏 ) remains close to one, verifying the key parallel condition for the gradient ﬂow: 𝜕𝜏 𝑝𝜏 ∝ −∇d  (𝑝𝜏 ). These two quantities are numerically computed as in ‖Φ𝑘+1 −Φ𝑘 ‖𝐿2 (𝜌 ) ‖[𝑝𝑘+1 ]◦Φ𝑘+1 ‖𝐿2 (𝜌0 ) +‖[𝑝𝑘 ]◦Φ𝑘 ‖𝐿2 (𝜌0 ) 0 . and −∇d  (𝑝𝜏𝑘 ) ≈ Eqn. (30) as follows 𝜕𝜏 𝑝𝜏𝑘 ≈ 𝜏𝑘+1 −𝜏𝑘 2 To compare with the true gradient ﬂow, we recover the physical times 𝑡0 , 𝑡1 , … , 𝑡𝐾−1 using Algorithm 3 from the geometric path 𝑝𝜏𝑘 (0 ≤ 𝑘 ≤ 𝐾), thereby aligning the numerical solution with the true solution as a function of time. Fig. 2(c) shows the contours of the density snapshots at two selected physical times, comparing the numerical and true solutions. Excellent agreement is observed, demonstrating that our geometric GenWGP accurately captures the dynamics even when evaluated in terms of physical time. C. Liu and X. Zhou: Preprint submitted to Elsevier

Page 15 of 39

Generative Wasserstein Gradient Path Method

4 4

0.4800 0.4795 0.4790 0.4785 0.4780 0.4775 0.4770 0.4765 0.4760

1.0 0.8 Cosine alignment

Arc length

Figure 1: Transport map learned by Algorithm 2 with 𝐾 = 9 layers for the 2D isotropic Gaussian case. The ﬂow map (indicated by the arrows) transports the initial density toward equilibrium while preserving isotropic structure.

0.6 0.4 0.2

4 5 Layer

0.0

4 5 Layer

(a) Arc-length (segment norm) between each pair of neigh- (b) Cosine alignment between tangent and negative gradient bouring layers at each layer Exact contour

Layer 0, t=0.000

Learned contour

Learned samples

Layer 4, t=0.144

Layer 9, t=2.099

0.014

0.055

0.067

0.114

0.028

0.198

0.00

0.09 5

0.094

2 4 4

Figure 2: Validation of the learned geometric path for the 2D isotropic Gaussian example. (a): nearly constant segment lengths indicate approximate arc-length parametrization; (b): cosine alignment close to one is consistent with the gradientﬂow direction; and (c) the density snapshots at two selected times agree well with the exact solution at the recovered physical times.

C. Liu and X. Zhou: Preprint submitted to Elsevier

Page 16 of 39

Generative Wasserstein Gradient Path Method

We also compare the geometric formulation (Algorithm 2) with its physical-time counterpart (Algorithm 1) (using 𝑇 = 1) under the same network architecture (𝐾 = 9) and training setup, demonstrating the consistency of the two numerical gradient paths while highlighting their distinct characteristics. Fig. 3a presents that the physical-time method uses a uniform discretization on [0, 𝑇 ] with 𝑇 = 1 which is suﬃciently large here to approach the equilibrium, whereas the geometric method recovers a non-uniform time mesh but adopts uniform in Wasserstein arc-length. The results in the panel (b)(c) in Fig. 3 are the decay of free energy  (𝑝) in terms of 𝑡 and 𝜏 respectively, for the numerical paths from these two methods and the truth WGF. In particular, more images are placed in the early stage where the free energy decays rapidly, leading to a more balanced distribution of resolution along the relaxation path. The gap of free energy between two neighboring discrete layers is more even in the geometric approach than the uniform physical-time approach. The accuracy of the two methods measured in d error is validated by Fig. 3c.

1.0 0.5 0.0

4 5 Layer

Exact Physical-time geometric + recovery

0.035

0.030

20 15 10 0

0.5

1.0 1.5 Physical time t

2.0

(a) Recovered physical time vs. (b) Free energy vs. physical time layer

Physical-time geometric + recovery

0.025 0.020 0.015

5 0.0

0.040

Physical-time geometric + recovery

d -distance

Free energy

Physical time tk

1.5

35 30 25 20 15 10 5 0

Free energy

Physical-time geometric + recovery T=1

2.0

4 5 Layer

0.010

0.0

0.5

1.0 1.5 Physical time t

2.0

(d) d error vs. physical time

Figure 3: Comparison between the physicaltime formulation and the geometric formulation for the 2D isotropic Gaussian example. (a): the uniform time mesh for each layer in Algorithm 1 and the recovered physical time mesh in Algorithm 2; (b): the decay of free energy plot in physical time; (c) the decay of free energy plot in layers; (d) the d errors.

2D anisotropic diﬀusion. For 𝜇 = [3, 3]⊤ and Σ = diag(1, 0.25), the exact solution remains Gaussian with 𝜇(𝑡) = [3(1 − 𝑒−𝑡 ), 3(1 − 𝑒−4𝑡 )]⊤ ,

Σ(𝑡) = diag(1, 0.25 + 0.75𝑒−8𝑡 ).

Compared with the isotropic case, this example is a bit more challenging because the two coordinate variables evolve on distinct time scales to take longer time to reach equilibrium. As in the isotropic case, we compare in Fig. 4 the geometric formulation (blue curve) against the matched physical-time approach (red curve) with 𝑇 = 1 under the same architecture and training setup. The recovered physical time from the geometric path is now much longer than 𝑇 = 1. Consequently, the panel (c) shows that the geometric path achieves a lower free energy value. The comparison of free energy decays up to 𝑇 = 1 conﬁrms that the geometric path resolves both the fast initial transient and the slower remaining relaxation. The d error curves in the panel (d) show that, on the interval [0, 1], the recovered geometric method has slightly less accurate than the physical-time path method, owing to the fewer discrete points available on the geometric path within [0, 1]. This accuracy can be straightforwardly improved, as discussed in Remark 5. Finally, the density snapshots in the panel (e) further conﬁrm the accuracy of the geometric WGF path when benchmarked against the true solution.

10D diﬀusion. Let 𝜇 = (1, 1, 0, 0, 1, 2, 0, 0, 2, 3)⊤ and Σ = diag(Σ𝐴 , 𝐼2 , Σ𝐵 , 𝐼2 , Σ𝐶 ), where [

Σ𝐴 =

] 5∕8 −3∕8 , −3∕8 5∕8

Σ𝐵 =

[ 1 0

] 0 , 0.25

Σ𝐶 = 0.25𝐼2 .

With the initial distribution 𝜌0 =  (0, 𝐼10 ), the exact solution is  (𝜇(𝑡), Σ(𝑡)), where 𝜇(𝑡) =(1 − 𝑒−𝑡 , 1 − 𝑒−𝑡 , 0, 0, 1 − 𝑒−𝑡 , 2(1 − 𝑒−4𝑡 ), 0, 0, 2(1 − 𝑒−4𝑡 ), 3(1 − 𝑒−4𝑡 ))⊤ , Σ(𝑡) =diag(Σ𝐴 (𝑡), 𝐼2 , Σ𝐵 (𝑡), 𝐼2 , Σ𝐶 (𝑡)), ] [ 3+3𝑒−4𝑡 5+3𝑒−4𝑡 − 8 8 , with Σ𝐴 (𝑡) = −4𝑡 5+3𝑒−4𝑡 − 3+3𝑒 8 8 C. Liu and X. Zhou: Preprint submitted to Elsevier

[ Σ𝐵 (𝑡) =

]

1 1+3𝑒−8𝑡 4

[ ,

Σ𝐶 (𝑡) =

1+3𝑒−8𝑡 4

] 1+3𝑒−8𝑡 4

Page 17 of 39

Generative Wasserstein Gradient Path Method

3 2

15 10

4 5 Layer

0 0

Physical-time geometric + recovery

d -distance

Exact Physical-time geometric + recovery

20 Free energy

Physical time tk

Physical-time geometric + recovery T=1

Free energy

3 4 Physical time t

(a) Recovered physical time vs. (b) Free energy vs. physical time layer Exact contour

Layer 0, t=0.000

4 5 Layer

Learned contour

Physical-time geometric + recovery

2 3 4 Physical time t

(d) d error vs. physical time

Learned samples

Layer 4, t=0.266

0.050 0.045 0.040 0.035 0.030 0.025 0.020 0.015 0.010

Layer 9, t=5.937

0.10

0.05

0.050

0.1 0.010

0.08 9

0.023

0.008

0.0

2 4

(e) Density snapshots

Figure 4: Comparison in the anisotropic Gaussian case.

This example examines whether the method remains accurate in a moderate-dimensional setting with coupled and anisotropic substructures. The learned ﬂow from Algorithm 2 captures the expected rotated and anisotropic components in the selected two-dimensional projections; see Fig. 5. For accuracy, Fig. 6 reports the errors in the mean and covariance against the exact Gaussian solution over time. The errors remain small throughout the evolution, indicating that the geometric formulation retains good accuracy in this higher-dimensional but still exactly solvable setting. Projection on 0-1 plane

Projection on 4-5 plane

0 2

5.0

2.5 0.0 2.5 Dimension 0

5.0

Dimension 9

4 Dimension 5

Dimension 1

Projection on 8-9 plane

0 2 4

5.0

2.5 0.0 2.5 Dimension 4

5.0

2.5 0.0 2.5 Dimension 8

5.0

Figure 5: 2D projections of the terminal distribution for the 10D (dimension 0 to 9) block-structured Gaussian example.

5.1.2. Non-Convex Potential: 10D Styblinski-Tang Potential We next consider the 10D Styblinski-Tang potential, given by a sum of identical one-dimensional potentials over each coordinate: ( 𝑑 ) 3 ∑ 4 𝑉 (𝑥) = 𝑥 − 16𝑥2𝑖 + 5𝑥𝑖 , 𝑥 = (𝑥1 , … , 𝑥10 ) ∈ ℝ10 . 50 𝑖=1 𝑖 The initial is the standard Gaussian measure. Because of permutation symmetry, all one-dimensional marginals 𝑝(𝑖) 𝜏 (𝑥𝑖 ) (1) are statistically identical, thus 𝑝𝜏 (𝑥) = Π𝑖 𝑝𝜏 (𝑥𝑖 ). C. Liu and X. Zhou: Preprint submitted to Elsevier

Page 18 of 39

Generative Wasserstein Gradient Path Method 0.6

Mean Frobenius Norm Error Covariance Frobenius Norm Error

0.5

Error

0.4

0.3

0.2

0.1

0.0

0.00

0.25

0.50

0.75

1.00

1.25

1.50

1.75

2.00

Figure 6: Absolute error of the mean and Frobenius norm of the covariance error vs. recovered time in the 10D blockstructured Gaussian example.

Parameterization and visualization. We parameterize the path by a non-uniform B-spline Flow [24] with two hidden layers (width 100, SiLU activation), which provides smooth 𝐶 2 -diﬀeomorphic transports with controlled regularity. Fig. 7 shows a representative two-dimensional projection (𝑥5 , 𝑥6 ) of the particle evolution along the path, illustrating the transition from a unimodal Gaussian to a complex multimodal distribution. Layer 1

Layer 3

Layer 4

0 2

0 2 4

0 2 5-th component

0 2 4

Layer 5

0 2 5-th component

6-th component

0 2 4

Layer 6

0 2 5-th component

0 2 4

Layer 7

0 2 5-th component

0 2 4

0 2 5-th component

0 2 4

0 2 5-th component

6-th component

0 2 4

0 2 5-th component

Layer 9

Layer 8

6-th component

Layer 2

6-th component

Layer 0 4

0 2 4

0 2 5-th component

Figure 7: Sample points projected onto the (𝑥5 , 𝑥6 )-plane at each layer along the learned geometric path for the 10D Styblinski-Tang potential.

Reference solution and comparison. Each one-dimensional marginal evolves independently according to the onedimensional Fokker-Planck equation, equivalently the over-damped Langevin SDE d𝑋𝑡 = −𝑉1′ (𝑋𝑡 )d𝑡 +

√

2d𝐵𝑡 ,

𝑉1 (𝑋) =

) 3 ( 4 𝑋 − 16𝑋 2 + 5𝑋 , 50

𝑋𝑡 ∈ ℝ.

Unlike the OrnsteinUhlenbeck process, this SDE has no analytical expression of the density evolution. We therefore simulate 5000 Euler-Maruyama trajectories with time step 10−3 , and compare the resulting empirical marginal density with the learned one-dimensional marginals at matched physical times. Fig. 8 shows close agreement, including pronounced non-Gaussian and multimodal features. This provides quantitative evidence that the learned transport C. Liu and X. Zhou: Preprint submitted to Elsevier

Page 19 of 39

Generative Wasserstein Gradient Path Method

geometry remains accurate in a high-dimensional nonconvex setting, at least at the marginal level made accessible by the separable structure. Layer 0 | t = 0.0000

1.0

Layer 1 | t = 0.1054

1.0

Layer 2 | t = 0.2188

1.0

Layer 3 | t = 0.3435

1.0

0.8

0.6

0.4

0.2

0.0

1.0

2 2 4 Layer 5 | t0 = 0.6530

1.0

2 2 4 Layer 6 | t0 = 0.8616

1.0

2 2 4 Layer 7 | t0 = 1.1424

1.0

2 2 4 Layer 8 | t0 = 1.5713

0.0

0.8

0.6

0.4

0.2

0.0

1.0

0.8

Layer 4 | t = 0.4853

1.0

0.0

2 2 4 Layer 9 | t0 = 2.3374

Figure 8: Comparison of marginal densities from the learned path (colored for each component) and from 1D SDE simulation (black) at each layer, with the recovered physical times indicated.

5.2. Interacting Particle Dynamics We next consider WGFs driven by nonlocal interaction energies with the pairwise term 12 ∫ℝ𝑑 ×ℝ𝑑 𝑊 (𝑥 − 𝑦)𝜌(𝑥)𝜌(𝑦)d𝑥d𝑦. Such models arise in aggregation, swarming, and mean-ﬁeld dynamics. In contrast to the FokkerPlanck examples above, these systems typically do not admit explicit transient solutions and may even possess compactly supported equilibrium states. Our validation therefore focuses on problem-adapted quantities: exact steadystate information whenever available, recovered physical-time comparisons under matched training setups, and structural diagnostics that test whether the learned path remains consistent with Wasserstein gradient-ﬂow behavior.

5.2.1. Pure Aggregation The ﬁrst example is a two-dimensional pure aggregation model where  (𝜌) =

1 𝑊 (𝑥 − 𝑦)𝜌(𝑥)𝜌(𝑦)d𝑥d𝑦, 2 ∫ℝ2 ×ℝ2

𝑊 (𝑥) =

1 ‖𝑥‖2 − log ‖𝑥‖. 2

This kernel balances quadratic attraction with Newtonian repulsion, and its steady state is the uniform distribution on the unit disk. We initialize from  (0, 0.25𝐼2 ) and use a seven-layer non-uniform B-spline Flow with 𝑁 = 5000 particles. Fig. 9 displays the particle transport and density evolution along the learned path. The solution becomes radially symmetric, develops compact support, and approaches the analytical minimizer at terminal time. There is no closed-form transient solution for this example. For quantitative validation, we therefore compute an independent numerical reference using the primal-dual scheme developed in [9] which is grid-based and of timemarching type. After reparameterizing the learned geometric path to physical time via Algorithm 3, we compare the resulting densities with the primal-dual solution at the corresponding discrete times. Fig. 10 shows that the relative 𝐿2 diﬀerence between the two solutions remains uniformly small throughout the evolution. This example therefore provides a direct quantitative comparison against an existing numerical solver for a nontrivial interacting-particle dynamics.

5.2.2. Aggregation-Drift Equation We next consider an aggregation model with an additional singular conﬁning potential 𝑉 , where the free energy is  (𝜌) =

𝑉 (𝑥)𝜌(𝑥)d𝑥 +

1 𝑊 (𝑥 − 𝑦)𝜌(𝑥)𝜌(𝑦)d𝑥d𝑦, 2 ∬ℝ2 ×ℝ2

1 ‖𝑥‖2 − log ‖𝑥‖, 2

𝛼 𝑉 (𝑥) = − 1 log ‖𝑥‖. 𝛼2

∫ℝ2

with 𝑊 (𝑥) =

C. Liu and X. Zhou: Preprint submitted to Elsevier

Page 20 of 39

Generative Wasserstein Gradient Path Method Layer 0

1.5

Layer 1

1.5

Layer 2

1.5

1.0

0.5

0.0

0.5

1.0

1.5

1.0

0.5

0.0

0.5

1.0

1.5

Layer 4

1.5

1.0

0.5

0.0

0.5

1.0

1.5

Layer 5

1.5

1.0

0.5

0.0

0.5

1.0

1.5

Layer 6

1.5

1.0

0.5

0.0

0.5

1.0

1.5

1.0

0.5

0.0

0.5

1.0

Layer 0 | t = 0.0000

1.5 1.0 0.5

1.5

2.0

1.0

1.0 1.5

1.0

0.5

0.0

0.5

1.0

1.5

0.5

1.0

0.0

1.5

0.5

1.0

0.4

0.5

0.0

0.3

0.0

0.5

0.2

0.5

1.0

0.1

1.0

0.0

1.5

1.0

0.5

0.0

0.5

1.5

1.0

1.5

0.0

0.5

1.0

1.5

0.5

0.0

0.5

1.0

1.5

1.4 1.2 1.0 0.8 0.6 0.4 0.2 0.0

1.0

0.5

0.0

0.5

1.5

1.0

1.5

0.0

0.5

1.0

1.5 1.0 0.8

0.5

0.6

0.0

1.5

0.5 1.0

0.0

1.5

0.4

1.0

0.3

0.5

0.0

0.5

1.0

1.5

Layer 6 | t = 0.8394

0.0 0.5

0.1

1.0

0.0

1.5

1.0

0.5

0.0

0.5

1.5

1.0

1.5

1.0

1.5

0.5

0.0

0.5

1.0

1.5

Layer 3 | t = 0.2037

1.5

1.0

0.5

0.0

0.5

1.0

1.5

1.0 0.5 0.0 0.5 1.0 1.5

0.7 0.6 0.5 0.4 0.3 0.2 0.1 0.0

Layer 7 | t = 1.6598

1.5 0.35 0.30 0.25 0.20 0.15 0.10 0.05 0.00

0.5

0.0

0.2 1.0

0.0

Layer 7

0.5

1.0 1.5

0.5

1.0

0.4

1.5

1.0

1.5

0.5

1.5

1.0

0.5

Layer 2 | t = 0.1160

1.0

0.2

1.5

1.5 1.5

Layer 5 | t = 0.5070

1.5

0.5

1.5

0.5

Layer 1 | t = 0.0502

0.5

1.0

1.5

1.0

0.0

Layer 4 | t = 0.3250

1.5

0.5

1.0

0.5

1.5

2.5

1.5

0.0

1.5

1.0

1.5

Layer 3

1.5

1.0

0.5

0.0

0.5

1.0

1.5

0.35 0.30 0.25 0.20 0.15 0.10 0.05 0.00

Figure 9: Pure aggregation: particle transport (top) and density evolution (bottom) along the learned geometric path. The terminal state approaches the uniform distribution supported on the unit disk.

Relative L2 Error

0.08 0.06 0.04 0.02 0.00 0

Layer

Figure 10: Relative 𝐿2 diﬀerence between the learned densities and the reference solution computed by the primal-dual method.

√𝛼 1 The explicit steady state [5] is the uniform distribution on the annulus with inner and outer radii 𝑅𝑖 = , 𝑅𝑜 = 𝛼2 √ 𝑅2𝑖 + 1. Starting from a non-radially symmetric initial datum composed of ﬁve Gaussians, the learned geometric path restores radial symmetry and converges to the annular steady state; see Fig. 11.

C. Liu and X. Zhou: Preprint submitted to Elsevier

Page 21 of 39

Generative Wasserstein Gradient Path Method Layer 0

Layer 1

Layer 2

Layer 3

Layer 4

1.5

1.0

0.5

0.0

0.5

1.0

1.5

1.0

0.5 0.0

0.5

1.0

1.5

1.0

0.5 0.0

Layer 5

0.5

1.0

1.5

1.0

Layer 6

0.5 0.0

0.5

1.0

1.5

1.5 1.5

1.0

0.5 0.0

Layer 7

0.5

1.0

1.5

1.0

0.5

0.0

0.5

1.0

1.5

1.0

0.5 0.0

0.5

1.0

1.5

1.0

0.5 0.0

0.5

1.0

1.5

1.0

0.5 0.0

0.5

1.0

1.5

0.5

1.0

1.5

0.5

1.0

1.5

Layer 9

1.5

1.0

Layer 8

1.5 1.5

1.0

0.5 0.0

0.5

1.0

1.5

1.0

0.5 0.0

Figure 11: Aggregation-drift: particle evolution from an asymmetric initial distribution toward the annular steady state.

We use this example to show how the geometric path obtained by Algorithm 2 can provide a better non-uniform time mesh as well as a good terminal time for the physical-time path minimization algorithm 1, as discussed in Remark 5. More precisely, we ﬁrst compute a geometric path, then recover the corresponding physical times by Algorithm 3, and ﬁnally use this recovered non-uniform time mesh back in Algorithm 1. By comparing the numerical optimal paths under the uniform time mesh grid and this recovered non-uniform time mesh grid, we show the improvement of the accuracy Fig. 12(b), where Fig. 12(a) presents the diﬀerence of these two time meshes in terms of each layer. Fig. 12(b) presents the cumulative MAM loss along the discrete path. The recovered-time discretization yields a uniformly smaller cumulative loss, indicating that the recovered mesh provides a more eﬀective physical-time representation of the same relaxation process. 0.025

Uniform time Recovered time

2.00 1.75

0.020 Cumulative MAM loss

1.50 Physical time

Uniform time Recovered time

1.25 1.00 0.75 0.50

0.015 0.010 0.005

0.25 0.00

0.000

4 Layer

(a) Physical time vs. layer

4 Layer

(b) Cumulative MAM loss vs. layer

Figure 12: Comparison between two physical-time discretizations for the aggregation-drift equation: the standard uniformtime mesh and the recovered time mesh obtained from the geometric path. The recovered time allocates more layers to the fast transient regime and produces a smaller cumulative MAM loss.

To further examine the structure of the two learned paths, we report in Fig. 13 three intrinsic diagnostics: the free-energy proﬁle, the cosine alignment between the discrete path velocity and the negative Wasserstein gradient, and the Wasserstein distance between neighboring layers. The recovered-time path oﬀers a more uniform drop of the free energy along the discrete path, a cosine value closer to one, and a more balanced distribution of inter-layer distances. This example illustrates the recommendation of Remark 5 where the geometric method serves as an eﬀective preprocessing step to further ﬁne tune the physical-time discretization for more accurate path.

C. Liu and X. Zhou: Preprint submitted to Elsevier

Page 22 of 39

Generative Wasserstein Gradient Path Method 1.00

Uniform time Recovered time

Cosine alignment

Free energy

0.65 0.60 0.55 0.50 0.45

0.150

0.90

0.125

0.85 0.80 Uniform time Recovered time

0.70 2

4 Layer

(a) Free energy vs. layer

0.100 0.075 0.050

0.75

Uniform time Recovered time

0.175

0.95

Arc length

0.70

0.025

4 Layer

(b) Cosine alignment vs. layer

4 Layer

Figure 13: Structural diagnostics for the aggregation-drift equation under uniform time and recovered time discretizations. From left to right: free energy, cosine alignment with −∇d  , and the Wasserstein distance between neighboring layers.

5.2.3. Aggregation-Diﬀusion Equation We conclude this section with an aggregation-diﬀusion model that combines nonlinear diﬀusion and nonlocal attraction:  (𝜌) =

𝜈 1 𝑊 (𝑥 − 𝑦)𝜌(𝑥)𝜌(𝑦)d𝑥d𝑦, 𝜌𝑚 (𝑥)d𝑥 + ∫ℝ2 𝑚 − 1 2 ∬ℝ2 ×ℝ2

where the interaction kernel is 2 1 𝑊 (𝑥) = − 𝑒−|𝑥| . 𝜋

This smooth radially symmetric kernel induces short-range attraction and is frequently used in models of biological aggregation and collective behavior. The corresponding evolution equation is 𝜕𝑡 𝜌 = ∇ ⋅ (𝜌∇𝑊 ∗ 𝜌) + 𝜈Δ𝜌𝑚 . We choose 𝑚 = 2. To probe nontrivial intermediate dynamics toward the equilibirum, we take as initial datum the characteristic function on [−3, 3]2 thus the constant mass is 9. These settings produce a compactly supported constantdensity proﬁle that undergoes highly nontrivial transient dynamics before converging to equilibrium. We employ an equal arc-action penalty (i.e.,  (𝑝𝑘+1 )− (𝑝𝑘 ) is constant) instead of the arc-length parametrization, so that distributed discrete images (layers) along the path oﬀer better quality with only 11 layers. Fig. 14 successfully discovers the complex evolution of this system: the initial density ﬁrst splits into four localized clusters, merging into a single radially symmetric steady state, reﬂecting the balance between attraction and diﬀusion in the underlying dynamics. Fig. 15 compares the free-energy proﬁles computed from two methods under two diﬀerent parameterizations. The left panel shows the energy evolution produced by a conventional primal-dual scheme [9] with ﬁxed physical time step Δ𝑡 = 0.5, which exhibits distinct dynamical phases: a relatively slow initial descent, a rapid drop during bump formation, and then slower relaxation. The right panel shows the energy proﬁle along the learned geometric path with arc-action parametrization, where the energy decreases much more evenly along the gradient ﬂow path. This example is included as a stress test of the central geometric claim of the paper: when the relaxation contains strongly nonuniform dynamical phases, an equal arc-action parametrization yields a more informative distribution of path images than a ﬁxed physical-time mesh.

6. Conclusion We have proposed an eﬃcient and robust computational method for the entire Wasserstein gradient ﬂows from a global pathoptimization perspective. The introduced GenWGP is a generative framework that represents the entire probability trajectory through a Lagrangian normalizingﬂow parametrization of transport maps, enabling to eﬃciently approximate the longtime relaxation process toward equilibrium without relying on sequential time marching. Our proposed pathﬁnding approach is based on the transport-ﬂow-based numerical scheme for the Dawson-Gartner action functional in the large deviation theory, which is presented in two forms: one is for the path in physical C. Liu and X. Zhou: Preprint submitted to Elsevier

Page 23 of 39

Generative Wasserstein Gradient Path Method

1.0

1.5

2.0

Energy

Figure 14: Aggregation-diﬀusion: density snapshots along the learned geometric path. Multi-bump transient states gradually merge into a smooth radial steady state.

2.5

3.0

3.5

4.0

4.0 0

10 Physical Time t

5 6 Layer

Figure 15: Free-energy proﬁles along the gradient ﬂow path. Left: the evolution computed by a conventional primal-dual scheme with a constant time step size 0.5. Right: evolution along the learned geometric path with only 11 images.

time within any speciﬁed ﬁnite time horizon, and the other is the parametrizationfree geometric formulation. The physicaltime formulation provides a horizonbased variational description together with a trainable discrete objective built from Monte Carlo particle approximation and a CrankNicolsontype discretization. The geometric formulation furthermore automatically determines the ﬁnal time and is capable of capturing the entire gradient path up to equilibrium, yielding a practical arclength or freeenergy parametrized curve. Both methods achieve wellapproximated paths for which timestepping methods would require a signiﬁcantly larger number of steps to attain comparable accuracy. Our analysis shows that these variational formulations and algorithms exhibit mathematically controlled errors. In particular, we derived an a priori KLdivergence estimate for the physicaltime formulation, established a trajectoryerror decomposition for its discrete scheme, and proved consistency of the discrete geometric objective. As demonstrated by numerical results on representative FokkerPlanck and interacting aggregation particle models, these ﬁndings indicate that GenWGP can approximate both relaxation trajectories and terminal equilibrium states in a stable and computationally eﬃcient manner. The leastaction principle underlying our method naturally extends beyond Wasserstein gradient ﬂows to more general nonequilibrium systems in probability measure space, even in the absence of a free energy or in the presence of additional nonconservative forces. The corresponding action functionals retain a similar structure, and both the C. Liu and X. Zhou: Preprint submitted to Elsevier

Page 24 of 39

Generative Wasserstein Gradient Path Method

physicaltime and geometric formulations developed in this paper remain applicableparticularly in scenarios where the terminal state is prescribed and the noiseinduced optimal transition path is of primary interest. We leave this promising application for future investigation.

References [1] Adams, S., Dirr, N., Peletier, M., and Zimmer, J. (2013). Large deviations and gradient ﬂows. Philosophical Transactions of the Royal Society A: Mathematical, Physical and Engineering Sciences, 371(2005):20120341. [2] Ambrosio, L., Gigli, N., and Savaré, G. (2008). Gradient ﬂows: in metric spaces and in the space of probability measures. Springer Science & Business Media. [3] Boﬃ, N. M. and Vanden-Eijnden, E. (2023). Probability ﬂow solution of the Fokker–Planck equation. Machine Learning: Science and Technology, 4(3):035012. [4] Boltzmann, L. (1872). Weitere studien über das wärmegleichgewicht unter gasmolekülen, volume 66. Aus der kk Hot-und Staatsdruckerei. [5] Byun, S.-S. (2024). Planar equilibrium measure problem in the quadratic ﬁelds with a point charge. Computational Methods and Function Theory, 24(2):303–332. [6] Cai, Z., Cao, Y., Huang, Y., and Zhou, X. (2026). Weak generative sampler to eﬃciently sample invariant distribution of stochastic diﬀerential equation. SIAM Journal on Scientiﬁc Computing (to appear. arXiv2405.19256 ). [7] Cai, Z., Liu, C., and Zhou, X. (2025). Weak generative sampler for stationary distributions of mckean-vlasov system. arXiv preprint arXiv:2509.12841. [8] Carrillo, J. A., Chertock, A., and Huang, Y. (2015). A ﬁnite-volume method for nonlinear nonlocal equations with a gradient ﬂow structure. Communications in Computational Physics, 17(1):233–258. [9] Carrillo, J. A., Craig, K., Wang, L., and Wei, C. (2022). Primal dual methods for wasserstein gradient ﬂows. Foundations of Computational Mathematics, pages 1–55. [10] Carrillo, J. A., McCann, R. J., and Villani, C. (2003). Kinetic equilibration rates for granular media and related equations: entropy dissipation and mass transportation estimates. Revista Matematica Iberoamericana, 19(3):971–1018. [11] Chen, R. T., Rubanova, Y., Bettencourt, J., and Duvenaud, D. K. (2018). Neural ordinary diﬀerential equations. Advances in neural information processing systems, 31. [12] Dawson, D. A. (1983). Critical dynamics and ﬂuctuations for a mean-ﬁeld model of cooperative behavior. Journal of Statistical Physics, 31(1):29–85. [13] Dawson, D. A. and Gärtner, J. (1987). Large deviations from the McKean-Vlasov limit for weakly interacting diﬀusions. Stochastics, 20(4):247–308. [14] Dawson, D. A. and Gärtner, J. (1989). Large deviations, free energy functional and quasi-potential for a mean ﬁeld model of interacting diﬀusions, volume 78. Memoirs of the American Mathematical Society. [15] Dinh, L., Sohl-Dickstein, J., and Bengio, S. (2016). Density estimation using Real NVP. In International Conference on Learning Representations. [16] Durkan, C., Bekasov, A., Murray, I., and Papamakarios, G. (2019). Neural spline ﬂows. Advances in neural information processing systems, 32. [17] E, W., Ren, W., and Vanden-Eijnden, E. (2002). String method for the study of rare events. Phys. Rev. B, 66:052301. [18] E, W., Ren, W., and Vanden-Eijnden, E. (2004). Minimum action method for the study of rare events. Comm. Pure Appl. Math., 57:637–656. [19] Feng, J. and Kurtz, T. G. (2006). Large Deviations for Stochastic Processes, volume 131 of Mathematical Surveys and Monographs. American Mathematical Society, Prividence, RI. [20] Freidlin, M. I. and Wentzell, A. D. (2012). Random Perturbations of Dynamical Systems. Grundlehren der mathematischen Wissenschaften. Springer-Verlag, New York, 3 edition. [21] Han, J., Wu, Z., Gu, S., and Zhou, X. (2026). StringNET: Neural Network based Variational Method for Transition Pathways. Communications in Computational Physics. [22] Heymann, M. and Vanden-Eijnden, E. (2008a). The geometric minimum action method: A least action principle on the space of curves. Communications on Pure and Applied Mathematics: A Journal Issued by the Courant Institute of Mathematical Sciences, 61(8):1052–1117. [23] Heymann, M. and Vanden-Eijnden, E. (2008b). The geometric minimum action method: a least action principle on the space of curves. Comm. Pure Appl. Math., 61:1052–1117. [24] Hong, S. and Chun, S. Y. (2023). Neural diﬀeomorphic non-uniform B-spline ﬂows. In Proceedings of the AAAI Conference on Artiﬁcial Intelligence, volume 37, pages 12225–12233. [25] Hu, Z., Liu, C., Wang, Y., and Xu, Z. (2024). Energetic variational neural network discretizations of gradient ﬂows. SIAM Journal on Scientiﬁc Computing, 46(4):A2528–A2556. [26] Huang, H., Yu, J., Chen, J., and Lai, R. (2023). Bridging mean-ﬁeld games and normalizing ﬂows with trajectory regularization. Journal of Computational Physics, 487:112155. [27] Huang, Y., Liu, C., and Zhou, X. (2026). Levy Score Function and Score-Based Particle Algorithm for Nonlinear Levy–Fokker–Planck Equations. SIAM Journal on Numerical Analysis (to appear), arXiv 2412.19520. [28] Hyvärinen, A. and Dayan, P. (2005). Estimation of non-normalized statistical models by score matching. Journal of Machine Learning Research, 6(4). [29] Jordan, R., Kinderlehrer, D., and Otto, F. (1998). The variational formulation of the Fokker–Planck equation. SIAM journal on mathematical analysis, 29(1):1–17. [30] Kobyzev, I., Prince, S. J., and Brubaker, M. A. (2020). Normalizing ﬂows: An introduction and review of current methods. IEEE transactions on pattern analysis and machine intelligence, 43(11):3964–3979.

C. Liu and X. Zhou: Preprint submitted to Elsevier

Page 25 of 39

Generative Wasserstein Gradient Path Method [31] Laﬀerty, J. D. (1988). The density manifold and conﬁguration space quantization. Transactions of the American Mathematical Society, 305(2):699–741. [32] Lee, W., Wang, L., and Li, W. (2024). Deep JKO: time-implicit particle methods for general nonlinear gradient ﬂows. Journal of Computational Physics, page 113187. [33] Li, L., Hurault, S., and Solomon, J. (2023). Self-consistent velocity matching of probability ﬂows. In Thirty-seventh Conference on Neural Information Processing Systems. [34] Liu, S., Li, W., Zha, H., and Zhou, H. (2022). Neural parametric fokker–planck equation. SIAM Journal on Numerical Analysis, 60(3):1385– 1449. [35] Lu, J., Wu, Y., and Xiang, Y. (2024). Score-based transport modeling for mean-ﬁeld Fokker-Planck equations. Journal of Computational Physics, 503:112859. [36] Nurbekyan, L., Lei, W., and Yang, Y. (2023). Eﬃcient natural gradient descent methods for large-scale PDE-based optimization problems. SIAM Journal on Scientiﬁc Computing, 45(4):A1621–A1655. [37] Onsager, L. (1931). Reciprocal relations in irreversible processes. i. Physical review, 37(4):405. [38] Otto, F. (2001). The geometry of dissipative evolution equations: the porous medium equation. Comm. Partial Diﬀerential Equations, 26:101–174. [39] Papamakarios, G., Nalisnick, E., Rezende, D. J., Mohamed, S., and Lakshminarayanan, B. (2021). Normalizing ﬂows for probabilistic modeling and inference. Journal of Machine Learning Research, 22(57):1–64. [40] Peletier, M. A. (2014). Variational modelling: Energies, gradient ﬂows, and large deviations. arXiv preprint arXiv:1402.1990. [41] Raissi, M., Perdikaris, P., and Karniadakis, G. E. (2019). Physics-informed neural networks: A deep learning framework for solving forward and inverse problems involving nonlinear partial diﬀerential equations. Journal of Computational physics, 378:686–707. [42] Rezende, D. and Mohamed, S. (2015). Variational inference with normalizing ﬂows. In International conference on machine learning, pages 1530–1538. PMLR. [43] Rousset, M., Stoltz, G., and Lelievre, T. (2010). Free energy computations: a mathematical perspective. World Scientiﬁc. [44] Shen, Z. and Wang, Z. (2024). Entropy-dissipation informed neural network for Mckean-Vlasov type PDEs. Advances in Neural Information Processing Systems, 36. [45] Simonnet, E. (2023). Computing non-equilibrium trajectories by a deep learning approach. Journal of Computational Physics, 491:112349. [46] Song, Y., Sohl-Dickstein, J., Kingma, D. P., Kumar, A., Ermon, S., and Poole, B. (2021). Score-based generative modeling through stochastic diﬀerential equations. In International Conference on Learning Representations. [47] Sun, Y. and Zhou, X. (2018). An improved adaptive minimum action method for the calculation of transition path in non-gradient systems. Communications in Computational Physics, 24(1):44–68. [48] Tang, K., Wan, X., and Liao, Q. (2022). Adaptive deep density approximation for Fokker-Planck equations. Journal of Computational Physics, 457:111080. [49] Vanden-Eijnden, E. and Heymann, M. (2008). The geometric minimum action method for computing minimum energy paths. J. Chem. Phys., 128:061103. [50] Vázquez, J. L. (2007). The porous medium equation: mathematical theory. Oxford university press. [51] Villani, C. (2009). Optimal transport: old and new, volume 338 of Grundlehren der Mathematischen Wissenschaften [Fundamental Principles of Mathematical Sciences]. Springer-Verlag, Berlin. [52] Villani, C. (2021). Topics in optimal transportation, volume 58. American Mathematical Soc. [53] Wan, X. (2015). A minimum action method with optimal linear time scaling. Communications in Computational Physics, 18(5):1352–1379. [54] Xie, H., Li, Z.-H., Wang, H., Zhang, L., and Wang, L. (2023). Deep variational free energy approach to dense hydrogen. Physical Review Letters, 131(12):126501. [55] Xu, C., Cheng, X., and Xie, Y. (2023). Normalizing ﬂow neural networks by JKO scheme. In Thirty-seventh Conference on Neural Information Processing Systems. [56] Yu, B. et al. (2018). The deep Ritz method: a deep learning-based numerical algorithm for solving variational problems. Communications in Mathematics and Statistics, 6(1):1–12. [57] Zhou, X., Ren, W., and E, W. (2008). Adaptive minimum action method for the study of rare events. J. Chem. Phys., 128(10):104111.

A. Proofs in Section 3 A.1. Proof of Theorem 1 Assumption 1. 1. The domain Ω is a bounded domain of ﬁnite measure (in particular, 𝕋 𝑑 ), and the boundary conditions are periodic or no-ﬂux. 2. The initial distribution 𝜌0 is absolutely continuous with respect to the Lebesgue measure (we still denote its density as 𝜌0 ) and there exists a positive constant 𝐶0 such that (𝐶0 )−1 ≤ 𝜌0 (𝑥) ≤ 𝐶0 for all 𝑥 ∈ Ω, and 𝜌0 ∈ 𝐶 2 (Ω). 3. For every 𝑇 ≥ 0, the solution 𝑝̂𝑡 ∈ 𝐶 3 (Ω) and there is a positive constant 𝐶𝑇∗ such that (𝐶𝑇∗ )−1 ≤ 𝑝̂𝑡 (𝑥) ≤ 𝐶𝑇∗ for all 𝑥 ∈ Ω, 𝑡 ∈ [0, 𝑇 ]. 4. The velocity ﬁeld 𝐟𝑡 (𝑥) = 𝐟 (𝑥, 𝑡) ∈ 𝐶 2,1 (Ω × ℝ, ℝ𝑑 ). For every 𝑇 ≥ 0, there is a positive constant 𝐶𝑇 such that sup(𝑡,𝑥)∈Ω×[0,𝑇 ] |∇ ⋅ 𝐟 (𝑥, 𝑡)| ≤ 𝐶𝑇 . C. Liu and X. Zhou: Preprint submitted to Elsevier

Page 26 of 39

Generative Wasserstein Gradient Path Method

5. The kernel 𝑊 ∈ 𝐶 3 (ℝ𝑑 ) and there exists a positive constant 𝐶𝑊 such that |∇𝑊 (𝑥)| ≤ 𝐶𝑊 for all 𝑥 ∈ Ω. 6. (Regularity of Generated Density) The density 𝑝𝑡 induced by the ﬂow Φ satisﬁes the following regularity conditions for 𝑡 ∈ [0, 𝑇 ]: • There exists a constant 𝐶𝑇𝑓 such that (𝐶𝑇𝑓 )−1 ≤ 𝑝𝑡 (𝑥) ≤ 𝐶𝑇𝑓 for all 𝑥 ∈ Ω. • The score function is bounded: there exists 𝐶𝑠 > 0 such that sup𝑥∈Ω ‖∇ log 𝑝𝑡 (𝑥)‖2 ≤ 𝐶𝑠 . We ﬁrst present a lemma that establishes the upper bound on the squared 𝐿2 -norm of two probability density functions in terms of their KullbackLeibler divergence. Lemma 1. Suppose 𝑝 and 𝑞 are two probability densities on Ω, and there exists a positive constant 𝐶MAX such that 0 < 𝑝(𝑥), 𝑞(𝑥) < 𝐶MAX for all 𝑥 ∈ Ω. Then we have ∫Ω

|𝑝(𝑥) − 𝑞(𝑥)|2 d𝑥 ≤

2𝐶MAX 𝐷 (𝑝‖𝑞). 1 − log 2 KL

𝑝(𝑥) PROOF. Deﬁne 𝜁 (𝑥) ∶= 𝑞(𝑥)−𝑝(𝑥) for 𝑥 ∈ Ω. Then 𝐷KL (𝑝‖𝑞) = ∫Ω 𝑝(𝑥) log 𝑞(𝑥) d𝑥 = − ∫Ω 𝑝(𝑥) log(1 + 𝜁 (𝑥)) d𝑥. 𝑝(𝑥)

Deﬁne two Borel sets: 𝐴 ∶= {𝑥 ∣ 𝜁 (𝑥) > 1} and 𝐵 ∶= {𝑥 ∣ 𝜁 (𝑥) ≤ 1}; then for 𝑥 ∈ 𝐴, 1 + 𝜁 (𝑥) ≤ 𝑒𝛼𝜁 (𝑥) 2 where 𝛼 = log 2 > 0; for 𝑥 ∈ 𝐵, 1 + 𝜁 (𝑥) ≤ 𝑒𝜁 (𝑥)−𝛽𝜁 (𝑥) where 𝛽 = 1 − log 2 > 0. Note that ∫Ω 𝑝(𝑥)𝜁 (𝑥) d𝑥 = ∫Ω (𝑞(𝑥) − 𝑝(𝑥)) d𝑥 = 0, which implies ∫𝐴 𝑝(𝑥)𝜁 (𝑥) d𝑥 = − ∫𝐵 𝑝(𝑥)𝜁 (𝑥) d𝑥. Thus 𝐷KL (𝑝‖𝑞) = −

∫𝐴

≥ −𝛼

𝑝(𝑥) log(1 + 𝜁 (𝑥)) d𝑥 −

∫𝐴

= (1 − 𝛼)

𝑝(𝑥)𝜁 (𝑥) d𝑥 − ∫𝐴

∫𝐵

𝑝(𝑥) log(1 + 𝜁 (𝑥)) d𝑥

𝑝(𝑥)𝜁 (𝑥) d𝑥 + 𝛽

∫𝐵

𝑝(𝑥)𝜁 (𝑥) d𝑥 + 𝛽 𝑝(𝑥)𝜁 (𝑥)2 d𝑥 ∫𝐵 ( (

= (1 − log 2)

∫𝐴

|𝑞(𝑥) − 𝑝(𝑥)| d𝑥 +

∫𝐵

𝑝(𝑥)

𝑝(𝑥)𝜁 (𝑥)2 d𝑥

𝑞(𝑥) − 𝑝(𝑥) 𝑝(𝑥)

) d𝑥 .

MAX

Lemma 2. Let 𝑝(𝑥) and 𝑞(𝑥) be probability densities deﬁned on a domain Ω satisfying 0 < 𝐶min ≤ 𝑝(𝑥), 𝑞(𝑥) ≤ 𝐶MAX < ∞. Assume the score function ∇ log 𝑝(𝑥) is bounded with 𝐶𝑠 = sup𝑥∈Ω ‖∇ log 𝑝(𝑥)‖2 < ∞. Then, for any real 𝑘 ≠ 0, the following inequality holds: [

]( ‖ ‖2 𝑝(𝑥) ) ‖∇ log 𝑝(𝑥) ‖ 𝑝(𝑥) d𝑥+𝐾 𝜎 𝐷KL (𝑝‖𝑞), 𝑝𝑘 (𝑥)∇ log 𝑝(𝑥)−𝑞 𝑘 (𝑥)∇ log 𝑞(𝑥) ⋅ ∇ log 𝑝(𝑥) d𝑥 ≥ 𝐾1𝜎 ‖ 2 ∫Ω ∫Ω ‖ 𝑞(𝑥) 𝑞(𝑥) ‖ ‖ (A.1.1) where 𝜎 = sign(𝑘) ∈ {+, −}, and for any 𝜆 > 0, 𝑘 𝐾1+ = 𝐶min −

2𝑘−2 , 𝐶 2𝑘−2 } 2 𝑘2 𝐶MAX max{𝐶min MAX

𝜆 𝐶, 2 𝑠

𝐾2+ = −

𝜆 𝐶, 2 𝑠

𝐾2− = −

(1 − log 2)𝜆

< 0,

(𝑘 > 0),

and 𝑘 𝐾1− = 𝐶MAX −

2𝑘−2 2 𝑘2 𝐶MAX 𝐶min

(1 − log 2)𝜆

< 0,

(𝑘 < 0).

In the applications below we take 𝑝 = 𝑝𝑡 and 𝑞 = 𝑝̂𝑡 . Then the existence of positive constants 𝐶min and 𝐶MAX is guaranteed by Assumptions 3 and 6. C. Liu and X. Zhou: Preprint submitted to Elsevier

Page 27 of 39

Generative Wasserstein Gradient Path Method

PROOF. Decompose the integrand on the LHS as 𝐼1 + 𝐼2 : [ 𝑝] 𝑝 LHS = (𝑝𝑘 − 𝑞 𝑘 )∇ log 𝑝 + 𝑞 𝑘 ∇ log ⋅ ∇ log 𝑝 d𝑥 ∫Ω 𝑞 𝑞 [ ] ‖ 𝑝 𝑝 ‖2 = (𝑝𝑘 − 𝑞 𝑘 )∇ log 𝑝 ⋅ ∇ log 𝑝 d𝑥 + 𝑞𝑘 ‖ ∇ log ‖ 𝑝 d𝑥 . ‖ ∫Ω ∫Ω ‖ 𝑞 𝑞‖ ‖ ⏟⏞⏞⏞⏞⏞⏞⏞⏞⏞⏞⏞⏞⏞⏞⏞⏞⏞⏞⏞⏞⏞⏞⏞⏞⏞⏟⏞⏞⏞⏞⏞⏞⏞⏞⏞⏞⏞⏞⏞⏞⏞⏞⏞⏞⏞⏞⏞⏞⏞⏞⏞⏟ ⏟⏞⏞⏞⏞⏞⏞⏞⏞⏞⏞⏞⏞⏞⏟⏞⏞⏞⏞⏞⏞⏞⏞⏞⏞⏞⏞⏞⏟ 𝐼1

𝐼2

To estimate 𝐼1 , we apply the mean value theorem to 𝑓 (𝑠) = 𝑠𝑘 . For each 𝑥, there exists 𝜉(𝑥) ∈ [min{𝑝, 𝑞}, max{𝑝, 𝑞}] such that 𝑝𝑘 − 𝑞 𝑘 = 𝑘𝜉 𝑘−1 (𝑝 − 𝑞). Applying Young’s inequality with any constant 𝜆 > 0: [ ] 𝑝 𝑘−1 𝐼1 = 𝑘𝜉 (𝑝 − 𝑞) ∇ log 𝑝 ⋅ ∇ log 𝑝 d𝑥 ∫Ω 𝑞 [ ]2 [ 𝑘−1 ]2 𝑝 1 𝜆 ≥− ∇ log 𝑝 ⋅ ∇ log 𝑘𝜉 (𝑝 − 𝑞) 𝑝 d𝑥 − 𝑝 d𝑥. 2𝜆 ∫Ω 2 ∫Ω 𝑞 Using CauchySchwarz and the bound 𝐶𝑠 : [ ]2 2 ‖ ‖ ‖2 𝑝 𝑝‖ ‖ ≤ 𝐶𝑠 ‖∇ log 𝑝 ‖ . ∇ log 𝑝 ⋅ ∇ log ≤ ‖∇ log 𝑝‖2 ‖ ∇ log ‖ ‖ 𝑞 𝑞‖ 𝑞‖ ‖ ‖ ‖ ‖

(A.1.2)

Also, since 𝜉(𝑥) ∈ [𝐶min , 𝐶MAX ], we have 𝑘−1 𝑘−1 |𝜉(𝑥)𝑘−1 | ≤ max{𝐶min , 𝐶MAX }.

Hence, using Lemma 1, [ ∫Ω

]2 2𝑘−2 2𝑘−2 𝑘𝜉 𝑘−1 (𝑝 − 𝑞) 𝑝𝑑𝑥 ≤ 𝑘2 max{𝐶min , 𝐶MAX }𝐶MAX ≤

∫Ω

2𝑘−2 , 𝐶 2𝑘−2 } 2 2𝑘2 𝐶MAX max{𝐶min MAX

1 − log 2

(𝑝 − 𝑞)2 𝑑𝑥

𝐷KL (𝑝‖𝑞).

For 𝐼2 , we distinguish two cases. If 𝑘 > 0, since 𝑞(𝑥) ≥ 𝐶min , we have 𝑘 𝐼2 ≥ 𝐶min

‖2 ‖ ‖∇ log 𝑝 ‖ 𝑝𝑑𝑥. ‖ ∫Ω ‖ 𝑞‖ ‖

If 𝑘 < 0, since 𝑞(𝑥) ≤ 𝐶MAX and 𝑠 ↦ 𝑠𝑘 is decreasing on (0, ∞), we have ‖2 ‖ ‖∇ log 𝑝 ‖ 𝑝𝑑𝑥. ∫Ω ‖ 𝑞‖ ‖ ‖ The result follows by combining the bounds for 𝐼1 and 𝐼2 . 𝑘 𝐼2 ≥ 𝐶MAX

PROOF (PROOF OF T HEOREM 1). The proof is inspired by the methodologies presented in [3, Proposition 1] and [44, Appendix E]. However, dealing with the non-entropy internal energy term introduces substantial complexity, requiring the introduction of novel techniques for a thorough analysis. First, using the deﬁnition of KL divergence, together with the facts that 𝑝𝑡 satisﬁes (16), 𝑝̂𝑡 satisﬁes (7) and ∫ 𝜕𝑡 𝑝𝑡 d𝑥 = 0, we derive that ( ( )) ( ) 𝑝 𝑝 𝑝 𝑝 d 𝐷KL (𝑝𝑡 ‖̂ 𝑝𝑡 ) = 𝜕𝑡 𝑝𝑡 log 𝑡 + 𝑝𝑡 𝜕𝑡 log 𝑡 d𝑥 = 𝜕𝑡 𝑝𝑡 log 𝑡 − 𝜕𝑡 𝑝̂𝑡 𝑡 d𝑥 ∫Ω ∫Ω d𝑡 𝑝̂𝑡 𝑝̂ 𝑝̂𝑡 𝑝̂𝑡 (𝑡 ) ( ) 𝑝𝑡 𝑝𝑡 = 𝐟(𝑥, 𝑡) ⋅ ∇ log 𝑝𝑡 d𝑥 − [̂ 𝑝𝑡 ] ⋅ ∇ 𝑝̂𝑡 d𝑥 ∫Ω ∫Ω 𝑝̂𝑡 𝑝̂𝑡 ( ) ( ) 𝑝𝑡 = 𝐟 (𝑥, 𝑡) − [𝑝𝑡 ] + [𝑝𝑡 ] − [̂ 𝑝𝑡 ] ⋅ ∇ log 𝑝𝑡 d𝑥. ∫Ω 𝑝̂𝑡 C. Liu and X. Zhou: Preprint submitted to Elsevier

Page 28 of 39

Generative Wasserstein Gradient Path Method

Using the expression of [̂ 𝑝𝑡 ] by (8), we decompose the above as: ( ) ( ) 𝑝𝑡 d 𝐷 (𝑝 ‖̂ 𝑝)= 𝐟 (𝑥, 𝑡) − [𝑝𝑡 ](𝑥, 𝑡) ⋅ ∇ log 𝑝𝑡 d𝑥 ∫Ω d𝑡 KL 𝑡 𝑡 𝑝̂𝑡 ⏟⏞⏞⏞⏞⏞⏞⏞⏞⏞⏞⏞⏞⏞⏞⏞⏞⏞⏞⏞⏞⏞⏞⏞⏞⏞⏞⏞⏞⏞⏞⏞⏞⏟⏞⏞⏞⏞⏞⏞⏞⏞⏞⏞⏞⏞⏞⏞⏞⏞⏞⏞⏞⏞⏞⏞⏞⏞⏞⏞⏞⏞⏞⏞⏞⏞⏟ Perturbation

−

) 𝑝𝑡 𝑝𝑡 d𝑥 ∫Ω 𝑝̂𝑡 ⏟⏞⏞⏞⏞⏞⏞⏞⏞⏞⏞⏞⏞⏞⏞⏞⏞⏞⏞⏞⏞⏞⏞⏞⏞⏞⏞⏞⏞⏞⏞⏞⏞⏞⏞⏞⏞⏞⏟⏞⏞⏞⏞⏞⏞⏞⏞⏞⏞⏞⏞⏞⏞⏞⏞⏞⏞⏞⏞⏞⏞⏞⏞⏞⏞⏞⏞⏞⏞⏞⏞⏞⏞⏞⏞⏞⏟ 𝛽

( −1

) ∇𝑈𝑚′ (𝑝𝑡 ) − ∇𝑈𝑚′ (̂ 𝑝𝑡 ) ⋅ ∇ log

(

Internal

(

−

) 𝑝𝑡 𝑝𝑡 d𝑥 ∫Ω 𝑝̂𝑡 ⏟⏞⏞⏞⏞⏞⏞⏞⏞⏞⏞⏞⏞⏞⏞⏞⏞⏞⏞⏞⏞⏞⏞⏞⏞⏟⏞⏞⏞⏞⏞⏞⏞⏞⏞⏞⏞⏞⏞⏞⏞⏞⏞⏞⏞⏞⏞⏞⏞⏞⏟ (∇𝑉 − ∇𝑉 ) ⋅ ∇ log Potential

(

) 𝑝𝑡 − 𝑝𝑡 d𝑥 . ∫Ω 𝑝̂𝑡 ⏟⏞⏞⏞⏞⏞⏞⏞⏞⏞⏞⏞⏞⏞⏞⏞⏞⏞⏞⏞⏞⏞⏞⏞⏞⏞⏞⏞⏞⏞⏟⏞⏞⏞⏞⏞⏞⏞⏞⏞⏞⏞⏞⏞⏞⏞⏞⏞⏞⏞⏞⏞⏞⏞⏞⏞⏞⏞⏞⏞⏟ (

) ∇𝑊 ∗ (𝑝𝑡 − 𝑝̂𝑡 ) ⋅ ∇ log Interaction

The bounds for these terms are studied below in four steps. (Step 1.) For the perturbation part, by Young’s inequality with any constant Λ > 0: ( ) ( ) 𝑝𝑡 𝐟 (𝑥, 𝑡) − [𝑝𝑡 ](𝑥, 𝑡) ⋅ ∇ log 𝑝𝑡 d𝑥 ∫Ω 𝑝̂𝑡 ( )‖ 2 ‖ 𝑝𝑡 ‖ 1 Λ ‖ ‖𝐟 (𝑥, 𝑡) − [𝑝 ](𝑥, 𝑡)‖2 𝑝 d𝑥. ≤ ∇ log ‖ 𝑝𝑡 d𝑥 + ‖ 𝑡 ‖ ‖ 𝑡 ‖ ∫ 2Λ ∫Ω ‖ 2 𝑝 ̂ Ω 𝑡 ‖ ‖ (Step 2.) For the internal energy part, there are two cases about diﬀusion (𝑚 = 1) and power-law case (𝑚 ≠ 1) for the internal energy 𝑈𝑚 . We discuss each case respectively. (Entropy case ∶ 𝑈𝑚 (𝜌) = 𝑈1 (𝜌) = 𝜌 log 𝜌) ( ) ( )‖2 ‖ ) ( 𝑝𝑡 𝑝𝑡 ‖ 1 ‖ ′ ′ −1 𝑝𝑡 d𝑥 = − 𝑝𝑡 ) ⋅ ∇ log − ∇𝑈𝑚 (𝑝𝑡 ) − ∇𝑈𝑚 (̂ 𝛽 ‖ 𝑝𝑡 d𝑥. ‖∇ log ‖ ∫ ∫Ω 𝛽 Ω‖ 𝑝̂𝑡 𝑝̂𝑡 ‖ ‖ 1 (Power-law case ∶ 𝑈𝑚 (𝜌) = 𝑚−1 𝜌𝑚 , 𝑚 ≠ 1) Set 𝑘 = 𝑚 − 1. Then 𝑘 ≠ 0, and by Lemma 2,

−

∫Ω

𝛽 −1

) 𝑚 ( ∇(𝑝𝑡 )𝑚−1 − ∇(̂ 𝑝𝑡 )𝑚−1 ⋅ ∇ log 𝑚−1 𝜎

( 𝜎

𝑝𝑡 𝑝̂𝑡

)

2 ‖ 𝑝𝑡 ‖ 𝑚 𝜎 𝑚 𝜎 ‖ ‖ 𝑝𝑡 𝑑𝑥 ≤ − 𝐾1 𝑚 𝑝𝑡 ), ‖∇ log ‖ 𝑝𝑡 𝑑𝑥 − 𝐾2 𝑚 𝐷KL (𝑝𝑡 ‖̂ ‖ ∫Ω ‖ 𝛽 𝛽 𝑝 ̂ 𝑡‖ ‖

where 𝜎𝑚 = sign(𝑚 − 1) ∈ {+, −}, and 𝐾1 𝑚 , 𝐾2 𝑚 are the constants from Lemma 2 evaluated at 𝑘 = 𝑚 − 1. To ensure 𝜎 𝐾1 𝑚 > 0, it is suﬃcient to choose { 𝑚−1 , 𝐶min 𝑚 > 1, 2𝑎𝑚 0<𝜆< , 𝑎𝑚 = 𝑚−1 𝐶𝑠 𝐶MAX , 0 < 𝑚 < 1. (Step 3.) The potential term vanishes trivially. (Step 4.) Due to Assumption 5 and CsiszárKullbackPinsker inequality: ( ) ( ) 𝑝𝑡 − ∇𝑊 ∗ (𝑝𝑡 − 𝑝̂𝑡 ) ⋅ ∇ log 𝑝𝑡 d𝑥 ∫Ω 𝑝̂𝑡 ( )2 ( )‖ 2 ‖ 𝑝𝑡 ‖ 1 Λ ‖ 2 |𝑝 (𝑦) − 𝑝̂𝑡 (𝑦)| d𝑦 ≤ ‖∇ log ‖ 𝑝𝑡 d𝑥 + (𝐶𝑊 ) ∫Ω 𝑡 2Λ ∫Ω ‖ 2 𝑝̂𝑡 ‖ ‖ ‖ ( )‖ 2 ‖ 𝑝𝑡 ‖ 1 ‖ 2 ≤ 𝐷KL (𝑝𝑡 ‖̂ 𝑝𝑡 ). ‖∇ log ‖ 𝑝𝑡 d𝑥 + Λ𝐶𝑊 ‖ 2Λ ∫Ω ‖ 𝑝 ̂ 𝑡 ‖ ‖ C. Liu and X. Zhou: Preprint submitted to Elsevier

Page 29 of 39

Generative Wasserstein Gradient Path Method

Finally, we collect these bounds. (Entropy Case) Setting Λ = 𝛽, we obtain: ( 2 ) 𝛽 d ‖𝐟 (𝑥, 𝑡) − [𝑝𝑡 ](𝑥, 𝑡)‖2 𝑝𝑡 d𝑥. 𝑝𝑡 ) ≤ 𝛽𝐶𝑊 𝐷KL (𝑝𝑡 ‖̂ 𝑝𝑡 ) + 𝐷KL (𝑝𝑡 ‖̂ ‖ d𝑡 2 ∫Ω ‖ (Power-law Case) d 𝐷 (𝑝 ‖̂ 𝑝)≤ d𝑡 KL 𝑡 𝑡

Setting Λ =

(

1 𝑚 𝜎𝑚 − 𝐾 Λ 𝛽 1

)

( ( )‖2 ) ‖ 𝑝𝑡 ‖ 𝑚 𝜎 Λ ‖ 2 ‖𝐟(𝑥, 𝑡) − [𝑝𝑡 ](𝑥, 𝑡)‖2 𝑝𝑡 d𝑥. − 𝐾2 𝑚 𝐷KL (𝑝𝑡 ‖̂ 𝑝𝑡 )+ ‖∇ log ‖ 𝑝𝑡 d𝑥+ Λ𝐶𝑊 ‖ ‖ ‖ ∫Ω ‖ ∫ 𝛽 2 𝑝 ̂ Ω 𝑡 ‖ ‖

𝛽 𝜎 , the gradient term vanishes, and therefore 𝑚𝐾1 𝑚

(

d 𝐷 (𝑝 ‖̂ 𝑝)≤ d𝑡 KL 𝑡 𝑡

𝛽

𝑚 𝜎𝑚 2 𝐾 𝜎 𝐶𝑊 − 𝛽 2 𝑚𝐾1 𝑚

By Gronwall’s inequality, we obtain ( sup 𝐷KL (𝑝𝑡 ‖̂ 𝑝𝑡 ) ≤ exp{𝛾𝑇 } 𝛼 𝑡∈[0,𝑇 ]

)

𝑇

∫0 ∫Ω

𝐷KL (𝑝𝑡 ‖̂ 𝑝𝑡 ) +

𝛽

𝜎

2𝑚𝐾1 𝑚 ∫Ω

‖𝐟(𝑥, 𝑡) − [𝑝𝑡 ](𝑥, 𝑡)‖2 𝑝𝑡 d𝑥. ‖ ‖

) ‖𝐟 (𝑥, 𝑡) − [𝑝𝑡 ](𝑥, 𝑡)‖2 𝑝𝑡 𝑑𝑥𝑑𝑡 ,

𝑇 ‖2 where 12 ∫0 ∫Ω ‖ ‖𝑓 − [𝑝𝑡 ]‖ 𝑝𝑡 d𝑥 d𝑡 = 𝐽 [Φ] represents the ﬂow matching loss (17), and the constants in the above bound are deﬁned as: ⎧𝛽 , ⎧𝛽𝐶 2 , 𝑚 = 1, 𝑚 = 1, ⎪ 𝑊 𝛼 ⎪2 = (A.1.3) 𝛾=⎨ 𝛽 𝑚 𝜎𝑚 𝛽 2 2 ⎨ 𝐾2 , 𝑚 ≠ 1, 𝜎 𝑚 𝐶𝑊 − ⎪ ⎪ 𝜎𝑚 , 𝑚 ≠ 1, 𝛽 ⎩ 2𝑚𝐾 ⎩ 𝑚𝐾1 1

𝜎

where 𝜎𝑚 = sign(𝑚 − 1) and 𝐾1 𝑚 , 𝐾2 𝑚 are the constants from Lemma 2 with 𝑘 = 𝑚 − 1.

A.2. Proof of Theorem 2

PROOF (PROOF OF T HEOREM 2). For notational simplicity, we write 𝑋𝑘∗ (𝑧) ∶= 𝑋 ∗ (𝑡𝑘 , 𝑧) and 𝑋𝑘𝑁 (𝑧) ∶= Φ𝑘 (𝑧) for 𝑘 = 0, … , 𝐾, where 𝑡𝑘 = 𝑘Δ𝑡. We deﬁne the (expected) trajectory error at the grid points by ‖ ‖ 𝜖𝑡𝑘 ∶= 𝔼𝑧∼𝜌0 ‖𝑋𝑘∗ (𝑧) − 𝑋𝑘𝑁 (𝑧)‖. ‖ ‖ Step 1: Average-velocity formulation at one time step. Fix 𝑘 ∈ {1, … , 𝐾} and 𝑧. Using the integral form of the exact ODE, 𝑡𝑘 ( ) ∗ (𝑧) = [̂ 𝑝𝑠 ] 𝑋 ∗ (𝑠, 𝑧) d𝑠, 𝑋𝑘∗ (𝑧) − 𝑋𝑘−1 ∫𝑡𝑘−1 we deﬁne the exact average velocity over the interval [𝑡𝑘−1 , 𝑡𝑘 ] by 𝑣̄ ∗𝑘 (𝑧) ∶=

∗ (𝑧) 𝑋𝑘∗ (𝑧) − 𝑋𝑘−1

Δ𝑡

𝑡𝑘 ( ) 1 [̂ 𝑝𝑠 ] 𝑋 ∗ (𝑠, 𝑧) d𝑠. Δ𝑡 ∫𝑡𝑘−1

By deﬁnition of the discrete trajectory 𝑋𝑘𝑁 (𝑧) = Φ𝑘 (𝑧), we introduce the discrete kinematic velocity 𝑣𝑁 𝑘 (𝑧) ∶=

𝑁 (𝑧) 𝑋𝑘𝑁 (𝑧) − 𝑋𝑘−1

. Δ𝑡 We also introduce the ideal Crank-Nicolson velocity (on the continuous model) and its empirical counterpart: ( ( ∗ ) ( ∗ )) 1 [̂ 𝑝 𝑣CN,∗ (𝑧) ∶= ] 𝑋 (𝑧) + [̂ 𝑝 ] 𝑋𝑘 (𝑧) , 𝑡 𝑡 𝑘−1 𝑘 𝑘−1 𝑘 2 ( ( 𝑁 ) ( 𝑁 )) 1  [̂ 𝑝 ] 𝑋 (𝑧) +  [̂ 𝑝 ] 𝑣CN,𝑁 (𝑧) ∶= 𝑁 𝑡𝑘 𝑋𝑘 (𝑧) . 𝑘−1 𝑘 2 𝑁 𝑡𝑘−1 By deﬁnition, ‖ ‖ CN,𝑁 sup ‖𝑣𝑁 (𝑧)‖ ≤ 𝜀. 𝑘 (𝑧) − 𝑣𝑘 ‖ 𝑘,𝑧 ‖ C. Liu and X. Zhou: Preprint submitted to Elsevier

Page 30 of 39

Generative Wasserstein Gradient Path Method

Step 2: Local Crank-Nicolson discretization error. Deﬁne ( ) 𝑓𝑘 (𝑠, 𝑧) ∶= [̂ 𝑝𝑠 ] 𝑋 ∗ (𝑠, 𝑧) ,

𝑠 ∈ [𝑡𝑘−1 , 𝑡𝑘 ].

Then 𝑣̄ ∗𝑘 (𝑧) =

𝑡

𝑘 1 𝑓 (𝑠, 𝑧) d𝑠, Δ𝑡 ∫𝑡𝑘−1 𝑘

𝑣CN,∗ (𝑧) = 𝑘

𝑓𝑘 (𝑡𝑘−1 , 𝑧) + 𝑓𝑘 (𝑡𝑘 , 𝑧) . 2

Under the smoothness assumptions in Appendix 1, 𝑓𝑘 (⋅, 𝑧) is twice continuously diﬀerentiable in 𝑡, with uniformly bounded second derivative. Applying the classical error estimate of the trapezoidal rule (which is the timediscretization underlying Crank-Nicolson) yields ‖ ∗ ‖ (𝑧)‖ ≤ 𝐶Δ𝑡2 , ‖𝑣̄ 𝑘 (𝑧) − 𝑣CN,∗ 𝑘 ‖ ‖ for some constant 𝐶 > 0 independent of 𝑘, Δ𝑡, and 𝑁. Equivalently, 𝑣̄ ∗𝑘 (𝑧) = 𝑣CN,∗ (𝑧) + (Δ𝑡2 ). 𝑘

(A.2.1)

Step 3: Ideal CN velocity vs. empirical CN velocity. We next bound the diﬀerence 𝑣CN,∗ (𝑧) − 𝑣CN,𝑁 (𝑧). Let 𝑘 𝑘 ( ) 𝑎𝑗 (𝑧) ∶= [̂ 𝑝𝑡𝑗 ] 𝑋𝑗∗ (𝑧) ,

Then 𝑣CN,∗ (𝑧) − 𝑣CN,𝑁 (𝑧) = 𝑘 𝑘

( ) 𝑏𝑗 (𝑧) ∶= 𝑁 [̂ 𝑝𝑡𝑗 ] 𝑋𝑗𝑁 (𝑧) ,

𝑗 = 𝑘 − 1, 𝑘.

( ) 1 (𝑎𝑘−1 − 𝑏𝑘−1 ) + (𝑎𝑘 − 𝑏𝑘 ) , 2

and hence ( ) ‖ CN,∗ ‖ 1 (𝑧) ≤ ‖𝑎 (𝑧) − 𝑏 (𝑧)‖ + ‖𝑎 (𝑧) − 𝑏 (𝑧)‖ . ‖𝑣𝑘 (𝑧) − 𝑣CN,𝑁 ‖ 𝑘−1 𝑘−1 𝑘 𝑘 𝑘 ‖ ‖ 2 For a generic index 𝑗 ∈ {𝑘 − 1, 𝑘}, we decompose )‖ ) ( ( ‖ 𝑝𝑡𝑗 ] 𝑋𝑗𝑁 (𝑧) ‖ ‖𝑎𝑗 (𝑧) − 𝑏𝑗 (𝑧)‖ = ‖[̂ 𝑝𝑡𝑗 ] 𝑋𝑗∗ (𝑧) − 𝑁 [̂ ‖ ‖ ( ∗ )‖ ( ∗ ) ‖ 𝑝𝑡𝑗 ] 𝑋𝑗 (𝑧) ‖ ≤ ‖[̂ 𝑝𝑡𝑗 ] 𝑋𝑗 (𝑧) − 𝑁 [̂ ‖ ‖ ⏟⏞⏞⏞⏞⏞⏞⏞⏞⏞⏞⏞⏞⏞⏞⏞⏞⏞⏞⏞⏞⏞⏞⏞⏞⏞⏞⏟⏞⏞⏞⏞⏞⏞⏞⏞⏞⏞⏞⏞⏞⏞⏞⏞⏞⏞⏞⏞⏞⏞⏞⏞⏞⏞⏟

(A.2.2)

ﬁnite-sample approximation of the ﬁeld

)‖ ) ( ( ‖ 𝑝𝑡𝑗 ] 𝑋𝑗𝑁 (𝑧) ‖ . + ‖𝑁 [̂ 𝑝𝑡𝑗 ] 𝑋𝑗∗ (𝑧) − 𝑁 [̂ ‖ ‖ ⏟⏞⏞⏞⏞⏞⏞⏞⏞⏞⏞⏞⏞⏞⏞⏞⏞⏞⏞⏞⏞⏞⏞⏞⏞⏞⏞⏞⏞⏟⏞⏞⏞⏞⏞⏞⏞⏞⏞⏞⏞⏞⏞⏞⏞⏞⏞⏞⏞⏞⏞⏞⏞⏞⏞⏞⏞⏞⏟ Lipschitzinspace × trajectory error

Since 𝑁 [̂ 𝑝𝑡𝑗 ](𝑥) is an unbiased empirical average of i.i.d. random vectors with uniformly bounded variance 5, it follows from the standard Monte Carlo estimate that ‖ ‖ 𝔼‖[̂ 𝑝𝑡𝑗 ](𝑥) − 𝑁 [̂ 𝑝𝑡𝑗 ](𝑥)‖ ≤ 𝐶1 𝑁 −1∕2 . ‖ ‖ In particular, ( ) ( )‖ ‖ 𝑝𝑡𝑗 ] 𝑋𝑗∗ (𝑧) ‖ ≤ 𝐶1 𝑁 −1∕2 . (A.2.3) 𝔼𝑧∼𝜌0 ‖[̂ 𝑝𝑡𝑗 ] 𝑋𝑗∗ (𝑧) − 𝑁 [̂ ‖ ‖ Moreover, by the Lipschitz continuity of 𝑁 [̂ 𝑝𝑡𝑗 ](⋅) in space (Assumption 1), there exists 𝐿𝑥 > 0 such that ( ) ( ) ‖ ‖ ‖ ‖ 𝑝𝑡𝑗 ] 𝑋𝑗∗ (𝑧) − 𝑁 [̂ 𝑝𝑡𝑗 ] 𝑋𝑗𝑁 (𝑧) ‖ ≤ 𝐿𝑥 ‖𝑋𝑗∗ (𝑧) − 𝑋𝑗𝑁 (𝑧)‖. ‖𝑁 [̂ ‖ ‖ ‖ ‖ Taking expectation in 𝑧 ∼ 𝜌0 and recalling the deﬁnition of 𝜖𝑡𝑗 , we obtain ( ) ( )‖ ‖ 𝔼𝑧∼𝜌0 ‖𝑁 [̂ 𝑝𝑡𝑗 ] 𝑋𝑗∗ (𝑧) − 𝑁 [̂ 𝑝𝑡𝑗 ] 𝑋𝑗𝑁 (𝑧) ‖ ≤ 𝐿𝑥 𝜖𝑡𝑗 . (A.2.4) ‖ ‖ Combining (A.2.2), (A.2.3), and (A.2.4), we deduce that ‖ ‖ 𝔼𝑧∼𝜌0 ‖𝑣CN,∗ (𝑧) − 𝑣CN,𝑁 (𝑧)‖ ≤ 𝐶1 𝑁 −1∕2 + 𝐶2 max{𝜖𝑡𝑘−1 , 𝜖𝑡𝑘 }, (A.2.5) 𝑘 ‖ 𝑘 ‖ for some constant 𝐶2 > 0 independent of 𝑁 and Δ𝑡. C. Liu and X. Zhou: Preprint submitted to Elsevier

Page 31 of 39

Generative Wasserstein Gradient Path Method

Step 4: Consistency residual relative to the exact Crank-Nicolson driving force. By the deﬁnition of 𝜀 in (22), ‖ ‖ CN,𝑁 sup ‖𝑣𝑁 (𝑧)‖ ≤ 𝜀, 𝑘 (𝑧) − 𝑣𝑘 ‖ 𝑧 ‖ and thus, in particular, ‖ ‖ 𝔼𝑧∼𝜌0 ‖𝑣𝑁 (𝑧) − 𝑣CN,𝑁 (𝑧)‖ ≤ 𝜀. 𝑘 ‖ 𝑘 ‖

(A.2.6)

Step 5: One-step error recursion in average velocity form. For each 𝑧, ( ) ∗ 𝑁 𝑋𝑘∗ (𝑧) − 𝑋𝑘𝑁 (𝑧) = 𝑋𝑘−1 (𝑧) − 𝑋𝑘−1 (𝑧) + Δ𝑡 𝑣̄ ∗𝑘 (𝑧) − 𝑣𝑁 (𝑧) , 𝑘 hence ‖ ‖ ‖ ‖ ∗ ‖ ‖ ∗ 𝑁 (𝑧)‖ + Δ𝑡‖𝑣̄ ∗𝑘 (𝑧) − 𝑣𝑁 . (𝑧) − 𝑋𝑘−1 ‖𝑋𝑘 (𝑧) − 𝑋𝑘𝑁 (𝑧)‖ ≤ ‖𝑋𝑘−1 𝑘 (𝑧)‖ ‖ ‖ ‖ ‖ ‖ ‖ Taking expectation and using the triangle inequality, ‖ ‖ 𝜖𝑡𝑘 ≤ 𝜖𝑡𝑘−1 + Δ𝑡𝔼𝑧∼𝜌0 ‖𝑣̄ ∗𝑘 (𝑧) − 𝑣𝑁 . 𝑘 (𝑧)‖ ‖ ‖ We now split the average-velocity discrepancy into the three components derived above: CN,∗ ∗ (𝑧) + 𝑣CN,∗ (𝑧) − 𝑣CN,𝑁 (𝑧) + 𝑣CN,𝑁 (𝑧) − 𝑣𝑁 𝑣̄ ∗𝑘 (𝑧) − 𝑣𝑁 𝑘 (𝑧) . 𝑘 (𝑧) = 𝑣̄ 𝑘 (𝑧) − 𝑣𝑘 𝑘 𝑘 𝑘 ⏟⏞⏞⏞⏞⏞⏞⏞⏞⏟⏞⏞⏞⏞⏞⏞⏞⏞⏟ ⏟⏞⏞⏞⏞⏞⏞⏞⏞⏞⏞⏞⏟⏞⏞⏞⏞⏞⏞⏞⏞⏞⏞⏞⏟ ⏟⏞⏞⏞⏞⏞⏞⏞⏞⏞⏟⏞⏞⏞⏞⏞⏞⏞⏞⏞⏟ CN discretization error

ﬁnite-sample + trajectory error

consistency residual

Using (A.2.1), (A.2.5), and (A.2.6), we obtain ‖ ‖ 𝔼𝑧∼𝜌0 ‖𝑣̄ ∗𝑘 (𝑧) − 𝑣𝑁 ≤ (Δ𝑡2 ) + 𝐶1 𝑁 −1∕2 + 𝐶2 max{𝜖𝑡𝑘−1 , 𝜖𝑡𝑘 } + 𝜀. 𝑘 (𝑧)‖ ‖ ‖ Therefore, ( ) 𝜖𝑡𝑘 ≤ 𝜖𝑡𝑘−1 + Δ𝑡 (Δ𝑡2 ) + 𝐶1 𝑁 −1∕2 + 𝐶2 max{𝜖𝑡𝑘−1 , 𝜖𝑡𝑘 } + 𝜀 .

(A.2.7)

Using max{𝜖𝑡𝑘−1 , 𝜖𝑡𝑘 } ≤ 𝜖𝑡𝑘−1 + 𝜖𝑡𝑘 and absorbing the resulting terms, we can rewrite (A.2.7) as ( ) 𝜖𝑡𝑘 ≤ (1 + 𝐶Δ𝑡)𝜖𝑡𝑘−1 + Δ𝑡 𝐶Δ𝑡2 + 𝐶𝑁 −1∕2 + 𝐶𝜀 , for some constant 𝐶 > 0 independent of 𝑁, Δ𝑡, and 𝑘 (provided Δ𝑡 is suﬃciently small so that 1−𝐶2 Δ𝑡 > 0). Applying the discrete Grönwall inequality and using the shared initial condition 𝑋0∗ = 𝑋0𝑁 (i.e. 𝜖𝑡0 = 0), we obtain ( ) 𝜖𝑡𝑘 ≤ 𝐶 ′ 𝑡𝑘 Δ𝑡2 + 𝑁 −1∕2 + 𝜀 , for all 𝑘 with 𝑡𝑘 ≤ 𝑇 , where 𝐶 ′ > 0 is a constant depending only on 𝑇 and the regularity constants of . This proves the claimed bound (23) and completes the proof.

A.3. Proof of Theorem 3 Assumption 2. 1.  (𝜌𝑎 ),  (𝜌𝑏 ) < +∞. 2. For ∀𝑇 ≥ 0 and every path 𝑝 ∈ 𝐴𝐶𝜌𝑎 ,𝜌𝑏 ,𝑇 , the map 𝑡 ↦  (𝑝𝑡 ) is absolutely continuous on [0, 𝑇 ] and satisﬁes ⟨ ⟩ 𝑑  (𝑝𝑡 ) = ∇d  (𝑝𝑡 ), 𝜕𝑡 𝑝𝑡 𝑑𝑡 −1,𝑝𝑡 C. Liu and X. Zhou: Preprint submitted to Elsevier

for a.e. 𝑡 ∈ (0, 𝑇 ). Page 32 of 39

Generative Wasserstein Gradient Path Method

Lemma 3. Under the assumptions 2 ,if 𝑆𝑇 [𝑝] < +∞, then ‖𝜕𝑡 𝑝𝑡 ‖−1,𝑝𝑡 ∈ 𝐿2 (0, 𝑇 ),

‖∇d  (𝑝𝑡 )‖−1,𝑝𝑡 ∈ 𝐿2 (0, 𝑇 ),

and 𝑆𝑇 [𝑝] =

𝑇

1 1 ‖𝜕𝑡 𝑝𝑡 ‖2−1,𝑝 d𝑡 + ‖∇d  (𝑝𝑡 )‖2−1,𝑝 d𝑡 +  (𝜌𝑏 ) −  (𝜌𝑎 ). 𝑡 𝑡 ∫ 2 0 2 ∫0

(A.3.1)

PROOF. Set 𝑢𝑡 ∶= 𝜕𝑡 𝑝𝑡 ,

𝑣𝑡 ∶= ∇d  (𝑝𝑡 ).

Since 𝑆𝑇 [𝑝] < +∞, the deﬁnition of 𝑆𝑇 yields ‖𝑢𝑡 + 𝑣𝑡 ‖2−1,𝑝 ∈ 𝐿1 (0, 𝑇 ). 𝑡

Moreover, by absolute continuity of 𝑡 ↦  (𝑝𝑡 ) and the chain rule, 𝑇

⟨𝑣𝑡 , 𝑢𝑡 ⟩−1,𝑝𝑡 ∈ 𝐿1 (0, 𝑇 ),

∫0

⟨𝑣𝑡 , 𝑢𝑡 ⟩−1,𝑝𝑡 d𝑡 =  (𝜌𝑏 ) −  (𝜌𝑎 ).

Using the pointwise identity ‖𝑢𝑡 ‖2−1,𝑝 + ‖𝑣𝑡 ‖2−1,𝑝 = ‖𝑢𝑡 + 𝑣𝑡 ‖2−1,𝑝 − 2⟨𝑣𝑡 , 𝑢𝑡 ⟩−1,𝑝𝑡 , 𝑡

𝑡

we conclude that ‖𝑢𝑡 ‖2−1,𝑝 + ‖𝑣𝑡 ‖2−1,𝑝 ∈ 𝐿1 (0, 𝑇 ). 𝑡

𝑡

Since both terms on the left-hand side are nonnegative, each belongs to 𝐿1 (0, 𝑇 ) separately, i.e. ‖𝜕𝑡 𝑝𝑡 ‖−1,𝑝𝑡 ∈ 𝐿2 (0, 𝑇 ),

‖∇d  (𝑝𝑡 )‖−1,𝑝𝑡 ∈ 𝐿2 (0, 𝑇 ).

Integrating the same identity and using the chain rule again gives 𝑇

∫0

‖𝑢𝑡 + 𝑣𝑡 ‖2−1,𝑝 d𝑡 = 𝑡

𝑇

‖𝑢𝑡 ‖2−1,𝑝 d𝑡 +

∫0

𝑡

𝑇

∫0

( ) ‖𝑣𝑡 ‖2−1,𝑝 d𝑡 + 2  (𝜌𝑏 ) −  (𝜌𝑎 ) , 𝑡

which is exactly (A.3.1). Lemma 4. Under the assumptions 2, let 𝑎𝑝 (𝜏) ∶= ‖𝜕𝜏 𝑝𝜏 ‖−1,𝑝𝜏 ,

𝑏𝑝 (𝜏) ∶= ‖∇d  (𝑝𝜏 )‖−1,𝑝𝜏 ,

⟨ ⟩ 𝑐𝑝 (𝜏) ∶= ∇d  (𝑝𝜏 ), 𝜕𝜏 𝑝𝜏

−1,𝑝𝜏

̂ < +∞, then 𝑐𝑝 ∈ 𝐿1 (0, 1), ∫ 𝑐𝑝 (𝜏) d𝜏 =  (𝜌𝑏 ) −  (𝜌𝑎 ), and 𝑎𝑝 𝑏𝑝 ∈ 𝐿1 (0, 1). In particular, If 𝑆[𝑝] 0 1

̂ = 𝑆[𝑝]

∫0

𝑎𝑝 (𝜏)𝑏𝑝 (𝜏) d𝜏 +

∫0

𝑐𝑝 (𝜏) d𝜏.

PROOF. Set 𝑔𝑝 (𝜏) ∶= 𝑎𝑝 (𝜏)𝑏𝑝 (𝜏) + 𝑐𝑝 (𝜏). By Cauchy-Schwarz, 𝑐𝑝 (𝜏) ≥ −𝑎𝑝 (𝜏)𝑏𝑝 (𝜏)

for a.e. 𝜏 ∈ (0, 1),

C. Liu and X. Zhou: Preprint submitted to Elsevier

Page 33 of 39

Generative Wasserstein Gradient Path Method

̂ < +∞, the deﬁnition of 𝑆̂ implies that 𝑔𝑝 ∈ 𝐿1 (0, 1). On the other hand, the absolute hence 𝑔𝑝 (𝜏) ≥ 0 a.e. Since 𝑆[𝑝] continuity of 𝜏 ↦  (𝑝𝜏 ) and the chain rule give 1

𝑐𝑝 ∈ 𝐿1 (0, 1),

∫0

𝑐𝑝 (𝜏) d𝜏 =  (𝜌𝑏 ) −  (𝜌𝑎 ).

Therefore 𝑎𝑝 𝑏𝑝 = 𝑔𝑝 − 𝑐𝑝 ∈ 𝐿1 (0, 1), ̂ follows. and the claimed representation of 𝑆[𝑝] PROOF (PROOF OF T HEOREM 3). We prove the two inequalities separately. Step 1. Fix 𝑇 > 0 and 𝑝 ∈ 𝐴𝐶𝜌𝑎 ,𝜌𝑏 ,𝑇 . If 𝑆𝑇 [𝑝] = +∞, the claim is immediate. Assume 𝑆𝑇 [𝑝] < +∞. By Lemma 3, ‖𝜕𝑡 𝑝𝑡 ‖−1,𝑝𝑡 , ‖∇d  (𝑝𝑡 )‖−1,𝑝𝑡 ∈ 𝐿2 (0, 𝑇 ). Deﬁne 𝑝̃𝜏 ∶= 𝑝𝑇 𝜏 on [0, 1]. Then 𝑝̃ ∈ 𝐴𝐶𝜌𝑎 ,𝜌𝑏 ,1 and, with 𝑡 = 𝑇 𝜏, 𝜕 𝑡 𝑝𝑡 =

1 𝜕 𝑝̃ 𝑇 𝜏 𝜏

for a.e. 𝜏 ∈ (0, 1).

Substituting this relation into (10) gives 𝑆𝑇 [𝑝] =

1 ‖1 ‖2 d𝜏 𝑇 ‖ 𝜕𝜏 𝑝̃𝜏 + ∇d  (𝑝̃𝜏 )‖ ‖ ‖−1,𝑝̃𝜏 ∫ 2 0 𝑇 1

1 𝑇 ‖𝜕𝜏 𝑝̃𝜏 ‖2−1,𝑝̃ d𝜏 + ‖∇d  (𝑝̃𝜏 )‖2−1,𝑝̃ d𝜏 + 𝜏 𝜏 ∫0 2𝑇 ∫0 2 ∫0 √ 𝐴 𝐵𝑇 = + + 𝐶 ≥ 𝐴𝐵 + 𝐶, 2𝑇 2 =

1⟨

∇d  (𝑝̃𝜏 ), 𝜕𝜏 𝑝̃𝜏

⟩ −1,𝑝̃𝜏

d𝜏

where 𝐴 ∶=

∫0

‖𝜕𝜏 𝑝̃𝜏 ‖2−1,𝑝̃ d𝜏,

𝐵 ∶=

𝜏

∫0

‖∇d  (𝑝̃𝜏 )‖2−1,𝑝̃ d𝜏, 𝜏

𝐶 ∶=

1⟨

∫0

∇d  (𝑝̃𝜏 ), 𝜕𝜏 𝑝̃𝜏

⟩ −1,𝑝̃𝜏

d𝜏.

By Cauchy-Schwarz, √

𝐴𝐵 ≥

∫0

‖𝜕𝜏 𝑝̃𝜏 ‖−1,𝑝̃𝜏 ‖∇d  (𝑝̃𝜏 )‖−1,𝑝̃𝜏 d𝜏.

Therefore ̂ 𝑝]. 𝑆𝑇 [𝑝] ≥ 𝑆[ ̃ Since this estimate holds for every 𝑇 > 0 and every 𝑝 ∈ 𝐴𝐶𝜌𝑎 ,𝜌𝑏 ,𝑇 , inf

inf

𝑇 >0 𝑝∈𝐴𝐶𝜌𝑎 ,𝜌𝑏 ,𝑇

𝑆𝑇 [𝑝] ≥

inf

𝑞∈𝐴𝐶𝜌𝑎 ,𝜌𝑏 ,1

̂ 𝑆[𝑞].

̂ = +∞, then Step 2. Fix 𝑝 ∈ 𝐴𝐶𝜌𝑎 ,𝜌𝑏 ,1 . If 𝑆[𝑝] inf

inf

𝑇 >0 𝑞∈𝐴𝐶𝜌𝑎 ,𝜌𝑏 ,𝑇

̂ 𝑆𝑇 [𝑞] ≤ 𝑆[𝑝]

̂ < +∞. is immediate. Assume 𝑆[𝑝] C. Liu and X. Zhou: Preprint submitted to Elsevier

Page 34 of 39

Generative Wasserstein Gradient Path Method

Since the integrand in (25) is positively homogeneous of degree one in 𝜕𝜏 𝑝𝜏 , the functional 𝑆̂ is invariant under absolutely continuous monotone reparameterizations. Thus, without loss of generality, we may assume that 𝑝 is parametrized with constant speed 𝑎(𝜏) ∶= ‖𝜕𝜏 𝑝𝜏 ‖−1,𝑝𝜏 ≡ 𝓁 for a.e. 𝜏 ∈ (0, 1), for some 𝓁 ≥ 0. If 𝓁 = 0, then 𝑝 is constant, hence 𝜌𝑎 = 𝜌𝑏 , and both sides of (24) vanish. We therefore restrict to the case 𝓁 > 0. Deﬁne ⟨ ⟩ 𝑏(𝜏) ∶= ‖∇d  (𝑝𝜏 )‖−1,𝑝𝜏 , 𝑐(𝜏) ∶= ∇d  (𝑝𝜏 ), 𝜕𝜏 𝑝𝜏 . −1,𝑝𝜏

By Lemma 4, we have 𝑎𝑏 ∈ 𝐿1 (0, 1), 𝑐 ∈ 𝐿1 (0, 1), and ̂ = 𝑆[𝑝]

∫0

𝑎(𝜏)𝑏(𝜏) d𝜏 +

∫0

𝑐(𝜏) d𝜏.

For 𝜀 > 0, deﬁne 𝜈𝜀 (𝜏) ∶=

𝑎(𝜏) , 𝑏(𝜏) + 𝜀

𝑡𝜀 (𝜏) ∶=

𝜏

∫0

𝜈𝜀 (𝑠) d𝑠.

Since 𝑎(𝜏) = 𝓁 > 0 a.e. and 𝑏(𝜏) + 𝜀 ≥ 𝜀, we have 0 < 𝜈𝜀 (𝜏) ≤

𝓁 𝜀

for a.e. 𝜏 ∈ (0, 1),

so 𝜈𝜀 ∈ 𝐿1 (0, 1) and 𝑡𝜀 is absolutely continuous and strictly increasing. Set 𝑇𝜀 ∶= 𝑡𝜀 (1), and let 𝜏𝜀 ∶ [0, 𝑇𝜀 ] → [0, 1] denote the inverse map of 𝑡𝜀 . Then 𝜏𝜀 is absolutely continuous and 𝜏̇ 𝜀 (𝑡) =

1 𝜈𝜀 (𝜏𝜀 (𝑡))

for a.e. 𝑡 ∈ (0, 𝑇𝜀 ).

Deﬁne 𝑝𝜀𝑡 ∶= 𝑝𝜏𝜀 (𝑡) ,

𝑡 ∈ [0, 𝑇𝜀 ].

Then 𝑝𝜀 ∈ 𝐴𝐶𝜌𝑎 ,𝜌𝑏 ,𝑇𝜀 . Moreover, ( ) 𝑎(𝜏)2 = 𝑎(𝜏) 𝑏(𝜏) + 𝜀 ∈ 𝐿1 (0, 1), 𝜈𝜀 (𝜏)

𝑏(𝜏)2 𝜈𝜀 (𝜏) =

𝑎(𝜏)𝑏(𝜏)2 ≤ 𝑎(𝜏)𝑏(𝜏) ∈ 𝐿1 (0, 1), 𝑏(𝜏) + 𝜀

and 𝑐 ∈ 𝐿1 (0, 1), so in particular 𝑆𝑇𝜀 [𝑝𝜀 ] < +∞. Using the chain rule and the change of variables 𝑡 = 𝑡𝜀 (𝜏), we obtain 𝑇

𝜀 1 ‖ 𝜀 ‖2 ‖𝜕𝑡 𝑝𝑡 + ∇d  (𝑝𝑡𝜀 )‖ 𝜀 d𝑡 ‖−1,𝑝𝑡 2 ∫0 ‖ ) 1( 1 2 𝑎(𝜏) 1 2 = + 𝑏(𝜏) 𝜈𝜀 (𝜏) d𝜏 + 𝑐(𝜏) d𝜏. ∫0 2 ∫0 𝜈𝜀 (𝜏)

𝑆𝑇𝜀 [𝑝𝜀 ] =

Substituting 𝜈𝜀 (𝜏) = 𝑎(𝜏)∕(𝑏(𝜏) + 𝜀) gives ( ) ( ) 1 𝑎2 1 𝑎𝑏2 𝜀2 𝑎 + 𝑏2 𝜈𝜀 = 𝑎(𝑏 + 𝜀) + = 𝑎𝑏 + . 2 𝜈𝜀 2 𝑏+𝜀 2(𝑏 + 𝜀) Therefore ̂ + 𝑅𝜀 [𝑝], 𝑆𝑇𝜀 [𝑝𝜀 ] = 𝑆[𝑝]

𝑅𝜀 [𝑝] ∶=

1 𝜀2 𝑎(𝜏) 1 d𝜏. 2 ∫0 𝑏(𝜏) + 𝜀

C. Liu and X. Zhou: Preprint submitted to Elsevier

Page 35 of 39

Generative Wasserstein Gradient Path Method

Since 0≤

𝜀2 𝑎(𝜏) ≤ 𝜀𝑎(𝜏) 𝑏(𝜏) + 𝜀

for a.e. 𝜏 ∈ (0, 1),

and 𝑎 ∈ 𝐿1 (0, 1), it follows that 0 ≤ 𝑅𝜀 [𝑝] ≤

𝜀𝓁 𝜀 𝑎(𝜏) d𝜏 = ←←←←←←→ ← 0. 2 ∫0 2 𝜀↓0

Hence ̂ lim 𝑆𝑇𝜀 [𝑝𝜀 ] = 𝑆[𝑝]. 𝜀↓0

Since 𝑝𝜀 ∈ 𝐴𝐶𝜌𝑎 ,𝜌𝑏 ,𝑇𝜀 for every 𝜀 > 0, inf

inf

𝑇 >0 𝑞∈𝐴𝐶𝜌𝑎 ,𝜌𝑏 ,𝑇

𝑆𝑇 [𝑞] ≤ 𝑆𝑇𝜀 [𝑝𝜀 ]

for all 𝜀 > 0.

Passing to the limit 𝜀 ↓ 0 yields inf

inf

𝑇 >0 𝑞∈𝐴𝐶𝜌𝑎 ,𝜌𝑏 ,𝑇

̂ 𝑆𝑇 [𝑞] ≤ 𝑆[𝑝].

Taking the inﬁmum over all 𝑝 ∈ 𝐴𝐶𝜌𝑎 ,𝜌𝑏 ,1 gives inf

inf

𝑇 >0 𝑞∈𝐴𝐶𝜌𝑎 ,𝜌𝑏 ,𝑇

𝑆𝑇 [𝑞] ≤

inf

𝑝∈𝐴𝐶𝜌𝑎 ,𝜌𝑏 ,1

̂ 𝑆[𝑝].

Combining Steps 1 and 2 proves (24).

A.4. Proof of Theorem 4

PROOF. Let 𝐽̂𝐾,∗ denote the minimum value of the Lagrangian problem (30), and 𝑆̂ 𝐾,∗ denote the minimum value of the Eulerian problem (29). First, we compare the minimum values of the two problems. For any admissible sequence of maps Φ, let 𝑝 be the induced density sequence deﬁned by 𝑝𝑘 = (Φ𝑘 )# 𝜌0 . By deﬁnition of the Wasserstein-2 distance, for each 𝑘, the coupling (Φ𝑘−1 , Φ𝑘 )# 𝜌0 between 𝑝𝑘−1 and 𝑝𝑘 yields 𝑊22 (𝑝𝑘−1 , 𝑝𝑘 ) ≤ 𝔼‖Φ𝑘 − Φ𝑘−1 ‖2 . Therefore, 𝑆̂ 𝐾 [𝑝] ≤ 𝐽̂𝐾 [Φ].

(A.4.1)

Taking the inﬁmum over all admissible Φ, we obtain 𝑆̂ 𝐾,∗ ≤ 𝐽̂𝐾,∗ . Conversely, let 𝑝∗ be a minimizer of (29), so that 𝑆̂ 𝐾 [𝑝∗ ] = 𝑆̂ 𝐾,∗ . Let Ψ𝑘 be the optimal transport map from 𝑝∗𝑘−1 to 𝑝∗𝑘 , and deﬁne the Lagrangian map sequence by Φ∗0 = id,

Φ∗𝑘 = Ψ𝑘 ◦Φ∗𝑘−1 = Ψ𝑘 ◦ ⋯ ◦Ψ1 .

Then (Φ∗𝑘 )# 𝜌0 = 𝑝∗𝑘 for every 𝑘, so Φ∗ is admissible for (30). By construction, each transport cost in 𝐽̂𝐾 [Φ∗ ] agrees exactly with the corresponding Wasserstein distance in 𝑆̂ 𝐾 [𝑝∗ ]. Hence 𝐽̂𝐾 [Φ∗ ] = 𝑆̂ 𝐾 [𝑝∗ ] = 𝑆̂ 𝐾,∗ . It follows that 𝐽̂𝐾,∗ ≤ 𝑆̂ 𝐾,∗ . C. Liu and X. Zhou: Preprint submitted to Elsevier

Page 36 of 39

Generative Wasserstein Gradient Path Method

Combining the two inequalities, we conclude that 𝐽̂𝐾,∗ = 𝑆̂ 𝐾,∗ . Proof of (a): Let Φ∗ be a minimizer of (30), and let 𝑝∗ be the induced density sequence deﬁned by 𝑝∗𝑘 = (Φ∗𝑘 )# 𝜌0 . By (A.4.1), 𝑆̂ 𝐾 [𝑝∗ ] ≤ 𝐽̂𝐾 [Φ∗ ] = 𝐽̂𝐾,∗ = 𝑆̂ 𝐾,∗ . Since 𝑆̂ 𝐾,∗ is the minimum value of the Eulerian problem, we also have 𝑆̂ 𝐾,∗ ≤ 𝑆̂ 𝐾 [𝑝∗ ]. Therefore, 𝑆̂ 𝐾 [𝑝∗ ] = 𝑆̂ 𝐾,∗ , and 𝑝∗ is a minimizer of (29). Proof of (b): Suppose 𝑝∗ is a minimizer of (29), so 𝑆̂ 𝐾 [𝑝∗ ] = 𝑆̂ 𝐾,∗ . Let Ψ𝑘 be the optimal transport map from ∗ 𝑝𝑘−1 to 𝑝∗𝑘 , and deﬁne Φ∗0 = id, Φ∗𝑘 = Ψ𝑘 ◦Φ∗𝑘−1 . Then, as shown above, 𝐽̂𝐾 [Φ∗ ] = 𝑆̂ 𝐾 [𝑝∗ ] = 𝑆̂ 𝐾,∗ = 𝐽̂𝐾,∗ . This proves that the composite map Φ∗ is a minimizer of the Lagrangian problem (30).

A.5. Proof of Theorem 5

𝐾 (Φ)|, which arises from the temporal discretization of the PROOF. We analyze the discretization error |𝐽̂(Φ) − 𝐽̂𝑁 geometric action and the Monte Carlo approximation of the 𝐿2 (𝜌0 )-norms. The continuous action functional is

𝐽̂(Φ) =

‖[𝑝𝜏 ](Φ(𝜏, ⋅))‖ 2 ‖𝜕𝜏 Φ(𝜏, ⋅)‖ 2 d𝜏. ‖𝐿 (𝜌0 ) ‖ ‖𝐿 (𝜌0 ) ∫0 ‖

(A.5.1)

Let 𝐿(𝜏) ∶= 𝐴(𝜏)𝐵(𝜏), where ‖ 𝐴(𝜏) ∶= ‖ ‖[𝑝𝜏 ](Φ(𝜏, ⋅))‖𝐿2 (𝜌 ) , 0

‖ 𝐵(𝜏) ∶= ‖ ‖𝜕𝜏 Φ(𝜏, ⋅)‖𝐿2 (𝜌 ) . 0

Let 𝜏𝑘 ∶= 𝑘Δ𝜏, with Δ𝜏 = 1∕𝐾, and Φ𝑘 ∶= Φ(𝜏𝑘 , ⋅). In accordance with (32), deﬁne 𝐴̂ 𝑁 𝑗 ∶= 𝑣𝑗 (Φ),

𝑑 (Φ) 𝐵̂ 𝑘𝑁 ∶= 𝑘 = Δ𝜏

(

𝑁

‖ 1 ∑‖ ‖ Φ𝑘 (𝑧𝑖 ) − Φ𝑘−1 (𝑧𝑖 ) ‖ ‖ 𝑁 𝑖=1 ‖ Δ𝜏 ‖ ‖

)1∕2 ,

where {𝑧𝑖 }𝑁 are i.i.d. samples drawn from 𝜌0 . Then 𝑖=1 𝐾 𝐽̂𝑁 (Φ) =

𝐾 𝐴̂ 𝑁 + 𝐴̂ 𝑁 ∑ 𝑘 𝑘−1

⋅ 𝐵̂ 𝑘𝑁 Δ𝜏, 2 𝑘=1 ⏟⏞⏞⏞⏞⏟⏞⏞⏞⏞⏟ ⏟⏟⏟ ≈𝐴(𝜏𝑘−1∕2 )

(A.5.2)

≈𝐵(𝜏𝑘−1∕2 )

with 𝜏𝑘−1∕2 ∶= (𝜏𝑘−1 + 𝜏𝑘 )∕2. We assume directly that 𝐴, 𝐵 ∈ 𝐶 2 ([0, 1]) with uniformly bounded derivatives. We also assume that Φ ∈ 3 𝐶 ([0, 1]; 𝐿2 (𝜌0 )). Furthermore, for ‖ Φ(𝜏𝑘 , ⋅) − Φ(𝜏𝑘−1 , ⋅) ‖ ‖ 𝐵̃𝑘 ∶= ‖ ‖ ‖ 2 , Δ𝜏 ‖ ‖𝐿 (𝜌0 ) we assume the Monte Carlo estimators satisfy the second-moment bounds [ ] [ ] 2 2 −1 ̃𝑘 |2 ≤ 𝐶 2 𝑁 −1 , 𝔼 |𝐴̂ 𝑁 𝔼 |𝐵̂ 𝑘𝑁 − 𝐵 𝑗 − 𝐴(𝜏𝑗 )| ≤ 𝐶𝐴 𝑁 , 𝐵 C. Liu and X. Zhou: Preprint submitted to Elsevier

Page 37 of 39

Generative Wasserstein Gradient Path Method

together with [ ] [ 𝑁 2] 2 ̂ sup 𝔼 |𝐴̂ 𝑁 𝑗 | + sup 𝔼 |𝐵𝑘 | ≤ 𝐶, 𝑗

𝑘

for some constant 𝐶 > 0 independent of 𝐾 and 𝑁. We ﬁrst estimate the deterministic quadrature error. Since 𝐿 ∈ 𝐶 2 ([0, 1]), the midpoint rule yields 𝜏𝑘

∫𝜏𝑘−1

𝐿(𝜏) d𝜏 = 𝐴(𝜏𝑘−1∕2 )𝐵(𝜏𝑘−1∕2 )Δ𝜏 + ((Δ𝜏)3 ),

(A.5.3)

where the constant depends only on a uniform bound on 𝐿′′ . We next estimate the two discrete factors. Term 𝐴̂ 𝑁 . By Taylor expansion around 𝜏𝑘−1∕2 , 𝐴(𝜏𝑗 ) = 𝐴(𝜏𝑘−1∕2 ) ±

Δ𝜏 ′ 𝐴 (𝜏𝑘−1∕2 ) + ((Δ𝜏)2 ), 2

𝑗 ∈ {𝑘 − 1, 𝑘},

and therefore 𝐴(𝜏𝑘 ) + 𝐴(𝜏𝑘−1 ) = 𝐴(𝜏𝑘−1∕2 ) + ((Δ𝜏)2 ). 2 It follows that 1∕2

|2 ⎞ ⎛ || 𝐴̂ 𝑁 + 𝐴̂ 𝑁 | 𝑘−1 ⎜𝔼 | 𝑘 − 𝐴(𝜏𝑘−1∕2 )|| ⎟ | ⎜ | 2 | ⎟ ⎝ | | ⎠

= ((Δ𝜏)2 ) + (𝑁 −1∕2 ).

(A.5.4)

Term 𝐵̂ 𝑘𝑁 . Since Φ ∈ 𝐶 3 ([0, 1]; 𝐿2 (𝜌0 )), the centered diﬀerence at the midpoint satisﬁes Φ(𝜏𝑘 , ⋅) − Φ(𝜏𝑘−1 , ⋅) = 𝜕𝜏 Φ(𝜏𝑘−1∕2 , ⋅) + ((Δ𝜏)2 ) Δ𝜏

in 𝐿2 (𝜌0 ).

Because the 𝐿2 (𝜌0 )-norm is 1-Lipschitz, we obtain 𝐵̃𝑘 = 𝐵(𝜏𝑘−1∕2 ) + ((Δ𝜏)2 ). Hence (

| |2 𝔼|𝐵̂ 𝑘𝑁 − 𝐵(𝜏𝑘−1∕2 )| | |

)1∕2

( ≤

2 | ̃𝑘 || 𝔼|𝐵̂ 𝑘𝑁 − 𝐵 | |

)1∕2

| | + |𝐵̃𝑘 − 𝐵(𝜏𝑘−1∕2 )| = ((Δ𝜏)2 ) + (𝑁 −1∕2 ). | |

(A.5.5)

Deﬁne 𝐿disc ∶= 𝑘

𝐴̂ 𝑁 + 𝐴̂ 𝑁 𝑘 𝑘−1 2

𝐵̂ 𝑘𝑁 .

Set Δ𝐴𝑘 ∶=

𝐴̂ 𝑁 + 𝐴̂ 𝑁 𝑘 𝑘−1 2

− 𝐴(𝜏𝑘−1∕2 ),

𝐴𝑚 ∶= 𝐴(𝜏𝑘−1∕2 ),

𝐵𝑚 ∶= 𝐵(𝜏𝑘−1∕2 ).

Then ( 𝑁 ) ̂𝑁 ̂ 𝐿disc 𝑘 − 𝐴𝑚 𝐵𝑚 = Δ𝐴𝑘 𝐵𝑘 + 𝐴𝑚 𝐵𝑘 − 𝐵𝑚 . C. Liu and X. Zhou: Preprint submitted to Elsevier

Page 38 of 39

Generative Wasserstein Gradient Path Method

Therefore, | | | | | | 𝔼|𝐿disc − 𝐴𝑚 𝐵𝑚 | ≤ 𝔼|Δ𝐴𝑘 𝐵̂ 𝑘𝑁 | + |𝐴𝑚 |𝔼|𝐵̂ 𝑘𝑁 − 𝐵𝑚 |. | 𝑘 | | | | | By Cauchy-Schwarz and (A.5.4), )1∕2 )1∕2 ( | ( | 𝔼|𝐵̂ 𝑘𝑁 |2 = ((Δ𝜏)2 ) + (𝑁 −1∕2 ), 𝔼|Δ𝐴𝑘 𝐵̂ 𝑘𝑁 | ≤ 𝔼|Δ𝐴𝑘 |2 | | while (A.5.5) gives )1∕2 ( | | |𝐴𝑚 |𝔼|𝐵̂ 𝑘𝑁 − 𝐵𝑚 | ≤ |𝐴𝑚 | 𝔼|𝐵̂ 𝑘𝑁 − 𝐵𝑚 |2 = ((Δ𝜏)2 ) + (𝑁 −1∕2 ). | | Hence | | 𝔼|𝐿disc − 𝐴(𝜏𝑘−1∕2 )𝐵(𝜏𝑘−1∕2 )| = ((Δ𝜏)2 ) + (𝑁 −1∕2 ). | 𝑘 |

(A.5.6)

Let 𝑘 ∶= 𝐿disc 𝑘 Δ𝜏 −

𝜏𝑘

∫𝜏𝑘−1

𝐿(𝜏) d𝜏

denote the local error on [𝜏𝑘−1 , 𝜏𝑘 ]. Then ( ) 𝑘 = 𝐿disc 𝑘 − 𝐴(𝜏𝑘−1∕2 )𝐵(𝜏𝑘−1∕2 ) Δ𝜏 +

( 𝐴(𝜏𝑘−1∕2 )𝐵(𝜏𝑘−1∕2 )Δ𝜏 −

𝜏𝑘

∫𝜏𝑘−1

) 𝐿(𝜏) d𝜏

| ∑ | 𝐾 𝔼|𝑘 | = 𝐾((Δ𝜏)3 ) + 𝐾(Δ𝜏𝑁 −1∕2 ) = ((Δ𝜏)2 ) + (𝑁 −1∕2 ) = (𝐾 −2 ) + (𝑁 −1∕2 ). (Φ) − 𝐽̂(Φ)| ≤ 𝔼|𝐽̂𝑁 | | 𝑘=1

This proves the claimed estimate.

C. Liu and X. Zhou: Preprint submitted to Elsevier

Page 39 of 39

Record · ID 10359 · SHA-256 f3bd362dcb68dac1
Conceptio Open Knowledge Archive — every document is proof-bundled with source, license, and retrieval metadata.