Finite-Sample Approximation of Hessian-Guided Perturbed Wasserstein Gradient Flows

Ryotaro Kawata; Atsushi Nitanda; Taiji Suzuki

cs.LGOct 7, 2026

Finite-Sample Approximation of Hessian-Guided Perturbed Wasserstein Gradient Flows

Authors: Ryotaro Kawata, Atsushi Nitanda, Taiji Suzuki

Organizations: The University of Tokyo · Agency for Science, Technology and Research (A*STAR) · Nanyang Technological University · RIKEN Center for Advanced Intelligence Project

Abstract

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

\kappa_{t}^{Y,N}=1

(positive curvature) if

t\in[0,1)\cup[2,3]

and

\kappa_{t}^{Y,N}=-1

(negative curvature) if

t\in[1,2)

. The factor

\mathfrak{C}_{t}^{Y}

measures the remaining effect of unit ODE forcing, while

\mathfrak{J}_{t}^{Y}

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

-\mu_{1}^{o}

denotes the reflected measure

(-\mathrm{Id})_{\#}\mu_{1}^{o}

. 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 $A_{\rm sp}=1$ in lemma A.34 (v)	—
Curvature-response mechanism	Nonlinear stability in theorem A.1 ; $N^{\upsilon}$ -scaled curvature-response bounds in corollary A.4	High-probability response bounds in lemma A.35 ; corollary A.4 applies with $\upsilon=\zeta=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

\theta<\pi

from

y

corresponds to the tangent vector

\mathrm{LOG}_{y}w

of length

\theta

;

\mathrm{EXP}_{y}

maps it back to

w

. Vectors are drawn with their tails at their base points. (b) Decreasing

s

from

1

0

C_{s}(w)

moves each point to its assigned pole and preserves its cap. (c) At the endpoint

w

q_{\star}(w)=\nabla_{\mathbb{S}}c_{\star}(w)=-\mathrm{LOG}_{w}\pi_{\star}(w)

points outward; its negative points toward the pole. The cost is

c_{\star}(w)=d_{\mathbb{S}}(w,\pi_{\star}(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

s_{t}

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

r_{t}

. Middle:

\lvert F_{X}(t)-F_{P}(t)\rvert

, where

F_{X}(t)=F(\widehat{\mu}_{t}^{N})

and

F_{P}(t)=F(\mu_{t})

. Bottom: the running maximum of the recorded ratio

s_{t}

, labeled

A_{\rm sp}(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(\mu_{t})

(black) and

F(\widehat{\mu}_{t}^{N})

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

Explore similar work

Wasserstein Residuals: Learning Gradient Flows from Population Dynamics
Jul 6, 2026Markus Heinonen, Yair Shenfeld, Ricardo Baptista +4Wasserstein Gradient Flows Wasserstein Distance
Basis Research Institute · Aalto University · Brown University +2
Stochastic Gradient Descent over P2
Sep 11, 2026Maria Oprea, Qin Li, Yunan YangStochastic Gradient Descent Stochastic Differential Equations
Center for Applied Mathematics, Cornell University · Department of Mathematics, University of Wisconsin–Madison, USA · Department of Mathematics, Cornell University, USA
Preservation of Log-Concavity and Convergence of Wasserstein-Fisher-Rao Gradient Flows
Sep 16, 2026Francesca Romana Crucinio, Sahani PathirajaWasserstein Gradient Flows Log-Concave Distributions
ESOMAS, University of Turin, Italy & Collegio Carlo Alberto, Turin, Italy · School of Mathematics & Statistics, UNSW Sydney, Australia

Jul 6, 2026, stat.MLJ/K move · Enter open · S save

Wasserstein Residuals: Learning Gradient Flows from Population Dynamics

Markus Heinonen, Yair Shenfeld, Ricardo Baptista, Daniel Waxman+3

Basis Research Institute · Aalto University · Brown University+2

Wasserstein Gradient Flows Wasserstein Distance Population Dynamics Gradient Trajectory Generation

Open paper AskPDF arXiv

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.

Abstract

Explore similar work

Wasserstein Residuals: Learning Gradient Flows from Population Dynamics

Stochastic Gradient Descent over P2

Preservation of Log-Concavity and Convergence of Wasserstein-Fisher-Rao Gradient Flows