Research map / Optimization / Nonconvex and smooth optimization
Continuous-time Riemannian SGD and SVRG Flows on Wasserstein Probabilistic Space
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Derives continuous Riemannian SGD and SVRG flows on the Wasserstein space for KL minimization by taking the step-size limit of discrete Riemannian methods and applying Fokker-Planck equations, with convergence rates matching Euclidean results.
Problem
Continuous-time counterparts (flows) of Riemannian SGD and SVRG in Wasserstein space for minimizing KL divergence to a target measure are not established, and existing continuous Wasserstein methods do not handle stochastic dynamics.
Approach
Construct SDEs in Euclidean space as the step-size limit of discrete Riemannian SGD/SVRG (via their SGLD/SVRG-Langevin equivalents), then map them to Wasserstein curves via the Fokker-Planck equation, yielding deterministic flows whose convergence rates are proved for non-convex and log-Sobolev (PL-type) settings.
Key results
- Riemannian SGD flow converges at O(1/sqrt(T)) for non-convex KL minimization and O(1/T) under log-Sobolev inequality, matching Euclidean SGD flow rates.
- Riemannian SVRG flow achieves O(N^{2/3}/T) non-convex stationarity rate and exponential convergence e^{-gamma T/N^{2/3}} under log-Sobolev, with complexity O(N^{2/3} epsilon^{-1}).
- Experiments on Gaussian and mixture-of-Gaussian targets (N=100 and N=5) show Riemannian SVRG and GD converge faster per update step than SGD, and SVRG is sharpest per gradient computation.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Use Riemannian GD or Riemannian SVRG instead of Riemannian SGD at the same number of update steps | Riemannian sampling algorithms · Gaussian and Mixture Gaussian | KL divergence | Helped |
| Use Riemannian GD or Riemannian SVRG instead of Riemannian SGD at the same number of update steps | Riemannian sampling algorithms · Gaussian and Mixture Gaussian | Fisher divergence | Helped |
| Use Riemannian SGD or Riemannian SVRG instead of Riemannian GD at the same number of gradient computations | Riemannian sampling algorithms · Gaussian and Mixture Gaussian | KL divergence | Helped |
| Use Riemannian SGD or Riemannian SVRG instead of Riemannian GD at the same number of gradient computations | Riemannian sampling algorithms · Gaussian and Mixture Gaussian | Fisher divergence | Helped |
Surprising findings
- In the Gaussian example the SVRG noise covariance vanishes, so the SVRG flow reduces to the Riemannian GD flow.
Limitations
- Convergence analysis relies on strong regularity assumptions such as Lipschitz log-densities and log-Sobolev inequality.
- Flows are not explicitly computable for mixture targets, so experiments use discrete approximations with density estimation.
- Stationarity measured via Fisher divergence does not imply convergence in total variation.
Experimental setup
Metrics: KL divergence Fisher divergence (Riemannian gradient norm)
Baselines: Riemannian GD flow Riemannian SGD flow
Builds on: Riemannian gradient flow in Wasserstein space SGLD SVRG Langevin dynamics Riemannian SVRG
Closest papers by meaning
- Acceleration via silver step-size on Riemannian manifolds with applications to Wasserstein space (NeurIPS 2025 · similarity 0.89)
- SERENA: A Unified Stochastic Recursive Variance Reduced Gradient Framework for Riemannian Non-Convex Optimization (ICML 2025 · similarity 0.88)
- Hessian-guided Perturbed Wasserstein Gradient Flows for Escaping Saddle Points (NeurIPS 2025 · similarity 0.87)
- Improved Finite-Particle Convergence Rates for Stein Variational Gradient Descent (ICLR 2025 · similarity 0.86)
- Stochastic variance-reduced Gaussian variational inference on the Bures-Wasserstein manifold (ICLR 2025 · similarity 0.86)
- Learning-Rate-Free Stochastic Optimization over Riemannian Manifolds (ICML 2024 · similarity 0.86)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.