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

Authors: Wei Guo, Molei Tao, Yongxin Chen
Primary area: probabilistic methods (Bayesian methods, variational inference, sampling, UQ, etc.)
Keywords: MCMC, Annealed Langevin Monte Carlo, Non-log-concave sampling, Non-asymptotic analysis

Abstract:
We consider the outstanding problem of sampling from an unnormalized density that may be non-log-concave and multimodal. To enhance the performance of simple Markov chain Monte Carlo (MCMC) methods, techniques of annealing type have been widely used. However, quantitative theoretical guarantees of these techniques are under-explored. This study takes a first step toward providing a non-asymptotic analysis of annealed MCMC. Specifically, we establish, for the first time, an oracle complexity of $\widetilde{O}\left(\frac{d\beta^2{\cal A}^2}{\varepsilon^6}\right)$ for the simple annealed Langevin Monte Carlo algorithm to achieve $\varepsilon^2$ accuracy in Kullback-Leibler divergence to the target distribution $\pi\propto{\rm e}^{-V}$ on $\mathbb{R}^d$ with $\beta$-smooth potential $V$. Here, ${\cal A}$ represents the action of a curve of probability measures interpolating the target distribution $\pi$ and a readily sampleable distribution.

Decision: accept (poster)
Reviewer scores: 6, 8, 6, 8

Reviews and discussion:
<thread>
creation date:
1730188369746

thread_type:
official_review

<thread_object>
creator: reviewer_wbhj

<content>
summary:
The paper consider the problem of sampling from an unnormalized density that may be non-log-concave and multimodal using annealed MCMC.  The authors present a non-asymptotic analysis of Annealed Langevin Monte Carlo, bypassing the commonly assumed log-concavity constraint in literature. The quantitative bounds derived are widely applicable and either comparable to or better than previous analyses for isoperimetry-free sampling methods.

strengths:
The paper provides, as far as I know, the first thorough and non-asymptotic analysis of the annealed Langevin Monte Carlo algorithm.  The results obtained advance the understanding and analysis of annealed MCMC, potentially setting the stage for further research in this area. The intuition behind the algorithm and the proofs are well explained.

weaknesses:
The discussion on the impact of the annealing schedule on the convergence rate is not sufficiently clear in the paper. How the choice of the annealing schedule influences the derived bounds? Specifically, in the example of mixture of Gaussians,  $\lambda(\theta) \propto (1-\theta)^\gamma$ is used,   how varying $\gamma$ affects the action $\mathcal{A}$? whether there is an optimal value of $\gamma$ that could minimize the convergence rate?

questions:
Please see the weakness part.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730401377460

thread_type:
official_review

<thread_object>
creator: reviewer_wqfj

<content>
summary:
The paper proposes a novel analysis of Annealed LMC which uses a specific scheduling system. The latter allows to benefit from another analysis scheme which uses the action on curves instead of isoperimetric inequalities.

strengths:
The paper provides a concise analysis with a clever mathematical trick of using the action of the curve instead of the isoperimetry constants. Assuming gradient Lipschitzness of the potential function allows to easily (rejection-acceptance) sample from an initial distribution, which then allows to sample from the target distribution (non-convex, gradient-Lipschitz) using an annealing schedule for LMC.

weaknesses:
- Assumption 1 is not well discussed. One cannot verify beforehand whether the curve $\pi_{\theta}$ is AC. How should it be done? See the respective question in the section below. 
- Assumption 2 is not the classic smoothness assumption. Instead it is the gradient Lipschitz continuity, which is stronger. This should be explicitly mentioned in the paper. 
- The notation is slightly confusing. There are too many indices for $\pi$ and the usage is incoherent. The reading would be rather simplified, if the notation was more developed.  
- Simple experiments on sampling from a mixture of Gaussians using the proposed and existing methods would significantly boost the paper. 
- In the abstract of the paper, the authors claim that their analysis is non-asymptotic, however it is not true. Theorem 2 (main result) and related results are given with $\tilde{O}(\cdot)$. 
### Mathematical comments 

- *The equation after line 1077* in the proof of Example 2 is incorrect. The density of the sum of two independent random variables(vectors) is the convolution of their densities rather than their product. Thus, the variable $X$ is not distributed as $\hat{\pi}_{\lambda}$. This serves as the cornerstone of the remainder of the proof, which means that the proof is incorrect. 
- *The third line in the proof of Example 2* is incorrect. After recentering the Gaussians, the authors erase the terms $\exp(-\frac{2\lambda \beta}{\lambda + \beta})\|y_i\|^2$ with the $\propto$ sign. Since we have different $y_i$'s, the resulting mixture weights will be different from $p_i$. Nevertheless, this could be fixed by replacing $p_i$ by some $q_i$'s. 
- *Paragraph starting at line 364* The authors motivate the use of annealing as the schedules start at a distribution $\pi_0$, which is easy to sample from. The reason for the easy sampling according to the authors is the strong-convexity of the negative log-density $-\log\pi_0(x) = \eta(0)V(x) + \lambda_0\|x\|^2/2$. The latter is true when $\lambda_0$ is large enough.
- *Continuation of the previous point. Line 831* The strong convexity constant $(\lambda_0 - \eta_0 \beta)$ might be negative, unless $\lambda_0$ is chosen to be large. 
- *Beginning of the proof of Theorem 1* $v_t$ is chosen to satisfy the Fokker-Planck equation on line 294. However, later on line 301 it is chosen to minimize the norm $\|v_t\|_{L_2(\tilde{\pi}_t)}$  using Lemma 2. This is possible if $v_t$ also generates $\tilde{\pi}_t$. Is it possible to find such $v_t$? See the relevant question.

#### typos
- *line 956:* cure -> curve

questions:
- Is the analysis using the action only available for annealing? Can we gain anything by applying this technique to LMC in the standard setting?
- What are the estimates of the action parameter $\mathcal{A}$ for a more general class of distributions than the mixture of Gaussians?
- Is it possible to guarantee that for certain target distributions and schedules Assumption 1 is satisfied? This is important, as it is not easy to manually check for every setting. 
- Is there always a $v_t$ such that $X_t \sim \tilde{\pi}_t$ for all $t$ in (4) and it generates $\tilde{\pi}_t$ i.e. satisfies the Fokker-Planck equation?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730501499681

thread_type:
official_review

<thread_object>
creator: reviewer_krfs

<content>
summary:
The paper proves new bounds on efficacy of Annealed Langevin MC by creating an SDE which approximates the target and then analyzing the discretization error introduced by simulating it on a finite machine

strengths:
Computational bounds are generic- a smoothness constraint for target distribution, rather than isoperimetry assumptions.

Explanations are clear and well-illustrated, suiting a broad generalist audience by mixing intuition-building and commentary.

The direct impact of this paper is unclear, but the methods it introduces are interesting contributions

weaknesses:
The paper shares the weakness of most such oracle inequalities, which is that they reassure us of convergence of an algorithm under assumptions that are hard to verify on a particular problem

Note that I am not expert in all the components of the authors' proof, and have taken some steps on faith.

questions:
What does $\beta$-smoothness imply about the target distribution? Can we know this in general, especially where the density arises from a posterior? It seems like this quantity would be hard to calculate for even tractable likelihoods, let alone general posterior densities.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730720373511

thread_type:
official_review

<thread_object>
creator: reviewer_4z2b

<content>
summary:
The paper studies the convergence of annealed Langevin algorithms for general, possibly non-log concave, target distributions. The authors propose using an argument based on Girsanov's theorem to bound the distance between a reference process and a Langevin SDE that is based on an annealing schedule of target distributions, bridging the target to a base distribution. The authors also consider a practical implementation of the continuous time SDE, based on the exponential integrator scheme.

strengths:
The authors consider an interesting problem, which is actively researched. Most of the research has focused on obtaining convergence bounds for standard Langevin algorithms imposing strong assumptions on the target distribution, such as log-concavity or log-Sobolev inequalities. The approach in this paper seems original, since it differs from strategies that are typically used in the literature as explained by the authors. Moreover, it gives convergence bounds in KL for an annealed Langevin algorithm, for which few results are available, and under "weak" conditions. Overall, the paper is clear enough to follow what the authors are doing and has a cohesive story.

weaknesses:
- The authors consider a specific annealing schedule and it is not clear if their results give any kind of clarity on how the schedule affects the overall complexity of the algorithm. 

- In several occasions I would have appreciated clearer explanations. For instance, why considering the exponential integrator scheme instead of an Euler scheme? Why considering that specific annealing schedule? It is also not completely clear why one finds the topics in Section 2.4 at a first read. It would be beneficial to convince the reader that these mathematical concepts are needed, also giving an idea of how these are useful in what follows. 

- The sketch of the proof of Theorem 2 is not very clear to me and I find it should be heavily improved. It seems the authors compute the KL between two processes, neither of which is the discretisation scheme. Then the notation in the proof in the appendix seems inconsistent, denoting by \nu^{ALMC} the measure of another process. As I motivate in the Questions section below, I struggle to see how the authors use Lemma 1 in the proof of this result. This should be clarified if the paper is to be considered for acceptance.

- the contribution is in general interesting, but I would have liked to find analyses such as some (simple) numerical simulations showing that the convergence of the algorithm is for instance not exponential in the dimension.

- Table 1 gives an overview of results available in the literature, but these are for various algorithms which are denoted by corresponding  acronyms.  This is not a very clear comparison in my opinion, and it would be nice to have a short explanation at the differences between these algorithms.

- the intuitive explanations in lines 318-323 are rather obvious and are not really clarified by the results in the paper. This should be clear to the reader to avoid confusion.

questions:
- Theorem 2 uses Lemma 1, that is a consequence of Girsanov's theorem. As far as I understand it uses it to bound the KL between the SDE in Equation (6) with the "true" SDE, that is Equation (4). However, the SDE (6) clearly does not fit into Lemma 1 of the paper, since a term in the drift depends on a past state. The authors should explain their reasoning and not just refer to Lemma 1, which is not general enough.

- in Lemma 2, in what sense do the authors say the equality is "reachable"?

- in lines 373-374 the authors state that the Euler scheme is "non-optimal", without a supporting reasoning. What are their reasons for writing that?

- In Example 1, the authors say that \rho_1 is "readily sampleable". Is it then the case that \rho_0 is a simple distribution that can be sampled? Here the notation is not clear yet. Also, in point 2. of the example, are the LSI and PI constants exponentially dependent on d for the distribution \rho_t for any t? In general, this example could be made clearer.
</content>
</thread_object>

</thread>
<thread>
creation date:
1732162047078

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
General response

comment:
We would like to extend our sincere gratitude to all the reviewers for their insightful comments and valuable feedback. We are pleased that the reviewers recognized the importance of the annealed sampling problem addressed in the paper (reviewer 4z2B), appreciated the novelty of our technical contributions (reviewers 4z2B, KRFS, WqFj, WbhJ), and commended the clarity of our presentation (reviewers 4z2B, KRFS, WbhJ). Our analysis introduces a novel framework for understanding annealing algorithms through the action of the interpolation trajectory, setting our work apart from existing bounds by avoiding dependence on isoperimetric inequalities. We believe this framework is going to significantly advance the understanding of annealed sampling algorithms and pave the way for future research directions, including the design of annealing schedules.

While *all* reviewers agreed on the novelty of our approach, two primary concerns were raised:

1. The need to improving the writing in several parts of the paper to enhance both accuracy and readability for a general audience.
2. The lack of simple numerical results to validate the theoretical findings presented in the paper.

In our revised manuscript, we have made substantial improvements to address these concerns, providing clarifications for all raised questions. Notable changes in the main paper are marked $\color{red}{\rm red}$ and enumerated below:

- We have added descriptions of the algorithms compared in Tab. 1.
- The statement of Girsanov's theorem (Lem. 1) has been modified to include cases where the drift terms in the SDE depend on the past state.
- We have clarified the statement "the equality is reachable" in Lem. 2. and expanded the discussion on the motivation for introducing OT concepts in Sec. 2.4.
- An error in Ex. 1 related to the dimensional dependence of LSI/PI constants has been corrected, along with adjustments to the corresponding proof in App. A.
- We added a footnote in the proof of Thm. 1 for rigor.
- In Sec. 6, we added a numerical experiment to verify the findings in Ex.2, and showed that the dependence of the number of iterations required to achieve a certain accuracy in KL divergence depends polynomially on $r$.
- Finally, we slightly modified the layout of the paper to fit the additional content.

We have also addressed all concerns in the individual responses below. We hope that the changes we have made reflect our commitment to improving the quality of our work and that the revised manuscript meets your expectations.
</content>
</thread_object>

</thread>


Metareview:
metareview:
This paper provides a non-asymptotic analysis of annealed MCMC. Authors establish an oracle complexity of a simple variant of annealed Langevin Monte Carlo and provide convergence guarantees for this algorithm to the target in KL divergence. Target is assumed to have a , smooth potential and the rate is in terms of action of a curve of measures that interpolates the target, which is an intriguing parameter.


This paper was reviewed by four reviewers the following Scores/Confidence: 8/4, 8/3, 6/4, 6/4. I think the paper is studying an interesting topic and the results are relevant to ICLR community. The following concerns were brought up by the reviewers:

- It is not clear if the rate is tight (most likely it is not). As is, the result seems to be another upper bound on sampling from a class of targets.
  
- Assumptions are not discussed in detail and compared with standard ones in the literature. This should be improved.

- Even though this is a theoretical paper, a detailed experimentation section would make the paper much stronger.
  
- minor typos were pointed by the reviewers. These should be carefully addressed. One reviewer requested detailed explanation of the integrator.


Authors should carefully go over reviewers' suggestions and address any remaining concerns in their final revision. Based on the reviewers' suggestion, as well as my own assessment of the paper, I recommend including this paper to the ICLR 2025 program.