On the Mode-Seeking Properties of Langevin Dynamics

7- Justification for the proposed algorithm

Our theoretical results in Theorems 1 and 2 suggest that Langevin dynamics could take an exponentially growing number of iterations to draw samples from low-variance modes in a mixture distribution. This is intuitive because Langevin dynamics may need to explore the entire $\mathbb{R}^d$ space until approaching some low-variance modes. Therefore, we proposed Chained Langevin dynamics, which divides the sample into patches of constant size $Q\ll d$, thus the volume it needs to explore is reduced to $\mathbb{R}^Q$ for each patch, which allows faster convergence.

8-"Is the idea behind the chained algorithm based/influenced on previous works?"

To the best of our knowledge, a sampling algorithm with a Chained-LD-like structure has not been proposed in the existing literature. The closest algorithm we could find in the literature is “autoregressive diffusion models” which is cited and discussed in the paper. However, please note that the autoregressive diffusion algorithm in [3,4,5] is a learning algorithm and its goal is to learn a generative model from training data. On the other hand, Chained-LD is a sampling algorithm and not a learning algorithm as we discussed in the global response. We have added Remark 3 to discuss the related works.

9- Typo at line 170-171

We thank the reviewer for pointing it out, which we have corrected in the revision.

[1] Max Welling and Yee W Teh. Bayesian learning via stochastic gradient langevin dynamics. In Proceedings of the 28th international conference on machine learning (ICML-11), pp. 681–688. Citeseer, 2011.

[2] Yang Song and Stefano Ermon. Generative modeling by estimating gradients of the data distribution. Advances in neural information processing systems, 32, 2019.

[3] Emiel Hoogeboom, Alexey A. Gritsenko, Jasmijn Bastings, Ben Poole, Rianne van den Berg, and Tim Salimans. Autoregressive diffusion models. In International Conference on Learning Representations, 2022

[4] Tong Wu, Zhihao Fan, Xiao Liu, Hai-Tao Zheng, Yeyun Gong, Yelong Shen, Jian Jiao, Juntao Li, Zhongyu Wei, Jian Guo, Nan Duan, and Weizhu Chen. AR-diffusion: Auto-regressive diffusion model for text generation. In Thirty-seventh Conference on Neural Information Processing Systems, 2023

[5] Kashif Rasul, Calvin Seward, Ingmar Schuster, and Roland Vollgraf. Autoregressive denoising diffusion models for multivariate probabilistic time series forecasting. In International Conference on Machine Learning, pp. 8857–8868. PMLR, 2021.

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

1- Connections with score-based generative modeling

As we have explained in the global response, our work focuses on sampling from a known distribution using variants of Langevin dynamics. To address the reviewer’s comment, we have clarified in the introduction that our paper studies the sampling task exclusively and does not focus on generative modeling. We have relocated all the citations for related references on generative modeling to the related work section, and the revised introduction solely focuses on the sampling task addressed by Langevin dynamics and Chained-LD algorithms. We have also included Remarks 1 and 2 to clarify the sampling nature of our proposed Chained-LD algorithm.

2- Connections with score-based sampling and GANs

Langevin dynamics and their properties have been extensively studied in the literature. In our work, we specifically focus on the mode-seeking tendencies of Langevin dynamics. Therefore, in the related work section, we discussed previous works on mode-seeking tendencies of sampling and generative modeling algorithms, including mode-seekingness of GANs. We think our studied problem is related to these references, because they all study the mode-seeking tendencies of sample generation algorithms.

3- The term “stochastic gradient Langevin dynamics”

We thank the reviewer for pointing this out. We have clarified the terminology in the introduction, and replaced it with "Langevin dynamics" in the revision.

4- Assumptions on the score function

As discussed in the global response, we would like to clarify that the main focus of our work is the mode-seeking properties of Langevin dynamics as a sampling algorithm from the score function of a given target distribution $P$. In our theoretical results for Langevin dynamics (Theorems 1 and 2), we assume that the sampling algorithm has access to the given score function and follows the Langevin dynamics update to generate a sample from the underlying distribution. Regarding Chained-LD, Algorithm 1 describes the steps of Chained-LD for a given score function (we do not learn or estimate it). Similarly, Proposition 1 analyzes its convergence by assuming access to the score function of the target distribution. We have added Remarks 1 and 2 to clarify this point in the revision.

5- Stationary distribution of Chained-LD sampler

Chained-LD’s stationary distribution is the same as the target distribution $P$. This statement follows from the chain rule in probability. Note that if every patch is sampled from the conditional distribution $P(x^{(q)}|x^{(1)}, \cdots, x^{(q-1)})$, the whole sequence follows from the target distribution $P$ due to the chain rule:

$P(x) = \prod_{q=1}^{d/Q} P(x^{(q)}|x^{(1)}, \cdots, x^{(q-1)})$

Also, the conditional densities used in Chained-LD do not require any extra information compared to the target distribution $P(x)$ assumed in Langevin dynamics. In both cases, the algorithms have complete knowledge of the distribution, where Langevin dynamics has access to the joint distribution $P$ and Chained-LD has access to the conditional distributions that, based on chain rule, have the same information as joint distribution $P$. We have added Remark 2 to make this point clear in the revision.

6- Interpreting the numerical results

In the synthetic Gaussian mixture experiments, we consider a mixture of three Gaussian modes. As shown in Figure 2, vanilla Langevin dynamics (VLD) cannot find the other two modes within $10^6$ iterations if the sample is initialized in mode 0, while chained vanilla Langevin dynamics (Chained-VLD) can find the other two modes in 1000 steps and correctly recover their frequencies as gradually increasing the number of iterations, which demonstrates the mode-seeking properties of VLD and verifies our theoretical results in Theorem 1 and Proposition 1.

In the MNIST and Fashion-MNIST experiments, we consider a mixture of two modes – the original samples from the datasets and the flipped samples. We followed the training algorithm proposed by [2] to estimate the score function from samples. Considering the resulting score function as a given score function, we ran the sampling algorithm of annealed Langevin dynamics (ALD) and Chained-ALD. As shown in Figures 3 and 4, if we initialize the sampling algorithms in one mode, ALD with bounded noise levels could not find the other missing mode, while ALD with a larger initial noise level and Chained-ALD is able to find both modes, which demonstrates the effect of noise levels on the mode-seeking properties of Langevin dynamics and verifies our theoretical results in Theorem 3.

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

1- Justification for Assumptions

We would like to clarify that our work is a theoretical study of the mode-seeking properties of Langevin dynamics (LD) as a sampling algorithm from a known distribution. In Theorems 1 and 2, our goal is to provide a theoretical setting where we can prove the possible exponentially growing number of iterations for LD to capture all the modes. We prove this result under mixtures of Gaussian and sub-Gaussian, where the latter case applies to random vectors with a bounded support set.

To the best of our knowledge, the existing theoretical results on LD prove satisfactory convergence of the dynamics in sampling from unimodal distributions. This differentiates our work from the existing literature and highlights the setting studied in our work where LD has a provably slow convergence. Since our theoretical results show the computational hardness of Langevin dynamics, they motivate finding a sampling algorithm with a lower iteration complexity. We therefore proposed Chained-LD to achieve this goal.

2- Comparison between Assumption 2 and Assumption 1

Note that Assumption 2 relaxes the Gaussian mode condition in Assumption 1 to a sub-Gaussian mode. This condition is significantly more general which generalizes the theoretical result to more general cases, e.g., where the distribution mode has a bounded support set (known to satisfy the sub-Gaussian condition). As we explained in Item 1 of our response, we do not know any existing computational hardness results for the discretized Langevin dynamic algorithm, which highlights the paper’s theoretical results in Theorems 1 and 2 in revealing settings where standard Langevin dynamics suffers from provably slow convergence.

9- Computational Efficiency

We would like to highlight that the time complexity of chained Langevin dynamics was significantly lower than vanilla/annealed Langevin dynamics in our experiments. Our comparison between Chained-LD and vanilla/annealed Langevin dynamics in the paper is based on iteration complexity, i.e., the number of queries to the score function $\nabla \log P(x^{(q)}|x^{(1)}, \cdots, x^{(q-1)})$ or $\nabla \log P(x)$. However, since Chained-LD only updates one patch at every iteration while vanilla/annealed LD updates the whole image, Chained-LD’s time complexity is significantly lower than vanilla/annealed Langevin dynamics. We have explained this point in Appendix D in the revision.

[1] Yang Song and Stefano Ermon. Generative modeling by estimating gradients of the data distribution. Advances in neural information processing systems, 32, 2019.

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

1- Related Reference

We thank Reviewer 7ZLa for pointing out the related work on the KL convergence guarantees of score-based diffusion models. We added a discussion about this paper in the revised paper.

2- Real-World Applicability and Initial Distribution

In scenarios where the underlying target distribution is unknown, as illustrated in Appendix D.2 in the revision, we can train an estimator to approximate the score function from the training dataset. To test the mode-seeking tendencies of the sampling algorithms, in our experiments, we initialize the samples from one mixture component and examine whether they can mix with other modes within a limited number of iterations. The experimental results indicate that Chained-LD is able to mix between modes even if it is initialized inside a mode. Therefore, we believe that Chained-LD is also capable of outputting samples from underrepresented groups.

3- Handling of Balanced Modes

We are unsure if this comment applies to our results, because we have not assumed that the in-between mode $P^{(0)}$ has a higher probability mass (parameter $\omega_0$ in the mixture representation) than the other modes. In fact, our theoretical results on the iteration complexity of Langevin Dynamics (LD) remain true even if the probability mass of in-between mode $P^{(0)}$ is small. Despite a minor mass, the higher variance of $P^{(0)}$ can still mislead the LD in the in-between space, due to the weaker influence of the other modes with lower variance but higher probability.

4- Assumption 2 and Gaussian Mixtures

We would like to clarify that Assumptions 2.(i)-(ii) automatically hold for Gaussian component $P^{(i)} = \mathcal{N}(\mu_i, \sigma_i^2 I)$. Assumption 2.(iii) is also true since the score function of Gaussian distribution $\nabla \log P^{(i)}(x) = -\frac{x}{\nu_i^2}$ is $\frac{1}{\nu_i^2}$-Lipschitz. Assumptions 2.(iv) and (v) are specific assumptions on the mean and variance of $P^{(i)}$, which is similar to Assumption 1 for our Gaussian mixture setting. We have made this point more clear in the revision.

5- Selection of Patch Size $Q$

In Proposition 1, our theoretical results suggest that the iteration complexity of Chained-LD scales as $\mathcal{O} \left( \exp(\mathcal{O}(LQ)) \frac{d^3}{\varepsilon^2\sqrt{Q}} \log \frac{d^2\sqrt{Q}}{\varepsilon^2} \right)$, which takes its minimum value when the patch size $Q$ is a constant not growing with $d$. In practice, the numerical results are insensitive to moderate values of constant $Q$, while for overly large $Q$, Chained-LD has mode-seeking tendencies similar to LD. We added Figures 7-9 in our revision, which demonstrates the effect of $Q$ in the synthetic Gaussian mixture case. We can observe that for dimension $d=100$, $Q \in \{1,4,10\}$ has similar performance, $Q=20$ needs more steps to find the other two modes, $Q=50$ almost cannot find other modes, and $Q=100$ reduces to vanilla Langevin dynamics.

6- Score Function Estimation

We would like to clarify that the score estimator is computed before running Chaine-LD, and in our numerical analysis, Chained-LD applies to a fixed score estimator. More precisely, we follow reference [1] which provides an algorithm for estimating the score function. More details can be found in Appendix D.2 of the revised draft.

7- Impact on Input Geometry

We believe the structural relationships in the data are maintained in Chained-LD due to the conditional distribution $P(x^{(q)}|x^{(1)}, \cdots, x^{(q-1)})$. Although the sample is divided into patches, the structural information in the previous patches can be reflected by the conditional distribution, therefore Chained-LD faithfully models the target distribution according to the chain rule $P(x) = \prod_{q=1}^{d/Q} P(x^{(q)}|x^{(1)}, \cdots, x^{(q-1)})$.

8- Proposition 1 and Convergence Bounds

We note that Proposition 1 is obtained from a divide-and-conquer method. Since the sample is divided into $d/Q$ patches, to achieve $\varepsilon$ total variation error for the sample, each patch achieves $\frac{\varepsilon}{d/Q} = \varepsilon Q/d$ total variation error in expectation.

7- Related works on autoregressive diffusion models

To the best of our knowledge, a sampling algorithm with a Chained-LD-like structure has not been proposed in the existing literature. The closest algorithm we could find in the literature is “autoregressive diffusion models” which is cited and discussed in the paper. However, please note that the autoregressive diffusion algorithm in [2,3,4] is a learning algorithm and its goal is to learn a generative model from training data. On the other hand, Chained-LD is a sampling algorithm and not a learning algorithm as we discussed in the global response. We have added Remark 3 to discuss the related works.

8- Comparison of our theoretical contributions with the existing literature

In this work, we study sub-Gaussian mixtures with an in-between distribution component $P^{(0)}$, which we have not seen in the existing literature. We prove that under such a mixture model, Langevin dynamics tends to be mode-seeking and may only find a subset of modes. To the best of our knowledge, mode-seeking properties of diffusion models under imbalanced modes have not been investigated in the literature. We also note that our analysis can be extended to annealed Langevin dynamics with bounded noise levels as we showed in Theorems 3 and 4 in Appendix B, indicating the effect of annealing noise levels on the mode-seeking tendencies of Langevin dynamics.

9- Forward process of Langevin dynamics

As discussed in [5], Langevin dynamics is the discretization of variance-exploding SDE, so the forward process converges to a Gaussian distribution with infinity variance, which is infeasible in practice. On the other hand, [6] suggests that Langevin dynamics asymptotically mix between modes and converge to the target distribution, regardless of the initialization. In our experiments, as we aim to study the mode-seeking tendencies of Langevin dynamics, we initialize the samples from one mixture component and examine whether they can mix with other modes within a limited number of iterations.

[1] Yi-An Ma, Yuansi Chen, Chi Jin, Nicolas Flammarion, and Michael I Jordan. Sampling can be faster than optimization. Proceedings of the National Academy of Sciences, 116(42):20881–20885, 2019.

[2] Emiel Hoogeboom, Alexey A. Gritsenko, Jasmijn Bastings, Ben Poole, Rianne van den Berg, and Tim Salimans. Autoregressive diffusion models. In International Conference on Learning Representations, 2022

[3] Tong Wu, Zhihao Fan, Xiao Liu, Hai-Tao Zheng, Yeyun Gong, Yelong Shen, Jian Jiao, Juntao Li, Zhongyu Wei, Jian Guo, Nan Duan, and Weizhu Chen. AR-diffusion: Auto-regressive diffusion model for text generation. In Thirty-seventh Conference on Neural Information Processing Systems, 2023

[4] Kashif Rasul, Calvin Seward, Ingmar Schuster, and Roland Vollgraf. Autoregressive denoising diffusion models for multivariate probabilistic time series forecasting. In International Conference on Machine Learning, pp. 8857–8868. PMLR, 2021.

[5] Yang Song, Jascha Sohl-Dickstein, Diederik P Kingma, Abhishek Kumar, Stefano Ermon, and Ben Poole. Score-based generative modeling through stochastic differential equations. In International Conference on Learning Representations, 2020c.

[6] Max Welling and Yee W Teh. Bayesian learning via stochastic gradient langevin dynamics. In Proceedings of the 28th international conference on machine learning (ICML-11), pp. 681–688. Citeseer, 2011.

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

1- Connection between hardness results and Chained-LD sampling algorithm

We believe that the two parts of our paper, the first on hardness results on iteration complexity of Langevin dynamics (LD) and the second on Chained-LD, are strongly connected. Note that the theoretical results in Theorems 1 and 2 show that LD can have an exponentially growing iteration complexity in dimension $d$ for the characterized mixture distribution. The complexity result implies that if the LD sampler could draw the coordinates sequentially, the exponential iteration complexity would be addressed. A natural way to implement this idea is to apply the Chain Rule in probability:

$P([x_1,\ldots , x_d]) = P(x_1)\cdot P(x_2|x_1)\cdots P(x_d|x_1,\ldots x_{d-1})$

Therefore, we propose Chained-LD to separate the sampling process across the $d$ coordinates and avoid the exponentially growing complexity of LD shown in Theorems 1 and 2. As we explained, the two parts of the paper are highly connected, and Chained-LD would not be motivated without showing the exponentially growing iteration complexity of LD in sampling from mixture distributions.

2- "I think the paper should make a clearer distinction between sampling from a known target distribution and generative modeling."

As we have explained in the global response, our work focuses on sampling from a known distribution using variants of Langevin dynamics. To address the reviewer’s comment, we have clarified in the introduction that our paper studies the sampling task exclusively and does not focus on generative modeling. We have relocated all the citations for related references on generative modeling to the related work section, and the revised introduction solely focuses on the sampling task addressed by LD and Chained-LD algorithms. We have also included Remarks 1 and 2 to clarify the sampling nature of our proposed Chained-LD algorithm.

3- “Firstly, chained LD is introduced in Algorithm 1 as a training algorithm of a generative model trained on data.”

We think this point is related to the reviewer’s previous comment. We believe the comment is due to a misinterpretation of Chained-LD as a learning algorithm. Note that the Chained-LD in Algorithm 1 is not a learning algorithm to learn a model from training samples. As we also discussed in the global response, Chained-LD is a sampling algorithm from an already known distribution $P$, which aims to address the exponentially growing iteration complexity of regular Langevin dynamics in sampling from mixture distributions (Theorems 1 and 2). We have clarified the sampling nature of Algorithm 1 in Remark 2 in the revised text.

4- “However, the paper does not explain why vanilla score-based generative models would fail on mixtures of sub-Gaussian distributions in the same way as SGLD”

As we responded to the previous points, Chained-LD is not a learning algorithm and it is only a sampling method similar to standard Langevin dynamics. Therefore, we did not compare it with generative modeling methods that learn the distribution from data. We hope that this clarification helps with the reviewer’s comment.

5- Log-concavity assumption in Proposition 1

We would like to clarify that as suggested by the reference [1], Proposition 1 only requires $P$ to be $m$-strongly log-concave outside a region of radius $R$, i.e., for $||\mathbf{x}^{(q)}|| > R$. Therefore, sub-Gaussian mixtures satisfying Assumption 2 would also satisfy this assumption for $R \gtrsim \max_i ||\mu_0-\mu_i||$. We have made this point clear in the revision.

6- Selection of patch size $Q$

In Proposition 1, our theoretical results suggest that the iteration complexity of Chained-LD scales as $\mathcal{O} \left( \exp(\mathcal{O}(LQ)) \frac{d^3}{\varepsilon^2\sqrt{Q}} \log \frac{d^2\sqrt{Q}}{\varepsilon^2} \right)$, which takes its minimum value when the patch size $Q$ is a constant not growing with $d$. In practice, the numerical results are insensitive to moderate values of constant $Q$, while for overly large $Q$, Chained-LD has mode-seeking tendencies similar to LD. We added Figures 7-9 in our revision, which demonstrates the effect of $Q$ in the synthetic Gaussian mixture case. We can observe that for dimension $d=100$, $Q \in \{1,4,10\}$ has similar performance, $Q=20$ needs more steps to find the other two modes, $Q=50$ almost cannot find other modes, and $Q=100$ reduces to vanilla Langevin dynamics.

We thank Reviewer Egfk for his/her time and feedback on our response. Regarding the raised points in the feedback,

1- "my concern is that the denoising score $\nabla \log P_\sigma (x)$ is typically intractable when only the target distribution is known"

First, we would like to clarify that in our experiments on the Gaussian mixture in Figure 2, we applied Algorithm 1 with $\sigma=0$, i.e. chained version of vanilla Langevin dynamics. Therefore, the improvement of Chained-LD over vanilla Langevin dynamics in Figure 2 does not require a choice of positive $\sigma>0$.

Also, in the specific case of Gaussian mixtures, such as the experiment in Figure 2, the score function $\nabla \log P_\sigma (x)$ is indeed tractable for a general $\sigma\ge 0$. To compare the performance between $\sigma=0$ and non-zero $\sigma \le 1$, we have run the experiment of Figure 2 with both vanilla LD (where $\sigma=0$) and annealed LD with $\sigma_{\max}=1$, as well as the chained version of both the algorithms. As shown in Figure 10 in the revised Appendix, the non-chained versions of both LD and ALD (with $\sigma_{\max}=1$) did not find the missing Gaussian modes until iteration $10^6$. On the other hand, their chained versions, i.e. Chained-VLD and Chained-ALD, could find the missing modes in only 1000 iterations.

As the experiments on the Gaussian mixture setting suggest, the gain by chaining Langevin dynamics is present under both zero and non-zero $\sigma$ parameters. We agree with the reviewer that the annealed LD with a general $\sigma>0$ can lead to an intractable score function in a general case; however, it does not affect our experiments on Gaussian mixtures where the score function under any $\sigma$ value can be efficiently computed.

2- "I would expect that a sampling algorithm that uses $\nabla \log P_\sigma$ instead of $\nabla \log P_0$ would already improve standard Langevin significantly in theory and practice, even without chaining."

As we explained in our response to Item 1, we ran the experiment on Gaussian mixtures and observed that the convergence speed gain by chaining Langevin dynamics is present under both zero and non-zero $\sigma$ parameter. Please refer to Figure 10 of the revised Appendix for the numerical results. Note that “vanilla” in the figure means $\sigma=0$.

3- "Are you proposing chained Langevin as a method to sample from pre-trained diffusion models?"

Please note that our work is a theoretical study of Langevin dynamics and their mode-seeking properties. Also, we would like to clarify that our proposed Chained-LD is intended to be a counterpart of the well-known Langevin dynamics. Therefore, as in practice, Langevin dynamics are not intended to be used for drawing samples from pre-trained models, we do not primarily intend to apply the Chained-LD algorithm to sample from pre-trained models.

To be clear about our theoretical study, our main point is that the (vanilla) Langevin dynamics can theoretically lead to exponentially-growing iteration complexity to sample from mixture distributions, while the counterpart Chained-LD could circumvent the exponential complexity in dimension $d$ by chaining the sampling process using the Chain Rule.

4- "If not, can you implement chained LD even when you don't have access to data sampled from the target distribution?"

If we correctly understood the reviewer’s question, the reviewer asks whether we can use Chained-LD if we only have access to $\nabla \log P_0(x)$ (where $\sigma =0$). The answer is yes, because as we explained the chaining approach is applicable to vanilla Langevin dynamics. For example, in Figure 2, the "Chained-VLD" refers to the case where $\sigma=0$, which we refer to as “Vanilla Langevin dynamics”.