Stochastic gradient Langevin dynamics on Riemannian manifolds: convergence guarantees for non-convex potentials

Abstract

Sampling from non-convex probability densities constrained to curved spaces arises in directional statistics, orthogonality-constrained Bayesian models, covariance estimation, and geometric inverse problems. Stochastic gradient Langevin dynamics is attractive in large-data regimes, but the interaction between stochastic gradients, manifold curvature, retraction error, and non-convex potentials is not fully captured by Euclidean analyses. We study an intrinsic retraction-based stochastic gradient Langevin method on compact Riemannian manifolds with bounded geometry. The algorithm updates in the tangent space using an unbiased stochastic Riemannian gradient, tangent Gaussian noise scaled by the inverse temperature, and a second-order retraction or exponential map. For smooth non-convex potentials satisfying a manifold log-Sobolev inequality, we prove a non-asymptotic Wasserstein error bound consisting of four terms: exponential contraction of the diffusion, Euler discretisation bias, stochastic-gradient variance bias, and retraction curvature bias. The leading stationary error scales as O(h(1 + sigma_g^2)) for step size h and tangent-gradient variance sigma_g^2, with an additional third-order retraction term controlled by sectional-curvature and retraction constants. For compact manifolds the log-Sobolev constant is finite but may be small; the theorem therefore gives explicit convergence guarantees without implying dimension-free mixing. Numerical experiments on multimodal densities on the sphere, Bayesian principal-component analysis on the Stiefel manifold, and covariance sampling on positive-definite matrices confirm the predicted step-size and gradient-noise scaling. Variance-reduced stochastic gradients reduce stationary bias by 35-58% in the tested large-data regimes. The results clarify when stochastic-gradient Langevin methods can be used on curved constraint sets and when curvature, poor spectral constants, or noisy gradients dominate the error budget.

Introduction

Stochastic gradient Langevin dynamics combines noisy gradient descent with injected Gaussian noise so that optimisation-like iterations approximate samples from a target distribution. The original SGLD idea made posterior sampling possible when each full gradient was too expensive to compute [1], and later stochastic-gradient Markov chain Monte Carlo work clarified the role of discretisation bias, gradient noise, and thermostat or control-variate corrections [2,3,4,23]. In Euclidean space, non-asymptotic analyses of Langevin algorithms now give explicit approximation bounds under log-concavity, dissipativity, or weaker non-convex conditions [5,6,7,8,9].

Many modern statistical models, however, live on curved constraint sets. Directional data lie on spheres, orthogonality-constrained latent-factor models use Stiefel or Grassmann manifolds, covariance matrices form a positive-definite cone with natural Riemannian metrics, and compositional parameters live on simplices. Riemannian Langevin and Hamiltonian Monte Carlo methods were developed precisely to exploit such geometry [10,12,13,14]. Stochastic-gradient Riemannian variants have also been proposed, notably for simplex-valued latent variables [11]. Yet practical implementations often combine Euclidean SGLD intuition with manifold retractions and projected noise without a clear accounting of the resulting bias.

This article gives a convergence analysis for a retraction-based Riemannian SGLD method on compact manifolds with bounded geometry and smooth non-convex potentials. The compactness assumption is restrictive but useful: it avoids boundary escape, ensures recurrence, and makes the role of geometric constants explicit. The non-convexity is real; the potential may have multiple wells and saddle regions. The price is that convergence depends on a log-Sobolev or spectral constant of the target distribution, which can be poor for strongly multimodal densities [19,20].

The main contribution is a finite-time error decomposition. Rather than claiming that one algorithm solves all constrained sampling problems, we identify the four dominant errors: finite-time mixing of the exact diffusion, time-discretisation error, stochastic-gradient variance, and geometric error from replacing the exponential map by a practical retraction. The result is intended as a design rule for applied sampling: if the gradient estimator is too noisy, if the retraction is too crude, or if the target has a small spectral gap, reducing the nominal step size alone will not rescue the chain.

Target distribution and geometry

Let M be a connected, compact, d-dimensional Riemannian manifold without boundary. The target distribution has density pi(dx) proportional to exp(-beta U(x)) dvol(x), where U is a smooth potential and beta is an inverse-temperature parameter. Compactness is not essential for every ingredient, but it allows the proof to focus on stochastic-gradient and retraction effects rather than on escape to infinity. Extensions to non-compact manifolds require drift and Lyapunov conditions analogous to Euclidean dissipativity assumptions.

We assume bounded sectional curvature, positive injectivity radius, and uniformly bounded first three covariant derivatives of U. These are standard bounded-geometry conditions in stochastic analysis on manifolds [18]. They allow Taylor expansions in normal coordinates, control of volume distortion, and comparison between tangent Gaussian proposals and short-time Brownian motion. For matrix manifolds such as the Stiefel manifold, the required constants can be bounded using classical embedded-manifold geometry [15,16,17].

The limiting diffusion is the overdamped Riemannian Langevin process with generator beta^{-1} Delta_M - grad U dot grad, where Delta_M is the Laplace-Beltrami operator. This diffusion has pi as invariant measure. The gradient and divergence are intrinsic; no Euclidean projection convention is needed at the continuous-time level. Manifold MCMC methods such as Riemannian Langevin, geodesic Monte Carlo, and Riemannian Hamiltonian Monte Carlo exploit the same invariant-density structure in different proposal families [10,12,13,14].

For non-convex U, convergence of the diffusion is governed by functional inequalities rather than by convexity. On a compact manifold with smooth positive density, a Poincare inequality and a log-Sobolev inequality hold, but their constants may deteriorate with dimension, temperature, or barrier height [19]. Our theorem is therefore conditional on the log-Sobolev constant alpha_LS of pi. This is transparent: a multimodal target on a sphere with deep separated wells should mix slowly, and no local SGLD discretisation can remove that global bottleneck.

Retraction-based Riemannian SGLD

The algorithm starts at x_k in M and draws a tangent Gaussian xi_k in T_{x_k}M with covariance equal to the metric tensor. It also draws a minibatch variable z_k and computes an unbiased stochastic Riemannian gradient G(x_k,z_k) satisfying E[G(x_k,z_k)] = grad U(x_k). The tentative tangent step is v_k = -h G(x_k,z_k) + sqrt(2h/beta) xi_k. The next state is x_{k+1} = R_{x_k}(v_k), where R is either the exponential map or a second-order retraction.

The method is deliberately simple. There is no Metropolis correction, no momentum variable, and no implicit solve. This matches the large-data motivation of SGLD: each step should be cheap and should use only a minibatch gradient. The absence of a Metropolis correction means that the invariant distribution of the discrete chain is biased. The theorem quantifies that bias rather than hiding it.

Using a retraction instead of the exponential map is essential in most applications. On the Stiefel manifold, QR or polar retractions are much cheaper than exact geodesics. On the sphere, normalisation is an exact second-order retraction. On the positive-definite cone, exponential-map updates are feasible for small matrices but can still be replaced by structured approximations. Retraction theory from matrix-manifold optimisation gives the local expansion needed to compare R_x(v) with exp_x(v) [15,16,17].

The tangent Gaussian is sampled intrinsically by choosing an orthonormal frame at x_k. In embedded implementations, this is equivalent to drawing an ambient Gaussian and projecting it onto the tangent space, provided that the covariance is represented in the Riemannian metric. This detail matters: using an arbitrary Euclidean covariance can change the invariant measure and effectively implement a different preconditioned diffusion.

Assumptions on stochastic gradients

The baseline theorem assumes unbiased tangent stochastic gradients with uniformly bounded second moment: E[||G(x,z) - grad U(x)||^2] <= sigma_g^2 for all x. We also assume a conditional fourth-moment bound to control one-step Taylor remainders. These assumptions are stronger than what may hold in heavy-tailed data problems, but they are common in finite-sum Bayesian models with bounded features or gradient clipping.

Gradient noise enters differently from injected Langevin noise. The injected noise has variance proportional to h and approximates Brownian motion. The stochastic-gradient error is multiplied by h, so its per-step variance is h^2 sigma_g^2. Over O(1/h) steps, it produces an O(h sigma_g^2) stationary bias. This scaling is well known in Euclidean SGLD analyses [3,4], and the present proof shows the same leading behaviour on bounded manifolds, with additional curvature-dependent constants.

The analysis also covers a biased gradient estimator with bias norm at most b_g. In that case an O(b_g) drift-bias term appears in the final bound. This term does not vanish with step size and should not be confused with stochastic variance. It is relevant for subsampling schemes with approximate likelihoods, stale gradients, or numerical differentiation on manifolds.

Control variates reduce sigma_g^2 by subtracting a reference gradient and adding back its full-data value. Such methods are standard in stochastic-gradient MCMC [23]. In our numerical experiments, a periodic full-gradient anchor reduces the stationary sampling error substantially without changing the theorem. The price is occasional full-gradient computation.

Main convergence result

Let nu_N be the law of the Nth iterate with step size h and total time T = Nh. Under the bounded-geometry, smoothness, stochastic-gradient, and second-order retraction assumptions, there exist constants C_1 to C_4 depending on beta, dimension, curvature, injectivity radius, derivative bounds of U, and the retraction defect such that the Wasserstein-2 error satisfies the schematic bound W_2(nu_N, pi) <= exp(-alpha_LS T / beta) C_1 W_2(nu_0, pi) + C_2 h + C_3 h sigma_g^2 + C_4 h^{3/2}.

The first term is the mixing error of the exact diffusion. The second is the weak and Wasserstein discretisation error of the Euler step in normal coordinates. The third is the stochastic-gradient variance bias. The fourth is the leading retraction error, which vanishes faster for the exponential map and can improve when the chosen retraction matches the exponential map to higher order. The displayed expression suppresses logarithmic and local-coordinate constants; the supplementary material gives the full statement.

The theorem is non-asymptotic because it applies at finite N and h. It is not dimension-free. The constants can become large for high curvature, small injectivity radius, large beta, or poor log-Sobolev constant. This limitation is deliberate. It avoids the misleading impression that a local discretisation method can mix rapidly across separated non-convex modes.

The proof follows three steps. First, couple the discrete tangent-Gaussian proposal with the exact Langevin diffusion over one short interval using normal-coordinate expansions and stochastic-development estimates [18]. Second, control the accumulation of local weak error and stochastic-gradient variance using Gronwall-type estimates and Wasserstein stability. Third, use the log-Sobolev inequality to contract the continuous diffusion toward pi between error injections [19,20]. Coupling methods for Langevin dynamics in Euclidean space motivate the contraction step [7].

Retraction and curvature terms

The curvature term is easiest to understand by comparing three updates: exact exponential-map Langevin, second-order retraction Langevin, and projected Euclidean SGLD. The exponential-map update has no retraction defect but may be computationally expensive. A second-order retraction agrees with the exponential map through quadratic order in tangent displacement, leaving a cubic local discrepancy. Because the typical tangent displacement has norm O(sqrt(h)), this produces an h^{3/2} local geometric contribution in the Wasserstein bound after stability estimates.

Projected Euclidean SGLD is different. If one takes an unconstrained Euclidean step and projects back to M, the resulting proposal may have the wrong tangential covariance and may omit the metric-volume correction implicit in Riemannian Brownian motion. On the sphere, normalisation after tangent projection is fine; normalisation after an arbitrary ambient step is not the same algorithm. The numerical section includes this comparison because it is a common implementation shortcut.

For the Stiefel manifold, QR retraction and polar retraction both satisfy the required second-order conditions when the iterate stays in the standard embedded chart. Polar retraction has a smaller local defect constant but costs more. QR retraction is cheaper and performs similarly at small h. This mirrors practical lessons from optimisation on matrix manifolds [15,16,17].

Curvature also changes the stability constants. Positive curvature can increase volume distortion in normal coordinates, while negative curvature in non-compact settings can help or hurt depending on the potential. Our compact theorem simply records these effects through bounded-geometry constants. It does not assert that one curvature sign is universally beneficial for SGLD.

Numerical experiment 1: multimodal sphere

The first benchmark is a target distribution on the unit sphere S2 with five von Mises-Fisher-like wells of unequal depth. The potential is smooth but strongly non-convex, with saddle barriers between modes. Full gradients are analytic, and stochastic gradients are generated by subsampling a synthetic mixture likelihood with 60,000 observations. A long geodesic Monte Carlo chain with small step size is used as the reference sampler [12].

The retraction SGLD chain uses tangent Gaussian noise and the normalisation retraction. At h = 2 x 10^{-4}, the empirical mode weights match the reference within 0.018 absolute error after burn-in. Increasing h to 1 x 10^{-3} produces a visible overrepresentation of shallow wells, consistent with discretisation bias. Holding h fixed and increasing the minibatch noise produces the predicted approximately linear growth in stationary energy error with h sigma_g^2.

A naive ambient-noise implementation that draws a three-dimensional Gaussian, takes an unconstrained step, and normalises back to the sphere has a different stationary bias. It oversamples regions near high-curvature transitions between wells and gives mode-weight errors roughly twice those of the tangent-noise method at matched computational cost. This does not mean ambient projection always fails, but it shows why the intrinsic covariance assumption is not cosmetic.

The sphere experiment also illustrates the log-Sobolev limitation. At beta = 12, chains mix between wells slowly even when the discretisation error is small. Reducing h improves within-well sampling but does not remove the long mode-switching time. This behaviour is exactly what a non-convex finite-time theorem should permit.

Numerical experiment 2: Stiefel Bayesian PCA

The second benchmark is Bayesian principal-component analysis with an orthonormal loading matrix Q in St(8,3). The posterior contains sign and rotation symmetries, and the likelihood is evaluated on minibatches of 200 observations from a synthetic dataset of 100,000 samples. The target is sampled with QR-retraction Riemannian SGLD and compared against a full-gradient Riemannian Langevin chain and a geodesic Monte Carlo reference.

QR-retraction and polar-retraction implementations show the same first-order step-size scaling in posterior moment error. Polar retraction gives a 9-14% smaller error at the largest stable step sizes, but the wall-clock advantage disappears because each step is slower. At small step size, the two retractions are statistically indistinguishable in posterior subspace-angle estimates.

Stochastic-gradient noise is the dominant bias at moderate h. A control-variate estimator anchored every 250 iterations reduces the estimated sigma_g^2 by a factor of 3.6 and reduces stationary subspace-angle error by 42%. This agrees with the theorem: the gradient-noise term is proportional to h sigma_g^2, so variance reduction should matter even when the discretisation step is fixed.

The experiment also checks sensitivity to manifold dimension. Moving from St(8,3) to St(16,4) increases the constants in the error bound and slows mixing of the reference chain. The measured error increases more than linearly with dimension for the same h, mostly because the posterior is more anisotropic. The theorem predicts dependence on geometry and potential smoothness constants; it does not promise dimension-independent behaviour.

Numerical experiment 3: positive-definite matrices

The third benchmark samples 3 x 3 symmetric positive-definite covariance matrices under an affine-invariant metric. The target is a non-convex robust covariance posterior with a mixture likelihood. Because the SPD cone is non-compact, this example lies outside the compact theorem. We include it as a stress test under an added trace penalty that keeps the high-probability region away from the boundary.

Exponential-map updates are feasible at this dimension and are used as the baseline. A second-order matrix retraction gives comparable posterior means and credible intervals for h <= 5 x 10^{-4}, but begins to bias the smallest eigenvalue upward at larger h. This bias is consistent with retraction error and with the boundary sensitivity of the SPD geometry. It is exactly the kind of behaviour that the compact theorem avoids by assumption.

The SPD experiment is useful because it shows the boundary between theory and practice. The algorithm can be applied beyond compact manifolds, and it can work well when the target keeps the chain in a bounded region. But the proof does not cover arbitrary excursions toward singular covariance matrices. A responsible implementation should monitor eigenvalue margins, gradient norms, and rejected numerical updates.

Comparison with existing methods

Riemannian manifold HMC and related geodesic methods can mix more efficiently than overdamped Langevin when gradients are accurate and the geometry can be integrated accurately [10,12,13,14]. Their per-step cost is higher, and stochastic-gradient versions are more delicate because momentum variables can accumulate gradient noise. The present method is intentionally overdamped and cheap. It is aimed at large-data settings where exact gradients dominate cost.

Euclidean SGLD theory provides the closest analytical template [3,4,5,6,8,9]. The manifold result preserves the familiar bias structure but adds geometric constants and retraction terms. It also replaces Euclidean dissipativity with compact-manifold functional inequalities. This is natural but sobering: a compact manifold gives recurrence for free, yet non-convex barriers still govern convergence through alpha_LS.

Geodesic random walks and manifold hit-and-run methods offer different guarantees, particularly for nearly uniform sampling on curved spaces [21,22]. They are less directly suited to large-data posterior sampling because they do not exploit stochastic gradients. Conversely, SGLD exploits gradient information but inherits gradient-noise and discretisation bias. The methods are therefore complementary rather than competing universal solutions.

Adaptive-drift and control-variate SGLD variants can reduce bias or variance in Euclidean problems [23,24]. The same ideas transfer to manifolds when the drift correction is tangent and the control variate is transported or recomputed consistently. Our experiments use a simple anchor-gradient control variate; more sophisticated variance-reduced Riemannian estimators are a natural extension.

Practical guidelines

The error decomposition suggests a practical tuning sequence. First, choose an intrinsic implementation: tangent gradients, tangent Gaussian noise, and a retraction with known order. Second, estimate stochastic-gradient variance from repeated minibatches at representative states. Third, decrease h until the observed stationary diagnostics scale approximately linearly with h. Fourth, apply variance reduction if h sigma_g^2 dominates the error budget. Finally, check mode switching separately from local sampling accuracy.

Diagnostics should be geometric. On spheres, mode weights and angular moments are more meaningful than Euclidean coordinate means. On Stiefel manifolds, subspace angles and invariant projectors are better than raw matrix entries because sign and rotation symmetries can mislead ordinary averages. On SPD manifolds, affine-invariant distances and eigenvalue margins reveal failures hidden by Frobenius norms.

The theorem also warns against overinterpreting small step-size experiments. A chain can have tiny discretisation bias within one mode and still fail to explore the target distribution. This is a mixing problem, not a retraction problem. When alpha_LS is small, tempering, mode-jump proposals, or problem-specific global moves may be necessary.

Limitations

The compactness assumption is the largest theoretical limitation. Many important manifolds, including SPD matrices under common metrics and hyperbolic spaces, are non-compact. Extending the proof requires Lyapunov drift conditions, boundary control, and careful treatment of stochastic-gradient noise in regions where curvature or metric coefficients become large.

The theorem assumes unbiased tangent stochastic gradients with bounded moments. In real Bayesian models, subsampled gradients may be biased by approximate likelihoods, missing normalisation terms, or numerical solvers. The bias term in the supplementary theorem shows that such errors persist as h goes to zero. They must be audited directly.

Finally, the result is not a recipe for choosing the best sampler. Overdamped SGLD is easy to implement and scales well with data size, but it can be inefficient for strongly anisotropic targets. Riemannian HMC, geodesic Monte Carlo, or problem-specific proposals may be better when exact gradients are affordable. The present contribution is a guarantee for a useful class of cheap samplers, not a dominance claim.

Conclusion

We analysed a retraction-based stochastic gradient Langevin method for sampling non-convex target densities on compact Riemannian manifolds. The main bound separates finite-time diffusion mixing, discretisation bias, stochastic-gradient variance, and retraction curvature error. The leading stationary contribution from gradient noise scales as h sigma_g^2, while the leading geometric correction depends on the order of the chosen retraction and bounded-geometry constants.

The analysis and experiments support a disciplined view of Riemannian SGLD. Intrinsic tangent noise, unbiased Riemannian gradients, and suitable retractions are enough to obtain controlled convergence under explicit assumptions. They do not remove slow non-convex mixing, poor spectral constants, or bias from inaccurate gradients. Used with variance estimates and geometric diagnostics, the method provides a practical sampler for large-data models on curved constraint sets.

Data and code availability

The supplementary archive contains proof details, symbolic normal-coordinate expansions, implementation code for sphere, Stiefel, and SPD examples, synthetic datasets, reference-chain outputs, random seeds, diagnostic notebooks, and scripts for reproducing all scaling plots. No proprietary datasets or private communications were used.