# Hessian-guided Perturbed Wasserstein Gradient Flows for Escaping Saddle Points

arXiv: https://arxiv.org/abs/2509.16974 · HTML: https://mlsanity.com/paper/2509.16974/

## Key facts

- Cited by 2 papers in the map (data through 14 Oct 2026).
- Builds on 28 of its 35 references that are in the map.
- First posted on arXiv 21 Sep 2025.

## Builds on (most cited first, 28 of 28)

- [Stein Variational Gradient Descent: A General Purpose Bayesian Inference Algorithm](https://mlsanity.com/paper/1608.04471/) (2016)
- [A Mean Field View of the Landscape of Two-Layers Neural Networks](https://mlsanity.com/paper/1804.06561/) (2018)
- [On the Global Convergence of Gradient Descent for Over-parameterized Models using Optimal Transport](https://mlsanity.com/paper/1805.09545/) (2018)
- [How to Escape Saddle Points Efficiently](https://mlsanity.com/paper/1703.00887/) (2017)
- [Escaping From Saddle Points --- Online Stochastic Gradient for Tensor Decomposition](https://mlsanity.com/paper/1503.02101/) (2015)
- [No Spurious Local Minima in Nonconvex Low Rank Problems: A Unified Geometric Analysis](https://mlsanity.com/paper/1704.00708/) (2017)
- [Maximum Mean Discrepancy Gradient Flow](https://mlsanity.com/paper/1906.04370/) (2019)
- [On the geometry of Stein variational gradient descent](https://mlsanity.com/paper/1912.00894/) (2019)
- [Gradient Descent Can Take Exponential Time to Escape Saddle Points](https://mlsanity.com/paper/1705.10412/) (2017)
- [SVGD as a kernelized Wasserstein gradient flow of the chi-squared divergence](https://mlsanity.com/paper/2006.02509/) (2020)
- [Mean-Field Langevin Dynamics: Exponential Convergence and Annealing](https://mlsanity.com/paper/2202.01009/) (2022)
- [Convex Analysis of the Mean Field Langevin Dynamics](https://mlsanity.com/paper/2201.10469/) (2022)
- [A Stein variational Newton method](https://mlsanity.com/paper/1806.03085/) (2018)
- [Transformers Learn Nonlinear Features In Context: Nonconvex Mean-field Dynamics on the Attention Landscape](https://mlsanity.com/paper/2402.01258/) (2024)
- [Accelerated Information Gradient flow](https://mlsanity.com/paper/1909.02102/) (2019)
- [Policy Optimization as Wasserstein Gradient Flows](https://mlsanity.com/paper/1808.03030/) (2018)
- [SSRGD: Simple Stochastic Recursive Gradient Descent for Escaping Saddle Points](https://mlsanity.com/paper/1904.09265/) (2019)
- [Information Newton's flow: second-order optimization method in probability space](https://mlsanity.com/paper/2001.04341/) (2020)
- [Convergence of mean-field Langevin dynamics: Time and space discretization, stochastic gradient, and variance reduction](https://mlsanity.com/paper/2306.07221/) (2023)
- [Vector-valued Gaussian Processes on Riemannian Manifolds via Gauge Independent Projected Kernels](https://mlsanity.com/paper/2110.14423/) (2021)
- [First-order Conditions for Optimization in the Wasserstein Space](https://mlsanity.com/paper/2209.12197/) (2022)
- [Regularized Stein Variational Gradient Flow](https://mlsanity.com/paper/2211.07861/) (2022)
- [Accelerated Flow for Probability Distributions](https://mlsanity.com/paper/1901.03317/) (2019)
- [Probability Functional Descent: A Unifying Perspective on GANs, Variational Inference, and Reinforcement Learning](https://mlsanity.com/paper/1901.10691/) (2019)
- [On Wasserstein Reinforcement Learning and the Fokker-Planck equation](https://mlsanity.com/paper/1712.07185/) (2017)
- [A Pontryagin Maximum Principle in Wasserstein Spaces for Constrained Optimal Control Problems](https://mlsanity.com/paper/1810.13117/) (2018)
- [Online Learning to Transport via the Minimal Selection Principle](https://mlsanity.com/paper/2202.04732/) (2022)
- [Understanding and Accelerating Particle-Based Variational Inference](https://mlsanity.com/paper/1807.01750/) (2018)

## Built on by (most cited first, 2 of 2)

- [FlashSinkhorn: IO-Aware Entropic Optimal Transport on GPU](https://mlsanity.com/paper/2602.03067/) (2026)
- [Trust-Region Optimization for Smooth Potential-Interaction Energies in Wasserstein Space](https://mlsanity.com/paper/2610.08883/) (2026)

## Abstract (arXiv)

Wasserstein gradient flow (WGF) is a common method to perform optimization over the space of probability measures. While WGF is guaranteed to converge to a first-order stationary point, for nonconvex functionals the converged solution does not necessarily satisfy the second-order optimality condition; i.e., it could converge to a saddle point. In this work, we propose a new algorithm for probability measure optimization, perturbed Wasserstein gradient flow (PWGF), that achieves second-order optimality for general nonconvex objectives. PWGF enhances WGF by injecting noisy perturbations near saddle points via a Gaussian process-based scheme. By pushing the measure forward along a random vector field generated from a Gaussian process, PWGF helps the solution escape saddle points efficiently by perturbing the solution towards the smallest eigenvalue direction of the Wasserstein Hessian. We theoretically derive the computational complexity for PWGF to achieve a second-order stationary point. Furthermore, we prove that PWGF converges to a global optimum in polynomial time for strictly benign objectives.