NeurIPS2022

Efficient Sampling on Riemannian Manifolds via Langevin MCMC

Xiang Cheng, Jingzhao Zhang, Suvrit Sra

13 citations

Abstract

We study the task of efficiently sampling from a Gibbs distribution dπ * = e -h dvol g over a Riemannian manifold M via (geometric) Langevin MCMC; this algorithm involves computing exponential maps in random Gaussian directions and is efficiently implementable in practice. The key to our analysis of Langevin MCMC is a bound on the discretization error of the geometric Euler-Murayama scheme, assuming ∇h is Lipschitz and M has bounded sectional curvature. Our error bound matches the error of Euclidean Euler-Murayama in terms of its stepsize dependence. Combined with a contraction guarantee for the geometric Langevin Diffusion under Kendall-Cranston coupling, we prove that the Langevin MCMC iterates lie within ε-Wasserstein distance of π * after Õ(ε -2 ) steps, which matches the iteration complexity for Euclidean Langevin MCMC. Our results apply in general settings where h can be nonconvex and M can have negative Ricci curvature. Under additional assumptions that the Riemannian curvature tensor has bounded derivatives, and that π * satisfies a CD(•, ∞) condition, we analyze the stochastic gradient version of Langevin MCMC, and bound its iteration complexity by Õ(ε -2 ) as well. An illustrative example: sampling from a sphere To give some intuition about Assumptions 1 -4, we present a simple example for sampling from S d-1 , the unit d-sphere in R d , which is a positively curved Riemannian manifold. 2 The subsequent bounds for the S d case also generalize with minor modifications to other closely-related manifolds such as SO(d). Let U (x) : R d → Re be a potential function, and suppose that we wish to sample from dp(x) ∝ e -U (x) dvol g (x) defined over S d-1 . It is known that for any x ∈ S d-1 , v ∈ T x M , Ric(v, v) > 0,