[Paper Review] Convergence Analysis of Riemannian Stochastic Approximation Schemes

Alain Durmus, Pablo Jiménez+3 more|arXiv (Cornell University)|May 27, 2020

Stochastic Gradient Optimization Techniques31 references4 citations

TL;DR

This paper presents a global convergence analysis for Riemannian stochastic approximation schemes on manifolds, allowing for biased vector fields, controlled Markovian sampling, and retraction maps in place of exponential maps. It establishes convergence rates of 𝒪(b∞ + log n / √n) for the expected squared norm of the mean field, even without prior boundedness or geodesic convexity assumptions.

ABSTRACT

This paper analyzes the convergence for a large class of Riemannian stochastic approximation (SA) schemes, which aim at tackling stochastic optimization problems. In particular, the recursions we study use either the exponential map of the considered manifold (geodesic schemes) or more general retraction functions (retraction schemes) used as a proxy for the exponential map. Such approximations are of great interest since they are low complexity alternatives to geodesic schemes. Under the assumption that the mean field of the SA is correlated with the gradient of a smooth Lyapunov function (possibly non-convex), we show that the above Riemannian SA schemes find an ${\mathcal{O}}(b_\infty + \log n / \sqrt{n})$-stationary point (in expectation) within ${\mathcal{O}}(n)$ iterations, where $b_\infty \geq 0$ is the asymptotic bias. Compared to previous works, the conditions we derive are considerably milder. First, all our analysis are global as we do not assume iterates to be a-priori bounded. Second, we study biased SA schemes. To be more specific, we consider the case where the mean-field function can only be estimated up to a small bias, and/or the case in which the samples are drawn from a controlled Markov chain. Third, the conditions on retractions required to ensure convergence of the related SA schemes are weak and hold for well-known examples. We illustrate our results on three machine learning problems.

Motivation & Objective

• To analyze the convergence of Riemannian stochastic approximation schemes on general Riemannian manifolds under relaxed assumptions.
• To extend convergence guarantees beyond the standard exponential map to retraction-based schemes, which are computationally more feasible.
• To handle biased mean-field estimators and Markovian sampling in the stochastic approximation framework.
• To remove the need for a priori boundedness of iterates and geodesic convexity, enabling global convergence results.
• To provide non-asymptotic convergence rates for both geodesic and retraction-based schemes under mild regularity conditions.

Proposed method

• Formulates the root-finding problem on a Riemannian manifold as finding θ such that h(θ) = 0, where h is the mean vector field defined via integration over a state space.
• Proposes a Riemannian stochastic approximation scheme using either the exponential map or retraction maps to update iterates: θ_{n+1} = Retr_{θ_n}(η_{n+1} H_{θ_n}(X_{n+1})).
• Analyzes convergence under weak assumptions: no a priori boundedness of iterates, no geodesic convexity, and allowance for biased or Markovian sampling.
• Applies the ODE method adapted to Riemannian geometry to study the limiting behavior of the stochastic process.
• Uses tools from Riemannian geometry, including parallel transport, Hessian bounds, and curvature control, to derive stability and convergence properties.
• Establishes convergence rates via martingale techniques and moment bounds, leveraging the structure of the manifold and the properties of retractions.

Experimental results

Research questions

• RQ1 Can Riemannian stochastic approximation converge globally without assuming bounded iterates or geodesic convexity?
• RQ2 How do biases in the vector field estimator affect the convergence rate of Riemannian SA schemes?
• RQ3 What convergence guarantees hold when the sampling distribution depends on the current iterate via a controlled Markov chain?
• RQ4 To what extent can retraction maps replace the exponential map in Riemannian SA while preserving convergence?
• RQ5 What non-asymptotic convergence rates can be derived for Riemannian SA under weak assumptions on the mean field and sampling mechanism?

Key findings

• The proposed Riemannian stochastic approximation scheme achieves a convergence rate of 𝒪(b∞ + log n / √n) for the expected squared norm of the mean field, where b∞ is the asymptotic bias.
• Convergence is established globally without requiring the iterates to be a priori bounded or the manifold to be compact.
• The analysis holds under weak assumptions on retractions, which are satisfied in many practical settings, including common retractions like the QR-based or polar retraction.
• The convergence result extends to cases where the samples are generated from a controlled Markov chain, broadening applicability to online learning and reinforcement learning settings.
• The framework applies to non-convex problems such as principal component analysis and Riemannian barycenter computation, with theoretical guarantees on convergence to critical points.
• The paper provides a theoretical foundation for using computationally cheaper retractions instead of exponential maps in Riemannian optimization, without sacrificing convergence guarantees.