Stochastic Modified Flows for Riemannian Stochastic Gradient Descent
This work provides theoretical insights into the convergence behavior of RSGD on manifolds, which is incremental for optimization in machine learning and related fields.
The paper tackles the problem of approximating Riemannian stochastic gradient descent (RSGD) by deriving a diffusion process called Riemannian stochastic modified flow (RSMF) to account for random fluctuations, showing that RSMF provides a higher-order approximation than deterministic gradient flow in the small learning rate regime. It proves quantitative bounds for the weak error of this approximation under assumptions on retraction maps, manifold geometry, and gradient estimators.
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.