Finite-Sample Approximation of Hessian-Guided Perturbed Wasserstein Gradient Flows
Ryotaro Kawata, Atsushi Nitanda, Taiji Suzuki
cs.LG
Oct 7, 2026 · v1
stat.ML
TL;DR
An accompanying Lean 4 development of 4,274 declarations formalizes the paper's finite-particle tracking theorems, the variance-plus-cosine example, and the local-geometry results, with an axiom audit.
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.
Problem
Hessian-guided perturbed Wasserstein gradient flow (PWGF) escapes saddle points in nonconvex optimization over probability measures, but existing guarantees cover only the population flow. The question is when a finite system of interacting particles accurately tracks the flow over growing time horizons.
Approach
The analysis tracks the curvature accumulated along the population reference path using a curvature propagator that amplifies or damps errors. A population-first coupling preserves the conditional independence of the reference particles and reduces Gaussian jump errors to a comparison of covariances. The conditions are verified in a variance-plus-cosine model, and the local geometry of a regularized matrix-factorization model is studied. The main results are formalized in Lean 4.
Results
Particle and objective-value tracking bounds are proved on a high-probability event under a prescribed jump schedule. A growing-horizon guarantee holds in the variance-plus-cosine model, and local attraction, transverse descent, and positive second variation are shown for the matrix-factorization model. The Lean build covers 4,274 declarations in 661 modules; well-posedness of the flow and the numerical experiments are not formalized.