Continuous-time Riemannian SGD and SVRG Flows on Wasserstein Probabilistic Space
Authors: Mingyang Yi
Bohan Wang
Venue: NeurIPS 2025 poster ()
Keywords: Optimization on Manifold; SGLD; Continuous Gradient Flow
Primary area: optimization
TL;DR: In this paper, we propose the continuous SGD/SVRG flow of minimizing KL divergence on Wasserstein space

Abstract: Recently, optimization on the Riemannian manifold have provided valuable insights to the optimization community. In this regard, extending these methods to to the Wasserstein space is of particular interest, since optimization on Wasserstein space is closely connected to practical sampling processes. Generally, the standard (continuous) optimization method on Wasserstein space is Riemannian gradient flow (i.e., Langevin dynamics when minimizing KL divergence). In this paper, we aim to enrich the family of continuous optimization methods in the Wasserstein space, by extending the gradient flow on it into the stochastic gradient descent (SGD) flow and stochastic variance reduction gradient (SVRG) flow. By leveraging the property of Wasserstein space, we construct stochastic differential equations (SDEs) to approximate the corresponding discrete Euclidean dynamics of the desired Riemannian stochastic methods. Then, we obtain the flows in Wasserstein space by Fokker-Planck equation. Finally, we establish convergence rates of the proposed stochastic flows, which align with those known in the Euclidean setting.

Continuous-time Riemannian SGD and SVRG Flows on Wasserstein Probabilistic Space
Mingyang Yi!; Bohan Wang²
Renmin University of China 2 Alibaba Group yimingyangōrue. edu bhwangfy0gmail.com com
Abstract
Recently, optimization on the Riemannian manifold have provided valuable insights the optimization community. In this regard, extending these methods to the Wasserstein space is of particular interest, since optimization on Wasserstein space is closely connected to practical sampling processes. Generally, the standard (continuous) optimization method on Wasserstein space is Riemannian gradient flow (i.c., Langevin dynamics when minimizing KL divergence). In this paper, we aim to enrich the family of continuous optimization methods in the Wasserstein space, by extending the gradient flow on it into the stochastic gradient descent (SGD) flow and stochastic variance reduction gradient (SVRG) flow. By leveraging the property of Wasserstein space, we construct stochastic differential equations (SDEs) to approximate the corresponding discrete Euclidean dynamics of the desired Riemannian stochastic methods. Then, we obtain the flows Wasserstein space by Fokker-Plane equation. Finally, we establish convergence rates of the proposed stochastic flows, which align with those known in the Euclidean setting.
Introduction
As a valuable extension of Euclidean space optimization, generalizing the parameter space to Riemannian manifold-such as a matrix manifold [ a probability measure space [10]-has greatly enriched the toolbox of the optimization community. Technically, optimization on manifold can be derived from Euclidean techniques by defining corresponding Riemannian gradients and transport rules [1]. For example, several widely used gradient- -ased optimization methods have been generalized to this setting, including gradient descent (GD) [53, stochastic gradient descent (SGD)
[4) and stochastic variance reduced gradient (SVRG) 52
On the other hand, optimization over probability measure spaces has attracted considerable attention due to its connection to sampling processes [17]. Interestingly, for a specific type of probability measure space-thh the second-order Wasserstein technique from Riemannian optimization can be directly applied [9]. owing to its manifold-like geometric structure. For example, minimizing the KL divergence via Riemannian gradient flow is equivalent to employing Langevin diffusion 35 standard method for sampling from target distribution. Therefore, advancing Riemannian optimization in the Wasserstein space may lead the development of novel sampling techniques.
To this end, we focus on Riemannian SGD 25 20 and SVRG 33 in Wasserstein space, which standard stochastic optimization methods known for their lower computational complexity in Euclidean settings [5]. Specifically, we investigate the continuous counterparts (flows) of these two previously unexplored Riemannian methods, as continuous formulation provide powerful
Conresponning to: Mingyang Yi and Bohan Wang
39th Conference on Neural Information Processing Systems (NcurIPS 2025).

---

framework for analyzing the properties of optimization algorithms 13 14]. In Euclidean space. such continuous stochastic flows are characterized by SDEs [2 20 not readily generalize to Riemannian manifolds, as their definitions rely on Brownian motion-
33 However, these SDEs do
concept that is considerably more complex to define on Riemannian manifolds [38]
In general. continuous optimization methods are derived by taking the limit of the step size in corresponding discrete optimization dynamics (e.g., transitioning from GD to gradient flow [40]). Naturally, we seek to apply this approach to discrete Riemannian SGD and SVRG 52). Unfor- tunately, such an extension nontrivial for two reasons: (1) the linear structure that underpins the aforementioned continuous dynamics does not necessarily exist on manifolds; (2) describing the randomness inherent in stochastic methods challenging in a manifold setting. Fortunately, the dynamics of probability measures in Wasserstein space can be equivalently described by the dynamics of random vectors Euclidean space. By taking the step size to zero, the corresponding discrete dynamics in Euclidean space converge to an SDE, which characterizes the evolution of probability measures Wasserstein space through the Fokker-Planck (F-P) equation [32]
Through this approach, we successfully establish continuous Riemannian SGD and SVRG flows in the Wasserstein space for minimizing the KL divergence intended. Notably, the existing MCMC geethods-stochsiie gradient Langevin dynamics [49] and stochastic variance reduction Langevin dynamics 7 are precisely the discrete counterparts of the two proposed flows. Furthermore, under appropriate regularity conditions, we prove convergence rates for these continuous stochastic flows. Specifically, for non-convex problems, the convergence rates (measured by the first-order stationarity criterion
( of the Riemannian SGD flow and Riemannian SVRG flow are O(1 ((1/VT) and ((N2/3 /T), respectively, where N denotes the number of component functions. Additionally, under an extra Riemannian Polyak-Lojae Lojjsi--jjsiiiiizz (PL) inequality [23] [ equivalent to the log-
Sobolev inequality [44 the two methods achieve global convergence rates of orders ( /T) and
O(e D(TN2//N2/8) respectively, matching their Euclidean counterparts (33
2 Related Work
Riemannian Optimization. Unlike continuous methods Riemannian manifolds (.. discrete methods have been extensively studied. Examples include Riemannian GD [] (3). Riemannian
Nesterov-type methods [54] 28 26). Riemannian SGD 4 33), Riemannian SVRG [52), and some other Algorithms 3 55
50 12 As established in the literature, the standard approach in Euclidean space for linking discrete dynamics their continuous counterparts involves taking the limit to the
step size in the discrete dynamics, which yields a differential equation [43] 40 (7]. However, this
derivation does not extend directly to arbitrary Riemannian manifolds. Such extrapolation has only been achieved in specific spaces, such as the Wasserstein space, where the resulting curve satisfies the Fokker-Planck equation 40), thereby establishing connection between the Wasserstein and Euclidean spaces. Nevertheless, only a limited number of discrete optimization methods in the Wasserstein space have been generalized to their continuous counterparts for example, gradient flow [. Nesterov accelerated flow [48), and Newton flow [47.. Unfortunately. these extrapolation techniques have not been applied to stochastic dynamics. Moreover, their methodologies are restricted to specific algorithms, in contrast to the more general framework proposed in this paper.
Stochastic Sampling. Standard sampling methods, such as MCMC 22 typically construct (stochastic) dynamics that converge to the target distribution, with convergence measured by probability distance or divergence. Consequently, sampling can be viewed as an optimization prob- lem in the space of probability measures (10]. However, existing literature has primarily focused on
discrete Langevin dynamics (35) and their stochastic variants 449
57 24], particularly ana- lyzing their convergence rates under different criteria-such KL divergence 8 R' e'eyii divergence
15 56
7
I1 30), or Wasserstein distance [66.. By contrast, exploring their connections with continuous Riemannian optimization methods, as undertaken in this paper, has received relatively little attention.
3 Preliminaries
supported on X' =Rd, where P the Wasserstein metric space endowed with the second- order Wasserstein distance 45) (abbreviated as Wasserstein distance), defined as W (T, ())

---

(nf f || dII x. y), where [TT,, () denotes the set of joint probability measures with marginals π and μ., respectively. Notably, the Wasserstein space possesses a geometric structure analogous to that of a Riemannian manifold (9] allowing us to apply Riemannian optimization techniques. We now introduce several key definitions. As in TR. the Wasserstein space (similar to Riemannian manifold) is equipped with an "inner product", resulting the Riemannian metric (.r) : P x T + IR. Here, T₂ P is the tangent space of P at т (see [9]. Section 1.3). Using this Riemannian metric, we define the Riemannian (Wasserstein) gradient gradF(I) E Tr. of function F : P - IR. as the unique element satisfying lim F(mt)= (mo) (grad (mo), 1o) for every curve in P with tangent vector E TTo To. The transportation of Tú E P in T.P determined by the exponential map Exp, TrP - P. Then, the discrete Riemannian GD )m) with learning rate 17 is defined as
Tn+1
Exp=.. |- (gn)d
(1)
to minimize F(x). We refer readers for more details about this dynamies to illustration of Riemannian GD is in Figure
6
53
An
1x+1 Expx, ngrad F F(X,)) ]
Next, we instantiate these definitions in the context of the Wasserstein space, First, a curve ne in P is characterized by the continuity equation (also known as the Fokker Plancck equation) (3):
grad P(
Omg 01
+ V. (TTTV)) = O,
(2)
T
where nH E P and v, E Tx P 9 vector field mapping X' + X. The Fokker-Planck equation implies that T. describes the probability distribution of the stochastic ordinary differential equation (SODE) dx, = un ddt, where the randomness originates from the initial condition æg- Consequently, the tangent vector to the curve Tr, i.e., the direction of curve is given by v, in the F-P equation (2). Besides, for u, v E T=P the Riemannian metric [10] in Wasserstein space (u, v)= [u,,, d') and || 2 = (u, u)),, where (::+) is the inner product in Euclidean space.
P
Figure I: Riemannian GD
Finally, the exponential map in Wasserstein space is
[xp,[u] (w + +id)
where (v - id) R the probability distribution of random variable æ + v(xr) with æ क (9) page 44). In this paper, we mainly explore minimizing the KL divergence [44] to a target probability measure f
min F(x) min (- 1²) TeP / log di
dm.
(4)
du
Here, the target distribution μ assumed to satisfy μ (x))) for some potential function (x)) [10]. For the minimization problem 4 global convergence results have been established under
specific regularity conditions, such as the Pog-Sobolev inequality 9 44 For the KL divergence F(T)) = (TT(TT μ), this inequality takes the form
DT 11) S 27 V di
1 2
du
Igrad DKL (T 11)1
(5)
du
for some 0 and all π, where the equality follows from Proposition 2 In what follows, we write grad, || μ³ as gKLd || (u) to simplify notation. In fact, the log-Sobolev inequality in the Wasserstein space generalizes the PL inequality [23 [1]. thereby guarantecing global convergence of optimization methods. Further details can be found in Section of the Appendix.
In this paper, we need the following lemma from 29 42 which connects the SODE and SDE.
Lemma 1. 29 42
The SDE dx; = b( t)dt + G( t)dW,, has the same density with SODE
1 [G(a,,t) (a))))(((.))))) o) 1/2 G( (x,, (aot)))))(( log (xx)ddt
da: b( t)
(6)
where me is the corresponded probability measure of xz-
As shown by this lemma and (2), the direction of n. can be directly linked SDE, For further details on these preliminaries, we refer readers to [ 6 9
We simplify log (dm,/ (x:)) log ITe (æ:) if' there is no obfuscation in sequel.
3

---

4 Riemannian Gradient Flow
In this section, we investigate the continuous gradient flow for minimizing the KI. divergence in the Wasserstein space through the framework of Riemannian manifold optimization. Although this continuous optimization method has been studied previously [28] 40 2 systematic analysis from the perspective of manifold optimization remains lacking. We demonstrate that our approach provides valuable insights for the subsequent development of Riemannian SGD and SVRG flows.
4.1 Constructing Riemannian Gradient Flow
We first calculate the Riemannian gradient of KI. divergence (proved in Appendix A Proposition 1. The Riemannian gradient of ((T) = Dxz ( |l p)) in Wasserstein space is
(radF(#) =u 14) V log dm
(7)
dp
With the defined Riemannian gradient of the KL divergence and exponential map we can implement the discrete Riemannian gradient descent as in (].
As noted in Section 1 we aim to establish correspondence between discrete dynamics and their continuous counterparts. In Euclidean space, the GD dynamics (In+1 2n,/n1 = ((x,) leads to an ODE (gradient flow) dæ; - -VF(x;) in the limit as 1) - 0. However, this approach does not directly extend to the manifold setting, since the Riemannian GD dynam- ies n in ①T) does not induce a linear structure. Fortunately, the probability measure Fn corre- sponds to random vectors Consequently, the dynamics induced by Riemannian GD (1) can be used to construct a differ- ential equation the limit 7 - O. The corresponding Fokker-Planck equation for this differ- ential equation in the Wasserstein space then yields the gradient flow. The conceptual frame- work is illustrated in Figure 2 and the formal result is stated in the following proposition.
E Rdd such that
Tn and In+1 In v log du
des (æx)).
Assumption For probability measure μu, the log μ and V log all Lipschitz continuous with coeffi- cient L1 and L2, respectively.
equstion
O
F-P equation
π discrete dynamics In
X
variable in 8*
variable
dynamics
Proposition 2. Under Assumption 7 and 1 Snn
Figure 2: Our idea to bridge the discrete dynamics to its continuous counterparts.
((11nn), the discrete Riemannian GD mates continuous Riemannian gradient flow
① approxi-
d
Ōt मe VT.rggradDDa ll pJ)
ne V log
(8)
du
by [| 2nn 12] S On). where T. nr in ① Enn
Tan in 8 for x. Zo-
This proposition shows that the discrete Riemannian GD approximates the flow (8) for any finite
product ny in the limit as 17 + O. Therefore, the flow (8) indeed corresponds to the Riemannian gradient flow. Moreover, by combining Lemma with the F-P equation (2). the measure n, (8) is
the distribution of Langevin dynamics, which serves standard stochastic sampling algorithm 335
given by dæ, V log p(x:)dl + v2dW,. This connection is also discussed in [0 40 26
Remark Notably, existing literature (40) has also established that ( corresponds the Rieman- nian gradient flow in the Wasserstein space for minimizing the KL divergence. However, their results are not derived from limiting procedure on the learning rate discrete Riemannian optimization methods, unlike our approach. The authors [40] obtain the curve n. by solving the variational prob-
lem Tn+11 arg min, [DKI. || (1) + W2 (T, /2n/21] in the limit rj - Unfortunately, unlike our method, this approach cannot be extended to the stochastic optimization setting in Section 5 as the aforementioned minimization problem does not readily incorporate stochastic gradients.
4.2 Convergence of Riemannian GD Flow
In this section, we prove the convergence of the Riemannian gradient flow 8 For non-convex problems in Euclidean space, convergence is typically measured by the first-order stationarity criterion
Since μ ex)(-V(xr)). the condition means ((x) and ( are all Lipschitz continuous.

---

Algorithm Discrete Riemannian SGD
Input: Exponential map Exp, initialized To, learning rate 7,, steps M.
1: for O,..... .M do
2: Sample En ~ ク independent with Tni
3: Update Tx+1 Expx,, [-ngrad (e..));
4: end for
5: Return: A.M-
19 6 [1] Accordingly, the Wasserstein space, we analogously analyze the convergence rate of the Riemannian gradient norm Igrad [)]m,²² following the approach of Boumal et al. 6). Balasubramanian et al. Moreover, the Pl. inequality (3 ensures global convergence in Euclidean space; for instance, under this condition, the gradient flow converges exponentially to a global minimum [23). As discussed Section 3 in the Wasserstein space, the Riemannian PL inequality is generalized by the log-Sobolev inequality [4], which likewise guarantees exponential global convergence rate for the Riemannian gradient flow, as stated in the following theorem.
Theorem 1. [[9]] Let 74 follows the Riemannian gradient flow ( then for any T >> 0. we have
T 1."
DKL(TO (b)
IgradD p)ll
(9)
T
Moreover, if the log-Sobolev inequality
5 (Riemannian PL inequality) is satisfied for p,, then
DKL( p) 2x1 (TToooTT || p).
(10)
We prove it in Appendix AA] The convergence rate Riemannian gradient flow is also proved 9 expected. We prove this theorem here to illustrate the criteria convergence rates under different conditions.
and it matches the results in Euclidean setting 43 23
Riemannian Stochastic Gradient Flow
In practice, SGD is preferred over GD due to its lower computational complexity. However, unlike in Euclidean space 20 25), the continuous SGD flow in Wasserstein space has not been explored yet. Next, our goal is to generalize the Riemannian gradient flow (8) to the Riemannian SGD flow.
5.1 Constructing Riemannian Stochastic Flow
The stochastic algorithm is developed to minimize the stochastic optimization problem such that min Rg (m)) min Re[DRz(= [t)],
(11)
where the expectation is taken over & parameterizes set of probability measures jlg (TT), we can get the optima of it by the following proposition proved in Appendix B
For objective
Proposition 3. The global optima of problem
x exp (Eg [log Mg])
Since we can explicitly get the optima μ defined in Proposition 3 the goal of Riemannian SGD flow
should be moving towards Thus, the target distribution becomes y Proposition (3) sequel.
To solve the problem (T)), one may use the standard discrete method, Riemannian SGD
as outlined in Algorithm| 1 By analogy with our derivation of the Riemannian GD flow in Section 4 the Riemannian SGD flow is naturally posited as its continuous counterpart. To derive this flow, we first construct the Euclidean-space dynamics 2,, of Algorithm I Then, following a procedure similar to that in Proposition 2 F-P equation then yields the continuous Riemannian SGD flow. Below is the result.
we approximate the dynamics of x, by a continuous SDE. The corresponding
Assumption 2. For any & and probability measure μig, log jg and V log Me are Lipschitz continuous with coefficient L1 and La respectively,
gKT. Fe ll 11)2 + 0 does not imply D(T,, μ4) + 0 for other probability distances or divergences D(. ), such as the total variation distance. Further details can be found in Balasubramanian et al. (e()))) is assumed to Msg (m) x exp(-Ve(x)).
As be sx)(-Ve), the assumption implies Ve2 and its gradient Lipschitz continuous.
5

---

Proposition 4. Under Assumption 2 let fę(T) DKL(T || pg). the discrete Riemannian SGD Algorithm 7 with 1 I n v ((11/n)) approximates the continuous Riemannian stochastic gradient
flow
0t nt V. ne v log
d " ESCD
" ESOD V log mi
(12)
du
2
by E In Bnn |?| - (n)) where In ~ Fn in Algorithm
1 Bnn ~ aan in for Zn = To. Here EsGD(#) (a) = Eg log (a) VI [g (V (e VI [log Me (x)1)"
(13)
As can be seen, similar to Proposition 2 we approximate the discrete Riemannian SGD with continuous flow (12). so that it is Riemannian SGD flow as desired.
One may observe that the Riemannian "stochastic" gradient flow 12) is deterministic curve in the Wasserstein space P. This is not inconsistent with the randomness in the index &n for n. in discrete Riemannian SGD. This is because the stochasticity introduced by &, when obtaining on (with x,, ~ TT) in Algorithm[]] i implicitly captured in the corresponding &n. In other words, the randomnesses from the random vector æ, itself and the random index εn are fully incorporated into 2n and manifested in its continuous approximation Tnn Kan. where ग.. E P is the deterministic curve [2) in the Wasserstein space. Thus, all randomness in the Riemannian SGD is accounted in this formulation.
Besides, can find that implementing the discrete Riemannian SGD is non-trivial, since it requires Riemannian gradient Vlog (dw,/ /dj).. However, during proving Proposition 4 we show the discrete stochastic gradient Langevin dynamics (SGLD) 49
an+1 Ir 7/V log Me. (#x) 2ne
(14)
approximates D the corresponded dynamics of } in discrete Riemannian SGD (Algorithml
I
deEn (w.).
Cn+1 = wn + v log
(15)
dmn
Thus, in practice, discrete SGLD can be implemented to approximate discrete Riemannian SGD. Furthermore, based on this approximation and Proposition the continuous counterpart æz of (15).
as governed by the Riemannian SGD flow 12 (Lemma satisfies the SDE
dp (xx) + (T ddW, drt.
dx: = V log
(16)
is also the continuous limit of the discrete SGLD (4). This establishes a connection between the Riemannian SGD flow and discrete SGLD-that is, the Riemannian SGD flow in the Wasserstein space corresponds precisely to continuous SGLD. Consequently, our Riemannian SGD flow provides powerful framework for analyzing discrete SGLD or Riemannian SGD. For example, combining Proposition 4 and Theorem 2 implies the convergence rate of discrete Riemannian SGD,
5.2 Convergence of Riemannian SGD Flow
Next, we examine the convergence rate of Riemannian SGD flow. As in Section 4 our analyses are respectively conducted with/without log-Sobolev inequality.
Theorem 2. Let na follows the Riemannian SGD flow (1) and μ defined in Proposition 3
Under we have
Dge( (To m)) DiL then by taking 1 T(4d1 (444:1b++d++1)1777
Assumption 2 ifT Z 2a12
Iggra H))
4DKL(To ll ())
1
(17)
TT
T/T
VT
Besides that, if
5 is satisfied for p. 17 1/7/Ta with 0 < a < 1, and T 2 (8L7/7)
1/a
then
1 (RI (A) S TT [AdL] (d + 1) 133
(18)
T
"The approximation is verified by noting E., [V V log || = -E. |a for continuous test function f. and combining Taylor's expansion, please check Appendix B for more details.

---

The proof of this theorem is provided in Appendix B.2 Notably, the presence of in the SDE leads to convergence rate of order D(1/vT) for the Riemannian SGD flow Under the log- Log-Sobolev inequality, a global convergence rate of (1//T) can be established (by taking a - 1). Therefore, the convergence rates proved in Theorem 2 align with those of the continuous SGD flow in Euclidean
space 19 (3), as demonstrated in Appendix
B..
Remark 2. During the proof to Theorem 2 we assume log Me is Lipschitz continuous. The assumption in 122. which is the standard "bounded variance" assumption in optimization on Euclidean space [19].
can be relaxed as [,,,, log u(xx)) -
Ee V log [(*+) S o², constant and
In the remainder of this section, we further demonstrate the tightness of the derived convergence rate for the Riemannian SGD flow through the following example. Although this example does not satisfy the Lipschitz continuity condition on log jgg in Assumption [2] it does satisfy the condition of bounded variance discussed in the preceding remark. Consequently, the results in Theorem 2
rremain valid.
Example 1. Let Fi ~ (E, I), with & (E1.... .EN}, maxisiSN ||E,|| S C for a constant C.
E, 5,/N, which is the target measure of Riemannian SGD flow. Then, we have ESGD t Σ レーー (E, (E; E(E)
Due to Proposition 33. we have μ ~ NE, T) with F((] = E Σ,
Var(E).
The assumptions (including log-Sobolev inequality) in Theorem 2 are all satisfied (see Lemma in Appendix). Then, the corresponding SDE of Riemannian SGD flow (12 is
dæ: = (a: E) + ( (vTVar² (E), v2
(19)
which has the following closed-form solution æ; = E + e (00 - E) + VTar ( -1 + e-(w(() -1 where w((2) for any Øo.
2 and w!2)) are standard independent Brownian motions. By taking t = nJT,
ZeT N E + e "T (2o E), (1 - ( ()Va(((()))))) 11)
(20)
Here, e -- N 0 for large T, and 1) = O(L/TT) with 0 < <a<< < Then, 2nTT Ne N Var(E)) I). The Riemannian gradient grad D ( |l Ma) ( (,,
E), which indicates
BgrradDx1(t ( Il ق)ه On the other hand, the KL divergence between Gaussian measures 36 TTT, / can be calculated as
= ur ((I v/Var(E)22) - - v)Va((()22) = 0 (T
(21)
3 [og II
(RL II (4)
(22)
As can be seen, the convergence rates in and (2) are consistent with the proved convergence rates in Theorem 2 under = 1/2 and convergence rates Theorem are sharp.
1, respectively. Thus, the example ( indicates the
Riemannian Stochastic Variance Reduction Gradient Flow
Next, we extend discrete algorithm SVRG 21 to its continuous counterpart in Wasserstein space.
6.1 Constructing Riemannian SVRG Flow
In practice, the objective in [I] often takes the form of finite sum. where & is uniformly distributed over -..... EEN. so that the objective becomes ming F(T) = min, Ee (T)) = minz 1V Σラ Je, (T). Clearly, computing gradF(m) requires O(N) operations. The convergence rates in Theorems and 4 imply that achieving Igrad (#t 4)) < for some that is, reaching an e-stationary point- requires computational complexities of ONNe=)) and O 3), re- spectively") Therefore, for large N, Riemannian SGD flow offers improvement over Riemannian GD flow.
*We emphasize that constraints imposed the learning rate 7) in Theorem hence, the convergence rate of Riemannian GD remains as stated in Theorem 'The computational complexity of continuous optimization is evaluated by implementing the corresponding discrete algorithm. Further details are provided in Appendix B.2
7

---

Algorithm 2 Discrete Riemannian SVRG
Input: Exponential map Expx, initialized no. learning rate 7, epoch I, steps M of each epoch.
1: Take mg m0:
2: for i = -0 0,..,, 1 do
3: Compute gradF(# N Ej= grad (m))
4: for = 0..... .M- do
5: Uniformly sample En E ffi,... .N3 independent with n
6:
Update Thh11 Expx; [-n[grad (T//) - r ab
(gradfe; (Tr) gradF ((r))))
7: end for
8:
TO ++1
-hhi
9:
end for
10:
Return: TN.
In Euclidean space, considerable efforts have been devoted to further improving computational complexity, with methods such as SVRG 21 [56), SPIDER [8 55), and SARAH [31]. Among these, the double-loop structure of SVRG represents core idea shared across several variance-reduced approaches. For this reason, we focus specifically on SVRG in this paper.
As presented in Section we begin with the discrete Riemannian SVRG method [52] outlined in Algorithm 2 Analogous to the Euclidean setting. line 6 of Algorithm Algorithm@employ eemploy the tangent vector
grad c4 (mh) r (grad (mr) gra)F( which serves as a variance-reducee estimator of the Riemannian gradient gradF(m;,. This constitutes the core idea of SVRG-type algorithms. In this context, the mapping I denotes parallel transport
Specifically, for the loss function (<(T) DR.(T .4g), the difference of Riemannian gradients grad (mp) grad FraaF((ri) (mr)) lies in tangent space T, (. while the exponential map Exp,4, is defined on T P. To reconcile this discrepancy, the parallel transport 1 Tsp P Tmj P is applied, which maps vectors from Tx:
to Txg P while preserving the Riemannian metric, i.e., || (w)llz; ||u|| for any E T
In Wasserstein space, we define I (w) = u o Tn for μ. v E P and u E Tuo where o is the composition operator, and Tx-tH : + satisfies T (x) ~ μ for x ~ π. Thus,
|| | = / || - ||u(Zi=a dm) / ||u T=
(u))dl,
(23)
which is consistent with the definition of parallel transport. Based on the preceding notations, we now proceed construct the continuous Riemannian SVRG flow derived from Algorithm 2 The underlying rationale aligns with the approaches established in Propositions 2 and 4 namely. approximating the discrete Riemannian SVRG updates with deterministic flow in the Wasserstein space. The formal result is stated in the following proposition.
Proposition 5. Under Assumption 2 let fT( = DKL(T || (e), the discrete Riemannian SVRG Algorithm 2 with 1 S S ((LL/n)) approximates the Riemannian SVRG flow
0 01 (ax) V. (m))(a))) log
(Vlog
die du
(w:)
/
(T, V. (y,
(w)
(24)
Te, E))RGG x Va log ग., ( x mdyy
2 (x))
in St t t++1, (#t. xx)
Te.. i. Ti since we have E|||"', 2 (;M+))n0
O1)) where x", ~ T; in Algorithm for x = 2o. Here
EsvnG (U, x) = Eg (V log uu (x) - V log ug (3) VE, (o
ug - VEg log (z()) (g) + VEg [log un [v) VE, (og
(25)
log
As in ④ the Riemannian SVRG flow defined in 5 also constitutes a deterministic curve in the Wasserstein space, which can be explained through similar discussion following Proposition 4

---

Consequently, all three continuous flows in the Wasserstein space, (8. 12 and (24) describe deterministic trajectories, irrespective of any randomness in their discrete counterparts.
Moreover, similar to the discussion after Proposition f although the discrete Riemannian SVRG method presents implementation challenges, it can be effectively approximated through both the discrete SVRG Langevin dynamics [56] and the proposed Riemannian SVRG flow in (24). Conse-
quently, our formulation presents valuable analytical framework for both SVRG Langevin dynamics and discrete Riemannian SVRG algorithms. Further details are in Appendix C.I
6.2 Convergence of Riemannian SVRG Flow
In this subsection, analyze the convergence rate of the Riemannian SVRG flow defined in equation (24). The main result, stated in Theorem below, adopts the same notations as those introduced in Proposition 5
Theorem 3. Let T4 follows Riemannian SVRG flow 24 Ia defined in Proposition 33. for sequences 11 11-1 0(1 vq), and n/T TA to run Riemannian SVRG flow for I epochs. Then for any T, if the V log Pe is Lipschitz continuous with coefficient for all g, under proper no, we have
{ti} with A t1 to
2Dxr( p))
Ilgra 11 dt S
(26)
T
T/
By taking n O(N 2/3). the computational complexity of Riemannian SVRG flow of order O(N2/S/ /e) to make minosesat II || 11)) e.
Furthermore, when the log-Sobolev inequality 5 is satisfied for 14, we have
DRr 1) S - DKz(#e ll p),
(27)
and it takes O((N + + (2/23) computational complexity to make DKI(FTT |l (l) S e.
The proof of this theorem is deferred to Appendix C.I As shown, for non-convex problems, the computational complexity required to reach e-stationary point of the Riemannian SVRG flow is O(N2/3 This complexity can be lower than that of Riemannian GD flow, which is O(Ne and Riemannian SGD flow, which is O(r=)), when e is sufficiently small (as clarified in Section 6.1
On the other hand, under the Riemannian PL-inequality (i.e., log-Sobolev inequality), the compu- tational complexities required to achieve (TTT (4) S are O N log e and O - for Riemannian GD and SGD flows. These can be improved by the Riemannian SVRG flow when >- > 1 and Oog 7 < O((NN + +N Besides that, it is worth noting that, matter with/without the log- Sobolev inequality, the proved convergence rates match the results in Euclidean space (33], and computational complexities match the discrete SVRG in Euclidean space (37]
In Theorem 3 the initial distribution mo required to satisfy ET,,, [tr(V2 log (dz:/dp)) (p)))) <
A [Ir(Z [IsvRG)] where A is polynomial function of l. Since no constraint imposed on the order of A₁ which may grow arbitrarily large, this assumption can be readily satisfied in practice. Further technical details are provided in Appendix C.
As in Section 5.2 the convergence results in Theorem do not require Lipschitz continuous log Mç. so that Example is suitable to be analyzed as follows.
As discussed in Section 6.1 the fundamental principle of SVRG lies in reducing the variance of stochastic gradient estimates at each update, thereby accelerating convergence. Interestingly, in Example []] the induced noise covariance ΣsVRG 0, which causes the Riemannian SVRG flow (24) degenerate into Riemannian GD flow (8). This indicates that the variance reduction technique completely eliminates gradient variance in this case. Notably, since V log Me (x) (x {) and v log (u())= = (x {), the corresponding random vector et E Rd in the discrete Riemannian SVRG (line 6 of Algorithm 2 satisfies.
2 = x ( Voog
dpet dri,
ditc dmb (mb) + log (ab) drb
de
dp
(ar)) V log
x + V log d
(a)),
which is exactly the discretion of Riemannian GD flow, but with (1) (instead of (NN as we do not compute & for each n) computational complexity for each update step in line 6 of Algorithm

---

Tay.
Sosile
Scale To
Mia
Mil
Figure 3:
The convergence rates measured by KL divergence and Fisher divergence (Riemannian gradient norm) under different optimization methods.
2 This explains the improved computational complexity of Riemannian SVRG flow. Note that the corresponded SDE of (24) is dæ, -(x, E + 2dW,, with closed-form solution 2nyT È e (xo - }) + e e- -- Then, we can prove the convergence rates of (x || μ)) (ZarT ~M T and its first order criteria are all of order ((e where the global convergence rate matches the result in Theorem
7 Experiments
In this section, we empirically evaluate the proposed sampling algorithms on two examples.
Gaussian. In this case, we set N 100, d 2, Me, (E:r I) with each & NNO,, I). Then, we get the target distribution μ (N,E,,, I). Within this framework, we report the KL divergence DKI(Ty) || μ)) and the Fisher divergence ||grad x (Tyr p)ll² (first order criteria), where the Tyr obtained by Riemannian GD flow, Riemannian SGD flow SVRG flow. In this case, due to Example T. the two criteria can be explicitly estimated. The results are summarized in Figure 3
Mixture Gaussian. In this case, we set N 5, d with Pe, JN( I) + (E I). where & ~ (O, I). Unfortunately, the proposed flows can not be explicitly computed. Therefore, we implement their discrete versions as in Algorithms 22 Then we report the KL divergence (xI(TTTT (1) and Fisher divergence ||grad ( / with स., approximated by the discrete Riemannian GD, SGD, and SVRG Algorithms. Here, the KL divergence and Riemannian gradient are estimated by density estimation in with 1000 independent samples. The results are summarized in Figure 3
As can be seen, under the same update steps, Riemannian SVRG and Riemannian GD have better convergence results than Riemannian SGD (Theorem [ 2 and while under the same gradient computations, Riemannian SGD and Riemannian SVRG (especially Riemannian SVRG) are sharper (see discussion after Theorem 3.. These results are consistent with theoretical conclusions.
8
Conclusion
Based on the principles of Riemannian manifold optimization, this paper investigates continuous Riemannian SGD and SVRG flows for minimizing KL. divergence within the Wasserstein space. We establish convergence rates for these stochastic flows that align with known results in Euclidean settings. Our technical approach involves constructing SDEs in Euclidean space by taking the limit of vanishing step sizes in discrete Riemannian optimization methods, where the corresponding Fokker- Planck equations characterize the desired curves in Wasserstein space. This framework provides new theoretical insights into continuous stochastic Riemannian optimization and demonstrates the utility of continuous methods as analytical tools for studying discrete optimization algorithms.
10

---

References
[1] Absil, P.-A., Mahony, R., and Sepulchre, R. (2009). Optimization algorithms on matrix manifolds. Princeton University Press.
[2] Balasubramanian, K., Chewi, S., Erdogdu, M. A., Salim, A., and Zhang, S. (2022). Towards a theory of non-log-concave sampling: first-order stationarity guarantees for langevin monte carlo. Conference on Learning Theory.
[3] Becigneul, G. and Ganea, O.-E. (2018). Riemannian adaptive optimization methods. In Interna- tional Conference on Learning Representations.
[4] Bonnabel, S. (2013). Stochastic gradient descent on riemannian manifolds. JEEE Trunsactions on Automatic Control, 58(9):2217-2229,
[5] Bottou, L.,, Curtis, F. E., and Nocedal, J. (2018). Optimization methods for large-scale machine learning, SIAM review, 6022):22-311.
[6] Boumal, N., Absil, P.-A., and Cartis, C. (2019). Global rates of convergence for nonconvex optimization on manifolds. IMA Journal of Numerical Analysis, 39(1):1-33.
[7] Chatterji, N., Flammarion, N., Ma, Y., Bartlett, P., and Jordan, M. (2018). On the theory of variance reduction for stochastic gradient monte carlo. In International Conference on Machine Learning.
[8] Cheng, X. and Bartlett, P. (2018). Convergence of langevin meme in kl-divergence. In Algorithmic Learning Theory.
[9] Chewi, S. (2023a). Log-concave sampling. Lecture Notes.
[10] Chewi, S. (2023b). An optimization perspective on log-concave sampling and beyond. PhD) thesis, Massachusetts Institute of Technology.
[11] Chewi, S., Erdogdu, M. A., Li, M., Shen, R., and Zhang. S. (2022). Analysis of langevin monte carlo from poincare to log-sobolev, In Conference on Learning Theory.
[12] Cho, M. and Lee, J. (2017). Riemannian approach to batch normalization. In Advances in Neural Information Processing Systems.
[13] Du, S., Lee, J., Li, H., Wang. L., and Zhai, X. (2019). Gradient descent finds global minima of deep neural networks. In International conference machine learning.
[14] Du, S. S., Hu, W., and Lee, J. D. (2018). Algorithmic regularization in learning deep homoge- neous models: Layers are automatically balanced. Advances in Neural Information Processing Systems.
[15] Dubey, K. A., Reddi, S., Williamson, A., Poczos, B., Smola, J. and Xing, E. P. (2016). Variance reduction in stochastic gradient langevin dynamics. Advances in Neural Information Processing Systems.
[16] Durmus, A. and Moulines, É. (2019). High-dimensine bayesian inference via the unadjusted langevin algorithm. Bernoulli, 25(4A):8854 2882.
[17] Dwivedi, R., Chen, Y., Wainwright, M.J., and Yu, B. (2018). Log-concave sampling: Metropolis- hastings algorithms are fast! In Conference on learning theory.
[18] Fang, C., Li, C. J,, Lin, Z., and Zhang, T. (2018). Spider: Near-optimal non-convex optimization via stochastic path-integratt differential estimator.
[19] Ghadimi, S. and Lan, G. (2013). Stochastic first-and zeroth-order methods for nonconvex stochastic programming. SIAM Journal Optimization, 23(4):2341-236.
[20] Hu, W., Li, C. J., Li, I.., and Liu, J.-G. (2019). On the diffusion approximation of nonconvex stochastic gradient descent. Annals of Mathematical Sciences and Applications, 4(1).

---

[21] Johnson, R. and Zhang, T. (2013). Accelerating stochastic gradient descent using predictive variance reduction. In Advances in neural information processing systems, pages 315-323.
[22] Karatzas, and Shreve, S. (2012). Brownian motion and stochastic calculus, volume 1 113. Springer Science & Business Media.
[23] Karimi, H., Nutini, J., and Schmidt, M. (2016). Linear convergence of gradient and proximal- gradient methods under the polyak-tojasiewicz condition. In Machine Learning and Knowledge Discovery in Dutabasex: European Conference.
[24] Kinoshita, Y. and Suzuki, T. (2022). Improved convergence rate of stochastic gradient langevin dynamics with variance reduction and its application to optimization.
[25] Li, Q., Tai, C., and Weinan, (2017). Stochastic modified equations and adaptive stochastic gradient algorithms. In International Conference Machine Learning.
[26] Liu. C., Zhuo, Cheng. P.. Zhang, R., and Zhu, (2019). Understanding and accelerating particle-based variational inference. In International Conference on Machine Learning.
[27] Liu, Q. (2017). Stein variational gradient descent as gradient flow.
[28] Liu, Y., Shang. F., Cheng, J., Cheng, H., and Jiao, L. (2017). Accelerated first-order methods for geodesically convex optimization on riemannian manifolds. In Advances in Neural Information Processing Systems.
[29] Maoutsa, D., Reich, S., and Opper, M. (2020). Interacting particle solutions of fokker fokker-planck equations through gradient- density estimation. Entropy, 22(8):802.
[30] Mossavi-Hosseini, A., Farghly, T., He, Y., Balasubramanian, K., and Erdogdu, M. A. (2023). Towards a complete analysis of langevin monte carlo: Beyond poincar') 'e inequality. In Conference on learning theory.
[31] Nguyen, L. M., Liu, J., Scheinberg, K., and Takáč, M. (2017). Sarah: novel method for machine learning problems using stochastic recursive gradient. In International Conference on Machine Learning.
[32] Oksendal, B. (2013). Stochastic differential equations: an introduction with applications. Springer Science & Business Media.
[33] Orvieto, A. and Lucchi, A. (2019). Continuous-time models for stochastic optimization algorithms. Advances in Neural Information Processing Systems.
[34] Otto, F. and Villani, C. (2000). Generalization of an incquality by talagrand and links with the logarithmic sobolev inequality. Journal of Functional Analysis, 173(2):361-400.
[35] Parisi, G. (1981). Correlation functions and computer simulations. Nuclear Physics B. 180(3):378-384
[36] Petersen, K. B., Pedersen, M. S., et al. (2008). The matrix cookbook. Technical University of Denmark, 7(15):110.
[37] Reddi, S. J., Hefny, A., Sra, S., Poczos, B., and Smola, A. (2016). Stochastic variance reduction for nonconvex optimization. In International Conference on Machine Learning.
[38] Ren, P. and Wang, F.-. (2024). Ornstein- uhlenbeck type processes on wasserstein spaces. Stochastic Processes and their Applications, 172:104339.
[39] Sam Patterson, Y. W. T. (2013). Stochastic gradient riemannian langevin dynamies on the probability simplex.
[40] Santambrogio, F. (2017). (Euclidean, metric, and Wasserstein) gradient flows: overview. Bulletin of Mathematical Sciences, 7:87-154.
[41] Shiryaev, A. N. (2016). Probability-1, volume 95. Springer.
12

---

[42] Song, Y., Sohl-Dickstein, J., Kingma, D. P., Kumar, A., Ermon, S., and Poole. B. (2020). Score- based generative modeling through stochastic differential equations. In International Conference on Learning Representations.
[43] Su, W., Boyd, S., and Candes, E. (2014). A differential equation for modeling nesterov's accelerated gradient method: theory and insights.
[44] Van Handel, R. (2014). Probability in high dimension. Lecture Notes (Princeton University).
[45] Villani, C. al. (2009). Optimal transport: old and new, volume 338. Springer.
[46] Von Neumann, (1937). Some matrix-inequalitee and metrization of matric space.
[47] Wang, Y. and Li, W. (2020). Information newton's flow: second-order optimization method in probability space. Preprint #tXiv:2001.04341
[48] Wang, Y. and Li, W. (2022). Accelerated information gradient flow. Journal of Scientific Computing, 90:1-47.
[49] Welling, M. and Teh, Y. W. (2011). Bayesian learning via stochastic gradient langevin dynamics. International Conference on Machine Learning.
[50] Yi, M. (2022). Accelerating training of batch normalization: A manifold perspective. In Uncertainty Artificial Intelligence.
[51] Yi, M., Wang, R., and Ma, Z.-M. (2022). Characterization of excess risk for locally strongly convex population risk.
[52] Zhang, H., Reddi, S., and Sra, S. (2016). Riemannian svrg: Fast stochastic optimization riemanmian manifolds.
[53] Zhang, H. and Sra, S. (2016). First-order methods for geodesically convex optimization. In Conference on Learning Theory, pages 1617-1638.
[54] Zhang, H. and Sra, S. (2018). Towards riemannian accelerated gradient methods. Preprint arXiv: wew:10806.2222.
[55] Zhang, J., Zhang, H., and Sra, S. (2018). R-spider: fast riemannian stochastic optimization algorithm with curvature independent rate. Preprint 1it::1811104199
[56] Zou, D., Xu, P., and Gu, Q. (2018). Subsampled stochastic variance-reduced gradient langevin dynamics. In International Conference on Uncertainty in Artificial Intelligence.
[57] Zou, D., Xu, P., and Gu, Q. (2019). Sampling from non-log-concave distributions via variance- reduced gradient langevin dynamics. In International Conference on Artificial Intelligence and Statistics.
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:
Justification:
The main contributions this paper have been clarified Abstract.
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made in the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made the paper and important assumptions and limitations. A No or NA answer to this question will not be perceived well by the reviewers.
13

---

The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize 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 of the work performed by the authors?
Answer: [Yes]
Justification:
The limitation of this paper discussed in the main paper.
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 encouraged to create separate "Limitations" section in 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 on 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, c.g., if the approach was only tested few datasets 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 or 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 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 be used by reviewers as grounds for rejection, worse outcome might be that reviewers discover limitations that aren't 't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an impor- tant role in developing norms that preserve the integrity of the community, Reviewers penalize honesty concerning limitations.
will be specifically instructed
3. Theory Assumptions and Proofs
Question: For each theoretical result, does the paper provide the full set of assumptions and a complete (and correct) proof?
Answer: [Yes]
Justification:
All results are proofed.
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 proofs 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 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 be properly referenced.
4. Experimental Result Reproducibility
14

---

Question: Does the paper fully disclose all the information needed to reproduce the main ex- perimental results 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)?
Answer: [NA]
Justification:
Guidelines:
The answer NA means that the paper does not include experiments. If the paper includes experiments, a 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 a novel architecture, describing the architecture fully might suffice, or if the contribution
specific model and empirical evaluation, may be necessary to either make it possible for others to replicate the model with the same dataset, or provide access to model. In general. releasing code and data often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (c.g., in the case of large language model), releasing of a model checkpoint, other means that are appropriate to the research performed.
While NeurlPS 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 to reproduce that algorithm.
(b) If the contribution is primarily new model architecture, the paper should describe the architecture clearly and fully.
(c) If the contribution is a 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 or instructions for how to construct the dataset).
(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome to describe 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., to registered users), but it should be possible for other researchers have some path 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?
Answer: [NA]
Justification:
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurIPS code and data submission guidelines (https://nips.cm cc/ gablic/guides. Coooe/Sumiiisiooooooo
for more details.
While we encourage the release of code and data, we understand that this might not be possible, so "No" is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for new open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurlPS code and data submission guidelines (https! //nips http://cccppcceeeeeeee/////aaaaa
for more details.
15

---

The authors should provide instructions on data access and preparation, including how access the raw data, preprocessed data, intermediate data, and generated data, 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 in supplemental material (appended the paper) is recommended, but including URLs to data and code is permitted.
6. Experimental Setting/Details
Question: Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Answer: [NA]
Justification:
Guidelines:
The answer NA means that the paper does not 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:
[NA]
Justification:
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, or statistical significance tests, at least for the experiments that support the main claims of the paper.
The factors of variability that the error 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).
It should be clear whether the error bar is the standard deviation the standard error of the mean.
It is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2iggma 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 error 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 the text.
8. Experiments Compute Resources
For each experiment, does the paper provide sufficient information the com- puter resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?
Question:
Answer: [NA]
16

---

Justification:
Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate the type of compute workers CPU 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 as 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 or failed experiments that didn't make into the paper).
9. Code Of Ethics
Does the research conducted the paper conform, in every respect, with the NeurlPS Code of Ethics
Question:
Answer:
[Yes]
Justification:
Guidelines:
The answer NA means that the authors have not reviewed the NeurlPS Code of Ethics.
If the authors answer No, they should explain the special circumstances that require a deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g.. if there 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:
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 or 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.
The conference expects that many papers will be foundational research and tied to 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 is being used as intended but gives incorrect results, and harms following from (intentional unintentional) misuse of the technology.
If there negative societal impacts, the authors could also discuss possible mitigation strategies (c.g., gated release 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 ML).
11. Safeguards
Question: Does the paper describe safeguards that have been put in place for responsible release of data models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
17

---

Answer: [NA]
Justification:
Guidelines:
The answer NA means that the paper poses no 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 require this, but we encourage authors to take this into account and make best faith effort.
12. Licenses 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:
Guidelines:
The answer NA means that the paper does not use existing assets.
The authors should cite the original paper that produced the code package dataset.
The authors should state which version the asset used and, possible, include URL.
The name of the license (e.g., CC-BY 4.0) should be included for each asset.
For scraped data from a particular source (e.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 re-packaged, both the original license and the license of the derived asset (if has changed) should be provided.
If this information is not available online, the authors are encouraged to reach out to the asset's s creators.
New Assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [NA]
Justification:
Guidelines:
The answer NA means that the paper does not release new assets.
Researchers should communicate the details the dataset/codeemoodll as part their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent was obtained from people whose asset 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. Crowdsourcing and Research with Human Subjects
18

---

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 as details about compensation (if any)?
Answer: [NA]
Justification:
Guidelines:
The answer NA means that the paper does not involve crowdsourcing 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 in the main paper.
According to the NeurIPS Code of Ethies, 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 equivalent approval/review based the requirements of your country or institution) were obtained?
Answer: [NA]
Justification:
Guidelines:
The answer NA means that the paper does not involve crowdsourcing research with human subjects.
Depending on 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 we expect authors to adhere to the NeurIPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.
19

---

A Proofs in Section 4
Proposition 1. The Riemannian gradient of F(T)) DxI(TT || p)) in Wasserstein space
= grad (x Fx (4) = V log du
(7)
dp
Proof. For any curve Th E P with To nT, we have,
lim (Te | n4) Drz 1)) DevDRz po(
① DKI(TT |l (4) lt-
Øt
n, 11-0 + / 0t
(1 + log TT)
log pdx I==0
(28)
0t.
(V log π
V log u. volda
vo: log dm dp
Thus we prove our conclusion due to the definition of Riemannian gradient.
Proposition 2. Under Assumption and S S [(L1/n]), the discrete Riemannian GD
approximates continuous Riemannian gradient flow Tt.
d ne log
ōt ne V (FograDDxi(T Il (4)
(8)
du
by Ex||=n - Bng |2] S (n), where un ~ Tn in ① Tnn ~ Tnn in
8 for wo = :00
Proof. For any S, due to the definition of Tn+1 in (3. we have Em [[(w)] = E., [f ngradf (x))))
(29)
so that for In ~ nn and 1n ~ Tn+ we must have
ao+1 on ngraaF( (wx)) on 1/V log dun (æn). dj
(30)
On the other hand, let us define
Bn+l Rn V (().) / 2mm
(31)
where en ~ (N0 I), and 2o = xo- Next, let us show æ,, approximates T., For any test function C with (spectral norm) bounded Hessian, we have
E [[[(=++1)) ) B[/(=.)) + V log V log n. (#x (
Blf(z.)) + 7 Va (ox,)))))))
(B[Af(x.)) (11²)
(32)
B[S(a.)) n/ (<x,),, Vn log ((e.)) )) + WB cn v s(x)). + O(7))
=B [ (a.. V ((..) 2nje.. O(r²)) due to the definition of and D DKI (4) < c0. Then, let us define the following
||==|2
||=|| < C 6;
fa,c(x)
(wx)
C < | C:
(33)
C> ||,||
where 5 > 0, C. > and (x: a quadratic function of x to make the above fa.c smooth. Then
[f mn+1)] = T.. fo.c (n + V= ((n,)()))) + 22n an+1) + O(n²)
BBo.cc (xo, + vv log ((x,)) - En. - TV log O(1²)
(34)
(1 + 1)BB fo.c (x.. xx)) 1
,BB [fa.c (V ((xx)))))))--))))))))))) V~ log O(7²)
77
< (1 +()) 16.00 (a. 2n)] + +11))
20

---

where the last two inequalities are respectively from Young's inequality ||a + bll < (1+ + n) || +
(1 1/m)]|b]] for any rj > 0, the Lipschitz continuity of V log μa, and the definition of fs.c- Then, by recursively using the above inequality, we have
(a an)) - O(n)
(35)
when n O(1/n). By taking C o0, and applying Fatou's Lemma [11] we get
B Ie.. x. [1]] = E8
lim 8-10.C
fa.c ]n)]
lim [ff,c(.. - -2.)] O()).
(36)
6-00,C
Next, we should show the nn (so that x,,) is discretion of the following SDE (continuous Langevin dynamics). Let us define
dā:: V log ((t))dt + V2dW,
(37)
For any given To = C.. similar to 33 we can prove
|| ( Rn+1 IP] S Ønn + n log (()n)))))))))))))))))))))))))( Be+ 11 +(1))
$ + 1)B [|
Tinll²] + 1
E V (lggm( V Vlog +07)
(38)
1 17
(1 + n)BB [||tan ==.ll"] 12n
+
E [l|tao +
= O(n) [1 - *x,l1 O17²))
By iteratively applying this inequality, we get E[[| by the triangle inequality
((n) for any S O(1/n). Thus,
I [ - x.. [|] s 2F [| - a..11]] + [|| a..1 s O(n).
(39)
Thus, prove conclusion by applying Lemma to 37p.
Remark 3. It is worthy to note that during our proof, we introduce the auxiliary sequence {@n} which is the discrete Langevin dynamics. We construct its continuous counterpart instead of an. because the drift term of is V log p which has verifiable continuity. However, if we directly analyze nn with drift term V log u/T
Ew its continuity is non-verifiable.
Theorem 1. [19]] Let ne, follows the Riemannian gradient flow
8 then for any T > O, we have
1 +
1 Iggra p)ll dt S DxI(mT ()
(9)
Moreover, the log-Sobolev inequality
5 (Riemannian PL inequality) is satisfied for p. then
DKu( p)) < -2:t (> ll p).
(10)
Proof. Similar 228
have
(Te ()) (m g)l- (#e |l μ²] Ilgrad 1Il...
(40)
Thus we know the Riemannian gradient flow resulted is monotonically decreased w.r.t. t Taking integral and from the non-negativity of KL divergence implies the conclusion.
On the other hand, if the Riemannian PL inequality equality further implies
(12 holds with coefficient 7, then the above
DxL(me () < DRy 14) 0t
(41)
So that taking integral implies the second conclusion.
B Proofs in Section S
Proposition 3. The global optima of problem
11 is exp (IEg [log Hel)
21

---

Proof. The result is directly proved by Langrange's 's multiplier theorem. Due to the definition of KL divergence and [T.. the optimal π should satisfies
Teg [T (e)) - ( }
Öm [()) dn
(42)
which results
f P(E) / (L+ log x (x)) - log se (4x) Adædg 0
(43)
for some À > 0, where p(E) is the density of ç. Change the order of integral, we know that for any æ,
(og)(w) + (1 - X) = p(E) los (e(c))dE
(44)
which indicates conclusion under the condition of f =(()dxx = 1.
B.1 Proofs of Proposition 4
Proposition 4. Under Assumption let fe(m) DKL(T pc). the discrete Riemannian SGD
Algorithm with 1 S n - O((1/n]) approximates the continuous Riemannian stochastic gradient
flow
dt nt V.
me V
log
d
V ESGD
" ESGD V log m
(12)
du
2
by B[[ Bnyl [2] < O1)), where U.
nn Algorithml (x) = [ log (a) VE [log (=)))
1 Bnn
Tan in 12 .or 26 vo. Here
(V log (xr) VR (z)))
(13)
Proof. Similar the proof of Proposition we can show that for
1ut1 n. + V log (#..) + 21jen,
(45)
with En -N, we have n, approximates F,.- That says, similar to 32 when 3o prove
Eo. can
[| - ā. 11 s Onn²)
(46)
Next, we show T. in (55 is the discretion of stochastic differential equation
dæ. VEg (e::)) (VT V2I dW,
(47)
b(#,) ( V21 dW,
where Escn (æ;) the covariance matrix
EscD(a.) = Eg (V log (mc) VEg [log (#+))) (V lg (ā.) - VE, [log me (ax))))
(48)
To check this, for any function C2 with bounded gradient and Hessian, due to Bo So. Dynkin's 'sssssssnnn formula 32]. we have
(n+1)n
B (+())))))) B (f(trr))
B [(Cf((()))
1 (a++in
S
(49)
E [f((rr)) + m/B ( ) + 2 mn
B ((x))) dadt nn
E [f(mrr) a/B [C.f + O17²),
where the last equality is due to the Lipschitz continuity of V log Me (z( indicates 2 log (ę( have an upper bounded spectral norm, and log (z(x) itself is upper bounded. Noting that
((Cf( = B.., (((().)), (f((mn)))) +
2
7B [tr (EsaD ( V ((t))) + (f)())))) (50)
Plugging these into 44 we get
E f (m+))))))
= R. [f((@n))))
+ Pan, ((b(tnn)), vf((nn))) 2 [tr (EsGD V2 f(tnn)))
(51)
+ E (AA()))))))) O
22

---

On the other hand, we can similarly prove that
((tn + v/V log Me,, (Zmy)) + /2nje.) = 13 (1mn)) + +n (b(((((@m)))) (.((nn)) + n/5 (dtn))
n E
ir (s b(@nn)b (ūemn) v² s( + 0
2
(52)
Thus we get
sup
[((((+1)a))
E= [(āng + 1/V log Men (#mn) + v 2mtjm) = =(1")
(53)
due the Lipschitz continuity of log ulş, where E [f():)] (lf( 30 x). Let fa,c (z) be the ones in (3 then similar to (34).
[f.. (#(441jn - mn+1)] E fo.c (inn + 1/ loggss un,, (inn) + +/ 2nen - mn-1) + 0
[fo.c (omn Cn v log Men (fmmy) - n/V log Me.. (ax ))] + (nn²)
< O(n)) [To.c (ānn - an). +
1 1 1
rf E [fa.c (Va log Men (Znn) - V Me, (an)) + O(1²)
1)
S (1 +(+)((n)) [fa.c (#ny kn)] + 0172)
IA WA ...
O(n).
(54)
The above inequality holds for any n S 1/1.. Similar to (36). by taking ō - C - xo, and applying Fatou's Lemma, we get
B | - क.. I|]] < O(r)-
(55)
Then combining (46) with the above inequality, and by triangle inequality, we have
E [[ | - O()).
(56)
Thus we prove our conclusion by combining Lemmal 1
B.2 Convergence of (Riemannian) SGD Flow
In this subsection, we first give proof of the convergence rate of the SGD flow in Euclidean space, by transferring it into its corresponding stochastic ordinary equation.
Similar to Proposition 4 we can prove, in Euclidean space, the SGD flow of minimizing F(x) = [e[fe [x)] takes the form of
dæ;
VF(..) (w. &dW,
(57)
with (x,) Eg [[ (e]f(( (x) Then by Lemma | it's
corresponded stochastic ordinary equation is
dæ: (=()) - V Lsc DSSD V log (x.)ddt.
(58)
Before providing theorem, we clarify the definition of our computational complexity to continuous optimization methods. Owing to the connection between the continuous method and its discrete counterpart as in Proposition 4 and It requires T discrete update steps to arrive ă.,r. Therefore, the computational complexity of running T discrete steps said to be the computational complexity of continuous optimization methods measured under T,T.
Next, let us check the convergence rates of (58) for general non-convex optimization or with PL inequality, It worth noting that for this problem, the convergence rate is measured by [E[] FF [
or E[F((;)) - inf, F(x), with/without PL. inequality.
Theorem 4. Let xt defined in (8). then if E [tr (z: <
by taking πj =
[[F( ) infa F(w)]
have
To
TT I Em, [IVF(())) dt 2(F(wo) inf F F(x))
(59)
nT
Thhis can be satisfied when fx( and its gradient are Lipschitz continuous. Because we have Ejtr(>ss (æx) ((x))) & Amax (V= ( F(xx ([s())(Ss due to the definition of EsGp.
23

---

On the other hand, if F(x:) satisfies PL inequality 1288 by taking 1//T" with 0 < a < 1, we have
E F(())) inl F
(60)
Besides that, (x:) is in the form of finite sum, the computational complexity is order O(c-2) to make Eeg. || (xx,) [xs))]]] for some i, and O(e1/a) make Exr, (F(x)) inf, F(x:)] under
PL inequality.
Proof. Due the definition of æt.. (F.,)) dt
(V((.)), -V(()) v (yEcn(())) " (lgg..)))))))) (61)
On the other hand, for æt ne.
(V)()))) ESaD
(mx(( (a) da
Em [VFF(z)), V Escn((:)) (Escn
(62)
Plugging this into 61 we get JE [ax))]
F(x
+
2
E [tr (Escp(z)) s -E [IVFF())))) + tn=
(63)
Thus
n /'
E (()()))
F(ETT))) +no²
Ez, (dt
1/
E. ((ro)) F(aTrn)or
= =2
(64)
T
2E
(xxu) F(wrr))
1/T
by taking 17
Flco) - Info F(e)]]
Then we prove our first conclusion.
Next, let us consider the global convergence rate under PL inequality. By applying 128 to get
63 we
a [x(x)) infa F(x) < (22(EE(((x))) inf F(x)) + noo
(65)
Øt
which implies
(((,))))) inf F(x) S RTl (B - inf
inf F(x)) + no² (I (1- (1-))) 2m1)),
(66)
by Gronwall inequality. Thus by taking 17 = 1//1"" with 1, we get second conclusion.
Under finite sum objective, when VT minogesoi Ex, [[,V((,,)) < Thus, due to Proposition it takes M - O(nT/n) O(e-2) steps in Algorithm to get the "e-stationary point".
Ol77))
O(e we have
Furthermore, under PL. inequality, due to (66). it takes 7 = O steps to make E (F(xTT))
inf, ((x) < which results in computational complexity of order O(e-1/a).
The convergence rate of Riemannian SGD flow can be similarly proven as the above theorem. Next, let us check it. We need lemma termed as von Neumann's 's trace inequality [46] to prove Theorem 4
Lemma 2 (von Neumann's trace inequality). For systematic matrices A and B, let (M(A)) and (A:(B) respectively be their eigenvalues with descending orders. Then
t=(AB) & (A)A)((B).
(67)
Theorem 2. Let n. follows the Riemannian SGD flow 12 and defined in Proposition 3 Under
Assumption 2 ifT Z 641, DKu then by taking
(( we have
Dxc(
1 Tn
4DRI( ll ()) = (1T
Igzra H))
(17)
nT Jo
TT
24

---

Besides that, if 5 is satisfied for p. 17 1/7/T" with 0 < < and T
(8L3/7)
1/a then
1 T [4dL]22 (d+ 1) 0
[ T
(Rr (1)
(18)
Proof. During the proof, we borrow the notations of H1, Hg. H3 in Lemma Next, we prove the results as in Theorem That is
(Te(Te ll (4) = (rradDKu(T ll ()) -graddDz I 1))+22 V EsGD + 2 ESOD V log
" EsGD + " EsOD log m IggraDDx ll 1)]l (erad u), 2
(68)
To begin with, we have
(gradDKa(Te 1l n): V EscD)
=
/((
log
de ? v
EsGD di
2
nl3
Igrad ( 11)=
"
B.,
[I EscDl]] (69)
8Hs
nHs
nHi
Igrad
m)
871,
where the first inequality is due to Young's s inequality and last one is from the Lemma have
Then, we
gradDRL(Te 1l H). ESaD V log mI ) = - = log
due 7 ESOD V log dm
du 2
log me, ESOD V log mt) du - / (v / log ESCD v log कt d
H3 /
| log doe dm / (v log T NESOD log mi dr
(70)
Ils
IggraDD<I w)ll,, - (v (Vlog ur 2
EseD V log mT dxc
nH3
| log 2l| dxc + nHl / (v loggH, V
log loggo)) dr..
In the r.h.s of the above inequality, the sum of the second and the forth terms can be bounded as
log H.
7 Esen log me dm,
nHa / (VVlog
v log क. dm.
n log n, ESED log mi d nHH
/ v Voog 14, " v log m dm
2
2
/ v loggH, " Esen) dm + / (v log ?
" Esen) ) dets
2
ʃ ir (v² log pISDD) dm nHs / (V² log dat
(71)
+ logg, ui,
"
MsCD) dite
2
1HH1H Hy <
nH,
yl |IV log log ull drc
2
2
nH,Ha
nH,Hs
nHs / | log all? dm
nHz
2
8Hs
where the inequality a is from Lemma 2 and the semi-positive definite property of ESGD. By
plugging 699 (70) and (7I) into (68). we have
nHZ
DaI( Il) g -(1 nyHa) Igra ll 1)1 + 7HH,H +
(72)
4Ha
25

---

due to the value of 17 makes nHs S 2- Taking integral w.r.t t as in 64 implies
/,"
nr
DKI( u)
nHă
Iggrad w)l dt
qtT, H
2T/T
nT
1/1a
(To p)
H2
2/
H, +
(73)
4Ha/
1 1/T
2D (To 11)))
so that we prove our conclusion due to the value of H1, H2, H3. The second result can be similarly obtained as in Theorem 4 by inequality
DeL. (TT |l μ) S e mi (xe II μ) +
H, +
Hi 4H,
(1-c
(74)
which is obtained by applying log- Sog-Sobolev inequality to 773 and Gronwall inequality. Thus, taking
1//T"" with 0 < a < implies our conclusion.
Similar to the proof of Theorem 2 we can get the computational complexity of Riemannian GD and respectively for non-convex problem to arrive e-stationary
SGD flow, i.e., O(Ne- and O point, but ON log and O under (log-Sobolev inequality) Riemannian PL inequality.
The following is the lemma implied by Assumption 2
Lemma 3. Under Assumption 2 holds E,T, [tr* (V² log u)] I dL2 HI
supp || -ESGDl
|| S
4(d + 1)L1 = H2, and supx Amax (aax(n (x)) S 4L3 = H3.
Proof. Due Assumption
2 we have
E tr' (V2 log PiTe [tr (Eg V log [e]]]
P.Te [d | V2 log pe
(75)
< dL2,
where ||||| here is the spectral norm of matrix. On the other hand, we notice
ESGD
Eg [tr (v2 log
E3 log (V log ue E log
(76)
+ E (V² log E& 2 log ME) (V logiue - E V rog)
Then by Assumption
sup
V EsGD(x)] < 4dL1 + 4L1 = 4id + 1)L1
(77)
Finally, due to Assumption
2
we have
sup Amax (E))D(x)) S 4L2.
(78)
C Proofs in Section 6
Proposition 5. Under Assumption 2 let fe(m) ( Il Mg). the discrete Riemannian SVRG
Algorithm 2 with < n S O(L1/n)) approximates the Riemannian SVRG flow
0 01 (.()) V (m)))))
V Vlog
di
(a:)
") (x)(
Te, T (y) V.
Esvno a)
dp
(24)
Te, x. E x) V log म., (y. m)yy
2 (a=)
for = 1, S S li+1 (dt,, wx) ~ Ti De ~ Ti in (24).
since we have [i[l t t
(1), where xt ~ T in Algorithm 2 for a! = dn. Here
ESVRG (y. = Eg [V log μe (x) - V log ( + VEg llog (3)l - VE, log (z()))
(25)
(V log ug (x) log + VEg [log 2s (v) VE, (o))
11xx(() (+) stands for the sum of absolute value of ''s eigenvalues.
26

---

Proof. The proof similar the one of Proposition Firstly, we choose r as the transportation that preserves the correlation between zá and an which means that for any function u E T we have
(u(z))) u(x)),
(79)
so that To x n is the union distribution of (x) v;,,
defined as below. Concretely, we know that for - L, the corresponded oy, of discrete SVRG in Algorithm satisfies
S n
S M 1, S
i
ducs. (xr V log dr
d/ssh (x)) + V log dig
dp
+++1 x +
Vog
V
(xb)
(80)
dmgg
In the rest of this proof, we neglect the subscript to simplify the notations. Similar to Proposition 4 can prove #1 is approximated by
Lo+1 T. + (V log Me,. (z..) V log ue,, (2o) + V log + 2nj..
(81)
with En (0, I) by noting the Lipchitz continuity of V log uę. As in Proposition the approxi- mation error is similarly proven as
E [||e. æ. 11
æo. mo] I O(r) + O[ - =oll²).
(82)
Next, goal is showing that the discrete dynamics H, is approximated by the following stochastic differential equation
dæ. = VF& [log (dt)) + (v #c), V2T dW:
(83)
= b( + (vTSšvno ic), /21 dWe:
= Sts (i+ 1)Mr,
where M is steps for each epoch, and S (āt,, xz)) is defined (25 Similar (49 for (n + iM = t E [ti, t.+1). We write e (n+iif) as &ny in the rest of this proof to simplify the notation. We can prove that for any f C2 with bounded gradient and Hessian under the condition of given a.,, peo [f(((1+1))]]
=T.0, [f(#n))] ") ((b(d)n)) Vf((fnn))) + 17 [Af(
+ n72 R² [tr [(SsvRc amn))) v ( (
(84)
On the other hand
R f (mnn 77 (V log Men (man) 1og (#te.) (o + / 21
2 [f(#mn)] n. [(b( V V ((mn)))) + mE #4 [ (m)))
(85)
1/
マーレ g [rr [ (Esvio (ār,, 2nn) + b6 4 f(Eer) O()"),
2
where we use the fact that
t [og Hea (ian) - v log Men. (wt,) + V log ((+;,)] = 18 V log uen (@mn)) = b((.n).
(86)
Then similar to (54) in the proof of Proposition 4 by taking f as fa. in 33 and combining (82) we prove
E [|l 12 i.. _o] O()) O (l xo 2o1l
(87)
By noting that do = =g. recursively using the above inequality over i, and taking expectation, we prove that
E ||
æ ||| < On))
(88)
On the other hand, we know that given Ct, y. the conditional probability i; Lt. follows conditional density स. e, (x) satisfies
01 Tele (x v) = V. Tele, (ee w) (vlog log dit (w) du
0
V. EsvRo(y.
II x) log (y.
" 2
2
(89)
Then, multiplying ne. (v) to the above equality, applying equality Va (o, V. x log Teje, (x y). and taking integral over y implies conclusion.
27

---

C.1 Convergence of Riemannian SVRG Flow
Similar the results of Riemannian SGD in Section B.2 we first give the convergence rate of continuous SVRG flow (stochastic ODE) in the Euclidean space, which helps understanding the
results in Wasserstein space. First, to minimize F(x) I/N - fc, (æ), by generalizing the
formulation in 24). the SVRG flow in Euclidean space is
dæ: (x.) " V ( 0s)
EsvRa(( 2. a.) log
4 kstSke (90)
2
with ΣSVRG defined as
EsvRo .²) Ei [V]((()) V - VF()) VF(=))(V/((()) Vi((v) - VF VF(=))
(91)
and ne is the density of Ht, E is uniform distribution over {E.). Then the convergence rate and computational complexity of (90) is presented in the following Theorem.
Theorem 5. For in (90), learning rate n O(N-2/3), A
1, 4-11
0(1 /1), and n/T = Tá, if fe(x) is Lipschitz continuous with coefficient I, and has Lipschitz
continuous gradient with coefficient Lzg, then
R [[IVF(())))) 2(E(F(aa)] inf F()))
(92)
n/
On the other hand, by properly taking hyperparameters, the computational complexity of SVRG flow of order O(N2/3/e /e). when minosesor" E [[VF(( I e. Further more, when F(x) satisfies
PL. inequality (288)
with coefficient 7, we have
E[F(x,T) inf F(x)) se
E F(xx) inf F(x)
(93)
The computational complexity of SVRG flow is of order O(N + N-133 log loge=)) to make E [[(())))) inf F(x)] Se.
Proof. Let define the Lyapunov function, for E [ti, li+ i] (with out loss of generality, let 0),
Re(=) F(x)
@ || 2tp
(94)
2
Then we have for x, defined in 900
d |., Fo d a (o:: Dtoy dxa).
dRe( (V (xx), dx.))
(95)
2
Due to
90
E [| E[a:] [ < E || ||] and Lemma 2 we have
(x()), daye) ) --
" [tr (Zssva 2i)) V F(x)]]
Mmmmxx 2 (<=)))))) E (Esvna az))]
s E [x=){()))
Amax 2 F(xx)) E [[V]e(> Vee( + VF(w))) VF(=)]] (96)
[IVF((=))]
7/4 F. [( V fe(,,)))1
S (x:))]
n22 2
< E [IVF(*)1))
E |||T Ft. 11
On the other hand, by Young's 's inequality, for some 3 >
E
(
xc
Dio. dt
dac )
EI- (xx Ztp. VF(x,)) ) 2
"E [tr (Esvno(to,
(97)
E [|
2ro 12
[IVF((=)))
nL2 [| - 2z, 1!
2
2,3
By plugging (96) and 997
into 95
we have
0 E [R(() < (1
(B +
nL2)et
nLil2 di E [|| 10, (98)
E
V
(xx)))
+
at
23
2
2
28

---

By making
(B nLaae .TL 4= =0,
(99)
we get
Co2 (0+n)2)(()) 10) cto
(1- ( (0toba)(t-100)
(100)
3+ n12
By taking Coo
/n. Ct 0, v get
t to log (1 + vn/ 1/2 + 1/111/1
O
(101)
v1 n2
when 17 + O. On the other hand, due to c/ < 0 and (99)
we have
2(t1 ta) (") [IVF((=)1*) dt =
Cta
E
[[IVF( (x:F(()))11))) )11] dt
f1 to
1g
2
Ct
E [[V
to
28
(102)
E Ri, (a
E RRg (xt,))
11 to F F )
E (2((1)))
11 to
Thus, for any 7 = t and A = 1 to =.. t - 1 defined in 101). we have
E [[VF(2))11 di 2(B (xx))) - inf F(x))
(103)
IA
which leads to the required convergence rate. Next, let us check the computational complexity of it. Due to /A = n/T' O(e A = O11/ /T), and Proposition we should running the SVRG for I (Vvi/c) epochs with .M = /n steps each epoch in Algorithm 2 Besides, note that for
each epoch, the computational complexity is of order O(M + N) = (D//// N). Then, by taking
(N ),3), the computational complexity of SVRG flow is of order
N2
+ =
10(M + N) = 10(A/1 - N) =
(104)
Next, let check the results when F(x) satisfies PL inequality with coefficient q. We will follow the above notations in the rest of this proof. For any E [lo 41.. we can reconstruct the er in Lyapunov function (55) such that
(2 + B nL2ca + nL1L2 + d =,
(105)
which implies
C₁ (2)+
27 + n2 (1-
nT. 1.2
(106)
Similarly, by taking CtD = -1 = = /17, we get to log (1+ 27/
1-22% VTLL /T/ 1/L2 min o (
(107)
27 /T n2
v
when 7) Plugging this into 988. combining PL inequality and the monotonically decreasing property of ct. we have
Øt ((( inf F(x)
a s -2
(B(F(xe)) - inf F(x)) - - [||are - -tte-x,,lll 11,
23
5> [ (xx)))
inf (())
- - [|| - -200
(108)
= -E [xx))) (xx) - inf F())
Hence
[F(()) inf [xx) E 7011 -1o--t) ( innf
(109)
Then for any T = f. and = f - 14 = ... = t1 - to. we have
E F(WT) - inf F(w) & e TTE F((o)) - inl F(x))
(110)
29

---

which implies the exponential convergence rate of SVRG flow under PL inequality.
On the other hand, let check the computational complexity of it under PI. inequality. As can be seen, to make [ (2,, - inf, F(x:]] I 6, we should take 7A = 77? 00)
11) 1o
Due to
1177 the SVRG flow should be conducted for / = max (O( 1og c with M = A/n steps in each epoch, and the computational complexity for each epoch is of order ( + N) O /n + N). So that the total computational complexity is of order
(log } epochs
10 / N) max ( -1 loge "), Ollog - "mmin (0(1 ((17)))(1))))))))) /n N), (1II)
If y 2 N1/3, we take = ( the above equality becomes O((N +/=1N2/3) loge
On the other hand, when y & N1/3, and = O(N 2/3), the above equality is also O + y 1 22/3) loge
Next, we prove the convergence rate of Riemannian SVRG flow as mentioned in main body of this paper. The following theorem is the formal statement of Theorem 3
Theorem 6. Let n. in 24 A 4 11-t
t 1 -1 O(1/ MT) and n/ 1 for I epochs. [tr(V2 log (tt//d/)EsvRc) I A. , [tr(ESvRa)) holds for any
Then, if Assumption 2 and E Ex,,,
t and dy is a polynomial of t, then
E/ Iggra ll 10 dt S 20x1( ll m)) T
1
(112)
1/
On the other hand, by taking 1) = OON-2(()))))) the computational complexity of Riemannian SVRG flow is of order ((N2/3 /e) when minosesa7 11 m)]l Se.
Furthermore, when F(T) (x |l p)) satisfies log- Sobolev inequality 5). we have
(RI (4) S - DKI(mo ll p))
(113)
Besides that, it takes O(NN + - log e 1) computational complexity to make D (FaT)
(1) S.
Proof.
Let us consider the Lyapunov function for union probability measure Tip.t with 40 S15h
(T(T (Fe p) + &e /
ʃ lIw 21 To. xx) DxI( ll u)
a ((mo.))
(114)
Then
& ((T.))))) + G(T.):))
Ca 8
Rd (TTo.)) (er ((m))) 01
(115)
2
2 Đt
where 0 = V (T,H(T;)) is used to simplify the notations. Then due to 24 Lemma similar induction (66)
and
(gradDxx (mee ll 11, h(Te)) = Dgz (#e Il m))l
log dm
(x))
Tep (V, V. Dyvno dydæ
dy
log dmc (x))
+
Ttp. x) ESVRG x) Van log rey, x))
dyda
du
= Ilgrad (TTe |l [1)1 + ! B ir 2 log
dm
dy (w
(x))svGG( (v.
< - IggradDi Il mill + n
[|| log (x(,,)
log 10)))()))11
2
< - ligradDzz Il will + n22de
G(Too..),
(116)
30

---

where the last inequality is due to the Lipschitz continuity of V Hz))) Similarly, by Fokker- Planck equation, we get
① G(TTO..))
0 F. ||
Tto [1=]
0t
0t
dxt #t u di )
R (at Do (TT)))))))))
(117)
1
gra.Dxx(mm || p)ll
G(TT)))) +
Bro [tr (E ax)]]
23
2
2
1
n2
$ gradDRr
|| m)ll
+
G(TTo.))-
23
2
2
Combining
(15).
(17 and (166 we get
[8 + (1+nn
ndce2 d G(T (118)
Ri(TT S -
IgradDRi Il m)ll. +
+
Øt
2/
2
2
2
By taking
+ n/2)e +
(119)
which implies
Ct = - (8+n22)(1 to) to oo f nladre -(+tL2)((-
(120)
Without loss of generality, let
20 1'h
(t)
(t to (8+gLa)(t-to) (
Op n/2
nz2
n1,2
(11 1o) - 1222 1(1 pll - 10)0-1e 1-p-le-0ttttalt-- di
8 112,
1122
(t1 to -e(0te-l -1000-)(0-60
(121)
p0p-11
8 8+1
11L2 (1 e -8t to:p))
1
where Poly(t1 to p) is p-th order polynomial of to. W.l.o.g, let
de Aeg +Plyy(t tu, A10 + Σ bi(t to)
(122)
Then we have
$1
aLzde (84r02)(1-10) = Ae ao + Σ aib
nL2do
(1- (8+012)0110))
nL
Σ by
-04t0t2N1--00 l0.p)
+ 11
3 + 12
11 3 112
21o
Poly(4 - lo p)c (8+122)(1
i=1
(123)
By invoking Coo = -/n, C2
= 0, 3
/n. and the above equality into (20)) we get
v7l2 L2Poly( to.p) L2(Ao Ei-,bb)
to / nT2
log
(124)
which implics t - to = 011 7 as in (OOT) (note that the value of b, are automatically adjusted to make the above equality meaningful), by taking + O. Thus, similar to (02).. get
2(41 (v) f( lggradDKI m)1 dt (RI ll μ) DrL( ())
(125)
Then the two conclusions are similarly obtained in Theorem by taking 7 2/33)33).
31

---

Notably, the imposed "proper" condition on can be implied by the polynomial upper bound to the spectral norm of Hessian. This because, from Lemma and semi-positive definite property of ESVRG, we know
B [V' log (dvv/d)) s Eree,, [max (V² log (d=+/du)) tr( tr(EsvnG)).
(126)
Then, we may take
E Amax (ur(V2 log (dm)/du) tr(Esvra)]
A
(127)
E (t(ESvRG)
Due to the formulation of Riemannian SVRG 24 the density He only depends on To and μ. Therefore, under properly chosen To, the obtained ni can satisfy the imposed condition the spectral norm. Besides that, we do not impose any restriction on the order of X while the polynomial function class can approximate any function that it can be extremely large. Therefore, the imposed polynomial order of A can be easily satisfied.
D More than Riemannian PL inequality
We observe that Theorems the convergence rates into global ones. Therefore, we briefly discuss the condition in this section. As mentioned in Section 3 the log-Sobolev inequality is indeed the Riemannian PI. inequality in Riemannian manifold. To see this, for function f defined on Rd the PI. inequality that for any global minima xx of it, we have
and demonstrate the nice log- Sobolev inequality (5) improves
[=))
PI.
(128)
holds for any The PL inequality indicates that all local minima global minima. It is nice property that guarantees the global convergence Euclidean space [23]. Naturally, we can generalize
it into Wasserstein space. To this end, let F(x) D || ()), the sole global minima is / with F (()) = 0. Then, in Wasserstein space, the PL inequality (128) is generalized to
2>Dxz(= II μ4) < IlgradD (xr |))l
Riemannian PL.
(129)
which is log-Sobolev inequality ( due to 7 Thus, our Theorems T and indicate that
Riemannian PL inequality guarantees the global convergence manifold.
Moreover, if (128) has Lipschitz continuous gradient, the properties of quadratic growth (QG) and error bound (EB) are equivalent to the PL inequality (23
f(x) - f(x) 2 || x x²1|²
QG,
(130)
[ 2 >/1 - x'll
EB,
(131)
so that global convergence. In Wasserstein space, the two properties are generalized as
DKL(M Il p)) (T,
Riemannian QG,
(132)
dm di
Vlog
WW( n4)
Riemannian EB.
(133)
Then, the relationship between the three properties in Wasserstein space is illustrated by the following proposition. This proposition is from [34], and we prove it to make this paper coffccontiined. Proposition 6. [Otto-Vallani] Vallani][34) Let ((T) be (KL(T) μ)) then Riemannian PL = Riemannian QG, Riemannian PL > Riemannian EB.
This proposition indicates that Riemannian PL inequality implies the other two conditions, but not vice-versa versa. This is because the Lipschitz continuity of Riemannian gradient i.e., D (T 2 < LW3 (1) does not hold as in Euclidean space 45 (Wasserstein distance is weaker than the KL divergence).
Proof. If Riemannian PL + Riemannian QG then Riemannian PL Riemannian EB is naturally proved. Next, we prove the first claim. Let G(T) = (F(T), then by chain-rule and Riemannian PL inequality,
l = gza ()
7
(134)
2D (T 11)
2
32

---

Then let us consider
0
Expx, gradG(Tt)]
(m))r
0t
(135)
n0 = π
where the last equality can be similarly proved in Then
G( - G(Tr)
G(m.)dtt
DevC(TT) gradl)mm)
|| g-gradG(mm)l.
(136)
2 dt
Y'''
Due to this, and G(x)) 2 there exists some T such that कr μ.. Since и is curvature satisfies To = x and Wy = na, due to the definition of Wasserstein distance geodesic distance, we have
W2( Wa(To, mT) (." [1)r)d
(137)
Thus,
f 11 g-grad 1 1 1/ 2 IlrradG(m))ll,dt 2
/ Wa ())
G(Ta) G(mo))
(138)
which implies our conclusion.
33


Reviews:
[{"review_id": "w5dHr7Pz41", "reviewer": "Reviewer_mLox", "summary": "This paper establishes the continuous Rimennaian SGD and SVRG flows in the Wasserstein probability space from the optimization perspective. The key idea is to link them to a dynamics in the Euclidean space.", "questions": "* For the Riemannian SGD flow, it seems $\\xi$ can be any distribution. Would be better to provide more discussion on some special case. For example, is there a special case wether the Riemannian SGD flow reduces to Riemannian GD flow from this perspective?\n* In Lemma 1, better be ``Lemma 1 ([29,42]).'', similar for Theorem 1.\n* What are the key differences from the existing proof for the establishment of Theorem 1?\n* Eq. (25) too small?", "limitations": "yes", "rating": 4, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 2, "originality": 3, "strengths_and_weaknesses": "Strength: Riemannian SGD and SVRG flows are etablished which seems to be new.\nWeakness:\n* Somehow lack numerical validations on the efficiency of Riemannian SGD and SVRG  on sampling\n* It seems the overall  technical details differ not too much from the Euclidean case on the convergence of the flows", "was_revised": true, "final_justification": "The authors have addressed my issues, but it seems the paper needs quite a few modifications. Thus, I will maintain the score.", "rebuttal": "To Reviewer mLox:\n\nWe thank for your valuable comment and address your concerns as follows.\n\n**Q1**: “Somehow lack numerical validations on the efficiency of Riemannian SGD and SVRG on sampling”\n\n**A1**: Thanks for your valuable advice. We have made an example (Example 1 in our paper) of sampling on Gaussian distributions to illustrate the tightness of our theoretical results. \n\n**However, we recognize that numerical results will make our conclusions more convincing**. We evaluate different sampling algorithms on two experiments. **Since the NeurIPS rebuttal do not allow us to convey these results by image**, we present them in the following tables (results are averaged by 5 independent runs). \n\nConcretely, we set the each target distribution $\\mu_{\\xi}$ as **Gaussian or mixture-Gaussian**, and conduct sampling by the proposed algorithms. Notably, our convergence rates are respectively evaluated by update steps ($T$) and gradient computation (Riemannian SGD and Riemannian SVRG has less gradient computations than Riemannian GD). The results are summarized in the following tables, the KL divergence and Fisher divergence are estimated by combining MCMC and an estimated density as in [1]. \n\nAs can be seen, **under the same update steps**, Riemannian SVRG and Riemannian GD have better convergence results than Riemannian SGD (see Theorem 1, 2, 3), while **under the same gradient computations**, Riemannian SGD and Riemannian SVRG (the two algorithms have less gradient computations in each update step) have better convergence results (see discussion after Theorem 3). These results are consistent with our conclusion in paper. **We promise to add these numerical results in the revised version**, and we sincerely hope you may reconsider the concern on numerical results.   \n\n**Table1. KL and Fisher divergences on Gaussian measured by update steps**\n| **KL divergence \\ Update Steps** | 0    | 20   | 50     | 80      | 100     | **Fisher divergence \\ Update Steps** | 0    | 20    | 50    | 80   | 100   |\n| -------------------------------- | ---- | ---- | ------ | ------- | ------- | -------------------------------- | ---- | ----- | ----- | ------ | ------ |\n| LD           | 6.12 | **0.43** | 0.0214 | **1.06e-3** | **1.44e-4** | LD      | 11.9 | **0.259** | **0.102** | **0.0655** | **0.0424** |\n| SGLD                 | 6.12 | 0.51 | 0.0274 | 1.26e-3 | 1.62e-4 | SGLD         | 11.9 | 0.279 | 0.142 | 0.0723 | 0.0512 |\n| SVRGLD               | 6.12 | 0.67 | **0.012** | 1.16e-3 | 1.52e-4  | SVRGLD            | 11.9 | 0.244 | 0.119 | 0.0679 | 0.0482 |\n\n\n**Table 2. KL and Fisher divergences on Gaussian measured by gradient computation** \n| **KL divergence \\ Gradient Computations** | 0 | 2000 | 5000 | 8000 | 10000 | **Fisher divergence \\ Gradient Computations** | 0 | 2000 | 5000 | 8000 | 10000 |\n| -------------------------------- | -------- | ------- | ------ | ------- | ------- | ---------------------------- | ------ | ------- | ------- | ------ |------ |\n| LD                    | 6.12 | **0.43** | 0.0214 | 1.06e-3 | 1.44e-4 | LD  | 11.9 | 0.259 | 0.102 | 0.0655 | 0.0424 |\n| SGLD                  | 6.12 | 1.76e-4 | 8.99e-6 | 1.71e-6 | 7.56e-7 |SGLD   | 11.9 | 3.64e-4 | 1.88e-5 | 3.73e-6 | 1.87e-6 |\n| SVRGLD                | 6.12 | **<1e-16** | **<1e-16** | **<1e-16** | **<1e-16** | SVRGLD   | 11.9 | **<1e-16** | **<1e-16** | **<1e-16** | **<1e-16** |\n\n **Table 3. KL divergence and Fisher divergenceson on mixture Gaussian measured by update steps** \n|  **KL divergence \\ Update Steps**  | 0 | 200 | 500 | 800 | 1000 |  **Fisher divergence \\ Update Steps**     | 0 | 200 | 500 | 800 | 1000 |\n| -------------------------------- | -------- | ------- | ------ | ------- | ------- | ---------------------------- | ------ | ------- | ------- | ------ |------ |\n| LD            | 0.57 | **0.33** | **0.07** | **0.082** | **0.076** |LD                | 341.2 | 0.252 | 0.248 | **0.014** | **0.018** |\n|SGLD        | 0.57 | 0.53 | 0.33 | 0.127 | 0.112 |SGLD            | 341.2 | 0.362 | **0.143** | 0.019 | 0.022 |\n|SVRGLD        | 0.57 | 0.39 | 0.16 | 0.092 | 0.082 |SVRGLD           | 341.2 | **0.135** | 0.036 | 0.017 | 0.021 |\n\n\n**Table 4. KL divergence and Fisher divergenceson on mixture Gaussian measured by gradient computations**\n| **Gradient Computations** | 0 | 1000 | 2500 | 4000 | 5000 | **Gradient Computations** | 0 | 1000 | 2500 | 4000 | 5000 |\n| -------------------------------- | -------- | ------- | ------ | ------- | ------- | ---------------------------- | ------ | ------- | ------- | ------ |------ |\n| LD                   | 0.57 | 0.33 | 0.07 | 0.082 | 0.076 | LD                    | 341.2 | 0.252 | 0.248 | 0.014 | 0.018 |\n| SGLD                 | 0.57 | 0.23 | 0.124 | 0.074 | 0.062 |SGLD                | 341.2 | **0.219** | **0.237** | 0.009 | 0.006 |\n|SVRGLD                | 0.57 | **0.21** | **0.062** | **0.072** | **0.032** |SVRGLD               | 341.2 | 0.659 | 0.241 | **0.0081** | **0.0049** |\n\n**Q2**: “The overall technical details differ not too much from the Euclidean case on the convergence of the flows”\n\n**A2**: Actually, our theoretical techniques are quite different from the ones in Euclidean space. Firstly, the construction of Riemannian SGD, SVRG i.e., Propositions 2, 4, 5 are mainly based on the techniques from stochastic analysis, which are quite new in machine learning community. On the other hand, we admit the convergence analysis **borrow the ideas** from Euclidean space, while their **proof techniques are more complicated**. For example, the construction of Lyapunov function (the core idea of proving convergence results) for Riemannian SVRG (114), which takes us a lot of efforts (introducing the distance penalization).   \n\n**Q3**: A discussion on the random noise $\\xi$ in objective, can Riemannian reduced into Riemannian GD.\n\n**A3**: Firstly, as you said, our conclusion is independent of $\\xi$ for Riemannian SGD (to apply Riemannian SVRG $\\xi$ should be discrete). We use **a constant $\\xi$** as a special case to address your concern about the connection between Riemannian SGD and GD. As can be seen in Riemannian SGD (12), when $\\xi$ is constant, the variance $\\Sigma_{SGD}$ becomes zero, which makes the Riemannian SGD reduce to Riemannian GD as you predicted. We will add this discussion in the revised version. \n\n**Q4**: In Lemma 1, better be ``Lemma 1 ([29,42]).'', similar for Theorem 1\n\n**A4**: Thanks for your suggestion, we will revise it.  \n\n**Q5**: What are the key differences from the existing proof for the establishment of Theorem 1?\n\n**A5**: Our proof technique is a little bit difference with the existing result. Concretely, we prove the theorem by borrowing the notations in Riemannian derivate (see (28) and (40)), and prove the theorem **under the idea of optimization on manifold**. However, the existing literature e.g., [1] prove this theorem by computing the derivate to $t$ directly. \n\n**Q6**: Eq. (25) too small?\n\n**A6**: Thanks for pointing out this, we will make (25) larger in the revised version. \n\nRef: \n[1] Kinoshita et al., 2022. Improved Convergence Rate of Stochastic Gradient Langevin Dynamics with Variance Reduction and its Application to Optimization."}, {"review_id": "5FsrLSehlq", "reviewer": "Reviewer_koyb", "summary": "The paper proposes continuous-time Riemannian stochastic gradient descent (SGD) and stochastic variance reduction gradient (SVRG) flows on the Wasserstein space, extending discrete optimization methods to their continuous counterparts. By leveraging stochastic differential equations (SDEs) and the Fokker-Planck equation, the authors construct these flows and prove their convergence rates, which align with results in Euclidean space. The work connects Riemannian optimization with practical sampling processes, offering theoretical insights into stochastic methods like Langevin dynamics and SVRG Langevin dynamics.", "questions": "- If I understand you correctly, the technology for constructing gradient flow (including SGD and SVRG flows) mentioned in this paper largely depends on the special properties of exponential mapping, to some extent, which is linear, and thus taking  $\\eta\\rightarrow \\infty$ gives the desired gradient flow which is essentially similar to the Euclidean case. Right?\n- The exponential mappings are different when using different metrics. Hence, if using another Riemannian metric instead of the inner product in the Euclidean space, the resulting exponential may not possess the \\`\\`linearity’’. Then, how do we construct the gradient flow in that case?\n- How do the results extend to non-smooth or non-log-concave distributions, which are common in practice?", "limitations": "Limitations are properly addressed.", "rating": 4, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 3, "originality": 3, "strengths_and_weaknesses": "**Strengths:**\n\n- The paper introduces continuous-time Riemannian SGD and SVRG flows, bridging discrete stochastic optimization and continuous stochastic dynamics in the Wasserstein space, a relatively unexplored area.\n- The convergence rates of proposed Riemannian SGD and SVRG flows are provided under both non-convex and Riemannian Polyak-Lojasiewicz (PL) conditions, matching Euclidean-space results.\n- The connection to sampling methods (e.g., Langevin dynamics) highlights potential applications in machine learning and statistics.\n\n**weaknesses:**\n\nThe paper is entirely theoretical and lacks numerical experiments or simulation results to verify the practical effectiveness of the proposed method. For example, no performance comparison in specific sampling problems is shown", "was_revised": true, "final_justification": "I have updated my score, and my questions are all answered.", "rebuttal": "**Q1**: “No **numerical performance comparison** in specific sampling problems is shown”\n\n**A1**: Thanks for your valuable advice. We have made an example (Example 1 in our paper) of sampling on Gaussian distributions to illustrate the tightness of our theoretical results. \n\n**However, we recognize that numerical results will make our conclusions more convincing**. We evaluate different sampling algorithms on two experiments. **Since the NeurIPS rebuttal do not allow us to convey these results by image**, we present them in the following tables (results are averaged by 5 independent runs). \n\nConcretely, we set the each target distribution $\\mu_{\\xi}$ as **Gaussian or mixture-Gaussian**, and conduct sampling by the proposed algorithms. Notably, our convergence rates are respectively evaluated by update steps ($T$) and gradient computation (Riemannian SGD and Riemannian SVRG has less gradient computations than Riemannian GD). The results are summarized in the following tables, the KL divergence and Fisher divergence are estimated by combining MCMC and an estimated density as in [1]. \n\nAs can be seen, **under the same update steps**, Riemannian SVRG and Riemannian GD have better convergence results than Riemannian SGD (see Theorem 1, 2, 3), while **under the same gradient computations**, Riemannian SGD and Riemannian SVRG (the two algorithms have less gradient computations in each update step) have better convergence results (see discussion after Theorem 3). These results are consistent with our conclusion in paper. **We promise to add these numerical results in the revised version**, and we sincerely hope you may reconsider the concern on numerical results.   \n\n**Table1. KL and Fisher divergences on Gaussian measured by update steps**\n| **KL divergence \\ Update Steps** | 0    | 20   | 50     | 80      | 100     | **Fisher divergence \\ Update Steps** | 0    | 20    | 50    | 80   | 100   |\n| -------------------------------- | ---- | ---- | ------ | ------- | ------- | -------------------------------- | ---- | ----- | ----- | ------ | ------ |\n| LD           | 6.12 | **0.43** | 0.0214 | **1.06e-3** | **1.44e-4** | LD      | 11.9 | **0.259** | **0.102** | **0.0655** | **0.0424** |\n| SGLD                 | 6.12 | 0.51 | 0.0274 | 1.26e-3 | 1.62e-4 | SGLD         | 11.9 | 0.279 | 0.142 | 0.0723 | 0.0512 |\n| SVRGLD               | 6.12 | 0.67 | **0.012** | 1.16e-3 | 1.52e-4  | SVRGLD            | 11.9 | 0.244 | 0.119 | 0.0679 | 0.0482 |\n\n\n**Table 2. KL and Fisher divergences on Gaussian measured by gradient computation** \n| **KL divergence \\ Gradient Computations** | 0 | 2000 | 5000 | 8000 | 10000 | **Fisher divergence \\ Gradient Computations** | 0 | 2000 | 5000 | 8000 | 10000 |\n| -------------------------------- | -------- | ------- | ------ | ------- | ------- | ---------------------------- | ------ | ------- | ------- | ------ |------ |\n| LD                    | 6.12 | **0.43** | 0.0214 | 1.06e-3 | 1.44e-4 | LD  | 11.9 | 0.259 | 0.102 | 0.0655 | 0.0424 |\n| SGLD                  | 6.12 | 1.76e-4 | 8.99e-6 | 1.71e-6 | 7.56e-7 |SGLD   | 11.9 | 3.64e-4 | 1.88e-5 | 3.73e-6 | 1.87e-6 |\n| SVRGLD                | 6.12 | **<1e-16** | **<1e-16** | **<1e-16** | **<1e-16** | SVRGLD   | 11.9 | **<1e-16** | **<1e-16** | **<1e-16** | **<1e-16** |\n\n **Table 3. KL divergence and Fisher divergenceson on mixture Gaussian measured by update steps** \n|  **KL divergence \\ Update Steps**  | 0 | 200 | 500 | 800 | 1000 |  **Fisher divergence \\ Update Steps**     | 0 | 200 | 500 | 800 | 1000 |\n| -------------------------------- | -------- | ------- | ------ | ------- | ------- | ---------------------------- | ------ | ------- | ------- | ------ |------ |\n| LD            | 0.57 | **0.33** | **0.07** | **0.082** | **0.076** |LD                | 341.2 | 0.252 | 0.248 | **0.014** | **0.018** |\n|SGLD        | 0.57 | 0.53 | 0.33 | 0.127 | 0.112 |SGLD            | 341.2 | 0.362 | **0.143** | 0.019 | 0.022 |\n|SVRGLD        | 0.57 | 0.39 | 0.16 | 0.092 | 0.082 |SVRGLD           | 341.2 | **0.135** | 0.036 | 0.017 | 0.021 |\n\n\n**Table 4. KL divergence and Fisher divergenceson on mixture Gaussian measured by gradient computations**\n| **Gradient Computations** | 0 | 1000 | 2500 | 4000 | 5000 | **Gradient Computations** | 0 | 1000 | 2500 | 4000 | 5000 |\n| -------------------------------- | -------- | ------- | ------ | ------- | ------- | ---------------------------- | ------ | ------- | ------- | ------ |------ |\n| LD                   | 0.57 | 0.33 | 0.07 | 0.082 | 0.076 | LD                    | 341.2 | 0.252 | 0.248 | 0.014 | 0.018 |\n| SGLD                 | 0.57 | 0.23 | 0.124 | 0.074 | 0.062 |SGLD                | 341.2 | **0.219** | **0.237** | 0.009 | 0.006 |\n|SVRGLD                | 0.57 | **0.21** | **0.062** | **0.072** | **0.032** |SVRGLD               | 341.2 | 0.659 | 0.241 | **0.0081** | **0.0049** |\n\n**Q2**: “Does the technology for constructing gradient flow (including SGD and SVRG flows) mentioned in this paper largely depends on the special properties of exponential mapping”\n\n**A2**: Yes, you are correct, our conclusion highly depends on the linear structure of exponential map. Roughly speaking, which makes the corresponded “update moving” of random vectors in Euclidean space is linear. Based on this, we can take $\\eta\\to \\infty$ and derive the SDE and corresponded flows in Wasserstein space. \n\n**Q3**: “The exponential mappings are different when using different metrics. Hence, if using another Riemannian metric instead of the inner product in the Euclidean space, the resulting exponential may not possess the ``linearity’’. Then, how do we construct the gradient flow in that case?”  \n\n**A3**: This is a good question. Technically, the Riemannian metric induced the distance in Wasserstein space, when the inner product is replaced (the definition above (3)), the distance in probability measure space is no longer Wasserstein space. As in Euclidean space the inner product $\\langle a, b\\rangle = a^{\\top}b$ is replaced by $\\langle a, b\\rangle_{H} = a^{\\top}Hb$ for some positively definite matrix $H$, then the distance between vectors $a$ and $b$ becomes $||a - b||_{H}^{2}$. \n\nBack to your question, when the inner product is changed and specifying the corresponded exponential map, we can similarly construct the Riemannian GD flow by following the idea in Figure 2. However, as you said, the new exponential map may not possess the ``linearity’’. Therefore, it may take some efforts to construct some differential structure in the new update rule. Fortunately, when replacing the inner product $\\langle u, v\\rangle_{\\pi} = \\int \\langle u, v\\rangle d\\pi$ with $\\langle u, v\\rangle_{\\pi} = \\int \\langle u, v\\rangle_{H} d\\pi$ for some positively definite matrix, the resulting exponential map still poses the linearity. \n\n**Q4**: How do the results extend to non-smooth or non-log-concave distributions, which are common in practice?\n\n**A4**: Firstly, we prove the convergence rates under non-log-concave distributions in the first parts of Theorem 1, 2, 3, which address your second concerns. \n\nSecondly, the smoothness Assumptions 1 and 2 are used to bridge the gap between discrete optimization methods and the continuous Riemannian flows, i.e., Propositions 2, 4, 5, which is unavoidable due to the theory of numerical ODE. However, without the smoothness Assumptions **does not influence** the convergence rates of our proposed flows. The conclusions in Theorems 1, 2, 3 **do not** rely on the smoothness Assumptions. \n\n**Q5**: Small size of equation (25)\n\n**A5**: Thanks for pointing out this, we will properly resize it. \n\nRef: \n[1] Kinoshita et al., 2022. Improved Convergence Rate of Stochastic Gradient Langevin Dynamics with Variance Reduction and its Application to Optimization."}, {"review_id": "0sHlgt7DGz", "reviewer": "Reviewer_ahtD", "summary": "This paper studies the problem of minimizing the KL divergence where the target distribution is stochastic. A solution to solve this problem is for instance the Stochastic Gradient Langevin dynamic (SGLD). In this work, this problem is studied from the Wasserstein gradient flow point of view. It is proposed to solve this problem through stochastic Riemannian gradient descent, and it is shown that this algorithm approximates the trajectory of a flow satisfying a continuity equation given in closed-form. Moreover, it is also shown that SGLD can be seen as a discretization of this continuity equation. Then, the convergence of the continuous flow is shown under mild assumption, and the PL inequality. Then, the same work is done for the stochastic Riemannian variance reduction gradient flow.", "questions": "In Section 3 and 4 giving preliminaries on the Wasserstein space and on Wasserstein gradient of the KL divergence, there are no assumptions of regularity on densities. For instance, I think the KL divergence has only subgradients on the Wasserstein space, and only assuming enough regularity, see e.g. [1] (Theorem 10.4.9) or [2]. I believe this should be clarified.\n\nIn Remark 1, it is stated \"Unfortunately, unlike our method, this method is not applied to the stochastic optimization problem in Section 5, because the aforementioned minimization problem can not involve stochastic gradient as expected.\" However, stochastic proximal algorithm also exist, and could be used to solve the same problem as Section 5, equation (11). While it may not have been done yet, I don't think the statement is true.\n\nIn this work, you analyze only the convergence of the continuous dynamic. Would it be possible to study also the convergence of the discrete schemes?\n\nOther objectives are solved using a stochastic Wasserstein gradient descent, e.g. barycenters [3,4] or the Sliced-Wasserstein distance [5]. Could the analysis done in this work be extended to these objectives?\n\nIt could be nice to add experiments demonstrating the rates, even in the Gaussian case where everything may be known in closed-form? \n\n**Typos**\n- Line 113: \"As can be seen, with this lemma and (2), we can directly get the direction of $\\pi_t$ corresponded to a SDE\": sentence not very clear, \"corresponded\"\n- Line 126: \"we seek to corresponding the discrete dynamics\"\n- Line 132-133: \"the corresponded F-P\" -> corresponding\n- Line 181: \"Thus, the target distribution becomes $\\mu$ Proposition (3) in the sequel.\"\n- Equation (23): lack a square on $\\|\\Gamma_\\mu^\\pi(u)\\|$\n\n\n[1] Ambrosio, L., Gigli, N., & Savaré, G. (2008). Gradient flows: in metric spaces and in the space of probability measures. Springer Science & Business Media.\n\n[2] Xu, Y., & Li, Q. (2024). Forward-Euler time-discretization for Wasserstein gradient flows can be wrong. arXiv preprint arXiv:2406.08209.\n\n[3] Backhoff-Veraguas, J., Fontbona, J., Rios, G., & Tobar, F. (2022). Bayesian learning with Wasserstein barycenters. ESAIM: Probability and Statistics, 26, 436-472.\n\n[4] Chewi, S., Maunu, T., Rigollet, P., & Stromme, A. J. (2020, July). Gradient descent algorithms for Bures-Wasserstein barycenters. In Conference on Learning Theory (pp. 1276-1304). PMLR.\n\n[5] Liutkus, A., Simsekli, U., Majewski, S., Durmus, A., & Stöter, F. R. (2019, May). Sliced-Wasserstein flows: Nonparametric generative modeling via optimal transport and diffusions. In International Conference on machine learning (pp. 4104-4113). PMLR.", "limitations": "yes", "rating": 5, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 3, "originality": 3, "strengths_and_weaknesses": "The point of view of analyzing the minimization of the KL divergence with a stochastic target is a nice contribution, which is complementary to works done on SGLD and SVRG. The paper is pretty clear, with a first section analyzing the minimization of the KL divergence under no stochasticity, and then diving into the stochastic versions of the problems. Moreover, there are examples on Gaussian allowing to verify the rates. Overall, I believe this is an interesting article, but it also has some weaknesses that I will describe below.\n\n\n**Strengths**:\n- A clear analysis of the minimization of the KL with and without stochasticity\n- Analysis of stochastic gradient descent and stochastic variance reduction gradient on the Wasserstein space, with convergence results\n- The strategy to get a continuity equation in the stochastic case is interesting\n- Example on Gaussian\n\n\n**Weaknesses**:\n- I believe some technical assumptions are lacking (e.g. the analysis is restricted to absolutely continuous distributions, but it is never mentioned...)\n- There are no numerical results demonstrating the theoretical results", "was_revised": true, "final_justification": "This works derive continuity equations for the continuous counterpart of the stochastic gradient descent and stochastic variance reduction gradient descent on the Wasserstein space. In my opinion, these results are quite interesting, and the authors made connections with known algorithms such as SGLD and SVRGLD, showed the convergence of the flows under mild assumptions, and provided an example in the Gaussian setting.\n\nThe original submission is however not very clear, not always very rigorous, and lacks numerical experiments. Nonetheless, I think some of the unclear things have been clarified in the rebuttal, and experiments were provided. Therefore, I believe that the main concerns have been adressed, and I would be inclined to suggest acceptance.", "rebuttal": "We thank for your valuable comment and address your concerns as follows.\n\n**Q1**: I believe some technical assumptions are lacking (e.g. the analysis is restricted to absolutely continuous distributions, but it is never mentioned...)\n\n**A1**: Thanks for pointing it. In this paper we assume the target distribution $\\mu\\propto \\exp(-V(x))$ in line 102. Then the KL $D_{KL}(\\pi\\parallel \\mu)\\propto \\int V(x)d\\pi(x) - \\int \\log{\\pi(x)}d\\pi(x)$, so that the KL divergence exists whenever $\\pi(x)$ is absolutely continuous to Lebesgue measure. This is true when the initialization distribution $\\pi_{0}$ satisfies this condition, since the discussed $\\pi_{t}$ satisfies the F-P equation (2). We will add this discussion in the revised version to make our conclusion more precises. \n\n**Q2**: There are no **numerical results** demonstrating the theoretical results\n\n**A2**: Thanks for your valuable advice. We have made an example (Example 1 in our paper) of sampling on Gaussian distributions to illustrate the tightness of our theoretical results. \n\n**However, we recognize that numerical results will make our conclusions more convincing**. We evaluate different sampling algorithms on two experiments. **Since the NeurIPS rebuttal do not allow us to convey these results by image**, we present them in the following tables (results are averaged by 5 independent runs). \n\nConcretely, we set the each target distribution $\\mu_{\\xi}$ as **Gaussian or mixture-Gaussian**, and conduct sampling by the proposed algorithms. Notably, our convergence rates are respectively evaluated by update steps ($T$) and gradient computation (Riemannian SGD and Riemannian SVRG has less gradient computations than Riemannian GD). The results are summarized in the following tables, the KL divergence and Fisher divergence are estimated by combining MCMC and an estimated density as in [1]. \n\nAs can be seen, **under the same update steps**, Riemannian SVRG and Riemannian GD have better convergence results than Riemannian SGD (see Theorem 1, 2, 3), while **under the same gradient computations**, Riemannian SGD and Riemannian SVRG (the two algorithms have less gradient computations in each update step) have better convergence results (see discussion after Theorem 3). These results are consistent with our conclusion in paper. **We promise to add these numerical results in the revised version**, and we sincerely hope you may reconsider the concern on numerical results.   \n\n**Table1. KL and Fisher divergences on Gaussian measured by update steps**\n| **KL divergence \\ Update Steps** | 0    | 20   | 50     | 80      | 100     | **Fisher divergence \\ Update Steps** | 0    | 20    | 50    | 80   | 100   |\n| -------------------------------- | ---- | ---- | ------ | ------- | ------- | -------------------------------- | ---- | ----- | ----- | ------ | ------ |\n| LD           | 6.12 | **0.43** | 0.0214 | **1.06e-3** | **1.44e-4** | LD      | 11.9 | **0.259** | **0.102** | **0.0655** | **0.0424** |\n| SGLD                 | 6.12 | 0.51 | 0.0274 | 1.26e-3 | 1.62e-4 | SGLD         | 11.9 | 0.279 | 0.142 | 0.0723 | 0.0512 |\n| SVRGLD               | 6.12 | 0.67 | **0.012** | 1.16e-3 | 1.52e-4  | SVRGLD            | 11.9 | 0.244 | 0.119 | 0.0679 | 0.0482 |\n\n\n**Table 2. KL and Fisher divergences on Gaussian measured by gradient computation** \n| **KL divergence \\ Gradient Computations** | 0 | 2000 | 5000 | 8000 | 10000 | **Fisher divergence \\ Gradient Computations** | 0 | 2000 | 5000 | 8000 | 10000 |\n| -------------------------------- | -------- | ------- | ------ | ------- | ------- | ---------------------------- | ------ | ------- | ------- | ------ |------ |\n| LD                    | 6.12 | **0.43** | 0.0214 | 1.06e-3 | 1.44e-4 | LD  | 11.9 | 0.259 | 0.102 | 0.0655 | 0.0424 |\n| SGLD                  | 6.12 | 1.76e-4 | 8.99e-6 | 1.71e-6 | 7.56e-7 |SGLD   | 11.9 | 3.64e-4 | 1.88e-5 | 3.73e-6 | 1.87e-6 |\n| SVRGLD                | 6.12 | **<1e-16** | **<1e-16** | **<1e-16** | **<1e-16** | SVRGLD   | 11.9 | **<1e-16** | **<1e-16** | **<1e-16** | **<1e-16** |\n\n **Table 3. KL divergence and Fisher divergenceson on mixture Gaussian measured by update steps** \n|  **KL divergence \\ Update Steps**  | 0 | 200 | 500 | 800 | 1000 |  **Fisher divergence \\ Update Steps**     | 0 | 200 | 500 | 800 | 1000 |\n| -------------------------------- | -------- | ------- | ------ | ------- | ------- | ---------------------------- | ------ | ------- | ------- | ------ |------ |\n| LD            | 0.57 | **0.33** | **0.07** | **0.082** | **0.076** |LD                | 341.2 | 0.252 | 0.248 | **0.014** | **0.018** |\n|SGLD        | 0.57 | 0.53 | 0.33 | 0.127 | 0.112 |SGLD            | 341.2 | 0.362 | **0.143** | 0.019 | 0.022 |\n|SVRGLD        | 0.57 | 0.39 | 0.16 | 0.092 | 0.082 |SVRGLD           | 341.2 | **0.135** | 0.036 | 0.017 | 0.021 |\n\n\n**Table 4. KL divergence and Fisher divergenceson on mixture Gaussian measured by gradient computations**\n| **Gradient Computations** | 0 | 1000 | 2500 | 4000 | 5000 | **Gradient Computations** | 0 | 1000 | 2500 | 4000 | 5000 |\n| -------------------------------- | -------- | ------- | ------ | ------- | ------- | ---------------------------- | ------ | ------- | ------- | ------ |------ |\n| LD                   | 0.57 | 0.33 | 0.07 | 0.082 | 0.076 | LD                    | 341.2 | 0.252 | 0.248 | 0.014 | 0.018 |\n| SGLD                 | 0.57 | 0.23 | 0.124 | 0.074 | 0.062 |SGLD                | 341.2 | **0.219** | **0.237** | 0.009 | 0.006 |\n|SVRGLD                | 0.57 | **0.21** | **0.062** | **0.072** | **0.032** |SVRGLD               | 341.2 | 0.659 | 0.241 | **0.0081** | **0.0049** |\n\n**Q3**: In Section 3 and 4 giving preliminaries on the Wasserstein space and on Wasserstein gradient of the KL divergence, there are no assumptions of regularity on densities.\n\n**A3**: Yes, you are correct, the Wasserstein space $\\mathcal{P}$ in this paper should be $\\mathcal{P_{2,ac}}$ which means the set of probability absolutely continuous  measures with finite $L_{2}$-norm. These conditions can be satisfied under an absolutely continuous $\\pi_{0}$ as mentioned in A1. We will revise it as you suggested. On the other, the definition of Riemannian gradient is similar to the definition of Wasserstein gradient in Section 1.4.1 of [1] and Riemannian gradient in Page 27 of [2]. Mathematically, it can be viewed as sub-gradient. But it is usually named as Riemannian gradient in optimization on manifold community. We will clarify this in the revised version. \n\n**Q4**: In Remark 1, it is stated \"Unfortunately, unlike our method, this method is not applied to the stochastic optimization problem in Section 5, because the aforementioned minimization problem can not involve stochastic gradient as expected.\" However, stochastic proximal algorithm also exist, and could be used to solve the same problem as Section 5, equation (11). While it may not have been done yet, I don't think the statement is true.\n\n**A4**: Our expression in Remark may leads to a little bit misunderstanding. We say solving $\\pi_{n + 1} = \\arg\\min_{\\pi}D_{KL}(\\pi\\parallel \\mu) + W_{2}^{2}(\\pi, \\pi_{n}) / 2\\eta$, when $\\eta\\to 0$ \"can not generalize to stochastic problem\" means **the method can not induce a continuous flow in stochastic version**. Let us check this in Euclidean space will be more clear. \n\nConsider $x_{n + 1} = \\arg\\min_{x} f(x) + \\frac{||x – x_{n}||^{2}}{2\\eta}$, we get $(x_{n + 1} – x_{n}) / \\eta = -\\nabla f(x_{n})$ so that $\\eta\\to 0$ implies the gradient flow $\\frac{dx_{t}}{dt} = -\\nabla f(x_{t})$. However, when the objective is stochastic i.e., $E_{\\xi}[f_{\\xi_{\\eta}}]$, the sub-problem becomes $ x_{n + 1} = \\arg\\min_{x} f_{\\xi_{n}}(x) + \\frac{||x – x_{n}||^{2}}{2\\eta}$, then $(x_{n + 1} – x_{n}) / \\eta = -\\nabla f_{\\xi_{n}}(x_{n})$, which implies $\\frac{dx_{t}}{dt} = -\\nabla f_{\\xi_{t}}(x_{t})$ with $\\xi_{t}$ is a sample of $\\xi$. As can be seen, the drift term has randomness, which is not expected in Wasserstein space (a **deterministic metric space consists of probability measures**). Because the direction $v$ of curves $\\frac{\\partial{\\pi_{t}}}{\\partial{t}} = -\\nabla\\cdot(\\pi_{t}v)$ in Wasserstein space should be deterministic. We will make this more clear in the revised version. \n\n**Q5**: In this work, you analyze only the convergence of the continuous dynamic. Would it be possible to study also the convergence of the discrete schemes?\n\n**A6**: The proof techniques of our continuous methods can be generalized to discrete schemes, by reusing our proposed Lyapunov function. A more straightforward extrapolation is directly combining our convergence rates of continuous methods with the proved gap between discrete and continuous methods in Propositions 2, 4, 5. Though this can be done when log-Sobolev inequality is satisfied i.e., $\\frac{\\gamma}{2}W_{2}^{2}(\\pi, \\mu) \\leq D_{KL}(\\pi\\parallel \\mu)$ (see our Proposition 6 in Appendix). \n\n**Q7**: Other objectives are solved using a stochastic Wasserstein gradient descent. Could the analysis done in this work be extended to these objectives?\n\n**A7**: Yes, the key idea is using our transformation in Figure 2 to link the constructed SDE and curve in Wasserstein space. \n\n**Q8**: It could be nice to add experiments demonstrating the rates, even in the Gaussian case where everything may be known in closed-form?\n\n**A8**: Thanks for pointing this, we have added the numerical experiments as clarified in Q2. \n\nRef:\n\n[1] Kinoshita et al., Improved Convergence Rate of Stochastic Gradient Langevin Dynamics with Variance Reduction and its Application to Optimization"}, {"review_id": "LhrJAABqYL", "reviewer": "Reviewer_oUnb", "summary": "This paper introduces a continuous‐time analogue of stochastic gradient descent (SGD) and stochastic variance–reduced gradient (SVRG) methods formulated on the manifold of probability measures endowed with the second‐order Wasserstein metric.  Building on the well‐known correspondence between Riemannian gradient flow in this space and Langevin diffusion, the authors construct Riemannian SGD and SVRG flows.  By interpreting the discrete Riemannian SGD and SVRG updates as transformations of random vectors in Euclidean space, and then taking the learning‐rate limit, they derive the corresponding stochastic differential equations (SDEs).  Applying the Fokker–Planck equation translates these SDEs into deterministic “probability‐measure flows” on the Riemannian manifold. The authors establish some convergence guarantees under Lipschitz‐gradient assumptions and they prove that the Riemannian SGD flow attains a first‐order stationary rate of $O(1/\\sqrt{T})$ (and $O(1/T)$ under a log‐Sobolev/PL inequality), while the Riemannian SVRG flow achieves $O(N^{2/3}/T)$ (and exponential convergence under PL).  These match the best-known Euclidean continuous‐time rates The authors show that standard discrete SGLD and variance‐reduced Langevin MCMC algorithms coincide as time‐discretizations of their proposed flows, thereby linking the continuous‐time structure underlying these sampling methods.", "questions": "* If the authors are able, then adding some numerical examples to the paper could significantly improve this work and help to solidify the application of this approach. \n* The language and grammar used throughout the paper is very challenging to read. This would need to be significantly improved before this work is published in NeurIPS.  \n* Could the authors comment on how their theoretical results could (or perhaps could not) be applied to functionals other than KL?", "limitations": "There are no ethical or societal issues that the authors need to address.", "rating": 3, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 2, "clarity": 1, "significance": 2, "originality": 1, "strengths_and_weaknesses": "Quality\n\nThe paper exhibits a high level of mathematical rigour: all constructions and convergence proofs are presented with precise assumptions—such as Lipschitz continuity and log-Sobolev or Polyak–Łojasiewicz inequalities—and are thoroughly detailed in the appendices. It also provides a comprehensive treatment of both nonconvex regimes and the faster rates attainable under PŁ conditions, unifying stochastic gradient and variance-reduced flows in a single framework. However, its reliance on strong technical conditions may be a small limit of its wider applicability. Verifying Lipschitz-gradient or log-Sobolev constants for complex target measures is often challenging in practice. Moreover, the absence of any empirical validation leaves open questions about the practical sharpness of the theoretical rates and the behaviour of the flows outside simplified settings.\n\nClarity\n\nWhile the manuscript is structured in a logical sequence, from stating the important preliminaries through to SGD and SVRG flows, the presentation suffers from awkward phrasing, grammatical errors, and convoluted sentence structures that permeate pretty much throughout every section. Definitions of measure-flow notation and details about Riemannian manifolds are technically sound, but the poor use of English frequently forces the reader to pause and re-read passages to extract the intended meaning.  These linguistic challenges, combined with the dense theoretical material and very little intuitive exposition, presents a significant barrier to comprehension for all but the most specialised experts in geometric measure theory and continuous-time stochastic dynamics.\n\nSignificance\n\nBy proposing a continuous-time analogue of stochastic gradient and variance-reduced methods on the Wasserstein manifold, the authors provide a new theoretical bridge between optimisation flows and Langevin-type sampling algorithms. This unification could lead to new algorithmic designs that blend optimisation and sampling perspectives. Yet, the paper stops short of connecting these continuous flows to concrete algorithmic implementations or demonstrating improvements over existing discrete methods. As a result, its immediate impact on practical machine-learning tasks remains speculative until accompanied by empirical studies. \n\nOriginality\n\nMethodologically, the use of the Fokker–Planck framework to lift Euclidean SDE limits into deterministic manifold flows represents an interesting idea for the Riemannian setting (but not unique for the Euclidean setting), and the insight that SGLD and variance-reduced Langevin samplers arise as discretisations of the same underlying flows gives a new view on this relationship. While these theoretical connections are interesting, many of the resulting convergence rates mirror those already known in the Euclidean setting, making some of the performance guarantees anticipated by analogy. The focus on KL-divergence minimisation also narrows the scope, leaving open the question of how readily the framework extends to other divergences or objective functions. Furthermore, there has been a lot of recent work in this area, and it is not entirely clear how much of this work is original in the context of other recent papers.", "was_revised": true, "final_justification": "The authors have addressed the questions that I asked, and if these changes are included in the paper, then I believe it warrants an increase in my score. However, as I have noted to the authors, I am still concerned about the readability of the paper and the limited range of experimental results provided.", "rebuttal": "We deeply grateful for your careful reviewing and valuable comments. We address your concerns as below. \n\n**Q1**: Strong technical conditions (Lipschitz-gradient or log-Sobolev constants).\n\n**A1**: For the Lipschitz-gradient Assumption 1 and 2, they are **only applied to bridge the gap between discrete and continuous methods** (Proposition 2, 4, 5), while our convergence results in Theorem 1, 2, 3 **do not require** them. Besides, the Lipschitz assumption is standard to bridge the discrete and continuous differential equation (as we did) in numerical analysis (see section of convergence in [4]). \n\nFor log-Sobolev inequality, we also characterize the convergence rates of our methods without it in Theorems 1, 2, 3, evaluated by the norm of Riemannian gradient as in **standard non-convex optimization**. Since log-Sobolev inequality roughly implies \"strongly-convexity\" as clarified in our paper. \n\n**Q2**: The absence of any empirical validation.\n\n**A2**: Thanks for your valuable advice. We have made an example (Example 1) of sampling on Gaussian to illustrate the tightness of our theoretical results. \n\n**However, we recognize that numerical results will make our conclusions more convincing**. We evaluate different sampling algorithms on two experiments. **Since the NeurIPS rebuttal do not allow us to convey these results by image**, we present them in the following tables (results are averaged by 5 independent runs). \n\nConcretely, we set the each target distribution $\\mu_{\\xi}$ as **Gaussian or mixture-Gaussian**, and conduct sampling by the proposed algorithms. Notably, our convergence rates are respectively evaluated by update steps ($T$) and gradient computation (Riemannian SGD and SVRG has less gradient computations than Riemannian GD). The results are summarized in Tables 1-4, the KL divergence and Fisher divergence are estimated by combining MCMC and an estimated density as in [1]. \n\nAs can be seen, **under the same update steps**, Riemannian SVRG and Riemannian GD have better convergence results than Riemannian SGD (see Theorem 1, 2, 3), while **under the same gradient computations**, Riemannian SGD and Riemannian SVRG (the two algorithms have less gradient computations in each update step) have better convergence results (see discussion after Theorem 3). These results are consistent with our conclusion in paper. **We promise to add these numerical results in the revised version**, and we sincerely hope you may reconsider the concern on numerical results.   \n\n**Table1. KL and Fisher divergences on Gaussian measured by update steps**\n| **KL divergence \\ Update Steps** | 0    | 20   | 50     | 80      | 100     | **Fisher divergence \\ Update Steps** | 0    | 20    | 50    | 80   | 100   |\n| -------------------------------- | ---- | ---- | ------ | ------- | ------- | -------------------------------- | ---- | ----- | ----- | ------ | ------ |\n| LD           | 6.12 | **0.43** | 0.0214 | **1.06e-3** | **1.44e-4** | LD      | 11.9 | **0.259** | **0.102** | **0.0655** | **0.0424** |\n| SGLD                 | 6.12 | 0.51 | 0.0274 | 1.26e-3 | 1.62e-4 | SGLD         | 11.9 | 0.279 | 0.142 | 0.0723 | 0.0512 |\n| SVRGLD               | 6.12 | 0.67 | **0.012** | 1.16e-3 | 1.52e-4  | SVRGLD            | 11.9 | 0.244 | 0.119 | 0.0679 | 0.0482 |\n\n\n**Table 2. KL and Fisher divergences on Gaussian measured by gradient computation** \n| **KL divergence \\ Gradient Computations** | 0 | 2000 | 5000 | 8000 | 10000 | **Fisher divergence \\ Gradient Computations** | 0 | 2000 | 5000 | 8000 | 10000 |\n| -------------------------------- | -------- | ------- | ------ | ------- | ------- | ---------------------------- | ------ | ------- | ------- | ------ |------ |\n| LD                    | 6.12 | **0.43** | 0.0214 | 1.06e-3 | 1.44e-4 | LD  | 11.9 | 0.259 | 0.102 | 0.0655 | 0.0424 |\n| SGLD                  | 6.12 | 1.76e-4 | 8.99e-6 | 1.71e-6 | 7.56e-7 |SGLD   | 11.9 | 3.64e-4 | 1.88e-5 | 3.73e-6 | 1.87e-6 |\n| SVRGLD                | 6.12 | **<1e-16** | **<1e-16** | **<1e-16** | **<1e-16** | SVRGLD   | 11.9 | **<1e-16** | **<1e-16** | **<1e-16** | **<1e-16** |\n\n **Table 3. KL divergence and Fisher divergenceson on mixture Gaussian measured by update steps** \n|  **KL divergence \\ Update Steps**  | 0 | 200 | 500 | 800 | 1000 |  **Fisher divergence \\ Update Steps**     | 0 | 200 | 500 | 800 | 1000 |\n| -------------------------------- | -------- | ------- | ------ | ------- | ------- | ---------------------------- | ------ | ------- | ------- | ------ |------ |\n| LD            | 0.57 | **0.33** | **0.07** | **0.082** | **0.076** |LD                | 341.2 | 0.252 | 0.248 | **0.014** | **0.018** |\n|SGLD        | 0.57 | 0.53 | 0.33 | 0.127 | 0.112 |SGLD            | 341.2 | 0.362 | **0.143** | 0.019 | 0.022 |\n|SVRGLD        | 0.57 | 0.39 | 0.16 | 0.092 | 0.082 |SVRGLD           | 341.2 | **0.135** | 0.036 | 0.017 | 0.021 |\n\n\n**Table 4. KL divergence and Fisher divergenceson on mixture Gaussian measured by gradient computations**\n| **Gradient Computations** | 0 | 1000 | 2500 | 4000 | 5000 | **Gradient Computations** | 0 | 1000 | 2500 | 4000 | 5000 |\n| -------------------------------- | -------- | ------- | ------ | ------- | ------- | ---------------------------- | ------ | ------- | ------- | ------ |------ |\n| LD                   | 0.57 | 0.33 | 0.07 | 0.082 | 0.076 | LD                    | 341.2 | 0.252 | 0.248 | 0.014 | 0.018 |\n| SGLD                 | 0.57 | 0.23 | 0.124 | 0.074 | 0.062 |SGLD                | 341.2 | **0.219** | **0.237** | 0.009 | 0.006 |\n|SVRGLD                | 0.57 | **0.21** | **0.062** | **0.072** | **0.032** |SVRGLD               | 341.2 | 0.659 | 0.241 | **0.0081** | **0.0049** |\n\n**Q3**: The language and grammar used throughout the paper is very challenging to read.\n\n**A3**: We really grateful for your carefully reviewing all these materials with heavy mathematical notations. We **promise** we will revise our paper as follows to make our paper more readable. \n\nWe will carefully check the paper word by word to fix all unproper phrasing, gramma issues, and convoluted sentence structures. \n\nWe will rewrite the part of theoretical background, add more intuitive exposition to make these materials friendly to all audiences.    \n\n\n**Q4**: Yet, the paper stops short of connecting these continuous flows to concrete algorithmic implementations or demonstrating improvements over existing discrete methods.\n\n**A4**: By combining our proved convergence rates of continuous methods and the proved gap between discrete and continuous methods can address your concerns. Roughly speaking, by Proposition 6, we know the Wasserstein distance is upper bounded by KL divergence (proved to be converge) $\\frac{\\gamma}{2}W_{2}^{2}(\\pi, \\mu) \\leq D_{KL}(\\pi\\parallel \\mu)$. Since Propositions 2, 4, 5 imply $W_{2}^{2}(\\hat{\\pi}_{n}, \\pi_{n\\eta}) \\leq E[||x_{n} - \\hat{x}_{n\\eta}||^{2}]\\leq \\mathcal{O}(\\eta)$ ($\\hat{\\pi}_{n}$ and $\\pi_{n\\eta}$ are respectively obtained by proposed discrete and continuous algorithms), we get the convergence rates of the discrete methods, which further demonstrates the improvements over existing discrete methods. **We promise to add these analysis to in the revised version.**\n\n**Q5**: The resulting convergence rates mirror those already known in the Euclidean setting.  \n\n**A5**: Though the convergence rates match the ones in Euclidean space, **but the proof techniques are more complicated**. For example, it takes us a lot efforts to construct the Lyapunov function in (114), which is critical to prove the convergence rate. The construction relies techniques from both optimization and stochastic analysis, which are more sophisticated compared with the scheme in Euclidean space. \n\n**Q6**: Extensions to other objective functions more than KL divergence. \n\n**A6**: Our framework has the potential to be extended to other objectives. As can be seen, our core idea in Figure 2 **does not depend on the optimization objective**, which guarantees the extensions of flows to the other objectives. However, it worth noting that a **practical algorithm** may only be obtained under KL divergence. \n\nTaking Riemannian gradient as an example, it takes the formulation of $\\frac{\\partial{\\pi_{t}}}{\\partial{t}} = \\nabla\\cdot(\\pi_{t}\\mathrm{grad} F(\\pi_{t}))$. The flow can be implemented in practice only if $\\mathrm{grad} F(\\pi_{t})$ is efficiently computed. However, this is not true for many objectives. For example, then $F(\\pi) = W_{2}^{2}(\\pi, \\mu)$, $\\mathrm{grad} F(\\pi) = -2(T_{\\pi\\rightarrow \\mu} - \\mathrm{id})$ (see Theorem 1.4.11 in [1]), the $ T_{\\pi\\rightarrow \\mu}$ is the optimal transportation map between $\\pi$ and $\\mu$, which is $\\mathcal{O}(e^{d})$ complex in general. \n\nHowever, for KL divergence, $\\frac{\\partial{\\pi_{t}}}{\\partial{t}} =  -\\nabla\\cdot(\\pi_{t}\\nabla\\log{\\mu}) + \\Delta\\pi_{t}$, corresponds to SDE $dx_{t} = \\nabla\\log{\\mu(x_{t})}dt + \\sqrt{2}dW_{t}$, which makes it practical. This explains why the existing gradient flow in Wasserstein space is mainly discussed under KL divergence [1, 2, 3]. \n\n**Q7**: Furthermore, there has been a lot of recent work in this area, and it is not entirely clear how much of this work is original in the context of other recent papers.\n\n**A7**: As clarified in line 72, the existing literature focus on discrete Langevin dynamics and its convergence rates under different criteria [1, 2, 3]. The main contributions in this paper **the continuous Riemannian SGD and SVRG flows** are not explored yet. \n\n**Q8-10**: The numerical verification; Grammar issues; The extensions to other divergence.\n\n**A8-10**: The answers refer to A2, A3, and A6 in above. \n \n Ref:\n\n[1] Kinoshita et al., Improved Convergence Rate of Stochastic Gradient Langevin Dynamics with Variance Reduction and its Application to Optimization\n\n[2] Cheng et al., Convergence of langevin mcmc in kl-divergence\n\n[3] Chewi et al., Log-concave sampling\n\n[4] https://en.wikipedia.org/wiki/Numerical_methods_for_ordinary_differential_equations"}]

Author rebuttal:

Meta-review: All reviewers have stated that their concerns and questions have been addressed, and I recommend acceptance. That said, the authors will need to incorporate the promised changes and experiments to the camera ready version, including improvements to the presentation that several reviewers deemed crucial.
Decision: Accept (poster)