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

Thank you for your appreciation for our work.

Weakness: 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 $\cal A$? whether there is an optimal value of $\gamma$ that could minimize the convergence rate?

Response: In this example, we obtain a uniform upper bound of the action $\cal A$ for all $\gamma\in[1,O(1)]$. We choose $\gamma\ge 1$ as intuitively, when $\theta$ is close to $1$, the distribution $\pi_\theta$ is more diffucult to sample from, and we should move more slowly. We agree that for each problem, there should be an optimal value of the parameter $\gamma$ that reaches the best sampling quality, yet studying the influence of $\gamma$ on the action $\cal A$ is highly non-trivial. This will be the focus of our study in the future.

We hope that the response above could successfully address your concerns.

Thank you for your careful reading of our paper and your valuable, constructive feedback.

Weakness 1 and Question 3: One cannot verify beforehand whether the curve is AC. Is it possible to guarantee that for certain target distributions and schedules Assumption 1 is satisfied?

Response: Absolute continuity is a fairly weak regularity condition on the curve of probability measures. Loosely speaking, as long as the curve $(\pi _\theta) _{\theta\in[0,1]}$ has a density $\pi _\theta(x)>0$ that is jointly $C^1$ w.r.t both $x\in\mathbb{R}^d$ and $\theta\in[0,1]$, then it is AC.

For instance, we can intuitively prove the proposition above in one-dimension. Let $F _\theta(x)=\int _{-\infty}^x\pi _\theta(u){\rm d}u$ be the c.d.f. of $\pi _\theta$, which is strictly increasing, and its inverse $F _\theta^{-1}$ is well-defined. $F^{-1} _{\cdot}(\cdot)$ is also jointly $C^1$. We know from standard textbook of OT that $W _2^2(\pi _\theta,\pi _{\theta+\delta})=\int _0^1|F _{\theta+\delta}^{-1}(q)-F _\theta^{-1}(q)|^2{\rm d}q.$ The above quantity should be $O(\delta^2)$ as $\delta\to0$ due to the $C^1$ property, and hence the curve is AC.

Weakness 2: 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.

Response: We agree that smoothness has several different definitions, such as being $C^1$ or $C^\infty$. The definition we used in this paper, the Lipschitz continuity of the gradient (see the "notations and definitions" part in Sec. 1), is also a common one that appears in many standard textbooks in optimization and sampling, such as Sec. 3.2 of Bubeck's Convex Optimization: Algorithms and Complexity and Chap. 4 of Chewi's Log-concave Sampling. We appreciate your comment for the potential ambiguity.

Weakness 3: There are too many indices for $\pi$ and the usage is incoherent.

Response: Sorry about the confusion the paper might have caused you. Throughout this paper, $(\pi _\theta) _{\theta\in[0,1]}$ is an AC curve bridging $\pi _0$, a simple distribution, and $\pi _1$, the target distribution, while $(\widetilde{\pi} _t=\pi _{t/T}) _{t\in[0,T]}$ is the curve after time-reparameterization. $0=\theta _0<\theta _1<...<\theta _M=1$ are the discrete time points on $[0,1]$, and $T _\ell=T\theta _\ell$'s are the discrete time points on $[0,T]$, with $h _\ell=T(\theta _\ell-\theta _{\ell-1})$ being the step size. We hope that this clarifies the usage of these notations.

Weakness 4: Simple experiments on sampling from a mixture of Gaussians using the proposed and existing methods would significantly boost the paper.

Response: Please see Sec. 6 for the experimental results, where we verified the findings in Ex. 2.

Thank you for your appreciation of our work and instructive feedbacks.

Weakness: 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.

Response: We agree that verifying these assumptions can be challenging in specific applications. However, our goal in establishing these conditions is to provide a broadly applicable theoretical framework that characterizes the behavior of annealed sampling in a general setting. These assumptions are intended to capture typical scenarios that arise in sampling from complex, multimodal distributions, allowing our results to extend to a wide range of problems.

Question:

1. What does $\beta$-smoothness imply about the target distribution?

   Response: Consider a target distribution $\pi\propto{\rm e}^{-V}$, where the potential $V$ is $\beta$-smooth. By definition, it means $-\beta I\preceq\nabla^2V\preceq\beta I$, which is equivalent to

   \begin{cases}V(y)\le V(x)+\left\langle\nabla V(x),y-x\right\rangle+\frac\beta2|y-x|^2\\ V(y)\ge V(x)+\left\langle\nabla V(x),y-x\right\rangle-\frac\beta2|y-x|^2\end{cases}

   Intuitively, this means the potential can be upper and lower bounded by quadratic functions at every point, meaning that it does not have regions of extremely high or low curvature. This assumption on the regularity of the target distribution aids in controlling the convergence of ALMC and in establishing the oracle complexity bounds.

2. 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.

   Response: Our paper focuses on the theoretical analysis of sampling, where we assume the potential is smooth to simplify the problem. In practice, the smoothness of the target distribution's potential requires case-by-case verification, and is calculable in some target distributions.

   For instance, consider the Bayesian logistic regression problem: the parameter has prior distribution $\beta\sim{\cal N}(0,\eta^2I)$, and the likelihood is $Y|X,\beta\sim{\rm Bernoulli}(\sigma(\beta^{\rm T}X))$, where $\sigma(t)=\frac{1}{1+{\rm e}^{-t}}$ is the sigmoid function. Then the posterior distribution given paired data points ${\cal D}=\{(x_i,y_i)\}_{i=1,2,...,N}$ is $p(\beta|{\cal D})\propto\exp\left(-\frac{\|\beta\|^2}{2\eta^2}-\sum_i\beta^{\rm T}x_i(1-y_i)-\sum_i\log(1+{\rm e}^{-\beta^{\rm T}x_i})\right)=:\exp(-V(\beta)).$ It is easy to verify that $\nabla^2V(\beta)=\frac{1}{\eta^2}I-\sum_ix_ix_i^{\rm T}\frac{{\rm e}^{\beta^{\rm T}x_i}}{(1+{\rm e}^{\beta^{\rm T}x_i})^2}.$ So for any unit vector $u$, $u^{\rm T}\nabla^2V(\beta)u=\frac{1}{\eta^2}-\sum_i(u^{\rm T}x_i)^2\frac{{\rm e}^{\beta^{\rm T}x_i}}{(1+{\rm e}^{\beta^{\rm T}x_i})^2}$, which implies $u^{\rm T}\nabla^2V(\beta)u\le\frac{1}{\eta^2}$ and $u^{\rm T}\nabla^2V(\beta)u\ge\frac{1}{\eta^2}-\sum_i(u^{\rm T}x_i)^2\ge\frac{1}{\eta^2}-\sum_i\|x_i\|^2$. Hence, $V$ is smooth with constant $\max\left(\frac{1}{\eta^2},\sum_i\|x_i\|^2-\frac{1}{\eta^2}\right)$.

   We agree that verifying smoothness may be hard in some general cases, let alone computing the constant $\beta$. Nevertheless, this assumption is common in non-asymptotic analyses, as it allows us to develop a clearer understanding of the algorithm’s behavior under idealized conditions, even if such conditions may be approximations in practical scenarios.

We hope that our response has successfully addressed your concerns.

Thank you for your careful reading of our paper and your valuable, constructive feedback.

Weakness 1: 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.

Response: The continuous-time analysis we presented in the paper is general and applies to any AC interpolation. In practice, the affect of annealing schedule to the complexity is a crucial question.

In Lem. 3(1), we showed that for fixed $\rho _0$ and $\rho _1$, the curve $\rho$ that has the minimial action ${\cal A}=\int _0^1|\dot\rho| _t^2{\rm d}t$ is the constant-speed Wasserstein geodesic. However, this does not have close form except in some trivial cases such as Gaussian distribution.

To realize constant-speed Wasserstein geodesic, one may do time-reparameterization on the Wasserstein gradient flow. For example, consider the OU process ${\rm d}Y _t=-Y _t{\rm d}t+\sqrt2{\rm d}B _t$, $Y _0\sim\pi$ (target distribution) and $Y _t\sim p _t$. The OU process is the Wasserstein gradient flow of the functional ${\rm KL}(\cdot\|{\cal N}(0,I))$, and after time-reparameterization of the curve $(p _t) _{t\in[0,\infty)}$, one may obtain a curve close to the constant-speed Wasserstein geodesic between $\pi$ and ${\cal N}(0,I)$. This curve is widely used in diffusion model, but we do not have closed form of the scores $\nabla\log p _t$ in general cases. As our goal is to obtain the complexity of learning-free sampling algorithms, we do not study this approach in the paper. Instead, we focus on the curve $\pi _\theta\propto\exp\left(-\eta(\theta)V-\frac{\lambda(\theta)}{2}\|{\cdot}\|^2\right),~\theta\in[0,1]$ in this paper. The effect of general annealing schedule to the sampling complexity will be the focus of our future study.

Weakness 2:

1. Why considering the exponential integrator scheme instead of an Euler scheme?

   Response: Since exponential integrator scheme reduces discretization error compared with Euler scheme. Recall that the SDE of ALD is ${\rm d}X _t=\left(-\eta\left(\frac{t}{T}\right)\nabla V(X _{t})-\lambda\left(\frac{t}{T}\right)X _t\right){\rm d}t+\sqrt{2}{\rm d}B _t,~t\in[0,T].$ Assume the discretization time points are $0=T _0<T _1<...<T _M=T$. On $[T _{\ell-1},T _\ell]$, the Euler discretization scheme regards the linear part in the drift term $-\lambda\left(\frac{t}{T}\right)X _t$ as unchanged, but actually the integral of $-\lambda\left(\frac{t}{T}\right)X _t$ on $[T _{\ell-1},T _\ell]$ can be computed in closed form, and by doing this, the exponential integrator scheme reduces discretization error. This is also a standard approach in accelerating the sampling of diffusion models.

2. Why considering that specific annealing schedule?

   Response: This specific family of interpolation curves is a popular choice in practice (see the last paragraph of section 5.1 for discussion). Moreover, another reason for choosing this specific family of interpolation curves is that the score functions $\nabla\log\pi _\theta$ have closed form, and therefore, unlike many other curves such as the Gaussian convolution curve $\pi _\theta=\pi*{\cal N}(0,\sigma^2(\theta)I)$, we do not need further training for approximating the score.

3. It is also not completely clear why one finds the topics in Section 2.4 at a first read.

   Response: Section 2.4 provides some essential technical tools in optimal transport for establishing our results. For instance, Lem. 2 is used in the proof of Thms. 1 and 2 to minimize the integral of squared $L^2$ norms, and the action will turn out to be a crucial property of the annealing curve (see, e.g., Ex. 2). Lem. 3 is on the relationship between action and isoperimetric inequalities, which sets this approach apart from the traditional analysis based on LSI or PI. We have added several linking sentences to inform the readers of the purpose of introducing these concepts and properties.

I thank the authors for their clarifications and modifications. I have changed my grade to 6.