Hessian-guided Perturbed Wasserstein Gradient Flows for Escaping Saddle Points
Authors: Naoya Yamamoto
Juno Kim
Taiji Suzuki
Venue: NeurIPS 2025 poster ()
Keywords: probability measure optimization
Wasserstein gradient flow
non-convex optimization
Primary area: 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.

Hessian-guides Perturbed Wasserstein Gradient Flows for Escaping Saddle Points
Naoya Yamamoto
Juno Kim UC Berkeley junokim@berkeley.ccmm
The University of Tokyo
-amamoto-naaya25510 wcw.:ee2......................ee www.::/..............e. uct
Taiji Suzuki
The University of Tokyo, RIKEN AIP
caiji@mist... uu w-tokyooocccco ac.
Abstract
Wasserstein gradient flow (WGF) is common method to perform optimization over the space 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., could converge to a saddle point. In this work, we propose 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-basee baseed scheme. By pushing the measure forward along 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 Wasser- stein Hessian. We theoretically derive the computational complexity for PWGF achieve second-order stationary point. Furthermore, prove that PWGF converges to a global optimum in polynomial time for strictly benign objectives.
Introduction
We consider the general problem of optimizing probability measure: min F(J),, where (R(R)) denotes the space of Borel probability measures on Rd with finite second moment and F (RR + E. is given functional which is not necessarily convex. Optimization of measures appears extensively in machine learning and statistics either explicitly or implicitly, such as in sampling and variational inference (Jordan et al., 1998; Liu & Wang, 2016), generative models (Arbel et al., 2019; Chu et al., 2019) and training neural networks (Mei et al., 2018: Nitanda et al., 2022), and has garnered significant attention both theoretically and practically. To solve such problems, the Wasserstein gradient flow (WGF) is extensively employed:
(-v
(-v
SF (pe)tee Sf
Oz44 + V
0, (mL) C P2(RR).
WGF is the continuity equation of the velocity field V éF and can be interpreted as gradient descent dyc with respect to the 2- 2-Wasserstein metric (Jordan et al., 1998). Moreover, (1) is equivalent to the limit of first-order optimization algorithms such as gradient
continuous-time and infinite-partic descent, and serves as the foundation for numerous machine learning methods; see Appendix A for discussion of related works and applications. It is thus of paramount importance understand the convergence of WGF, and to develop algorithms which guarantee well-behaved solutions.
This work primarily conducted while the author the University of Tokyo and RIKEN AIP.
39th Conference on Neural Information Processing Systems (NcurIPS 2025).

---

Convergence of WGF. A prominent line of work on the convergence of WGF is the study of neural network optimization using mean-field theory (Mei al.. 2018). Meen-fieed theory models the evolution of an interacting particle system as the number of particles tends to infinity, such as the training dynamies of an anfinite-widtte neural network, through the WGF of the limiting distribution. For simple problem settings such as regression with two-layer networks, the corresponding loss functional is shown to be (linearly) convex, allowing for global convergence analysis of WGF (Chizat & Bach, 2018; Mei et al., 2018). Furthermore, Nitanda et al. (2022); Chizat (2022); Suzuki et al. (2023) obtained linear convergence under the log-Sobolev inequality (LSD) condition using mean-field Langevin dynamics. This adds isotropic noise via Brownian motion, corresponding an additional entropy regularization which effectively makes the objective strongly convex.
Non-convex objectives. While these works primarily rely on convexity, the vast majority of objectives arising in deep learning (such as the loss function for neural networks with three or more layers) are non-convex, even when lifted the space of measures. However, due the inherent difficulty of infiniteddimensional non-convex optimization, convergence guarantees for WGF in this regime has been extremely limited. Recently Kim & Suzuki (2024) studied transformer model combining mean-field neural network with linear attention for in-context learning, resulting in a non- convex optimization problem. Through landscape analysis, they showed that the objective possesses a desirable benignity property: all second-orde optima are either saddle points global optima. A similar result has been demonstrated for energy kernels such as the MMD functional (Boufadene Vialard, 2024). These findings suggest the necessity of probability measure optimization algorithms that account for second-order optimality. While the first-order optimality of WGF has been shown by Lanzetti et al. (2025), there are few studies that rigorously consider second-order conditions.
Perturbation methods. The main difficulty of non-convex optimization is due to the existence of saddle points. A first-order method with naive initialization may end up converging saddle point. Moreover, Du et al. (2017) showed that gradient descent can be significantly delayed near saddle points, taking exponential time to converge. For optimization in ditee-iimeniooal Euclidean space, Ge et al. (2015); Jin et al. (2017); Li (2019) proposed methods to overcome this issue by introducing perturbations in the vicinity of saddle points to 'fall off" the saddle and escape efficiently. These methods, often referred to as perturbed gradient descent (PGD), are particularly useful as they achieve second-order optimality in polynomial time while primarily relying on first-order techniques. In particular, PGD guarantees global convergence for problems satisfying a strict benignity property, such matrix factorization (Jin et al., 2017) and tensor decomposition (Ge al., 2015).
With these issues in mind, we ask the following question:
Can we develop a perturbative version of Wasserstein gradient flow for non-convex objectives which converges efficiently second-order optimal points?
Our Contributions. In this study, we propose a perturbative modification to WGF that efficiently avoids saddle points. Considering the tangent bundle structure of (2 induced by the Wasserstein distance, it is natural to extend the notion of perturbation in Euclidean space to measure space by defining perturbations of the drift function through randomly generated vector fields. Such a perturbative method was conjectured to improve convergence by Kim & Suzuki (2024), but without any theoretical guarantees. Our contributions are summarized as follows:
We propose a new implementable algorithm, perturbed Wasserstein gradient flow (PWGF), that guarantees second-order optimality for general smooth and non-convex functionals. Unlike the method proposed by Kim Suzuki (2024), which injects isotropic noise near saddles into the WGF, we guarantee improvement by directing the noise using the (Wasserstein) Hessian. Specifically,, PWGF pushes the measure along a random velocity field generated from a Gaussian process whose covariance is constructed from the Hessian of the objective.
We prove that PWGF effectively avoids saddle points and reaches second-order optimal points in time that depends polynomially on precision parameters, enabling the optimization of non-convex distributional objectives. Compared to the finite-dimensional setting (Li, 2019), the analysis requires a careful treatment of an infinite dimensional objective. For this purpose, we utilize techniques from Wasserstein geometry, optimal transport and the theory of Gaussian processes.

---

Organization. The paper is organized as follows. Section 2 provides theoretical preliminaries, supplemented in Appendix Second-order optimality conditions are presented in Section 3. In Section 4, we introduce the proposed PWGF algorithm. Section presents the convergence analysis, along with rough sketch of the proof. Numerical experiments are provided Appendix
Preliminaries
In this paper, consider the optimization problem (i F(u) over the space of probability distributions, where F ((R)) + R is real-valued lower-bounded functional defined on P2( This section introduces the problem of probability measure optimization and reviews key concepts in optimal transport and Wasserstein geometry. See Appendix for further details.
Notation. Let Id be the identity map on E The canonical projection to the ith coordinate is denoted by ps. The Euclidean inner product and operator norm are (t- |I-I|- The Frobenius norm |1-1p-- The of real-valued functions Rd that infinitely differentiable with compact support is C8 (R.). The inner product and norm operator norm in T.2 (p)) (s)) - -I (2 The trace norm an operator Il-)),,,, and the Hilbert-Schmidt norm is ||- ||-ll
Z indicates that an operator is positive semi-definite. The smallest eigenvalue of a compact operator T is denoted by A' The exponential of an operator denoted by ex or exp We use to denote big O notation ignoring logarithmic factors.
The set of all Borel probability measures on Rd with finite second moments is denoted by P2 (R.). and the subset of measures absolutely continuous with respect to Lebesgue measure is denoted by P5 (Rd). The Dirac measure on E R 82 E P2(R)). f denotes the pushforward of u P2 by a measurable map IRd -
2.1 Wasserstein Geometry
Definition 2.1 (Wasserstein metric). The 2- W-sssersteinn metric between μa, E (2( is defined as
/IL= (dxxdd TED(MM))
W2(u, ()( min
where L'(H, v) represents the of all transport plans (or the of couplings) of μu, v, that is, all joint distributions on Rd x Rd whose marginal distributions μ and We denote the set of all optimal transport plans by (1 p) and the optimal transport map by Tp-
fundamental dynamics in Wasserstein space is the continuity equation with velocity field et.
dep + (Dipe) 0, μt E P2( e
Intuitively, (3) describes how particle distribution μr evolves along vector field In particular, 5F
WGF (1) moves particles according to the vector field e = - the continuity equation with velocity field such that
8u (pt) (see Section 2.3). Moreover,
(Vp C. (.7())
Ut E Tan Pa (R.)
can be locally approximated by pushforward along d (Proposition B.9). Therefore the continuity equation is computationally approximated by the pushforward representation:
Heti + (Id + Atv, #
The absolute continuity of the curve b4 with respect to the Wasserstein distance is equivalent to satisfying (3) for some U E (pe) (Ambrosio et al. (2008), Theorem 8.3.1). In this sense, the continuity equation provides concept of differentiation consistent with the Wasserstein metric. For
further background on optimal transport theory, see Appendix B.
2.2 Wasserstein Gradient
The Wasserstein gradient is the fundamental quantity for first-order analysis in Wasserstein space.
3

---

Definition 2.2 (Wasserstein gradient). The Wasserstein gradient of F at p E P2(R2) ) is vector field F
IRd Rd such that for any v E (2((R") and (T,,,,)
F(D) - F(n)
V (()) x)m((ddy)) + (W2(),
The first variation also frequently appears in the context of probability measure optimization (cf. proximal Gibbs measure (Nitanda et al., 2022)).
Definition 2.3 (First variation). The first variation
P2((R.) Rd R is defined a functional
8y
satisfying for any P2(R").
/
6F (1, x) ())x)))))
F(u+ h(v 1)) dh h=0
naive computation might suggest that F V however, this not generally true without 5p
additional conditions. Nevertheless, we do not distinguish between the two, see Appendix B.2.
First-order optimality. The Wasserstein gradient allows to construct first-order approximations of functionals. Furthermore, Lanzetti et al. (2025) demonstrated two analogies to unite-dimensiona optimization. The first is that v 0 serves necessary condition for local optimality. The second is that F O becomes a sufficient condition for global optimality in the convex case. Based these considerations,
define the following.
Definition 2.4 (First-order stationary point). Supoose that functional F P2(R)) + R satisfies sufficient smoothness. We say that f E (2((Rd) is a first-order stationary point, if p satisfies V V μ.a.e.
2.3
Wasserstein Gradient Flow (WGF)
As counterpart of gradient descent in Euclidean space, the WGF in Wasserstein space is defined as
O,t4 + ((v) 0, (14) C P2(R.).
From previous observations, the direction e, (p,)(p) yields the steepest descent direction of the objective F. Furthermore, by the chain rule (Proposition C.3) it holds that
F(o) = - [V,
00
dt
This indicates that the WGF monotonically decreases the objective function. Indeed, the WGF can be interpreted as gradient descent method in the space of probability measures (Jordan al., 1998). Moreover, WGF becomes stationary iff the solution is at a first-order stationary point (Definition 2.4).
3 Second Order Optimality on Probability Space
In order study second-order behavior, define suitable class of sufficiently regular functionals. Definition 3.1 (Sufficient smoothness). functional F (P(( + R is sufficiently smooth
F admits (p(()--iteerrbbl Wasserstein gradient v at all μ E PR()).
F(,, further admits Wasserstein gradient V P2(Rd) x Rd Rdd - Rdxd and is differentiable with respect to the second coordinate for any E P2(Rd) and u-ace. x. Furthermore, (()) is L2 p)) integrable and μ-ess supl| VV ((, x))] oc.
Assumption 1. The objective F (R()) + R is a sufficiently smooth functional.
Building upon the discussion of first-order optimality, we extend the analysis to second-order conditions. For simplicity of notation, we define the following operators Hp, H', for I (RR)),
Hxf(z)) =
H²f(z)) VV (V.((),))((z).
We establish the following proposition.

---

Proposition 3.2. For F: (2((R)) + IR a sufficiently smooth functional. for all (3()))
d²
F((Id + hw) = (u, (H + HH,))) LP(N))
dhf
h=0
This demonstrates that the second order term of vector field perturbation is characterized by the operator Hμ + Hy For more general version, see Proposition C.7. In particular, when examining stability at first-order stationary points, we have ((,)( = O, which implies VV, (()))))) = 0, that is, H4 = 0. Consequently, the change in F due to perturbations along vector fields is determined solely by the integral operator Hμ up to second order. Note that this does not necessarily hold for Φ P2 (Rd); for further details, refer to Appendix C.
Next, paralleling the work by Lanzetti et al. (2025), we demonstrate that Hp ł O serves sufficient condition for local stability, establishing the analogy with second-order conditions in Euclidean space.
Proposition 3.3 (second-orde necessary condition). Let F: P2( - RR be sufficiently smooth. If μ4" E P2 (RR) is a local minimum of F, then it holds that H₁. č O.
Based on this observation, we define the second-order optimality condition for probability measures. Definition 3.4 (second-order stationary point). For F (2((R - R a sufficiently smooth functional and μ E (2
We say that is a second-order stationary point if is a first-order stationary point and H₂ > O.
We say that is a saddle point if μ is a first order stationary point and satisfies H₃ X O, i.e. the smallest eigenvalue of H,, is strictly negative: Amin H < 0.
Since we seek approximate solutions for F, we also define approximate second order stationary points and saddle points.
Definition 3.5 (approximate second-order stationary point). Suppose that F P2(RR)) + IR is sufficiently smooth and E PP2(R.).
We say that μ is an (e, 8) -ttationary point, if || (V S and AminH 2 -8.
We say that is an (e, 8) saddle point, if (V,F())) and Amin Hy < -8
While this definition is useful for convergence analysis, it is not fully appropriate in light of the second-order expansion (7). This is because the condition (()))l S E only implies that
V ((4) is small in the sense of T norm and does not indicate that F(u)) is close to zero Therefore, in that case, second order optimality should be determined by H₃ + H", rather than Hp- Consequently, this paper assumes that the L norm Wasserstein gradients being small implies that the supremum norm of their gradients is also small. Assumption 2. For any E (2((R)), it holds that
Ju-ess =ERd
R2 || VV R2IV
(8)
Under this assumption, a small | a (())l implies that u-esssup / F(,
[Hll) is also small, justifying the definition of (c, -)statioonyyy points.
3.1 Global Convergence for Strictly Benign Objectives
is known that certain non-convex optimization problems possess a desirable property called benignity, i.c all local minima must be global minima. In Euclidean spaces, examples include tensor decomposition (Ge et al., 2015; Jin et al., 2017). With our definitions of (approximate) second-order optimality, similar property can be considered for the Wasserstein space.
Definition 3.6 (Strict benignity). The functional (2(Rd - IR is said to be (8, atssstictlly benign if at least one the following conditions holds for any E (IR)))
1. V F())lE
5

---

2. Amin H < -6
3. W2( nO) S a for some global optima μ4".
We provide examples non-convex objective functions that exhibit strict benignity, For details on the properties of each objective function, refer to the appendix and the cited papers. Example (Matrix decomposition). This example is inspired by the fact that finiteddimensional tensor decomposition exhibits strict benignity. We use mean-field two-layer neural network
1n(=)= = = f aa =)
to learn rank-one matrix induced by the target measure μº, where is data input that follows certain distribution, the parameter is x = (a, w) E R and a is an activation function, which we set as the sigmoid function. The objective functional can be expressed as follows:
F(n) F.z E= (=) (2) [[
We defer detailed analysis of this objective to Appendix G.
Example (3-layer neural network). Consider three-layer neural network model consisting of a mean-field two-layer network followed by linear layer. Assuming realizability, the L loss with respect to a teacher network T h (z) can be written as follows:
F(u,T) Eg [11 hy
hu (=) (x)l1
(9)
In this case, the optimization problem with respect to T can be explicitly solved. Defining nº (TT x Idg # and E,., Ez (=) the optimal T satisfies T Exp D and (9)
k.p
reduces to the following probability measure optimization problem:
F(u) E. [[ (2) (=
(10)
Kim & Suzuki (2024) analyze this optimization problem and essentially show strict benignity, assuming E, is bounded away from degeneracy.
Example 3 (Coulomb MMD). The maximum mean discrepancy (MMD) with Coulomb kernel can be expressed as:
F(H) - / (μ 12) 1 ( (0) (drdy)
d-
This energy functional has been studied by Boufadene & Vialard (2024), who showed that any absolutely continuous stationary point must be global optimum. Therefore, F' becomes benign in regions that do not involve singular distributions.
Perturbed Wasserstein Gradient Flow
In this section, introduce our proposed algorithm that incorporates perturbations in measure space along random velocity fields to escape saddle points.
4.1 Perturbations in Wasserstein Space
In Euclidean spaces, perturbed gradient descent (PGD) is first-order optimization method capable of escaping saddle points efficiently by adding perturbations (Ge et al., 2015; Jin et al., 2017: Li, 2019). Despite relying solely first order information, this method effectively achieves second order optimality, making it highly practical especially for problems with known benign structure. A typical perturbation technique in Euclidean spaces involves adding vectors sampled uniformly from ball of small radius. Intuitively, such perturbations can uniformly explore all directions, making it likely to include the unstable direction corresponding the smallest eigenvalue of the Hessian, thereby quickly 'falling off' the saddle.
To extend this idea to the space of probability measures, two issues must be addressed: (1) how to induce perturbations Wasserstein space, and (2) whether the perturbation includes the unstable direction that maximally reduces the objective. For (1), Kim & Suzuki (2024) proposed introducing

---

Algorithm PWGF (continuous-time)
set hyperparameter 7/p 0 A ) Tihres = Ō(t). and Fihree = 0(6%)
initialize f and Ep = - -hiree for t > do
if VV ((e)))))))) S and t tp > fehres then E GP pw) ue (Id + mpE) tp t- t end if'
if' t Ep + Thhres and F(uep) F(Pe) & Fihres then
return 14g end if
Oep + (-V = 0
end for
perturbations via multivariate Gaussian process GP( K) with fixed positive semi-definite kernel K : Rd x Rd + Rdxd transforming μ. into (Id + n Leveraging vector fields to represent infinitesimal changes in probability measures is both natural and practical. However, they did not provide theoretical justification for the effectiveness of this method.
We resolve this issue and guarantee (2) by constructing measure-dependen kernel K K₃ based on the Wasserstein Hessian of the objective:
= J FV,
(11)
The integral operator defined by this kernel coincides with the squared Hessian operator H Specifi- cally, it holds for f E (1)) that
(x)())= = J
(I 2. u(dz) = ((.,, (()))
v)u(d2)) f(y))(dy)
= _ (v))
The Hessian-basse basiinn-baaeddde kernel K. symmetric, positive semi-definite, and meets the trace-class condition due to the integrability of V F(u). This means that K₁, satisfies the requirements for a Gaussian process kernel. More importantly, the Gaussian process is 'directed' by Hμ., ensuring that the perturbation is likely to include the direction corresponding to the smallest eigenvalue of H thereby achieving reduction in the objective (Proposition 5.4). Details are provided in Appendix D.
4.2
Algorithm of PWGF
Based on the above considerations, we propose perturbed Wasserstein gradient flow (PWGF) as method for solving non-convex optimization problems of probability measures. PWGF alternates between perturbing near saddle points using a Gaussian process with the kernel defined (11) and evolving via WGF when not near saddle points.
We present the specific algorithm Algorithm 1, including the saddle detection mechanism. In the space of probability measures, the Hessian H, is operator on I.2 (2))., and determining whether its smallest eigenvalue is below a certain threshold computationally challenging. Drawing on Jin et al. (2017), we propose practical criterion: Proposition 5.4 says that if perturbation is introduced at an (s, -saddee point, the objective function decreases with high probability after a certain period of WGF Consequently, by always introducing perturbations at first-order stationary points, we can determine whether a given stationary point is a saddle point based on the decrease on the objective within a fixed time threshold.
The dime-discretized version PWGF is provided in Algorithm 2. To implement PWGF numerically, the optimal probability measure is approximated by the ensemble average of N particles as μ Σソ= and the pushforward and descent steps are directly applied to each particle
(perturbation)
Ij + zj + 1/p {(x;),
(gradient descent)
xj xi - n/ F( 2xz).
7

---

Algorithm 2 PWGF (discrete-time)
set hyperparameter 7/p = Ō Ō ) - 0(1), Aehres ((t) and Fitres O(8%)
initialize y and kpp = kibres do 11 S and ht hip > ketree then
for 0.
if #
{ ~~ GO, K (ww)) 1(()) t (Id + 1p #1
kp k
end if
if = kp + lithres and F(u((t)) F(ult)) S Fthres then
return (k()
end if
p(k+1) + (Id-ryV F1 (((())
-
end for
Convergence Analysis
In this section, we provide the theoretical guarantee for the PWGF algorithm. In addition to the assumptions introduced the previous sections, we impose the following Lipschitz continuity of the Wasserstein gradient and Hessian.
Assumption 3. F satisfies the following smoothness:
VuF is L1 Li-iiscciitz iee. for any E v),
(...))))) s f || - >dxddd)
(12)
V2F is Lz-Lipschitz iee. for any V),,
JIV (1, Vi, y2 1| 32 (drddy) S L3 || xdrdyd
(13)
. VV fl is L3- Lipschitz, i.e. for any E l'H,
7 p--sssup.ere (V,
ov, L3 | (
(14)
Remark 5.1. Similar gradient and Hessian Lipschitz continuity assumptions are made the conver- gence analysis of perturbed gradient descent methods in Euclidean spaces (Ge al., 2015; Jin et al., 2017; Li, 2019). Additionally, the first-order Lipschitz assumption is similar to works on convex functionals (Chizat, 2022; Suzuki al., 2023).
5.1 Convergence Results for Continuous-time PWGF
The following is the main theorem this paper, asserting that PWGF terminates polynomial time and reaches an (e, 8) -taationary point with high probability.
Theorem 5.2. Let e, 8, Ç 0 be chosen such that (L2 La) < 82 and e, < Õ(1). Let the initial point be po E P2( and AF F(Po) inf peP F(R). For hyperparameters
lp δ 8² 5) Tehres Ō()). and Fihres Ō(88), PWGF halts after
t= o AF
)
1 8z
time steps and reaches an (e.oo-second-orrd stationary point with probability 1-
To prove Theorem 5.2, we present several supporting results. Detailed proofs are deferred to Appendix E. Lemma 5.3 is fundamental property of WGF, providing a lower bound the decrease in the objective for non-stationary points.
In the analysis of PGD Euclidean spaces. pe < 8 is often assumed, where p is the Lipschitz constant of the Hessian (Jin et al., 2017; Li, 2019). We adopt this assumption to the probabilistic measure space setting. where the Hessian Lipschitz constant becomes p Le La.

---

Lemma 5.3. For curve of probability measures μr following the WGF, the following holds:
F(UO) - F(14) f. 1
Proposition 5.4 the crucial step of our analysis, showing that perturbing an (E. -0-addle point using Gaussian process with kernel (11) enables WGF to move along the unstable direction and decrease the objective function. other words, allows us to efficiently escape the saddle point.
Proposition 5.4. Set
O(1) and let e, 6, Up: Tthres: Fehres be chosen as in Theorem 5.2.
Suppose j² E P2 (R)) satisfies V F(Fu)))llaaa < E and Xo == Xuin pt S -5. Generating
& &~ GO, kp) and setting po = (Id + mpE)tte as the initial point of the WGF, we have with probability
F(H3) F(UTT.)..))
Fihres.
Combining these results establishes convergence.
Proof of Theorem 5.2. Eet 6, <> be chosen arbitrarily chosen such that (L2 L3) S 82, and set C such that C' is polynomial in and up to logarithmic factors.
From Proposition 5.4, perturbations occur at most m = AF times. Thus, the probability of Fitres
(') < mç'. Setting 5" = & ensures that the failure after m perturbations is at most 1 1- - (1- -
algorithm reaches an (s, (second order stationary point with probability at least l
PWGF consists of gradient descent phase and an evaluation phase, where the decrease in the objective is assessed after applying a perturbation. Then we define the gradient descent phase as State 0, and the evaluation phase as Stute 1. Let To denote the total time State O, where (1))(()) 2 By Lemma 5.3, the decrease in the objective during State 0 is at least e AF Moreover, the total time T1 in State I is bounded as T1 S mTTures AFTaze
implying To < e2 Ö (() Hence, PWGF halts in To + Ti = Ó + ))
Fihres
5.2
Convergence Analysis for Discrete-time PWGF
We also prove convergence to second-order stationary point for the time-discretized PWGF (Algo- rithm 2).
Theorem 5.5. Let the initial point be Ho E P2()R) and denote AF= F(Po) "nf F(μ). Set 17 = O() and let E, 6, Tp. Tihres: Fihres be chosen as in Theorem 5.2. Then, discrete time PWGF halts after
k k= AF k
)) +))
81
steps and reaches (E,6)-second-ore stationary point with probability c-
The proof is similar the continuous- case; details are deferred Appendix E.2. Since second- order stationary points are global optima for a strictly benign objective, the convergence of PWGF to global solution is also guaranteed.
Corollary 5.6. Under the same setting as Theorem 5.5, discrete-time PWGF for (e, 6, ot)-strictly benign objective F halts after Ō ( + ft)) steps and reaches a-close to some global optima 4²; W2(u, p2) S a with probability -CC.
Conclusion
We proposed new method for non-convex probability optimization, perturbed Wasserstein gradient flows (PWGF), which alternates between perturbing near saddle points using a deessinnguidee Gaussian process and evolving via WGF. We have established that PWGF efficiently achieves second- order optimality with high probability. A potential avenue future work is to reduce computational cost by using stochastic approximation of the Hessian as the kernel of the Gaussian process, analogous stochastic gradients. Another direction to provide method for analyzing benignity of non-convex distributional objectives, thereby broadening the range of applications of PWGF.

---

Acknowledgments
NY and JK were partially supported by JST CREST (JPMICR2015). TS was partially supported by JSPS KAKENHI (24K0290S) and JST CREST (JPMICR2115) 15). This research is supported by the National Research Foundation, Singapore, Infocomm Media Development Authority under its Trust Tech Funding Initiative, and the Ministry of Digital Development and Information under the AI Visiting Professorship Programme (award number AIVP-2024-004), Any opinions, findings and conclusions or recommendations expressed in this material are those of the author(s) and do not reflect the views of National Research Foundation, Singapore, Infocomm Media Development Authority, and the Ministry of Digital Development and Information.
References
Ambrosio, L.., Gigli, N., and Savaré, G. Gradient flows: in metric spaces and in the space of probability measures. Springer Science Business Media, 2008
Arbel, M., Korba, A., SALIM, A., and Gretton, A. Maximum mean discrepancy gradient flow. In Advances in Neural Information Processing Systems, 2019.
Bonnet, B. Pontryagin Maximum Principle in Wasserstein spaces for constrained optimal control problems. ESAIM: Control, Optimisation and Calculus of Variations, 25:52, 2019.
Boufadene, S. and Vialard, F.-X. On the global convergence of Wasserstein gradient flow of the Coulomb discrepancy. arXiv preprint arXiv:23/2.00800, 2024.
Chen, Z., Fan, J., and Wang. K. Multivariate Gaussian processes: definitions, examples and applications. Metron, 81(2):81-191, 2023.
Chewi, S., Le Gouic, T., Lu, C., Maunu, T., and Rigollet, P. SVGD as kernelized Wasserstein gradient flow of the chi-squared divergence. Advances in Neural Information Processing Systems, 33:2098-2109, 2020.
Chizat, L. Mean-field Langevin Dynamics: Exponential Convergence and Annealing. arXiv preprint arXiv:2202.01009 2022.
Chizat, L. and Bach, F. On the Global Convergence of Gradient Descent for Over ever-parameterized Models using Optimal Transport. Advances in Neural Information Processing Systems, 31, 2018.
Chu, C., Blanchet, J. H., and Glynn, P. W. Probability functional descent: A unifying perspective on GANs, variational inference, and reinforcement learning. International Conference on Machine Learning, 2019.
Detommaso, G., Cui, T., Marzouk, Y., Spantini, A., and Scheichl, R. A Stein variational Newton method. Advances in Neural Information Processing Systems, 31, 2018
Du, S. S., Jin, C., Lee, J. D., Jordan, M. I., Singh, A., and Poczos, B. Gradient Descent Can Take Exponential Time to Escape Saddle Points. Advances in Neural Information Processing Systems, 30, 2017.
Duncan, A., Nüsken, N., and Szpruch, L. On the geometry of Stein variational gradient descent. Journal of Machine Learning Research, 24(56):1-39, 2023.
Ge, R. Huang, F., Jin, C., and Yuan, Escaping From Saddle Points -Online Stochastic Gradient for Tensor Decomposition. In Conference on Learning Theory, pp. 797-842. PMLR, 2015.
Ge, R., Jin, C., and Zheng, Y. No Spurious Local Minima in Nonconvex Low Rank Problems: A Unified Geometric Analysis. In International conference on machine learning, pp. 1233-1242. PMLR, 2017.
Guo, W., Hur, Y., Liang. T., and Ryan, C. Online Learning to Transport via the Minimal Selection Principle. In Conference on Learning Theory, pp. 1085-4109 PMLR, 2022.
Han, J. Wasserstein gradient descent for online learning. Available SSRN 4946060, 2024.
10

---

He, Y., Balasubramanian, K., Sriperumbudur, B. K., and Lu, J. Regularized Stein Variational Gradient Flow. Foundations of Computational Mathematics, pp. 1-59, 2024.
Hutchinson, M., Terenin, A., Borovitskiy, V., Takao, S., Teh, Y., and Deisenroth, M. Vector-valued Gaussian Processes Riemannian Manifolds via Gauge Independent Projected Kernels. Advances in Neural Information Processing Systems, 34:17 14:17660-17169, 2021.
Jin, C., Ge, R., Netrapalli, P., Kakade, S. M., and Jordan, M. I. How to Escape Saddle Points Efficiently. In International Conference Machine Learning, 2017.
Jordan, R., Kinderlchrer, D., and Otto, F. The Variational Formulation of the Fokker-Planck Equation. SIAM journal on mathematical analysis, 29(1):1-17, 1998.
Kent, C., Li, J., Blanchet, J., and Glynn. P. W. Modified Frank Wolfe in Probability Space. Advances in Neural Information Processing Systems, 34:14448- 14462, 2021.
Kim, and Suzuki, T. Transformers Learn Nonlinear Features In Context. ICLR 2024 Workshop Mathematical and Empirical Understanding of Foundation Models, 2024.
Lanzetti, N., Bolognani, S., and Dörfter, F. First-Order Conditions for Optimization in the Wasserstein Space. SIAM Journal on Mathematics of Data Science, 7(1):274-300, 2025.
Li, Z. SSRGD: Simple Stochastic Recursive Gradient Descent for Escaping Saddle Points. Advances in Neural Information Processing Systems, 32, 2019.
Liu, C., Zhuo, J., Cheng, P., Zhang, R., Zhu, and Carin, L. Accelerated First-order Methods on the Wasserstein Space for Bayesian Inference. stat, 1050:4, 2018.
Liu, Q. and Wang, D. Stein Variational Gradient Descent: A General Purpose Bayesian Inference Algorithm. Advances in Neural Information Prucessing Systems, 29, 2016.
Mei, S., Montanari, A., and Nguyen. P.-M. A mean field view of the landscape of two-layer neural networks. Proceedings the National Academy of Sciences, 115(33) 2018.
Nitanda, A., Wu, D., and Suzuki, T. Convex Analysis of the Mean Field Langevin Dynamics. In International Conference on Artificial Intelligence and Statistics, pp. 9741-9757. PMLR, 2022.
Richemond, P. H. and Maginnis, B. On Wasserstein Reinforcement Learning and the Lokker-Plancee Equation. arXiv preprint 17 185, 2017.
Santambrogio, F. Optimal Transport for Applied Mathematicians. Birkäuser, NY, 55(58-63):94, 2015.
Suzuki, T., Wu, D., and Nitanda, A. Mean-field Langevin dynamics: Time and space discretization, stochastic gradient, and variance reduction. Advances in Neural Information Processing Systems, 2023.
Taghvaei, A. and Mehta, P. Accelerated Flow for Probability Distributions. In International Conference on Machine Learning, 2019.
Wang, Y. and Li, W. Information Newton's 's flow: second-order optimization method in probability space. arXiv preprint arXiv:200J.0434 044341, 2020.
Wang, Y. and Li, W. Accelerated Information Gradient Flow. Journal of Scientific Computing, 90: 1-47, 2022.
Zhang, R., Chen, C., Li, C., and Carin, L. Policy Optimization Wasserstein Gradient Flows. In International Conference on Machine Learning, 2018.
11

---

NeurIPS Paper Checklist
1. Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification: The proposed PWGF algorithm is detailed in Section 4. The convergence
analysis is given in Section 5, Proof details are provided throughout the appendix.
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or
NA answer to this question will not be perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.
is fine to include aspirational goals as motivation as long as is clear that these goals not attained by the paper.
2. Limitations
Question: Does the paper discuss the limitations the work performed by the authors?
Answer: [Yes]
Justification:
See the conclusion, well Appendix B.2.
Guidelines:
The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors are encouraged to create separate "Limitations" section their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect how these assumptions might be violated in practice and what the implications would be.
The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect on the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution is low images are taken in low lighting. Or speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.
The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness.
While the authors might fear that complete honesty about limitations might used by reviewers as grounds for rejection, worse outcome might be that reviewers discover limitations that aren't 't acknowledged judgment and recognize that individual actions in favor of transparency play impor- tant role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.
the paper. The authors should use their best
3. Theory assumptions and proofs
Question: For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?
Answer: [Yes]
12

---

See Assumptions 1,2,3. All complete proofs are provided throughout the appendix.
Justification:
Guidelines:
The answer NA means that the paper does not include theoretical results.
All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated or referenced in the statement of any theorems.
The proof's can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged to provide a short proof sketch provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material.
Theorems and Lemmas that the proof relies upon should properly referenced.
4. Experimental result reproducibility
Does the paper fully disclose all the information needed to reproduce the main ex- perimental results of the paper to the extent that affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?
Question:
Answer: [Yes]
Justification:
Experimental details provided in Appendix F.
Guidelines:
The answer NA means that the paper does include experiments. If the paper includes experiments, No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
If the contribution is dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is novel architecture, describing the architecture fully might suffice, or if the contribution is a specific model and empirical evaluation, it may be necessary to either make possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g.. in the case of large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.
While NeurIPS does not require releasing code, the conference does require all submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) If the contribution is primarily a new algorithm, the paper should make it clear how reproduce that algorithm. (b) If the contribution is primarily the architecture clearly and fully.
model architecture, the paper should describe
(c) If the contribution is new model (e.g., a large language model), then there should either be way to access this model for reproducing the results or way to reproduce the model (e.g., with an open-source dataset instructions for how to construct the dataset).
(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome deseribe the particular way they provide for reproducibility, In the case of closed-source models, it may be that access to the model is limited in some way (e.g., registered users), but it should be possible for other researchers to have some path to reproducing or verifying the results.
5. Open access to data and code
Question: Does the paper provide open access to the data and code, with sufficient instruc- tions to faithfully reproduce the main experimental results, as described in supplemental material?
13

---

Answer: [NA]
Justification:
The conducted experiments are toy simulations.
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurIPS code and data submission guidelines (https://nips cc/
for more details.
While we encourage the release of code and data, we understand that this might not be possible, "No" is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (c.g., for new open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurIPS code and data submission guidelines (https: /nips. wwtppccccp/caeeecccccccccccceaae guwdee/e////ee /edeee/eeeebdmiai aassonoooicyyy for more details.
The authors should provide instructions on data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. If only subset of experiments are reproducible, they should state which ones are omitted from the script and why.
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
Providing as much information as possible supplemental material (appended the paper) is recommended, but including URLs to data and code is permitted.
6. Experimental setting/details
Does the paper specify all the training and test details (c.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary understand the results?
Question:
Answer: [Yes]
Justification: Experimental details are provided in Appendix F. Guidelines:
The answer NA means that the paper does include experiments.
The experimental setting should be presented in the core the paper to a level of detail that is necessary appreciate the results and make sense of them.
The full details can be provided either with the code, appendix, as supplemental material.
7. Experiment statistical significance
Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Question:
Answer: [Yes]
Justification: We conducted experiments under the same conditions, plotting the mean and standard deviation.
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results are accompanied by error bars, confi- dence intervals, statistical significance tests, least for the experiments that support the main claims of the paper.
The factors of variability that the enror bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall with given experimental conditions).
The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)
The assumptions made should be given (e.g., Normally distributed errors).
14

---

It should be clear whether the error bar is the standard deviation or the standard error of the mean.
is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out range (e.g. negative emror rates).
If error bars reported in tables plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables in the text.
8. Experiments compute resources
For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time of execution) needed reproduce the experiments?
Question:
Answer: [Yes]
Justification:
Experimental details are provided in Appendix F.
Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage.
The paper should provide the amount compute required for each of the individual experimental runs well estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary failed experiments that didn't make into the paper).
9. Code ethics
Question: Does the research conducted the paper conform, in every respect, with the NeurlPS Code of Ethics https :tap::://neuue /ttp///eurrpp.eeeee
Answer: [Yes]
Justification: The authors have checked that the research conforms with the NeurIPS Code Ethics.
Guidelines:
The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.
If the authors answer No, they should explain the special circumstances that require deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g., there a special consid- eration due to laws or regulations in their jurisdiction).
10. Broader impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [NA]
Justification:
The paper is primarily theoretical and no immediate societal impact is expected.
Guidelines:
The answer NA means that there no societal impact of the work performed.
If the authors answer NA or No, they should explain why their work has no societal impact why the paper does not address societal impact.
Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
15

---

The conference expects that many papers will be foundational research and not tied particular applications, let alone deployments. However, if there is direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology being used intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.
If there negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms monitor how system learns from feedback over time, improving the efficiency and accessibility of MI..
11. Safeguards
Does the paper describe safeguards that have been put in place for responsible release of data models that have high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Question:
Answer: [NA]
Justification: The paper is primarily theoretical.
Guidelines:
The answer NA means that the paper poses such risks.
Released models that have high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere usage guidelines or restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make best faith effort.
12. Lieenses for existing assets
Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper. properly credited and are the license and terms of use explicitly mentioned and properly respected?
Answer: [NA]
Justification: The paper does not use existing assets.
Guidelines:
The answer NA means that the paper does existing assets.
The authors should cite the original paper that produced the code package or dataset.
The authors should state which version of the asset is used and, if possible, include URL.
The name of the license (c.g., CC-BY 4.0) should be included for each asset.
For scraped data from a particular source (c.g., website), the copyright and terms of service of that source should be provided.
If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswi thcode. com datasets has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.
16

---

If this information is not available online, the authors are encouraged to reach out to the asset's creators.
13. New assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [NA]
Justification: The paper does not release new assets.
Guidelines:
The answer NA means that the paper does not release new assets.
Researchers should communicate the details of the dataset/code/mmoel as part of their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent obtained from people whose asset is used.
At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.
14. Crowdsoureing and research with human subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well details about compensation (if any)?
Answer: [NA]
Justification: The paper does not involve crowdsourcing nor research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Including this information in the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then as much detail as possible should be included the main paper.
According the NeurIPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional review board (IRB) approvals or equivalent for research with human subjects
Question:
Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?
Answer: [NA]
Justification:
The paper does not involve crowdsourcing nor research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Depending the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
We recognize that the procedures for this may vary significantly between institutions and locations, and expect authors to adhere to the NeurlPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if applicable), such the institution conducting the review.
16. Declaration of LLM usage
17

---

Question: Does the paper describe the usage of LLMs if it is important, original, or non-standard component of the core methods this research? Note that if the LLM used only for writing, editing, or formatting purposes and does not impact the core methodology, scientific rigorousness, or originality of the research, declaration is not required.
Answer: [NA]
Justification: The core method development in this research does not involve LLMs.
Guidelines:
The answer NA means that the core method development in this research does not involve LLMs any important, original, non-standard components.
Please refer to our LLM policy (https httppp////eripp....ooofeeeeeeeee c0tt.c/eeeeeeene/222//////////// for what should or should not be described.
18

---

A Applications of Wasserstein Gradient Flow
A.1 Mean-Field Analysis
Mean-field analysis can be applied to optimization problems where the objective function is expressed function of the ensemble average of some underlying functions. We consider the optimization problem
minimize
G
Σ ==1
kxzz)
s.t. 21,
EN E IR",
(15)
N
where Rd + IRd, G: Fy is computed follows:
Rd 7 RR and = 1 for simplicity. The gradient direction the variable
N
N h(zz)
G
Σ j=1
h(zz))
Vh(r,).
(16)
N
N
N
Taking the mean field limit N of (15), we lift this problem to the space of measures:
((fio)u(as)
minimize
s.t. 1 E P2(R2).
(17)
The Wasserstein gradient is computed as follows:
SF VaG(H) V 81 [ (u,))
=( ) Vh(z).
(18)
Consequently, from (16) and (18), the update rules for gradient descent applied to the original problem (15) and WGF applied to the lifted problem (17) are equivalent, up to constant scaling factor. However, properties differ significantly between the two formulations. For example, when F is the identity function, the original problem (15) is not necessarily linear, while the mean-field problem (17) is linear with respect to μ. Similarly, for commonly encountered loss functions with convex losses, (15) is not necessarily convex with respect to the variables 21,--.., but in the mean-field setting (17) becomes convex with respect to μ..
The properties of functions defined on the space of probability measures facilitate the design of WGF- based algorithms with improved performance. The works of Nitanda et al. (2022) and Chizat (2022) have explored the ability of optimization of mean-field Langevin dynamics (MFLD), an approach that augments WGF with Brownian noise, based on the convexity of the objective function. Similarly, Kim & Suzuki (2024) have focused the benignity of objective functions regarding in-context learning of certain transformer models and, leveraging this insight, have proposed a birth-death modification and also isotropic perturbation scheme. Our proposed algorithm contributes this line of research by providing the first convergence guarantees for strictly benign problems on the space of probability measures.
A.2 Additional Related Works
WGF has important applications in Bayesian inference, where posterior distributions are usually esti- mated from data via variational inference (VT). VI formulates posterior estimation an optimization particular, particle-based VI founded on the idea of evolving empirical measures formed by particles using WGF. This idea originates from the work of Jordan et al. (1998), which established a connection between diffusion processes and gradient descent in 2- 2-Wasserstein space with entropy regularization.
problem of the KL divergence.
Stein variational gradient descent (Liu Wang, 2016; He et al., 2024) circumvents computational difficulties in calculating descent directions through kernel methods. This approach can be seen as version of WGF where an integral operator acts on the descent direction (Chewi et al., 2020; Dunean et al., 2023). In this context, derivative methods such as Newton's 's method on the space of probability measures (Detommaso et al., 2018; Wang & Li, 2020) and accelerated methods (Liu et al., 2018: 8; Taghvaei & Mehta, 2019; Wang & Li, 2022) have also been proposed. Furthermore, optimization of measures also appears in contexts such as online optimization (Guo al., 2022; Han, 2024) and reinforcement learning (Richemond & Maginnis, 2017: Zhang et al., 2018).
19

---

B Wasserstein Geometry
B.1 Wasserstein Space
In this section, we provide basic aspects Wasserstein geometry and propositions used in this paper. We refer Ambrosio al. (2008) for a comprehensive review.
We assume that all curves fee E (2(( that appear in this section are absolutely continuous and satisfy the continuity equation with respect vector field 2 Tan P2((),) The existence condition for solutions to the continuity equation corresponding to vector field is provided, for example, by Ambrosio et al. (2008), Section 8.2.
Definition B.1 (Wasserstein distance). The 2- 2-sasserseeinn distance between two points p,, V P2(R2) defined as
/ || rdlddd)
W2(n, v)
inf TET(MM)
(19)
Here, v) represents the of all transport plans:
F(,,v) {> E P2 R² PI#t μt p2#t = =v}
Moreover, 7 E T(H, 2) is called the optimal transport plan, if the infimum in (19) is attained by y- The set all optimal transport plans is denoted by To (o(,, 1). Furthermore, a measurable map f: IRd Rd satisfies (Id f) E (w, v), then f is called the optimal transport map between jr and v. In this case, denote the optimal transport map T
W2 (2()) P2( R defines metric structure on P2 (IRd). Therefore, consider P2 (I as a metric space equipped with the W₂ metric. Convergence in the W₂ metric is equivalent to the weak convergence of measures plus uniform integrability of second moments (Ambrosio al., 2008).
An optimal transport plan is guaranteed to exist for any V.V. E P2(Rd), as stated in the following proposition.
Proposition B.2. For P. P2(R.), there exists an optimal transport plan between and v. That is, there exists 7 E (H,,v) such that
W)(,, -1l11
|10-- (la-vll
((d=yy))(dddy)) yddxdy)
Proof. Refer to Ambrosio et al. (2008), Chapter 6.
An optimal transport map, not a plan, does not necessarily exist. Specifically, there are cases where even transport maps themselves do not exist. For example, when d = 1, μ = ôg, and y no transport map exists, H# * for any T R R. However, The following proposition establishes that under certain conditions, the existence and uniqueness of an optimal transport map are guaranteed.
4
2.
Propesition B.3 (Brenier's theorem). For any μ (R( and P2(Rd), there exists optimal transport map TV Rd + Rdd Specifically, Ty satisfies
W2(,v /IITI [1T - ((dx).
Furthermore, the following hold:
(o( = ((Id x TH t,), that is, the unique optimal transport plan from μ to v is induced by TH-
The optimal transport map be expressed as Ty Vip, where ip is convex function defined p-almost everywhere.
If Py (IR) as well, then TV o TT² -===a Id p.a.e and Ty o T = ld v-ae.
Proof. See Ambrosio et al. (2008), Chapter 6.
20

---

P2(RR)) forms a geodesic metric space due to the existence of optimal transport plans and the pushforward property:
Proposition B.4. Let p.v E P2((R)), and E To(H, Then μ = ((1 - t)p1 + tp2)ty defines geodesic between p and i.e.
W2(p, (W2(uD) (VI [0. 1])
In particular, in case where Y is induced by an optimal transport map THi Mr = ((1- and Tht (1 1)ld + UTp hold.
Proof. See Ambrosio et al. (2008), Lemma 7.2.1.
The following continuity equation describes how particle distribution μ4 evolves along time- dependent vector field Uz.
Definition B.5 (Continuity equation, Ambrosio et al. (2008)). Let tr P2 (t I) be curve in Wasserstein space, and let L( (L E 7) be the corresponding vector field. The curve nr is said to satisfy the continuity equation with respect to the vector field by if the distribution equation:
O.j4 + (w14) =0
(20)
holds. Equation (20) means that for all &p E C8 (R²),
d dt I ((x)))((x)) = / Vp((z) (.())
The absolute continuity of the curve μr with respect the Wasserstein distance is equivalent the satisfaction of the continuity equation for some E L (px) (Ambrosio et al. (2008) Theorem 8.3.1). In this sense, the continuity equation provides concept of differentiation consistent with the distance structure induced by the Wasserstein distance.
Proposition B.6 (Ambrosio et al. (2008), Theorem 8.3.1). Let I R be an open interval, and ler He I (R be continuous curve. μr is absolutely continuous, if and if only pr satisfies continuity equation (3) for some vector field vt L²)))..
For the vector field that gives the continuity equation (3),
1t w4 (Vp p Co (RR) LL2(() = X
V. we p.:)
equivalent to the fact that the continuity equation gives the same curve μ4. Noting the orthogonal decomposition (1²) = Xz,. @ X the vector field with the minimal L norm among those that give the same curve must have no component in the subspace X,., and it follows that E X p From these observations, we define the tangent space representing infinitesimal changes the space of probability measures (R))) as follows:
Definition B.7 (Tangent bundle, Ambrosio et al. (2008)). We define the tangent space Po C (p) at E (2( as
Tan, P2(R X = (V p C (Rd) (2))))
The following proposition, referred the Cenamou-Brenise formula, demonstrates that the Wasserstein distance is characterized by the minimal action among absolutely continuous curves connecting two given probability measures.
Proposition B.8 (Benamou-Brent) Brenner formula). For any E (2(R))) and T > O, the following holds:
//" (1velB22(u)) di Oes4 + V (ven)) = (1 (0,7)), n = PT
u) = inf {T),
Proof. From Ambrosio al. (2008), Theorem 8.3.1,
2) = inf { |1 [|velizung dt 0e + V (..p)) = 0 (t (0, Ao = 41 -
By changing the variable t + $. the claim follows.
21

---

The following proposition establishes that the infinitesimal behavior of an absolutely continuous curve can be expressed by the pushforward (Id + he, hpp
Proposition B.9. Let M4 be an absolutely continuous curve satisfying continuous equation with vector field it E Tan,, P2( Then.
W2(pe+h: (Td hvv)# 0.
lim
h-0
h
Proof. See Ambrosio al. (2008), Proposition 8.4.6.
The following propositions provide sufficient conditions for a map to be an optimal transport map. Proposition B.10 (Santambrogio (2015), Theorem 1.48). Suppose that μ E P2(TT() and that convex and p-a.e. differentiable function with Vę L2)))). Then, the map Vip provides the optimal transport map from (Vy)iμ.
Proof. Let 10 be a Legendre- Fenchel transformation of p. Then, it holds that ((x) - (1) 2 (x, y) ((x) (y()) (a, y)
Vx, E R, (Id #paee. (z,z).
For any transport plan E T(4, V)# it holds that
2 /
(z. v) x(dxdy)
2 (p((x) + + (1))
= -- / p())) - f 40 (V((x))))((z)
-2 / (1, (Id x W.)
By adding we obtain
[[11==11 2 ([1]==1++1llv0) +ldx (()) both sides of the inequality,
/ / / yll (ddxd)) 2 / I 2 Vip)H)(dz))
Then Vep a optimal transport map.
Proposition B.11 (Lanzetti al. (2025), Lemma 2.4). Let μ. E P2(RR)) and & E (RR) Then, there exists h > 0 such that ld + AV is an optimal transport map from to (Id + 1V ( fo -h, h). Furthermore,
W2(n, (Id + (V))#m) = 41l VolVlllam
Proof. From sb E C C (IR*), we have supwEr [+) xo, supZER ||V2 co. Then, by taking
suP=ER= [>)(())] we verify that $ || + hnb(z) is convex and (Id + (x) is (²) integrable for all h E [-h.h]. Noting that V ( [=1l=11 + hup(z)) (Id + (z) and applying Proposition B.10, obtain that Id + /V4 yields an optimal transport map. Consequently, the Wasserstein distance between and (Id + 4 computed as:
V (Hll=e/1² hvb(z)
& (fll+ x|| cll2((de)) =h[lV 1
W2(u, (ld AV
Corollary B.12. Let p. P2(Rd) and & C=(Rd) satisfy / oo.
u-esssupreRd (2(((z)]| < cc. Then, there exists h > 0 such that 1d+hV<p is an optimal transport map from u to (Id+hVv)) E [-h, h). Furthermore, Wa(u, (Id+hVp))#w) = h|/V holds.
22

---

B.2 Wasserstein Gradient
This section provides additional discussion on the relationship between the Wasserstein gradient and the first variation.
Definition B.13 (Wasserstein gradient, Lanzetti al. (2025); Bonnet (2019)). The Wasserstein gradient V F P2 (IR") x Rd - R at P2((R)) is defined as vector-field such that for any E P2 (R") and E To( v), it holds that
F(U) F(1)
(v- (00-2))(ddd))) + O
(21)
In particular, there exists a unique element that satisfies V F(u) E (2(R))((RR) (Lanzetti et al. (2025) Proposition 2.5), and we take this as the Wasserstein gradient.
The first variation
oF appearing in the WGF equation (1) is defined follows. Ayz
Definition B.14 (First variation). The first variation 8F öμ P2 P2((R) xR > defined as a functional
satisfying for any P2(")..
d ah h=0
8F (u 8p.
F(u h(v - n1) ()))=
(22)
The first variation, if exists, is unique up to a constant difference.
Recall the definition of the total variation:
Definition B.15 (Total variation). Suppose that ps, P2( The total variation between fe and is:
vV(,,) = sup BEB(R4)
IvBB)-(B)]-
Σ (u(B)~v(L))
CCB(RE). Rd BEC
As is evident from the definition, the mixture (1 - h) + hv, defines the constant-speed geodesic between μI, v in the sense of total variation. Thus, the first variation can be interpreted as the coefficient of differentiation along the geodesic with respect the total variation distance. On the other hand, as observed in the previous section, the geodesic of W2 is represented as ((1 h)p1 + hp2) #h for E To (p, v). Therefore, the Wasserstein gradient can be understood as the coefficient of differentiation along the geodesic with respect to the Wasserstein distance."
Through a naive but mathematically non-rigorous calculation, we have
F(V) - F(P) ~
x) (v)((z))
/
6F
5F 8y (n, u)
ōu (1, x)) r(dxdy)
5F
SF 8h (1,, 21, - x + 2 x,
(H,x) (()(v-x)))) =)) (ddrdy)
)
22
5p
8F
x), -x) (drddy) + v()
This points V v F
(R). This equation does hold without certain conditions, but is valid and implicitly assumed hold many practical cases. Kent al. (2021) make brief mention of this frustration. The equivalence V, V by specifically provides the following correspondence.
6F
V [p = SF x)
ōp (M,
8²F
V2 (,, = V. V N 8n
(1, xt v),
vV,F(u, 2 8F
dfz x).
"hhere is strict dominance subordination between total variation and Wasserstein distance in ged
23

---

C Optimality Conditions for Functionals on Probability Space
In this section, discuss the details of the optimality conditions for functionals probability measures based on the Wasserstein gradient. The first-order optimality conditions have been studied in Bonnet (2019) and Lanzetti et al. (2025), and we begin by reviewing these works. We then extend these results to second-order conditions.
C.1 First-order Condition
The first order condition for probability measure optimization is studied by Lanzetti et al. (2025), including constrained optimization problem. Here, we review the results of the first-order condition for unconstrained problem. In the next section, we extend these results and obtain the second-order condition.
The following proposition shows that Equation (21) holds even when y is not optimal. Proposition C.1. Let F P2 (IR") IR be functional probability space, and differentiable at μ.. The for any transport plan E [(u, v) (not necessarily an optimal one).
F(v) ((v)=F(M) = /
V x), [wy x)d(ddyy) + +0 ( || - yl vl(r)d
Proof. See Lanzetti et al. (2025).
The following two propositions are useful for computing changes in the objective function. The next proposition provides a first-order Taylor expansion for infinitesimal perturbations induced by vector field.
Proposition C.2. For sufficiently smooth functional F P2(RR)) + R and vector field v (²) (n E (2 it holds that F((Id ho)# F(e) = (V F(), 12()) + +
Proof. Consider the situation of Proposition C.1 where (Id #1, 7 (Id x (Id hb))#H Then,
/ v (F((,,,))))) æ)) y(dzdy) = / (V (V=F(H,)), (x + ho(x) - x) u(dx)
= (V, v) 2)
and also
/ x - (dxdd)) = -f [1 ho(z)) - x =((r)(dr) = 111110122000
holds. So, by Proposition C.I. the claim follows.
The following proposition provides a formula for differentiating the objective along an absolutely continuous curve. This corresponds to the chain rule for differentiable curves in Euclidean space. Proposition C.3 (Chain rule). For sufficiently smooth finctional F :P2(Rd) + R and absolutely continuous curve be satisfying continuous equation O.j4. V (Mep4)) = 0, + F(M) is differentiable, and the following holds:
d
F(u)) = (V, F(u)), (a())) dt
Proof. By applying the definition of Wasserstein gradient (Definition B.13) for μ = pethi V (Id #1, and y To(u44,, (Id hwr Hte it holds that
F( F([Id + #))
1
<
(V VuF(ue x) x)) (dwdyy) + o(1)
h
Woluetho (Id + her # + o(1)
< V (me))L
h
+ 0,
24

---

as h + O, where Cauchy Schwarz inequality is used in second line and Proposition B.9 is used in
third line. Then by Proposition C.2, we have
d
F(U++h) F(ue)
F(m))
lim
h-+0
F((Id + Hw) F(1e)
lim
h-+0
h
<v ((t)), ve) (o(H)))
The above proposition shows that the Wasserstein gradient plays crucial role for first-order pertur- bations. In particular, v. P(n) 0, the objective function does not change under any first-order perturbation. This suggests that is reasonable to define first-order stationary points points satisfying V (u) = Furthermore, supporting this observation, Lanzetti et al. (2025) established the following proposition.
Proposition C.4 (First-order necessary condition). Let J¹ E P22(R) be local minimizer of differentiable functional F P2( - IR i.c. holds that there exists a constant r >> O such that
W2(u, k*) <t = F(R) S F(u").
Then the Wasserstein gradient of vanishes at μ²:
VuF(u", =0 p -ae.x,
Lee. V [uF("²)= = (²
Proof. See Lanzetti et al. (2025), Theorem 3.1.
Proposition C.5 (First-order sufficient condition). Suppose that F is differentiable and a- geodesically convex with a 0, i.e. it holds that
F(v) - F())) f (V) - 2) r(dzdy) + a
(a Wz(u,)) Vy E To(,,))
Then, F(x") p-a.e. implies that pº is global minimizer of i.e. F(J) 2 F(μ") holds for any μ E (R
Proof. See Lanzetti et al. (2025), Theorem 3.3.
C.2 Second-order Condition
Building on the results from the previous section, we discuss second-order optimality in measure optimization. To obtain second-order terms, we first prove the following lemma. The key point is that we obtain the derivative F(4)) not only at f = =0 0, but also for < < 1. By obtaining this, we compute the second-order coefficient."
Lemma C.6. Let F be sufficiently smooth functional, μ, v E P2 and Ph be constant geodesic induced by y E L'(,,v) Then, for any S <
d F(w)) dt
/ (V,F(u,, (1- t)x ty). x) (ddzyy).
Proof. By Proposition C.I. obtain the following statement (+): let p. E P2(Rd), 2 F(p,)) and let f4 ((1 - t)y + (p2)# #7, then
d F(M) / (V,, - x) (dxdy))
d
"Similar to this proposition, higher-order order terms (third order and beyond) can be computed in the same manner. It is conjectured that higher-order terms also be expressed as the action of V, V on F.
25

---

Since we suppose t < 1, d
F()++0) F(14)
F(4) = lim
dt
5-0
8
1
F(H++(1-1)s) - F((4)
lim
1 1 s-+0
d
F(U4+(2-1)0))
1 ds 5=0
Here, it holds that
M++(1-1)= ((1 - (L ()s))) (1 ()11-)))]]) ()
((1 - s)((1 - t)pp + tp2) + #p
((1 s)p1 + +sp2) (((1 t)p1 + tp2) x #2)
Then, by (*) for + Me: 4 v + ((1 (1)1 + lp2 ld) we have
d di
1
F()(+(-1)))) - F(4)
F(ee)
lim
I-
(V F) x)) ((1 t)p1 tp2 x Id)##(dzdyy)
1-t
/ (VnF(( (1 t) + ty). y- - ((1 - -t)x + ty)) (dzdy)
1
V F(uu, (1 t)æ ty). v- 2) (dzdy))
Proposition C.7. Suppose F' P2( is sufficiently smooth. Let u. P2( and let ps be constant geodesic induced by y E L'olu, v). Then, as h + 0,
F(un) F(x)= = /
- x) (dxxdy)
(yi xτ, 22 ((dx h? /
+ h22
/ (y VV F x) =)) (drdy o(h²).
Proof. By Lemma C.6,
d²
d
F(Pn)
(V (V,F(Hn (1- (1- hy). y x) (dxdy)
dh?
dh
h=0
1=0
= =
d
Vn (1 x) r(drdy)
dh
h=0
= /
x2 ) 7do2dy2
AV (1, x)) (,) (
= -// (1) 21 22F(,,2,, x2)) y (dx) 7 dy / (VV)F(,, 2) - x) (dxdy
Remark C.8. Except for the term VV, F, this expression can interpreted in a manner similar to Taylor expansion in Euclidean space. The term VV, v arises from the change in the metric μ of
26

---

the tangent space. This phenomenon is similar to what occurs Riemannian manifolds, where the metric of the tangent spaces necessarily constant. is worth noting that similar calculation performed in Appendix 5 of Bonnet (2019). As stated in the main text, at first-order stationary points, we have v, F
which implies vV,,F F Therefore, second-order optimality at stationary points can be understood terms of the integral operator property of V.F.F.
By setting v = (Td + hw) in Proposition C.7, we obtain the following: Proposition C.9. Suppose F. P2(IR)) + IR is sufficiently smooth. Then, it holds that for all v E (p)),
1² F((Id HW) - F(4) = (V,F()) 12)))) + (H, H,)u) (2(N) + o(h?) (h - 0).
2
For the next lemma, denote the interior and boundary of a set A Rd A and AA, respectively. Lemma C.10. Suppose that (TR() and that ʃ TR + R is of class C² (RR). Then, it holds that
= p.a.e. + ff p--.e
Proof. Let B C IRd be a closed set B for Rd f() = 0}. For all B C B, there exists r > 0 such that
1y 1- x|| <r = E B = f(D) = O.
Then we have Vf(x) = Thus,
U((ZER'VV f(z) = =0}) 2 w(B))
= μd () ((OB)
1 ((B)).
Since is absolutely continuous with respect to Lebesgue measure and B is Jordan measurable set, (aB) holds. Hence, fx()) = fı-a.e.
Proposition C.11 (Second-order necessary condition, restatement of Proposition 3.3). Let F P2(RR) - R be sufficiently smooth. 1f E P be a local minimum of F, i.e. holds that there exists constant such that
Wa(u, <T = F(u) F(u").
Then it holds that Hμ- 20.
Proof. According to Proposition C.4, V F(p")) = p.a.e. Then follows from Lemma C.10 that VV FV,F(H")) = 0 a.e.
For any vector field E (H²)).. there exists a constant h > 0 such that h < h = Wz(n", (Id + hu)##' h||b|| (cP(W))) S r. By applying Proposition C.9, we have
(v, H v) (P())) F((1d + hv) F(n") + o(1) 2 o(1).
By letting h + 0, we have (v, Ha* v) (3((m) 2 0ie. Hş. 2 O.
27

---

D Kernels and Gaussian Processes
The purpose of this section to present series of propositions regarding positive semi-definite kernels, the integral operators they define, and Gaussian processes. Additionally, we aim to derive an inequality that evaluates the tail probability of the (p) norm of a Gaussian process.
In the following propositions, the kernel function K R x I& + Rdxd of interest corresponds to the integral kernel K₂ of the squared Hessian defined in (11). However, we first examine the properties of general kernels and Gaussian processes, applying these results to K₁. defined in (11) at the end of this section.
In this paper, we consider multivariate Gaussian processes. Kim & Suzuki (2024) were the first to propose introducing random perturbations probability measures using a multivariate Gaussian process. Another application of multivariate Gaussian processes is modeling vector fields on Rieman- nian manifolds (Hutchinson et al., 2021). For detailed definition and properties, please refer to Chen et al. (2023). Definition D.1 (Multivariate Gaussian process). The vector-valued function & R + RR is said to follow multivariate Gaussian process if any finite collection of variables ((x1).... {(rN) are jointly normally distributed. This process is determined by the vector-valued mean function m Ed Rd and the matrix-valued covariance function K I2 x Fd Idddd
m(r) E[{(x)) (a R K(x = (m(=)(((e)) m(=)) (e,
In this case, we denote & ~ K).
Proposition D.2. Suppose K TRd x Rd - Rdxd satisfies J [I[ y)ll 1 2 (drdy) < oo and o all x, E R.d. Then, there exists a sequence {ka C RR which is finite
K K<,,y) K(y, for or converges to satisfies ||Tx||Te (2()) En>i [Nial co, and is non-increasing. Furthermore, there exists an orthonormal basis (vo 1/ni211 C L (RRdd such that
K(x,y) Σ Kn (1)
n21
where the infinite sum converges in the (2( norm.
Proof. This follows from the eigenvalue expansion theorem for compact self-adjoint operators Hilbert spaces. The trace-class property Σn21 [Kal < oo follows from T₃ being trace-class.
Proposition D.3. Let p. (2(R)) and vp E L2( Suppose K Rd x Rd gaxd positive semi-definite kernel satisfying f ||(x, μ (dzdy) < co. Then, & ~ GP(O, K) satisfies & (2) almost surely.
Proof. Using the integrability of ie and equivalence of matrix norms:
ElE(x) [((x)] p(dr)= / tr(K(2, u(dr) < co.
|
Setting A, 1 p 2 n} for n E N, it holds that
P(IEE113 oo) = 1 P(nz1 Am) =1 1- lim P(An) lim sup E|||[| = 1. n=020
n- n
Thus, & E (2) almost surely.
Proposition D.4. Let J E (RR) and E (p)),)) Suppose K : Ted Rdd Rdxdd is positive semi-definite kernel satisfying |K (I, a)ll u (drdy) < co. Then, for any GP(O, K), we have (b, () (0, (v, (P)))
Proof. The proof follows Kim & Suzuki (2024), Lemma E.9, First, we show that (w), L3w) is normally distributed.
28

---

We define the closed subspace of qquareinntegrbbl real-valued random variables by E span(>(x) {(x) J. For any Z E Et, the following holds:
(1) ( [RIZ
u(dx) O.
Thus, (9) ()) E (E')'= E holds, meaning that (vb, 1))) the (P) limit, and therefore. the law convergence limit of of normally distributed random variables. As will be shown later, the mean of ($, ( is and its variance is (v', Tx Vxv) =²(w)' Therefore, the characteristic function converges pointwise to the characteristic function of a normal distribution, implying that (V, <) LP(W))
follows normal distribution.
Moreover, the mean and variance are computed as:
/ "[[(()) p(dx) O,
E[ (()())
(i (((b)))))(()) ())) = f ((x)) ((((())))))) (E)(() ()v(v) 2 (dxdy)
- (()) u(dx) (1b, Tk Tx (2()))
Here, Fubini's theorem and the definition of Gaussian processes are used.
Proposition D.5 (Karhunen-Loève expansion). Suppose K Rd x Rd + Raxd be a positive semidefinite kernel and satisfy J (1, v)ll (drdy) < co. Let { C C220 and {Vn Wnin21 C (p)) be as in Lemma D.2. The sequence of the random variables Xn == (1n,)) ~ ~0,0ka) is mutually independent. Furthermore, the Gaussian process & GP(O, K) is represented as follows:
5(x) = Σ (xt
a21
where the right-hand infinite means convergence in LP v).
Proof. First, we show that the sequence of the random variables Xn (W () (n 2 1) is mutually independent.
/ vn(()) ((e)a) ((x) / E(v) ((y)TTm (u)
E X - E
(((=))T(((=)E()) J (x) @ p(dzdy)
/ (x) ( K(xx g) (((s)))llaaa u(dr)
/ Vin (x) (Nm)m(z)))lee (x))
Km 8nmm
where Fubini's theorem is used in the second and fourth line. The fifth line follows from the fact that an is the eigenvector the integral operator Tk-
Thus * implies the covariance of X,, and Xm is equal to 0. This implies that X,, and X,, are mutually independent, because they follow the normal distribution. Hence we obtain the mutual independency of the sequence X,, (vh, (P())) (n 2 1).
29

---

The Karhunen-Loëve expansion follows from
E < Σ 2
=B
|1E 1 (w))
- Σ
E Xn (5, (.) (P(m))
(x()))
n=1
(N'n vm I2(N)) mm
=B [1
Σ n=1
x
m
N N-000
Kp
Proposition D.6 (extended Markov's s inequality). For any random variable X and non-decreasing positive- positive-valued valued measurable function yp IR Ro the following holds
Fl(((X) (MM)
P(X 2 M) -
(VM E R).
Proof. By applying Markov's inequality, we have
[lp(X)] p(M)
P(X 2 M) < (r((((X) 2 ((M))
(AM ER.
With the above preparations, we obtain an upper bound for the tail probability of the I,2 norm of the Gaussian process.
Proposition D.7. Suppose K Rdd Tgd Rdxd be a positive semi-defiite symmetric kernel and 6 (dxdy) < oo holds. Let (Kn)22 C B2zo and (thn) n21 C (1)22) be as in Proposition D.2. The Gaussian process ~ GP(O, K) with kernel function K satisfies the following:
satisfy f [K(x (x. y) I| 62
e-1 M2 + Ensi
P(Isllesws 2 M) S exp
(-e-l
(VM > l (T
2ck1
2k1
Proof. First, we calculate the characteristic function of [|11222 Since GP(O.K). K), by Lemma D.5, express ||
En X2 Y (Xn
(O Kn) Vn 1). Let YN En Xn 2 Then for any C, we have
v1 2kn
and so, by the independence of (Xm)n2))
N
1
YN
YN
VI
2
n=1
In partieular, setting iu (u RR, the characteristic function of YN
log Euc
log(1 2in, u).
(23)
2 m=
30

---

Let On = Arg(1 - 2ini u) (-2 2 < 0, < ). Then for the real and imaginary parts,
1
Re
(log(1 2(xs
-
log( + 4n u²) S
4k, uu²
I
2u²
Σ
Kn I 2u
Σ
00,
n21
Im Σ (log(1 u))
Σ
1 S
tan 0ml
Σ n=1
2km |1| S 2/u
Σ x21
Kn Co.
Thus, the characteristic function of YN converges pointwise as N oc. By Levy's continuity theorem, Yv weakly converges to limv- Yv Enz1 X, and its characteristic function is log E[e""Y] 1/2 En21 log(1 u). Now, we will proceed to
Show that for some t > O, replacing iu with i makes the series on the right-hand side of (23) converge, This allows us to show the existence of the moment generating function Ele "V for such I.
Use Lemma D.6 with (px) - e.x, X t- Y to obtain the upper bound for the tail probability.
Adjust t >> 0 (within the range where the moment generating function exists) to obtain the best possible upper bound for the tail probability.
The function t + - log(1 - 2knt) is convex, and the tangent at is + Knt. Thus for a > 1, the equation for
atett 2 log(1 - 2k,xt)
has solution bn for > 0, and holds that
0 S t S tn
log(1 2krtt) S axnt.
Since f,, is constant for 2 1 and {ni 121 is non-increasing, the sequence [tn]n21 is non- decreasing. Hence, for any n 2 1, E [0, 41), we have - log(1 2x,,t) S ax,t, so
log Ele s at Kn 00.
Therefore, for t E [0, t1), the moment generating function E[e[|e)) exists, and E[c"[[l]]] < aat Engi Ka. Hence, for the upper bound of the tail probability of the Gaussian process norm, we have
E(eAHEI1,
p (( M1 s eLM2
(1112 + at
S exp
By optimizing the right-hand side of this expression with respect to f, a, we obtain the best upper bound. For 1, atk; log(1 2, considering large M, solve
minimize f(t) -t + at
Kns
(24)
t 0, a 2xgl log 1 2xyt) 3) 1.
(25)
31

---

The derivative of (24) is
dt d( [(] -2 + En Mi
0,
(26)
- 2xst
which gives
E Enzi
1" 2k1
(27)
M²
If M Enzi Ni: then 1" 0, 1, satisfying the condition (25). The upper bound becomes
P P(IlEILa@ 2 M) S exp
En21 Kin
exp
-2"
log(1 - -2kat")
2k1
M
Emzl Kin
Enzi
Km
M2
exp
+
log
2k1
2k1
2k1
Σn>1 Kin
where in the third line we used (26), (27). Finally, applying the inequality -r + log S 6-1 at
2
we obtain
En21 Fe
Σn>1 Kin
1 e-1 M
f(t") <
2
Σnzi Kin
So, we conclude
( e-- M²
Enz1 Kin
P(lelleew 2 2M) exp
2n1
We apply the results from the previous section.
Lemma D.8. Let K2 be the Hessian-basse kernel introduced in (11), Then, the Gaussian process & GP(O, K2) satisfies the following:
(wx (erE) (O,,),,)))
P-almost surely, it occurs that & E R(H,) P2(R.)). In particular, from the assumptions, Id+ gives the optimal transport map from to (Id+ sufficice small пp > 0 P-a.s.
For any constant M > Rs (Enz1 2n²)=
the following holds:
En>1 An 2
P (IlllIZa()) 2 M) S exp
e-1 M² +
2c>1
2x1 2
Proof.
Given the assumption / V3 x, v) 1 (drdy) < cc. Lemma D.2 can be applied to
K = (()). Specifically, there exists sequence of real numbers {Am)a C R \ {0) satisfying || = Σn21 An oo and an orthonormal basis {dn)n2 of (2()). such that
F2
v.F(Lz V V
(I, Σ Andn
- (z.v).
(28)
n2l
Regarding the kernel of the squared Hessian Kn,
K K2 9)
Σ (2),)))))) 1 (I.
32

---

Here, K2 satisfies the assumptions of Lemma D.5, corresponding to the case where Kx = X in that lemma. Then, from Proposition D.4, (War (2()) ~ NO, X holds. Finally, it follows from
[7 (1 Tr. (Enzi>n)) dn²)² 7
Proposition D.7 that for any M 2 R₂ 2 ||
M2 + Σn>1 An 2
P(IleILAaw 2 M) s sexp
e-l
2eA1
2X1 2
33

---

E
Proofs for Convergence Analysis
The lemmas for the main theorem are re-stated and the proofs are provided here.
E.1
Continuous-time Convergence Analysis
E.1.1
Proof Sketch for Lemmas
The proof of Lemma 5.3 is relatively straightforward using the chain rule. In contrast, the proof of Lemma 5.4 is lengthy and relies on two sub-lemmas inspired by the analysis of SSRGD (Li, 2019) in Euclidean spaces.
Lemma E.1:
small increase by perturbation. This lemma ensures that the increase in the objective function due perturbation is upper bounded by Fihres with high probability.
Lemma E.2: large decrease by WGF. Put simply, this lemma asserts that a small deviation in the direction of the eigenvector corresponding to the smallest eigenvalue of the Hessian results in exponential decrease in the objective under WGF. This result is derived from the fact that the dynamics of slightly deviated two points evolve approximately under Hgt, causing their "distance" to grow exponentially over time.
By combining these results, namely that the objective function does not increase significantly and decreases substantially after perturbation, the lemma is proven.
E.1.2 Proof Lemma 5.3
Lemma 5.3. For a curve of probability measures f4 following the WGF, the following holds:
F(Ho) - F(m)) / 1 (vrr(H)) dr.
Proof. We have
( - F(ur)dr / [V dr
F(Po) F(())
due to the chain rule.
E.1.3 Proof of Proposition 5.4
Tehres: Pehres be chosen as in Theorem 5.2. Suppose jt E (Rd) satisfies V F(F((t)))) ((ut))lraa)) < E and No == >uin Hu nt S -5. Generating
Proposition 5.4. Set O(1) and let E, 8, Up:
& ~~ GP(O, kp) and setting μo (Id + npS)tut the initial point of the WGF, we have with probability
F(12) - F(IT....)) 2 Fthres-
Intuitively, Proposition 5.4 indicates that when perturbation is applied near the saddle point, the objective function decreases over period of time Thrre with high probability, which corresponds to escaping the saddle point. As discussed above, the approach is based on the argument by Li (2019) and utilizes Lemma E. and Lemma E.2. These two lemmas postulate that the 7.2-norm ( of the Gaussian process is uniformly bounded. This corresponds to the condition that perturbations Euclidean space are sampled from spheres with a fixed radius, thus the perturbation size was uniformly bounded. In diment-oniiesionnlll spaces, such uniform sampling cannot be used, which one of the reasons why the Gaussian process is employed. In this case, the norm the Gaussian process take arbitrarily large values even though with low probability. Therefore, it is necessary to exploit tail probability estimates the norm of the Gaussian process, as the following.
Lemma D.8. Let K2 be the Hessian-basee kernel introduced in (n(1)) (11). Then, the Gaussian process E (0, K₂)) satisfies the following
34

---

Wn. (w (0, 2,
P-almost surely, it occurs that & E R(Hy) с Tany P2(R In particular, from the assumptions, 1d+ gives the optimal transport map from μ (Id+ p sufficiently small Upp > 0 P-a.s.
For any constant M Rs 2 (Enz1 in >n²),
the following holds
e- 2 M² + Enz1 An
P ((llll2a00 2 M) S exp
2
2
Let 5
eRi² c-1
11+ log 2))
3
4v2
M
v V2R1
log
Ö(1).
(29)
5¹
From Assumption 1,
R² 2
/
k
=
Σ
An 1,22>>² X1
n21
so that
2 eR1² 2
2
1 2log
5"
2 e
An
+
2e
A1 2
log
e
6
5
n2l
holds. Then follows that
En>1 dn
c c-1 1
M2
exp
2eX1
2\1²
Therefore, from Lemma D.8, IKll S M Ō(1) occurs with probability least 1
We now take the hyperparameters 7p, Fihvess Tthres follows:
T Tehres
2 log
16L, A
=
Vedēr
log
77
log
8
2heres
63
"p M(E Ve² + 2L, Fihros)
A
(30)
0
83
log 57
1
-1 7(108
log
log & 832(
-3
A log 5)
1
-1/
log 3" log
832C'
Tehres -3
Fihres
log = Ō(83)
18(L2 + Lz)
2
1 log
5' 0 83 log 8tZ'
It should be noted that Fihres = Me + 412 11p
2 2 holds.
'The righ side implies the upper bound exp of H.. This result is utilized later in the proof of Proposition 5.4.
A22
< 1. where Ao is the smallest eigenvalue
4A0
35

---

Lemma E.1. Consider the situation of Proposition 5.4, and let pt be an (E, 5--dddle point, let the initial point be Pa = (Id + mp5) and the hyperparameters Up, Fihres be taken as in (30). If the -.-normm of the Gaussian process & sstisfies ([C]| I M, then it holds that
F(po) - F(u") S Fihres.
Proof.
Let V/ = (Id + hnp))tul for E [0, 1], joining y and no = (Id + mp{)#pl From Lemma
C.6,
d dh
F(un) = (V ((n)) (Id - knp,), 3, 4) (2(10)
- 1p v. o (Id hnp (
nh (IV ((2(())))(())))) V ( (Id + hnpE) VuF((t)) (l (cll))
# (11)) 12(Ut)
S uh e Li (f
||0 yll (Id (Id knnp)) (drdy)) ||5|
= mp upllEll2c
Ling 1111122
Here, the Cauchy-Schwarr -Schwazz inequality was used in the second line, the conditions | , FF)))l < and the gradient's Lipschitz continuity were applied in the fourth line, and Proposition B.4 was invoked in the fifth line. Therefore,
F(Po) F(x))= = (''
d F(un) dh
dh
+ L/ (illelli
L1
I no r
S MpMe +
I1
2 112-²
Fihree-
The third line follows from [[ < M and the fourth line follows from the choice of Tp in (30).
Lemma E.2. Consider the situation of Proposition 5.4. That is, let n² be (e, - point,
Ao = (ld + #pp)# and the hyperparameters 7p Fihres be taken as (30). Furthermore, we set R is constant, and th the eigenvector of the Hessian operator Hp corresponding to the smallest eigenvalue Xg-
Ao (Id + 7pt) where & { + rbu,
Letting Pt, Pa. be WGF initialized at po, io respectively, then [1 S M and Dolc' S
|r| < 2M implies that there exists t E [0, Tthres] satisfying
(F(po) - F(j4)) V (F(po) F(j))) 2 2Fthres-
Proof. We prove this lemma by contradiction. We assume that for any E (0, Tthres), it holds that
(F(po) - F(J4) v (F()o) F(R)) 2 Fthhres
We denote the characteristics of the vector fields - -V F(u)) and V F(Te)) as X1, X, respectively.
That is, Mr X.# ji, No and
d X1 = F(,(( Xc,
d X, = -v,F(i)
dt
dt
36

---

It is worth noting here that
=
= -11%' (-V o X,
| - Id [d
$ - V (VLF(())llLA
13 (L 11V (1
$ (fthres(F(Ho)) ( F(
where the first line follows from the fact that X2 is the characteristic of (u)),,), the second from the properties of the Bochner integral, the third from Jensen's 's inequality, and the fourth from Proposition 5.3. Similarly, ||Xx, - Id]l I (Tthres F(jo)) - F())) holds. Furthurmore, we set Y == X² (Id mpE) and Y Xx, (Id + mpE). which satisfy
||Y Id (=())
|| (Id + 7pp (1 + 7)llEll I (Tihres (Ho) - - ((t))))² + no 1ell6ll2((t)
22 F Pitres? + M,
(31)
thrres
and
1 Tdll S (Tthres (F(oo) F(ū+)))) 1p (2()))
< (Tures (F(o)) - F())))) + + + 11pl111 (ell) Ap
< V2T Farcs + M + t/p". thres
(32)
We analyze the vector w. = y, - Y, The goal is to obtain a contradiction by confirming that |/ ( becomes large. To achieve this, we investigate how wi evolves with respect to time:
d
d Y. Y,
d df
dt wi
dt
d
d X, (Id+ + 1pE)
N, (Id mnp2)
dl
di
((p(((tt) o X, o (Id + ) + v F(F( X₂ (ld+
I'' d dh (V, Yn)ch,
(33)
where vn is connecting μ4 and po 24 ((1 h)y: + hY;)#ut (h [0, μr and in = ü. hold. Moreover, its direction vector is always we V, Yt, the integrand of (33) is obtained by operating wi- Specifically,
d dh
(V F(un) (((l))ox (x)
- Yan,
(dy) VV Yx (x)
holds. Let
dt 14 = - H Hat ( ALndth )us - HH, Aewe,
d
(34)
37

---

where A is an operator on (xt)), which satisfies
(())(()))
- -/ YV))) ) f()))² (dy)
(V
/ ( (z) Vi,F(ul f()) (.)
+ (vn Mx(x)) VV,Fu,, VV VV 2)f(z).
(35)
Moreover, the operator A, is defined as A. So Aendd, the norm of which is bounded as Ō (Thhee F2 three + 1/p + e) from (31), (32), and (35). It should be noted that from (34), wy evolves according to Hpt unless | is small. Since
d (d wi) p (a. dt ue + H.rwe)
- Dews
dt
holds from (34), by integrating we have
( c =H A,w,dr,
IB At We
and so
W/ 1, wp - S (x where we used the equation wi Yo Yy =o==ppE -E) 7p Then, it holds that
A-W,dT = 10 S-- (T A,
[[welle Aot npT
= || (Lq(u)) - ||e not - reoll
|e -= m (oll
L ( 12())) | (ell
< {'¹ Salt- 11 (lw
(36)
Here we use []e(T-() (|12(11)) (o) which is implied by do Amin (Hin(H and S From (35) and the inequality [1'', (12,ill12(1)) S R2 1 )ll ca S Rze, follows from (31) and (32) that
| S (I2 + Lx)l|V, - Td]lza(ut) + Ras
=(1 () L3)l|y, Id| LP(n0)
+ thLL2 + Lz) Y1 ld
(2(xt) + +Rys
< (12 + La) v 22T p3
thres chres + M + hmprr + R28.
Therefore, it holds that
| (ell 5 L/ |1 [Auallleewwn
$ (L2 + La) (2I ihres
F Thres + M + ) + Rog.
< (L2 + La) (2Tihree Fihres + +mppMM Roz
A.
38

---

By manipulating the inequality (36), the following is obtained:
- ["ol|lelll - 7pr = A L. e ( (w
The application of Gronwall's inequality to the upper bound yields e || (w S Apre Then holds that
Soare-Aoe Aut Adtt (" @ dr
| (wwt)) 2 More
mpre Aut (1 (1-1 {-" car dr)
mpre -ppre-Aol22-
(37)
The left side of (37) is upper-bounded as
[|welle2(nt) <
-ld 12
ld
+ ||Y - -Idll
2222Tihre Fihress
$ + 2np + T/p
< 4v 4V22ihrees Fthres*.
where in the last line it follows from Tehres Pthres = Ō(), M o()), 1pp = = On the other hand, using 7p MTeres o(1), eTthres o(1) and the definition of Fihress have
t S Tehres
((L2 (3)(V2Tihee))
Fibres + 21p-M + R5E RggeYTehree
V2(L2 + La)res $ Fihres 1/2 + 2 +
+ RasTihres
3 log
3
3
log
log
3
log
Then the right side of (37) is lower-bounded as
Thre Aot 1o 22 ((2r))
Tpp
2
> 7p
2 (côt)
where we used the second line the inequality eă S Veri as x = 8t. Letting Tehres and transforming (37) yields
8/2 Fihres
Thaee
Ve 82 d2mprr
16L, 1
Ve 5žr
where the second line is implied by
2. 2Fthres
Mp.M
+ Ve² + 2hg Fehres
2F
2
28 + 22T, Fehres
22Ferees
2 L
Here it follows from the definition of and the assumption 82 2 (L2 + Lz)8. Since we set
16L, 1
Tihres
log
this leads to a contradiction.
Vestr
39

---

Proof of Proposition 5.4. From the discussion provided after Lemma D.8, by setting M
2og 2. = Ö(1), it holds that || I M = 0(1) with probability - 5 2 (=S (1
At this point, by choosing the hyperparameters as in (30), we have ?p = ) Theres" = O Fthres Ö(5)). and Lemma E. Lemma E.2 can be applied. There exists & [0, T Tthres) such that
P ((F(Ho) F(u)) S 2Fthres) 2 P [(Ho) F(e))= S 2F
(w'b, E) L3(U)))
S M
x²
dr (/0b):))////22
exp
RV STElAnIS
VER Aolc'
2 Aol
2200
V Ao
41a V2e
1
-2
4R1
2 Aal
5'
5'
21
4
2'
where several previously established results are used. We provide a detailed explanation below:
The first line is straightforward.
The second line is followed by Lemma E.2 shown below. The key point of Lemma E.2 is that when considering two points respectively perturbed but the vo-ddrrction of perturbation differs by fixed amount, the WGF dynamics lead to least one point having an decrease of the objective greater than 2Fthees
Let J be the sample space associated with the randomness of the Gaussian process and define random variable X as = (1/o ()
Take W1, We S such that (x(w)) V [(w2)] s M, M(2(w) {w1) + rvu., where r is real constant satisfying the assumption in Lemma E.2. We consider applying Lemma E.2 in this setting.
Since ((n)) & M, (X S M implies |r| = (X(w)) - X(w2) s 2M, from Lemma / Aol'' then least one the samples w1, w2
E.2, we obtain that if |r| [nol, vo
satisfies F(Ho) F (())) 2 2
Based on this, consider the following two cases:
There exists a point un [M-M, M] such that X(w) 300
F(uo) F(u))
2Furres-
For all points wo [-M, (w) zo = F(wo) - F(u+) 2 2F
In the former case, by (+), x [-M, M) (ro ro, 3o + ro) = F(Uo) F(u)
2 tirres holds. Therefore, either case, the following holds:
P (F - F(w)) S 2F Fhtheess (Wo. () s M)
2P - \ (zo ru. Ru -rr))
1
- dxc
(MM.M) 22T
25 dx.
2xp dx
B((30-ro.2otrm)
V2TAA0
JR
2T
Here, the first term is minimized at zo = 0, which justifies the bound used in the second line above.
The third line The lower bound of the first term follows from the fact that the Gaussian PDF reaches its maximum
22n at x = 0.
40

---

The lower bound of the second term comes from the Gaussian tail bound. (This follows from Proposition D.6 by taking (vo, L)(tt)) (p(x) = exp 4x6
The fourth line follows from the definition of M. Thus, with probability
2 we have
F(Po) - F(4) 2 2Fthres-
Finally, combining with Lemma E.I, the following holds:
F(1)) - F(UTT = F(ut) F(Ho) + F(wo) - F(
2 Fihree 2Fehres
Fthres-
This occurs with probability more than 1- ()=-)-1-0
E.2
Discrete-time Convergence Analysis
this section, we prove the convergence discrete-time PWGF (Theorem 5.5). For the discretization, the following Lipschitz continuity of the Wasserstein gradient is utilized to bridge the gap with continuous time.
Lemma E.3. Let F Pe(IR)) + R be sufficiently smooth and suppose that the Wasserstein gradient V F be Lx-Lipschitz continuous." Then, it holds that for any μ1, PP2(RR) 7 (
F(v) F{H)+ - / VFF((,)).. z) (ddddy) +
L1 / || ?dxdd)
2
Proof. Let ((1 - t)p1 + #y E P2(R"), = =Idxx (1 t)pi + + E I(, Pt). The application of Lemma C.6 implies
- F(uo) = / (V (V,F(lr (1 - t)x tw), x) (dxdyy) di
d
/ (1, - (drry
x) (drrdy) (V,F(ue (1- f)x + tw) - V F(F((x))
$ / (V,, xx) 7drryy)
L1 (f || - v11"2 ) 1 ( || ||x- v1 (drdd))))))) :
v (V, F,, (drdy + Li / || "
In the third line, we use the Cauchy- Schwarz inequality and Lipschitz continuity. Then, we have
F(v) - F(u) = //'' d
F(
dt
ʃ (V,F(,,)) (p,x), - 2) (ddrdy) 2 / || vdl3d(d
In
The following proposition discrete version of Gronwall's inequality, which will be used in subsequent proofs.
SSee Assumption 3.
41

---

Proposition E.4, Let faxt be real-valued sequence and b + 0, CE R. Then.
k-1
0k c S
a (vk 20) = as a0 S c(b 1)*.
1=0
Proof. Letting d4 = Ei=o ak2 we have
dx + 1 s ag + d + b $ (b+ 1)d4-1 + c
c
b b"
= (b + +1(((d + 9 s (b+ 1
Then it holds that
ak c bdz-1 S bc 1)a - c(b+1)*
The following proposition corresponds to Lemma 5.3 in continuous time and serves prove two key points: the evaluation of the decrease in when the gradient is large, and the ability to reduce the objective function near saddle points.
Proposition E.5. Let {n()) 1t= be a sequence of probability measures generated by discrete-timg time
PWGF (Algorithm 2) with step size n S L,"
J. Then it holds that
k
F() F(u()))
V F (1(0)) 13(u00)
1=0
Proof. By Lemma E.3, it holds that for any 1 0,... kk- 1,
F(y(+1)) - F) V / (V (v x). y - x) (Id x (Id-
- 7/ v)F(((()) # (dxdy)
L1
/
+
||0 y1 ld x (Id - F( #1 (dody)
2
|/ F(N 1|2
2
L2(UE0
17
V F(u)))))))
S
2
L
Then we have
k-1
F(y(t)) - ())
Σ
( -(())))
=
1=0
17 h-1
s
Σ
V ) Lo(W00"
42

---

In discrete-time PWGF, we take hyperparameters Up. Finres, Rehuess 17 as follows:
M- = ell (1
1 + 2 log 2 ))
t V2 log
4v2 $
- Ö(1),
5
7 S L1 0(1)
2
16 n
- - -0(
kthres
log(1 1 + 76) log
(38)
Ver log² (1 1n)))
Fihres 11 3
3 = Ō(5²) log"
=((3),
18(L2 + Lz)
2
FFehres
{ ^ 1
δ 83
T/p = M(E + VE2 F
Since Up and M are set in the same way as continuous time (30), Lemma E.1 holds in discrete time as well. It should also be noted that Pthres is defined to correspond to Tihres and rykthress
Next, we present the discrete-time counterpart of Lemma E.2.
Proposition E.6. Let p¹ be an (e, -ssaddle point, p! (Id + #p and the hyperparameters
Tho Fihres be taken as in (38). Furthermore, we set i(()) (p)) where ? = & + rvo., T E IR
is constant, and tho is the eigenvector of the Hessian operator H₁ corresponding to the smallest
eigenvalue 2o-
Letting (,,,,, ((()) be WGF initialized at ())))(())) respectively. then l|Ellaa()) S M and
AnlCA
S |m| S 2M implies that there exists = 0,. kehress satisfying
F( 2 2F Fthhes (F(x F V (F(
Proof. We give a proof by contradiction. Assume that for any S S kihress
(F(x - F() (F(E))) )( F (p (1)) < 22 Fihres.
The following vector fields are used to evaluate how A AF( (F ()) u (k) and A (o7(0) ((() evolve:
y(k) (Td V, v F(((t ................................ (Td v (1110)) o (Id + 1755),
(4)) = (Id F ..... (Id 7 F F(z (Id + 5)
Here, we note that
(1()) ld
S < +
y(+1) - -
L2(((0)
S <mbllll1> +
V (()()))))))
= = + n/
V p F (14
k
1=0
< <mplell>> + n/
Σ
V
F (1)
k
1=0
= = + 22
1/22
1
k-1
V
((()))))))))) 112
1-1
<nblell<>>(>) + Ven (F() (F(x()) (1)
S M + 2ny 2 pd
(39)
kihres thres'
43

---

The third line is due to Jensen's 's inequality and the fifth line is due to Proposition E.5. Similarly,
| Ial La(Ht) S Mp-M + nHr + 2n thres thres'
(40)
We demonstrate that either ((()) or j((k) significantly decreases the objective function. To this end, we define the following vector field as measure of the "distance" between (x()) and ji(k;;
w(k) == yr(k) y((k)
(k) follows the recurrence relation given by the following:
w (b+ (A) = - v. FT o 8(()) V,V (u((l) y(b)
= =- I/' d
dh ((u)))) v )dh
nHH,. -n (k) w)(()
(41)
where we set (k(k
(1 h) + hT(()) Vn = (n #ut ((1 - -h)y + +hr(0) and
A A (n)
dh,
A f(x) / (V (1 - V2r("' - x. v)) s(y) (dy)
(VV,
- ( -V...(,)))))
+ [V,FFul,x)(00
The recurrence formula (41) yields
(k) (1 1-1)HHt))) k (0) 17
Σ (1 n/1111) )
A (0
1=0
k-1
= (1 - n7o) km (0) - 11
Σ
(1 - n
| (1) (1
1=0
Here, we use the fact that (0) (0()) - - noro Then, have
|| (wscao) LP( - (1 - 77700= 1 (k) - (1 - -1720)* (0)
La(at)
|| HH nuر (1 (1)
(w(0)
17
(2t)))
(2(
- 17120) 1 ) 111000)
7
12(11)'
::. =111 n720) - -
< A
- 170 -1 La(att"
L2(Ut)
where the constant A upper bounds the norm of (R) and set in the same manner in continuous time:
A = (L2 + Lz) kitures thres + M) + Rus 2 | Lz(at)
Using the discrete version of Gronwall's inequality (Proposition E.4) with as (1 (7)p)= -k (R)
(e(n)) b b= nd
1- c= mpt' c=m,,, we obtain the following:
(1 (ddn) -10
I 1pr
1+
7
(P(a))
1 - adu
44

---

k2 From this, we have
k-1
(1
10
2 upr
7
-1-
1-11000
(P(
1=0
n
n
7pp
1
+
1 17A00
17200
t=0
(1
n
nA
1-120
Apt
- 17A0
nd
1-mAa
-
n
Opr
2
1 1 1-1200
1120
9pp (2 (1 + ,/)))*
2 rp exp (npithres
1p"
2
where the last line holds as pletures A S log 3 in the same manner as in continuous time,
2
mksehres A S nphttress (L2 L3) (2n thres F + 2T + Rgg)
< 22L2 + (n)hithne
)a F Pirres + 2(L2 + L3) miichves + R
3
3
$ log 2
3 log
Then have
1o(k)
2 7pp""
(1 11A0)
2
2 7pp""
(1 + n1)x
2
2 hp
Vek (1 + 110 log (1 + 16).
(42)
2
On the other hand,
[ < (a)) ld (x²)
+ | ld (2(1t)
41j kihres Fir + 20p + Tppp"
8np thres Fiaces
(43)
where in the second line (39) and (40) are used, and in the third line n7 kahres Firres Ō()), MpAM = o(8), dpr o()). Letting = kthres. it follows from (42), (43) and nn M 2
2F that V la
16 F thres Vetpp log (1 n6)
(1 + n7) Ayge
16 26 17 M Ver log (1 188)
This leads contradiction, we had set Ratres in (38).
The following proposition corresponds to the discrete-time version of Proposition 5.4.
45

---

Proposition E.7. Lete, 5, Ç > O be chosen such that (L2 + L3) S 8.. Suppose μtt E P2 (Rd) satisfies [ (12())))l00 (3 < & and X == >min Hpp S -6. x Considering & ~ GP(O, ky) and
setting (Ao = (ld + ppp) as the initial value of the discrete time PWGF flow 4(*), with parameters
= 0(1), 1pp = Ô 63 A ) ketress = =O))), and Fehres = Ö the following holds with
probability 1-
F(u') F(u!herrrw)) 2 Fthres.
Proof. From the discussion at the beginning of the previous section, by setting M
cC 11+2log )) Ó(1), it holds that || S M = =(1) with probability 1 2
By choosing the hyperparameters in (38), have 7 0(1), 7р = 0 ) kttres
( Fthres - Ô(8*) and Lemma E. and Proposition E.6 can be applied. similar manner to the proof of Proposirion 5.4 (continuous version), it follows that there exists 0 S & S kitres such that
2 21-5 S
P (F(x -F(((() 2 FFihres
2
Thus, with probability S, we have
1)) F( F(u(0)) - F(u(t)) 2 2 F
Combining with Lemma E. the following holds:
F(w²) F(u F(u (Munce)) F(ut) 1) + F(u( F(H(hrenc))
2 Pthres + 2FFerres
Fibres.
+)+
This occurs with probability more than 1
With the above preparations, we finally prove the convergence of discrete-time time PWGF to second- order stationary point.
Theorem 5.5. Let the initial point be μf P2((R) and denote AF = F(Po) (R F()).
Set r O() and let E, 8, up. Tihres, Fihres be chosen as in Theorem 5.2. Then, discrete time PWGF halts after
k k- Ar (
steps and reaches (eee)))sssooddddeeeeeeeeo -eecond-oor stationary point with probability -
Proof. Let 5, <> 0 be chosen arbitrarily chosen such that (I1 + LL)) E S 52, and set Ç' >> 0 such that Ç' is polynomial in 1 and ( up logarithmic factors, provided later. By the settings of m,, Tpp-
kehres: and Fibres as in Proposition E.7, we have 1) = 0(1). 7p Ô ))) kithres Ö()). and
Fihres )
AF times. Thus, the probability of From Proposition F.7, perturbations occur at most m = Fehres
failure after m perturbations at most 1 - (1 - C²) S mC' Setting ç' = S ensures that the algorithm reaches an (s, -secoond order stationary point with probability at least -.
m
Discrete time PWGF determines whether the objective decreases by least Fibres after a certain number of iterations kibres following the application of a perturbation. Then, we define the period perturbation and this evaluation as State and all other times as State O. Let Ko denote the total time spent in State O, where (V,
between the application
2 E. By Lemma 5.3, the decrease in the objective function during this time is at least c² ko. implying ko S 22 AFkiusa Ō ())) Hence, the Fibres
AF Moreover,
the total time h in State 1 is upper bounded by hı I mTThres
algorithm halts in kn + ki Ó(AF(#+ + )+) iterations.
46

---

Numerical Experiments
ICFL Functional
2.00
static
isotropic
hessian
epochs (1e4)
Figure 1: Trajectories of the training loss for no-noise ("static")), isotropic noise ("isotropic") and Hessian guided noise ("hessian") settings.
We conducted numerical experiments using the loss functional of in-context learning of Transformers from Kim & Suzuki (2024) as the objective functional. See below for details.
We compared the dynamics of the loss function under three variants of WGF; WGF without noise (static), WGF with isotropic noise (isotropic). and WGF with noise guided by the Hessian (hessian).
As Figure shows, the loss decreases gradually in the "static" case, whereas the "isotropic" and "hessian" cases exhibit significant reduction in loss, leading saturation. Furthermore, the Hessian- based noise demonstrates a more efficient decrease in loss.
Experimental details. provide a minimal explanation of the loss function used in our numerical experiments for in-context learning in Transformers. For an in-depth exposition the problem setup, derivation the loss functional, and an analysis of the loss landscape, we refer to the work by Kim & Suzuki (2024).
We consider a mean field two-layer neural network with sigmoid activation function, which takes 1-dimensional data inputs and k-dimensional data outputs
hz(2)) = ha (z))((d))) =
ad(wT 2)p(dz) (z = (a, w) E IR x R²),
where RC is given data and follows certain distribution. We also define the following matrices
Σn = E, (x) (
We consider performing in-context learning using hy as feedforward layer, followed by reparametrized linear self-attention mechanism which can be described by a single attention matrix W E & The optimal value of W is determined to satisfy Eye, W En - with given Thus, the objective with optimal is derived follows:
D-1
F(n) - [)) Ene, E M (z)
(44)
where 112 E (PR is the true feature and Ga-,, (+()) = (1)(2) - Enoo i (z) The Wasserstein gradient of the objective (44) is computed as
- E E o Eg [a(w" 2)Gpo,,(())
V FnF w)
2n ('(ww' (wT =))n),
2
47

---

Algorithm 3 PWGD timee/ space discrete)
initialize (0)
--. N 10(0)
10)
E 80000 fo>
for k 0,1,... do
if
N F F S E and h: - kip > kittree then
ç ~~ GP(O0 kytts)
(f1, ((.) t (E(x)-) (1))
((., (1)))
oj - zj + +n (j = 1,,,,, N)
(k)
((k) + N E 0.g0
ki 6-
end if
if = kp + kitres and F(u(tt))) F(((()) S Fihres then
return je (kp)
end if
2 (R+1)
t 2. (6)
F (M))), (j = 1...., N), N) t N E=== (att
DN
end for
Furthermore, Hessian V2F at a first-order optimal point μ is computed as
Hı
Hı
a,, w. b,
H21
H22
where
H1, a, w,
= VaV, F²F
6n2 (114,, w,
= =(Flo((w 2)0(1 (2) 2)] - E, 2)h,(2)] [o) Emp +Efoo(w 2) Te)Gar [x)) (=)] EL(o(v" 2) (1 (2))
F uip una Dno is u.l
and
H12(H, a, w. w), = Hn
= V.. F²F
8uz (M, u,
= -E Ee E
(Es[o(w) 2)0 (1 2)2 - E, [o [x)] =))
(1)))
[o("" 2) [Er (TT)2=]
+ (ow (2)h)(2) E - m Eg (z))(
as well
a,
(a Epph Ene n Expbb
(Ez[za' (w 2)0 (o 2)z 1 - Eg [20 (w Iz) -- u. Ez (z)))(o=)))))
(a ExT b) Eg [2o 2) (2) T JE [S (2)) (2)0 (w 2)²
+ Eg [2
(2)TT²=) ba 2n Ep Es (Sa (2)0 2)=
+ [z (wo 2)G (2) "E Dn E. (4) =)=
The experimental setup is as follows. compared the dynamics of the loss function under three variants of WGF; WGF without noise (static), WGF with isotropic noise (isotropic), and PWGF
48

---

(hessian). To ensure fair comparison of the three algorithms,, no stopping criteria were incorporated into the algorithms.
The input and output dimensions were set to / = 20, 5. We approximated the measure using 400 neurons and generated 800 i.i.d. input data points from the standard normal distribution (0, 1) for each coordinate. The optimization probability measure was randomly initialized and we conducted fiveive experiments under the same conditions, plotting the mean and standard deviation. We used parameters as " = 0.015, kthres 100. In addition, SGD was used in the optimization process with the learning rate 77 = 10
F.2
MatroiiDecoomosition Functional
Next, we conducted experiments using the matrix decomposition setting presented in Example 1. Details of the objective function, including analytical expressions for the gradient and Hessian, as well as proposition suggesting strict benignity of the matrix decomposition objective, are provided in Appendix G.
25
25
static
isotropic
20
20
hessian
15-
10
Liylo صباا
static isolropit hessian
60
0.0
0.5
3.0
0.0
0.5
1.0
Figure Trajectories of the training loss and the norm of the gradient for no-noise ("static"), isotropic noise ("isotropic") and Hessian guided noise ("hessian"" settings.
Due to the stochastic nature of the algorithms, report the mean and standard deviation over 10 runs, with the corresponding error bars. As shown in Figure 2, the Hessian and isotropic noise injection methods achieve faster objective reduction and exhibit earlier peaks in the gradient norm, demonstrating a more efficient escape from the initial critical point. In contrast, the perturbation-free method tends stagnate for longer periods. The Hessian method shows the best performance, although the performance of isotropic noise is comparable. The effectiveness of isotropic noise can be attributed the fact that the infinite-dimensional nature of the problem has not yet manifested due to the number of particles used in the approximation.
In practice, points where the gradient norm is small, but the point is not saddle, adding noise may hinder the gradient descent. This issue is particularly pronounced for the method with Hessian noise, where the magnitude of the noise depends on the local curvature. Consequently, whether to inject noise should be determined adaptively in combination with the criteria discussed above.
Experimental details. The input and output dimensions are / 15, A 5. We approximate the measure using 400 neurons and generate 800 i.i.d. input data points z from the standard normal dis- used in the optimization Fehres = 10-2.
tribution (0, 1) for each coordinate. Similarly the ICFL case, SGD process and the learning rate 7) = 10-6, We set kehres 100, Up = 3 x 10-3
49

---

G
Landscape Analysis of Matrix Decomposition
We analyze matrix factorization (Example 1) as an example of non-convex stochastic optimization:
F(u)
Es
ku (2) (z)
2
1
Es (2) - M
2
E.
(=)
2
where we set
(2) f Na (a),(()))dddd
na = ha)
aww 2)
M₁(2) (2)
(2, = Mu(2) - M
The Wasserstein gradient computed as:
F, a1, d2, w) 2E IV Va,wha,w (
2F [o(w =))
(45)
2 [2o'(w 2))x(2)) AM,(z)) a1
The Hessian is computed as:
V F x, V,FFl,
(46)
= Eg V (2M + (1h (x)( (V (V))
(47)
From (45) and (47),
For any μ P2(R'), n 60 is strict saddle point.
. 4= == (=Idg) x (Idg' 4t is global optima.
Furthermore, by similar argument to Ge et al. (2017), we can deduce the following proposition. This proposition asserts that e-stationary point μ which is not global minimizer, the objective function can be strictly decreased. This suggests that F possesses strict benignity.
Proposition G.1. Let P(R*) be e-stationary and not a global optima, i.e; satisfy F() S and F(J) # 0. (=) 2 0 a.s. 7, Assumption and Wz(n, nP) S C hold, then there exists curve μg s.t. Pl = μ and at t = 0.
d² F(4)
-Ezil (x)) + (O(().
dta
Proof. We define curve br by μr = (1 (1-t)p- t)p + tu. Then we obtain t = 0
d
F(m) = Eg tr ((a) d
hi - hyo (lamh - )1) (2)
dt
dt
((1
d ti + hp, de 1o (Mire ha 1 [x)
T
Fg tr
(48)
dt
This condition can be regarded extension of non-negative matrix factorization.
50

---

d2
= Ex tr (
d²
d²
F(u))
K pip + hue dtz Mr (hize hi - h k (T))
di3
dt²
d
dt kmn (hunting - ( ]
d
+ Eg
tr
2
hus
dt
(a)) Mu Mipe + live d
di (1
ti + luke dle "thh Ni (x)
+ Eg
tr
di
cll
E.
(4
3)tr (1ima
4Eg
tr
hus ct ( (Buathte - te dt
d
d
E2
tr
hipe
hy dt
tepe J Ni (T)(2)
dt
dt
dt
d
d tu ai
-h (e (2)
Es
tx
Mu dl "e di
Ti
3E.
[ur
(Iige Ar. h()) (l
h Mo )
d
F(ue).
(49)
dt
where we used, in the sixth line, the equations;
d²
hyg (2) 0,
dt2
t=0
d
dl
hyc (2) + hus
hx.(2) (2) (h2 - h )h (2) + ha(h )
dt
1==
t=0
= -hiphi h (h1 - h1p - 1mw )(2)
d d
-h 1 dt ==
Mye dt 1=0 (x),
and in the ninth line;
d
F(ee)
-E (tr(((uttii te (hallip hn ( (2))]
dt
1=0
E& tr ( d
d
Nug गi pp (huo hic (z)
di
which are derived from the definition of μ4- Noting that
hus
(z)
hp - -hook (1)
(50)
dt
dt
F
1==
which is obtained by straightforward calculation and the assumption hp 2 0 a.s.;
2 App NT, -np (2)
- dt
d
Aye
d
(z)
dt
me
1==
2111 - 1 (2) - [1 - hn)
hm,)) T 1| (x)
[1)) + | ( (2) 2111 Ayw (²(2) 11 1 (2(2) (2)
(11bml1(()) || 11(()) 4h hw 11,, k (z) (z)
2 0
51

---

Then we obtain, from (49) and (50),
d²
F(m)
- -E huth - h - dl 1=
F(M))
(51)
di2
t=
Finally we will show that
F(me)
= =2 [1h -o,)"
dt ==
- ((s) for some Ö > For any y E (()
S , ((W)l=/
(52)
F(ue)
dt
1=0
E [hy - h1p AM,
=
- - AMaa
2
E.
2 (s 1) TT,EE E.
= 2 / ((x 3) V2F(u,)) + (2 - vV (u,
0 - x)) v)) 7ddryy))
s "", s 2 II
+ W2(H, n ) sup VV F(,,
where Taylor's expansion is used and 0, Ø [0, 1] in the third and fourth lines, the Cauchy-Schwarz inequality and the definition of the 2-sasserttein distance in the fifth line, and the assumptions W2(n, v²) S C' and Assumption 2 in the sixth line, Setting Ö = 2(C R2C²), we obtain (52) and hence, from (51),
2
F(w)) S -F (I hi [+] + +
dt
52


Reviews:
[{"review_id": "oBVsGoVYJm", "reviewer": "Reviewer_LaL6", "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.", "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?\n\n2. 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?\n\n3. Do you have any intuition on the choice of the kernel? Could you explain maybe by showing the analogy in euclidean space?", "limitations": "yes", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 4, "significance": 3, "originality": 3, "strengths_and_weaknesses": "Strengths\n\nThe 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.\n\n\nWeaknesses\n\nJust 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.", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for the thorough review and positive evaluation of our contributions! Our responses are as follows.\n\n**(Weakness) I would suggest adding more empirical demonstrations if possible just to show the potential empirical values.**\n\nWe have **added new experiments and a detailed landscape analysis** of the matrix factorization problem using a mean-field neural network (Example 1) in Appendix F.2 and G which suggest the presence of strict benignity and show that noise injection helps escape the initial critical point. Since NeurIPS this year does not allow updating the submission pdf nor adding anonymous links during rebuttal, we summarize the results in text here.\n\n* Experimental details: The input and output dimensions are $l=15,k=5$. We approximate the measure using $400$ neurons and generate $800$ i.i.d. input data points $z$ from the standard normal distribution $\\mathcal{N}(0,1)$ for each coordinate. Similarly to the ICFL case, SGD was used in the optimization process and the learning rate $\\eta = 10^{-6}$. We set $k_{\\mathrm{thres}} = 100, ~ \\eta_p = 3\\times 10^{-3}, ~ F_{\\mathrm{thres}} = 10^{-2}$. We report the mean and standard deviation over 10 runs, with the corresponding error bars.\n* We find that both Hessian and isotropic noise injection methods achieve faster objective reduction and exhibit earlier peaks in the gradient norm, demonstrating a more efficient escape from the initial critical point. In contrast, the perturbation-free method tends to stagnate for longer periods. In particular, the Hessian method shows the best performance, although the performance of isotropic noise is comparable.\n\n**(Q1) 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?**\n\nUnfortunately, it seems difficult to extend the analysis to backward Euler in the discrete time setting. The discrete-time analysis is discussed in Appendix E. Using the definition of the JKO scheme, we can derive a result similar to Proposition E.5. However, it is difficult to establish Proposition E.6, which is key to the proof of saddle point escape. This is because the transport map $Y^{(k)}$ (defined by $\\mu^{(k)} = Y^{(k)} \\sharp \\mu_0$) is generally not available in implicit schemes such as the backward method, and unless $Y^{(k)}$ is expressed explicitly, it is difficult to carry out any Hessian-based arguments. \n\nIn terms of stability, implicit schemes such as the JKO scheme are generally considered to be more stable than the explicit scheme used in this paper. However, for the class of smooth objective functionals we consider, stability does not pose a problem, and the theoretical analysis remains valid. \n\n**(Q2) 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.**\n\nClassical optimization problems such as matrix factorization and tensor decomposition are known to be affected by saddle points. These problems arise in practical scenarios such as image reconstruction, object detection via matrix sensing, and recommendation systems [ZHBL10, ZKLR13, Y08]. In Euclidean space, the landscape of the objective function for matrix factorization has been extensively studied [GHJY15, GJZ17]. In particular, it is known that saddle points exist, and that all non-saddle stationary points are global optima for these types of objectives. A major concern in such situations is that gradient updates can get trapped near saddle points. For example, passing through neighborhoods of saddle points using gradient-based methods can take arbitrarily long time [DJL+17]. Therefore, escaping saddle points is crucial in such settings. \n\nIt is highly likely that analogous problems in Wasserstein space (e.g. mean-field formulations of the above problems) suffer from the same issues, as we demonstrated in our new neural matrix factorization analysis and numerical experiments. Moreover, recent distributional landscape analyses (including the MMD and ICFL settings introduced in the paper) show that new types of saddle points or benignity properties can manifest in fundamentally distributional problems. On the other hand, existing mean-field studies have been mostly constrained to the linearly convex setting, such as shallow neural networks; our secondary goal is to more deeply study second-order behavior of WGF to understand wider families of deep learning methods.\n\n**(Q3) Do you have any intuition on the choice of the kernel? Could you explain maybe by showing the analogy in euclidean space?**\n\nThe intuition is that the kernel is chosen so that points move in directions where the objective function is unstable and tends to decrease. At a saddle point, the minimal eigenvalue of the Hessian is negative, and the objective function is locally concave along the corresponding eigenvector direction. For example in Euclidean space, one can look at the function $f(x,y) = x^2 - y^2$ around the origin, where the $-y$ direction exhibits this concavity. Therefore, by using noise to perturb the point in that direction, we aim to move it where the gradient is steep and the objective function decreases. In the above example, we would want the noise to contain the $(0,1)$ direction, which is the eigenvector corresponding to the negative eigenvalue of $\\text{Hess}(f)$. \n\nIn finite $d$-dimensional settings, it is known that adding isotropic noise is sufficient and $\\text{polylog}(d)$ updates suffice to escape via a geometric argument: the region of points which fail to escape is shaped like a pancake and thus has small volume. However, in infinite-dimensional curved spaces, this naive argument does not work and the noise must be modified into a more tractable form. To this end, we consider Gaussian noise defined via a Hessian-based kernel. This \"directed\" Gaussian noise is fundamentally connected to the eigenstructure of the Hessian as discussed in Appendix D, and faithfully reflects the above intuition even in Wasserstein space.\n\n**References**\n\n[ZHBL10] Bo Zhao, Justin P. Haldar, Cornelius Brinegar, Zhi-Pei Liang. Low rank matrix recovery for real-time cardiac MRI. In 2010 IEEE International Symposium on Biomedical Imaging: From Nano to Macro, pp.996–999. IEEE, 2010.\n\n[ZKLR13] Wenbin Zou, Kidiyo Kpalma, Zhi Liu, Joseph Ronsin. Segmentation driven low-rank matrix recovery for saliency detection. In 24th British machine vision conference (BMVC), pp.1–13, 2013. \n\n[Y08] Koren, Yehuda. Factorization meets the neighborhood: a multifaceted collaborative filtering model. In Proceedings of the 14th ACM SIGKDD international conference on knowledge discovery and data mining. p.426-434. 2008. \n\n[GHJY15] Rong Ge, Furong Huang, Chi Jin, Yang Yuan. Escaping From Saddle Points—Online Stochastic Gradient for Tensor Decomposition. COLT, 2015.\n\n[GJZ17] Rong Ge, Chi Jin, Yi Zheng. No Spurious Local Minima in Nonconvex Low Rank Problems: A Unified Geometric Analysis. ICML, 2017.\n\n[DJL+17] Simon S Du, Chi Jin, Jason Lee, Michael I Jordan, Aarti Singh, Barnabas Poczos. Gradient Descent Can Take Exponential Time to Escape Saddle Points. NeurIPS, 2017."}, {"review_id": "y5DJW8WDON", "reviewer": "Reviewer_FN5C", "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.", "questions": "1. Can you provide matching lower bounds (even in simpler settings) to argue that the derived rate is unimprovable?\n\n2.  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?\n\n3. Given the cost of sampling from the Gaussian process, are there efficient approximations that preserve the escape guarantees?", "limitations": "Yes.", "rating": 5, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 4, "clarity": 4, "significance": 3, "originality": 3, "strengths_and_weaknesses": "Strengths\n\n1. 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.\n\n2. Unlike isotropic noise, the Hessian-guided GP perturbations are directed along the most negative curvature, mirroring perturbed-gradient-descent strategies in Euclidean space.\n\nWeakness:\n\n1. It would be nice if the paper could touch base on some practical issues as detailed in the question/limitation box below.", "was_revised": true, "final_justification": "See my reply to the rebuttal.", "rebuttal": "Thank you for your detailed review and insightful questions! Below are our responses.\n\n**(Q1) Can you provide matching lower bounds (even in simpler settings) to argue that the derived rate is unimprovable?**\n\nOur primary goal was to initiate a theoretical investigation of perturbative methods for nonconvex probability distribution optimization and achieve a polynomial rate which is already first-of-its-kind, so we did not focus on optimality. Naive WGF can of course take arbitrarily long time to escape a saddle depending on how close it was initialized to the saddle point, and so global rates are impossible to obtain as in finite dimensions.\n\nIn fact, with $\\delta^2\\sim\\epsilon$, the derived rate is $\\widetilde{O}(\\epsilon^{-2})$ which we believe is potentially suboptimal. In finite dimensions, sophisticated second-order methods such as ARC [CGT11a] which minimizes an adaptive cubic-regularized local second-order approximation of the loss, or TRACE [CRS17] which solves a regularized trust region problem at each step, achieve an iteration complexity of $O(\\epsilon^{-3/2})$ for nonconvex optimization, which has been shown to be optimal in certain settings [CGT11b]. These results (neither lower bounds nor the proposed methods) cannot be translated directly to the distributional optimization problem, however, so we anticipate more work is needed to rigorously establish optimal bounds and develop matching algorithms.\n\n**(Q2) 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?**\n\nExisting analyses of $N$-particle approximation error (the so-called mean-field propagation of chaos) either blow up exponentially in time via a naive Gronwall-type bound [MMN18, MMM19, DBDFS20], or rely on uniform log-Sobolev arguments to prove uniform-in-time bounds, which only hold for strongly convex systems with Langevin noise and generally require exponential runtime [CLRW24, SWN23, KZC+24]. We can easily obtain a bound of the former type by the same arguments; however, none of these approaches are desirable as we aim to achieve polynomial iteration complexity for general nonconvex problems, for which little is known besides the recent work of [GWB25], which requires a local strong convexity assumption as well as other highly technical conditions. Nonetheless, we agree that studying the particle discretization is a natural and important direction of future work.\n\n**(Q3) Given the cost of sampling from the Gaussian process, are there efficient approximations that preserve the escape guarantees?**\n\nOne natural strategy is to approximate the Hessian stochastically, for instance by computing it using a subsample of the data or particles. In the finite-dimensional setting, perturbation-based methods near saddle points have been extended to SGD [L19]. Under suitable assumptions on the variance, such stochastic approximations may retain the desired escape guarantees.\n\nIn addition, in finite dimensions, the Hessian-vector product can be used to run negative curvature descent methods [LY17,CR19] without computing the full Hessian inverse as in the usual second-order methods. Such methods can also be converted into first-order algorithms such as Neon2 [AZL18]. These ideas may be leveraged to develop a more efficient or even fully first-order version of PWGF.\n\n[GWB25] Margalit Glasgow, Denny Wu and Joan Bruna. Propagation of Chaos in One-hidden-layer Neural Networks beyond Logarithmic Time. arXiv:2504.13110, 2025.\n\n[MMN18] Song Mei, Andrea Montanari, and Phan-Minh Nguyen. A mean field view of the landscape of two-layer neural networks. Proceedings of the National Academy of Sciences, 115(33):E7665–E7671, 2018.\n\n[MMM19] Song Mei, Theodor Misiakiewicz, and Andrea Montanari. Mean-field theory of two-layers neural networks: dimension-free bounds and kernel limit. COLT, 2019.\n\n[DBDFS20] Valentin De Bortoli, Alain Durmus, Xavier Fontaine, and Umut Simsekli. Quantitative propagation of chaos for sgd in wide neural networks. NeurIPS, 2020.\n\n[CLRW24] Fan Chen, Yiqing Lin, Zhenjie Ren, and Songbo Wang. Uniform-in-time propagation of chaos for kinetic mean field langevin dynamics. Electronic Journal of Probability, 29:1–43, 2024.\n\n[SWN23] Taiji Suzuki, Denny Wu, and Atsushi Nitanda. Convergence of mean-field langevin dynamics: Time and space discretization, stochastic gradient, and variance reduction. NeurIPS, 2023.\n\n[KZC+24] Yunbum Kook, Matthew S Zhang, Sinho Chewi, Murat A Erdogdu, and Mufan (Bill) Li. Sampling from the mean-field stationary distribution. arXiv preprint arXiv:2402.07355, 2024.\n\n[AZL18] Zeyuan Allen-Zhu and Yuanzhi Li. NEON2: Finding Local Minima via First-Order Oracles. NeurIPS, 2018.\n\n[CGT11a] Cartis, C., Gould, N.I.M., Toint, Ph.L. Adaptive cubic regularisation methods for unconstrained optimization. Part I: Motivation, convergence and numerical results. Math. Program. 127, 245–295, 2011.\n\n[CGT11b] Cartis, C., Gould, N.I.M., Toint, Ph.L. Optimal Newton-type methods for nonconvex smooth optimization problems. Technical report ERGO 11-009, School of Mathematics, University of Edinburgh.\n\n[CR19] F. E. Curtis and D. P. Robinson. Exploiting negative curvature in deterministic and stochastic optimization. Mathematical Programming: Series A and B, Volume 176, Issue 1-2, 69-94, 2019.\n\n[CRS17] F. E. Curtis, D. P. Robinson, and M. Samadi. A trust region algorithm with a worst-case iteration complexity of $o(\\epsilon^{−3/2})$ for nonconvex optimization. Mathematical Programming, 162(1-2):1–32, 2017.\n\n[L19] Zhize Li. SSRGD: Simple Stochastic Recursive Gradient Descent for Escaping Saddle Points. NeurIPS, 2019. \n\n[LY17] Mingrui Liu, Tianbao Yang. On Noisy Negative Curvature Descent: Competing with Gradient Descent for Faster Non-convex Optimization. arXiv:1709.08571, 2017."}, {"review_id": "8KWfZkRBGI", "reviewer": "Reviewer_hSh6", "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.", "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.", "rating": 5, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 2, "originality": 2, "strengths_and_weaknesses": "**strength**:\n1. 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.\n2. The writing is clear and well-structured.\n\n**weakness**:\n1. The proposed algorithm requires computing Hessian or Hessian-vector products for each particle or for the distribution, which is computationally expensive.\n2. 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.", "was_revised": true, "final_justification": "The authors have provided detailed responses and additional experiments, which largely addressed my concerns.", "rebuttal": "Thank you for the insightful review! Our responses are as follows.\n\n**(Weakness 1) The proposed algorithm requires computing Hessian or Hessian-vector products which is computationally expensive.**\n\nOur primary goal is to initiate a theoretical investigation of perturbative methods for nonconvex probability distribution optimization, analogous to how perturbed GD provably achieves second-order optimality in finite dimensions, and hence we did not focus on computational optimality. Indeed, the use of Hessian-informed perturbations is natural: PGD in $d$ dimensions works with uniformly random perturbations since there will be a nonzero correlation with the desired escape direction with high probability, resulting in a $\\text{polylog}(d)$ rate. However, this argument breaks down when $d=\\infty$ (and also due to the non-Euclidean nature of Wasserstein space/WGF), and we need directed noise to efficiently escape the saddle. Our algorithm provides a first-of-its-kind construction of a concrete Gaussian process-based perturbation which provably attains second-order optimality in the challenging infinite-dimensional setting.\n\nNonetheless, there are ways to improve the computational cost of the process. In finite dimensions, algorithms such as NEON2 [AZL18] can convert Hessian-vector product computations to first-order approximations. This idea may be leveraged to derive a fully first-order WGF algorithm with provable second-order optimality.\n\n**(Weakness 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.**\n\nWe have **added new experiments and a detailed landscape analysis** of the matrix factorization problem using a mean-field neural network (Example 1) in Appendix F.2 and G which suggest the presence of strict benignity and show that noise injection helps escape the initial critical point. Since NeurIPS this year does not allow updating the submission pdf nor adding anonymous links during rebuttal, we summarize the results in text here.\n\n* Experimental details: The input and output dimensions are $l=15,k=5$. We approximate the measure using $400$ neurons and generate $800$ i.i.d. input data points $z$ from the standard normal distribution $\\mathcal{N}(0,1)$ for each coordinate. Similarly to the ICFL case, SGD was used in the optimization process and the learning rate $\\eta = 10^{-6}$. We set $k_{\\mathrm{thres}} = 100, ~ \\eta_p = 3\\times 10^{-3}, ~ F_{\\mathrm{thres}} = 10^{-2}$. We report the mean and standard deviation over 10 runs, with the corresponding error bars.\n* We find that both Hessian and isotropic noise injection methods achieve faster objective reduction and exhibit earlier peaks in the gradient norm, demonstrating a more efficient escape from the initial critical point. In contrast, the perturbation-free method tends to stagnate for longer periods. In particular, the Hessian method shows the best performance, although the performance of isotropic noise is comparable.\n\n**(Question) Given that it frequently computes the Hessian, have the authors considered incorporating second-order methods?**\n\nIt seems nontrivial to directly apply second-order methods to nonconvex distributional optimization as it is unclear how to minimize even the local regularized second-order approximation of the loss functional; the kernel we compute for PWGF cannot be directly used for this purpose, and we are unaware of any existing works which study this problem. (A Wasserstein Newton's flow is studied in [WL20], but only obtain local convergence guarantees assuming strong convexity, and do not consider the nonconvex/saddle problem.) Nonetheless, second-order methods represent a potential alternative approach from our paper. We have provided a review of Hessian-based second-order methods in Euclidean space in our response to Reviewer cKxd, which may be leveraged for this purpose.\n\n[AZL18] Zeyuan Allen-Zhu and Yuanzhi Li. NEON2: Finding Local Minima via First-Order Oracles. NeurIPS, 2018.\n\n[WL20] Yifei Wang and Wuchen Li. Information Newton’s flow: second-order optimization method in probability space. arXiv:2001.04341, 2020."}, {"review_id": "60LEsYJIh6", "reviewer": "Reviewer_cKxd", "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.", "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?\n\nI will raise my score if more experiments can be provided.", "limitations": "No", "rating": 5, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 2, "originality": 2, "strengths_and_weaknesses": "Strengths:\n\n- The paper is well written and presents a detailed theoretical framework.\n\n- The convergence analysis is technically sound. \n\nWeaknesses:\n\n- Literature on the benefit of the Hessian-based saddle-point escape method, even in the Euclidean spaces, is missing. \n\n- Although the authors provide examples of non-convex objective functions that exhibit strict benignity, they are far from practice.", "was_revised": true, "final_justification": "The authors provided a detailed literature review on Hessian-based methods and interesting experiments on the matrix factorization problem using a mean-field neural network.", "rebuttal": "Thank you for the insightful suggestions, they have helped us greatly to improve the manuscript! Below are our responses.\n\n**Weakness 1:**\n\nWe have added an overview of Hessian-based methods in finite dimensions and related works in the manuscript.\n\n* **Hessian-based methods:** The most basic second-order optimization method that utilizes the Hessian of the loss is Newton's method, which however may not converge in general. Global complexity guarantees were first established in [NP06], which studied a cubic regularization of Newton method. More sophisticated methods such as ARC [CGT11a,CGT11b] minimize an adaptive cubic-regularized second-order approximation of the loss and achieve an iteration complexity of $O(\\epsilon^{-3/2})$ for nonconvex optimization, which has been shown to be optimal in certain settings. More recently, [CRS17] proposed the TRACE algorithm which solves a regularized trust region problem at each step, also achieving an iteration complexity of $O(\\epsilon^{-3/2})$. These have been adapted to constrained optimization in e.g. [NR20].\n* Alternatively, the Hessian-vector product can be used to run negative curvature descent methods [LY17,CR19] without computing the full Hessian inverse. Such methods can also be converted into first-order algorithms such as Neon2 [AZL18]. See the discussion in the cited works for a more comprehensive overview. We remark that it is unclear how to adapt these methods to the distributional setting as optimizing even a single step of the second-order approximation of the loss is in general intractable. Nevertheless, these represent potential alternative approaches to achieve second-order optimization in Wasserstein space.\n\n* **Perturbation methods:** since PWGF is primarily first-order gradient descent combined with random perturbations, the most relevant algorithms in finite dimensions are perturbed gradient descent and its variants, which we cover in the introduction.\n\n* **Noisy methods:** noise-injected updates such as stochastic gradient Langevin dynamics are also effective in optimizing particle systems with nonconvex objectives. However, the randomness vanishes in the mean-field limit since the dynamics converges to the *deterministic* WGF w.r.t. the lifted objective with entropy regularization, which does not help if the resulting landscape is nonconvex on Wasserstein space. Our method can be seen an effective way to inject noise in infinite dimensions using directed Gaussian processes.\n\n\n**Weakness 2/Question:** \n\nWe have **added new experiments and a detailed landscape analysis** of the matrix factorization problem using a mean-field neural network (Example 1) in Appendix F.2 and G which suggest the presence of strict benignity and show that noise injection helps escape the initial critical point. Since NeurIPS this year does not allow updating the submission pdf nor adding anonymous links during rebuttal, we summarize the results in text here.\n\n* Motivation: Classical optimization problems such as matrix factorization and tensor decomposition are known to be affected by saddle points. These problems arise in practical scenarios such as image reconstruction, object detection via matrix sensing, and recommendation systems [ZHBL10, ZKLR13, Y08]. In Euclidean space, the landscape of the objective function for matrix factorization has been extensively studied [GHJY15, GJZ17]. In particular, it is known that saddle points exist, and that all non-saddle stationary points are global optima for these types of objectives. (We perform the corresponding mean-field landscape analysis in the revised Appendix G.) A major concern in such situations is that gradient updates can get trapped near saddle points. For example, passing through neighborhoods of saddle points using gradient-based methods can take arbitrarily long time [DJL+17]. Therefore, escaping saddle points is crucial in such settings. \n* Experimental details: The input and output dimensions are $l=15,k=5$. We approximate the measure using $400$ neurons and generate $800$ i.i.d. input data points $z$ from the standard normal distribution $\\mathcal{N}(0,1)$ for each coordinate. Similarly to the ICFL case, SGD was used in the optimization process and the learning rate $\\eta = 10^{-6}$. We set $k_{\\mathrm{thres}} = 100, ~ \\eta_p = 3\\times 10^{-3}, ~ F_{\\mathrm{thres}} = 10^{-2}$. We report the mean and standard deviation over 10 runs, with the corresponding error bars.\n* We find that both Hessian and isotropic noise injection methods achieve faster objective reduction and exhibit earlier peaks in the gradient norm, demonstrating a more efficient escape from the initial critical point. In contrast, the perturbation-free method tends to stagnate for longer periods. In particular, the Hessian method shows the best performance, although the performance of isotropic noise is comparable.\n\nWe thank the reviewer again for their time and humbly ask to consider updating their score in light of the added results.\n\n**References**\n\n[AZL18] Zeyuan Allen-Zhu, Yuanzhi Li. NEON2: Finding Local Minima via First-Order Oracles. NeurIPS, 2018.\n\n[CGT11a] Cartis, C., Gould, N.I.M., Toint, Ph.L. Adaptive cubic regularisation methods for unconstrained optimization. Part I: Motivation, convergence and numerical results. Math. Program. 127, 245–295, 2011.\n\n[CGT11b] Cartis, C., Gould, N.I.M., Toint, Ph.L. Adaptive cubic regularisation methods for unconstrained optimization. Part II: Worst-case function- and derivative-evaluation complexity. Math. Program. 130, 295–319, 2011.\n\n[CR19] F. E. Curtis and D. P. Robinson. Exploiting negative curvature in deterministic and stochastic optimization. Mathematical Programming: Series A and B, Volume 176, Issue 1-2, 69-94, 2019.\n\n[CRS17] F. E. Curtis, D. P. Robinson, and M. Samadi. A trust region algorithm with a worst-case iteration complexity of $o(\\epsilon^{−3/2})$ for nonconvex optimization. Mathematical Programming, 162(1-2):1–32, 2017.\n\n[LY17] Mingrui Liu, Tianbao Yang. On Noisy Negative Curvature Descent: Competing with Gradient Descent for Faster Non-convex Optimization. arXiv:1709.08571, 2017.\n\n[NP06] Y. Nesterov and B. T. Polyak. Cubic regularization of newton method and its global performance. Mathematical Programming, 108(1):177–205, 2006.\n\n[NR20] M. Nouiehed and M. Razaviyayn. A Trust Region Method for Finding Second-Order Stationarity in Linearly Constrained Nonconvex Optimization.\nSIAM Journal on Optimization 30(3):2501-2529, 2020.\n\n[ZHBL10] Bo Zhao, Justin P. Haldar, Cornelius Brinegar, Zhi-Pei Liang. Low rank matrix recovery for real-time cardiac MRI. In 2010 IEEE International Symposium on Biomedical Imaging: From Nano to Macro, pp.996–999. IEEE, 2010.\n\n[ZKLR13] Wenbin Zou, Kidiyo Kpalma, Zhi Liu, Joseph Ronsin. Segmentation driven low-rank matrix recovery for saliency detection. In 24th British machine vision conference (BMVC), pp.1–13, 2013. \n\n[Y08] Koren, Yehuda. Factorization meets the neighborhood: a multifaceted collaborative filtering model. In Proceedings of the 14th ACM SIGKDD international conference on knowledge discovery and data mining. p.426-434. 2008. \n\n[GHJY15] Rong Ge, Furong Huang, Chi Jin, Yang Yuan. Escaping From Saddle Points—Online Stochastic Gradient for Tensor Decomposition. COLT, 2015.\n\n[GJZ17] Rong Ge, Chi Jin, Yi Zheng. No Spurious Local Minima in Nonconvex Low Rank Problems: A Unified Geometric Analysis. ICML, 2017.\n\n[DJL+17] Simon S Du, Chi Jin, Jason Lee, Michael I Jordan, Aarti Singh, Barnabas Poczos. Gradient Descent Can Take Exponential Time to Escape Saddle Points. NeurIPS, 2017."}]

Author rebuttal:

Meta-review: The reviewers found the contribution solid and were positive about the paper. They also made several suggestions that can improve the paper, such as discussing the motivation of the work and including more numerical results. Please take the reviewers' suggestions into account as you prepare the camera-ready version of the paper.
Decision: Accept (poster)