On the Mode-Seeking Properties of Langevin Dynamics

We would like to thank Reviewer nEdB for his/her time and constructive feedback on our work. Below is our response to the questions and comments in the review.

1- Guarantees in terms of standard distance measures between probability models

Re: Please refer to global rebuttal #1.

2- Comparison between vanilla and annealed Langevin dynamics

Re: We would like to clarify that our theoretical results do not imply that annealed Langevin dynamics is worse than the vanilla Langevin dynamics, since the hidden constant in the notation $\Omega(d)$ in Theorem 2 is smaller than Theorem 1. Note that the hidden constant $c_1$ in $\Omega(d)$ for Theorem 1 is

c_1 = \min \left\\{ \frac{1}{2} \left(\frac{\nu_0^2-\nu_{\max}^2}{8\nu_0^2}\right)^2 , \frac{1}{8} \left( \log \left( \frac{\nu_{\max}^2}{\nu_0^2} \right) - \frac{\nu_{\max}^2}{2\nu_0^2} + \frac{\nu_0^2}{2\nu_{\max}^2} \right), \frac{1}{32}, \frac{(\nu_0^2-\nu_{\max}^2)^2}{32\nu_0^2(\nu_0^2+\nu_{\max}^2)} \right\\},

while the constant $c_2$ for Theorem 2 is

c_2 = \min \left\\{ \frac{1}{2} \left(\frac{\nu_0^2-\nu_{\max}^2}{8(\nu_0^2+c_{\sigma}^2)}\right)^2 , \frac{1}{8} \left( \log \left( \frac{\nu_{\max}^2+c_{\sigma}^2}{\nu_0^2+c_{\sigma}^2} \right) - \frac{\nu_{\max}^2+c_{\sigma}^2}{2(\nu_0^2+c_{\sigma}^2)} + \frac{\nu_0^2+c_{\sigma}^2}{2(\nu_{\max}^2+c_{\sigma}^2)} \right), \frac{1}{32}, \frac{(\nu_0^2-\nu_{\max}^2)^2}{32(\nu_0^2+c_{\sigma}^2)(\nu_0^2+\nu_{\max}^2+2c_{\sigma}^2)} \right\\}.

3- Effect of dimension reduction on mode covering hardness

Re: As mentioned in Theorem 5, we prove a linear reduction from learning the high-dimensional variable to a constant-dimensional variable. We define $\tau(\varepsilon/d)$ as the iteration complexity for Langevin dynamics to learn a $Q$-dimensional distribution (for constant $Q$) within $Q \cdot \varepsilon/d$ total variation distance. We note that Theorem 1 of [1] implies that for a $Q$-dimensional distribution $P(\mathbf{x}^{(q)} \mid \mathbf{x}^{(1)}, \cdots, \mathbf{x}^{(q-1)})$ with smoothness and local nonconvexity assumptions on the log-pdf (specified in Appendix A of [1]), we have

$$\tau(\varepsilon/d) \le c \cdot \frac{d^2}{\varepsilon^2} \log \left( \frac{d^2}{\varepsilon^2} \right)$$

for some constant $c > 0$. Therefore, Theorem 5 shows that Chained Langevin dynamics can achieve $TV(\hat P(\mathbf{x}), P(\mathbf{x})) \le \varepsilon$ in $\frac{c}{Q} \cdot \frac{d^3}{\varepsilon^2} \log \left( \frac{d^2}{\varepsilon^2} \right)$ iterations.

[1] Ma, Y. A., Chen, Y., Jin, C., Flammarion, N., & Jordan, M. I. (2019). Sampling can be faster than optimization. Proceedings of the National Academy of Sciences, 116(42), 20881-20885.

I thank the authors for their response. However, two points are still quite unclear to me.

Explaining the limitation of Annealed Langevin dynamics. I thank the authors for their clarification on the bounds for Vanilla and Annealed Langevin Dynamics. I understand these bounds have hidden constants and therefore one cannot easily claim that Annealed Langevin Dynamics performs worse. That being said, the authors seem to claim a serious limitation of Annealed Langevin Dynamics in Theorem 2, which I interpret to be that particles following this process will always be "far enough" from some mode and this hinders convergence. As I stated above: this is surprising, given the widespread success of Annealed Langevin Dynamics for efficiently sampling from multimodal target distributions. Reviewer Ave7 also picked up on this, saying "unlike what was believed, annealed Langevin also fails." Because this is surprising, Theorem 2 deserves some commentary, some high-level argument explaining why this popular method may fail. Pointing to the proof is not enough. As it stands, Theorem 2 is formally stated at the end of section 4.1 with no explanation.

Explaining the benefits of Chained Langevin dynamics. While I understand the authors' result, I still find it unclear why using the conditionals of the target $p(x_i | x_{-i})$ makes the sampling problem easier than sampling from the joint target $p(x)$. The authors mention Theorem 1 of [1] which they apply to each conditional $p(x_i | x_{-i})$, but couldn't this Theorem be applied directly to the joint target directly $p(x)$? The authors also mention that the joint target has a multi-dimensional input, while the conditionals have a one-dimensional input: still, that is not a satisfying explanation. Sampling from a one-dimensional distribution is not necessarily easier. Again, could the authors provide some high-level arguments on why sampling from the conditional distributions helps, versus sampling from the joint distribution?

[1] Ma, Y. A., Chen, Y., Jin, C., Flammarion, N., & Jordan, M. I. (2019). Sampling can be faster than optimization. Proceedings of the National Academy of Sciences, 116(42), 20881-20885.

We thank Reviewer Ave7 for his/her time and constructive feedback on our work. Below is our response to the questions and comments in the review.

1- Insights behind Theorems 1,2 and the role of mode $P^{(0)}$

Re: Please note that in Theorem 1, the mode $P^{(0)}$ plays the role of a large-variance Gaussian mode surrounding the other modes $P^{(i)}$’s with a significantly lower variance. Our intuition is that if the Langevin Dynamics gets initialized at a sample from $P^{(0)}$, then the score function will be dominated by the mode $P^{(0)}$, where the PDF of the other modes is expected to have a significantly less impact on the PDF of the mixture distribution (assuming a high dimension $d$). Therefore, the Langevin Dynamics is expected to randomly explore a large area in $\mathbb{R}^d$ (due to the high variance of $P^{(0)}$) which makes finding the remaining low-variance mode $P^{(i)}$’s overly expensive, requiring an exponential time (in terms of $d$) to find the missing modes. Theorems 1,2 formalize this intuition and prove the iteration complexity of finding the low-variance modes will indeed exponentially grow with the dimension.

We note that the above result can be further extended to a hardness result under the mean separation. Here, we again assume a high-variance mode $P^{(0)}$ filling in the space between the support sets of the low-variance mode $P^{(1)},\ldots , P^{(k)}$ (with bounded support sets with Euclidean radius $r$). This time, we suppose the Langevin Dynamics is initialized at a sample drawn from $P^{(1)}$ whose mean vector is sufficiently separated from the mean of the other modes. Then, there would be two possibilities for the Langevin Dynamics. Either the dynamics will remain in the support set of $P^{(1)}$ which cannot capture the other modes, or the dynamics would exit $P^{(1)}$‘s support set which will reduce to the case in Theorem 1, requiring the exploration of a large subset of $\mathbb{R}^d$ due to the high variance of the surrounding mode $P^{(0)}$. We will add the remark explaining the implications of Theorem 1 in such a setting where the support sets of $P^{(1)},\ldots , P^{(k)}$ are bounded with a small radius and have sufficiently distant means.

2- "Do similar results hold when $p_0$ is initialized at $\mathcal{N}(\mathbf{0}_d, \mathbf{1}_d)$?"

Re: We note that our theoretical result holds as long as, with high probability, the initial sample $\mathbf{x}_0$ will be far from the vector space of \left\\{ \mathbf{\mu}\_i \right\\}\_{i \in [k]}. If we assume the sample is initialized as $\mathbf{x}_0 \sim \mathcal{N}(\mathbf{0}_d, \mathbf{1}_d)$ and $P^{(0)}$ satisfies $\nu_0 < 1$, similar to Proposition 2 we know \left\\| \mathbf{n}\_0 \right\\|^2 \ge \frac{3\nu_0^2+\nu\_{\max}^2}{4} d with high probability. Therefore by Proposition 3, we obtain the same result as Theorem 1.

We would like to thank Reviewer 7G8Z for his/her time and constructive feedback and suggestions on our work. Below is our response to the questions and comments in the review.

1- The order of patches in Chained Langevin Dynamics

Re: In our analysis, the convergence rate of chained Langevin dynamics does not change with the patch ordering, as long as each patch can be accurately sampled according to the conditional distribution given the previous patches (an assumption that is supposed to hold for running Langevin Dynamics). To test this, we have performed experiments by (uniformly) randomly choosing the order of patches. As suggested in Figure 1 of rebuttal PDF, the numerical results looked similar to the results of our original implementation.

2- Selection of the patch size hyperparameter

Re: In the paper’s experiments on MNIST and Fashion-MNIST datasets, we chose patch size $Q=14$. To address the reviewer’s question, we tested different values of $Q \in \\{1,7,14,28\\}$. As suggested by Figure 2 in rebuttal PDF, the experimental results look insensitive to a moderate (not overly large) choice of patch size.

3- Selection of the neural network architecture

Re: We chose different architectures for vanilla Langevin dynamics and chained Langevin dynamics due to the difference in the learning objectives. In vanilla and annealed Langevin dynamics, we used a U-Net to jointly estimate the score function of every dimension of the sample due to its high capacity. In chained Langevin dynamics, we applied a Recurrent Neural Network (RNN) to memorize information about the previous inputs and estimate the conditional distribution of the next patch.

4- The reference on the reconstruction of images

Re: We thank the reviewer for introducing the related work. Our intuition of sequential sampling is echoed by the idea in the related work about reconstructing the true signal from its compression at a high level. We will discuss the work in the revised text.

We thank Reviewer 3VKE for his/her time and detailed feedback on our work. Below is our response to the questions and comments in the review.

1- Guarantees in terms of standard distance measures between probability models

Re: Please refer to global rebuttal #1.

2- Insights of the lower bounds' dependence on the covariance difference

Re: We understand the reviewer’s comment on the impact of covariance separation on our results in Theorems 1,2. In the following, we first explain our intuition behind the setting in these theorems, and next we argue how this result can be extended to the case with mean separation.

Please note that in Theorem 1, the mode $P^{(0)}$ plays the role of a large-variance Gaussian mode surrounding the other modes $P^{(i)}$’s with a significantly lower variance. Our intuition is that if the Langevin Dynamics gets initialized at a sample from $P^{(0)}$, then the score function will be dominated by the mode $P^{(0)}$, where the PDF of the other modes is expected to have a significantly less impact on the PDF of the mixture distribution (assuming a high dimension $d$). Therefore, the Langevin Dynamics is expected to randomly explore a large area in $\mathbb{R}^d$ (due to the high variance of $P^{(0)}$) which makes finding the remaining low-variance mode $P^{(i)}$’s overly expensive, requiring an exponential time (in terms of $d$) to find the missing modes. Theorems 1,2 formalize this intuition and prove the iteration complexity of finding the low-variance modes will indeed exponentially grow with the dimension.

We note that the above result can be further extended to a hardness result under the mean separation. Here, we again assume a high-variance mode $P^{(0)}$ filling in the space between the support sets of the low-variance mode $P^{(1)},\ldots , P^{(k)}$ (with bounded support sets with Euclidean radius $r$). This time, we suppose the Langevin Dynamics is initialized at a sample drawn from $P^{(1)}$ whose mean vector is sufficiently separated from the mean of the other modes. Then, there would be two possibilities for the Langevin Dynamics. Either the dynamics will remain in the support set of $P^{(1)}$ which cannot capture the other modes, or the dynamics would exit $P^{(1)}$‘s support set which will reduce to the case in Theorem 1, requiring the exploration of a large subset of $\mathbb{R}^d$ due to the high variance of the surrounding mode $P^{(0)}$. We will add the remark explaining the implications of Theorem 1 in such a setting where the support sets of $P^{(1)},\ldots , P^{(k)}$ are bounded with a small radius and have sufficiently distant means.

3- Hidden constants in the $\Omega$ notation

Re: In Theorem 1, notation $\Omega(d)$ means $\Omega(d) \ge cd$, for the following constant $c$

c = \min \left\\{ \frac{1}{2} \left(\frac{\nu_0^2-\nu_{\max}^2}{8\nu_0^2}\right)^2 , \frac{1}{8} \left( \log \left( \frac{\nu_{\max}^2}{\nu_0^2} \right) - \frac{\nu_{\max}^2}{2\nu_0^2} + \frac{\nu_0^2}{2\nu_{\max}^2} \right), \frac{1}{32}, \frac{(\nu_0^2-\nu_{\max}^2)^2}{32\nu_0^2(\nu_0^2+\nu_{\max}^2)} \right\\},

when $d$ is greater than

\max \left\\{ 8 \left( \log \left( \frac{\nu_{\max}^2}{\nu_0^2} \right) - \frac{\nu_{\max}^2}{2\nu_0^2} + \frac{\nu_0^2}{2\nu_{\max}^2} \right)^{-1} \log \left(\frac{3\nu_0^3}{w_0 \min_{i\in[k]}\nu_i^2}\right), \frac{8\nu_0^2(3\nu_0^2+\nu_{\max}^2)}{\pi(\nu_0^2-\nu_{\max}^2)^2} \right\\}.

For example in the Gaussian mixture in our synthetic experiments (Section 6), $\nu_0=\sqrt{3}$, $\nu_1=\nu_2=1$ and $w_0=0.2$, therefore $\Omega(d) \ge \frac{1}{288} d$ for any $d \ge 149$. The constant in Theorem 2 can be obtained by substituting $\nu_i^2$ with $\nu_i^2+c_{\sigma}^2$. We will include these constants in the main text.

4- Running time of Chained Langevin Dynamics

Re: In Theorem 5, we define $\tau(\varepsilon/d)$ as the iteration complexity for Langevin dynamics to learn a $Q$-dimensional distribution (for constant $Q$) within $Q \cdot \varepsilon/d$ total variation distance. We note that Theorem 1 of [1] implies that for a $Q$-dimensional distribution $P(\mathbf{x}^{(q)} \mid \mathbf{x}^{(1)}, \cdots, \mathbf{x}^{(q-1)})$ with smoothness and local nonconvexity assumptions on the log-pdf (specified in Appendix A of [1]), we have

$$\tau(\varepsilon/d) \le c \cdot \frac{d^2}{\varepsilon^2} \log \left( \frac{d^2}{\varepsilon^2} \right)$$

for some constant $c > 0$. Therefore, Theorem 5 shows that Chained Langevin dynamics can achieve $TV(\hat P(\mathbf{x}), P(\mathbf{x})) \le \varepsilon$ in $\frac{c}{Q} \cdot \frac{d^3}{\varepsilon^2} \log \left( \frac{d^2}{\varepsilon^2} \right)$ iterations.

5- Motivation for studying annealed Langevin dynamics with exact score function

Re: We note that the key difference between Langevin dynamics and denoising diffusion models (DDPM) is that the DDPM’s update rule scales the sample $\mathbf{x}_{i-1}$ with a factor of $\frac{1}{\sqrt{1-\beta_i}}$ at every iteration while the Langevin dynamics do not. The difference is referred to as the variance exploding property of Langevin dynamics and the variance preserving property of DDPM [2]. We think the scaling of samples is an important factor in analyzing the mode-seeking properties of Langevin dynamics.

6- Results of Lee et al. (2018) [3] regarding isotropic Gaussian mixtures

Re: Please refer to Theorem K.1 in Appendix K (Lower bound when Gaussians have different variance) of Lee et al. (2018) [3].

[1] Ma, Y. A. et al (2019). Sampling can be faster than optimization. PNAS, 116(42), 20881-20885.

[2] Song, Y. et al. (2020c). Score-based generative modeling through stochastic differential equations. In ICLR.

[3] Lee, H. et al. (2018). Beyond log-concavity: Provable guarantees for sampling multi-modal distributions using simulated tempering langevin monte carlo. Advances in neural information processing systems, 31.