Stochastic Modified Flows for Riemannian Stochastic Gradient Descent
Benjamin Gess, Sebastian Kassing, Nimit Rana
TL;DR
This work studies stochastic optimization on Riemannian manifolds by tying Riemannian SGD (RSGD) to continuous-time limits. It first proves a weak order-1 approximation of RSGD by the Riemannian gradient flow $\dot z_t = -\mathrm{grad}\, f(z_t)$ under uniform first-order retractions and bounded geometry. It then introduces the Riemannian stochastic modified flow (RSMF), a diffusion on $M$ driven by a cylindrical Wiener process, and shows a weak order-2 approximation of RSGD by RSMF under stronger regularity and second-order retractions. The results are supported by a robust geometric framework (uniform retractions, BG-manifolds, and SDEs on manifolds with cylindrical noise) and illustrated through PCA on Grassmann/Stiefel, weight normalization in neural nets, hyperbolic space, and statistical manifolds, highlighting practical implications for manifold-constrained learning and geometry-aware optimization.
Abstract
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.