Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Riemannian Proximal Sampler for High-accuracy Sampling on Manifolds

Yunrui Guan, Krishnakumar Balasubramanian, Shiqian Ma

TL;DR

This work introduces the Riemannian Proximal Sampler for densities on manifolds, built around two oracles: Manifold Brownian Increments and the Riemannian Heat Kernel. Under a Log-Sobolev inequality, the algorithm achieves high-accuracy sampling with iteration complexity of $\tilde{O}(\log(1/\varepsilon))$ in KL divergence for exact oracles, and $\tilde{O}(\log^2(1/\varepsilon))$ in TV distance for inexact oracles, with curvature-dependent contraction rates. The paper provides practical oracle implementations via heat-kernel truncation and Varadhan's asymptotics, and connects the RHK step to entropy-regularized JKO schemes on Wasserstein space, offering a diffusion-transport perspective. Empirical illustrations on manifolds such as spheres and SPD matrices demonstrate the approach's effectiveness and robustness. The framework opens avenues for high-precision sampling on curved spaces and invites further exploration of oracle design and broader manifold classes.

Abstract

We introduce the Riemannian Proximal Sampler, a method for sampling from densities defined on Riemannian manifolds. The performance of this sampler critically depends on two key oracles: the Manifold Brownian Increments (MBI) oracle and the Riemannian Heat-kernel (RHK) oracle. We establish high-accuracy sampling guarantees for the Riemannian Proximal Sampler, showing that generating samples with $\varepsilon$-accuracy requires $O(\log(1/\varepsilon))$ iterations in Kullback-Leibler divergence assuming access to exact oracles and $O(\log^2(1/\varepsilon))$ iterations in the total variation metric assuming access to sufficiently accurate inexact oracles. Furthermore, we present practical implementations of these oracles by leveraging heat-kernel truncation and Varadhan's asymptotics. In the latter case, we interpret the Riemannian Proximal Sampler as a discretization of the entropy-regularized Riemannian Proximal Point Method on the associated Wasserstein space. We provide preliminary numerical results that illustrate the effectiveness of the proposed methodology.

Riemannian Proximal Sampler for High-accuracy Sampling on Manifolds

TL;DR

This work introduces the Riemannian Proximal Sampler for densities on manifolds, built around two oracles: Manifold Brownian Increments and the Riemannian Heat Kernel. Under a Log-Sobolev inequality, the algorithm achieves high-accuracy sampling with iteration complexity of in KL divergence for exact oracles, and in TV distance for inexact oracles, with curvature-dependent contraction rates. The paper provides practical oracle implementations via heat-kernel truncation and Varadhan's asymptotics, and connects the RHK step to entropy-regularized JKO schemes on Wasserstein space, offering a diffusion-transport perspective. Empirical illustrations on manifolds such as spheres and SPD matrices demonstrate the approach's effectiveness and robustness. The framework opens avenues for high-precision sampling on curved spaces and invites further exploration of oracle design and broader manifold classes.

Abstract

We introduce the Riemannian Proximal Sampler, a method for sampling from densities defined on Riemannian manifolds. The performance of this sampler critically depends on two key oracles: the Manifold Brownian Increments (MBI) oracle and the Riemannian Heat-kernel (RHK) oracle. We establish high-accuracy sampling guarantees for the Riemannian Proximal Sampler, showing that generating samples with -accuracy requires iterations in Kullback-Leibler divergence assuming access to exact oracles and iterations in the total variation metric assuming access to sufficiently accurate inexact oracles. Furthermore, we present practical implementations of these oracles by leveraging heat-kernel truncation and Varadhan's asymptotics. In the latter case, we interpret the Riemannian Proximal Sampler as a discretization of the entropy-regularized Riemannian Proximal Point Method on the associated Wasserstein space. We provide preliminary numerical results that illustrate the effectiveness of the proposed methodology.

Paper Structure

This paper contains 45 sections, 31 theorems, 126 equations, 1 figure, 1 table, 5 algorithms.

Table of Contents

  1. Introduction
  2. Preliminaries
  3. Curvature
  4. Brownian Motion on Manifolds
  5. The Riemannian Proximal Sampler
  6. Step 1:
  7. Step 2:
  8. High-Accuracy Convergence Rates
  9. Rates with Exact Oracles
  10. Rates with Inexact Oracles
  11. Implementation of Inexact Oracles via Heat Kernel Trucation
  12. Implementation of RHK
  13. Verification of Assumption \ref{['Assumption_Oracle_TV_quality']}
  14. Implementation via Varadhan's Asymptotics and Connection to Entropy-Regularized JKO Scheme
  15. RHK as a proximal operator on Wasserstein space
  16. ...and 30 more sections

Key Result

Theorem 6

Let $M$ be a Riemannian manifold without boundary, i.e., $\partial M = \emptyset$. Assume $\pi^{X}$ satisfies $\alpha$-$\mathsf{LSI}$. Denote the distribution for the $k$-th iteration of Algorithm Manifold_Proximal_Sampler_Ideal as $x_{k} \sim \rho_{k}^{X}$. For any initial distribution $\rho_{0}^{X where $\kappa$ is the lower bound of Ricci curvature. In case of negative curvature, we have

Figures (1)

  • Figure 1: Frechét variance (i.e., E versus number of iterations. Left and Middle figure correspond to the implementation via Algorithm \ref{['Inexact_Rejection_Sampling']} and \ref{['Inexact_BM']}. Right figure corresponds to implementation via Algorithm \ref{['Manifold_Proximal_Sampler_Gaussian']}.

Theorems & Definitions (36)

  • Definition 1: TV distance
  • Definition 2: KL divergence
  • Definition 3: Log-Sobolev Inequality (LSI)
  • Example 5
  • Theorem 6
  • Lemma 7
  • Theorem 8
  • Proposition 9
  • Lemma 10
  • Theorem 11
  • ...and 26 more