BNEM: A Boltzmann Sampler Based on Bootstrapped Noised Energy Matching

## Abstract

Abstract Developing an efficient sampler capable of generating independent and identically distributed (IID) samples from a Boltzmann distribution is a crucial challenge in scientific research, e.g. molecular dynamics.
In this work, we intend to learn neural samplers given energy functions instead of data sampled from the Boltzmann distribution.
By learning the energies of the noised data, we propose a diffusion-based sampler, Noised Energy Matching , which theoretically has lower variance and more complexity compared to related works.
Furthermore, a novel bootstrapping technique is applied to NEM to balance between bias and variance.
We evaluate NEM and BNEM on a 2 2 2 2 -dimensional 40 40 40 40 Gaussian Mixture Model (GMM) and a 4 4 4 4 -particle double-well potential (DW-4). The experimental results demonstrate that BNEM can achieve state-of-the-art performance while being more robust.

## 1 Introduction

Given energy functions $\mathcal{E}(x)$, samplers for unnormalized target distributions i.e. Boltzmann distribution $\mu_{\text{target}}\propto\exp(-\mathcal{E}(x))$ are fundamental in probabilistic modeling and physical systems simulation.
For example, predicting the folding of proteins could be formalized as sampling from a Boltzmann distribution (Śledź and Caflisch,, ), where the energies are defined by inter-atomic forces (Case et al.,, ).
Efficient samplers for many-particle systems could have the potential to speed up drug discovery (Zheng et al.,, ) and material design (Komanduri et al.,, ).

However, most related works face challenges with scalability to high dimensions and are time-consuming.
To solve these problems, a previous work(Akhound-Sadegh et al.,, ) proposes Iterated Denoising Energy Matching (iDEM), which is not only computationally tractable but also guarantees good coverage of all modes.
iDEM proposes a bi-level training scheme that iteratively generates samples from the learned sampler and performs score matching to the Monte Carlo estimated target.
While Woo and Ahn, () proposes a variant, iEFM, that targets the MC estimated vector fields in a Flow Matching fashion, which we found can be linked through Tweedie’s formula (Efron,, ) shown in Appendix G. Nevertheless, both iDEM and iEFM require large numbers of MC samplings for score estimation to minimize the variance at large time steps, which holds disadvantages for complicated energy functions.

In this work, we propose an energy-based variant of DEM, Noised Energy Matching (NEM), which targets the MC estimated noised energies instead of scores.
Despite a need to differentiate the neural network when simulating the diffusion sampler, NEM is theoretically found to be targeting less noisy objectives compared with iDEM.
Our method demonstrates better performance on a $2$-dimensional $40$ Gaussian Mixture Model (GMM) and a $4$-particle double-well potential (DW-4) ().
To further reduce the variance, Bootstrap NEM (BNEM) is proposed in our work, which estimates high noise-level energies by bootstrapping from slightly lower noise-level estimators.
We theoretically found that BNEM trades bias to the variance of training target, which empirically leads to state-of-the-art performance.

## 2 Related works

Conventional methods for many-body system simulation are based on molecular dynamics (MD) Leimkuhler and Matthews, () techniques which require long simulation times. Others leverage Monte Carlo techniques e.g. Annealed Importance Sampling (AIS) (Lyman and Zuckerman,, ), which are computationally expensive.
To speed up this process, a promising remedy is training a deep generative model, i.e. a neural sampler.
Unlike data-driven tasks, simply minimizing the reverse Kullback-Leibler (KL) divergence for samplers can lead to mode-seeking behavior. Noé et al., () address this problem by using the Boltzmann Generator, which minimizes the combination of forward and reverse KL divergence.

Inspired by the rapid development of deep generative models, e.g. diffusion models (Song and Ermon,, ; Ho et al.,, ),
pseudo samples could be generated from an arbitrary prior distribution.
Then, we can train the neural samplers by matching these sample trajectories, as in Path Integral Sampler (PIS)(Zhang and Chen,, )
and Denoising Diffusion Sampler (DDS) (Vargas et al.,, ).
Midgley et al., () further deploy a replay buffer for the trajectories while proposing an $\alpha$-divergence as the objective to avoid mode-seeking.
However, these methods require simulation during training. DDS and PIS require the diffusion trajectories to minimize the KL divergence between prior distribution and target distribution over them, while FAB requires simulating the AIS to correct the weights. This poses challenges for scaling up to higher dimensional tasks.

## 3 Methods

### 3.1 Denoising diffusion based Boltzmann sampler

In this work, we consider training an energy-based diffusion sampler corresponding to a variance exploding (VE) noising process defined by $dx_{t}=g(t)dw_{t}$, where $t\in[0,1]$, $g(t)$ is a function of time and $w_{t}$ is a Brownian motion.
Then the corresponding reverse SDE with Brownian motion $\bar{w}_{t}$ is $dx_{t}=-g^{2}(t)\nabla\log p_{t}(x_{t})d_{t}+g(t)d\bar{w}_{t}$, where $p_{t}$ is the marginal distribution of the diffusion process that starts at $p_{0}=\mu_{target}$.
Given access to the system energy $\mathcal{E}(x)$ and the perturbation kernel $q_{t}(x_{t}|x_{0})=\mathcal{N}(x_{t};x_{0},\sigma_{t}^{2})$, where $\exp(-\mathcal{E}(x))\propto p_{0}(x)$, one can obtain the marginal distribution $p_{t}$

$$ $\displaystyle p_{t}(x_{t})$ $\displaystyle\propto\int\exp(-\mathcal{E}(x_{0}))\mathcal{N}(x_{t};x_{0}, \sigma_{t}^{2}I)\mathrm{d}x_{0}=\mathbb{E}_{\mathcal{N}(x;x_{t},\sigma_{t}^{2} I)}[\exp(-\mathcal{E}(x))]$ (1) $$

iDEM is proposed to train a score network $s_{\theta}$ to match an MC score estimator $S_{K}$

$$ $\displaystyle S_{K}(x_{t},t):=\nabla\log\frac{1}{K}\sum_{i=1}^{K}\exp(- \mathcal{E}(x_{0|t}^{(i)})),\quad x_{0|t}^{(i)}\sim\mathcal{N}(x;x_{t},\sigma_ {t}^{2}I)$ (2) $$

Instead of the score network, here we consider training an energy network, $E_{\theta}(x,t)$, to approximate the following noised energy, i.e. the energy of noised data

$$ $\displaystyle\mathcal{E}_{t}(x):=-\log\mathbb{E}_{\mathcal{N}(x;x_{t},\sigma_{ t}^{2}I)}[\exp(-\mathcal{E}(x))]$ (3) $$

where $\exp(-\mathcal{E}_{t}(x_{t}))\propto p_{t}(x_{t})$.
Then the score of marginal distribution at $t$ can be approximated by differentiating the energy network w.r.t its input $x$, i.e. $s_{\theta}(x,t)=-\nabla_{x}E_{\theta}(x,t)$.
We train NEM with a bi-level scheme by iterating:
(a) an outer-loop that simulates the energy-based diffusion sampler and updates the replay buffer;
(b) a simulation-free inner-loop that matches the noised energies of samples in the replay buffer.
Detailed implementation is discussed in Appendix A.

### 3.2 Noised Energy Matching

NEM targets a novel MC energy estimator that approximates the noised energies

$$ $\displaystyle E_{K}(x_{t},t):$ $\displaystyle=-\log\frac{1}{K}\sum_{i=1}^{K}\exp(-\mathcal{E}(x_{0|t}^{(i)})), \quad x_{0|t}^{(i)}\sim\mathcal{N}(x;x_{t},\sigma_{t}^{2}I)$ (4) $$

where Equation 4 can be implemented by the LogSumExp trick for stability. We characterize the bias of $E_{K}$ with the following proposition:

###### Proposition 1

If $\exp(-\mathcal{E}(x_{0|t}^{(i)}))$ is sub-Gaussian, then there exists a constant $\tilde{c}(x_{t})$ such that with probability $1-\delta$ over $x_{0|t}^{(i)}\sim\mathcal{N}(x_{t},\sigma_{t}^{2})$, we have

$$ $\|E_{K}(x_{t},t)-\mathcal{E}_{t}(x_{t})\|\leq\frac{\tilde{c}(x_{t})\sqrt{\log{ (1}/{\delta})}}{\sqrt{K}}$ (5) $$

with ${c}(x_{t})/\tilde{c}(x_{t})=2(1+\|\nabla\mathcal{E}_{t}(x_{t})\|)$, where

$$ $\|S_{K}(x_{t},t)-S_{t}(x_{t})\|\leq\frac{c(x_{t})\sqrt{\log{(1}/{\delta})}}{ \sqrt{K}}.$ (6) $$

Proposition 1 shows that the training target of NEM has a smaller bound of bias compared with the one of DEM, especially on regions with steep gradient, i.e. large $\|\nabla\mathcal{E}_{t}(x))\|$. The complete proof is given in Appendix B. Besides, we characterize the variance of
$S_{K}$
and
$E_{K}$
as follows:

###### Proposition 2

If $\exp(-\mathcal{E}(x_{0|t}^{(i)}))$ is sub-Gaussian and $\|\nabla\exp(-\mathcal{E}(x_{0|t}^{(i)}))\|$ is bounded, the total variance of the MC score estimator $S_{K}$ is consistently larger than that of the MC energy estimator $E_{K}$ in regions associated with low energies, with

$$ $\displaystyle\frac{\textbf{tr}\left(\text{Cov}[S_{K}(x_{t},t)]\right)}{\text{ Var}[E_{K}(x_{t},t)]}$ $\displaystyle={4(1+\|\nabla\mathcal{E}_{t}(x_{t})\|)^{2}}.$ (7) $$

In regions associated with high energies, $\text{Var}[E_{K}(x_{t},t)]<\text{Var}[S_{K}(x_{t},t)]$ holds when the target energy $\mathcal{E}(x_{t})$ is positively related to at least one element of the score $\nabla\mathcal{E}(x_{t})$.

It
demonstrates that the MC energy estimator can provide a less noisy training signal than the score one, showcasing the theoretical advantage of NEM compared with DEM. The complete proof is provided in Appendix C.

###### Proposition 1

###### Proposition 2

### 3.3 Bootstrap NEM: an improvement with bootstrapped energy estimation

Bootstrap NEM, or BNEM, targets a novel MC energy estimator at a high noise level that is bootstrapped from the learned energies at a slightly lower noise level. Intuitively, the variance of $E_{K}$ exploded at a high noise level as a result of the VE noising process; while we can estimate these energies from the ones at a low noise level rather than the system energy to reduce variance. Suppose a low-level noise-convolved energy is well learned, say $\mathcal{E}_{s}$, we can construct a bootstrap energy estimator at higher noise level $t$ by

$$ $\displaystyle E_{K}(x_{t},t,s;\phi):$ $\displaystyle=-\log\frac{1}{K}\sum_{i=1}^{K}\exp(-{E}_{\phi}(x_{s|t}^{(i)},s)) ,\quad x_{s|t}^{(i)}\sim\mathcal{N}\left(x;x_{s},(\sigma_{t}^{2}-\sigma_{s}^{2 })I\right)$ (8) $$

where $s<t$ and $\phi$ a well learned neural network for any $u\in[0,s]$. We show that this bootstrap energy estimator is trading variance to bias, in terms of training target, characterized in the following proposition:

###### Proposition 3

Given a bootstrap trajectory $\{s_{i}\}_{i=0}^{n}$ such that $\sigma^{2}_{s_{i}}-\sigma^{2}_{s_{i-1}}\leq\kappa$, where $s_{0}=0$, $s_{n}=1$ and $\kappa>0$ is a small number, suppose $E_{\theta}$ is incrementally optimized from 0 0 to $1$ as follows: if $t\in[s_{i},s_{i+1}]$, then $E_{\theta}$ targets an energy estimator bootstrapped from $s\in[s_{i-1},s_{i}]$ using Eq. 8.
Let $v_{yz}(x_{z})=\text{Var}_{\mathcal{N}(x_{z},(\sigma_{z}^{2}-\sigma_{y}^{2})I)}
[\exp(-\mathcal{E}(x))]$ and $m_{z}(x_{z})=\exp(-\mathcal{E}_{z}(x_{z}))$, where $0\leq y<z\leq 1$ and $\mathcal{E}_{z}$ is the ground truth noised energy at time $z$. The variance of the bootstrap energy estimator is given by

$$ $\displaystyle{\text{Var}[E_{K}(x_{t},t,s;\theta)]}$ $\displaystyle=\frac{v_{st}(x_{t})}{v_{0t}(x_{t})}\text{Var}[E_{K}(x_{t},t)]$ (9) $$

and the bias of $E_{K}(x_{t},t,s;\theta)$ is given by

$$ $\displaystyle\frac{v_{0t}(x_{t})}{2m_{t}^{2}(x_{t})K^{n+1}}+\sum_{j=1}^{n} \frac{v_{0s_{j}}(x_{t})}{2m_{s_{j}}^{2}(x_{t})K^{j}}.$ (10) $$

A detailed discussion and proof are given in Appendix D. Proposition 3 demonstrates that the bootstrap energy estimator, which estimates the noised energies by sampling in a smaller ball, can reduce the variance of the training target while this new target can introduce accumulated bias.

###### Proposition 3

## 4 Experiments

**Table 1: Neural sampler performance comparison for GMM-40 and DW-4 energy function. we measured the performance using data Wasserstein-2 distance (x-$\mathcal{W}_{2}$), Energy Wasserstein-2 distance ($\mathcal{E}$-$\mathcal{W}_{2}$), and Total Variation (TV). * indicates divergent training. Bold indicates the best values and underline indicates the second best ones.**
| Energy $\rightarrow$ | GMM-40 ($d=2$) | DW-4 ($d=8$) |  |  |  |  |
| --- | --- | --- | --- | --- | --- | --- |
| Sampler $\downarrow$ | x-$\mathcal{W}_{2}$$\downarrow$ | $\mathcal{E}$-$\mathcal{W}_{2}$$\downarrow$ | TV$\downarrow$ | x-$\mathcal{W}_{2}$$\downarrow$ | $\mathcal{E}$-$\mathcal{W}_{2}$$\downarrow$ | TV$\downarrow$ |
| DDS | 11.69 | 86.69 | 0.944 | 0.701 | 109.8 | 0.429 |
| PIS | 5.806 | 76.35 | 0.940 | * | * | * |
| FAB | 3.828 | 64.23 | 0.824 | 0.614 | 211.5 | 0.359 |
| iDEM | 8.512 | 562.7 | 0.909 | 0.532 | 2.109 | 0.161 |
| NEM (ours) | 5.192 | 85.05 | 0.906 | 0.489 | 0.999 | 0.145 |
| BNEM (ours) | 3.652 | 2.973 | 0.830 | 0.467 | 0.458 | 0.134 |

Figure: (a) DDS
Refer to caption: extracted/5916790/figs/gmm/dds.png

We evaluate our methods and baseline models on 2 potentials, the Gaussian mixture model (GMM) and the 4-particle double-well (DW-4) potential. We provide the full description of the experimental setting in Appendix H.

Baseline. We compare NEM and BNEM to following recent works: Denoising Diffusion Sampler (DDS)(Vargas et al.,, ), Path Integral Sampler (PIS)(Zhang and Chen,, ), Flow Annealed Bootstrap (FAB)(Midgley et al.,, ) and Iterated Denoising Energy Matching (iDEM)(Akhound-Sadegh et al.,, ).
For a fair comparison, we set the number of integration steps to $100$ and the number of MC samples to $100$.
For baselines, we train all samplers using an NVIDIA-A100 GPU.

Architecture. We implement the same network architecture (MLP for GMM and EGNN for DW-4) for all baselines.
To ensure a similar number of parameters for each sampler, if the score network is parameterized by $s_{\theta}(x,t)=f_{\theta}(x,t)$, the energy network is set to be $E_{\theta}(x,t)=\sum f_{\theta}(x,t)+c$ with a learnable scalar $c$.
Furthermore, this setting ensures SE(3) invariance for the energy network.

Metrics. We use x-$\mathcal{W}_{2}$, $\mathcal{E}{-}\mathcal{W}_{2}$, and TV as metrics. To compute $\mathcal{W}_{2}$ and TV, we use pre-generated samples as datasets: (a) For GMM, we sample from the ground truth distribution; (b) For DW-4, we use samples from .
.

Main results.
We report x-$\mathcal{W}_{2}$, $\mathcal{E}{-}\mathcal{W}_{2}$ and TV for GMM and DW-4 potentials in Table 1. It shows that when the reverse SDE integration step is limited to $100$, NEM can outperform iDEM on all metrics in each task, showcasing its robustness. Furthermore, BNEM can outperform NEM by targeting the bootstrap energy estimator, achieving state-of-the-art performance. In addition, the visualized results for GMM are provided in Figure 1, which shows that NEM can generate samples more concentrated to the GMM modes than iDEM and BNEM can further boost this concentration.

## 5 Conclusion

In this work, we introduce NEM and BNEM, neural samplers for Boltzmann distribution and equilibrium systems like many-body systems. NEM uses a novel Monte Carlo energy estimator with reduced bias and variance. BNEM builds on NEM, employing an energy estimator bootstrapped from lower noise-level data, theoretically trading bias for variance. Empirically, BNEM achieves state-of-the-art results on the GMM and DW4 benchmarks. Future work will focus on scaling these methods to higher-dimensional tasks like Lennard-Jones potential and alanine dipeptide.

## Acknowledgements

José Miguel Hernández-Lobato acknowledges support from a Turing AI Fellowship under grant EP/V023756/1. RKOY acknowledges resources provided by the Cambridge Service for Data Driven Discovery (CSD3) operated by the University of Cambridge Research Computing Service ([www.csd3.cam.ac.uk](www.csd3.cam.ac.uk)), provided by Dell EMC and Intel using Tier-2 funding from the Engineering and Physical Sciences Research Council (capital grant EP/T022159/1), and DiRAC funding from the Science and Technology Facilities Council ([www.dirac.ac.uk](www.dirac.ac.uk)).

## Appendix A Training Details

Algorithm 1 describes NEM in detail. The key difference in training the BNEM is only in the inner-loop.

To train the BNEM, it’s crucial that the bootstrap energy estimation at $t$ can be accurate only when the noised energy at $s$ is well-learned. Therefore, we need to sequentially learn energy at small t using the original
MC energy estimator then refines the estimator using bootstrap for large t, which is inefficient. To be more efficient, we adapt a rejection training scheme: (a) given $s$ and $t$, we first compute the loss of targeting the MC energy estimator (4), $l_{s}$ and $l_{t}$; (b) These losses indicate how does the energy network fit the noised energies at different times, and therefore compute $\alpha=l_{s}/l_{t}$ as an indicator; (c) with probability $\alpha$, we accept targeting a energy estimator bootstrapped from $s$ and otherwise we stick to target the original MC estimator. We provide a full description of the inner-loop of BNEM training in Algorithm 2.

Figure: Algorithm 1 Iterated training for Noised Energy Matching

Figure: Algorithm 2 Inner-loop of Bootstrap Noised Energy Matching training

## Appendix B Proof of Proposition 1

###### Proposition 1

If $\exp(-\mathcal{E}(x_{0|t}^{(i)}))$ is sub-Gaussian, then there exists a constant $\tilde{c}(x_{t})$ such that with probability $1-\delta$ over $x_{0|t}^{(i)}\sim\mathcal{N}(x_{t},\sigma_{t}^{2})$, we have

$$ $\|E_{K}(x_{t},t)-\mathcal{E}_{t}(x_{t})\|\leq\frac{\tilde{c}(x_{t})\sqrt{\log{ (1}/{\delta})}}{\sqrt{K}}$ (11) $$

with ${c}(x_{t})/\tilde{c}(x_{t})=2(1+\|\nabla\mathcal{E}_{t}(x_{t})\|)$, where

$$ $\|S_{K}(x_{t},t)-S_{t}(x_{t})\|\leq\frac{c(x_{t})\sqrt{\log{(1}/{\delta})}}{ \sqrt{K}}.$ (12) $$

Proof. We first introduce the error bound of the MC score estimator $S_{K}$, where $S_{K}=\nabla E_{K}$, proposed by Akhound-Sadegh et al., () as follows

$$ $\displaystyle\|S_{K}(x_{t},t)-S(x_{t},t)\|\leq\frac{2C\sqrt{\log(\frac{2}{ \delta})}(1+\|\nabla{\mathcal{E}}_{t}(x_{t})\|)\exp({\mathcal{E}}_{t}(x_{t}))} {\sqrt{K}}$ (13) $$

which assumes that $\exp(-\mathcal{E}(x_{0|t}^{(i)}))$ is sub-Gaussian. Let’s define the following variables

$$ $\displaystyle m_{t}(x_{t})$ $\displaystyle=\exp(-\mathcal{E}_{t}(x_{t}))$ (14) $\displaystyle v_{st}(x_{t})$ $\displaystyle=\text{Var}_{\mathcal{N}(x_{t},(\sigma_{t}^{2}-\sigma_{s}^{2})I)} [\exp(-\mathcal{E}(x))]$ (15) $$

By the sub-Gaussianess assumption, it’s easy to show that the constant term $C$ in Equation 13 is $C=\sqrt{2v_{0t}(x_{t})}$. Notice that $E_{K}$ is a logarithm of an unbiased estimator
. By the sub-Gaussian assumption, one can derive that $E_{K}$ is also sub-Gaussian. Furthermore, it’s mean and variance can be derived by employing a first-order Taylor expansion:

$$ $\displaystyle\mathbb{E}[E_{K}(x_{t},t)]$ $\displaystyle\approx\mathcal{E}_{t}(x_{t})+\frac{v_{0t}(x_{t})}{2m^{2}_{t}(x_{ t})K}$ (16) $\displaystyle\text{Var}[E_{K}(x_{t},t)]$ $\displaystyle\approx\frac{v_{0t}(x_{t})}{m^{2}_{t}(x_{t})K}$ (17) $$

And one can obtain its concentration inequality by incorporating the sub-Gaussianess

$$ $\displaystyle\left\|E_{K}(x_{t},t)-\mathbb{E}[E_{K}(x_{t},t)]\right\|$ $\displaystyle\leq\sqrt{2\frac{v_{0t}(x_{t})}{m^{2}_{0t}(x_{t})K}\log\frac{2}{ \delta}}$ (18) $$

By using the above Inequality 18 and the triangle inequality

$$ $\displaystyle\|E_{K}(x_{t},t)-\mathcal{E}_{t}(x_{t})\|$ $\displaystyle\leq\|\log E_{K}(x_{t},t)-\mathbb{E}[E_{K}(x_{t},t)]\|+\|\mathbb{ E}[E_{K}(x_{t},t)]-\mathcal{E}_{t}(x_{t})\|$ (19) $\displaystyle=\|\log E_{K}(x_{t},t)-\mathbb{E}[E_{K}(x_{t},t)]\|+\frac{v_{0t}( x_{t})}{2m^{2}_{t}(x_{t})K}$ (20) $\displaystyle\leq\sqrt{2\frac{v_{0t}(x_{t})}{m^{2}_{t}(x_{t})K}\log\frac{2}{ \delta}}+\frac{v_{0t}(x_{t})}{2m^{2}_{t}(x_{t})K}$ (21) $\displaystyle=\frac{C\sqrt{\log\frac{2}{\delta}}\exp(\mathcal{E}_{t}(x_{t}))}{ \sqrt{K}}+O(1/K)$ (22) $\displaystyle=\frac{c(x_{t})}{2(1+\|\nabla\mathcal{E}_{t}(x_{t})\|)}\frac{ \sqrt{\log\frac{1}{\delta}}}{\sqrt{K}}$ (23) $$

Therefore, we have $c(x_{t})=2(1+\|\nabla\mathcal{E}_{t}(x_{t})\|)\tilde{c}(x_{t})$. It demonstrates a less biased estimator, which, what’s more, doesn’t require a sub-Gaussianess assumption over $\|\nabla\mathcal{E}(x_{0|t}^{(i)})\|$.
$\square$

###### Proposition 1

## Appendix C Proof of Proposition 2

###### Proposition 2

If $\exp(-\mathcal{E}(x_{0|t}^{(i)}))$ is sub-Gaussian and $\|\nabla\exp(-\mathcal{E}(x_{0|t}^{(i)}))\|$ is bounded, the total variance of the MC score estimator $S_{K}$ is consistently larger than that of the MC energy estimator $E_{K}$ in regions associated with low energies, with

$$ $\displaystyle\frac{\textbf{tr}\left(\text{Cov}[S_{K}(x_{t},t)]\right)}{\text{ Var}[E_{K}(x_{t},t)]}$ $\displaystyle={4(1+\|\nabla\mathcal{E}_{t}(x_{t})\|)^{2}}.$ (24) $$

In regions associated with high energies, $\text{Var}[E_{K}(x_{t},t)]<\text{Var}[S_{K}(x_{t},t)]$ holds when the target energy $\mathcal{E}(x_{t})$ is positively related to at least one element of the score $\nabla\mathcal{E}(x_{t})$.

Proof. We split the proof into two parts: low-energy region and high-energy one. The proof in the low-energy region requires only the aforementioned sub-Gaussianess and bounded assumptions, while the one in the high-energy region requires an additional constraint which will be clarified later. Review that $S_{K}$ can be expressed as an importance-weighted estimator as follows:

$$ $\displaystyle S_{K}(x_{t},t)$ $\displaystyle=\frac{\frac{1}{K}\sum_{i=1}^{K}\nabla\exp(-\mathcal{E}(x_{0|t}^{ (i)}))}{\frac{1}{K}\sum_{i=1}^{K}\exp(-\mathcal{E}(x_{0|t}^{(i)}))}$ (25) $$

Let $\|\nabla\exp(-\mathcal{E}(x_{0|t}^{(i)}))\|\leq M$, where $M>0$. Since a bounded variable is sub-Gaussian, this assumption resembles a sub-Gaussianess assumption of $\|\nabla\exp(-\mathcal{E}(x_{0|t}^{(i)}))\|$. Then each element of $\|\nabla\exp(-\mathcal{E}(x_{0|t}^{(i)}))\|$, i.e. $\nabla\exp(-\mathcal{E}(x_{0|t}^{(i)}))[j]$, is bounded by $M$. And therefore $\nabla\exp(-\mathcal{E}(x_{0|t}^{(i)}))[j]$ are sub-Gaussian.

In low-energy regions. $\exp(-\mathcal{E}(x))$ is concentrated away from 0 0 as $\mathcal{E}(x)$ is small.
Then, there exists a constant $c$ such that $\exp(-\mathcal{E}(x_{0|t}^{(i)}))\geq c>0$ and thus for each element $j=1,..,d$:

$$ $\displaystyle\|S_{K}(x_{t},t)[j]\|$ $\displaystyle=\left\|\frac{\frac{1}{K}\sum_{i=1}^{K}\nabla\exp(-\mathcal{E}(x_ {0|t}^{(i)}))[j]}{\frac{1}{K}\sum_{i=1}^{K}\exp(-\mathcal{E}(x_{0|t}^{(i)}))}\right\|$ (26) $\displaystyle\leq\frac{\|\sum_{i=1}^{K}\nabla\exp(-\mathcal{E}(x_{0|t}^{(i)})) [j]\|}{Kc}\leq M/c$ (27) $$

therefore, the $j^{th}$ element of $S_{K}$, i.e.$S_{K}[j]$, is bounded by $M/c$, suggesting it is sub-Gaussian. While Inequality 13 can be expressed as

$$ $\displaystyle\sqrt{\sum_{j=1}^{d}(\mathcal{S}_{K}(x_{t},t)[j]-S_{t}(x_{t})[j]) ^{2}}\leq\frac{2\sqrt{2v_{0t}(x_{t})\log(\frac{2}{\delta})}(1+\|\nabla{ \mathcal{E}}_{t}(x_{t})\|)}{m_{t}(x_{t})\sqrt{K}}$ (28) $$

We can roughly derive a bound elementwisely

$$ $\displaystyle|\mathcal{S}_{K}(x_{t},t)[j]-S_{t}(x_{t})[j]|\leq\frac{2\sqrt{2v_ {0t}(x_{t})\log(\frac{2}{\delta})}(1+\|\nabla{\mathcal{E}}_{t}(x_{t})\|)}{m_{t }(x_{t})\sqrt{Kd}}$ (29) $$

which suggests that we can approximate the variance of $S_{K}(x_{t},t)[j]$ by leveraging its sub-Gaussianess

$$ $\displaystyle\text{Var}(S_{K}(x_{t},t)[j])\approx\frac{4v_{0t}(x_{t})({1+\| \nabla\mathcal{E}_{t}(x_{t})\|})^{2}}{m^{2}_{t}(x_{t})Kd}$ (30) $$

Therefore, according to Equation 17 we can derive that

$$ $\displaystyle\textbf{tr}(\text{Cov}[S_{K}(x_{t},t)])$ $\displaystyle=\sum_{j=1}^{d}\text{Var}[S_{K}(x_{t},t)[j]]$ (31) $\displaystyle=\frac{4v_{0t}(x_{t})({1+\|\nabla\mathcal{E}_{t}(x_{t})\|})^{2}}{ m^{2}_{t}(x_{t})K}$ (32) $\displaystyle=4(1+\|\nabla\mathcal{E}_{t}(x_{t})\|)^{2}\text{Var}[E_{K}(x_{t}, t)]$ (33) $$

In high-energy region. we assume that there exists a direction with a large norm pointing to low energy regions, i.e. $\exists j$ such that $\mathcal{E}(x)$ are positively related to $\nabla\mathcal{E}(x)[j]$. According to Section 9.2 in Owen, (), the asymptotic variance of a self-normalized importance sampling estimator is given by:

$$ $\displaystyle\mu$ $\displaystyle=\mathbb{E}_{q}[f(X)]$ (34) $\displaystyle\tilde{\mu}_{q}$ $\displaystyle=\frac{\sum_{i=1}^{K}w_{i}f_{i}}{\sum_{i=1}^{K}w_{i}}$ (35) $\displaystyle\text{Var}(\tilde{\mu}_{q})$ $\displaystyle\approx\frac{1}{K}\mathbb{E}_{q}[w(X)]^{-2}\mathbb{E}_{q}[w(X)^{2 }(f(X)-\mu)^{2}]$ (36) $$

By substituting $\tilde{\mu}_{q}=S_{K}(x_{t},t)[j]$, $f(X)=-\nabla\mathcal{E}(X)[j]$, $w(X)=\exp(-\mathcal{E}(X))$, $q=N(x;x_{t},\sigma_{t}^{2})$, $\mathbb{E}_{q}[w(X)]=m_{t}(x_{t})$ and $\mathbb{E}_{q}[w^{2}(X)]=v_{0t}(x_{t})+m^{2}_{t}(x_{t})$, as well as using $w(X)$ and $f(X)$ are positive related, we have:

$$ $\displaystyle\text{Var}[S_{K}(x_{t},t)[j]]$ $\displaystyle\geq\frac{1}{K}\mathbb{E}_{q}[w(X)]^{-2}\mathbb{E}_{q}[w^{2}(X)] \mathbb{E}_{q}[(f(X)-\mu)^{2}]$ (37) $\displaystyle=\frac{v_{0t}(x_{t})+m^{2}_{t}(x_{t})}{m^{2}_{t}(x_{t})K}\text{ Var}_{q}[\nabla\mathcal{E}(x)[j]]$ (38) $$

Therefore, if we further have a large variance over the system score at this region, i.e. $\text{Var}_{q}[\nabla\mathcal{E}(x)[j]]>1$, then we have

$$ $\text{Var}[S_{K}(x_{t},t)[j]]>\frac{v_{0t}(x_{t})}{m^{2}_{t}(x_{t})K}=\text{ Var}[E_{K}(x_{t},t)]$ (40) $$

and thus $\textbf{tr}(\text{Cov}[S_{K}(x_{t},t]))>\text{Var}[E_{K}(x_{t},t)]$ holds.
$\square$

###### Proposition 2

## Appendix D Proof of Proposition 3

###### Proposition 3

Given a bootstrap trajectory $\{s_{i}\}_{i=0}^{n}$ such that $\sigma^{2}_{s_{i}}-\sigma^{2}_{s_{i-1}}\leq\kappa$, where $s_{0}=0$, $s_{n}=1$ and $\kappa>0$ is a small number, suppose $E_{\theta}$ is incrementally optimized from 0 0 to $1$ as follows: if $t\in[s_{i},s_{i+1}]$, then $E_{\theta}$ targets an energy estimator bootstrapped from $s\in[s_{i-1},s_{i}]$ using Eq. 8.
Let $v_{yz}(x_{z})=\text{Var}_{\mathcal{N}(x_{z},(\sigma_{z}^{2}-\sigma_{y}^{2})I)}
[\exp(-\mathcal{E}(x))]$ and $m_{z}(x_{z})=\exp(-\mathcal{E}_{z}(x_{z}))$, where $0\leq y<z\leq 1$ and $\mathcal{E}_{z}$ is the ground truth noised energy at time $z$. The variance of the bootstrap energy estimator is given by

$$ $\displaystyle{\text{Var}[E_{K}(x_{t},t,s;\theta)]}$ $\displaystyle=\frac{v_{st}(x_{t})}{v_{0t}(x_{t})}\text{Var}[E_{K}(x_{t},t)]$ (41) $$

and the bias of $E_{K}(x_{t},t,s;\theta)$ is given by

$$ $\displaystyle\frac{v_{0t}(x_{t})}{2m_{t}^{2}(x_{t})K^{n+1}}+\sum_{j=1}^{n} \frac{v_{0s_{j}}(x_{t})}{2m_{s_{j}}^{2}(x_{t})K^{j}}.$ (42) $$

Proof. The variance of $E_{K}(x_{t},t,s;\phi)$ can be simply derived by leveraging the variance of a sub-Gaussian random variable similar to Equation 17. While the entire proof for bias of $E_{K}(x_{t},t,s;\phi)$ is organized as follows:

- 1.
we first show the bias of Bootstrap($1$) estimator, which is bootstrapped from the system energy
- 2.
we then show the bias of Bootstrap($n$) estimator, which is bootstrapped from a lower level noise convolved energy recursively, by induction.

###### Proposition 3

### D.1 Bootstrap( 1 1 1 1 ) estimator

The Sequential estimator and Bootstrap($1$) estimator are defined by:

$$ $\displaystyle E^{\text{Seq}}_{K}(x_{t},t):$ $\displaystyle=-\log\frac{1}{K}\sum_{i=1}^{K}\exp(-E_{K}(x_{s|t}^{(i)},s)), \quad x_{s|t}^{(i)}\sim\mathcal{N}(x;x_{t},(\sigma_{t}^{2}-\sigma_{s}^{2})I)$ (43) $\displaystyle=-\log\frac{1}{K^{2}}\sum_{i=1}^{K}\sum_{j=1}^{K}\exp(-\mathcal{E }(x_{0|t}^{(ij)})),\quad x_{0|t}^{(ij)}\sim\mathcal{N}(x;x_{t},\sigma_{t}^{2}I)$ (44) $\displaystyle E_{K}^{B(1)}(x_{t},t,s;\phi):$ $\displaystyle=-\log\frac{1}{K}\sum_{i=1}^{K}\exp(-E_{\phi}(x_{s|t}^{(i)},s)), \quad x_{s|t}^{(i)}\sim\mathcal{N}(x;x_{t},(\sigma_{t}^{2}-\sigma_{s}^{2})I)$ (45) $$

The mean and variance of a Sequential estimator can be derived by considering it as the MC estimator with $K^{2}$ samples:

$$ $\displaystyle\mathbb{E}[E^{\text{Seq}}_{K}(x_{t},t)]$ $\displaystyle=\mathcal{E}_{t}(x_{t})+\frac{v_{0t}(x_{t})}{2m^{2}_{0t}(x_{t})K^ {2}}\quad\text{ and }\quad\text{Var}(E^{\text{Seq}}_{K}(x_{t},t))=\frac{v_{0t} (x_{t})}{m^{2}_{0t}(x_{t})K^{2}}$ (46) $$

While an optimal network obtained by targeting the original MC energy estimator 4 at $s$ is
(^1^11We consider minimizing the $L_{2}$-norm, i.e. $\phi^{*}=\arg\min_{\phi}\mathbb{E}_{x_{0},t}[\|E_{\theta}(x_{t},t)-E_{K}(x_{t} ,t)\|^{2}]$. Since the target, $E_{K}$, is noisy, the optimal outputs are given by the expectation, i.e. $E_{\phi}^{*}=\mathbb{E}[E_{K}]$.)
:

$$ $E_{\phi^{*}}(x_{s},s)=\mathbb{E}[E_{K}(x_{s},s)]=-\log m_{s}(x_{s})+\frac{v_{0 s}(x_{s})}{2m^{2}_{s}(x_{s})K}$ (47) $$

Then the optimal Bootstrap($1$) estimator can be expressed as:

$$ $\displaystyle E_{K}^{B(1)}(x_{t},t,s;\phi^{*})$ $\displaystyle=-\log\frac{1}{K}\sum_{i=1}^{K}\exp\left(-\left(-\log m_{s}(x_{s| t}^{(i)})+\frac{v_{0s}(x_{s|t}^{(i)})}{2m^{2}_{s}(x_{s|t}^{(i)})K}\right)\right)$ (48) $$

Before linking the Bootstrap estimator and the Sequential one, we provide the following approximation which is useful. Let $a$, $b$ two random variables and $\{a_{i}\}_{i=1}^{K}$, $\{b_{i}\}_{i=1}^{K}$ are corresponding samples. Assume that $\{b_{i}\}_{i=1}^{K}$ are close to 0 0 and concentrated at $m_{b}$, while $\{a_{i}\}_{i=1}^{K}$ are concentrated at $m_{a}$, then

$$ $\displaystyle\log\frac{1}{K}\sum_{i=1}^{K}\exp(-(a_{i}+b_{i}))$ $\displaystyle=\log\frac{1}{K}\left\{\sum_{i=1}^{K}\exp(-a_{i})\left[\frac{\sum _{i=1}^{K}\exp(-(a_{i}+b_{i}))}{\sum_{i=1}^{K}\exp(-a_{i})}\right]\right\}$ (49) $\displaystyle=\log\frac{1}{K}\sum_{i=1}^{K}\exp(-a_{i})+\log\frac{\sum_{i=1}^{ K}\exp(-(a_{i}+b_{i}))}{\sum_{i=1}^{K}\exp(-a_{i})}$ (50) $\displaystyle\approx\log\frac{1}{K}\sum_{i=1}^{K}\exp(-a_{i})+\log\frac{\sum_{ i=1}^{K}\exp(-a_{i})(1-b_{i})}{\sum_{i=1}^{K}\exp(-a_{i})}$ (51) $\displaystyle=\log\frac{1}{K}\sum_{i=1}^{K}\exp(-a_{i})+\log\left(1-\frac{\sum _{i=1}^{K}\exp(-a_{i})b_{i}}{\sum_{i=1}^{K}\exp(-a_{i})}\right)$ (52) $\displaystyle\approx\log\frac{1}{K}\sum_{i=1}^{K}\exp(-a_{i})-\frac{\sum_{i=1} ^{K}\exp(-a_{i})b_{i}}{\sum_{i=1}^{K}\exp(-a_{i})}$ (53) $\displaystyle\approx\log\frac{1}{K}\sum_{i=1}^{K}\exp(-a_{i})-m_{b}$ (54) $$

where Approximation applies a first order Taylor expansion of $e^{x}\approx 1+x$ around $x=0$ since $\{b_{i}\}_{i=1}^{K}$ are close to 0 0; while Approximation uses $\log(1+x)\approx x$ under the same assumption. Notice that when $K$ is large and $\sigma_{t}^{2}-\sigma_{s}^{2}\leq\kappa$ is small
, $\{\frac{v_{0s}(x_{s|t}^{(i)})}{2m^{2}_{s}(x_{s|t}^{(i)})K}\}_{i=1}^{K}$ are close to 0 0 and concentrated at $\frac{v_{0s}(x_{t})}{2m^{2}_{s}(x_{t})K}$. Therefore, by plugging them into Equation 54, Equation 48 can be approximated by

$$ $\displaystyle E_{K}^{B(1)}(x_{t},t,s;\phi^{*})$ $\displaystyle\approx-\log\frac{1}{K}\sum_{i=1}^{K}m_{s}(x_{s|t}^{(i)})+\frac{v _{0s}(x_{t})}{2m^{2}_{s}(x_{t})K}$ (55) $$

When $K$ is large and $\sigma_{s}^{2}$ is small, the bias and variance of $E_{K}(x_{s|t}^{(i)},s)$ are small, then we have

$$ $\displaystyle-\log\frac{1}{K}\sum_{i=1}^{K}m_{s}(x_{s|t}^{(i)})$ $\displaystyle\approx-\log\frac{1}{K}\sum_{i=1}^{K}E_{K}(x_{s|t}^{(i)},s)=E_{K} ^{\text{Seq}}(x_{t},t)$ (56) $$

Therefore, the optimal Bootstrap estimator can be approximated as follows:

$$ $\displaystyle E_{K}^{B(1)}(x_{t},t,s;\phi^{*})$ $\displaystyle\approx E_{K}^{\text{Seq}}(x_{t},t)+\frac{v_{0s}(x_{t})}{2m^{2}_{ s}(x_{t})K}$ (57) $$

where its mean and variance depend on those of the Sequential estimator (46):

$$ $\displaystyle\mathbb{E}[E_{K}^{B(1)}(x_{t},t,s;\phi^{*})]$ $\displaystyle=\mathcal{E}_{t}(x_{t})+\frac{v_{0t}(x_{t})}{2m^{2}_{t}(x_{t})K^{ 2}}+\frac{v_{0s}(x_{t})}{2m^{2}_{s}(x_{t})K}$ (58) $\displaystyle\text{Var}[E_{K}^{B(1)}(x_{t},t,s;\phi^{*})]$ $\displaystyle=\frac{v_{0t}(x_{t})}{m^{2}_{t}(x_{t})K^{2}}$ (59) $$

### D.2 Bootstrap( n 𝑛 n italic_n ) estimator

Given a bootstrap trajectory $\{s_{i}\}_{i=1}^{n}$ where $s_{0}=0$ and $s_{n}=s$, and $E_{\phi}$ is well learned at $[0,s]$. Let the energy network be optimal for $u\leq s_{n}$ by learning a sequence of Bootstrap($i$) energy estimators ($i\leq n$). Then
the optimal value of $E_{\theta}(x_{s},s)$ is given by $\mathbb{E}[E_{K}^{B(n-1)}(x_{s},s)]$. We are going to show the variance of a Bootstrap($n$) estimator by induction. Suppose we have:

$$ $\displaystyle E_{\phi^{*}}(x_{s},s)$ $\displaystyle=\mathcal{E}_{s}(x_{s})+\sum_{j=1}^{n}\frac{v_{0s_{j}}(x_{s})}{2m _{s_{j}}^{2}(x_{s})K^{j}}$ (60) $$

Then for any $t\in(s,1]$, the learning target of $E_{\theta}(x_{t},t)$ is bootstrapped from $s_{n}=s$,

$$ $\displaystyle E_{K}^{B(n)}(x_{t},t)$ $\displaystyle=-\log\frac{1}{K}\sum_{i=1}^{K}\exp(-E_{\phi^{*}}(x_{s|t}^{(i)},s )),\quad x_{s|t}^{(i)}\sim\mathcal{N}(x;x_{t},(\sigma_{t}^{2}-\sigma_{s}^{2})I)$ (61) $\displaystyle=-\log\frac{1}{K}\sum_{i=1}^{K}\exp\left(-\mathcal{E}_{s}(x_{s|t} ^{(i)})-\sum_{j=1}^{n}\frac{v_{0s_{j}}(x_{s|t}^{(i)})}{2m_{s_{j}}^{2}(x_{s|t}^ {(i)})K^{j}}\right)$ (62) $$

Assume that $\sigma_{t}^{2}-\sigma_{s}^{2}$ is small and $K$ is large, then we can apply Approximation 54 and have

$$ $\displaystyle E_{K}^{B(n)}(x_{t},t)$ $\displaystyle\approx-\log\frac{1}{K}\sum_{i=1}^{K}\exp(-\mathcal{E}_{s}(x_{s|t }^{(i)}))+\sum_{j=1}^{n}\frac{v_{0s_{j}}(x_{t})}{2m_{s_{j}}^{2}(x_{t})K^{j}}$ (63) $$

In Bootstrap($n$) setting, $\sigma_{s}^{2}$ is not small and we can’t approximate $\mathcal{E}_{s}(x_{s|t}^{(i)})$ simply by a MC estimator $E_{K}$. However, we can sequentially estimate such energy by bootstrapping through the trajectory $\{s_{i}\}_{i=1}^{n}$, resembling a Sequential($n$) estimator which is equivalent to $E_{K^{n+1}}$,

$$ $\displaystyle E_{K}^{B(n)}(x_{t},t)$ $\displaystyle\approx E_{K^{n+1}}(x_{t},t)+\sum_{j=1}^{n}\frac{v_{0s_{j}}(x_{t} )}{2m_{s_{j}}^{2}(x_{t})K^{j}}$ (64) $$

therefore, the optimal output of the energy network at $(x_{t},t)$ by learning this Bootstrap($n$) estimator is

$$ $\displaystyle\mathbb{E}[E_{K}^{B(n)}(x_{t},t)]$ $\displaystyle\approx\mathcal{E}_{t}(x_{t})+\frac{v_{0t}(x_{t})}{2m_{t}^{2}(x_{ t})K^{n+1}}+\sum_{j=1}^{n}\frac{v_{0s_{j}}(x_{t})}{2m_{s_{j}}^{2}(x_{t})K^{j}}$ (65) $$

which suggests that the accumulated bias of a Bootstrap($n$) estimator is given by

$$ $\displaystyle\frac{v_{0t}(x_{t})}{2m_{t}^{2}(x_{t})K^{n+1}}+\sum_{j=1}^{n} \frac{v_{0s_{j}}(x_{t})}{2m_{s_{j}}^{2}(x_{t})K^{j}}$ (66) $$

$\square$

## Appendix E Incorporating Symmetry Using NEM

We consider applying NEM and BNEM in physical systems with symmetry constraints like $n$-body system. We prove that our MC energy estimator $E_{K}$ is $G$-invariant under certain conditions, given in the following Proposition.

###### Proposition 4

Let $G$ be the product group $\text{SE}(3)\times\mathbb{S}_{n}\hookrightarrow O(3n)$ and $p_{0}$ be a $G$-invariant density in $\mathbb{R}^{d}$. Then the Monte Carlo energy estimator of $E_{K}(x_{t},t)$ is $G$-invariant if the sampling distribution $x_{0|t}\sim\mathcal{N}(x_{0|t};x_{t},\sigma_{t}^{2})$ is $G$-invariant, i.e.,

$$ $\mathcal{N}(x_{0|t};g\circ x_{t},\sigma_{t}^{2})=\mathcal{N}(g^{-1}x_{0|t};x_{ t},\sigma_{t}^{2}).$ $$

Proof.
Since $p_{0}$ is $G$-invariant, then $\mathcal{E}$ is $G$-invariant as well. Let $g\in G$ acts on $x\in\mathbb{R}^{d}$ where $g\circ x=gx$. Since $x_{0|t}^{(i)}\sim\overline{\mathcal{N}}(x_{0|t};x_{t},\sigma_{t}^{2})$ is equivalent to $g\circ x_{0|t}^{(i)}\sim\overline{\mathcal{N}}(x_{0|t};g\circ x_{t},\sigma_{t}
^{2})$. Then we have

$$ $\displaystyle E_{K}(g\circ x_{t},t)$ $\displaystyle=-\log\frac{1}{K}\sum_{i=1}^{K}\exp(-\mathcal{E}(g\circ x_{0|t}^{ (i)}))$ (67) $\displaystyle=-\log\frac{1}{K}\sum_{i=1}^{K}\exp(-\mathcal{E}(x_{0|t}^{(i)}))= E_{K}(x_{t},t)$ (68) $\displaystyle x_{(0|t)}^{(i)}$ $\displaystyle\sim\overline{\mathcal{N}}(x_{0|t};x_{t},\sigma_{t}^{2})$ (69) $$

Therefore, $E_{K}$ is invariant to $G=\text{SE}(3)\times\mathbb{S}_{n}$. $\square$

Furthermore, $E_{K}(x_{t},t,s;\phi)$ is obtained by applying a learned energy network, which is $G$-invariant, to the analogous process and therefore is $G$-invariant as well.

###### Proposition 4

## Appendix F NEM for General SDEs

Diffusion models can be generalized to any SDEs as $dx_{t}=f(x_{t},t)dt+g(t)d{{{w_{t}}}}$,
where $t\geq 0$ and $w_{t}$ is a Brownian motion. Particularly, we consider $f(x,t):=-\alpha(t)x$, i.e.

$$ $\displaystyle dx_{t}$ $\displaystyle=-\alpha(t)x_{t}dt+g(t)d{w_{t}}$ (70) $$

Then the marginal of the above SDE can be analytically derived as:

$$ $\displaystyle x_{t}$ $\displaystyle=\beta(t)x_{0}+\beta(t)\sqrt{\int_{0}^{t}(g(s)\beta(s))^{2} \mathrm{d}s}\epsilon,\quad\beta(t):=e^{-\int_{0}^{t}\alpha(s)\mathrm{d}s}$ (71) $$

where $\epsilon\sim\mathcal{N}(0,I)$. For example, when $g(t)=\sqrt{\bar{\beta}(t)}$ and $\alpha(t)=\frac{1}{2}\bar{\beta}(t)$, where $\bar{\beta}(t)$ is a monotonic function (e.g. linear) increasing from $\bar{\beta}_{\min}$ to $\bar{\beta}_{\max}$, the above SDE resembles a Variance Preserving (VP) process (Song et al.,, ). In DMs, VP can be a favor since it constrains the magnitude of noisy data across $t$; while a VE process doesn’t, and the magnitude of data can explode as the noise explodes. Therefore, we aim to discover whether any SDEs rather than VE can be better by generalizing NEM and DEM to general SDEs.

In this work, we provide a solution for general SDEs (70) rather than a VE SDE.
For simplification, we exchangeably use $\beta(t)$ and $\beta_{t}$. Given a SDE as Equation 70 for any integrable functions $\alpha$ and $g$, we can first derive its marginal as Equation 71, which can be expressed as:

$$ $\displaystyle\beta_{t}^{-1}x_{t}$ $\displaystyle=x_{0}+\sqrt{\int_{0}^{t}(g(s)\beta(s))^{2}\mathrm{d}s}\epsilon$ (72) $$

Therefore, by defining $y_{t}=\beta_{t}^{-1}x_{t}$ we have $y_{0}=x_{0}$ and therefore:

$$ $\displaystyle y_{t}$ $\displaystyle=y_{0}+\sqrt{\int_{0}^{t}(g(s)\beta(s))^{2}\mathrm{d}s}\epsilon$ (73) $$

which resembles a VE SDE with noise schedule $\tilde{\sigma}^{2}(t)=\int_{0}^{t}(g(s)\beta(s))^{2}ds$. In fact, we can also derive this by changing variables:

$$ $\displaystyle dy_{t}$ $\displaystyle=(\beta^{-1}(t))^{\prime}x_{t}dt+\beta^{-1}(t)\mathrm{d}x_{t}$ (74) $\displaystyle=\beta^{-1}(t)\alpha(t)x_{t}dt+\beta^{-1}(t)(-\alpha(t)x_{t}dt+g( t)\mathrm{d}{{{w_{t}}}})$ (75) $\displaystyle=\beta^{-1}(t)g(t)\mathrm{d}{{{w_{t}}}}$ (76) $$

which also leads to Equation 73. Let $\tilde{p}_{t}$ be the marginal distribution of $y_{t}$ and $p_{t}$ the marginal distribution of $x_{t}$,
with $y_{0|t}^{(i)}\sim\mathcal{N}(y;y_{t},\tilde{\sigma}^{2}_{t}I)$ we have

$$ $\displaystyle\tilde{p}_{t}(y_{t})$ $\displaystyle\propto\int{\exp(-\mathcal{E}(y))}\mathcal{N}(y_{t};y,\tilde{ \sigma}_{t}^{2}I)\mathrm{d}y$ (77) $\displaystyle\tilde{S}_{t}(y_{t})$ $\displaystyle=\nabla_{y_{t}}\log\tilde{p}_{t}(y_{t})\approx\nabla_{y_{t}}\log \sum_{i=1}^{K}\exp(-\mathcal{E}(y_{0|t}^{(i)}))$ (78) $\displaystyle\tilde{\mathcal{E}}_{t}(y_{t})$ $\displaystyle\approx-\log\frac{1}{K}\sum_{i=1}^{K}\exp(-\mathcal{E}(y_{0|t}^{( i)}))$ (79) $$

Therefore, we can learn scores and energies of $y_{t}$ simply by following DEM and NEM for VE SDEs. Then for sampling, we can simulate the reverse SDE of $y_{t}$ and eventually, we have $x_{0}=y_{0}$.

Instead, we can also learn energies and scores of $x_{t}$. By changing the variable, we can have

$$ $\displaystyle p_{t}(x_{t})$ $\displaystyle=\beta_{t}^{-1}\tilde{p}_{t}(\beta_{t}^{-1}x_{t})=\beta_{t}^{-1} \tilde{p}_{t}(y_{t})$ (80) $\displaystyle S_{t}(x_{t})$ $\displaystyle=\beta_{t}^{-1}\tilde{S}_{t}(\beta_{t}^{-1}x_{t})=\beta_{t}^{-1} \tilde{S}_{t}(y_{t})$ (81) $$

which provides us the energy and score estimator for $x_{t}$:

$$ $\displaystyle\mathcal{E}_{t}(x_{t})$ $\displaystyle\approx-\log\beta_{t}^{-1}\frac{1}{K}\sum_{i=1}^{K}\exp(-\mathcal {E}(x_{0|t}^{(i)}))$ (82) $\displaystyle S_{t}(x_{t})$ $\displaystyle\approx\beta_{t}^{-1}\nabla_{x_{t}}\log\sum_{i=1}^{K}\exp(- \mathcal{E}(x_{0|t}^{(i)}))$ (83) $\displaystyle x_{0|t}^{(i)}$ $\displaystyle\sim\mathcal{N}(x;\beta_{t}^{-1}x_{t},\tilde{\sigma}^{2}(t)I)$ (84) $$

Typically, $\alpha$ is a non-negative function, resulting in $\beta(t)$ decreasing from $1$ and can be close to 0 0 when $t$ is large. Therefore, the above equations realize that even though both the energies and scores for a general SDE can be estimated, the estimators are not reliable at large $t$ since $\beta_{t}^{-1}$ can be extremely large; while the SDE of $y_{t}$ (76) indicates that this equivalent VE SDE is scaled by $\beta_{t}^{-1}$, resulting that the variance of $y_{t}$ at large $t$ can be extremely large and requires much more MC samples for a reliable estimator. This issue can be a bottleneck of generalizing DEM, NEM, and BNEM to other SDE settings, therefore developing more reliable estimators for both scores and energies is of interest in future work.

## Appendix G Linking iDEM and iEFM Through Tweedie’s Formula

In this supplementary work, we propose TWEEdie DEM (TweeDEM), by leveraging the Tweedie’s formula (Efron,, ) into DEM, i.e. $\nabla_{x}\log p_{t}(x)=\mathbb{E}_{p(x_{0}|x_{t})}\left[\frac{x_{0}-x_{t}}{
\sigma_{t}^{2}}\right]$. Surprisingly, TweeDEM can be equivalent to the iEFM-VE proposed by Woo and Ahn, (), which is a variant of iDEM corresponding to another family of generative model, flow matching.

We first derive an MC estimator denoiser, i.e. the expected clean data given a noised data $x_{t}$ at $t$

$$ $\displaystyle\mathbb{E}[x_{0}|x_{t}]$ $\displaystyle=\int x_{0}p(x_{0}|x_{t})dx_{0}$ (85) $\displaystyle=\int x_{0}\frac{q_{t}(x_{t}|x_{0})p_{0}(x_{0})}{p_{t}(x_{t})}dx_ {0}$ (86) $\displaystyle=\int x_{0}\frac{\mathcal{N}(x_{t};x_{0},\sigma^{2}_{t}I)\exp(- \mathcal{E}(x_{0}))}{\exp(-\mathcal{E}_{t}(x_{t}))}dx_{0}$ (87) $$

where the numerator can be estimated by an MC estimator $\mathbb{E}_{\mathcal{N}(x_{t},\sigma^{2}_{t}I)}[x\exp(-\mathcal{E}(x))]$ and the denominator can be estimated by another similar MC estimator $\mathbb{E}_{\mathcal{N}(x_{t},\sigma^{2}_{t}I)}[\exp(-\mathcal{E}(x))]$, suggesting we can approximate this denoiser through self-normalized importance sampling as follows

$$ $\displaystyle D_{K}(x_{t},t):$ $\displaystyle=\sum_{i=1}^{K}\frac{\exp(-\mathcal{E}(x_{0|t}^{(i)}))}{\sum_{j=1 }^{K}\exp(-\mathcal{E}(x_{0|t}^{(j)}))}x_{0|t}^{(i)}$ (88) $\displaystyle=\sum_{i=1}^{K}w_{i}x_{0|t}^{(i)}$ (89) $$

where $x_{0|t}^{(i)}\sim\mathcal{N}(x_{t},\sigma^{2}_{t}I)$, $w_{i}$ are the importance weights and $D_{K}(x_{t},t)\approx\mathbb{E}[x_{0}|x_{t}]$. Then a new MC score estimator can be constructed by plugging the denoiser estimator $D_{K}$ into Tweedie’s formula

$$ $\displaystyle\tilde{S}_{K}(x_{t},t):=\sum_{i=1}^{K}w_{i}\frac{x_{0|t}^{(i)}-x_ {t}}{\sigma_{t}^{2}}$ (90) $$

where $\frac{x_{0|t}^{(i)}-x_{t}}{\sigma_{t}^{2}}$ resembles the vector fields $v_{t}(x_{t})$ in Flow Matching. In another perspective, these vector fields can be seen as scores of Gaussian, i.e. $\nabla\log\mathcal{N}(x;x_{t},\sigma^{2}_{t}I)$, and therefore $\tilde{S}_{K}$ is an importance-weighted sum of Gaussian scores while $S_{K}$ can be expressed as an importance-weighted sum of system scores $-\nabla\mathcal{E}$. In addition, Karras et al., () demonstrates that in Denoising Diffusion Models, the optimal scores are an importance-weighted sum of Gaussian scores, while these importance weights are given by the corresponding Gaussian density, i.e. $S_{\text{DM}}(x_{t},t)=\sum_{i}\tilde{w}_{i}({x_{0|t}^{(i)}-x_{t}})/{\sigma_{t
}^{2}}$ and
$\tilde{w}_{i}\propto\mathcal{N}(x_{0|t}^{(i)};x_{t},\sigma^{2}_{t}I)$. We summarize these three different score estimators as follows:

$\bullet$ DEM, sum of system scores, weighted by system energies:

$$ $S_{K}(x_{t},t)=\sum_{i}w_{i}(-\nabla\mathcal{E}(x_{0|t}^{(i)}))$ $$

$\bullet$ DDM, sum of perturbation kernel scores, weighted by perturbation densities:

$$ $S_{\text{DM}}(x_{t},t)=\sum_{i}\tilde{w}_{i}\nabla\log\mathcal{N}(x_{0|t}^{(i) };x_{t},\sigma_{t}^{2}I)$ $$

$\bullet$ EFM-VE=DEM+Tweedie, sum of perturbation kernel scores, weighted by system energies:

$$ $\tilde{S}_{K}(x_{t},t)=\sum_{i}w_{i}\nabla\log\mathcal{N}(x_{0|t}^{(i)};x_{t}, \sigma_{t}^{2}I)$ $$

## Appendix H Additional Experimental Details

### H.1 Energy functions

GMM. A Gaussian Mixture density in $2$-dimensional space with $40$ modes, which is proposed by Midgley et al., (). Each mode in this density is evenly weighted, with identical covariances,

$$ $\displaystyle\Sigma=\begin{pmatrix}40&0\\ 0&40\end{pmatrix}$ (91) $$

and the means $\{\mu_{i}\}_{i=1}^{40}$ are uniformly sampled from $[-40,40]^{2}$, i.e.

$$ $\displaystyle p_{\mathrm{gmm}}(x)=\frac{1}{40}\sum_{i=1}^{40}\mathcal{N}(x;\mu _{i},\Sigma)$ (92) $$

Then its energy is defined by the negative-log-likelihood, i.e.

$$ $\displaystyle\mathcal{E}^{\mathrm{GMM}}(x)$ $\displaystyle=-\log p_{gmm}(x)$ (93) $$

For evaluation, we sample $1000$ data from this GMM with $TORCH.RANDOM.SEED(0)$ following Midgley et al., (); Akhound-Sadegh et al., () as a test set.

DW-4. First introduced by Köhler et al., (), the DW-4 dataset describes a system with $4$ particles in $2$-dimensional space, resulting in a task with dimensionality $d=8$. The energy of the system is given by the double-well potential based on pairwise Euclidean distances of the particles,

$$ $\displaystyle\mathcal{E}^{\mathrm{DW}}(x)$ $\displaystyle=\frac{1}{2\tau}\sum_{ij}a(d_{ij}-d_{0})+b(d_{ij}-d_{0})^{2}+c(d_ {ij}-d_{0})^{4}$ (94) $$

where $a$, $b$, $c$ and $d_{0}$ are chosen design parameters of the system, $\tau$ the dimensionless temperature and $d_{ij}=\|x_{i}-x_{j}\|_{2}$ are Euclidean distance between two particles. Following Akhound-Sadegh et al., (), we set $a=0$, $b=-4$, $c=0.9$ $d_{0}=4$ and $\tau=1$, and we use validation and test set from the MCMC samples in as the “Ground truth” samples for evaluating.

LJ-n.
This dataset describes a system consisting of $n$ particles in $3$-dimensional space, resulting in a task with dimensionality $d=3n$. Following Akhound-Sadegh et al., (), the energy of the system is given by $\mathcal{E}^{\mathrm{Tot}}(x)=\mathcal{E}^{\mathrm{LJ}}(x)+c\mathcal{E}^{
\mathrm{osc}}(x)$ with the Lennard-Jones potential

$$ $\displaystyle\mathcal{E}^{\mathrm{LJ}}(x)$ $\displaystyle=\frac{\epsilon}{2\tau}\sum_{ij}\left(\left(\frac{r_{m}}{d_{ij}} \right)^{6}-\left(\frac{r_{m}}{d_{ij}}\right)^{12}\right)$ (95) and the harmonic potential $\displaystyle\mathcal{E}^{\mathrm{osc}}(x)$ $\displaystyle=\frac{1}{2}\sum_{i}\|x_{i}-x_{COM}\|^{2}$ (96) $$

where $d_{ij}=\|x_{i}-x_{j}\|_{2}$ are Euclidean distance between two particles, $r_{m}$, $\tau$ and $\epsilon$ are physical constants, $x_{\mathrm{COM}}$ refers to the center of mass of the system and $c$ the oscillator scale. We use $r_{m}=1$, $\tau=1$, $\epsilon=1$ and $c=0.5$ the same as Akhound-Sadegh et al., (). We test our models in LJ-13 and LJ-55, which correspond to $d=65$ and $d=165$ respectively. And we use the MCMC samples given by as a test set.

### H.2 Evaluation Metrics

2-Wasserstein distance $\mathcal{W}_{2}$. Given empirical samples $\mu$ from the sampler and ground truth samples $\nu$, the 2-Wasserstein distance is defined as:

$$ $\displaystyle\mathcal{W}_{2}(\mu,\nu)=(\text{inf}_{\pi}\int\pi(x,y)d^{2}(x,y) dxdy)^{\frac{1}{2}}$ (97) $$

where $\pi$ is the transport plan with marginals constrained to $\mu$ and $\nu$ respectively. Following Akhound-Sadegh et al., (), we use the Hungarian algorithm as implemented in the Python optimal transport package (POT) (Flamary et al.,, ) to solve this optimization for discrete samples with the Euclidean distance $d(x,y)=\|x-y\|_{2}$. $x-\mathcal{W}_{2}$ is based on the data and $\mathcal{E}-\mathcal{W}_{2}$ is based on the corresponding energy.

Total Variation (TV). The total variation measures the dissimilarity between two probability distributions. It quantifies the maximum difference between the probabilities assigned to the same event by two distributions, thereby providing a sense of how distinguishable the distributions are. Given two distribution $P$ and $Q$, with densities $p$ and $q$, over the same sample space $\Omega$, the TV distance is defined as

$$ $\displaystyle TV(P,Q)=\frac{1}{2}\int_{\Omega}|p(x)-q(x)|dx$ (98) $$

Following Akhound-Sadegh et al., (), for low-dimentional datasets like GMM, we use $200$ bins in each dimension. For larger equivariant datasets, the total variation distance is computed over the distribution of the interatomic distances of the particles.

### H.3 Experiment Settings

We set the reverse SDE integration steps for iDEM, NEM, and BNEM in our main experiment as $100$. In experiments for TweeDEM, we set this integration step $1000$ for TweeDEM and iDEM.

GMM-40.
For the basic model $f_{\theta}$, we use an MLP with sinusoidal and positional embeddings which has $3$ layers of size $128$ as well as positional embeddings of size $128$. The replay buffer is set to a maximum length of $10000$.

During training, the generated data was in the range $[-1,1]$ so to calculate the energy it was scaled appropriately by unnormalizing by a factor of $50$. All models are trained with a geometric noise schedule with $\sigma_{\min}=1e-5$, $\sigma_{\max}=1$, $K=500$ samples for computing both the MC score estimator $S_{K}$ and MC energy estimator $E_{K}$, $K=400$ samples for computing the Bootstrap energy estimator $E_{K}^{B}$ and Bootstrap score estimator $S_{K}^{B}=\nabla E_{K}^{B}$ and we clipped the norm of $S_{K}$ and $S_{K}^{B}$ to $70$.
All models are trained with a learning rate of $5e-4$.

DW-4. All models use an EGNN with $3$ message-passing layers and a $2$-hidden layer MLP of size $128$. All models are trained with a geometric noise schedule with $\sigma_{\min}=1e-5$, $\sigma_{\max}=3$, a learning rate of $1e-3$, $K=1000$ samples for computing $S_{K}$ and $E_{K}$, $K=400$ samples for computing $S_{K}^{B}$ and $E_{K}^{B}$, and we clipped $S_{K}$ and $S_{K}^{B}$ to a max norm of $20$.

LJ-13. All models use an EGNN with $5$ hidden layers and hidden layer size $128$. All models are trained with a geometric noise schedule with $\sigma_{\min}=0.01$ and $\sigma_{\max}=2$, a learning rate of $1e-3$, $K=1000$ samples for $S_{K}$ and $E_{K}$, $K=100$ samples for $E_{K}^{B}$ and $S_{K}^{B}$, and we clipped $S_{K}$ and $S_{K}^{B}$ to a max norm of $20$.

For all datasets. We use clipped scores as targets for iDEM and TweeDEM training for all tasks. Meanwhile, we also clip scores during sampling in outer-loop of training, when calculating the reverse SDE integral. These settings are shown to be crucial
especially when the energy landscape is non-smooth and exists extremely large energies or scores, like LJ-13 and LJ-55. In fact, targeting the clipped scores refers to learning scores of smoothed energies. While we’re learning unadjusted energy for NEM and BNEM, the training can be unstable, and therefore we blue smooth the Lennard-Jones potential through the cubic spline interpolation, according to Moore et al., (). Besides, we predict per-particle energies for LJ-n datasets. It shows that this setting can significantly stabilize training and boost performance.

## Appendix I Supplementary Experients

### I.1 Comparing the Robustness of Energy-Matching and Score-Matching

In this section, we discussed the robustness of the energy-matching model(NEM) with the score-matching model(iDEM) by analyzing the influence of the numbers of MC samples used for estimators and choice of noise schedule on the sampler’s performance.

Figure: Figure 2: Comparison of the Energy Wasserstein-2 distance in DW4 benchmark between iDEM and NEM across varying numbers of MC samples.
Refer to caption: extracted/5916790/figs/mc_samples_Ew2.png

As in Figure 2, NEM consistently outperforms iDEM when more than 100 MC samples are used for the estimator. Besides, NEM shows a faster decline when the number of MC samples increases. Therefore, we can conclude that the low variance of Energy-matching makes it more beneficial when we boost with more MC samples.

Figure: (a) Geometric: 4.16
Refer to caption: extracted/5916790/figs/robust/geo_dem.png

Figure: (a) Geometric: 2.64
Refer to caption: extracted/5916790/figs/robust/geo_endem.png

Then, we evaluate the performance differences when applying various noise schedules. The following four schedules were tested in the experiment:

- •
Geometric noise schedule: The noise level decreases geometrically in this schedule. The noise at step $t$ is given by: $\sigma_{t}=\sigma_{0}^{1-t}\cdot\sigma_{1}^{t}$
where $\sigma_{0}=0.0001$ is the initial noise level, $\sigma_{1}=1$ is the maximum noise level, and $t$ is the time step.
- •
Cosine noise schedule: The noise level follows a cosine function over time, represented by:
$sigma_{t}=\sigma_{1}\cdot\cos(\pi/2\frac{1+\delta-t}{1+\delta})^{2}$,
where $\delta=0.008$ is a hyper-parameter that controls the decay rate.
- •
Quadratic noise schedule: The noise level follows a quadratic decay:$\sigma_{t}=\sigma_{0}t^{2}$
where $\sigma_{0}$ is the initial noise level. This schedule applies a slow decay initially, followed by a more rapid reduction.
- •
Linear noise schedule: In this case, the noise decreases linearly over time, represented as:
$\sigma_{t}=\sigma_{1}t$

The experimental results are depicted in Figure 3 and Figure 4. It is pretty obvious that for iDEM the performance varied for different noise schedules. iDEM favors noise schedules that decay more rapidly to 0 0 when $t$ approaches 0.
When applying the linear noise schedules, the samples are a lot more noisy than other schedules. This also proves our theoretical analysis that the variance would make the score network hard to train.
On the contrary, all 4 schedules are able to perform well on NEM. This illustrates that the reduced variance makes NEM more robust and requires less hyperparameter tuning.

### I.2 Scaling up to Lennard Jones Potential

Figure: (a) iDEM
Refer to caption: extracted/5916790/figs/LJ13/dem.png

Besides GMM and DW-4, we would also like to show preliminary experimental results on the LJ-13 potentials. The distributions of interatomic distances and energies are shown in Figure 5.

The performance comparison between iDEM and NEM for the LJ-13 potential demonstrates the superior accuracy of NEM in capturing both the interatomic distance distribution and energy distribution. These results suggest that NEM’s capabilities allow it to generalize better in complex many-body systems, like the Lennard-Jones potential, providing more reliable energy estimations and physical interpretations.

### I.3 Empirical Analysis of the Variance of E K subscript 𝐸 𝐾 E_{K} italic_E start_POSTSUBSCRIPT italic_K end_POSTSUBSCRIPT and S K subscript 𝑆 𝐾 S_{K} italic_S start_POSTSUBSCRIPT italic_K end_POSTSUBSCRIPT

To justify the theoretical results for the variance of the MC energy estimator (4) and MC score estimator (2), we first empirically explore a $2$D GMM. For better visualization, the GMM is set to be evenly weighted by $10$ modes located in $[-1,1]^{2}$ with identical variance $1/40$ for each component, resulting in the following density

$$ $\displaystyle p^{\prime}_{GMM}(x)=\frac{1}{10}\sum_{i=1}^{10}\mathcal{N}\left( x;\mu_{i},\frac{1}{40}I\right)$ (99) $$

while the marginal perturbed distribution at $t$ can be analytically derived from Gaussian’s property:

$$ $\displaystyle p_{t}(x)=(p^{\prime}_{GMM}*\mathcal{N}(0,\sigma_{t}^{2}))(x_{t}) =\frac{1}{10}\sum_{i=1}^{10}\mathcal{N}\left(x;\mu_{i},\left(\frac{1}{40}+ \sigma_{t}^{2}\right)I\right)$ (100) $$

given a VE noising process.

We empirically estimate the variance for each pair of $(x_{t},t)$ by simulating $10$ times the MC estimators. Besides, we estimate the expected variance over $x$ for each time $t$, i.e. $\mathbb{E}_{p_{t}(x_{t})}[\text{Var}(E_{K}(x_{t},t))]$ and $\mathbb{E}_{p_{t}(x_{t})}[\text{Var}(S_{K}(x_{t},t))]$.

Figure: (a) Ground truth energy across $t\in[0,1]$
Refer to caption: extracted/5916790/figs/var_study/case-study-ground-truth.png

Figure 6(a) shows that, the variance of both MC energy estimator and MC score estimator increase as time increases. In contrast, the variance of $E_{K}$ can be smaller than that of $S_{K}$ in most areas, especially when the energies are low (see Figure 6(c)), aligning our Proposition 2. Figure 6(b) shows that in expectation over true data distribution, the variance of $E_{K}$ is always smaller than that of $S_{K}$ across $t\in[0,1]$.

### I.4 Experiments for TweeDEM

Figure: (d) $\tilde{S}_{K}$ (ours)
Refer to caption: extracted/5916790/figs/efm=dem+tweedie/ground_truth/gt_tweedie_gmm.png

**Table 2: Sampler performance comparison for GMM-40 and DW-4 energy function. we measured the performance using data Wasserstein-2 distance (x-$\mathcal{W}_{2}$), Energy Wasserstein-2 distance ($\mathcal{E}$-$\mathcal{W}_{2}$) and Total Variation (TV). $\dagger$ We compare the optimal number reported by Woo and Ahn, () and Akhound-Sadegh et al., (). - indicates metric non-reported. Bold indicates the best values.**
| Energy $\rightarrow$ | GMM-40 ($d=2$) | DW-4 ($d=8$) |  |  |  |  |
| --- | --- | --- | --- | --- | --- | --- |
| Sampler $\downarrow$ | x-$\mathcal{W}_{2}$$\downarrow$ | $\mathcal{E}$-$\mathcal{W}_{2}$$\downarrow$ | TV$\downarrow$ | x-$\mathcal{W}_{2}$$\downarrow$ | $\mathcal{E}$-$\mathcal{W}_{2}$$\downarrow$ | TV$\downarrow$ |
| $S_{K}$ | 2.864 | 0.010 | 0.812 | 1.841 | 0.040 | 0.092 |
| $\tilde{S}_{K}$ (ours) | 2.506 | 0.124 | 0.826 | 1.835 | 0.145 | 0.087 |
| iDEM$\dagger$ | 3.98 | - | 0.81 | 2.09 | - | 0.09 |
| iDEM (rerun) | 6.406 | 46.90 | 0.859 | 1.862 | 0.030 | 0.093 |
| iEFM-VE$\dagger$ | 4.31 | - | - | 2.21 | - | - |
| iEFM-OT$\dagger$ | 4.21 | - | - | 2.07 | - | - |
| TweeDEM (ours) | 3.182 | 1.753 | 0.815 | 1.910 | 0.217 | 0.120 |

In Appendix G, we propose TweeDEM, a variant of DEM by leveraging Tweedie’s formula (Efron,, ), which theoretically links iDEM and iEFM-VE and suggests that we can simply replace the score estimator $S_{K}$ (2) with $\tilde{S}_{K}$ (90) to reconstruct a iEFM-VE.
We conduct experiments for this variant with the aforementioned GMM and DW-4 potential functions.

Setting. We follow the ones aforementioned, but setting the steps for reverse SDE integration $1000$, the number of MC samples $500$ for GMM and $1000$ for DW-4. We set a quadratic noise schedule ranging from 0 0 to $3$ for TweeDEM in DW-4.

To compare the two score estimators $S_{K}$ and $\tilde{S}_{K}$ fundamentally, we first conduct experiments using these ground truth estimations for reverse SDE integration, i.e. samplers without learning. In addition, we consider using a neural network to approximate these estimators, i.e. iDEM and TweeDEM.

Table 2 reports x-$\mathcal{W}_{2}$, $\mathcal{E}{-}\mathcal{W}_{2}$, and TV for GMM and DW-4 potentials. We reports the numbers from Akhound-Sadegh et al., () and Woo and Ahn, (). We also rerun the iDEM and evaluate it for a more reliable comparisom.
Table 2 shows that when using the ground truth estimators for sampling, there’s no significant evidence demonstrating the privilege between $S_{K}$ and $\tilde{S}_{K}$. However, when training a neural sampler, TweeDEM can significantly outperform iDEM (rerun), iEFM-VE, and iEFM-OT for GMM potential. While for DW4, TweeDEM outperforms iEFM-OT and iEFM-VE in terms of $x-\mathcal{W}_{2}$ but are not as good as our rerun iDEM.

Figure 7 visualizes the generated samples from ground truth samplers, i.e. $S_{K}$ and $\tilde{S}_{K}$, and neural samplers, i.e. TweeDEM and iDEM. It shows that the ground truth samplers can generate well mode-concentrated samples, as well as TweeDEM, while samples generated by iDEM are not concentrated on the modes and therefore result in the high value of $\mathcal{W}_{2}$ based metrics. Also, this phenomenon aligns with the one reported by Woo and Ahn, (), where the iEFM-OT and iEFM-VE can generate samples more concentrated on the modes than iDEM.

Above all, simply replacing the score estimator $S_{K}$ with $\tilde{S}_{K}$ can improve generated data quality and outperform iEFM in GMM and DW-4 potentials. Though TweeDEM can outperform the previous state-of-the-art sampler iDEM on GMM, it is still not as capable as iDEM on DW-4. Except scaling up and conducting experiments on larger datasets like LJ-13, combing $S_{K}$ and $\tilde{S}_{K}$ is of interest in the future, which balances the system scores and Gaussian ones and can possibly provide more useful and less noisy training signals. In addition, we are considering implementing a denoiser network for TweeDEM as our future work, which might stabilize the training process.

## Peer Reviews

confidence: 4
contribution: 3
initial_rating: 8
presentation: 3
questions: Why is the distribution of interatomic distances for NEM and BNEM less sharply peaked than the ground truth for  LJ-55?

Have the method been tried on any larger, more complicated systems? Even if the results are not impressive, it would be useful for the field to know when this method breaks down.
soundness: 4
strengths: The paper describes the method very clearly. In particular, figure 1 is quite concise. Moreover, it shows proof that the improvements suggested are indeed improvements. Section 3 is well written, and the score vs. energy section contribution is valuable.

The results are strong and show improvement over other methods for the toy problems.
summary: The paper presents a neural sampling method for known energy functions. The importance of this problem is well-established, and the distinction between this and methods learned from sample data is sound. The NEM method improves on previous neural samplers by using an energy-based training objective, using a lower noise target, and introducing a bootstrapping method. The results show improvements over other listed neural samplers.
weaknesses: While the paper clearly shows improvement over previous neural sampling methods, it still makes little or no progress on the fundamental challenge of scaling to larger systems. The authors acknowledge this, but it ultimately make this an incremental improvement, not a transformational one.

---

confidence: 4
contribution: 3
initial_rating: 8
presentation: 3
questions: See "Weaknesses" above.
soundness: 3
strengths: The MS attacks an important problem and improves upon the state of the art.  The solution is original (to the best of this reviewer's knowledge) and provides theoretical as well as empirical reasons to prefer their proposal over alternatives.
summary: The authors consider a setting in which we want to draw samples from a Boltzmann distribution, for which we have the ground truth energy function, but cannot normalize into a probability distribution (for the usual reasons).  A very recent proposal (Akhound-Sadegh et al., 2024, "iDEM") for this problem is to model the "score function" (in Hyvarinen's sense) with a neural network, as in a diffusion model; but, lacking data samples, to target not the empirical score (via Tweedie's formula) but rather a Monte Carlo estimate written in terms of the energy function.  The MC samples are drawn in an outer loop by simulating the reverse diffusion process, using the current score-estimating network.

The present MS proposes to alter iDEM slightly, by modeling the *energy* rather than the score--again by targeting a MC estimate, in this case of the energy itself rather than its gradient.  They prove that (under certain conditions--I did not read the proof in detail) the error of iDEM's score estimator is strictly larger than the error of their own energy estimator; and (empirically) than the error in the score function that is produced by differentiating their energy estimator (and which matters because it will be used for sampling).  To reduce variance in this estimator further still, the authors further propose a "bootstrapping" technique in which the target energy is not constructed from the known data energy, but instead estimated from a model energy at a smaller noise level.

Empirically, the authors show that their approach yields superior samples to iDEM as well as other recent proposals, particularly with the "bootstrap" improvement.
weaknesses: (1) On the GMM-40 task, the manuscript's version of iDEM (Fig. 2) is visibly inferior to the one in the original paper (Akhound-Sadegh et al, Fig. 3), which is much closer to ground truth.  Something similar is true of the energy histogram for LJ-55 (Fig. 3 in this MS, Fig. 4 in op. cit.).  Can the authors explain these discrepancies?

(2) For DW-4, LJ-13, and LJ-55, the improvement in Wasserstein-2 distance provided by NEM and BNEM over iDEM appears marginal (Table 1).  Considering the additional computational cost of evaluating the energy gradient, this may not justify the proposed approach.

(3) The reported W2-E for iDEM is high, potentially due to outlier samples as pointed out in the manuscript. It would be nice to re-evaluate W2-E removing outliers for comparison with NEM and BNEM.  

(4) Learning energy instead of score may improve estimation, but sampling requires costly repeated gradient evaluations as steps increase.  A table comparing sampling times for score-based (iDEM) and energy-based (NEM) training would be helpful.

(5) The descriptions of the problem setting and proposed solution could be much improved; this reviewer, at least, found the exposition in Akhound-Sadegh et al. (2024) much clearer.

(6) The Fisher divergence (Eq. 8) is a well understood quantity, with connections to KL divergence---e.g., in the context of diffusion models, it arises naturally under the standard ELBO loss.  Can the authors provide any similar theoretical appeal for the squared error in energy, Eq. 10?



MINOR  
Diffusion models were introduced by Sohl-Dickstein et al. in "Deep unsupervised learning using nonequilibrium thermodynamics" (ICML, 2015) and this paper should be referenced when they are introduced in the MS.

---

confidence: 4
contribution: 2
initial_rating: 3
presentation: 3
questions: See above.
soundness: 2
strengths: The paper is generally well-written and easy to follow. The authors provide enough background to make the paper easy to understand. The topic of diffusion-based samplers is also an important and interesting topic with applications in many domains. Additionally,  improving the MC estimator of the score when we have access to the energy and not data is, in my opinion, very well-motivated.
summary: In this work, the authors propose Noised Energy Matching (NEM) and Bootstrap NEM (BNEM), a diffusion-based neural sampler based on previous related work (iDEM). In particular, the authors propose to parametrise the energy with a network, instead of the score, to reduce the variance of the estimator. They also propose to bootstrap the energy estimator at high noise levels from learned estimators at lower noise levels to further reduce the variance.
weaknesses: **Novelty and Motivation**: 
Although the NEM algorithm is a simple reparametrization of the estimator in iDEM, the paper analyzes the implications of an energy or score parametrization and  shows that this simple change can lead to a better estimator. The more novel algorithm is BNEM, however, the paper does not provide enough experimental results to support the benefit of this method (see my notes under Experiments).

- Theoretical Limitations: The work makes several claims regarding the variance of the estimators which aren’t precise. For example, the claim that “the variances of EK and SK explode at high noise levels as a result of the VE noising.” This isn’t correct as the variance is a function of multiple factors, including the importance sampling proposal chosen for the estimator, as well as the energy landscape itself.
Additionally, the proof of proposition 2 considers low-energy regions and claims that $exp{(\mathcal{E}(x_{0|t}^{(i)})} \geq c$. However, it isn’t clear to me how these regions are defined as the energy is evaluated at the noised samples.

**Experiments**: 
- The work contains several experiments on the same datasets and benchmarks as the previously related paper (iDEM), showing better performance of the method, especially when the number of MC samples and integration steps are limited. However, I find that additional experiments, to show the benefit of lowering the variance at the cost of increasing the bias would strengthen the claims of the paper much more. For example, the performance gain from BNEM, at the cost of increasing the bias isn’t clear in my opinion. Generally, neural networks are fairly good at training in high-variance settings (e.g. diffusion models). On the other hand, for complex energy distributions and in high dimensional settings, the bias could result in significant problems. In general, I think the authors should provide more analysis of the trade-off of bias and variance when introducing the BNEM estimator.
- The paper also claims that the proposed approach is more robust compared to iDEM as it doesn’t rely on clipping the score. However, for the LJ potentials cubic spline interpolation was used, as the energy is very sharp. It isn’t clear to me how much of an advantage this provides if smoothing is still necessary and how applicable this is to other complex and realistic settings.
- The results for iDEM on LJ-13 and LJ-55 tasks are significantly different from what iDEM paper reported. The authors also indicate divergent training for DDS, PIS and FAB on those tasks, while iDEM reported competitive values for both FAB and PIS on LJ-13 and for FAB on LJ-55. Could the author elaborate on how and why they were unable to reproduce the results?
- Some metrics such as NLL and ESS which are commonly reported in relevant works are not reported.
- In Table 3, the standard deviations over the seeds isn’t reported, and the authors report the best value, making the results misleading. For example, they indicate that the mean E-W2 of LJ-55 over the 3 seeds is in fact even higher than the one they report for iDEM. This in fact confirms my concern that improving the variance of the estimator with bootstrapping at the cost of increasing the performance doesn’t necessarily lead to increased performance.

---

confidence: 4
contribution: 2
initial_rating: 3
presentation: 3
questions: 1. To point in the weaknesses, can you give an estimate of the additional overhead with the change from score to energy parameterization? Specifically, to better understand the practical utility, an empirical comparison of the it the decrease of the variance per compute unit would need to be needed.
2. The authors speculate on the use of a Metropolis-Hastings type correction in future work. How do they envision doing this within the current BNEM framework?
3. Is it possible to correct for distributional bias of the neural sampler e.g. through importance sampling? it is with many of the base-lines you compare against. If so, what is the computational cost of this compared to the baselines?
4. Why did the authors not try a standard molecular benchmark such as ALA2? Is it related to the smoothing that needs to be performed to get the LJ-n systems to train stably? If this is the case, how do you envision moving this methodology beyond toy-systems?
soundness: 3
strengths: 1. The paper is clearly written and addresses an important problem
2. The authors are transparent about the limitations of the work and in particular the challenges that remain for with regard to training stability etc.
summary: Neural sampler of Boltzmann distributions. A variation of iDEM with better variance. The central idea of iDEM and this paper is to fit a diffusion model which generates samples which are approximately Boltzmann distributed mu(x) \propto e^(-E(x)) where E(x) is the energy function. In iDEM model is learned by matching a score network to the gradient of a monte-carlo estimator of the noised energy at a given time, t. This estimator is simulation-free but very high variance, and a bi-level training scheme is introduced to stabilize training.

The key innovation in this paper is to match the energy instead of the gradient of the energy (force/score), which they show in Propositions 1 and 2, has smaller error and variance than training against the score. A second contribution in this paper using a bootstrapping scheme to further reduce variance, at the cost of introducing some bias.
weaknesses: 1. Apart from the memory overhead highlighted by the authors in their limitations section, the new methods also introduce significant computational overhead, by having to use autograd to compute the derivative of the energy during sampling. 
2. Experiments are performed on toy systems and not systems which reflect the actual use cases of methods like this, e.g. molecules. A simple case should be included, such as ALA2.
3. Missing comparisons to standard sampling baselines for non-normalized densities such as parallel tempering. In particular, considering the number of evaluations of E, and overall compute cost.

Meta-review: This paper proposes a method to learn a Boltzmann sampler by learning the noise energy function to eventually get an estimation of the noise score function to sample efficiently according to the Boltzmann distribution. They provide theory bounding the variance of the energy function and provide experiments on a toy Mixture of Gaussians datasets, LJ-13, and LJ-55.

The concerns of the reviewers were the following:
1. The theory bounds the variance of the stochastic target of the energy function. However, since the quantity of interest is the score, It is not clear that having a lower variance in the score target and then differentiate the estimate is better than havnig a larger variance directly estimating the gradients. Hence the claim 
> (NEM) has lower variance ... than previous related works. 

is somewhat misleading since $\nabla E_K = S_K$ (the gradient of the estimator proposed in this work is exactly the one proposed in iDEM so I would argue that we cannot really compare the variances of this work and the previous related works)

2. Secondly, two reviewers had some concerns about the experiments. In particular, I agree that the authors should be able to reproduce the baselines to reasonable performance or at least take some time to explain why they think they could not reproduce it after having tried a significant amount of time. 

3. The NEM methods introduce a computational overhead by requiring the computation of the gradient of the energy function during sampling. Even though the computational overhead is reasonable, it remains a disadvantage of the method. 

Overall, even though I recommend rejection, I believe that a future iteration of this paper with a more substantial experimental section (in particular, with respect to the baselines) would be a significant contribution to the community.

Decision: Reject