Title: Hessian-guided Perturbed Wasserstein Gradient Flows for Escaping Saddle Points
Venue: NeurIPS 2025
Accepted: yes
Keywords: probability measure optimization

Wasserstein gradient flow

non-convex optimization
Abstract: 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, \emph{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.


Peer reviews:
comments: SUMMARY:
The authors propose a novel algorithm perturbed Wasserstein Gradient flow (PWGF) to address the issue of how to escape saddle points in probability distribution optimization problems using Hessian perturbation. The idea is that, in order to escape a saddle point, it is more efficient to go along the direction of the smallest eigenvalue. Under reasonable assumptions, for problems satisfying a strict benignity propert, the authors establish polynomial-time convergence to second-order stationary points, both in continuous and discrete setups.

STRENGTHS AND WEAKNESSES:
Strengths:

- The paper is well written and presents a detailed theoretical framework.

- The convergence analysis is technically sound. 

Weaknesses:

- Literature on the benefit of the Hessian-based saddle-point escape method, even in the Euclidean spaces, is missing. 

- Although the authors provide examples of non-convex objective functions that exhibit strict benignity, they are far from practice.

QUESTIONS:
- Can the authors provide experiments demonstrating the advantage of Hessian-based PWGF on realistic tasks that do not necessarily satisfy the strict benignity condition?

I will raise my score if more experiments can be provided.

LIMITATIONS:
No
TITLE: Official Review
CLARITY: 3
IMPACT: 2
ORIGINALITY: 2
DATE: 2025-06-26

comments: SUMMARY:
This paper introduces Perturbed Wasserstein Gradient Flow (PWGF), an algorithm for optimizing nonconvex functionals over probability measures. PWGF enhances standard Wasserstein gradient flow by introducing a perturbation with a Hessian-based kernel and thus efficiently escaping saddle points. The authors also provide the computational complexity of the proposed PWGF method for finding the ε-second-order stationary point.

STRENGTHS AND WEAKNESSES:
**strength**:
1. A rigorous convergence analysis is provided. In particular, the authors show the computational complexity of the proposed PWGF method for finding the ε-second-order stationary point, under standard smoothness and strict-saddle assumptions.
2. The writing is clear and well-structured.

**weakness**:
1. The proposed algorithm requires computing Hessian or Hessian-vector products for each particle or for the distribution, which is computationally expensive.
2. The numerical experiments are only on the loss functional of in-context learning of Transformers. A few more complex examples or ablation studies would broaden the practical impact.

QUESTIONS:
I am not very familiar with Wasserstein Gradient Flow, as my background is primarily in continuous optimization. I have a question regarding the algorithm's design: Given that it frequently computes the Hessian, have the authors considered incorporating second-order methods to leverage this information for improved optimization performance?

LIMITATIONS:
See weaknesses.
TITLE: Official Review
CLARITY: 3
IMPACT: 2
ORIGINALITY: 2
DATE: 2025-07-02

comments: SUMMARY:
This paper introduces Perturbed Wasserstein Gradient Flow (PWGF), an algorithm for optimizing non-convex functionals over the space of probability measures. PWGF augments the usual Wasserstein gradient flow by injecting Hessian-guided Gaussian perturbations near saddle points, ensuring efficient escape and convergence to second-order stationary points. The main convergence result shows that PWGF reaches an $(\epsilon,\delta)$-second-order stationary point in $\tilde{O}(1/\epsilon^2 + 1/\delta^2)$ time steps with high probability, and under a strict “benignity” assumption, it further converges to a global optimum.

STRENGTHS AND WEAKNESSES:
Strengths

1. The paper provides a meticulous infinite-dimensional analysis, including (i) second-order optimality conditions in Wasserstein space (Prop. 3.2–3.3), (ii) a novel construction of the Hessian-based kernel, and (iii) a clear proof of polynomial-time escape from saddles.

2. Unlike isotropic noise, the Hessian-guided GP perturbations are directed along the most negative curvature, mirroring perturbed-gradient-descent strategies in Euclidean space.

Weakness:

1. It would be nice if the paper could touch base on some practical issues as detailed in the question/limitation box below.

QUESTIONS:
1. Can you provide matching lower bounds (even in simpler settings) to argue that the derived rate is unimprovable?

2.  Is there a way to extend the analysis to the N-particle system, quantifying how the error between particle flow and the true flow scales with N and time?

3. Given the cost of sampling from the Gaussian process, are there efficient approximations that preserve the escape guarantees?

LIMITATIONS:
Yes.
TITLE: Official Review
CLARITY: 4
IMPACT: 3
ORIGINALITY: 3
DATE: 2025-07-03

comments: SUMMARY:
The paper performs a theoretical analysis of escaping saddle points via perturbation in Wasserstein-2 space analogous to Jin et al. 2017. The key here is selecting a hessian guided kernel in the gaussian process.

STRENGTHS AND WEAKNESSES:
Strengths

The analysis is very neat and the writings are very structural. Although I cannot check every proof, I think key theoretical results are clear and valid.


Weaknesses

Just like WGF, from an algorithmic perspective, this is conceptual rather than empirical. The specific implementation depends a lot on what functional you use and how well it is estimated. I would suggest adding more empirical demonstrations if possible just to show the potential empirical values. Maybe using MMD as the functional could be easier.

QUESTIONS:
1. For the discrete time analysis, what if we use backward euler instead of forward as in JKO-scheme? What would the algorithm be like and can we get similar convergence?

2. Why is it important to escape saddle points in WGF? Is there a particular practical scenario that this makes a major issue or do you have any intuition on when saddle points occur frequently?

3. Do you have any intuition on the choice of the kernel? Could you explain maybe by showing the analogy in euclidean space?

LIMITATIONS:
yes
TITLE: Official Review
CLARITY: 4
IMPACT: 3
ORIGINALITY: 3
DATE: 2025-07-23