Organizations: The University of Tokyo · Agency for Science, Technology and Research (A*STAR) · Nanyang Technological University · RIKEN Center for Advanced Intelligence Project
Wasserstein gradient flow extends gradient descent to probability measures. Its Hessian-guided perturbed variant (PWGF) adds Gaussian perturbations to escape saddle points in nonconvex problems. We investigate when its approximation by finitely many interacting particles remains accurate over growing time horizons. Our analysis retains the curvature accumulated along the population-driven reference path: negative curvature can amplify approximation errors, while subsequent positive curvature can damp their influence. This captures favorable scenarios in which temporary instability is compatible with accurate tracking over growing horizons. Under regularity assumptions and a prescribed common perturbation schedule, we prove particle and objective-value tracking bounds on a high-probability event for reference paths satisfying explicit conditions on accumulated curvature. To handle state-dependent Gaussian jumps, we construct a population-first coupling that preserves the reference particles' conditional independence and reduces jump errors to covariance comparison. We verify the conditions in a variance-plus-cosine model, where curvature recovery yields a growing-horizon tracking guarantee. We also establish local attraction, transverse descent, and positive second variation in two regions of a regularized matrix-factorization model, motivating a positive-negative-positive curvature pattern.
Figures & tables
Figure 1: Illustration of the curvature-response factors in the tracking bound. We assume that κtY,N=1 (positive curvature) if t∈[0,1)∪[2,3] and κtY,N=−1 (negative curvature) if t∈[1,2) . The factor CtY measures the remaining effect of unit ODE forcing, while JtY records a unit jump at t=1.5 , followed by amplification under negative curvature and decay under positive curvature.
Figure 2: Schematic local geometry in Proposition 4.1 : (A)–(B) the stationary target and its local attracting branch (i) ; (C) annular positive curvature (ii) ; (D) transverse descent (iii) ; and (E) positive curvature near the teacher (iv) . The diagram’s −μ1o denotes the reflected measure (−Id)#μ1o . Pale arrows depict curvature, not gradient-flow directions. The arrows do not assert a connecting trajectory.
Appendix figures & tables6 assets
Supplementary material from the paper’s appendix.
Appendix
Condition or conclusion
General comparison
Variance–cosine example
Matrix-factorization case study
Moment structure and weighted smoothness
Assumed in assumption A.1
Verified in lemma A.34 (i)
Separate product-model conditions in assumption 4.1
Gaussian jumps and schedule
Prescribed by assumption A.2
One prescribed common jump, verified in lemma A.34 (ii)
—
Moment growth and ODE regularity
Assumed in assumptions A.3 and A.4
Verified in lemma A.34 (iii)–(iv)
—
Non-spikiness
Assumed in assumption A.5
Verified with Asp=1 in lemma A.34 (v)
—
Curvature-response mechanism
Nonlinear stability in theorem A.1 ; Nυ -scaled curvature-response bounds in corollary A.4
High-probability response bounds in lemma A.35 ; corollary A.4 applies with υ=ζ=1/16
Population-level local geometry in proposition 4.1 : attraction, transverse descent, and positive spherical product pushforward second variation in two regions
Finite-particle conclusion
Conditional tracking and objective comparison
Growing-horizon conclusion in proposition A.1
Open
Appendix
Table 1: Relation of the model sections to the comparison theory. The matrix-factorization column records a separate population curvature-response case study; — denotes items outside its scope, and “Open” denotes the bridge to finite-particle tracking.
Figure 3: Spherical logarithms, exponentials, and pole contraction ( definitions B.3 and B.2 ). (a) A point w at distance θ<π from y corresponds to the tangent vector LOGyw of length θ ; EXPy maps it back to w . Vectors are drawn with their tails at their base points. (b) Decreasing s from 1 to 0 in Cs(w) moves each point to its assigned pole and preserves its cap. (c) At the endpoint w , q⋆(w)=∇Sc⋆(w)=−LOGwπ⋆(w) points outward; its negative points toward the pole. The cost is c⋆(w)=dS(w,π⋆(w))2/2 . Each circle is a great-circle cross section of the unit sphere; these geometric paths do not represent a WGF trajectory.
Figure 4: The instantaneous cubic ratio st in eq. 401 , with all three trials and five particle counts. Curves end when either system stops; near-zero denominators are omitted. The ratio measures the distribution of paired errors, not particle positions.
Figure 5: Independently chosen adaptive schedules, showing every system. Circles mark jumps whose decrease test passes; triangles mark the saved pre-jump candidates returned after a failed test; crosses mark stopping times. The open square marks an unmatched empirical attempt. Coupling by attempt index does not force jump times to coincide.
Figure 6: Common-time diagnostics for all 15 conditions, with one column per trial. Top: paired RMS error rt . Middle: ∣FX(t)−FP(t)∣ , where FX(t)=F(μtN) and FP(t)=F(μt) . Bottom: the running maximum of the recorded ratio st , labeled Asp(t) , rather than its instantaneous value in fig. 4 . Dots in the top row mark empirical jumps; dotted vertical lines mark population jumps. Curves end when either system stops; vertical segments connect recorded pre/post-jump values.
Figure 7: Objective values F(μt) (black) and F(μtN) (colored) for every trial and particle count. Each curve continues to that system’s own stopping time. Stars mark the saved pre-jump candidates at their candidate times, rather than the final executed states. The minimum of eq. 397 is zero.
Reconstructing population dynamics is a central problem in the physical and data sciences. Often, the dynamics are modeled as a Wasserstein gradient flow (WGF): a curve of distributions driven by an energy functional. Though there are multiple mathematical characterizations of a WGF, the dominant algorithmic approach relies on the Jordan--Kinderlehrer--Otto (JKO) scheme. JKO-based methods are inflexible to time discretisation and require solving costly optimal transport problems. We take a residual approach, enforcing the continuity equations via a non-negative loss function whose minimum is the WGF. Combined with a data-fitting divergence, this gives a single global objective. This perspective unifies several existing methods and leads to a new particle-based method, stitching, that is simulation-free and robust to large gaps between observations. We demonstrate that the stitching method achieves state-of-the-art performance across trajectory inference benchmarks. For code see github.com/BasisResearch/wasserstein-residuals.
Markus Heinonen, Yair Shenfeld, Ricardo Baptista +4
Basis Research Institute · Aalto University · Brown University +2
Stochastic gradient descent (SGD) admits diffusion approximations that replace the complicated randomness of stochastic gradients by Gaussian noise, providing a powerful tool for understanding its dynamics and long-time behavior. We investigate whether an analogous approximation principle holds for optimization over probability measures, where the objective is a functional defined on the Wasserstein space P2. The nonlinear geometry and infinite-dimensional nature of P2 prevent a direct extension of the classical Euclidean theory. Using Lions differentiability, we lift the problem to a linear Hilbert space, where higher-order differential calculus becomes available. We then construct a Gaussian random-field approximation whose velocity field matches the mean and covariance of the original stochastic gradient. By exploiting this moment matching through higher-order Taylor expansions, we show that the Gaussian approximation captures the SGD dynamics with second-order weak accuracy. Our result provides a rigorous foundation for replacing sample-driven randomness by analytically tractable Gaussian fluctuations in stochastic optimization over probability measures.
Maria Oprea, Qin Li, Yunan Yang
Center for Applied Mathematics, Cornell University · Department of Mathematics, University of Wisconsin–Madison, USA · Department of Mathematics, Cornell University, USA
We study the convergence of Wasserstein-Fisher-Rao (WFR) gradient flows for sampling from probability distributions known up to a normalisation constant. By combining Wasserstein transport with Fisher-Rao birth-death dynamics, WFR flows balance exploration and selection. These flows have been recognised as a promising mechanism to accelerate convergence beyond Langevin dynamics. We show that for a class of strongly log-concave target distributions satisfying additional curvature conditions, WFR flows preserve strong log-concavity, in contrast to Wasserstein flows which enjoy this property only in the Gaussian setting. Exploiting this result, we derive explicit non-asymptotic convergence rates for the symmetrised Kullback-Leibler divergence, without requiring a warm-start as required in current estimates. In particular, we show that the convergence rate decomposes additively into Wasserstein and Fisher-Rao contributions, thereby confirming a recent conjecture within this setting. These results provide refined convergence guarantees and further develop the theoretical foundations of WFR gradient flows for sampling and Bayesian inference.
Francesca Romana Crucinio, Sahani Pathiraja
ESOMAS, University of Turin, Italy & Collegio Carlo Alberto, Turin, Italy · School of Mathematics & Statistics, UNSW Sydney, Australia