BNEM: A Boltzmann Sampler Based on Bootstrapped Noised Energy Matching

Dear reviewer 2YRA,

Thank you very much for your detailed feedback and your constructive comments. We have updated our manuscript and attached answers to all of your questions and concerns below. We believe that our responses and the revised manuscript could directly address your concerns. We hope these additions encourage you to adjust your rating accordingly. If any questions or concerns arise, please do not hesitate to let us know, and we will address them properly.

Weakness1: Balance between performance improvement and computational cost

For your concerns about computational overheads, this section is outlined below:

(1) Time complexity Analysis
(2) Empirical run time comparison
(3) Performance comparison with similar computation
(4) ME-NEM: an improved method that is more memory-efficient

(1) Time complexity Analysis

We include a time complexity comparison (see https://anonymous.4open.science/r/nem_rebuttal-0D92/time_complexity.png), which theoretically shows that during training NEM can be slightly faster than iDEM in the inner-loop while it uses double computation in the outer-loop (i.e. sampling) due to the neural network differentiation. While the inner-loop training time of BNEM depends on the relative complexity of evaluating the clean energy function versus the neural network, which can be relatively higher.

(2) Empirical run time comparison

We conduct time logging for iDEM, NEM, and BNEM (see https://anonymous.4open.science/r/nem_rebuttal-0D92/run_time_comparison.png). Empirically, it shows that the sampling time, i.e. outer-loop, of BNEM/NEM is approximately twice that of iDEM, which remains within the same magnitude as iDEM. In contrast, the inner-loop time of NEM is slightly faster than that of iDEM, and the difference becomes more pronounced for more complex systems. For BNEM, the sampling time is comparable to NEM, but the inner-loop time depends on the relative complexity of evaluating the clean energy function versus the neural network, which can be relatively higher.

Based on these observations, we recommend using NEM for complex target energy functions in large systems. When sufficient computational resources are available, BNEM can be employed to trade time for improved accuracy.

(3) Performance comparison with similar computation

We emphasize that, due to the differentiation of NN, the computational cost of both NEM and BNEM during sampling is approximately TWICE as that of iDEM. Therefore, we conduct experiments with a similar computation setting (1000 step iDEM v.s. 500 step NEM/BNEM) on LJ-13 and LJ-55. You can find the results in this link: https://anonymous.4open.science/r/nem_rebuttal-0D92/half_integration_step_comparison.png which shows that with similar computations NEM/BNEM can still outperform iDEM. We can also improve more by doubling the computation during sampling.

(4) ME-NEM: an improved method that is more memory-efficient

We also proposed another estimator by using Tweedie’s formula, which is able to bypass the memory issue of differentiating NN to evaluating NN twice. And the two times of NN evaluations can definitely be computed parallelly, resulting in a similar speed as iDEM but just doubling the memory usage. The proposed method can be found in (Appendix F: Memory-Efficient NEM). And experimental results are provided in https://anonymous.4open.science/r/nem_rebuttal-0D92/memory_efficient_nem.png

[FYI: the experiments on LJ-55 are still running. We’ll update them once they are finished.]

Weakness2: Experiencing our methods on larger systems

We extended our method to the Alanine Dipeptide (ALA2) system. It is a very challenging task without any known data from the target ALA2 distribution. We provide preliminary results available at https://anonymous.4open.science/r/nem_rebuttal-0D92/aldp.png.

iDEM exhibits mode-seeking behavior, fitting extremely sharp distributions to only a limited number of modes. In contrast, our method demonstrates slightly higher variance but achieves better mass coverage with reduced bias.

Scaling up to very large systems is currently beyond the reach of existing methods, and the ALA2 problem already presents substantial challenges for current techniques. Additionally, training diffusion models becomes significantly harder when non-clean data is provided, as the noise and irregularities in the data can exacerbate the complexity of learning the underlying distributions. As the field advances, larger and more complex systems will inevitably be explored.

Weakness3: Missing comparison of traditional sampling methods

The primary contribution of our paper is to improve upon existing neural samplers based on diffusion models. Therefore, we didn’t include traditional methods such as parallel tempering, HMC, and simulated annealing. While it is true that these traditional sampling methods could be applied directly, these methods typically require significantly more steps (e.g., >10k) to converge to the target distribution, as is common in molecular dynamics. In contrast, although samplers based on neural networks require an additional training procedure, they are far more efficient during inference and offer an option to trade accuracy for fewer SDE integration steps.

Q1: Estimate of additional overhead and qualitative results for the decrease of variance

The discussion about additional overhead can be found in our answer to “Weakness1: Balance between performance improvement and computational cost”.

For qualitative results for the decrease of variance, as we mentioned in the method section (Proposition 3), the variance of the training target contributes to the error of the predicted score (we also call it bias). BNEM decreases the variance of the training target and therefore reduces these biases. However, in bootstrapping, as bias can be accumulated from the previous time step, the bias can also be increased. Therefore, bootstrapping can reduce bias on one hand, but also increase it on the other hand. We empirically show such behavior on the GMM-40 task (see https://anonymous.4open.science/r/nem_rebuttal-0D92/bias_analysis.png). It shows that we can reduce bias within properly set bootstrapping parameters.

Q2: Envision of incorporating MH sampling

Both iDEM, NEM, and BNEM cannot generate samples with accurate weights, though they can find modes, which is not optimal in sampling settings. Therefore, improving the weights with tools like Metropolis-Hastings (MH) is necessary.

To incorporate MH, we have

(1) the learned marginal energy $E_\theta(x_t, t)\approx\mathcal{E}_t(x_t)$, and

(2) the transitions between adjacent noise levels are Gaussian, i.e. 1-step Langevin dynamics $x_{t-1}|x_t\sim\mathcal{N}(x;x_t - g(t)^2\nabla E_\theta(x_t, t), g(t)^2I)$

We can apply MH at each internal transition to get higher quality samples.

Probs:

(1) The model can generate samples with more accurate weights, instead of only discovering modes

(2) we might be able to decrease the integration steps a lot without sacrificing the performance by doing MH.

Cons:

(1) The marginal energy is approximated by the energy model, which might not be exact. Therefore, the imperfectness of model introduces bias in the MH step, which might result in some misleadings. However, we are able to replace the predicted energy with the clean energy to provide more accurate guidance, i.e. using $\mathcal{E}(x_0)$ instead of $E_\theta(x_t, t)$ for small $t$.

Q3: Implement importance sampling and the corresponding computational cost

We believe that implementing importance sampling (IS) is not feasible or necessary in our case, as we do not have access to $p_\theta(x)$ for the generated samples. For standard diffusion models, $p_\theta(x)$ is obtained by integrating the probability flow ODE (PF-ODE), but this approach is often inaccurate, especially when we do not have clean sampling. In the iDEM framework, a continuous normalizing flow (CNF) is proposed to compute $p_\theta(x)$, but this method is also computationally inefficient as it requires generating lots of data and training an additional flow-based model. For NEM/BNEM, one potential approach could be to use $\exp(-E_\theta(x, 0))$ as an unnormalized model density. Since, in principle, $p_0$ and $p_{\theta, 0}$ share the same partition function, the importance weight could then be approximated as $\exp(-\mathcal{E}(x) + E_\theta(x, 0))$. However, it might yield inaccurate estimations because the ground-truth model density should be integrating the PF-ODE from t=1 to t=0. Simply using $\exp(-E_\theta(x, 0))$ as model density can introduce mismatch, therefore biased.

Q4: Experiencing proposed methods on Alanine Dipeptide

We updated new experiments on ALA2 system with our method, where the preliminary results (https://anonymous.4open.science/r/nem_rebuttal-0D92/aldp.png) show that NEM can be more likely to succeed in ALA2 than iDEM. The discussion can be found in our response to Weakness2: Experiencing our methods on larger systems above.

End of Response

Thank you again for your valuable and constructive comments. We really appreciate your efforts. We would like to kindly ask if we have addressed all of your concerns. If so, could you increase your rating accordingly? If there are any remaining concerns or if you want more discussions on the presented methods, please do not hesitate to let us know.

Q5: Improvement of method presentation

Thank you for your suggestion on clarifying our manuscript. We have refined our method section to enhance clarity and readability. Specifically, we have added a more comprehensive introduction to the overall framework before diving into the detailed analysis of the variance and bias of the proposed Monte Carlo estimator for energy, as well as the bootstrapping technique. Revised contents are colored as blues

Q6: Theoretical appeal from KL div /Fisher div for Eq. 10

In fact, (Eq.10) doesn’t relate to any f-Divergence, but simply a $L_2$ error of the energy. But there’re still theoretical advantages behind optimizing this $L_2$ error.

A previous work [1] shows that for any two $\mathbb{R}^d$ GMM with same modes but different weights, the Fisher divergence between them can be arbitrary small but the KL divergence between them remains large when two modes are sufficiently separate (See Example 1 and Theorem 6).

In addition, their Figure 1 shows that (1) In left, the Fisher divergence can NOT distinguish two GMM with same modes but different weights, while KL divergence can and $L_2$ error can be even better.(2) In right shows that the perturbed score can not tell the difference of weights until the every end of the forward process. On the other hand, it is noticeable that the perturbed log-density distinguishes the weights well throughout the forward process (Figure 1, middle).

[1] Shi, Zhekun, et al. "Diffusion-PINN Sampler." arXiv preprint arXiv:2410.15336 (2024).

End of Response

Thank you again for your reviews and detailed comments. We hope that our responses have effectively addressed all of your concerns. If so, could you increase your score accordingly? If there are any issues or further clarifications needed, please let us know. We will make every effort to address them promptly.

Dear reviewer ajpL,

Thank you very much for your recognition for our methods, your detailed feedback, and your constructive comments. We believe that our responses and the revised manuscript could directly address your concerns. We hope these additions encourage you to adjust your rating accordingly. If any questions or concerns arise, please do not hesitate to let us know, and we will address them properly. Before starting, thank you for pointing out that we missed to reference the paper by Sohl-Dickstein et al , which is already referenced now.

Q1: Explanation of the performance discrepancies compared to the iDEM paper

It is true that our visualized performance is different from the baseline model. Our experiment is running in a more challenging setting. For the GMM experiment, we reduced both the number of integration steps and the number of Monte Carlo (MC) samples from 1000 to 100. If we adopt the same setting in the iDEM paper there will not be notable performance differences observed in visualizations. For the LJ-55 experiment, we found that training iDEM is highly sensitive to random seed selection. Consequently, reproducing iDEM’s results without the specific random seed provided is nearly impossible. To ensure fairness, we reran their code and selected the best-performing seed for all models when generating plots.

Q2: Balance between performance improvement and computational cost

We emphasize that, due to the differentiation of NN, the computational cost of both NEM and BNEM during sampling is approximately TWICE as that of iDEM. Therefore, we conduct experiments with a similar computation setting (1000 step iDEM v.s. 500 step NEM/BNEM) on LJ-13 and LJ-55. You can find the results in this link: https://anonymous.4open.science/r/nem_rebuttal-0D92/half_integration_step_comparison.png which shows that with similar computations NEM/BNEM can still outperform iDEM. We can also improve more by doubling the computation during sampling.

Besides, (Fig. 5, see: https://anonymous.4open.science/r/nem_rebuttal-0D92/lj13-100.png) in our manuscript shows that NEM and BNEM can even achieve better results by reducing their computations in sampling to 1/10, which means they are 5 times faster than iDEM.

Q3: Experimental results of removing outlier

Based on (Fig 4.), we set the maximum energy as (GMM: 100; DW4: 0; LJ13: 0; LJ55: -150). We removed outliers based on this threshold and recomputed the E-W2. We report the new E-W2 as well as the percentage of outliers. The results can be found in https://anonymous.4open.science/r/nem_rebuttal-0D92/remove_outliers.png, which shows that the order of performance (BNEM>NEM>iDEM) still holds in terms of better E-W2 value and less percentage of outliers. It also indicates that without the influence of outliers, samples generated by NEM and BNEM are much more similar to the ground-truth ones.

Q4: Comparison of training time and sampling time

We include a time complexity comparison (see https://anonymous.4open.science/r/nem_rebuttal-0D92/time_complexity.png), which theoretically shows that during training NEM can be slightly faster than iDEM in the inner-loop while it uses double computation in the outer-loop (i.e. sampling) due to the neural network differentiation. The inner-loop training time of BNEM depends on the relative complexity of evaluating the clean energy function versus the neural network, which can be relatively higher.

We conduct time logging for iDEM, NEM, and BNEM (see https://anonymous.4open.science/r/nem_rebuttal-0D92/run_time_comparison.png). Empirically, it shows that the sampling time, i.e. outer-loop, of BNEM/NEM is approximately twice that of iDEM, which remains within the same magnitude as iDEM. In contrast, the inner-loop time of NEM is slightly faster than that of iDEM, and the difference becomes more pronounced for more complex systems. For BNEM, the sampling time is comparable to NEM, but the inner-loop time depends on the relative complexity of evaluating the clean energy function versus the neural network, which can be relatively higher.

Based on these observations, we recommend using NEM for complex target energy functions in large systems. When sufficient computational resources are available, BNEM can be employed to trade time for improved accuracy.

We also proposed another estimator by using Tweedie’s formula, which are able to bypass the memory issue of differentiating NN to evaluating NN twice. And the two times of NN evaluations can be definitely computed parallelly, resulting in a similar speed as iDEM but just doubling the memory usage. The proposed method can be found in (Appendix F: Memory-Efficient NEM). Experimental results are provided in https://anonymous.4open.science/r/nem_rebuttal-0D92/memory_efficient_nem.png

[FYI: the experiments on LJ-55 are still running. We’ll update them once they are finished.]

Dear reviewer zWMt,

Thank you very much for pointing out your concerns. We sincerely apologize for replying late because of computational resource shortage for running experiments to address your concerns. We much appreciate your detailed feedback and the constructive comments. We hope that we carefully addressed your concerns and made significant updates to the manuscript based on your comments. We believe that our responses and the revised submission should address your concerns and demonstrate the strengths and contributions of our work. If so, we wish you could adjust your rating accordingly. If not, please let us know so that we can address any remaining concerns.

Q1: The definition of low-energy region in Proposition 2 and clarification on variance explosion

The “low energy regions” is defined as the regions that $\mathcal{E}(x_t) < c$, where $x_t\in\mathcal{X}$ is in the support of the target distribution and $c$ is some small constant. Therefore, the surrounding points $x_{0|t}^{(i)}$ of $x_t$ can be assumed also in the low energy region of the energy function. The estimated score in this region is bounded , so that we could derive $tr(Cov(S_K(x_t, t))) > Var(E_K(x_t, t))$

Thank you for stressing out that the variance of the estimator is of multiple factors including the importance sampling proposal. It is true that when the sampling proposal is proportional to the energy, as if we used clean samples to get $x_t$, the variance only comes from the energy landscape. However, when we prove Proposition 2, we already taken the proposal $x_{0|t}^{(i)}\sim\mathcal{N}(x_t, \sigma_t^2I)$ into account for comparing $S_K$ with $E_K$.

In addition to theoretical results, the variance explosion at large time $t$ is also true from experimental results. As in Figure 9, the variance of $S_K$ and $E_K$ are visualized for the GMM example. It is obvious that for $t$ close to 1, both estimators have high variance compared to when $t$ is close to 0. Beyond this toy example, we also observe much lower losses for NEM and BNEM than iDEM, which suggests much lower variance of training targets (ideally, the optimal loss would be the variance of your target when using a $L_2$ loss).

Q2: Clarification on the bias and variance trade-off

Thank you for providing valuable comments on the bias-variance trade-off. We would like to explain more intuitions about the this in the following answer. Please do not hesitate to reach out for more discussions if you have any concerns. This answer would be outlined as follows:

(1) Why is high variance of training target favorable in standard diffusion?
(2) Why is high variance of training target NOT favorable in iDEM, NEM, and BNEM?
(3) Why bootstrapping can improve performance?
(4) How to improve performance by bootstrapping?
(5)Answer to “for complex energy distributions and in high dimensional settings, the bias could result in significant problems”
(6) Empirical study

(1) Why is high variance of training target favorable in standard diffusion?

In data-driven diffusion, it is true that neural networks perform well when training on high-variance input. But that’s because normal diffusion is designed to fit the posterior mean. Generally, when targeting a neural network $f_\theta(x)$ to noisy targets $\hat{y}(x)$, the optimality would be obtained as the posterior mean of the noisy target, i.e. $f_{\theta^*}(x)=\mathbb{E}[\hat{y}(x)|x]$. Therefore, standard diffusion models can be fitted by optimizing a relatively simple objective, e.g. noise prediction. In such case, the optimal neural network outputs $f_\theta(x_t, t)=\mathbb{E}[\epsilon_t|x_t]$ and we therefore can simply undo the Gaussian noise injection via $\hat{x}_0=\frac{x_t - \sigma_t f\theta(x_t, t)}{\alpha_t}$

In particular, the optimal marginal scores at time $t$ are: $s_t(x_t)=\int p(x_0|x_t)\nabla\log\mathcal{N}(x_t;x_0, \sigma_t^2I)dx_0$ (DSI), then the optimal $s_\theta$ obtained by

$\min_{\theta} \mathbb{E}|| s_\theta(x_t, t)- \nabla_{x_t} \log \mathcal{N}(x_t;x_0, \sigma_t^2I)||^2$, with $x_0\sim p_0$ and $x_t|x_0\sim\mathcal{N}(x_0, \sigma_t^2I)$

is given by the above posterior mean (DSI).

(2) Why is high variance of training target NOT favorable in iDEM, NEM, and BNEM?

But in our case, the high variance of training target only means informative training signals and the well-fitted posterior means are not informative. In fact, learning the noisy score estimator can result in an incorrect score at $t$ when the noise scale surpasses the data scale.

(1) Variance: It can be reduced by the smaller variance of training target
(2) Bias: It can be increased because the training target has accumulated error due to bootstrapping

Q4: Elaboration on the reproduction of results

It is true that our visualized performance is different from the baseline model. Our experiment is running in a more challenging setting.

For the GMM experiment, we reduced both the number of integration steps and the number of Monte Carlo (MC) samples from 1000 to 100. If we adopt the sampling setting in the iDEM paper there will not be notable performance differences observed in visualizations, because it’s a very simple task if we have enough computational budget. For the LJ-55 experiment, we found that training iDEM is highly sensitive to random seed selection. Consequently, reproducing iDEM’s results without the specific random seed provided is nearly impossible. To ensure fairness, we reran their code and selected the best-performing seed for all models when generating plots. For other baseline methods, like FAB, DDS and PIS, we are not sure what is the exact setting that the authors of the iDEM paper conduct their experiment as they do not provide codes or random seeds, even the architecture of neural networks for them. Therefore, we have implemented all these methods and aligned them by applying the same architecture (3 layer MLP for GMM, EGNN for DW4, LJ13 and LJ55) with the same number of parameters. We have ensured a fair comparison between baseline methods in our paper.

Q5: Missing metrics of NLL and ESS

NLL and ESS metrics, used for evaluation in the iDEM paper, are missing in our experiments because they rely on an additional flow-matching model trained on sampled data points to compute the likelihood. While reusing the code from the public iDEM repository, we discovered that the ESS implementation contained mathematical errors. Moreover, we observed that both ESS and NLL on the validation set exhibit high variance depending on the choice of checkpoints, often overshadowing the differences between sampling methods.

Also, we believe that the energy Wasserstein-2 (E-W2) contains some information about the NLL and ESS. Given

1. samples $x_i \sim p_\theta$ ($i=1,...,n$) generated by the neural sampler and
2. samples $y_i\sim \mu_\mathrm{target}$ ($i=1,...,n$) from the ground-truth distribution,

the E-W2 is computed between $[\mathcal{E}(x_1), ..., \mathcal{E}(x_n)]$ and $[\mathcal{E}(y_1), ..., \mathcal{E}(y_n)]$, i.e. the log-target-density between ground-truth samples and generated samples.

1. When E-W2 is large, it indicates a fundamental mismatch between $p_\theta$ and $\mu_\mathrm{target}$. Therefore, it implies poor NLL and ESS.

2. When E-W2 is small, the samples $y_j\sim p_\theta$ have similar log-target-density as the samples $x_j\sim \mu_\mathrm{target}$. In other words, our model $p_θ$ is generating samples that have similar "typicality" under $\mu_\mathrm{target}$ as real samples from $\mu_\mathrm{target}$, implying a good NLL and potentially good ESS

For these reasons, we decided to exclude these metrics from our evaluation.

Q6: Clarification on missing random seeds reviation

We apologize for making confusion when presenting our main table. According to your suggestion, we would like to present the experimental values with their means and standard-deviations instead of just the best value. On the other hand, we realize that our previous implementation can introduce instability when training BNEM. We’ve improved our codes and rerun their experiments. The new table (see https://anonymous.4open.science/r/nem_rebuttal-0D92/main-table-new.png) is updated in our revised manuscript.

The results show that BNEM can outperform NEM in most cases. But it’s noticeable that for LJ-55, the $\mathcal{E}$-$\mathcal{W}_2$ metric of BNEM is worse than that of NEM in terms of mean and std. However, we also notice that the best trained BNEM can achieve the best value (45.61 (BNEM) v.s. 84.17 (NEM)). And that’s because bootstrapping highly depends on whether we learn the “base” energy (the one without bootstrapping) well. It turns out that we still need a better variance indicator for evaluating NEM on small $t$ to determine whether the energy network is ready for bootstrapping. Therefore, we will keep improving the training stability of BNEM in our future work.

According to our discussion in the above Q2 and the experiment results, we conclude that bootstrapping can theoretically reduce the error of predicted scores, while practically achieve better performance by improving NEM.

Additional: we would like to point out one more benefit of NEM/BNEM

Generating high-quality samples from the target distribution is a fundamental sub-task in sampling. Beyond this, it is often essential for the generated samples to be assigned correct weights under the target distribution. However, several studies have noted that score-based sampling methods can be problematic in this context (e.g., [1][2]). While [1] shows that Metropolis-Hastings can solve this problem. Notably, iDEM, NEM, and BNEM cannot achieve this (see experiments on DW-4 and additional experiments on MW32 https://anonymous.4open.science/r/nem_rebuttal-0D92/mw32_prelim.png ). However, NEM and BNEM enable the estimation of unnormalized marginal densities, i.e., noised energies. This allows more techniques to be employed, such as Metropolis-Hastings inside the sampling trajectory, which leverage the learned energy to yield better weight estimations.

Reference [1] Chen, Wenlin, et al. "Diffusive Gibbs Sampling." arXiv preprint arXiv:2402.03008 (2024). [2] Shi, Zhekun, et al. "Diffusion-PINN Sampler." arXiv preprint arXiv:2410.15336 (2024).

End of Response

We really appreciate your efforts for these reviews and all the comments. We would like to kindly ask if we have addressed all of your concerns. If so, could you increase your rating accordingly? If there are any remaining concerns on experiments or theoretical analysis, please let us know so that we can address them promptly.

Q2.2: Bias and Error

We apologize for making any confusion. We would like to explain why BNEM can improve performance in terms of less bias in a more intuitive way.

Firstly, we should notice that the “true” target of our energy model is the noise-energy $\mathcal{E}_t(x)$. However, we now only can get a biased estimator $E_K(x, t)$. Therefore, it would introduce the first mismatch, we named it estimation bias. And this estimation bias can be increased along time $t$ (see Proposition 1 and its proof in Appendix A).

Secondly, we tend to train an energy model to fit those energy estimators. The variance of the training target can introduce the second mismatch that the learned energy is NOT perfect. And typically in NEM, when $t$ increases the variance of training target increases, which results in larger mismatches. We name this mismatch between learned energy $E_\theta(x, t)$ and the original MC estimator $E_K(x, t)$ learning bias, which is positively correlated to the variance of the training target.

Therefore, in NEM, the true bias between our learned energy $E_\theta(x, t)$ and its “true” target $\mathcal{E}_t(x)$ is contributed by both estimation bias and learning bias.

On the other hand, BNEM bootstraps the learned energy at a slightly smaller time $s$ ($s<t$), resulting in smaller learning bias and slight larger estimation bias.

In Proposition 3, we theoretically compare the estimation bias between NEM and BNEM at time $t$, where the estimation bias at t of BNEM is contributed by all estimation bias at s and learning bias at s with some rescaling factors. Proposition 3 shows that, their estimation bias at t can be derived as:

NEM: $\frac{v_{0t}(x)}{2m_t(x)^2K_1}$

BNEM: $\frac{v_{0t}(x)}{2m_t^2(x)K_2^{i+1}} + \sum_{j=1}^i\frac{v_{0s_j}(x)}{2m_{sj}^2(x)K_2^j}$

For example, consider we learn noised energy at $s$ by NEM, and learn noised energy at $t$ by bootstrapping the learned energy at $s$, and we use the same $K_1=K_2=K$. Then the estimation bias between NEM and BNEM at $(x, t)$ become

$\text{Bias}^{\text{NEM}}(x, t)=\frac{v_{0t}(x)}{2m_t(x)^2K}$

$\text{Bias}^{\text{BNEM}}(x, t; s)=\frac{v_{0t}(x)}{2m_t^2(x_t)K^{2}} +\frac{v_{0s}(x)}{2m_{s}^2(x)K}$

It shows that $\text{Bias}^{\text{BNEM}}(x, t; s)=\frac{1}{K}\text{Bias}^{\text{NEM}}(x, t) + \text{Bias}^{\text{NEM}}(x, s)$. Typically, $\text{Bias}^{\text{NEM}}(x, s)$ is much smaller and $K$ can reduce $\text{Bias}^{\text{NEM}}(x, t)$ much, Therefore, we can have $\text{Bias}^{\text{BNEM}}(x, t; s)<\text{Bias}^{\text{NEM}}(x, t)$.

Our additional experiments also convince that BNEM can reduce such estimation bias to gain improvement. For learning bias, since the variance of BNEM’s target is always smaller than NEM’s (you can intuitively think that we’re sampling in a smaller neighbourhood region when estimating the bootstrapped energy estimation in BNEM), the learning bias of BNEM can be smaller. Therefore, we can also gain improvement in terms of reducing learning bias.

Q3: Worse E-W2 in BNEM and negative tasks for NEM

Thank you for pointing out the the previous E-W2 of BNEM is worse than NEM. We realize there’re few implementation issues a few days ago, so we rerun the experiments on LJ-55. Our rerun 1.0 version, the one we shown in our first response, still has some issues and we were still waiting for the rest of experiments at that time due to shortage of computational resources in this week. Fortunately, the remaining experiments are finished now, and we would like to provide you the latest results (2.0 version) on LJ-55: https://anonymous.4open.science/r/nem_rebuttal-0D92/main-table-new.png. It shows that BNEM achieves similar mean E-W2 (119.46 v.s. 118.58), while having less variance (77.92 v.s. 106.63). We believe that it can convince the improvement gained by BNEM.

We found that NEM can consistently outperform iDEM in the providing tasks, which are considered in [1]. We therefore cannot provide any negative example that NEM fails at the current stage. But we will conduct more experiments on larger systems like ALA2 in the future. Nevertheless, we believe that the improvement gained by NEM also convince our motivation and theoretical results, which suggests learning less variated targets and has more potential to improve performance e.g. bootstrapping or Metropolis-Hastings (future work).

[1] Tara Akhound-Sadegh, Jarrid Rector-Brooks, Avishek Joey Bose, Sarthak Mittal, Pablo Lemos, Cheng-Hao Liu, Marcin Sendera, Siamak Ravanbakhsh, Gauthier Gidel, Yoshua Bengio, Nikolay Malkin, and Alexander Tong. Iterated denoising energy matching for sampling from boltzmann densities, 2024.

Q4: Different architectures for baselines and Inreproducibility of iDEM in LJ13

For the baselines, FAB, PIS, and DIS, we change their model architecture to be the same as iDEM and ours, to ensure all models have the same number of parameters for fair comparison.

For LJ-13 and LJ-55, it’s true that our results cannot reproduce the ones provided in iDEM [1]. We didn’t change the setting of iDEM, rerun their experiments on both LJ-13 and LJ-55, and reported the results. Especially for LJ-55, we found that training iDEM is highly sensitive to random seeds. We also found that few people raised issues in iDEM's repo (https://github.com/jarridrb/DEM) about the reproducibility of iDEM on LJ tasks, which is hard to reproduce and can have a very unstable training process. We believe that, to the best of our knowledge, reproducing the results in [1] is a common challenge.

[1] Tara Akhound-Sadegh, Jarrid Rector-Brooks, Avishek Joey Bose, Sarthak Mittal, Pablo Lemos, Cheng-Hao Liu, Marcin Sendera, Siamak Ravanbakhsh, Gauthier Gidel, Yoshua Bengio, Nikolay Malkin, and Alexander Tong. Iterated denoising energy matching for sampling from boltzmann densities, 2024.

End of Response

Once again, thank you for your response and constructive comments. We sincerely appreciate your feedback. We kindly ask if our explanations, intuitions, and theories were conveyed clearly and whether we have adequately addressed your concerns. If you have any remaining questions or uncertainties, please do not hesitate to let us know.

Dear reviewer zWMt,

Thank you very much for the discussion! We would like to introduce a more intuitive way to show why bootstrapping can be beneficial. We hope this additional content can address your concerns.

Additional Explanation on Bias of Bootstrapping

Suppose for $0\leq s<t\leq 1$, we learn $E_\theta(x_s, s)$ by NEM, i.e. targeting $E_K(x_s, s)$ and learn $E_\theta(x_t, t)$ by bootstrapping from the learned values at $s$.

Let’s forget about the randomness of $E_K(x_s, s)$ first. Then the optimal network predicts $E_{\tilde\theta}(x, s)=E_K(x, s)$ for $\forall x$. In such case, the bootstrapping estimator becomes:

$E_K^{B(n)}(x_t, t, s; \tilde \theta) = -\log\frac{1}{K}\sum_{i=1}^K\exp\left(-E_K(x_{s|t}^{(i)}, s)\right)=-\log\frac{1}{K}\sum_{i=1}^K\frac{1}{K}\sum_{j=1}^K\exp(-\mathcal{E}(x_{0|t}^{(ij)}, t))$

where $x_{s|t}^{(i)}\sim\mathcal{N}(x;x_t, (\sigma_t^2-\sigma_2)I)$ and $x_{0|t}^{(ij)}\sim\mathcal{N}(x;x_{s|t}^{(i)}, \sigma_s^2I)$. Therefore, we can have

$E_K^{B(n)}(x_t, t, s; \tilde \theta)=E_{K^2}(x_t, t)$

which means that ideally the bootstrap(1) estimator (with bootstrapping once corresponding to a bootstrapping chain $0\rightarrow s\rightarrow t$) is equivalent to using quadratically more MC samples. By induction, if we keep bootstrapping for several times, e.g. via a chain $0\rightarrow s_1\rightarrow s_2\rightarrow...\rightarrow s_n=s\rightarrow s_{n+1}=t$, then the bootstrap(n) estimator at $t$ can reduce the variance and bias of $E_K$ polynomially w.r.t. $K$.

However, since $E_K$ is a random variable, the optimal neural network can only fit its mean. While $E_K$ is biased, its mean equals to the true noised energy $\mathcal{E}_t$ plus a bias-term, which is also half of its variance, i.e. $\mathbb{E}[E_K(x_t, t)]=\mathcal{E}_t(x_t)+\mathrm{Bias}[E_K(x_t, t)]$ and $\mathrm{Bias}[E_K(x_t, t)]=\mathrm{Var}[E_K(x_t, t)]/2$.

Therefore, bootstrapping can polynomially reduce the original variance and bias of $E_K$ but involves extra bias inherited from the last noise level, i.e.

For $t\in[s_1, s_2]$ and $s\in[0, s_1]$, $\mathrm{Bias}[E_K(x_t, t, s;\theta)]=\frac{v_{0t}(x_t)}{2m_t^2(x_t)K}=\mathrm{Bias}[E_K(x_t, t)]$

For $t\in[s_2, s_3]$ and $s\in[s_1, s_2]$, $\mathrm{Bias}[E_K(x_t, t, s;\theta)]=\frac{v_{0t}(x_t)}{2m_t^2(x_t)K^{2}}+\frac{v_{0s}(x_t)}{2m_{s}^2(x_t)K}$

For $t\in[s_3, s_4]$ and $s\in[s_2, s_3]$, $\mathrm{Bias}[E_K(x_t, t, s;\theta)]=\frac{v_{0t}(x_t)}{2m_t^2(x)K^{3}}+\frac{v_{0s}(x_t)}{2m_{s}^2(x_t)K^2}+\frac{v_{0s_1}(x_t)}{2m_{s_1}^2(x_t)K}$

We can control the bootstrapping parameters as well as $K$ such that $\mathrm{Bias}[E_K(x_t, t, s;\theta)]<\mathrm{Bias}[E_K(x_t, t)]$. Clearer mathematical discussions can be found in the revised 3.4 section and the referenced contents in our appendix.

End of Response

Again, thank you so much for discussing this with us. If you have any remaining concerns or questions, please do not hesitate to let us know!

Thank you for your response. I have carefully read your explanations and will maintain my score for the following reasons:

1. I believe that not providing a bound on the error of the score when learning about energy is a significant limitation of the work. In fact, the theoretical discussions the authors provide do not apply to what we are interested in, which is the score used for sampling. As far as I can tell, the current theoretical results cannot be used to bound the variance of the score. This is very concerning and limits the theoretical contributions of the paper significantly.
2. I think the authors have a misunderstanding regarding the effect of the variance in the data-based training of diffusion models. Reducing the variance in the estimated score is also important in diffusion models and can lead to an improvement in their performance. This is why researchers are interested in flow-matching models or optimal transport and many works have shown that reducing the variance of the score/flow estimator improves performance. Therefore, claiming that it is not a concern in data-based training is not accurate [1, 2, 3].
3. I understand the authors' point regarding the learning vs. estimation bias. Firstly, the term "learning bias" is misleading and they should refer to this as a learning error since the term bias refers to an irreducible error. Conversely, this "learning bias" will, in theory, go to 0 when trained for long enough. However, the estimation bias is irreducible and causes significant problems, especially in high-dimensional settings.
4. Can the authors clarify how they used the same architecture for FAB when it is a normalizing flow and hence requires an invertible architecture? Additionally, besides the results of iDEM, I am concerned about how the authors failed to produce results for PIS on LJ13.

[1] Daegyu Kim, Jooyoung Choi, Chaehun Shin, Uiwon Hwang, and Sungroh Yoon. Improving diffusion-based generative models via approximated Optimal Transport, 2024.

[2] Tong, A., Malkin, N., Huguet, G., Zhang, Y., Rector-Brooks, J., Fatras, K., Wolf, G., and Bengio, Y. Improving and generalizing flow-based generative models with minibatch optimal transport, 2023.

[3] Yiheng Li, Heyang Jiang, Akio Kodaira, Masayoshi Tomizuka, and Kurt Keutzer. Immiscible Diffusion: Accelerating Diffusion Training with Noise Assignment, 2024.

Dear reviewer MWcQ,

We would like to start by thanking you for highlighting our strengths, including the concise figure 1, the theories, and the results. We hope that our detailed responses and the updates in the revised manuscript reaffirm your satisfaction with the paper. We also hope these additions encourage you to increase your rating accordingly. If any questions or concerns arise, please do not hesitate to let us know, and we will address them properly.

Weakness1: Paper makes little or no progress on the fundamental challenge of scaling to large systems

Thank you for pointing this out. Firstly, we would like to say that NEM and BNEM are transformational rather than incremental. Here’s the reason: Generating high-quality samples from the target distribution is a fundamental sub-task in sampling. Beyond this, it is often essential for the generated samples to be assigned correct weights under the target distribution. However, several studies have noted that score-based sampling methods can be problematic in this context (e.g., [1][2]), where modes tend to be covered but are associated with inaccurate weights. While [1] can generate samples with correct weights once applying Metropolis-Hastings. Notably, iDEM, NEM, and BNEM cannot achieve this (see experiments on DW-4 and additional experiments on MW32 https://anonymous.4open.science/r/nem_rebuttal-0D92/mw32_prelim.png). However, NEM and BNEM enable the estimation of unnormalized marginal densities, i.e., exponential negative noised energies. This allows more techniques to be employed, such as Metropolis-Hastings, which leverage the learned nosied energies to potentially yield better weight estimations.

Secondly, we admit that we are testing models on toy problems. However, generating samples without any data is a very difficult problem, and only FAB [3] can achieve this on Alanine Dipeptide (ALA2) now. Our preliminary results show that NEM can be more likely to succeed in ALA2 than iDEM. Here’s the link for our preliminary results: https://anonymous.4open.science/r/nem_rebuttal-0D92/aldp.png, and we also include more discussions in the below Q2 answer.

[1] Chen, Wenlin, et al. "Diffusive Gibbs Sampling." arXiv preprint arXiv:2402.03008 (2024).

[2] Shi, Zhekun, et al. "Diffusion-PINN Sampler." arXiv preprint arXiv:2410.15336 (2024).

[3] Midgley, Laurence Illing, et al. "Flow annealed importance sampling bootstrap." arXiv preprint arXiv:2208.01893 (2022).

Q1: Why interatomic distances are less sharply peaked than ground truth

Thank you for highlighting this observation. The less sharply peaked distance distribution in our results may be influenced by the gradient clipping in iDEM and the energy smoothing employed in our work. Modeling extremely steep energy landscapes—where the energy approaches infinity as two particles come close in space—remains a challenge for model numeric stability. However, when both iDEM and BNEM sample from the same smoothed target energy function, our model consistently outperforms score-based methods. This demonstrates the robustness and effectiveness of our approach under comparable conditions.

Q2: Experiencing our methods on larger systems

We scale up our method to the ALA2 system, where preliminary results are provided in https://anonymous.4open.science/r/nem_rebuttal-0D92/aldp.png. DEM tends to exhibit mode-seeking behavior, fitting extremely sharp distributions to only a limited number of modes. In contrast, our method, while demonstrating slightly higher variance, shows a mass-covering behavior with reduced bias. Scaling up to very large systems remains beyond the capabilities of existing methods, and the ALA2 problem already poses significant challenges for current techniques. As the field progresses, larger and more complex systems will undoubtedly be considered. However, at this stage, it is premature to address such systems comprehensively.

End of Response

Thank you again for your valuable reviews. We hope that our responses have effectively addressed all of your concerns. Please feel free to comment if you have other concerns or suggestions. We hope you could increase the rating accordingly if we have addressed your concerns successfully.