Sampling by averaging: A multiscale approach to score estimation

Thank you for your thorough review and helpful comments. We also greatly appreciate your positive feedback.

We reply to your questions below, please let us know if anything is unclear.

As pointed out in the paper, the multiscale approach avoids the need for a nested algorithm. However, a nested algorithm that does not introduce SDE with different time scales is a natural competitor for the introduced multiscale methods. In such a nested approach, samples ~ $\rho$ in (2) would be drawn in the inner-loop by simulating the corresponding SDE. While some of the methods in the experimental section rely on such a nested approach it is not clear whether any of those corresponds to a variant of the introduced methods where the multiscale SDE simulation is replaced with a nested SDE simulation.

Reverse Diffusion Monte Carlo (RDMC) proposed by Huang et al. (2024) [1] corresponds to the framework that you describe when the averaged process is given by the reverse diffusion process (that is, the setting described in Section 3.2 of our work). We include comparisons of our algorithms with RDMC in Tables 1 and 2.

However, RDMC does not explore the cases where the averaged process is given by the time inhomogeneous Langevin diffusion (i.e., Equation (3) of our work) either for Gaussian diffusion paths, where the base distribution $\nu$ is a Normal distribution, or for heavy-tailed diffusions with a Student’s t base distribution.

The nested inner-outer loop algorithm that you suggest based on annealed (time-inhomogeneous) Langevin diffusion with a Gaussian base distribution will generally lead to a worse performance compared to Reverse Diffusion Monte Carlo. This is due to the bias inherent in annealed Langevin dynamics, as discussed in Section 4.2 of our work. Moreover, for completeness we show in the table below, that the nested inner-outer loop algorithm based on annealed Langevin diffusion with a heavy-tailed base distribution underperforms our multiscale approach for the mixture of Student's t distributions benchmarks, when the number of energy evaluations is set to be the same for all methods.

Table 1. Comparison of our MultDALMC method with a nested inner–outer loop algorithm, both using a heavy-tailed base distribution. Results are averaged over 30 random seeds. We evaluate two variants: one based on the modified Itô diffusion process (Equation (7) in our work) and another using an overdamped Langevin diffusion for the fast process and inner loop. Performance is measured by the entropy-regularised Wasserstein-2 distance.

How is it ensured that hyperparameter optimization was carried out with comparable computational resources for the compared sampling methods? What were the hyperparameter grids used for each method and benchmark?

Below, we provide the hyperparameter grids used for our proposed algorithms as well as for all baseline methods. The design of these grids was carefully guided by the original papers introducing each method. We allocated sufficient computational resources to all baselines to ensure that the selected combination of hyperparameters was optimal, some methods required more extensive tuning than others. More specifically, each hyperparameter configuration was evaluated over 30 independent random seeds to account for variability and ensure robustness. When a method exhibited high sensitivity to certain hyperparameters, we expanded the grid around high-performing regions for finer exploration.

Besides, it is important to remark that to ensure a fair comparison we used the same initialisation based on a standard Gaussian distribution for our methods and all baselines in Tables 1 and 2. However, this initialisation is not optimal for the SLIPS algorithm [2], which benefits from a more tailored initialisation. For completeness, we also report the performance of SLIPS with optimal initialisation in Table 7, Appendix D. This optimal initialisation assumes knowledge of the distribution’s scalar variance which is unavailable in practical scenarios.

Hyperparameter grids for the different algorithms

• SMC and AIS: The hyperparameter we tune is the number of MCMC steps $n_{\text{MCMC}}$, its value is selected from the interval $[20, 100]$ using a grid step of 4.
• ULMC: The step size is chosen within the following grid $h\in$0.001,\dots, 0.009, 0.01, \dots, 0.09, 0.10, 0.11, \dots, 0.20. The friction coefficients $\Gamma_{\min}$ and $\Gamma_{\max}$ are chosen within the grid 0.001, 0.005, 0.01, , \dots, 0.09, 0.1,\dots, 0.9, 1, 1.5, \dots, 10$$. The mass matrix is set to the identity.
• PT: The grid for the step size is the same as that of ULMC. The number of chains is selected from the interval $[1, 20]$ with a grid step of 2. The corresponding temperatures are chosen to be equally spaced on a logarithmic scale, with the minimum temperature fixed at 1 and the maximum temperature chosen from the grid $10^2, 10^3, 10^4, 10^5, 10^6, 10^7$.
• DiGS: The grid for the step size is the same as that of ULMC. The number of noise levels $K$ is selected within the interval $[1, 20]$ with a grid step of 2. The maximum and minimum noise levels, $\alpha_1$ and $\alpha_K$, respectively, are chosen from the grid $0.05, \dots, 0.09, 0.1, 0.2, \dots, 1$. Finally, the number of Gibbs sweeps $n_{\text{Gibbs}}$ is selected from the interval $[50, 500]$ using a grid step of 50.
• RDMC: It has a single hyperparameter $T$ which we search within the grid $T \in$0.001, 0.005, 0.01, 0.05, 0.10, 0.15, 0.20, 0.25, 0.30, 0.35, 0.40, 0.45, 0.5$$ as suggested in [1].
• SLIPS: Following [2], the number of MCMC steps is selected from the interval $[20, 100]$ with a grid step of 2. The value for $t_0$ is chosen by approximate equal log-SNR spacing in $[-5.0, 0.0]$. The value of $\eta$ is selected from the interval $[4,6]$ using a grid step of 0.2.
• Ours: The grid for the step size is the same as that of ULMC. The scale separation parameter $\varepsilon$ is chosen within the grid $0.001, 0.005, 0.01, \dots, 0.09, 0.10, 0.11, \dots ,0.20$. The friction coefficients $\Gamma_{\min}$ and $\Gamma_{\max}$ are selected within the grid $0.001, 0.005, 0.01, , \dots, 0.09, 0.1,\dots, 0.9, 1, 1.5, \dots, 10$. The mass matrix is set to the identity. The parameter $\tilde\lambda$ is chosen from the interval $[0.3, 0.7]$ using a grid step of 0.05.

We look forward to engaging further during the rebuttal period to address any remaining concerns you may have. Thanks again!

[1] Huang, X. et al. (2024) Reverse Diffusion Monte Carlo. The Twelfth International Conference on Learning Representations.

[2] Grenioux, L. et al. (2024) Stochastic Localization via Iterative Posterior Sampling. In Proceedings of the 41st International Conference on Machine Learning.

Thanks for this insightful and detailed response. I would ask the authors to include these detailed grids for the hyperparameter optimization in the appendix of the paper. I have no further questions for the authors.

Thank you for your thoughtful review and insightful comments. We also greatly appreciate your positive feedback on our work.

We reply to your concerns below, please let us know if anything is unclear.

I'd like to know the results on real-world data, at leaset like mnist or cifar10 because I find the performance on high-dimensional and complex data (table 2) doesn't has consistent improvement.

We would like to clarify that our algorithms are designed for sampling tasks where we have direct access to a complex unnormalised target distribution and its gradient, and the goal is to generate samples from it. This contrasts with typical generative modelling tasks where the objective is to generate new samples from an unknown data distribution, given only access to a dataset of existing samples. To make this distinction clearer, we provide the following table highlighting the key differences between the two tasks.

Therefore, direct sampling from MNIST or CIFAR10, given access to the dataset of images, does not fit naturally into our framework. To apply our methods to these datasets, one would first need to train an energy-based model [1,2] to learn the unnormalised target distribution corresponding to the dataset and then use our proposed method to draw new samples from the learnt distribution. However, in such a setup, the quality of the generated samples would depend not only on the performance of our sampler but also on how well the energy-based model captures the true data distribution. This would make it difficult to isolate and evaluate the effectiveness of the sampling algorithm itself.

That said, we have evaluated our methods on a wide range of established benchmarks for sampling algorithms [3,4,5] (some of them arising from real-world applications such as Bayesian inference) and included a real-world example involving a statistical physics model which has a dimension of 100 (see Figure 1a).

Lack of comparison of computational cost.

To ensure a fair comparison across different algorithms we use the same number of energy evaluations, as detailed in Table 3 (Appendix D.3). Thus, the computational cost is kept the same for all compared methods.

Additionally, in Appendix D.4, we report the computational resources used, along with the computation times (averaged over 30 random seeds) for each algorithm on the 40-component mixture of Gaussians benchmark in 50 dimensions (see Figure 2).

To further analyse the computational cost, we provide below a comparison in terms of energy evaluations and computation times required by each algorithm on the 40-MoG benchmark (dimension 50) to achieve comparable performance, as measured by the entropy-regularised Wasserstein-2 distance. Results are averaged over 30 random seeds.

Table 1. Comparison of the computational cost across methods on the 40-MoG benchmark in 50-dimensions. Results are averaged over 30 random seeds.

We note that our proposed algorithms consistently require fewer energy evaluations to achieve comparable performance than alternative methods. We plan to include this discussion in an updated version of the paper.

We hope that this addresses your concerns, and we sincerely request any further feedback that you may be able to provide us during the rebuttal period. Thank you for your continued feedback!

[1] Du, Y. et al. (2021) Improved Contrastive Divergence Training of Energy-Based Model. In Proceedings of the 38th International Conference on Machine Learning.

[2] Song, Y. and Kingma, D. (2021) How to Train Your Energy-Based Models. arXiv:2101.03288.

[3] Grenioux, L. et al. (2024) Stochastic Localization via Iterative Posterior Sampling. In Proceedings of the 41st International Conference on Machine Learning.

[4] Vargas, F. et al. (2024) Transport meets Variational Inference: Controlled Monte Carlo Diffusions. The Twelfth International Conference on Learning Representations.

[5] Richter, L. and Berner, J. (2024) Improved sampling via learned diffusions. The Twelfth International Conference on Learning Representations.

Thanks a lot for your thoughtful review and helpful comments. We also greatly appreciate your positive remarks on our work.

The practical viability of the approach to some real-world settings is not immediately clear given the constraints and assumptions placed within the theoretical analysis.

Our algorithm is specifically designed to sample from diffusion paths with the aim of targeting highly complex probability distributions in practice. In our experiments, we demonstrate that the proposed algorithms exhibit strong empirical performance in general settings beyond strong convexity. This includes a real-world application in statistical physics involving a 100-dimensional system (see Figure 1a), as well as benchmarks from Bayesian inference tasks such as Ionosphere and Sonar.

We acknowledge that the assumptions made in the theoretical section are arguably strong. Our objective in providing this theoretical guarantee was to demonstrate that the method is principled under suitable conditions. Moreover, the analysis serves as a template that can be extended to more general assumptions in future work.

Analysing multiscale systems under weaker assumptions presents significant technical challenges and constitutes a substantial line of research in its own right. For this reason, it falls outside the scope of the present work. Nevertheless, we hope that the connections we establish with the averaging literature will encourage applied mathematicians working on multiscale dynamics to engage with this problem. We are genuinely excited and eager to see what theoretical developments our initial approach may inspire in future work.

How does the empirical performance change with the dimensionality of the underlying signal? Do the number of steps and scale separation parameter have to be tuned very carefully as the dimensionality changes?

As the dimension of the underlying signal (i.e., the $(Y_t)\_{t\geq 0}$ process), and consequently the dimension of the $(X_t)_{t\geq 0}$ process, increase, convergence to the target distribution becomes slower, requiring a larger number of time steps. However, convergence still occurs. This is illustrated in Table 9 (Appendix D.4), where we report performance metrics (specifically, the entropy-regularised Wasserstein-2 distance) alongside the number of energy evaluations (which is in one-to-one correspondence with the number of steps) and computations times for our sampling methods on the 40-MoG benchmark across different dimensions.

On the other hand, we observed that our algorithms demonstrate robustness and stability across a broad range of hyperparameter values, as shown in the table below. We plan to include a plot (extending the results of this table) on a finer hyperparameter grid in an updated version of the paper. These results suggest that while some tuning is beneficial, our methods do not require highly sensitive or exhaustive hyperparameter optimisation.

Table 1. Ablation results on the 40-component Mixture of Gaussians benchmark in 50 dimensions for the scale separation parameter $\varepsilon$ and the number of steps. Results are averaged over 30 random seeds. Notably, the performance remains stable across a wide range of hyperparameter values.

Table 2. Ablation results on the Sonar benchmark for the scale separation parameter $\varepsilon$ and the number of steps. Results are averaged over 30 random seeds. The performance metric is the average predictive posterior log-likelihood on a test dataset. Notably, the performance remains stable across a wide range of hyperparameter values.

Is there any intuition on how to choose the hyper parameters of the proposed algorithm carefully without too much experimental tuning? What happens when the conditional distribution mixing time is set too slow or too fast?

The main hyperparameters in our algorithms are the step size and scale separation parameter ($\varepsilon$). The intuition for choosing the step size is similar to that used in classical Langevin-based algorithms (see [1,2]).In particular, this parameter typically scales inversely with the dimension of the problem. Regarding the scale separation parameter, we empirically observe that values in the range 0.01-0.11 yield good performance across a wide range of problem settings (see Table 8, Appendix D.3, where we report the optimal value of $\varepsilon$ for each benchmark distribution).

If the conditional distribution's mixing time is set too slow, then the obtained samples from the process $Y_t$ will not correspond to that of the correct conditional distribution. This results in the drift term of the slow process ($X_t$) being poorly estimated. That is, the drift term will not accurately reflect the drift term corresponding to the true moving target, which in turn fails to capture the correct dynamics of the underlying moving target. Intuitively, this leads to a lag in tracking the moving target distribution, which may cause some modes of the final target distribution to be missed.

On the other hand, if the conditional distribution mixing time is too fast, the particles from the slow process $(X_t)$ may exhibit overly noisy or "wiggly" trajectories, hindering convergence over time.

We hope that the additional discussions and ablation studies have addressed the reviewer’s questions, and we look forward to receiving further feedback!

[1] Chewi, S. et al. (2024) Analysis of Langevin Monte Carlo from Poincare to Log-Sobolev. Foundations of Computational Mathematics.

[2] Milstein, G. N. (1994) Numerical integration of stochastic differential equations. Springer.

Thank you for your review and insightful comments. We also sincerely appreciate your positive remarks on the quality and clarity of our work.

We reply to your concerns and questions below, please let us know if anything is unclear.

This is my main concern for the paper. Namely, it is by now fairly well understood how to sample from strongly log-concave distributions given this sort of oracle access, e.g. the papers that the authors cite of Chen-Chewi-Salim-Wibisono and Lee-Shen-Tian. Moreover, these papers yield non-asymptotic rates, as opposed to the asymptotic convergence demonstrated here.

We are not entirely sure whether, in your comment, oracle refers to access to the gradients of the target potential, i.e. $\nabla \log \pi = \nabla V_\pi$, or instead to what Chen–Chewi–Salim–Wibisono and Lee–Shen–Tian define as a restricted Gaussian oracle (RGO). For completeness, we address both interpretations below.

1. Oracle = access to $\nabla\log\pi=\nabla V_\pi$: Our algorithms are designed for sampling tasks, where the goal is to obtain samples from a complex target distribution $\pi$. In this setting, it is assumed in the literature that we have access to $\log \pi(x)$ (up to an additive constant) and its gradient $\nabla \log \pi(x) = \nabla V_\pi(x)$. This contrasts with typical generative modelling tasks where the objective is to generate new samples from an unknown data distribution, given only access to a dataset of existing samples. To make this distinction clearer, we provide the following table highlighting the key differences between the two tasks.

Having said this, when $\pi$ is strongly log-concave, standard Langevin Monte Carlo (LMC) algorithms can be readily applied, with known convergence guarantees in this setting [1], without requiring the use of the more sophisticated frameworks of Chen–Chewi–Salim–Wibisono or Lee–Shen–Tian, which rely on access to an RGO. However, we would like to emphasise that, when developing our algorithms, our main focus is on complex multimodal distributions, where plain LMC and related methods typically struggle. To address this challenge, we employ methods based on annealed Langevin diffusions, which are better suited for these harder sampling problems.

2. Oracle as in Chen-Chewi-Salim-Wibisono and Lee-Shen-Tian: We would like to clarify that our setting does not assume access to an oracle in this sense. The works of Chen-Chewi-Salim-Wibisono and Lee-Shen-Tian rely on access to what they refer to as a restricted Gaussian oracle (RGO), which corresponds in our notation to exact sampling from the conditional distribution $Y\sim\rho_{t,X}(y)$ (or $\pi^{X|Y}$ in their notation). Some of their theoretical results explicitly assume this oracle access and therefore do not account for the complexity of sampling from $Y\sim\rho_{t,X}(y)$, which in our case corresponds to the fast process. Our theoretical analysis does include the complexity of this sampling step.

Besides, while both aforementioned works also provide theoretical results when implementing the oracle approximately, via a rejection sampling scheme (Corollary 6 in Chen-Chewi-Salim-Wibisono and Corollary 1 in Lee-Shen-Tian), these approaches assume access to the minimiser of the target potential—an assumption that can be impractical in higher-dimensional settings.

Despite the paper’s claims, I think it is unlikely that this approach will work for complex real-world settings (as it doesn’t learn a score function). Indeed, the evaluations in the paper are on simple, synthetic settings. Thus, I think the main motivation must be theoretical, but the results are then largely superceded by these previous works.

We would like to highlight that we do estimate the drift term of the averaged process $\bar X_t$, Equation (3), (which corresponds to the score function of the true moving target distribution $\mu_t$). This estimation is made possible via the auxiliary fast process $Y_t$, which enables us to approximate the drift of the marginal $X_t$ without direct score learning.

Besides, our algorithms demonstrate strong performance across a variety of established benchmarks in the sampling literature that go beyond the strongly log-concave setting (some of them arising from real-world applications such as Bayesian inference). In addition, we include a real-world example given by a statistical physics model which has a dimension of 100 (see Figure 1a). Importantly, this example is not only high-dimensional but also inherently challenging due to the free-energy barrier between the modes [2].

While the authors briefly justify why this method could be preferred to the other works, these concerns do not really hold for the strongly log-concave setting that their theoretical results hold for. If the authors could justify why the results are interesting in comparison to prior work specifically in this strongly log-concave setting, that would be very helpful.

As we acknowlegde in the discussion section, the strongly-convex assumption required for deriving theoretical guarantees is indeed restrictive. However, the motivation for providing partial theoretical guarantees was to demonstrate that our method is principled and to provide an initial template for more general theory. This does not mean that the proposed methods are restricted to the assumption used in the theoretical analysis. Developing a full theoretical framework for multiscale systems under weaker assumptions remains technically challenging and constitutes a substantial line of research in its own right. For this reason, it falls outside the scope of the present work. Nevertheless, we hope that the connections we establish with the averaging literature will encourage applied mathematicians working on multiscale dynamics to engage with this problem. We are genuinely excited and eager to see what theoretical developments our initial approach may inspire in future work.

Conclusion

We sincerely appreciate your recognition that our work is “quite nontrivial and above the bar for acceptance” (as noted in your Quality comment). We also fully acknowledge your concerns regarding the scope of our theoretical results, which are currently confined to the strongly log-concave and asymptotic settings.

That said, we would like to emphasise that while our proof-of-concept theoretical guarantees focus on the strongly log-concave case, our method is explicitly designed for — and evaluated on — real-world scenarios that are non log-concave and high-dimensional. To this end, we have benchmarked our algorithms against a broad set of state-of-the-art methods on challenging, high-dimensional sampling problems demonstrating both practicality and competitive performance.

In light of this, we kindly ask you to consider reevaluating your score. We believe that current results in our manuscript provide a comprehensive evaluation of our samplers, including a practical (and highly challenging) statistical physics model.

We look forward to engaging with you further during the rebuttal period to address any remaining concerns and welcome your feedback on our responses. Thanks again!

[1] Durmus, A. et al. (2019) Analysis of Langevin Monte Carlo via Convex Optimization. Journal of Machine Learning Research, 20 (73).

[2] Gabrié, M. et al. (2022) Proc. Natl. Acad. Sci. U.S.A. 119 (10).