Sampling from multi-modal distributions on Riemannian manifolds with training-free stochastic interpolants

TL;DR

A sampling algorithm based on the simulation of a non-equilibrium deterministic dynamics that transports an easy-to-sample noise distribution toward the target, thereby extending recent diffusion-based sampling methodologies beyond the Euclidean setting.

Abstract

In this paper, we propose a general methodology for sampling from un-normalized densities defined on Riemannian manifolds, with a particular focus on multi-modal targets that remain challenging for existing sampling methods. Inspired by the framework of diffusion models developed for generative modeling, we introduce a sampling algorithm based on the simulation of a non-equilibrium deterministic dynamics that transports an easy-to-sample noise distribution toward the target. At the marginal level, the induced density path follows a prescribed stochastic interpolant between the noise and target distributions, specifically constructed to respect the underlying Riemannian geometry. In contrast to related generative modeling approaches that rely on machine learning, our method is entirely training-free. It instead builds on iterative posterior sampling procedures using only standard Monte Carlo techniques, thereby extending recent diffusion-based sampling methodologies beyond the Euclidean setting. We complement our approach with a rigorous theoretical analysis and demonstrate its effectiveness on a range of multi-modal sampling problems, including high-dimensional and heavy-tailed examples.

Figures (7)

• Figure 1: Illustration of the invertibility property of the map $x_t\mapsto\psi_t^{-1}(x_t|x_1)$ in the case $\mathsf{M}=\mathsf{S}^1$. Given $x_1\in\mathsf{M}$ and $t\in[0,1)$, the image set $\mathsf{O}_t(x_1)=\{x_t\in\mathsf{S}^1\,:\, \widehat{x_tx_1}<(1-t)\pi\}$ is shown in red, while the geodesic mapping $x_0$ to $x_1$ and passing through $x_t$ at time $t$ is depicted in blue. (Left) If $t=1/3$ and $x_{1/3}$ lies in the image, then $x_0=\psi_{1/3}^{-1}(x_{1/3};x_1)$ satisfies $x_{1/3}=\psi_{1/3}(x_0;x_1)$. The geodesic passing through $x_{1/3}$ at time $1/3$ is therefore minimizing. (Right) If $t=2/3$ and $x_{2/3}$ lies outside the image, the geodesic passing through $x_{2/3}$ at time $2/3$ is not minimizing; consequently, $x_0=\psi_{2/3}^{-1}(x_{2/3};x_1)$ satisfies $x_{2/3}\neq\psi_{2/3}(x_0;x_1)$.
• Figure 2: Illustration of the spherical coordinates with a base point $z$ on the manifold $\mathsf{M}=\mathsf{S}^1$. Apart from $z$ and $-z$, every point $x\in\mathsf{M}$ is associated with a unique pair $(t,\xi)$ such that $x=\mathrm{Exp}_z t\xi$, with $\xi$ (colored in red) being a unit vector of $T_z\mathsf{M}$ (here, there are only two such vectors), and $t=\operatorname{dist}(z,x)=\widehat{zx}$ (colored in green).
• Figure 3: Illustration of $x_t\mapsto p_{t|1}(x_t|x_1)\propto\operatorname{Jac}_t(x_t|x_1)\mathbbm{1}_{\mathsf{O}_t(x_1)}(x_t)$ for different values of $t$, in the case where $\mathsf{M}=\mathsf{S}^2$, with $x_1\in \mathsf{M}$ being displayed in red. (From left to right): $t=0.3,0.5,0.7,0.9$. While $p_{t|1}(\cdot|x_1)$ is close to the uniform distribution for small times, we observe that it concentrates more and more around $x_1$ as $t$ increases, since its support gradually shrinks.
• Figure 4: Performances of FRIPS for the bi-modal targets defined in \ref{['tab:weight']}, as a function of $t_0$. Blue: Average relative error on strongest mode weight estimation. Red: Wasserstein distance with respect to ground truth samples. For each setting, we observe a "sweet spot", i.e. a range of $t_0$ where both metrics reach a minimum. The vertical dashed line highlights the value $t_0$ leading to the best mode weight estimation, corresponding to the values reported in \ref{['tab:weight']}.
• Figure 5: Return accuracy $\tau_1$ for the bi-modal target \ref{['eq:mog_target']}, with instances for each dimension defined in \ref{['tab:weight']}. For each setting, the smaller plot corresponds to the $t_0$ range of \ref{['fig:weight']}, where the best value of $t_0$ (vertical dashed line) was found.

Theorems & Definitions (15)

• Example 1: The $d$-sphere example.
• Proposition 1
• proof
• Lemma 1
• proof
• Proposition 2
• proof
• Proposition 3
• proof
• Proposition 4