Zeroth-Order Sampling Methods for Non-Log-Concave Distributions: Alleviating Metastability by Denoising Diffusion

We sincerely thank the reviewer for their valuable advice and comments and greatly appreciate the positive evaluation.

The authors are clear about the limitations of their method, for example the complexity in Corollary 3.1. is exponential not only in the dimension d but also in the smoothness L of the target (this holds even in one dimension).

That is right. Just to further remark about this fact: For general target distributions as considered in our paper, it was shown in [1] that the diffusion posterior sampling is computationally intractable in the worse case, and that worse case is among the distributions considered in our paper. Therefore, we feel if the theoretical complexity needs to be improved, one possibility is to restrict the target distributions to have more structures, while possibly adapting our algorithm to these structures to get better smoothness/dimension-dependence.

[1] Gupta, Shivam, et al. "Diffusion posterior sampling is computationally intractable." arXiv preprint arXiv:2402.12727 (2024).

To support my claim 2 in the 'Weaknesses' section, could the authors provide (i) the value of the average score error between the true score and the approximated score near t=T, while varying the number of queries to the target distribution (ie oracle complexity per score evaluation) and (ii) results of rejection sampling directly applied on the target distribution with same oracle complexity as other methods ?

Re (i): we added numerical comparison among the score estimation errors at $T$ for ZOD-MC, RDMC, RSDMC and SLIPS across different oracle complexities. We consider two target distributions (one in $2d$ and one in $5d$). In both cases, score estimation error in ZOD-MC is smaller than that in RDMC and comparable to those in RSDMC and SLIPS.

Re (ii): Rejection Sampling (RS) requires to construct an upper envelope to generate an accurate sample, and it is not known how to do that for general distributions. RS also suffers from the metastability. For example, to sample GMMs, the complexity can exponentially depend on the distance between modes. However, when combined with diffusion model, this issue can be alleviated because the intermediate target $p_{0|t}$ has more concentrate modes and we don't require very accurate samples when $t$ is large. To provide more evidence of this we make an experiment with a simple target distribution and demonstrate that rejection sampling is unable to sample from it. The example is a GMM with means $(0,0), (5,5), (-6,-8)$ and covariances $.1 I$, we then construct an envelope using knowledge of the means and stds of each mode (notice that this information is not generally available). We then create $5 \cdot 10^7$ proposals and only $3$ get accepted. When using the same oracle complexity our method is able to correcly sample from the target distribution.

The main (and crucial) limitation of ZOD-MC is the fact that it can only be applied to small dimensional settings ($d\le 5$), as rejection sampling is notoriously known to suffer from the curse of dimensionality (which is the reason why it is not a popular sampling method in practice). Although the authors acknowledge this limitation in the abstract, it is not well highlighted at all in the main of the paper: this limitation appears explicitly in Remark 5 in Section 3.3 in the Appendix, but not in the main paper. As suggested by NeurIPS guidelines, a section/paragraph 'Limitations' should be given in the main.

Besides mentioning the curse dimensionality in the abstract, we also mention the exponential dimenison dependency explicitly in Remark 4 in the main part of the paper. We will highlight this in the updated manuscript. However, it is not true that we can only handle $d\leq 5$ in fact in the experiment in Figure 1b we sample a problem of dimenson $7$ with high accuracy.

On the other hand, ZOD-MC requires to know the location of local minima of the energy landscape, i.e., the location of the modes of the target distribution. As indicated by the authors, this can be done by applying Newton's method on the target potential (see Remark 1) before or while sampling; however, this is a 2nd-order method, which violates the framework chosen by the authors (0-th order method). Although it is a common assumption to have access to the location of the modes in realistic settings (which turns out to be another challenge of sampling procedures), this requirement should be much more highlighted. Otherwise the title of the paper is misleading.

We agree ZOD-MC also needs an oracle about the minimum value of the potential. It can be relaxed to a global lower bound of $V$. Nevertheless, we agree it still needs to be implemented. Our practical implementation was mentioned in Remark 1, where we get an approximation using Newton's method, and like the reviewer mentioned, this is where 2nd-order information starts being used. Although other methods can also approximate the oracle, we agree with the reviewer and will clarify the writing and highlight what exactly are needed.

[1] Louis Grenioux et al. "Stochastic Localization via Iterative Posterior Sampling." ICML'24.

The key reason is we don't need a high-quality sample to approximate the score well for large $t$. As explained in line 173, when $t=T$ (which is chosen to be $\Theta(\log(d/\varepsilon))$ in Corollary 3.1), we only need one low-quality ($d$-accuracy in $W_2$) sample from $p_{0|T}$ to ensure the score estimation error is $O(\varepsilon)$. Therefore, even though the dual-logconvavity arguement in [1] says sampling from $p_{0|T}$ is comparable hard to sampling from the target, the diffusion-based sampling approach make things easier because it only requires a low-quality sample from $p_{0|T}$. Regarding the reviewer's question on the theoretical side, our Corollary 3.1 actually proves a worst-case complexity: we run the rejection sampler until it generates an accurate sample ($\delta(t)=0$). This complexity strictly upper bounds what we need in practice: as we addressed previously, we don't need a high-quality sample from $p_{0|t}$ when $t$ is large. Therefore, in practice we can set a threshold on the number of rejections to get a relative low-quality sample: if no proposal is accepted within this threshold value, we simply approximate the score by $-\frac{x}{1-e^{-2t}}$ (inspired by Lemma 1 in the paper). In this way, we save the number of queries without hurting the score-estimation accuracy.

The section on numerical experiments is incomplete for several reasons: no indication of tuning of the hyperparameters (N, T, step-size) either for ZOD-MC or other approaches, no ablation study on these hyperparameters to exhibit robustness, display of the estimation error of the true score for a large (where it becomes harder) and a variety of levels of oracle complexity, the plots in Figure 1 and Figure 4.a do not display any variance, missing related work as stated in my claim 1.

For ZOD-MC, since N,T and the step-size are values that are theoretically and empirically well understood, we didn't even tune them. We focus instead on understanding the sample quanlity with different oracle complexities and different target distributions. For other methods, we used either the official codes, or our best tuned versions. The reviewer is right that we should have explicitly mentioned this.

Regarding the score estimation error plots in Fig 1 & 4a, we have added the comparison to SLIPS, as displayed in Fig 1 in rebuttal supplementary pdf. To add the variance takes considerable amount of time, and we will add it in the revised manuscripts.

I have a question about Algorithm 1: to obtain the sampling recursion with the MC estimator, you first discretize the SDE with exponential integration and then replace the score by its MC estimator. Would it bring less discretization error to first replace the score in the SDE by the term involving the conditional expectation and then apply exponential integration upon it ? It seems to me that exponential integration would still be tractable in this case.

Thanks for the interesting question. It is true that we can approximate the score before applying the exponential integrator scheme. However, we don't think this would necessary decrease the order of the discretization error. From the analytical perspective, the discretization error factor in line 506-507 will change from $\lVert \nabla \ln p_{T-t}( \bar{X_t} ) -\nabla \ln p_{T-t_k}( \bar{X_{t_k}} ) \rVert^2$ to $\lVert \nabla \ln p_{T-t}( \bar{X_t} ) -\nabla \ln p_{T-t}( \bar{X_{t_k}} ) \rVert^2$. Noticing that in our paper, instead of splitting the first error via triangle inequality which would make the error bound strictly bigger than the second error, we analyzed it by looking at the dynamics of $\mathrm{d} \nabla\log p_{T-t}(\bar{X_t})$. This finer analysis help us obtain an upper bound of order $O(t-t_k)$. We believe this upper bound order also applies to the second error.

We would like to kindly direct the reviewer to "Comment" for the response to the rest of the review. We sincerely apologize for exceeding the character limit but we eagerly hope to thoroughly address the reviewer's concerns.

We sincerely thank the reviewer for their valuable advice and comments.

The comparison to the related work is incomplete. The authors do not cite the work [1] (although it was published approximately at the same time as [2]), where the authors propose a 1st-order sampling method, SLIPS, based on stochastic localization, that can be linked to a denoising diffusion model. In SLIPS, samples from the target distribution are obtained by following a diffusion process with intractable drift, which is nothing less than an expectation of a posterior distribution; in practice, the authors estimate this term by MCMC method (Langevin). In contrast to RDMC or ZOD-MC, SLIPS may exhibit multiple advantages: (i) it can be applied to a flexible class of stochastic localization schemes, whereas RDMC and ZOD-MC only consider the standard Ornstein-Uhlenbeck diffusion process, (ii) it provides an exact finite-time setting for sampling (which is not the case for RDMC or ZOD-MC), (iii) it is shown to scale well in dimension in their experiments (up to dimension 100), (iv) it has a small oracle complexity per score evaluation compared to the numerics presented here (32 vs 200 at least in the numerics of the current paper). Since the current approach has the same spirit as SLIPS (diffusion-based approach with MC estimation), it should be compared with it. Moreover, no numerical comparison with AIS or SMC (which are considered gold-standard to sample from low-dimensional multi-modal distributions) is provided. For me, this justifies why this paper does not meet the standards of NeurIPS.

We apologize for not having cited the very related work [1]. We were not aware of it at the time our manuscript was submitted.

Similar to ZOD-MC, SLIPS is also based on the denoising diffusion model and an MCMC score estimation. We have added a comparison between ZOD-MC and SLIPS in our updated manuscript. While SLIPS is great work, we hope to summarize three major differences between ZOD-MC and SLIPS: (1) SLIPS relies on MALA to approximate the score. As a result, to analytically establish convergence of SLIPS requires log-concavity outside a ball assumption due to the difficulty in analyzing MALA. The goal of our paper is to investigate a non-logconcave sampling algorithm without any convexity-related assumption. Therefore, we choose to approximate the score via rejection sampler, which can be analyzed under mild smoothness conditions. (2) SLIPS uses MALA to generate initial point: the denoising process is simulated from the middle of the observation process. This adds extra difficulty on initialization: the initialization error is hard to control numerically and analytically in general. In contrast, the initialization error in ZOD-MC can be controlled by starting at a large time $T$ in the forward process. (3) ZOD-MC can sample from non-differentiable and even discontinuous densities.

Numerically, we have updated our experiments to include comparisons with SLIPS, AIS and SMC. We started by running these methods under the same set up as in the experiment in Figure 1. For SLIPS we used the same initialization as used in their code ($\mathcal{N}(0,5)$ + LMC steps) and initialize AIS and SMC with $\mathcal{N}(0,5)$. Under this set up we find that all of the methods perform quite well, despite this ZOD-MC still has the best performance out of all methods. We further demonstrate that our method is less sensitive to the initial condition by initializing the methods with $\mathcal{N}(0,1)$. We then show that these methods still suffer from metastability and are unable to sample from all modes even at different oracle complexities.

In my opinion, the proposed approach misses a fundamental point, which lies in the MC estimation of the denoiser (i.e. the conditional expectation over the posterior distribution): it is linked to the ability to be able to provide an accurate MC estimate near at $t_0=0$, equivalently $t=T$, namely when $p_t$ is fully noised. When $t\approx T$ , it is clear that the corresponding posterior distribution is approximately equal to the target distribution itself (since we have full noise, the quadratic term vanishes). Therefore, applying rejection sampling on the posterior distribution at the very beginning of the sampling process in ZOD-MC is equivalent to do rejection sampling on the target distribution itself! Hence, in practice, the whole diffusion-based approach presented in ZOD-MC to sample from the target distribution via an annealing scheme is actually as hard as directly sampling from the target distribution, which hurts its usefulness. This remark is intimately linked to the notion of duality of log-concavity presented in [1], which explains that it is required to start the diffusion process further such that the posterior distribution is "smoother", and easier to sample from. On the theoretical side, I have a concerns that is linked to my point: it seems to me that the authors make an assumption on the MC error $\delta(t)$ to estimate the denoiser (which is not explicited in the main) in order to derive the total complexity cost (hidden in Section B.5 in the appendix); as said above, this control is actually crucial in practice and it is not straightforward to verify an upper bound on it.

Apology for a serious misunderstanding of our main results (Proposition 3.1, Theorem 1 and Corollary 3.1), but we're afraid the comment "the whole diffusion-based approach presented in ZOD-MC to sample from the target distribution via an annealing scheme is actually as hard as directly sampling from the target distribution" is not true.

We would like to kindly direct the reviewer to "Comment" for the response to the rest of the review. We sincerely apologize for exceeding the character limit but we eagerly hope to thoroughly address the reviewer's concerns.

Thank you for your kind response. Given all the elements discussed above, I have raised my score from 3 to 5.

Note: I have edited my previous comment to add a last remark on the use of 1st/2nd order methods, but it does not change my final decision.

We sincerely thank the reviewer again, for continuously engaging with us and helpful suggestions!

Thank you once again for the precision, it makes more sense to me ! I am still a bit confused about the gap between the initialization with variance $5$ (given in the rebuttal) and the initialization with variance $\sigma^2t_0=8.75$ (given in your last reply). Could you explain it ?

We apologize for creating a confusion fully due to us (it has been a very long 2 weeks) and thank you very much for catching it. It was a typo in our rebuttal that the variance of the initialization was stated to be $5$. What we really meant was $\sigma=5$. We did use $\sigma=5,t_0=0.35$, which led to an initialization with variance $8.75$. Results and values stated in our last reply remain correct, and we will still use $t_0=0.001$ and $\sigma=5$ in our revision, leading to $W_2=5.108$ for SLIPS.

• I think that both the authors and I agree that SLIPS should be added to the related work and to the numerical experiments with precise tuning of its critical hyperparameter (although it does not seem obvious as explained above). I think we also agree that SLIPS and ZOD-MC have complementary strengths and weaknesses (they could be compared in the manuscript), and none of them should be said to be the best in an absolute way. I think we finally agree that AIS and SMC should be added to the competing algorithms, once again with careful tuning of their hyperparameters.
• I now understand better the low importance of the score estimation for large $t$ (which I did not find so well emphasized in the submission) and apologize to the authors if I was, at first sight, misconfused about this part. I really think that the authors should add remarks on this specific part to avoid confusion among the readers.
• On the methoddological side: I encourage the authors to include in the revised version of the manuscript the guidelines detailed in their rebuttal and their following comments (and for instance, also include the ablation studies presented above): namely, the threshold on the number of rejections, the values of the hyperparameters...

We appreciate the reviewer for summarizing the fruitful discussion, and agree with the summary. We truly believe that adding the details discussed in the rebuttal (the comparison to SLIPS, AIS and SMC, more remarks on ZOD-MC and the guidelines) will make the ZOD-MC work more comprehensive and improve the clarity and accessibility to broader audience. We will incorporate these details in the updated manuscript.

Thank you for giving the practical guidelines to set the hyperparameters of ZOD-MC in practice. Once again, I think it is of high interest to enable the use of this algorithm by practitioners.

Overall, I really want to thank again the authors for giving precise answers to my several questions, and bringing lots of additional details on the way ZOD-MC works. I would like to sum up those points:

• I think that both the authors and I agree that SLIPS should be added to the related work and to the numerical experiments with precise tuning of its critical hyperparameter (although it does not seem obvious as explained above). I think we also agree that SLIPS and ZOD-MC have complementary strengths and weaknesses (they could be compared in the manuscript), and none of them should be said to be the best in an absolute way. I think we finally agree that AIS and SMC should be added to the competing algorithms, once again with careful tuning of their hyperparameters.
• I now understand better the low importance of the score estimation for large $t$ (which I did not find so well emphasized in the submission) and apologize to the authors if I was, at first sight, misconfused about this part. I really think that the authors should add remarks on this specific part to avoid confusion among the readers.
• On the methoddological side: I encourage the authors to include in the revised version of the manuscript the guidelines detailed in their rebuttal and their following comments (and for instance, also include the ablation studies presented above): namely, the threshold on the number of rejections, the values of the hyperparameters...
• Last but not least: put more emphasis on the fact that ZOD-MC requires the use of 1st/2nd order methods to find the modes of the target distribution in order to have an adapted proposal distribution in the rejection sampling phase (without it, ZOD-MC could not work, it is crucial). This setting is not problematic (it is quite a standard assumption in a variety of sampling methods), but as such, the title and the abstract of the submission suggest that the proposed methodology is purely zero-th order...

I will wait the reply of the authors to my last comments to decide the change of my score. Thank you once again for the fruitful discussion !

We sincerely thank the reviewer for their valuable advice and comments.

Some related work is missed in the authors’ survey, e.g., [1]. Theorem 7 of [1] provides a similar zeroth-order complexity O(exp(d)log(1/\varepsilon)) for achieving a minimal optimal error. I suggest the authors compare the theoretical results with [1].

Thanks for mentioning this very interesting work. It is related to the zeroth-order sampling problem that we considered, and we will add the following discussion in our updated manuscript:

[1] provides upper bounds on convergence rates for a zeroth-order sampling algorithm that combines an approximation technique and rejection sampling. We compare their results to ours from the following two perspectives:

1. target distributions: the target distributions in [1] are restricted to have compact support and to be m-differentiable ($m>0$). Our considered target distributions are more general. For example, they can have full support and they can even have discontinuities as shown in our Section 4 Figure 5.

2. complexity bounds: we can compare our derived complexity upper bound to the results in [1] for target distributions in the restricted class considered in [1]. By using the fact that KL is smaller than $R_\infty$, theorem 12 in [1] can be written as an upper bound on the complexity to reach $\varepsilon$-accuracy in $R_\infty$: the complexity is of order $\Omega_{d}(\varepsilon^{-d/m})$. The $d$-dependency is implicit, therefore we only compare the $\varepsilon$-dependency. The $\varepsilon$-dependency in both complexities is polynomial in $\varepsilon$, with a linear d factor in the exponent. Since we consider a general class of targets, our result doesn't reflect the smoothness of the potential as their result does. Within the restricted class, if the target is differentiable up to an order greater than 2, their result is better than ours (in terms of $\varepsilon$-dependency). Otherwise, our result is better.

The implementation of Algorithm 3 requires the minimum of $V$, which is denoted as $V_*$. How do we obtain this minimum? Since the function is highly irregular, calculating the minimum may also be difficult.

This is an oracle that ZOD-MC relies on. What we need can be relaxed to a global lower bound of $V$. Nevertheless, we agree this oracle can still be challenging to implement. Our practical implementation was mentioned in Remark 1, where we get an approximation using Newton's method (although other methods like Gradient Descent or the proximal bundle can also provide good approximations). As the sampling process explores the space more, the estimated minimum is updated as necessary.

Although the experiments seem to be comprehensive, some details are not reported. For example, what is the meaning of oracle complexity? Does it include only zeroth-order oracle, first-order oracle, or both? If only counting the first-order oracle, is the comparison fair to first-order methods? What results are if the x-axis is set as the wall clock time? What are the hyper-parameters of ZOD-MC and baselines, and how do we choose these hyper-parameters? I hope these details will be included in the appendix.

Thanks for a great question. Our claim is at least fair and in fact an understatement -- As mentioned in line 274 we count the total number of first and zeroth-order queries. We count each of them as being the same. This actually puts our method at a disadvantage since evaluating first-order queries is more expensive than zeroth-order queries. Despite this, we still see improved counts. The only exception to this is in the dimension experiment (Figure 1b in the main paper) where we matched number of function evaluations. For instance, a first-order query requires $d$ evaluations. We include a figure displaying the clock time as a function of the oracle complexity under the same set up as that of Figure 1a in the main paper. As introduced at the beginning of Section 4, the baselines we considered includes RDMC, RSDMC, Proximal sampler, Parallel Tempering and unadjusted Langevin Monte Carlo. The hyperparameters N,T and the step-size are analytically explored and we chose them according to Corollary 3.1. We will highlight these information in the updated manuscripts.

In Corollary 3.1, we can find the zeroth-order query complexity will depend on $\varepsilon_{\text{KL}}^{-d}$. But I would be curious to know if we can achieve a complexity whose order of $\varepsilon$ is independent with $d$? What are the barriers to obtaining this complexity?

The exponent $d$ in the $\varepsilon$-dependency is due to the curse of dimensionality in the rejection sampler: when we use the current version of rejection sampler to sampler $p_{0|t}$, we are only able to prove that the expected number of rejections is of order $O(\exp(2dt))$. Since the largest $t$ in the denoising process is of order $\log(d/\varepsilon_{\text{KL}})$, we end up with the $\varepsilon_{\text{KL}}^{-d}$ complexity. To improve the $\varepsilon$-dependency requires a finer design of the rejection sampler with higher acceptance rate. We believe this is an interesting future work, and theoretically it could be done by considering non-logconcave targets with special structures.

All these scientific discussions are helpful and greatly appreciated. At this moment our method's strength still lies in being agnostic to multimodality in low (e.g., $\lesssim$ 10) dimensions.

[1]: Holzmüller, David, and Francis Bach. "Convergence rates for non-log-concave sampling and log-partition estimation." arXiv preprint arXiv:2303.03237 (2023).

We sincerely thank the reviewer for their valuable advice and comments and greatly appreciate the positive evaluation.

One possible drawback of the proposed sampler, just as the authors stated in the paper, is that its iteration complexity depends linearly on the data dimension. Hence, it will probably be appealing only for low-dimensional distributions. Therefore, it might be meaningful to do a numerical investigation to see how the sampler proposed in this paper behaves on high dimensional distributions.

We thank the reviewer for this suggestion. In Figure 1b of the main paper, numerically we've shown that the sample quality of ZOD-MC decreases a little as the dimension increases from 1 to 7. But ZOD-MC still generates relative high-quality samples and outperforms RDMC and RSDMC at a fixed number of function evaluations.

For the sake of completeness, it might be necessary for the authors to include a section of related work on derivative-free methods for optimization and sampling from unnormalized densities. Some classical examples include the Ensemble Kalman Filter (EnKF) and Ensemble Kalman Inversion (EKI).

Thanks for the suggestion. The Ensemble Kalman Filter (EnKF), Ensemble Kalman Inversion (EKI) and the Ensemble Kalman Sampler (EKS) are all derivative-free sampling algorithms that are based on moving a set of easy-to-sample particle according to certain dynamics. However, these methods seem to be mainly for data assimilation / Bayesian posterior sampling, which require (noisy) observations from the target distributions. That is different from our setting: sampling using queries on the target potential functions only. That's why we did not compare to them.

However, we will add a discussion on these important related algorithms.

The reviewer's main concern is that when using rejection sampling to implement the Restricted Gaussian Oracle (RGO) in the Alternating Sampling Framework (ASF), one can only obtain the expected number of executions for rejection sampling (Proposition B.3). Therefore, it might be fair to compare the expected time complexity of this sampler with the exact time complexity of other samplers listed in Table 1?

We appreciate the question on the comparison in Table 1. We used zeroth-order oracle complexities in expectation for ZOD-MC in Table 1. We believe the comparison is fair due to the following two reasons:

(1) lots of tasks, such as volume computing and Bayesian inference, require a large number of samples from the target distribution. Hence, we need to run the sampling algorithm many times. As the number of algorithm running gets large, the total oracle complexity will be approximately the total expected oracle complexity. From these perspective, the comparison in table 1 is fair. Similar comparison has been made in [1] to compare the complexity of the proximal sampler to the complexities of LMC and MALA.

(2) it is worth mentioning the complexities of RDMC and RSDMC are not exact either. The complexties in RDMC (Theorem 1) and RSDMC (Theorem 4.1) are proved only to ensure an $\varepsilon$-accurate sample with high probability, which means the complexities to obtain an $\varepsilon$-accurate sample almost surely for RDMC and RSDMC should be larger than the corresponding complexities presented in table 1.

[1] Chen, Yongxin, et al. "Improved analysis for a proximal algorithm for sampling." Conference on Learning Theory. PMLR, 2022.