Paper: Chain of Log-Concave Markov Chains (ICLR-2024)

Reviewer comment (Theory):
One thing that is missing from this paper is the final step of computing the bayes estimator $\mathbb{E}[X|y_{1:m}]$ to obtain samples from the true distribution. For that, it is needed to estimate the score function $\nabla \log p(y_{1:m})$ of the smoothed distribution. The authors acknowledge that this is a separate issue that needs to be addressed. I understand this might be beyond the scope of the current work, which is more focused on showing the benefits of smoothing. However, ultimately estimating the score function is a crucial step for this to be an end-to-end approach for sampling and it is not clear what is the difficulty of this problem compared to the original sampling problem (is it strictly easier in some sense?). In my opinion, this would make clear what the benefit of this approach over standard MCMC algorithms is.

Author response:
We have expanded Sec. 4.2.1 on the important problem of estimating $g(y; \sigma)$, however due to space we moved the details to the appendix (Appendix F in the revised manuscript). We have also added more experiments to Appendix I using Langevin MCMC for estimating the smoothed score functions.  The question the reviewer has raised on the difficulty of estimating $g(y ; \sigma)$ is very important (we had a brief discussion of it in Sec. 6) and theoretical analysis for this aspect of our algorithm is a great research direction.

Reviewer comment (Theory):
Related to a previous point about estimating the score function, how exactly do the authors arrive at the expression (4.8) for estimating the score function? In appendix B, a number of properties is proven about the score function, among which that it is related to $\mathbb{E}[X|y_{1:m}]$. However, I don't see how this is captured in the empirical definition of (4.8). What is the meaning of taking a weighted average of the $\epsilon_i$s?

Author response:
We added a new section to the appendix, Appendix F: Estimators for $\nabla \log p(y)$, giving the derivations for the estimators (F.1) (the same as (4.8) in our original submission) and (F.2) for the Langevin MCMC estimator. In that section, we also give the algorithm (Algorithm 3) for estimating $g(y; \sigma)$ using (overdamped) Langevin MCMC.

Reviewer comment (Theory):
Related to the previous question, if we care about the expected hessian of the conditional sampling, I'm wondering whether it's possible to have a version of Theorem 1 showing that the expected condition number of the conditional sampling problem converges to $1$ with rate $1/m$ without assuming that the original distribution contains some amount of Gaussian noise. The reason is the following: the proof for all the claims about the "increase" in log-concavity basically hinge on equation (D.2), which asserts that $\nabla^2 \log p(y_{1:m}) = -\sigma^{-2} I + \sigma^{-4} Cov(X|y_{1:m})$. The argument then proceeds by the observation that conditioning on more and more observations can only reduce the variance of $X$ in expectation, so the "positive" part on the right hand side can only become smaller as $m$ increases. The proof of Theorem 1 expresses $Cov(X|y_{1:m})$ as $O(1/ m) + O(1/m^2) Cov(Z|y_{1:m})$ and bounds the latter covariance by a constant, since $Z$ is supported on a bounded set (as the footnote in page 18 suggests). However, it seems conceivable to me that if we do not impose the assumption $X = Z + N_0$, we could instead show directly that $\mathbb{E}[Cov(X|y_{1:m})]$ decreases as $O(1/m)$ if $X$ is bounded. The reason I believe this is true is that, by Proposition 2, we know that the posterior of $X$ conditioned on $y_{1:m}$ converges in distribution to $X$ at a rate of $1/m$. Thus, the conditional variance of $X$ given $y_{1:m}$ should decrease as $m$ becomes larger. Another intuitive way to see this is that we could for example simply take the average $\overline{y_m}$, which is $O(\sigma/m)$ close to $X$. Thus, we gain more and more information about $X$ as $m$ increases, which suggests that $\mathbb{E}[Cov(X|y_{1:m})] = O(1/m)$ always. Have the authors considered this more general statement?

Author response:
We did consider strengthening the theorem along the lines of the reviewer's comments, in particular whether $1/m$ rate still holds when $\tau=0$. Please note that in Theorem 1 (as opposed to Lemma 2) we are interested in the conditional Hessian itself, not the expected conditional Hessian. The main challenge here is that Theorem 1 is currently too restrictive (almost surely) and we think a great future direction is to see under what conditions the $1/m$ rate holds with *high probability*. This discussion is also related to our Remark 3 in the paper on the hardness of strenghtening  Theorem 1 regarding the monotonicity (4.7). We appreciate the reviewer sharing their thoughts on this aspect of the problem: although the analysis of the expectation of the Hessian is interesting in its own terms, it cannot be directly used to provide some guarantees on progressively more log-concave sampling.

Reviewer comment (Theory):
One weakness is in the result presentation: The paper does not provide a direct performance guarantee on Algorithm 1. I understand the inner-loop method is the object of study here. Still, it should be feasible to make some assumptions around its properties and derive a theorem that captures the algorithm's performance as $m$ and $n_t$ increase. The error metric could ideally be a Wasserstain-type distance for consistency with the experiment section.

Author response:
That direct performance guarantee on Algorithm 1 is an important problem to tackle in understanding and extending the formalism developed in this work. We hope the reviewer reconsider this as a weakness if it has affected their score. The paper is already dense with theoretical results and we had limited space for supporting the theoretical formalism with experiments on demonstrating the effectiveness of smoothing and the log-concave sampling strategy which the reviewer has kindly highlighted.

Reviewer comment (Experiments):
The experiment evaluation partially supports and reflects the theoretical claims, but the settings could be more complex for a thorough evaluation. For example, it would be interesting to understand what magnitude the # sampling iterations must be for some rather complex distributions, e.g., $k$-mixture of Gaussians where $k$ increases. The densities used in Section 5 are well-designed to test specific hypotheses/claims. Yet, experiments should probably cover more than that, such as demonstrating the properties of the main algorithm(s), which can be critical for method adoption.

Author response:
The experiment on $k$-mixture of Gaussians as $k$ and $d$ increase is very interesting, but requires careful considerations in experimental design... We are currently working on extending our work to the Ising model on the Boolean hypercube {$-1,1$}$^d$, where, after smoothing, one can effectively think of as $k$-mixture of Gaussians with $k=2^d$ mixtures.

Reviewer comment (Presentation):
May I know the purpose of "4.1 Example" as it seems to be just deriving various quantities for a concrete/standard Gaussian mixture? I don't see how it helps the readers better understand Theorem 1.

Author response:
Thank you for highlighting this. The proof of Theorem 1 is somewhat involved and we found it useful to give an example where the reader can easily follow all the derivations. We also like Example 4.1 since it points to some stronger theoretical results than what we could prove in this paper, e.g., the upper bound in Theorem 1 becomes tight in the example we considered.

Reviewer comment (Theory):
In Section 5.2, is there an intuitive explanation for why "the optimal $\sigma$ here is in fact larger than the noise level needed to make $p(y_1)$ log-concave"?

Author response:
Our intuitive explanation is that larger $\sigma$ leads to better condition numbers which seems to help beyond the log-concavity threshold. This is however dependent on the MCMC$_\sigma$ used in Algorithm 1.

Reviewer comment (Reproducibility):
Is the actual code anonymously shared somewhere? It would be good to have it for result reproducibility and smooth adoption of the method.

Author response:
We have shared the code for this submission [at this anonymous link](https://anonymous.4open.science/r/solaris-1FC4) and will publicly release the code upon acceptance of the paper.

Reviewer comment (Writing):
The paper could have done a better job at exposition of not so well-known notions like walk-jump sampling. The current exposition is good enough for someone intimately familiar with the literature, but for others, like me, it takes some reading of cited papers to not see ideas as being "pulled out of a hat".

Author response:
We appreciate the time you took on the literature on empirical Bayes and walk-jump sampling. We will take the reviewer's comment into consideration in follow-up work.

Reviewer comment (Theory):
How to handle the case when the desired probability measure does not have a density?

Author response:
1. We are not sure if the reviewer is referring to discrete distributions here. We are currenly working on extending our work to sampling from the Ising model where the random variable takes values in the Boolean hypercube \{$-1,1$\}$^d$. In our formalism, after smoothing sampling becomes equivalent to sampling from a (smoothed) k-mixture of Gaussians with $k=2^d$ mixtures.

Reviewer comment (Experiments):
Score estimation is essential for the algorithm's efficiency. However, most of the experiments in the paper use the analytic score function, which is not realistic in practice. Moreover, the score has to be estimated at each step of ${\rm MCMC}_\sigma$, which will significantly increase the algorithm's complexity.

Author response:
We agree with the reviewer that estimating the score function using Langevin MCMC is better than using the estimator we experimented with for the results in the appendix. Following the submission of the paper we did conduct experiments where we used Langevin MCMC for estimating the smoothed score functions. We modified Sec. 4.2.1 and added a new section to the appendix (Appendix F in the revised manuscript). We have also updated the results in Appendix I with the new experiments.

We consider the new experiments added to the appendix (Fig. 7) on running Algorithm 1 using score estimation with Langevin MCMC as preliminary as we did not conduct exhaustive tuning of the underdamped Langevin MCMC used for the score estimation experiments.

Reviewer comment (Experiments):
The paper only compares the performance of accurate and approximate scores in appendix H, and the result is not very satisfactory. The estimated score performs well when $d=2$, but it deteriorates when $d=8$. These results suggest that the score estimate's precision could be much worse in high dimensions, which I believe is the paper's main weakness.

Author response:
We agree with the reviewer that estimating the score function using Langevin MCMC is better than using the estimator we experimented with for the results in the appendix. Following the submission of the paper we did conduct experiments where we used Langevin MCMC for estimating the smoothed score functions. We modified Sec. 4.2.1 and added a new section to the appendix (Appendix F in the revised manuscript). We have also updated the results in Appendix I with the new experiments.

We consider the new experiments added to the appendix (Fig. 7) on running Algorithm 1 using score estimation with Langevin MCMC as preliminary as we did not conduct exhaustive tuning of the underdamped Langevin MCMC used for the score estimation experiments.

Reviewer comment (Experiments):
In 4.2.1, the proposed method for score estimation is based on importance sampling (which transforms the expectation w.r.t. the probability density $\nu\propto\exp(-f-\frac{1}{2\sigma^2}||\cdot-y||^2)$ to $\mu=\mathcal{N}(y,\sigma^2I)$). However, (Chatterjee and Diaconis, 2018) has shown that the sample size needed for accurate estimation by importance sampling is usually $\exp({\rm KL}(\nu||\mu))$ for general target functions. For this reason, (Huang et al, 2023, section 3) use Langevin Monte Carlo to sample from $\nu$ and combine it with importance sampling in a similar task of score estimation. The experiments in your paper indicate that the score estimator also suffers from high variance. The authors should experiment with different score estimators, instead of saying that "studying the covariance of the plug-in estimator and devising better estimators is beyond the scope of this paper".

Author response:
We agree with the reviewer that estimating the score function using Langevin MCMC is better than using the estimator we experimented with for the results in the appendix. Following the submission of the paper we did conduct experiments where we used Langevin MCMC for estimating the smoothed score functions. We modified Sec. 4.2.1 and added a new section to the appendix (Appendix F in the revised manuscript). We have also updated the results in Appendix I with the new experiments.

Thank you for pointing out the reference (Chatterjee and Diaconis, 2018). This problem is also related to the literature on "Entropy SGD", e.g., (Chaudhari et al., 2017), where the estimating smoothed score functions (although it's not framed like that) is used in the context of optimizing neural network parameters. (Chatterjee and Diaconis, 2018) and (Chaudhari et al., 2017) are now cited -- please see Appendix F. (Huang et al, 2023) is a concurrent work, which we were not aware of. We would be happy to cite it in the camera-ready after a careful reading of the paper.

We consider the new experiments added to the appendix (Fig. 7) on running Algorithm 1 using score estimation with Langevin MCMC as preliminary as we did not conduct exhaustive tuning of the underdamped Langevin MCMC used for the score estimation experiments.