# Stochastic Modified Flows for Riemannian Stochastic Gradient Descent

arXiv: https://arxiv.org/abs/2402.03467 · HTML: https://mlsanity.com/paper/2402.03467/

## Key facts

- Cited by 2 papers in the map (data through 14 Oct 2026).
- Builds on 27 of its 64 references that are in the map.
- First posted on arXiv 2 Feb 2024.

## Builds on (most cited first, 27 of 27)

- [Poincar'e Embeddings for Learning Hierarchical Representations](https://mlsanity.com/paper/1705.08039/) (2017)
- [Weight Normalization: A Simple Reparameterization to Accelerate Training of Deep Neural Networks](https://mlsanity.com/paper/1602.07868/) (2016)
- [Stochastic modified equations and adaptive stochastic gradient algorithms](https://mlsanity.com/paper/1511.06251/) (2015)
- [Improving GANs Using Optimal Transport](https://mlsanity.com/paper/1803.05573/) (2018)
- [Stochastic Modified Equations and Dynamics of Stochastic Gradient Algorithms I: Mathematical Foundations](https://mlsanity.com/paper/1811.01558/) (2018)
- [What Happens after SGD Reaches Zero Loss? --A Mathematical Framework](https://mlsanity.com/paper/2110.06914/) (2021)
- [On the Validity of Modeling SGD with Stochastic Differential Equations (SDEs)](https://mlsanity.com/paper/2102.12470/) (2021)
- [Convergence rates for the stochastic gradient descent method for non-convex objective functions](https://mlsanity.com/paper/1904.01517/) (2019)
- [Riemannian SVRG: Fast Stochastic Optimization on Riemannian Manifolds](https://mlsanity.com/paper/1605.07147/) (2016)
- [Neural Embeddings of Graphs in Hyperbolic Space](https://mlsanity.com/paper/1705.10359/) (2017)
- [Projection Robust Wasserstein Distance and Riemannian Optimization](https://mlsanity.com/paper/2006.07458/) (2020)
- [Averaging Stochastic Gradient Descent on Riemannian Manifolds](https://mlsanity.com/paper/1802.09128/) (2018)
- [Escaping from saddle points on Riemannian manifolds](https://mlsanity.com/paper/1906.07355/) (2019)
- [Stochastic gradient descent with noise of machine learning type. Part II: Continuous time analysis](https://mlsanity.com/paper/2106.02588/) (2021)
- [Uniform-in-Time Weak Error Analysis for Stochastic Gradient Descent Algorithms via Diffusion Approximation](https://mlsanity.com/paper/1902.00635/) (2019)
- [On minimal representations of shallow ReLU networks](https://mlsanity.com/paper/2108.05643/) (2021)
- [On Riemannian Stochastic Approximation Schemes with Fixed Step-Size](https://mlsanity.com/paper/2102.07586/) (2021)
- [Robust Implicit Regularization via Weight Normalization](https://mlsanity.com/paper/2305.05448/) (2023)
- [Gradient descent in hyperbolic space](https://mlsanity.com/paper/1805.08207/) (2018)
- [Sinkhorn Natural Gradient for Generative Models](https://mlsanity.com/paper/2011.04162/) (2020)
- [Riemannian stochastic optimization methods avoid strict saddle points](https://mlsanity.com/paper/2311.02374/) (2023)
- [Stochastic Approximation on Riemannian manifolds](https://mlsanity.com/paper/1711.10754/) (2017)
- [Convergence Analysis of Riemannian Stochastic Approximation Schemes](https://mlsanity.com/paper/2005.13284/) (2020)
- [Inductive Bias of Gradient Descent for Weight Normalized Smooth Homogeneous Neural Nets](https://mlsanity.com/paper/2010.12909/) (2020)
- [Unlocking optimal batch size schedules using continuous-time control and perturbation theory](https://mlsanity.com/paper/2312.01898/) (2023)
- [Convergence of Riemannian Stochastic Gradient Descent on Hadamard Manifold](https://mlsanity.com/paper/2312.07990/) (2023)
- [Riemannian stochastic approximation algorithms](https://mlsanity.com/paper/2206.06795/) (2022)

## Built on by (most cited first, 2 of 2)

- [ODE approximation for the Adam algorithm: General and overparametrized setting](https://mlsanity.com/paper/2511.04622/) (2025)
- [Effective continuous equations for adaptive SGD: a stochastic analysis view](https://mlsanity.com/paper/2509.21614/) (2025)

## Abstract (arXiv)

We give quantitative estimates for the rate of convergence of Riemannian stochastic gradient descent (RSGD) to Riemannian gradient flow and to a diffusion process, the so-called Riemannian stochastic modified flow (RSMF). Using tools from stochastic differential geometry we show that, in the small learning rate regime, RSGD can be approximated by the solution to the RSMF driven by an infinite-dimensional Wiener process. The RSMF accounts for the random fluctuations of RSGD and, thereby, increases the order of approximation compared to the deterministic Riemannian gradient flow. The RSGD is build using the concept of a retraction map, that is, a cost efficient approximation of the exponential map, and we prove quantitative bounds for the weak error of the diffusion approximation under assumptions on the retraction map, the geometry of the manifold, and the random estimators of the gradient.