Paper: Recursive Score Estimation Accelerates Diffusion-Based Monte Carlo (ICLR-2024)

Reviewer comment (Theory):
The argument for efficient sampling is clear: the authors choose a "small enough" window length $S$, so that its uniform discretization into steps of length $\eta$, produces sampling distributions $q_{k, t}$ that are log-concave and can therefore be sampled using a classical ULA (Unadjusted Langevin Algorithm) with a polynomial number of steps. Is that right?

Author response:
Yes, your understanding is correct.  Just as you mentioned in your later question, the choice of “S” is subtle in theory. We set it to be “just small enough” such that sampling distributions $q_{k,t}$’s are log-concave, while in this case, $S$ is still in the constant level. If it is too small (e.g., $O(\epsilon)$),  RS-DMC may not be able to achieve a quasi-polynomial complexity since the recursion number ($\tilde{O}(1/S)$) will be too large and harm the efficiency.

Reviewer comment (Theory):
However, the argument for efficient estimation is not explicit to me. Is it that we are, in the first equation of section 3.1., essentially just computing the empirical mean of $q$, rescaled by a certain factor. So the error of that estimate is the standard error of the mean (SEM), which is proportional to the variance of $q$. Do we have a handle on the variance of $q$?

Author response:
We have controlled the variance of q in our proof in an implicit manner. As you mentioned before, the distribution q is strongly log-concave (given in Eq. 4) and we can get a lower bound on the strongly log-concave constant (see Lemma E.2). Therefore, it also satisfies the log-Sobolev inequality due to Lemma F.4, which can imply the variance upper bound (see Lemma F.11). Then in our proof, we directly make use of the Sobolev inequality of $q$ to derive the high-probability bound (or concentration results) for estimating the mean of $q$.

Reviewer comment (Reproducibility):
Could the authors reunite their recommendations on setting hyperparameters in a list? For example:  
- the diffusion length T should be "big enough". Practically, T should be at least on the order of $\log d / \epsilon$  
- the window length  S should be "small enough" for the sampled distributions to be log-concave but "big enough" to avoid extra discretization steps. Practically, S should be on the order of $\frac{1}{2} \log \frac{2L + 1}{2L}$.  
- the length of a step inside a window $\eta$ should be ...? is there any recommendation for that?

Author response:
Sorry for the confusion. We have listed all settings of hyper-parameters in Theorem B.1, which is a detailed version of the main theorem (Theorem 4.1) in our paper. Specifically, the step inside a window, i.e.,  $\eta$ is required to be
$$\eta = C_\eta(M+d)^{-1}\epsilon,$$
where $C_\eta$ is a constant shown in Table. 2.

Reviewer comment (Experiments):
Can the authors discuss which hyperparameters would be easier or harder to set? For example, T seems to be easier to set for two reasons. First, we can actually compute $\log d / \epsilon$. Second, if T is "too big", the authors' convergence rate would still apply. However, setting S seems to be trickier. We cannot actually compute $\frac{1}{2} \log \frac{2L + 1}{2L}$ given that we do not know the smoothness constant $L$. So we would set S somewhat heuristically, and S could be potentially "too big" or "too small". If I correctly understand the authors' argument, the more dangerous situation is if S is "too big" as the sampling distributions might not be log-concave and this would introduce an exponential number of steps in sampling.

Author response:
This is a good point. In particular, the setting $S$ in our algorithm is similar to setting stepsize in optimization, where theoretically people need to use the smoothness constant to determine the feasible choices, while in practice they are just hyperparameters that can be free to tune. We agree that from a theoretical perspective, $S$ needs to be properly chosen in order to achieve good gradient complexity, making it “too big” may lead to a bad gradient complexity in theory. However, in practice, just like trying different stepsize for optimization, we can tune different $S$ in our algorithms. We have conducted numerical experiments on a multi-mode Gaussian-mixture case (see Appendix G of the latest revision), where we find that using a constant $S$ (after a simple grid search) in our algorithm can achieve improved sampling performance than baseline algorithms (DMC and ULA). This implies that the hyperparameters of our algorithms, including $T$ and $S$, are easy to set in practice.

Reviewer comment (Theory):
The efficiency of the authors' method seems to rely on the assumption that the "error propagation is benign, i.e. $l_{k, r}(\epsilon) = \epsilon$". Could the authors discuss the plausibility of this assumption? Is this a common assumption in the literature? Or are there works where benign error propagation appears in another context, supporting the claim that error propagation can indeed be benign in certain cases?

Author response:
A good question! In our paper, supposing $l_{k,r}(\epsilon)=\epsilon $is only for explaining our intuition easily. It is actually not an assumption, and we have proved it in the Appendix. Actually, we prove $l_{k,r}(\epsilon)=C\epsilon$ by setting $C=\frac{e^{-2S}}{640}$, which will also preserve the quasi-poly gradient complexity. Specifically, we show that a good score estimation approximation can be ensured

$$\mathbb{P}\left[\left\Vert v^\gets_{k,r\eta}(x) - \nabla\log p_{k,S-r\eta}(x)\right\Vert^2\le \epsilon\right] \ge 1-\delta$$

with the following benign error propagation of score estimation, i.e., 

$$\left\Vert \nabla \log p_{k,0}(x^\prime) -v^\prime_{k,0}(x^\prime)\right\Vert^2\le C\epsilon\le  \frac{e^{-2(S-r\eta)}\epsilon}{640}.$$

This statement adapts to Lemma E.7 in our paper.

Reviewer comment (Theory):
Practical usefulness is not clear. "S" could be very small yielding very large "K", which implies the method would be practically difficult to perform even if it is of quasi-polynomial gradient complexity.

Author response:
We respectfully disagree with your comments. It should be emphasized that the length of segments will be $S=1/2\log(1+1/2L)$, where $L$ is the Lipschitz parameter of the score function. Then it can be seen that $S$ is not very small but can be in the constant level (if treating $L$ to be constant). Consequently, the recursion number $K= 2\log [(Ld+M)/\epsilon] \cdot S^{-1}$ will not be extremely large, which is only logarithmic in the target error $\epsilon$ and dimension $d$.

Moreover, in practice, “S” and “K” are hyperparameters that can be tuned. We have conducted numerical experiments on a multi-mode Gaussian mixture case (which we added in Appendix G of the latest version) with tuned $S$ and $K$ (where $K$ is typically in the constant level), from which we can find that the proposed RS-DMC algorithm can achieve much better performance than baselines (DMC and ULA), even with fewer gradient complexities. This supports our theory and justifies the practical usefulness of RS-DMC.

Reviewer comment (Experiments):
Some experimental evidence for the practical usefulness is desired.

Author response:
Moreover, in practice, “S” and “K” are hyperparameters that can be tuned. We have conducted numerical experiments on a multi-mode Gaussian mixture case (which we added in Appendix G of the latest version) with tuned $S$ and $K$ (where $K$ is typically in the constant level), from which we can find that the proposed RS-DMC algorithm can achieve much better performance than baselines (DMC and ULA), even with fewer gradient complexities. This supports our theory and justifies the practical usefulness of RS-DMC.

To verify the usefulness of RS-DMC, we conduct numerical experiments on 2D multi-mode Gaussian mixture targets, which are non-log-concave with bad isoperimetric properties. We compare ULA, DMC, and RS-DMC. We have added the experiment details and anonymous link to our code in the revision. We attached the error comparisons (measured by Maximum Mean Discrepancy) among them as follows:

| Grad Num| ULA   | DMC   | RS-DMC |
| ------------------- | ----- | ----- | ------ |
| 200                 | 1.386 | 0.101 | 0.087  |
| 400                 | 1.386 | 0.088 | 0.074  |
| 800                 | 1.386 | 0.080 | 0.057  |
| 1600                | 1.382 | 0.078 | 0.044  |
| 3200                | 1.382 | 0.060 | 0.038  |

 From our experimental results, we can observe:

- ULA will quickly fall into some specific modes with an extremely high sampling error, measured in Maximum Mean Discrepancy (MMD), and is significantly worse than DMC and RSDMC.
- DMC will quickly cover the different modes. However, the process of converging to their means will be slower than RS-DMC.
- RS-DMC will quickly cover the different modes and converge to their means. Compared to DMC, it requires much fewer gradient complexities to achieve a comparable or better sampling error. For instance, RS-DMC with 800 gradient complexity can outperform DMC with over 3200 gradient complexity.

Then it is clear that RS-DMC significantly outperforms ULA and DMC in terms of gradient complexities, which backs up our theory and demonstrates the superior empirical performance of RS-DMC.  More details and visualizations of the experiment can be found in Appendix G of the latest revision, where we include our source code with an anonymous link. We will further add more numerical experiments to justify the practical advantages of RS-DMC in the revision.

Reviewer comment (Experiments):
Can you show the usefulness of the proposed method by numerical experiments on some concrete examples?

Author response:
To verify the usefulness of RS-DMC, we conduct numerical experiments on 2D multi-mode Gaussian mixture targets, which are non-log-concave with bad isoperimetric properties. We compare ULA, DMC, and RS-DMC. We have added the experiment details and anonymous link to our code in the revision. We attached the error comparisons (measured by Maximum Mean Discrepancy) among them as follows:

| Grad Num| ULA   | DMC   | RS-DMC |
| ------------------- | ----- | ----- | ------ |
| 200                 | 1.386 | 0.101 | 0.087  |
| 400                 | 1.386 | 0.088 | 0.074  |
| 800                 | 1.386 | 0.080 | 0.057  |
| 1600                | 1.382 | 0.078 | 0.044  |
| 3200                | 1.382 | 0.060 | 0.038  |

 From our experimental results, we can observe:

- ULA will quickly fall into some specific modes with an extremely high sampling error, measured in Maximum Mean Discrepancy (MMD), and is significantly worse than DMC and RSDMC.
- DMC will quickly cover the different modes. However, the process of converging to their means will be slower than RS-DMC.
- RS-DMC will quickly cover the different modes and converge to their means. Compared to DMC, it requires much fewer gradient complexities to achieve a comparable or better sampling error. For instance, RS-DMC with 800 gradient complexity can outperform DMC with over 3200 gradient complexity.

Then it is clear that RS-DMC significantly outperforms ULA and DMC in terms of gradient complexities, which backs up our theory and demonstrates the superior empirical performance of RS-DMC.  More details and visualizations of the experiment can be found in Appendix G of the latest revision, where we include our source code with an anonymous link. We will further add more numerical experiments to justify the practical advantages of RS-DMC in the revision.

Reviewer comment (Presentation):
The presentation can use a lot of improvement. The figures are currently difficult to understand. The algorithm blocks are also difficult to interpret.

Author response:
We will simplify our figures and the interpretation of our algorithm, and clarify the correspondence between them to make people understand it easier in the next version.

Reviewer comment (Theory):
Quasi-polynomial is interesting, but I would be curious to know if you think polynomial is possible/what the barriers are. Would really appreciate it if you put something about this at the end of the paper.

Author response:
A good question! In general, we think the polynomial complexity for sampling from general non-log-concave distribution can hardly be achieved. Our understanding is drawn from the optimization perspective in [raginsky2017non], which uses the sampling algorithm for solving a nonconvex optimization problem. Then, if the polynomial complexity can be achieved, we can probably optimize a general non-convex objective in polynomial time. This potentially contradicts the well-known hardness result, i.e., finding a global minimum for general non-convex objectives is an NP-hard problem.  

If we seek better complexity results (not limited to polynomial complexity) for SDE-based sampling algorithms, the barrier could be **balancing the property of the SDE path and the computational cost for implementing this path**. On one hand, it can be observed that the cost of implementing Langevin-based methods is relatively low, i.e., one only needs to calculate the gradient $\nabla f(x)$. **However**, the property of SDE path of Langevin-based methods is not good, it may need exponential time to reach the target distribution. On the other hand, for diffusion-based MCMC, the computational cost for implementation is heavy since it involves heavy and complicated score estimations along the entire SDE path. However,  for exchange, the path property of DMC can be much better, it can reach the target distribution within only logarithmic time. Our proposed algorithm, RS-DMC,  can be regarded as a good trade-off between the property of SDE path and the computational cost for implementing this path and thus achieves much lower gradient complexity compared to Langevin-based algorithms. Then we suspect that if such a trade-off can be improved, i.e., either a better SDE is developed or the implementation is more efficient, one can design a faster sampling algorithm with improved complexities (potentially polynomial).

Reviewer comment (Theory):
More intuition about why the complexity is quasi-polynomial would be useful.

Author response:
The quasi-polynomial complexity stems from the recursive score estimation. In our theory, each recursion step is decomposed into solving a polynomial number of subproblems,  denoted by $N$, to ensure stable error propagation, which will be further transferred to the next recursion step. Besides, we have also shown that the number of recursive steps $K$ is logarithmic in the target error $\epsilon$ and dimension $d$. Therefore, the total complexity of the score estimation will be $O(N^K)$ as we will need to solve at most $N^K$ subproblems. Recalling that $N$ is polynomial in $\epsilon$ and $d$ and $K$ is logarithmic in $\epsilon$ and $d$, we can conclude that $N^K$ is quasi-polynomial in $\epsilon$ and $d$.

We have already provided a rough calculation in Section 3.3. We will further polish this section to improve the clarity.

Reviewer comment (Theory):
What are the barriers to obtaining polynomial complexity?

Author response:
A good question! In general, we think the polynomial complexity for sampling from general non-log-concave distribution can hardly be achieved. Our understanding is drawn from the optimization perspective in [raginsky2017non], which uses the sampling algorithm for solving a nonconvex optimization problem. Then, if the polynomial complexity can be achieved, we can probably optimize a general non-convex objective in polynomial time. This potentially contradicts the well-known hardness result, i.e., finding a global minimum for general non-convex objectives is an NP-hard problem.  

If we seek better complexity results (not limited to polynomial complexity) for SDE-based sampling algorithms, the barrier could be **balancing the property of the SDE path and the computational cost for implementing this path**. On one hand, it can be observed that the cost of implementing Langevin-based methods is relatively low, i.e., one only needs to calculate the gradient $\nabla f(x)$. **However**, the property of SDE path of Langevin-based methods is not good, it may need exponential time to reach the target distribution. On the other hand, for diffusion-based MCMC, the computational cost for implementation is heavy since it involves heavy and complicated score estimations along the entire SDE path. However,  for exchange, the path property of DMC can be much better, it can reach the target distribution within only logarithmic time. Our proposed algorithm, RS-DMC,  can be regarded as a good trade-off between the property of SDE path and the computational cost for implementing this path and thus achieves much lower gradient complexity compared to Langevin-based algorithms. Then we suspect that if such a trade-off can be improved, i.e., either a better SDE is developed or the implementation is more efficient, one can design a faster sampling algorithm with improved complexities (potentially polynomial).

Reviewer comment (Theory):
Can you give more intuition for where the quasi-polynomial complexity comes from? Currently, there is a small block that is a bit difficult to interpret.

Author response:
The quasi-polynomial complexity stems from the recursive score estimation. In our theory, each recursion step is decomposed into solving a polynomial number of subproblems,  denoted by $N$, to ensure stable error propagation, which will be further transferred to the next recursion step. Besides, we have also shown that the number of recursive steps $K$ is logarithmic in the target error $\epsilon$ and dimension $d$. Therefore, the total complexity of the score estimation will be $O(N^K)$ as we will need to solve at most $N^K$ subproblems. Recalling that $N$ is polynomial in $\epsilon$ and $d$ and $K$ is logarithmic in $\epsilon$ and $d$, we can conclude that $N^K$ is quasi-polynomial in $\epsilon$ and $d$.

We have already provided a rough calculation in Section 3.3. We will further polish this section to improve the clarity.

Reviewer comment (Experiments):
Of course, the main weakness of this paper is the lack of experiment. The main contribution here is methodological and so one would expect to have numerical experiments backing up the methodological and theoretical results. I find it very strange that the authors did not include numerical experiments, comparing for example their method to ULA or such, on non-log concave target distributions and with runtime comparisons. I am willing to raise my score to 8 if the authors provide such comparisons.

Author response:
Thanks for your suggestions. We have added numerical experiments in the latest version (see Appendix G), where we conduct numerical experiments on 2D multi-mode Gaussian mixture targets, which are non-log-concave with bad isoperimetric properties. We compare ULA, DMC, and our RS-DMC. We presented the sampling errors (measured by Maximum Mean Discrepancy) obtained by these three algorithms as follows:  

| Grad Num| ULA   | DMC   | RS-DMC |
| ------------------- | ----- | ----- | ------ |
| 200                 | 1.386 | 0.101 | 0.087  |
| 400                 | 1.386 | 0.088 | 0.074  |
| 800                 | 1.386 | 0.080 | 0.057  |
| 1600                | 1.382 | 0.078 | 0.044  |
| 3200                | 1.382 | 0.060 | 0.038  |

Then we can find that:

- ULA will quickly fall into some specific modes with an extremely high sampling error, measured in Maximum Mean Discrepancy (MMD), and is significantly worse than DMC and RS-DMC.
- DMC will quickly cover the different modes. However, the process of converging to their means will be slower than RS-DMC.
- RS-DMC will quickly cover the different modes and converge to their means. Compared to DMC, it requires much fewer gradient complexities to achieve a comparable or better sampling error. For instance, RS-DMC with 800 gradient complexity can outperform DMC with over 3200 gradient complexity.

Overall, the additional experimental results, though primitive, have already shown the superior empirical performance of RS-DMC over DMC and ULA. This also supports our theoretical results.

More details and visualizations of the experiment can be found in Appendix G of the latest revision, where we include our source code with an anonymous link. Besides, we will add more numerical experiments to justify the practical advantages of RS-DMC in the revision.

Reviewer comment (Writing):
I found the paper to be quite difficult to read due to some of the notations that seem to be unncessarily confusing. The figures do not help either. I think that it would have been easier to just explain the algorithm by fixing some timesteps $t_1, \dotsc, t_K$ such that $X_{t_k} | X_{t_{k+1}}$ is log concave (using the inequality (4)) and then running the backward diffusion using the same discretization, without further diving the segments $[t_k, t_{k+1}]$, leaving it for the appendix.

Author response:
Thank you for your suggestions. Your understanding of the high-level idea of our algorithm design is exactly correct. We will simplify the notations as well as our diagram figures according to your suggestions in the revision.

Reviewer comment (Theory):
The main theorem in the paper results in a high probability bound of the form $KL(P_*||P_{0,S}^\leftarrow) = \tilde O(\epsilon)$ with probability $1-\epsilon$. This means that the accuracy of the method scales linearly with the probability, so that we are forced to take $\epsilon$ very close to $0$. This would result that in practice many samples are needed to obtain high accuracy.

Author response:
We would like to explain that the probability does not need to scale linearly with the accuracy and we set the probability $1-\epsilon$ only for the presentation simplicity (as we do not want to involve additional notations). Actually, the convergence of KL  $\mathrm{KL}(p_\ast\Vert p_{0,S}^\gets)=\tilde{O}(\epsilon)$ can be achieved with probability $1-\delta$ for arbitrary small $\delta$, where $\delta$ can be chosen independently on $\epsilon$ shown as follows:

Under almost the same setting as in Theorem 4.1, with probability at least $1-\delta$, it holds that $\mathrm{KL}(p_\ast\Vert p_{0,S}^\gets)=\tilde{O}(\epsilon)$. Besides, the gradient complexity of RS-DMC is
$$\exp\left[\mathcal{O}\left( \max\left\(\left(\log \frac{Ld+M}{\epsilon}\right)^3, \log \frac{Ld+M}{\epsilon}\cdot \log \frac{1}{\delta}\right) \right)\right].$$

We provide the formal version in Corollary B.2 in the latest revision with blue text.

Moreover, we would like to point out that the high-probability argument can be merged with the KL-divergence sampling error bound, leading to a clean sampling error bound (high probability argument is no longer needed) in total variation distance. To show this, by Pinsker inequality, we have

$\mathrm{TV}(p_*,p_{0,S}^\gets) \le \mathbb{P}[\mathrm{KL}(p_*\Vert p_{0,S}^\gets)=\tilde{O}(\epsilon)]\cdot \sqrt{\frac{1}{2}\mathrm{KL}(p_*\Vert p_{0,S}^\gets)}+\mathbb{P}[\mathrm{KL}(p_*\Vert p_{0,S}^\gets)\not =\tilde{O}(\epsilon)]\cdot \mathrm{TV}(p_*,p_{0,S}^\gets) =\tilde O(\epsilon^{1/2}).$

Then there is no high probability argument in such a sampling error bound, i.e., “it holds with probability $1$’.’ Then, we can arbitrarily choose the number of samples and then the sampling error bound in total variation distance can be exactly formulated as the function of this number.

Reviewer comment (Experiments):
The paper lacks numerical examples to demonstrate their techniques. This is important as it would demonstrate if the method is actually an improvement from the DMC paper. One reason for the lack of experiments could be that the total number of samples needed to run this algorithm grows with $n_k * m_k = O(1/\epsilon^5)$ which can be computationally expensive. If this was the case then the algorithm is not implementable in practice despite its remarkable properties.

Author response:
First, we would like to clarify that the $\tilde{O}(\epsilon^{-5})$ requirement of the total sample number is for the theoretical guarantee. From the theoretical perspective, we have demonstrated that the gradient complexity of our algorithm (see Theorem 4.1 ) has outperformed that of the DMC algorithm.  We have clearly mentioned this after Theorem 4.1.

Second, in practice, both $n_{k,r}$ and $m_{k,r}$ are hyperparameters and we are free to tune them. Then, to compare with DMC regarding their empirical performances, we have performed numerical experiments on 2D multi-mode Gaussian mixture targets (which are non-log-concave with bad isoperimetric properties). We have demonstrated the superior performance of RS-DMC over DMC (as well as ULA) in sampling these hard targets in the experiment. The experimental results are summarized in Appendix G of the latest version and some key comparisons (measured by Maximum Mean Discrepancy) are also provided as follows: 

| Grad Num| ULA   | DMC   | RS-DMC |
| ------------------- | ----- | ----- | ------ |
| 200                 | 1.386 | 0.101 | 0.087  |
| 400                 | 1.386 | 0.088 | 0.074  |
| 800                 | 1.386 | 0.080 | 0.057  |
| 1600                | 1.382 | 0.078 | 0.044  |
| 3200                | 1.382 | 0.060 | 0.038  |

Then, it can be seen that in order to match the accuracy obtained by RSDMC with 800 gradient calculations, DMC requires over 3200 gradient calculations. Besides we also visualize the samples generated by ULA, DMC, and RS-DMC. We also observe that the distribution generated by RS-DCM is closer to the target compared to ULA and DMC. More details can be found in Appendix G. We also included our source code with an anonymous link. Besides, we will add more numerical experiments to justify the practical advantages of RS-DMC in future versions.

To sum up, the proposed RS-DMC algorithm is demonstrated to outperform DMC in both theory and experiment. We hope that our clarification and additional experimental results can address your concerns.

Reviewer comment (Reproducibility):
Despite being a theoretical paper, I think the key of the proposed method is that it tries to find a way to implement the problem for nonconvex problems, something that DMC would struggle with. Because of that I wonder if this method can be implementable in practice, considering the computational challenges that come with it. It seems that the recursion, although significantly simplifying the sampling tasks, results in very strong computational requirements, so addressing this would be very important for me.

Author response:
RS-DMC is more implementable than DMC in real practice.  Since RS-DMC will be equivalent to DMC when we choose the total recursion number $K =1$. In real practice, the recursion number $K$ is considered a hyperparameter that can be flexibly tuned. You can utilize this hyper-parameter to adjust the iteration number of the inner loops. For example, when you find the iteration number of the inner loop, i.e., $m$, is too large, you can try to reduce it and increase $K$, which may help you save the running time. In summary, RS-DMC will not be worse than DMC and can be much better than DMC when $K>1$ is properly chosen (see our response to your previous questions regarding the numerical experiments).