Provable Benefit of Annealed Langevin Monte Carlo for Non-log-concave Sampling

We study the task of efficient sampling from a probability distribution $\pi\propto{\rm e}^{-V}$ on $\mathbb{R}^{d}$. This fundamental problem is pivotal across various fields including computational statistics [42, 50, 5, 48], Bayesian inference [30], statistical physics [46, 37], and finance [17], and has been extensively studied in the literature [12]. The most common approach to this problem is Markov Chain Monte Carlo (MCMC), among which Langevin Monte Carlo (LMC) [20, 61, 13, 44] is a particularly popular choice. LMC can be understood as a time-discretization of a diffusion process, known as Langevin diffusion (LD), whose stationary distribution is the target distribution $\pi$, and has been attractive partly due to its robust performance despite conceptual simplicity.

Although LMC and its variants converge rapidly when the target distribution $\pi$ is strongly log-concave or satisfies isoperimetric inequalities such as the log-Sobolev inequality (LSI) [20, 61, 13], its effectiveness diminishes when dealing with target distributions that are not strongly log-concave or are multimodal, such as mixtures of Gaussians. In such scenarios, the sampler often becomes confined to a single mode, severely limiting its ability to explore the entire distribution effectively. This results in significant challenges in transitioning between modes, which can dramatically increase the mixing time, making it exponential in dimension $d$. Such limitations highlight the need for enhanced MCMC methodologies that can efficiently navigate the complex landscapes of multimodal distributions, thereby improving convergence rates and overall sampling efficiency.

To address the challenges posed by multimodality, techniques around the notion of annealing have been widely employed [29, 45]. The general philosophy involves constructing a sequence of intermediate distributions $\pi_{0},\pi_{1},...,\pi_{M}$ that bridge the gap between an easily samplable distribution $\pi_{0}$ (e.g., Gaussian or Dirac-like), and the target distribution $\pi_{M}=\pi$. The process starts with sampling from $\pi_{0}$ and progressively samples from each subsequent distribution until $\pi_{M}$ is reached. When $\pi_{i}$ and $\pi_{i+1}$ are close enough, approximate samples from $\pi_{i}$ can serve as a warm start for sampling from $\pi_{i+1}$, thereby facilitating this transition. Employing LMC within this framework gives rise to what is known as the annealed LMC algorithm, which is the focus of our study. Despite its empirical success [54, 55, 69, 68], a thorough theoretical understanding of annealed LMC, particularly its non-asymptotic complexity bounds, remains elusive.

In this work, we take a first step toward providing a non-asymptotic analysis of annealed MCMC. By leveraging the Girsanov theorem to quantify the differences between the sampling dynamics and a reference process, we establish an upper bound of the error of annealed MCMC by the summation of two terms. One term is equal to the ratio of an action integral in the Wasserstein geometry induced by optimal transport and the duration of the process and diminishes when the annealing is sufficiently slow. Another term quantifies the discretization error in implementation. Our approach challenges the traditional point of view that annealed MCMC is a cascade of warm start. Analysis based on this viewpoint still requires log-concavity or isoperimetry. Instead, our theoretical analysis of the annealed Langevin-based MCMC advances the field by eliminating the reliance on assumptions related to log-concavity or isoperimetry. Our strategy represents a paradigm shift in analyzing annealed MCMC. Even though a fully polynomial complexity bound has not been achieved yet. We hope our strategy can inspire future work towards this ambitious goal.

Our key technical contributions are summarized as follows.

1. 1.

   We discover a novel strategy to analyze the non-asymptotic complexity bound of annealed MCMC algorithms that can circumvent the need for log-concavity or isoperimetry.

2. 2.

   In Section 4, we investigate the annealed LD, which involves running LD with a dynamically changing target distribution and derive a surprising bound on the time required to simulate the SDE for achieving $\varepsilon^{2}$-accuracy in KL divergence.

3. 3.

   Building on the insights from the analysis of the continuous dynamic and accounting for discretization errors, we establish a non-asymptotic oracle complexity bound for annealed LMC in Section 5, which is applicable to a broad range of annealing schemes.

The quantitative results are summarized and compared to other sampling algorithms in Table 1. Notably, our approach operates under the least stringent assumptions and exhibits the most favorable $\varepsilon$-dependence among all isoperimetry-free sampling methods.

Related works and comparison. We provide a brief overview of the literature, mainly focusing on the algorithms for non-log-concave sampling and their theoretical analysis.

• •

  Samplers based on tempering. The fundamental concept of tempering involves sampling the system at various temperatures simultaneously: at higher temparatures, the distribution flattens, allowing particles to easily transition between modes, while at lower temperatures, particles can more effectively explore local structures. In simulated tempering [43, 64], the system’s temperature is randomly switched, while in parallel tempering (also known as replica exchange) [57, 38], the temperatures of two particles are swapped according to a specific rule. However, quantitative theoretical results for tempering are limited, and the existing results (e.g., [27, 28, 19]) apply primarily to certain special classes of non-log-concave distributions.

• •

  Samplers based on general diffusions. Inspired by score-based diffusion models [33, 53, 56], recent advances have introduced sampling methods that reverse the Ornstein-Uhlenbeck process, as detailed in [35, 34, 32]. These samplers exhibit reduced sensitivity to isoperimetric conditions, but rely on estimating score functions (gradients of log-density) via importance sampling, which poses significant challenges in high-dimensional settings. Concurrently, studies such as [66, 59, 60, 49] have employed neural networks to approximate unknown drift terms, enabling an SDE to transport a readily sampleable distribution to the target distribution. This approach has shown excellent performance in handling complex distributions, albeit at the expense of significant computational resources required for neural network training. In contrast, annealed LMC runs on a known interpolation of probability distributions, thus simplifying sampling by obviating the need for intensive score estimation or neural network training.

• •

  Non-asymptotic analysis for non-log-concave sampling. Drawing upon the stationary-point analysis in non-convex optimization, the seminal work [4] characterizes the convergence of non-log-concave sampling via Fisher divergence. Subsequently, [11] applies this methodology to examine the local mixing of LMC. However, Fisher divergence is a relatively weak criterion compared to more commonly employed metrics such as total-variational distance or Wasserstein distances. In contrast, our study provides a convergence guarantee in terms of KL divergence, which implies convergence in total-variation distance and offers a stronger result.

Notations and definitions. For $a,b\in\mathbb{R}$, we define $\left[\!\left[a,b\right]\!\right]:=[a,b]\cap\mathbb{Z}$, $a\wedge b:=\min(a,b)$, and $a\vee b:=\max(a,b)$. For $a,b>0$, the notations $a\lesssim b$, $a=O(b)$, and $b=\Omega(a)$ indicate that $a\leq Cb$ for some universal constant $C>0$, and the notations $a\asymp b$ and $a=\Theta(b)$ stand for both $a=O(b)$ and $b=O(a)$. $\widetilde{O}\left(\cdot\right)$ hides logarithmic dependence in $O(\cdot)$. A function $U\in C^{2}(\mathbb{R}^{d})$ is $\alpha(>0)$-strongly-convex if $\nabla^{2}U\succeq\alpha I$, and is $\beta(>0)$-smooth if $-\beta I\preceq\nabla^{2}U\preceq\beta I$. The total-variation (TV) distance is defined as $\operatorname{TV}\left(\mu,\nu\right)=\sup_{A\subset\mathbb{R}^{d}}\left|\mu(A)-\nu(A)\right|$, and the Kullback-Leibler (KL) divergence is defined as $\operatorname{KL}\left(\mu\middle\|\nu\right)=\operatorname{\mathbb{E}}_{\mu}\left[\log\frac{{\rm d}\mu}{{\rm d}\nu}\right]$. $\left\|\cdot\right\|$ represents the $\ell^{2}$ norm on $\mathbb{R}^{d}$. For $f:\mathbb{R}^{d}\to\mathbb{R}^{d^{\prime}}$ and a probability measure $\mu$ on $\mathbb{R}^{d}$, $\left\|f\right\|_{L^{2}(\mu)}:=\left(\int\left\|f\right\|^{2}{\rm d}\mu\right)^{1/2}$, and the second-order moment of $\mu$ is defined as $\operatorname{\mathbb{E}}_{\mu}\left[\left\|\cdot\right\|^{2}\right]$.

In this paper, we consider the following mild assumption on the target distribution $\pi\propto{\rm e}^{-V}$ on $\mathbb{R}^{d}$, which is widely used in the field of sampling (see, e.g., [20, 61, 4, 13]):

Recall that the rationale behind annealing involves a gradual transition from a distribution that is easy to sample from to the more complex target distribution, which is crucial to our algorithm. Throughout this paper, we define the following curve of probability measures on $\mathbb{R}^{d}$:

We use the term annealing schedule to refer to the functions $\eta(\cdot)$ and $\lambda(\cdot)$. They need to be differentiable and monotone, satisfying the boundary conditions $\eta_{0}=\eta(0)\nearrow\eta(1)=1$ and $\lambda_{0}=\lambda(0)\searrow\lambda(1)=0$. The values of $\eta_{0}\in[0,1]$ and $\lambda_{0}\in(1,{+\infty})$ will be determined later. This flexible interpolation scheme includes many general cases. For example, [6] and [26] used the schedule $\eta(\cdot)\equiv 1$, while [45] used the schedule $\lambda(\theta)=c(1-\eta(\theta))$. The key motivation for this interpolation is that when $\theta=0$, the quadratic term predominates, making the potential of $\pi_{0}$ strongly-convex with a moderate condition number, thus $\pi_{0}$ is easy to sample from; on the other hand, when $\theta=1$, $\pi_{1}$ is just the target distribution $\pi$. As $\theta$ gradually increases from $0$ to $1$, the readily sampleable distribution $\pi_{0}$ slowly transform into the target distribution $\pi_{1}$.

In Lemma 5, we prove that $\pi_{\theta}$ also has finite second-order moment, so the W2 distance between $\pi_{\theta}$’s makes sense. We further make a mild assumption on the absolute continuity of the curve:

To facilitate easy sampling from $\pi_{0}$, we consider two choices of the parameters $\eta_{0}$ and $\lambda_{0}$, which have the following complexity guarantees.

See Appendix A for the rejection sampling algorithm as well as the full proof of this lemma. The parameter $R$ reflects prior knowledge about global minimizer(s) of the potential function $V$. Unless it is exceptionally large, indicating scarce prior information about the global minimizer(s) of $V$, this $\widetilde{O}\left(1\right)$ complexity is negligible compared to the overall complexity of sampling. In particular, when the exact location of a global minimizer of $V$ is known, we can adjust the potential $V$ so that $0$ becomes a global minimizer, thereby reducing the complexity to $O(1)$.

Equipped with this foundational setup, we now proceed to introduce the annealed LD and annealed LMC algorithms.

To elucidate the concept of annealing more clearly, we first consider the annealed Langevin diffusion (ALD) algorithm, which samples from the $\pi\propto{\rm e}^{-V}$ by running LD with a dynamically changing target distribution. For the purpose of this discussion, let us assume that we can exactly simulate any SDE with known drift and diffusion terms.

Fix an annealing schedule $\eta(\cdot),\lambda(\cdot)$ and choose a sufficiently long time $T$. We define a reparametrized curve of probability measures $\left(\widetilde{\pi}_{t}:=\pi_{t/T}\right)_{t\in[0,T]}$. Starting with an initial sample $X_{0}\sim\pi_{0}=\widetilde{\pi}_{0}$, we run the SDE

and ultimately output $X_{T}\sim\nu^{\rm ALD}$ as an approximate sample from the target distribution $\pi$. Intuitively, when $\widetilde{\pi}_{t}$ is changing slowly, the distribution of $X_{t}$ should closely resemble $\widetilde{\pi}_{t}$, leading to an output distribution $\nu^{\rm ALD}$ that approximates the target distribution. This turns out to be true, as is confirmed by the following theorem, which provides a convergence guarantee for the ALD process.

Although we adopt a specific parametrization of the interpolating distribution (Equation 3), which is widely used in practice, the results in Theorem 1 are applicable to any curve of probability measures $(\pi_{\theta})_{\theta\in[0,1]}$ that bridge $\pi_{0}$ and $\pi_{1}=\pi$. This applicability extends as long as the closed forms of the drift and diffusion terms in the SDE, which precisely follow the path $(\pi_{\theta})_{\theta\in[0,1]}$, are known.

Let us delve deeper into the mechanics of the ALD. Although at time $t$ the SDE (Equation 4) targets the distribution $\widetilde{\pi}_{t}$, the distribution of $X_{t}$ does not precisely match $\widetilde{\pi}_{t}$. Nevertheless, by choosing a sufficiently long time $T$, we actually move on the curve $(\pi_{\theta})_{\theta\in[0,1]}$ sufficiently slowly, thus minimizing the discrepancy between the path measure of $(X_{t})_{t\in[0,T]}$ and the reference curve $(\widetilde{\pi}_{t})_{t\in[0,T]}$. By applying data-processing inequality, we can upper bound the error between the marginal distributions at time $T$ by the joint distributions of the two paths. In essence, moving more slowly, sampling more precisely.

Our analysis predominantly addresses the global error across the entire curve of probability measures, rather than focusing solely on the local error at time $T$. This approach is inspired by [18] (see also [12, Section 4.4]) and [9], and stands in contrast to the traditional isoperimetry-based analyses of LD (e.g., [61, 13, 4]), which emphasize the decay of the KL divergence from the distribution of $X_{t}$ to the target distribution and require LSI to bound the time derivative of this quantity. Notably, the total time $T$ needed to run the SDE depends solely on the action of the curve $(\pi_{\theta})_{\theta\in[0,1]}$, obviating the need for assumptions related to log-concavity or isoperimetry.

We also remark that the ALD plays a significant role in another important subject known as non-equilibrium stochastic thermodynamics [52]. Recently, a refinement of the fluctuation theorem was discovered in [10, 24], stating that the irreversible entropy production in a stochastic thermodynamic system is equal to the ratio of a similar action integral as in Theorem 1 and the duration of the process $T$, resembling Theorem 1.

It is crucial to recognize that the ALD (Equation 4) cannot be precisely simulated in practice. Transitioning from LD to LMC introduces additional errors due to discretization. This section will present a detailed convergence analysis for the annealed Langevin Monte Carlo (ALMC) algorithm, which is implementable in practical scenarios.

A straightforward yet non-optimal method to discretize Equation 4 involves employing the Euler-Maruyama scheme, i.e.,

However, considering that $\nabla\log\widetilde{\pi}_{t}(x)=-\eta\left(\frac{t}{T}\right)\nabla V(x)-\lambda\left(\frac{t}{T}\right)x$, the integral of the linear term can be computed in closed form, so we can use the exponential-integrator scheme [65, 67] to further reduce the discretization error.

Given the total time $T$, we define a sequence of points $0=\theta_{0}<\theta_{1}<...<\theta_{M}=1$, and set $T_{\ell}=T\theta_{\ell}$, $h_{\ell}=T(\theta_{\ell}-\theta_{\ell-1})$. The exponential-integrator scheme is then expressed as

where $t_{-}:=T_{\ell-1}$ when $t\in[T_{\ell-1},T_{\ell})$, $\ell\in\left[\!\left[1,M\right]\!\right]$. The explicit update rule is detailed in Algorithm 1, with $x_{\ell}$ denoting $X_{T_{\ell}}$, and the derivation of Equation 7 is presented in Section B.1. Notably, replacing $\nabla V(X_{t_{-}})$ with $\nabla V(X_{t})$ recovers the ALD (Equation 4), and setting $\eta\equiv 1$ and $\lambda\equiv 0$ reduces to the classical LMC iterations.

We illustrate the ALMC algorithm in Figure 1. The underlying intuition behind this algorithm is that by setting a sufficiently long total time $T$, the trajectory of the continuous dynamic (i.e., annealed LD) approaches the reference trajectory closely, as established in Theorem 1. Additionally, by adopting sufficiently small step sizes $h_{1},...,h_{M}$, the discretization error can be substantially reduced. Unlike traditional annealed LMC methods, which require multiple LMC update steps for each intermediate distribution $\pi_{1},...,\pi_{M}$, our approach views the annealed LMC as a discretization of a continuous-time process. Consequently, it is adequate to perform a single step of LMC towards each intermediate distribution $\widetilde{\pi}_{T_{1}},...,\widetilde{\pi}_{T_{M}}$, provided that $T$ is sufficiently large and the step sizes are appropriately small.

The subsequent theorem provides a convergence guarantee for the annealed LMC algorithm, with the complete proof detailed in Section B.2.

Once again, our analysis relies on bounding the global error between two path measures by Girsanov theorem. The annealing schedule $\eta(\cdot)$ and $\lambda(\cdot)$ influence the complexity exclusively through the action ${\cal A}_{\eta,\lambda}(\pi)$. Identifying the optimal annealing schedule to minimize this action remains a significant challenge, and we propose this as an area for future research.

We also note that our Assumption 1 encompasses strongly-log-concave target distributions. For sampling from these well-conditioned distributions via LMC, the complexity required to achieve $\varepsilon$- in TV distance scales as $\widetilde{O}\left(\varepsilon^{-2}\right)$ [12]. However, using Pinsker inequality ${\rm KL}\geq 2{\rm TV}^{2}$, our complexity to meet the same error criterion is $O\left(\varepsilon^{-6}\right)$, indicating a significantly higher computational demand. Refining the parameter dependence in our algorithm to reduce this complexity remains a key area for future work.

Finally, we present a conjecture concerning the action ${\cal A}_{\eta,\lambda}(\pi)$, leaving its proof or disproof as an open question for future research. Following this conjecture, we demonstrate its applicability to a specific class of non-log-concave distributions in the subsequent example, with the detailed proof provided in Appendix C.

To demonstrate the superiority of our theoretical results, let us consider a simplified scenario where $N=2$, $y_{1}=-y_{2}$, and $r^{2}\gg\frac{1}{\beta}$. We derive from Example 1 and Pinsker inequality that the total complexity required to obtain a sample that is $\varepsilon$-accurate in TV is $\widetilde{O}\left(d^{3}\beta^{2}r^{4}(r^{4}\beta^{2}\vee d^{2})\varepsilon^{-6}\right)$. In contrast, studies such as [51, 8, 19] indicate that the LSI constant of $\pi$ is $\Omega\left({\rm e}^{\Theta(\beta r^{2})}\right)$, so existing bounds in Table 1 suggest that LMC, to achieve the same accuracy, can only provide an exponential complexity guarantee of $\widetilde{O}\left({\rm e}^{\Theta(\beta r^{2})}d\varepsilon^{-2}\right)$.

In this paper, we have explored the complexity of annealed LMC for sampling from a non-log-concave probability measure without relying on log-concavity or isoperimetry. Central to our analysis are the Girsanov theorem and optimal transport techniques, thus circumventing the need for isoperimetric inequalities. While our proof primarily focuses on the annealing scheme as described in Equation 3, there is potential for adaptation to more general interpolations. Further exploration of these applications will be a focus of our future research. Technically, our proof methodology could be expanded to include a broader range of target distributions beyond those with smooth potentials, such as those with Hölder-continuous gradients [7, 21, 40, 47, 23, 22], or even heavy-tailed target distributions [44, 31]. Furthermore, eliminating the necessity of assuming global minimizers of the potential function would enhance the practical utility of our algorithm. Theoretically, proving Conjecture 1 or obtaining a polynomial bound of ${\cal A}_{\eta,\lambda}(\pi)$ under less restrictive conditions would further validate that our complexity bounds are genuinely polynomial. Finally, while this work emphasizes the upper bounds for non-log-concave sampling, an intriguing future avenue would be to verify the tightness of these bounds and explore potential lower bounds for this problem [16, 15, 14].