Title: Continuous-time Riemannian SGD and SVRG Flows on Wasserstein Probabilistic Space
Venue: NeurIPS 2025
Accepted: yes
Keywords: Optimization on Manifold; SGLD; Continuous Gradient Flow
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.


Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS AND WEAKNESSES:\nQuality\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.\n\nQUESTIONS:\n* 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?\n\nLIMITATIONS:\nThere are no ethical or societal issues that the authors need to address.", "TITLE": "Official Review", "RECOMMENDATION": 3, "REVIEWER_CONFIDENCE": 4, "CLARITY": 1, "IMPACT": 2, "ORIGINALITY": 1, "DATE": "2025-06-19"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS AND WEAKNESSES:\nThe 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\n\nQUESTIONS:\nIn 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.\n\nLIMITATIONS:\nyes", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 3, "CLARITY": 3, "IMPACT": 3, "ORIGINALITY": 3, "DATE": "2025-06-21"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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.\n\nSTRENGTHS AND WEAKNESSES:\n**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\n\nQUESTIONS:\n- 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?\n\nLIMITATIONS:\nLimitations are properly addressed.", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 4, "CLARITY": 3, "IMPACT": 3, "ORIGINALITY": 3, "DATE": "2025-06-30"}

comments: 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.

STRENGTHS AND WEAKNESSES:
Strength: Riemannian SGD and SVRG flows are etablished which seems to be new.
Weakness:
* Somehow lack numerical validations on the efficiency of Riemannian SGD and SVRG  on sampling
* It seems the overall  technical details differ not too much from the Euclidean case on the convergence of the flows

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?
* In Lemma 1, better be ``Lemma 1 ([29,42]).'', similar for Theorem 1.
* What are the key differences from the existing proof for the establishment of Theorem 1?
* Eq. (25) too small?

LIMITATIONS:
yes
TITLE: Official Review
CLARITY: 3
IMPACT: 2
ORIGINALITY: 3
DATE: 2025-07-02