Efficient Training of Energy-Based Models Using Jarzynski Equality
Authors: Davide Carbone
Mengjian Hua
Simon Coste
Eric Vanden-Eijnden
Venue: NeurIPS 2023 poster (2023)
Keywords: Energy-Based Model
contrastive learning
generative models
Jarzynski identity
ULA

Abstract: Energy-based models (EBMs) are generative models inspired by statistical physics with a wide range of applications in unsupervised learning.  Their performance is well measured by the cross-entropy (CE) of the model distribution relative to the data distribution. Using the CE as the objective for training is however challenging because the computation of its gradient with respect to the model parameters requires sampling the model distribution. Here we show how results for nonequilibrium thermodynamics based on Jarzynski equality together with tools from sequential Monte-Carlo sampling can be used to perform this computation efficiently and avoid the uncontrolled approximations made using the standard contrastive divergence algorithm.  Specifically, we introduce a modification of the unadjusted Langevin algorithm (ULA) in which each walker acquires a  weight that enables the estimation of the gradient of the cross-entropy at any step during GD, thereby bypassing sampling biases induced by slow mixing of ULA. We illustrate these results with numerical experiments on Gaussian mixture distributions as well as the MNIST and CIFAR-10 datasets. We show that the proposed approach outperforms  methods based on the contrastive divergence algorithm in all the considered situations.

Efficient Training of Energy-Based Models Using Jarzynski Equality
Davide Carbone
Dipartimento di Scienze Matematiche, Politecnico di Torino Istituto Nazionale di Fisica Nucleare, Sezione di Torino davide. carbone@polito..t
Mengjian Hua
Courant Institute of Mathematical Sciences, New York University mh51 aa51130nyu.edi
Simon Coste
LPSM, Université Paris-Cité simon. www.totsstaaiissso..ooou
Eric Vanden-Eijnden
Courant Institute of Mathematical Sciences, New York University eve20nyu. edu
Abstract
Energy-based models (EBMs) are generative models inspired by statistical physics with wide range of applications in unsupervised learning. Their performance is well measured by the cross-entropy (CE) of the model distribution relative to the data distribution. Using the CE as the objective for training is however challenging because the computation of its gradient with respect to the model parameters requires sampling the model distribution. Here we show how results for nonequilibrium thermodynamics based on Jarzynski equality together with tools from sequential Monte Carlo sampling can be used to perform this computation efficiently and avoid the uncontrolled approximations made using the standard contrastive divergence algorithm. Specifically, unadjusted Langevin algorithm (ULA) in which each walker acquires a weight that enables the estimation of the gradient of the cross-entropy at any step during GD, thereby bypassing sampling biases induced by slow mixing of ULA. We illustrate these results with numerical experiments well as the MNIST and CIFAR- 10 datasets. We show that the proposed approach outperforms methods based the contrastive divergence algorithm in all the considered situations.
introduce modification of the
Gaussian mixture distributions
Introduction
Probabilistic models have become key tool in generative artificial intelligence (AI) and unsupervised learning. Their goal is twofold: explain the training data, and allow the synthesis of new samples. Many flavors have been introduced in the last decades, including variational auto-encoders 2
generative adversarial networks [4, normalizing flows 66 7 g 10,, diffusion-based models
22
13
restricted Boltzmann machines [14] 15 [6], and energy-based models (EBMs) [1 18
37th Conference on Neural Information Processing Systems (NcurIPS 2023).

---

Inspired by physics, EBMs are unnormalized probability models, specified via an energy function U, with the underlying probability density function (PDF) defined as p exp(-U)/2, where Z is a normalization constant: this constant is often intractable but crucially it is not needed for data generation via Monte-Can Carlo sampling. In statistical physics, such PDFs are called Boltzmann- densities 20 21 22 and their energy function is often known in advance; in the context of EBMs, the aim is to estimate this energy function in some parametric class using the available data. Such models benefit from century of intuitions from computational physics to guide the design of the energy U, and help the sampling; hence, interpretability is a clear strength of EBMs, since in many applications the parameters of the model have direct meaning. Unfortunately, training EBMS is a challenging task as it typically requires to sample complex multimodal (i.e. non-log-concave distributions in high-dimension.
This issue arises because real-world datasets are often clustered into different modes with imbalanced weights. A skewed estimation of these weights can lead to harmful biases when it comes to applica- tions: for example, under-representation of elements from particular population. It is important to ensure that these biases are either avoided or compensated, leading methods that ensure fairness ((23); however, the training routines for EBMs often have hard time properly learning the weights of different modes, and can often completely fail at this task.
Two popular classes of methods are used to train an EBM. the first, the Fisher divergence between the model and the data distributions is taken as the objective to minimize, which amounts to learning the gradient of the energy function, VU- = V log p. Although computationally efficient due to score-matching techniques (24), methods in this class are provably unable to learn the relative weights of modes when they are separated by low-density regions.
In the second class of methods, the cross-entropy between the model and the data distributions is used as the objective. Unlike the Fisher divergence, the cross-entropy is sensitive the relative weights in multimodal distributions, but it unfortunately leads to training algorithms that are less justified theoretically and more delicate to use in practice. Indeed, gradient-based optimization procedures on the cross-entropy require sampling from the model distribution using Markov Markov-Chain Monte-Carlo (MCMC) methods or the unadjusted Langevin algorithm (ULA) until mixing, which in a high- -ithdimmeniionll context or without logcconcavity can be prohibitive, even considering the large computational power available today.
As a result, one typically needs to resort to approximations to estimate the gradient of the cross- entropy, for example by using the contrastive divergence (CD) or the persistent contrastive divergence (PCD) algorithms, see (25) 26 27.. Unfortunately, the approximations made in these algorithms are uncontrolled and they are known to induce biases similar to those observed with score-based methods (the CD algorithm reduces to gradient descent over the Fisher divergence in the limit of frequent resetting from the data (28). Many techniques have been proposed overcome this issue, for example, based on MCMC sampling (29 30, 33 or on various regularizations of CD [32 [3). In practice though, these techniques still do not handle well multimodal distributions and they come with little theoretical guarantees.
Main contributions. In this work, we go back to the original problem of training EBMS using the cross-entropy between the model and the data distributions objective, and:
We derive exact expressions for the cross-entropy and its gradient that involve expectations over an ensemble of weighted walkers whose weights and positions evolve concurrently with the model energy during optimization. Our main tool is Jarzynski equality (34. an exact relation
between the normalizing constants of two distributions linked by an out-of-equilibrium dynamics.
Based on these formulas, we design sequential Monte Carlo sampling procedures [35] to estimate the cross-entropy and its gradient in a practical way at essentially no additional computational cost compared to using the CD or PCD algorithms with ULA.
We show that reweighting the walkers is necessary in general and that training procedures based on the CD algorithm lead to uncontrolled biases whereas those based on using the PCD algorithm lead to mode collapse in general.
We illustrate these results numerically on synthetic examples involving Gaussian mixture models where these effects can be demonstrated.
2

---

We also apply our method on the MNIST and CIFAR-10 datasets. In the first, we intentionally bias the proportion of the different digits, demonstrate that our method allows the retrieval of these proportions whereas the other methods fail to do so.
Related works. How to train EBMs is a longstanding question, and we refer e.g. to 17 36 for
overviews on this topic. The use of modern deep neural architectures to model the energy was then proposed in [37], and a relation between EBMs and classifiers has been highlighted in [38] How to sample from unnormalized probability models is also an old and rich problem, see 13 40 for
general introductions.
Score-matching techniques and variants originate from 24 41 42]; their shortcoming in the context of EBM training is investigated in 44] and their blindness to the presence of multiple, imbalanced modes in the target density has been known for long: we refer to 44 45] for discussions. Contrastive
divergence (CD) algorithms originate from 25
These methods only perform one or a few sampling steps of the algorithm, with random walkers that are repeatedly restarted at the data points. Persistent contrastive divergence (PCD) algorithm, introduced in (29]. eliminates the restarts and evolves walkers using ULA. Unlike the approach proposed in this paper, these methods are known to give estimates of the gradient of the cross-entropy that have uncontrolled biases which are difficult to remove; there were many attempts in this direction, also in cooperation with other unsupervised
26 27
techniques, see e.g. [46 47 48
49
50
51 52..
It is worth noting that the original papers on CD proposed to use an objective which is different from the cross-entropy, and their proposed implementation (what we call the CD algorithm) does not perform gradient descent on this objective, due to some gradient terms being neglected [33]. This sparked some debate about which objective, if any, was actually minimized by CD algorithms: for example, [53] 54] showed that the CD and PCD algorithms are essentially adversarial procedures; introduced a way to approximate the missing term in the gradient of the CD objective; and showed that in the limit of small noise, CD is essentially equivalent to score matching. In contrast, there is no ambiguity about the objective used the method we propose: it is always the cross-entropy.
33
28
55
Jarzynski equality (JE) was introduced in [34] and gives an exact expression relating the normaliz- ing constants (or equivalently free energy 22 of two distributions linked by out-of-qquiibrium continuous- time dynamics. A discrete-time analog of JE is used in Neal's annealed importance sampling [56], which belongs to the framework of sequential Monte-Carlo methods (35]. These meth- ods have been used in the context of generative models based on variational autoencoders [57] (58).
normalizing flows [59] 60), and diffusion based models [6]] 62]. In contrast, here we use sequential
Monte-Carlo method to train EBM on the cross-entropy directly.
Energy-Based Models
Setup, notations, and assumptions. The problem we consider can be formulated as follows: assume that we are given n E N data points {x; [=))== in Rd drawn from an unknown probability distribution that is absolutely continuous with respect to the Lebesgue measure on Rd with a positive probability density function (PDF) p. (x) > 0 (also unknown). Our aim is to estimate this PDF via an energy-based model (EBM), ie. to find a suitable energy function a parametric class, Uo Rd + oc) with parameters A E O, such that the associated Boltzmann- PDF
pe(x) Z3 Uo(())
Ze
e Uo(u) dx
(I)
is an approximation of the target density (x The normalization factor Zg is known as the partition function in statistical physies [22]. This factor is hard estimate and one advantage of EBMs is that they provide generative models that do not require the explicit knowledge of Za since Markov Chain Monte-Carlo (MCMC) methods can in principle be used sample pe knowing only Up the design of such MCMC methods is an integral part of the problem of building an EBM.
To proceed will assume that the parametric class of energy we use is such that, for all 0 E O,
Ug (RR);;
3L E R4 |VV((()) S Vo WERR.
(2)
3a E and a compact set E Rd x (Ub(2) 2 2| Vz V C.
3

---

These assumptions guarantee that Ze < oo (i.e. we can associate PDF pe to Ug via ( for any 0 () and that the Langevin equations as well as their time-discretized versions we will use to sample Po have global solutions and are ergodic 63
64
65
We stress that (2) does not imply that
Up is convex (i.e. that po log-coneave): in fact, we will be most interested in situations where Up has multiple local minima so that Pe is multimodal. For simplicity we will also assume that is in the parametric class, i.c. 30. E : pg. = p.. Our aims are primarily to identify this 0, and to sample p9. in the process, we will also show how to estimate Ze..
Cross-entropy minimization. To measure the quality of the EBM and train its parameters one can the eross-entropy of the model density po relative to the target density p.
H (ps: P.) (pp. - -fe log po (=))) (x)dd log Zo + /
(z)d)
Rd
where we used the definition of in ① to get the second equality. The cross-entropy is related the Kullback-Leible entropy of p.. and its gradient with respect to the parameter 0 can be calculated using the identity
divergence via H (pe, p.) Hp.)) + DKL (px(px)lpe). where Hp..) is the
On log Ze = - Su OgU0(x) (xU (o())) to obtain
0 Hp P-) = [ (1U)(=)) - / (pU
= [.,Uo Es (0gUo.
The cross-entropy is stringent, and therefore better, than other objectives like the Fisher divergence: for example, unlike the latter, it is sensitive to the relative probability weights of modes on p. separated by low-density regions (36). Unfortunately, the cross entropy is also much harder to use in practice since evaluating it requires estimating Za. and evaluating its gradient requires calculating the expectation IEg [0gUa] (in contrast E., [Ug] and IE, [8gUe] can be readily estimated the data). Typical training methods, c.g. based on the CD or the PCD algorithms, give up estimating Ze and resort to various approximations for more discussion about these methods. While these approaches have proven successful in many situations, they are prone to training instabilities that limit their applicability. They also come with no theoretical guarantees in terms of convergence.
calculate the expectation Es [agUp]]
Appendix A..5
3 Training via sequential Monte-Carlo methods based on Jarzynski equality
In this section we use tools from nonequilibrium statistical mechanics sions for both Eg [0gU] and Za (Sec 3.1 that are amenable empirical estimation via sequential Monte-Carlo methods [35], thereby enabling gradient descent-typp type algorithms for the optimization of EBMs (Sec. 3.2p.
34 56 to write exact expres-
3.1 Jarzynski equality in discrete-time
Proposition 1. Assume that the parameters 0 are evolved by some time-discrete protocol {0a JReN and that 2 hold. Given any h E L), let X E I& and Ag IR be given by the iteration rule
Xx+1 Xx AVU,, (Xx) + 2hh
Xo p0u
(5)
AR+1 = A
Xx) + ax Xx-1); Ao = 0,
where (2) is the model energy, [ are independent N N(0,, Ia), and we defined
(e(xx))= = Uo. (x) + (y - x) VUo. (x) + U (x))²
(6)
Then, for all k E No.
E(D0Up, Ax Ble:An]
Es [0qUo_]
Zo. = La [e
(7)
where the expectations the right-hand side are over the law of the joint process (Xe,

---

Algorithm 1 Sequential Monte-Carlo training with Jarzynski correction
1: Inputs: data points {-² {==)== energy model Up: optimizer step opt(O, D) using 0 and the empirical CE gradient D; initial parameters lo: number of walkers N E No: set of walkers
(X sampled from Peg total duration K E N; ULA time step h; set of positive constants {CKE
2: Aj = for ,,... ...
3: for k = 0.... ......KK K do
4:
Pk = =exp(A)) Σj=1 exp(Ak))
normalized weights
5:
D = Σx=1 P/U (X - 2==100U0.
empirical CE gradient
6:
OK+I - opt(OR, Dz)
optimization step
7:
for i = 1, ..... .N do
8:
E NoOs,, A 4 = A ax++1 (X xx) au (x xx+1)
X = x NVU,, (
D ULA weight update
9:
10: end for
11: Resample the walkers and reset the weights if ESS < Ck+1, see
resampling step
12:
end for
13:
Outputs: Optimized energy Uox set weighted walkers (Xk:
k: A - sampling pox partition
function estimate Zox = Zo, N E exp(AAx); CE estimate log Zon + + Σ5= Uox (a!)
The proof of the proposition is given in Appendix A.2 for completeness also give continuous- time version of this proposition in Appendix [.I]] We stress that the inclusion of the weights in ( is key, E Ei0 Un, (Xx)] + Es, [0,Uo,] in general. We also stress that (7) holds exactly despite the fact that for h >> 0 the iteration step for X5 in (5) is that of the unadjusted Langevin algorithm (ULA) with Metropolis correction in MALA [66 667. That is, the inclusion of the weights Ag exactly corrects for the biases induced by both the slow mixing and the time-discretization errors in ULA.
Proposition D shows that can evolve the parameters by gradient descent over the cross-entropy by solving 5 concurrently with
D2 = -0g Pp) El0gUee (Xx))) [0gUa.]
0k+1 = 0x + DD:
(8)
EleAz]
where 7k is the learning rate and k E No with 6, given. We can also replace the gradient step for 0g in (8) by any update optimization step (via AdaGrad, ADAM, etc.) that uses as input the gradient Da of the cross-entropy evaluated 0₂ to get 0x+1. Assuming that know Zo, we can track the evolution of the cross-entropy via
(p P=) log E A[cAA] + log Zo, + (Uo(-
3.2 Practical implementation
Empirical estimators and optimization step. We introduce N independent pairs of walkers and
weights, {Xú, J which we evolve independently using (5) for each pair. To evolve 0₂ from some prescribed 0g we can then use the empirical version of 8
0k+1 = 6 + 7 Dr.
(10)
where D is the estimator for the gradient in A of the cross-entropy:
D3 ΣN N OoUg, (X exp(Ai)
n 3 (x!))
(11)
Σn (A
These steps are summarized in Algorithm I which is specific instance of a sequential Monte-Carlo algorithm. We can also use mini-batches of {Xū, A 1N and the data set {a4 In at every iteration (see Algorithm 2 in Appendix A.3), and switch to any optimizer step that uses 02 and Dj. input to get the updated U++ During the calculation, we can monitor the evolution of the partition function and the cross-entropy using as estimators
1 Zor Ze N
1 Σ
exp( (A))
H1 log Zo. +
Upx (x!)
(12)
5

---

where a,, b. IRd specify the means of the two modes and E R controls their relative weights P. 1/(1 + e=²" and 9. 1 p. =-= e= The values of a, b., &x are carefully chosen such that the modes are well separated and the energy barrier between the modes high These steps are summarized in Algorithm] [ which is specific instance of a sequential Monte-Carke algorithm. We adapted our routine to mini-batcher in Algorithm 2 without any explicit additional source of error: in fact, the particles outside the mini-batch have their weights updated too. The only sources of error in these algorithms come from the finite sample sizes, < 0o and n < oo. Regarding we may need to add regularization term in the loss to avoid overfitting: this is standard. Regarding N, we need to make sure that the effective sample size of the walkers remains sufficient during the evolution. This is nontrivial since the A 's will spread away from zero during the optimization. implying that the weights exp( A4)) will become non- non-unifone uniform, thereby reducing the effective sample size. This a known issue with sequential Monte-Carle algorithms that can be alleviated by resampling as discussed next.
Resampling step. A standard quantity to monitor the effective sample size [40] is the ratio between the square of the empirical mean of the weights and their empirical variance, L...
(N- Di= exp(AA))² 2 E 1
N ESS N EE exp()A
(13)
The effective sample size of the N walkers is ESS, Initially, since A O, ESS, = 1, but decreases with k. At each iteration k,- such that ESSk, < Ck,, where {cr JeN is set of predefined positive constants in (0, 1), we then:
Resample the walkers ky using ph. = Ahr /iz11 Ahr as probability to pick walker
2. Reset A = 0:
3. Use the update Ze. = Zee, N E exp(A)) for 2 k, until the next resampling step.
This resampling standard (35) and can be done with various levels of sophistication, as discussed in Appendix A.4 Other criteria, based e.g. on the entropy of the weights, are also possible, see e.g. [68)
Generative modeling. During the training stage, i.e. as the weights A evolve, the algorithm produces weighted samples X At any iteration, however, equal-weight samples can be generated by resampling. Notice that, even if we no longer evolve the model parameters 6z, the algorithm is such that removes the bias from ULA coming from because N is finite, but this can be controlled by increasing N at the stage when the EBM is used as a generative model.
this bias removal is not perfect, again
Numerical experiments
4.1 Gaussian Mixtures
In this section, we use a synthetic model to illustrate the advantages of our approach. Specifically, we assume that the data is drawn from the Gaussian mixture density with two modes given by
px(x)= z = (e +==--=-
+ -10-b611-22)
Z. = -22))d/2 (1+c-*)
(14)
enough such that jumps of the walkers between the modes are not observed during the simulation with ULA. Consistent with (14 we use an EBM with
Up(r) = log le
al +e - -))
(15)
where 0 - (a, b, =) are the parameters to be optimized. We choose this model as it allows to calculate the partition function of the model at any value of the parameters, 2. (2x)d/2 (1 + e=*). We use this information as benchmark to compare the prediction with those produced by our method.

---

Mass the first mode
Means
Kl. divergence
19
24
Algo PCD
62
Algo
Rescaled Iterations
107 Rescaled lterations
Rescaled Iterations
Figure 1: GMM experiments: Evolution of the parameters and the cross entropy during training by Algorithm 2 PCD, and CD. Average of 20 runs. Left panels: evolution of ps: = 1/(1 middle panel: cvolution of ag and by: right panel: evolution of the Dollback-Leibler divergence. All three methods capture the location of the modes accurately, but only ours get the relative weights of these modes accurately (whereas PCD leads to mode collapse, and CD to an inaccurate estimate). Our method is also the only one that allows for direct estimation of the cross-entropy during training, and the only one performing GD on this tross-entropy-forr for better visualization we subtract the entropy of the target H (px) and plot the Kullback-Leible divergence instead of the cross-entropy.
In our numerical experiments, set d = 50, use N 10 walkers with mini-batch of N' 10 and n = 10 data points. We initialize the model at 8g = (ag. bu, 20) with dg and by drawn from an N (0, Ig) with 0.1 and 20 O, meaning that the initial peo is close to the PDF of an N La)- 0.2 for
The training is performed using Algorithm with / = 0.1 and fixed learning rates and and 7% for zg. We perform the resampling step by monitoring ESS,, defined in (T3) with constant 1/ck- = 1.05 and using the systematic method. We also compare our results to those obtained using ULA with these same parameters (which akin to training with the PCD algorithm) and with those obtained with the CD algorithm: in the latter case, we evolve the walkers by ULA with = 0.1 for 4 steps between resets at the data points, and we adjust the learning rates by multiplying them by a factor 10. In all cases, we use the full batches of walkers, weights, and data points to estimate the empirical averages. We also use (12 to estimate the cross-entropy H(o P. during training by our method (CD and PCD do not provide estimates for these quantities), and in all cases compare the result with the estimate
n n E1iog log fe +e - -all (16)
Hg log (27))/22 (1 (1+ (1+0-2)))
The results are shown in Figure I As can be seen, all three methods learn well the values of a, and be specifying the positions of the modes. However, only our approach learns the value of specifying their relative weights. In contrast, the PCD algorithm leads to mode collapse, consistent with the theoretical explanation given in Appendix [C.1 and the CD algorithm returns a biased value of 2, the objective. The results
consistent with the fact that it effectively uses the Fisher divergence also show that the cross-entropy decreases with our approach, but bounces back up with the PCD) algorithm and stalls with the CD algorithms: this is consistent with the fact that only approach actually performs the GD on the cross-entropy, which, unlike the other algorithms, our approach estimates accurately during the training.
4.2 MNIST
Next, perform empirical experiments the MNIST dataset to answer the following question: when it comes to high- sigh-dimensional datasets with multiple modes, can our method produces an EBM that generates high-quality samples and captures the relative weights of the modes accurately?
To this end, we select a subset of MNIST consisting of only three digits: 2, 3, and 6. Then, we choose 5600 images of label 2, 2800 images of label 3, and 1400 images of label 6 from the training set (for a total of = 9800 data points), so that in this manufactured dataset the digits are in have respective weights 4/7, 2/7, and 1/7.
7

---

2
6
3
3
3
2
2
3
2
2
3
2
2
2
2
2
2
3
2
2
2
7
2
2
2
3
3
2
3
6
2
6
3
6
2
2
6
6
2
2
6
3
2
3
3
3
2
3
6
3
2
3
2
3
2
2
2
2
2
3
6
3
2
3
3
2
3
6
6
3 2
3
4
2
2
2
3
3
3
6
3
6
6
6
6
6
3
3
2
2
6
2
3
2
6
2
3
2
6
6
3
2
2
2
6
6
2
3
2
3
6
2 2
2
7
3
7
6
A
Generated samples (right resampling)
Images from MNIST set
Generated samples (without Jarzynski)
Figure 2: MNIST: Left panel: Examples of images generated by our method right after resampling in the last epoch. Middle panel: Images randomly selected from the test dataset of MNIST. Right panel: Examples of images generated by training using the persistent contrastive divergence (PCD) algorithm.
Error weight estimation (digt
Error the estimation (digit
the weight estimation
w/o jerzynski
jarzynski
Jerzyrski
w/ jarzynski
w/ jarzynsk
175
wi jarzynsk
Relative
1.50
Epoch
Epoch
Epoch
Figure 3: MNIST: Relative error of the weight estimation of the three modes (i.e. three digits). Our method outperforms the PCD algorithm in terms of recovery of the weight each mode.
We train two EBMs to represent this data set, the first using our Algorithm] 2 and the second using the PCD Algorithm We represent the energy using a simple six-layer convolutional neural network with the swish activation and about 77K parameters. We use the ADAM optimizer for the training with a learning rate starting from 10-4 and linearly decaying 10-10 until the final training step. The sample size of the walkers is set to N 1024 and it is fixed throughout the training.
and The results shown in Figure Both our Algorithm and the PCD Algorithm Aenerate images of reasonable quality; as discussed in Appendix [B.2] we observe that the Jarzynski weights can help us track the quality of generated images, and the resampling step is also helpful for improving this quality as well as the learned energy function.
However, the real difference between the methods comes when we look at the relative proportion of the digits generated in comparison to those in the data set. In Figure we show the relative error on the weight estimation of each mode obtained by using classifier pre-trained Appendix B.2 details). Although the EBM trained with the PCD algorithm can discover all the modes and generate images of all three digits present in the training set, it cannot accurately recover the relative proportions of each digit. In contrast, our method is successful in both mode exploration and weight recovery.
the data set (see
Unlike in the numerical experiments done on Gaussian mixtures with teacher-student models, for the MNIST dataset, we cannot estimate the KL divergence throughout the training normalization constant of the true data distribution. Nevertheless, we can still use the formula (16) to track the cross-entropy between the data distribution and the walker distribution plot of which is given in the right panel Figure 4 for the mini- mini-batched training algorithm (Algorithm 2)
do not know the
In these experiments with MNIST, we used the adaptive resampling scheme described after equa- tion 13 practice, we observed that few resamplings are needed during training, and they can often by avoided altogether if the learning rate is sufficiently small. For example, in the experiment reported in the left panel of Figure single resampling step was made. We also found empirically that the results are insensitive to the choice of of the parameters eg used for resampling: the rule of
8

---

Convargence of the cross antropy estmmte
Effective Sample Size
Estimated crocs entroyy
rross-entropy
Effective Sample Size
Rearpiie
Effective Sample Size
Epocts
Figure 4: MNIST dataset: Left Panel: Convergence of the cross-entropy estimated by using formulae with the mini-batched algorithm (Algorithm 2). Right Panel: Evolution of the effective sample size (ESS) defined in (13) here resampling was started after 240 epochs (with c₁ = 0 before and C 0.5 afterwards), and occurred only once immediately after being switched on.
thumb we found is to not resample at the beginning of training, to avoid possible mode collapse, and resampling towards the end, to improve the image quality.
4.3 CIFAR-10
perform an empirical evaluation our method on the full CIFAR-IO (32 x 32) image dataset. We use the setup of [47]. with the same neural architecture with n 128 features, and compare the results obtained with our approach with mini-batching (Algorithm 2 to those obtained with PCD and PCD with data augmentation of [33] (which consists of combination of color distortion, horizontal flip, rescaling, and Gaussian blur augmentations, to help the mixing of the MCMC sampling and stabilize training).
Generated CIFAR-10 samples with approach Generated CIFAR-10 samples with PCD
Figure 5: CIFAR-1O dataset: Left panel: images generated by training with Algorithm 2 Right
panel: Images generated from the PCD with mini-batches. Up to date images available at https 7/gi thub. com/ 1ataasiooee//eeee/eaeee

---

The code used to perform these new experiments is available in the anonymized GitHub referenced in our paper. The hyperparameters are the same in all cases: we take N = 4096 Langevin walkers with a mini-batch size N" = 256. We use the Adam optimizer with learning rate 1j 10-4 and inject a Gaussian noise of standard deviation o 3 x 10-2 to the dataset while performing gradient clipping in Langevin sampling for better performance. All the experiments were performed on single A100 GPU. Training for 600 epochs took about 34 hours with the PCD algorithm (w/ and w/o data augmentation) and about 36 hours with our method.
Some of the images generated by our method are shown in Figure 5 We also quantitatively evaluate the performance of our models with the commonly used metrics (e.g. FID and Inception Score): the results are given in Table They indicate that our method can achieve slightly better performance than the PCD algorithms w/ and w/o data augmentation at similar computational cost. Furthermore, these results on CIFAR-10 suggest that our method scales well to complicated training tasks on more realistic data sets.
Method
FID
Inception Score (IS)
PCD with mini-batches
38.25
5.96
PCD with mini-batches and data augmentation
36.43
6.54
Algorithm 2 with multinomial resampling 32.18
6.88
Algorithm with systematic resampling
30.24
6.97
Table I: textitCIFAR-10 dataset: Comparison of FID and Inception Score (IS) for PCD and Algorithm Experiments performed using the neural architecture in 47 to model the energy.
worth stressing that a key component of the success of our method the resampling step, as discussed in Appendix B.3 where we plot the Effective Sample Size (ESS) in Figure 13
5 Concluding Remarks
In this work, we proposed a simple modification of the persistent contrastive divergence algorithm which corrects its biases and provably allows to minimize the cross-entropy between the model and the target densities to train EBMs. Our approach rests on results from non-equilibrium ther- modynamics (in particular, Jarzynski equality) that show how to take exact expectations over probability density function (PDF) that is evolving in time-in the present context the PDF is that of the EBM, and it is changing as we train the model parameters. These formulas naturally lend themselves to practical implementation using sequential Monte Carlo sampling methods. The only difference with respect to training methods based on ULA is that the proposed approach maintains set of weights associated with every walker. These weights correct for the bias induced by the evolution of the energy, and they also give a way to assess the quality of the samples generated by the method. On synthetic examples using Gaussian mixture densities as well as on some simple data sets involving MNIST, we showed that our method dodges some common drawbacks of other EBM training methods. In particular, it is able to learn the relative weights of various high-density regions in a multimodal distribution, thus ensuring higher level of fairness than common EBM training techniques based on contrastive divergence.
To focus on the core idea, in the present study we put aside tuning considerations and did not sweep over the various hyperparameters and design choices, like learning rates, number of walkers, or resampling criteria and strategies. These considerations could greatly enhance the performance of our method. We believe that our results show that the method deserves deeper practical investigation on more realistic datasets and with more complex neural architectures.
10

---

Acknowledgements
D.C. has worked under the auspices of Italian National Group Mathematical Physies (GNFM). M.H. is supported by NYU McCracken doctoral fellowship. EVE is supported by the National Science Foundation under awards DMR- 1420-14,,73, DMS-2012510, and DMS-2134216 by the Simons Collaboration on Wave Turbulence, Grant No. 617006, and by a Vannevar Bush Faculty Fellowship.
References
[1] Diederik P Kingma and Max Welling. An introduction to variational autoencoders. Foundations and Trends in Machine Learning, 12(4):307-392, 2019.
[2] Carl Doersch. Tutorial on variational autoencoders. arXiv preprint arXiv: 1606. 05900, 2016.
[3] Shengjia Zhao, Jiaming Song. and Stefano Ermon. Infovae: Balancing learning and inference in variational autoencoders. In Proceedings of the aaai conference on artificial intelligence, volume 33, pages 5885-5892, 2019.
[4] Jonathan Ho and Stefano Ermon. Generative adversarial imitation learning. In D. Lee, M. Sugiyama, U. Luxburg, I. Guyon, and R. Garnett, editors, Advances in Neural Information Processing Systems, volume 29, 2016.
[5] Ian Goodfellow, Jean Pouget-Abadie, Mehdi Mirza, Bing Xu, David Warde-Farley, Sherjil Ozair, Aaron Courville, and Yoshua Bengio. Generative adversarial networks. Communications of the ACM, 63(11):139-144, 2020.
[6] Esteban G. Tabak and Eric Vanden-Eijnden. Density estimation by dual ascent of the log- likelihood. Communications in Mathematical Sciences, 8(1):21711111888111))))) 233,, 2010.
[7] E. G. Tabak and Cristina V. Turner. A family of nonparametric density estimation algorithms. Communications Pure and Applied Mathematics, 66(2):145-164 2013.
[8] Danilo Rezende and Shakir Mohamed. Variational inference with normalizing flows. In International conference machine learning, pages 1530-1538. PMLR, 2015.
[9] George Papamakarios, Eric Nalisnick, Danilo Jimenez Rezende, Shakir Mohamed, and Balaji Lakshminarayanan. Normalizing flows for probabilistic modeling and inference. The Journal of Machine Learning Research, 22(1):2617-2680, 2021.
[10] Emile Mathieu and Maximilian Nickel. Riemannian continuous normalizing flows. Advances in Neural Information Processing Systems, 33:2503-2515, 2020.
[11] Yang Song. Jascha Sohl-Dickstein, Diederik P Kingma, Abhishek Kumar, Stefano Ermon, and Ben Poole. score-based generative modeling through stochastic differential equations. In International Conference on Learning Representations, 2021.
[12] Jonathan Ho, Ajay Jain, and Pieter Abbeel. Denoising diffusion probabilistic models. In Advances in Neural Information Processing Systems, volume 33, pages 6840-6851, 2020.
[13] Yang Song, Conor Durkan, Iain Murray, and Stefano Ermon. Maximum likelihood training of score-based diffusion models. Advances Neural Information Processing Systems, 34: 1415-1428, 2021.
[14] Geoffrey E Hinton. A practical guide training restricted boltzmann machines. Neural Networks: Tricks of the Trade: Second Edition, pages 599-619, 2012.
[15) Ruslan Salakhutdinov and Hugo Larochelle. Efficient learning of deep boltzmann machines. In Proceedings of the thirteenth international conference on artificial intelligence and statistics, pages 693-700. JMLR Workshop and Conference Proceedings, 2010.
[16] Hannes Schulz, Andreas Müller, Sven Behnke, et al. Investigating convergence of restricted boltzmann machine learning. NIPS 2010 Workshop on Deep Learning and Unsupervised Feature Learning. volume 1, pages 6-1, 2010.

---

[17] Yann LeCun, Sumit Chopra, and Raia Hadsell. A tutorial on energy-based learning. In Gökhan BakIr, Thomas Hofmann, Alexander J Smola, Bernhard Schölkopf, and Ben Taskar, editors, Predicting structured data, chapter 10. MIT press, 2007.
[18] Michael Gutmann and Aapo Hyvärinen. Noise-con estimation: A new estimation principle for unnormalized statistical models. In Proceedings of the thirteenth international con- ference artificial intelligence and statistics, pages 297-304. JMLR Workshop and Conference Proceedings, 2010.
[19] Yang Song, Sahaj Garg, Jiaxin Shi, and Stefano Ermon. Sliced score matching: A scalable approach to density and score estimation. In Uncertainty in Artificial Intelligence, pages 574-584, PMLR, 2020.
[20] Ludwig Boltzmann. Weitere studien über das wärmegleichgewicht unter gasmolekülen. Kinetis- che Theorie II, pages 115-225, 1970.
[21] Josiah Willard Gibbs. Elementary principles in statistical mechanics: developed with especial reference to the rational foundations of thernodynamics. C. Scribner's sons, 1902.
[22] Evgenii Mikhailovich Lifshitz and Lev Petrovich Pitaevskii. Statistical physics: theory of the condensed state, volume 9, Elsevier, 2013.
[23] Ninareh Mehrabi, Fred Morstatter, Nripsuta Saxena, Kristina Lerman, and Aram Galstyan. A survey on bias and fairness in machine learning. ACM Computing Surveys (CSUR), 54(6):1-35, 2021.
[24] Aapo Hyvärinen and Peter Dayan. Estimation of non-normalized statistical models by score matching. Journal of Machine Learning Research, 6(4), 2005.
[25] Geoffrey Hinton. Training products of experts by minimizing contrastive divergence. Neural compulation, 14(8):1771-1800, 2002.
[26] Max Welling and Geoffrey E Hinton. A new learning algorithm for mean field boltzmann machines. In International conference on artificial neural networks, pages 351-357, 2002.
[27] Miguel A Carreira-Perpinaa and Geoffrey Hinton. contrastive divergence learning. In International workshop on artificial intelligence and statistics, pages 33-40, PMLR, 2005.
(28) Aapo Hyvarinen. Connections between score matching, contrastive divergence, and pseudo- likelihood for continuous-valuce variables. IEEE Transactions on neural networks, 18(5): 3529-1531, 2007.
[29] Tijmen Tieleman. Training restricted boltzmann machines using approximations to the likeli- hood gradient. International conference on Machine learning, pages 1064-1071, 2008
[30] Pierre Jacob, John O'Leary, and Yves Atchadé. Unbiased markov chain monte carlo methods with couplings. Journal of the Royal Statistical Society: Series B (Statistical Methodology), 82
(3), 2020.
[31] Yixuan Qiu, Lingsong Zhang, and Xiao Wang. Unbiased contrastive divergence algorithm for training energy-basse based latent variable models. In International Conference on Learning Representations, 2020.
[32] Yilun Du and Igor Mordatch. Implicit generation and modeling with energy based models. In Advances in Neural Information Processing Systems, 2019.
[33] Yilun Du, Shuang Li, Joshua Tenenbaum, and Igor Mordatch. Improved contrastive divergence training of easrg models. In International Conference on Machine Learning, pages 2837-2848. PMLR, 2021.
[34] C Jarzynski. Nonequilibrium equality for free energy differences. Physical Review Letters, 78 (14):2690, 1997.
[35] Arnaud Doucet, Nando De Freitas, Neil James Gordon, et Sequential Monte Carlo methods in practice, volume 1. Springer, 2001.
12

---

[36] Yang Song and Diederik P Kingma. How to train your energy energy-based models. arXiv preprint arXiv:2101.03288, 2021.
[37] Jianwen Xie, Yang Lu, Song-Chun Zhu, and Yingnian Wu. A theory of generative convnet. In International Conference on Machine Learning, pages 2635-2644, PMLR, 2016.
[38] Will Grathwohl, Kuan-Chich Wang, Joem--Herik Jacobsen, David Duvenaud, Mohammad Norouzi, and Kevin Swersky. Your classifier is secretly energy based model and you should treat it like one. In International Conference on Learning Representations, 2019.
[39] Steve Brooks, Andrew Gelman, Galin Jones, and Xiao-Li Meng. Handbook of Markov chain Monte Carlo. CRC press, 2011.
[40] Jun S Liu and Jun S Liu. Monte Carlo strategies in scientific computing, volume 75. Springer, 2001.
[41] Pascal Vincent. A connection between score matching and denoising autoencoders. Neural computation, 23(7):1661-1674, 2011.
[42] Kevin Swersky, Mare' Aurelio Ranzato, David Buchman, Nando D Freitas, and Benjamin M Marlin. On autoencoders and score matching for energy based models. In International conference on machine learning (ICML-11), pages 1201-1208, 2011.
[43] Li Kevin Wenliang. On the failure of variational score matching for VAE models. arXiv preprint arXiv:2210. 13390, 2022.
[44] Li Wenliang, Danica J Sutherland, Heiko Strathmann, and Arthur Gretton. Learning deep kernels for exponential family densities. In International Conference on Machine Learning. 2019.
[45] Yang Song and Stefano Ermon. Generative modeling by estimating gradients of the data distribution. Advances in neural information processing systems, volume 32, 2019.
[46] Jianwen Xie, Yang Lu, Ruiqi Gao, Song-Chun Zhu, and Ying Nian Wu. Cooperative training of descriptor and generator networks. IEEE transactions on pattern analysis and machine intelligence, 22(1):27-45, 2018.
[47] Erik Nijkamp. Mitch Hill, Song- Lang-Chun Zhu, and Ying Nian Wu. Learning non-convergent non- persistent short-run MCMC toward energy-based model. In Advances in Neural Information Processing Systems, volume 32, 2019.
[48] Ruiqi Gao, Erik Nijkamp, Diederik P Kingma, Zhen Xu, Andrew M Dai, and Ying Nian Wu. Flow contrastive estimation of energy-based models. In Proceedings of the IEEF/CVF Conference on Computer Vision and Pattern Recognition, pages 7518-7528, 2020.
[49] Zhisheng Xiao, Karsten Kreis, Jan Kautz, and Arash Vahdat. Vaebm: A symbiosis between variational autoencoders and energy-based models. In International Conference on Learning Representattons, 2020.
[50] Jianwen Xie, Zilong Zheng, and Ping Li. Learning energy-based model with variational auto- encoder amortized sampler. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 35, pages 10441-,0451, 2021.
[51] Jianwen Xie, Yaxuan Zhu, Jun Li, and Ping Li. tale of two flows: Cooperative learning of langevin flow and normalizing flow toward energy-basse model. In International Conference on Learning Representations, 2021.
[52] Hankook Lee, Jongheon Jeong, Sejun Park, and Jinwoo Shin. Guiding energy-based mod- els via contrastive latent variables. In The Eleventh International Conference on Learning Representations, 2022.
[53] Omer Yair and Tomer Michaeli. Contrastive divergence learning a time reversal adversarial game. In International Conference on Learning Representations, 2021.
13

---

[54] Carles Domingo-Enriche Alberto Bietti, Eric Hunden-Eijnden and Joan Bruna. On energy based models with overparametriiee shallow neural networks. In International Conference on Machine Learning, pages 2771-2782, 2021.
[55] Carles Domingo-Enrich, Alberto Bietti, Marylou Gabrié, Joan Bruna, and Eric Vanden-Eijnden. Dual training of energy-based models with overparametrized shallow neural networks. arXiv preprint arXiv:2 14,.05114, 2021.
[56] Radford M Neal. Annealed importance sampling. Statistics and computing, 11:125-133 139, 2001.
[57] Tuan Anh Le, Maximilian Igl, Tom Rainforth, Tom Jin, and Frank Wood. Auto-encoding sequential monte carlo. In International Conference on Learning Representations, 2018. URL
[58] Xinqiang Ding and David J. Freedman. Learning deep generative models with annealed importance sampling, 2020.
[59] Laurence Illing Midgley, Vincent Stimper, Gregor N. C. Simm, Bernhard Schölkopf, and José Miguel Hernádeez-Lobato Flow annealed importance sampling bootstrap. In The Eleventh International Conference on Learning Representations, 2023. URL https://openreviin
het:/forum72d-xCCVPPP...
[60] Yilun Du, Conor Durkan, Robin Strudel, Joshua B Tenenbaum, Sander Dieleman, Rob Fergus, Jascha Sohl-Dickstein, Arnaud Doucet, and Will Sussman Grathwohl. Reduce, reuse, recycle: Compositional generation with energy-based diffusion models and meme. International Conference Muchine Learning, pages 8489-8510. PMLR, 2023.
[61] Jascha
Eric Weiss, Niru Maheswaranathan, and Surya Ganguli. Deep unsu- pervised learning using nonequilibrium thermodynamics. In Francis Bach and David Blei, editors, Proceedings of the 32nd International Conference Machine Learning, volume 37 of Proceedings of Machine Learning Research, pages 2256-2265, Lille, France, 07 Jul 2015. 1htmmmmm11m.mmm1m11mmmmmm
PMLR. htttpt::////:ooceeeeee
mlr wtt....eeee///////ooo----- wateesss////shh-aaiaaaa
[62] Arnaud Doucet, Will Sussman Grathwohl, Alexander G. D. G. Matthews, and Heiko Strathmann. Score-based diffusion meets annealed importance sampling. In Advances in Neural Information Processing Systems, 2022.
[63] Bernt Oksendal. Stochastic Differential Equations. Springer- Verlag Berlin Heidelberg, 6 edition, 2003.
[64] J. C. Mattingly, A. M. Stuart, and D. J. Higham. Ergodicity for SDEs and approximations: locally Lipschitz vector fields and degenerate noise. Stochastic Processes and their Applications, 101(2):185-23) October 2002.
[65] Denis Talay and Luciano Tubaro. Expansion of the global error for numerical schemes solving stochastic differential equations. Stochastic Analysis and Applications, 8(4):483-509 1990.
[66] Gareth O. Roberts and Richard L. Tweedie. Exponential convergence of langevin distributions and their discrete approximations. Bernoulli, 2(4):341-363, 1996. ISSN 13507265. URL httpp://ww. jstor org/ stt///stteeee33/111////
[67] Gareth O Roberts and Jeffrey S Rosenthal. Optimal scaling of discrete approximations to langevin diffusions. Journal of the Royal Statistical Society: Series B (Statistical Methodology). 60(1):255-268, 1998.
[68] Pierre Del Moral, Arnaud Doucet, and Ajay Jasra. On adaptive resampling strategies for sequential Monte Carlo methods. Bernoulli, 18(1):252 278, 2012.
[69] Lawrence C Evans. Partial differential equations, volume 19. American Mathematical Society, 2022.
[70] Neil J Gordon, David J Salmond, and Adrian FM Smith. Novel approach to nonlinear/non- gaussian bayesian state estimation. IEE proceedings (radar and signal processing). volume 140, pages 107-113. IET, 1993.
14

---

[71] Genshiro Kitagawa. Monte carlo filter and smoother for non gaussian nonlinear state space models, Journal of computational and graphical statistics, 5(1):1-25, 1996.
[72] James Carpenter, Peter Clifford, and Paul Fearnhead. Improved particle filter for nonlinear problems. IEE Proceedings-Rader, Sonar and Navigation, 146(1):22 (46(1))2-7, 1999.
[73] Tiancheng Li, Miodrag Bolic, and Petar M Djurie. Resampling methods for particle filtering: classification, implementation, and strategies. IEEE Signal processing magazine, 32(3):70-86,
2015.
[74] Michael S Albergo. Nicholas M Boffi, and Eric Panden-Eiinden. Stochastic interpolants: A unifying framework for flows and diffusions. arXiv preprint arXiv:2303.08797, 2023.
15

---

A Additional theoretical results
A.1 Jarzynski equality in continuous-time
Before proving Proposition |
we give continuous-tine version of this result whose proof helps
guide the intuition.
Proposition 2. Assume that the parameters 0 are evolved according to some time-differentiable protocol O(1) such that 0(0) = Og- Given any a > 0, let Xt E IR and A, E RR be the solutions of
dXx: = - VV (Xi)) +
Xu ~ p0a
(17)
A. = -n (X.) Ó(1),
=
where (x) is the model energy and Wi E R is standard Wiener process. Then, for any 2 0,
Boco|OopUccol = ([00
Ze(c) Zo, Ble"),
(18)
BleAi]
where the expectations on the right-hand side are over the law of the joint process (X,, Ar).
The second equation in (88) can also be written in term of the free energy Fo
log Zo as Fa(e) = Fgg log E[e this is Jarzynski's ' equality [34]. We stress that it is key to include the weights (18) and, particular, E(0 (X.)] + Eetee [8gUs(e]] This is because the PDF of X, alone lags behind the model PDF Patey: the larger the smaller this lag, but it is always there if a < 0o, see the remark at the end of this section for more discussion on this point. The inclusion of the weights in (18) corrects exactly for the bias induced by this lag.
An immediate consequence of Proposition 2 is that we can evolve f(t) by the gradient descent flow over the cross-entropy by solving (17) concurrently with
@(t) E B(agUecss ece)
(X,)) (X))e E. (ooeee)
0(0) = 0
(19)
EleA,
since by 18) the right hand side of 19 is precisely - H (Po(e)- p.) Eo(e) (or E., (o)U) Assuming that we know Zoo we can also track the evolution of the cross-cntropy via
H(Po)) P-) log R[e + log Zoo (Usco)-
(20)
Proof. The joint PDF f(t, x. a) of the process (Xe, Ac) satisfying 17 solves the Tokker-Planc's equation (FPE)
0.ʃ aV. (V,Us(0) + vnf) OoUa()) (())a., 2(0,2,a) (n)
(21)
Let us derive an equation for
x(t,z)= _tta))/_"ef(t. f(t,
(22)
To this end, multiply 21 by 6a, integrate the result over a E (-00, 00), and use integration by parts for the last term at the right-hand side to obtain:
o,p = a., (V, (V2Ust) V.p) - ÜgUnce ((1)p
p(0, 2 Usg (x)
(23)
By general results for the solutions of parabolic PDES such as 23 (see 669, Chapter 7), we know that the solution to this equation is unique, and
check by direct substitution that it is given by
p(t,,)) = Z 1 (n)
(24)
This implies that
[ x) Z Zo(t)-
(25)
Since by definition E[eA:]= See S e" e((1,, a.,,),dddd-- = Sas p(l, r)de this establishes the second equation in (88) To establish the first notice that
E (e0) (X.)e
Jae I (z))" 2(t,
Sa doUoro o (x),
E3 AeAA
SEe S e" aa)da
Sa plt, ddx
(26)
Z Sax OgUocc (z)e Uno(s)dr
(z)ee-unll)dddr 1%o - Zo10
Eo(o) [0gUn())]
16

---

The need for Jarzynski's correction. equation (77
Suppose that the walkers satisfy the Langevin equation (first
dx, vo + 22 dW..
Xo ~ Pda
(27)
where ((t) is evolving according to some protocol. The probability density function p(t, of X then satisfies the Pokker-Plance equation (compare 23
3pp a (VUo(s) Vp), p(t= = Pd.
(28)
The solution to this equation is not available in closed form, and in particular p(t, # Pe(t) (a) p(t, is only close to Pa(t) (z) if let + oo, so that the walkers X, move much faster than the parameters ((t) but this limit is not easily achievable in practice (as convergence of the FPE solution to its equilibrium is very slow in general if the potential Up(e) is complicated). As result (() (X,) * Ea(s) [QUue implying the necessity include the weights in the expectation (88
A.2
Proof of Proposition
The iteration rule for Ag in implies that
d
A
Σ q=1
(a--1 1
ag(xg. Xg-1))
keN.
(29)
For E No, let
1 4h
Bu (2, v) = (4Th)- d/2
exp
+ (x))²
(30)
be the transition probability density of the ULA update in (5), i.e. Ba Xk+1) is the probability density of Xx+1 conditionally on Xx. By the definition of At, we have
exp(AA) = II exp (X - ag (Xg. Xg-1))
(31)
A II
(g X
- (Xa)+ (Xo)
3 (x, Xa) g=1
where in the second line added and subtracted IX, Xg-1 12/44 and used the definition of on (x, y) given in (6). Since the joint probability density function of the path XI,.,.,..., at
Xx)
any k N is
o( ,1,,.,, RK) Z002 ---e-Umn
Uno (wa) II B (2g
(32)
deduce from (31 and 32 that, given an f Rd R. express the expectation eex as the integral
= Sf qu
(xw)e- Uog (ow) (2o)
Bg (Iq, el Tı, ......, (x))do ... dxy
By- (Ii Fig)
(33)
b
= 2002 1 / - (o)) II
(x xy )ddo
dxa
Since Sea (a u)dy = 1 for & E No and all x Rd, we can perform the integrals over zg. then IT. etc. in this expression to be left with
([f( = Zo6 / f(xx)e Un (an dx
(34)
17

---

Setting f(x) = 1 in this expression gives
BleA+] 7.80
/
e-u Uor (2x) (xx) 7:80" 70.
(355
JR
which implies the second equation in 7 setting (xx) 00U'0 (2xx) in (3 gives
Bi ]
Zoo" Ra
OaUa, (I ( Un, (2x) dzy day
Zoo Zo By 12U
(36)
which can be combined with 335 to arrive at the first equation
A.3 Mini-batched version of Algorithm I
In Algorithm 2 we present a mini-batched version of Algorithm [ where we do not update the positions of every walker at every iteration. Instead, we evolve only a small portion of the walkers, while keeping the other walkers frozen and only updating their weights using the information from the model update (which uses only the model energy and does not require pock-propagation). This version is more computationally efficient and leads to convergence in training with much fewer steps of ULA. More importantly, the mini-batched version of the algorithm, as compared to the full-batched version, enlarges the total sample size and therefore improves the sample variety with even less computational cost.
Algorithm 2 Mini-batched Sequential Monte-Carlo training with Jarzynski correction
1: Inputs: data points {x;
[n Ji=1; energy model Us: optimizer step opt(O, D) using 0 and the empirical CE gradient D; initial parameters Oo; number walkers N E No: batch size N' E No with N' < N, set of walkers {X)) sampled from Peg total duration K E N; ULA time step
h; set of positive constants {ck JKeN.
2: Aj 0 for I.....N..
3: for = 1:
1 do
4:
Pk (xp(A E/- exp(Ak)
normalized weights empirical CE gradient
5:
Dj = Di= 1 (Xx) n Σド OgUi (a!))
6:
0k+1 = opt(Ak, D2)
optimization step
7:
Randomly select mini-batch (X with # #B = N' from the set of walkers {
8:
for E do
9:
X X AVU,, (Xx) + 22
c N La)
b ULA weight update
10:
A = A 1 Xx) +ax(X X
11:
end for
12:
for j & B do
13:
X x X
update of the walkers
14:
k+1 = A.e U (X.) + Uo (X{)
weight update
15: end for
16: Resample the walkers and reset the weights if ESS < Ck+1, see
13
resampling step
17:
end for
18:
Outputs: Optimized energy Uax: set of weighted walkers (Xi A ) sampling Pax partition
function estimate Zax No, EN exp(AAx): CE estimate log Zop +n " U@x (w!)
A.4 Resampling Routines
Resampling schemes are necessary tackle the decay of effective sample size (3). For the sake of
completeness, here we recall three of the most widely used routines: multinomial [70]. stratified [71].
and systematic resampling (72), and refer the reader the review (73) for more details.
Given a of normalized scalar weights (p:) E [0, 1] with M Pi associated to the N walkers, we define the cumulative sum
Pi
n I....NN
(37)
1=1
18

---

1. All three methods prescribe way to choose a set {un Jn=1 E (0, 1] used to perform the resampling in the following way: for every n, m E [I.....N] the m-th particle is chosen during the n-th extraction if
Pm-1 < Un < Pm
(38)
Let us now specify how the set of un is selected in the cases in study. We denote by U(a, bj the uniform probability distribution on the interval (a, b].
Multinomial resampling. Sample q ~ ((, 1], independently for every n E f1....,NN N). This approach leads large number of possible resampled configurations, which is not desirable in practice as it increases the variance of the estimator.
Stratified resampling. Partition the interval (0, 1J into N sub-intervals, or strata, of size I/N; then, sample u ~ [[n 1)/N, n/N] independently for each = (1,,.... ,I,...,A} This approach picks single in each stratus, thereby reducing the number of possible resampled configurations.
Systematic resampling. Partition the interval (0, 1] into N sub-intervals, or strata, of size 1/N; then, sample u ~ ((, 1/NN], and n. ui + ()/N for 1. This method also reduces the number of possible resampled configurations.
Note that there various modifications of these three methods (see the review (33). all of them meeting the so-called unbiasedness condition, that is, the é-particle expected to be sampled in average Np, times. However, these extensions do not lead to critical lowering of the number of possible resampled configurations compared to systematic resampling. For this reason in our numerical experimentals resampling routine one uses, a fundamental role is played by the variance of the weights: the cumulative sum is dominated by one or very few weights, i.e. resampling is triggered too late, it does not remedy the suppression of population variability.
only tested the three methods above. Note also that, regardless of the
A.5 Contrastive divergence and persistent contrastive divergence algorithms
For completeness, we give the CD and PCD algorithms we used to compare with our Algorithms ! 2 As written these algorithms use the full batch of data points at every optimization step, but they can easily be modified to do mini-batching.
Algorithm 3 Contrastive divergence (CD) algorithm
Inputs: data points {z )a in F4, energy model Ug: optimizer step opt(O, D) using 0 and the empirical gradient D; initial parameters 6o; number of walkers N No with N < nz; total duration K N: ULA time step h; P N.
2: for k = 1, K do
3: for 1, N do
4:
X = Sample M uulidmmmaati
5: for p = 0, P do
6:
X X
hVUs, (X+) + 22h
~ ~~Na00. Ia)
ULA
7:
end for end for
8:
9:
Dr = NN-111 ΣN
i= OoUo, (Xp) -n En OoUo, (xt)
D empirical gradient optimization step
O2+1 opt(Ox, Dz)
11: end for Outputs: Optimized energy Uex: of walkers (X))
Interestingly, we can write down an equation that mimics the evolution the PDF of the walkers in the CD algorithm, at least in the continuous-tima limit: this equation reads
0.p V (VUnte
Vp) r(p - p-),
p(t= 0)= = p+
(39)
where the parameter > controls the rate at which the walkers are reinitialized at the data points: the last term in 39 is birth-death term that captures the effect of these reinitializations. The solution
19

---

Algorithm 4 Persistent contrastive divergence (PCD) algorithm
{2; ii=1 in F2 energy model Ug: optimizer step opt(O, D) using 0 and the empirical CE gradient D; imitial parameters do: number of walkers N E No with N < n; total duration K E N: ULA time step h.
Inputs: data points
2: X Sample.M ultinomial(({) for = I.
3: for 1..... K do
4: Dr = N-1 ΣN OgUn, (Xi) n Ei= OoUo,
P empirical gradient i optimization step
5:
Ok+1 02, D2)
6: for i ,,,,, .N do
7:
X X AVU,, (X 2
~
ULA
8:
end for end for Outputs: Optimized energy U set of walkers (X IN
9:
10:
to this equation is not available in closed from (and p(t, 2) + Pa(c) (2) in general), but in the limit of large (i.e. with very frequent reinitializations), we can show (55) that
p(t,x) px(x)) (x) + (+U)))())) ())()))))))))))))))))))))))))))) O(v
(40)
As result
J/L (1) x))c
=- / (xU (U)
Vx.(x)) (p)))) dr + (2-
(41)
Rdd
v- [ (o)VUo())(t) VUnc (z) - Og (g (r)de + O
The leading order term at the right hand side precisely 1/ times the gradient with respect 9 of the Fisher divergence
1
VUo(x) + v log px) (x)d
2 Rd
(42)
1
/ [1V -2AU6)))))))) + log p. (x)) p. (x)dr
2 J
where A denotes the Laplacian and used Jau (Ue(x) V log (z)) fox VUa(x)- (p.(x)d2 Sga (U) This confirms the known fact that the CD algorithm effectively performs GD on the Fisher divergence rather than the cross-entropy (28)
B Additional numerical results
Here give additional details about numerical results. The codes are available at: https //github com/ wwt.ccom/muu/s/eeee submsusseee ssiesssosss oo111/1/1//// 2/222///BB.222 jarsynJaryyski
B.1 Gaussian mixture distributions
Here we give some additional numerical results for the teacher-student model discussed Section 4.1 in which the two wells of the teacher are aligned along the first dimension, fixing the first component of the means to be a = -10 and b. = 6, and a = b0 = for any & = 2..... d; we also set p. 1/(1 + e ==) 0.25 of the mode centered
2, 1og(3), corresponding to
All the simulations are performed number of data points is 10a The setup of the teacher is the same for every simulation we display here and the optimization step performed with full batch gradient descent with learning rate constant in time.
50, with a time step
0.1 for ULA update, The
20

---

Mass of the first mode
Means
KL divergence
0.
101
Algo PCD
18-1
184
Algo 1
(21)
259
300
790
150
Rescaled lterations
Rescaled Iterations
Rescaled Iterations
Figure 6: GMM experiments: Evolution of the parameters and the cross entropy during training by Algorithm|] I PCD, and CD. W.r.. A we display the results for five different thresholds in the resampling step. Left panels: evolution of Pk 1/(1 + e **) middle panel: evolution of ag and by: right panel: evolution the Dallback-Leibles divergence.
ES5
8333.3
7142.9
64.10
5585.6
5000.0
Rescaled Iterations
Figure 7: GMM experiments: Evolution of the effective sample size (ESS) for five different choices of threshold associated to c, constant in time. Bimodal student.
GMM for different choices of Ck. Results are shown in Figures 6 and 7 As initial conditions we select aj -00 6 10-1 a 10 10-1002220 (0, 1) and b8 10 (0 1) for any 2, ..... .,. d; this perturbation around = b 0 is prescribed to avoid numerical degeneracy. For z, we fix 2g = 0. We run Algorithm
as well as the CD and PCD Algorithms 3 and using N 10* walkers for K = 8 10 iterations. We use different learning rate for and a, b, namely Y₂ = 0.125 and %..6 %. Moreover, these values are multiplied by factor 10 in CD. With regard the resampling step in Algorithm| we choose threshold e = which is fixed in time. We display the result for five possible values of this hyperparameter, namely c = 10.1, 0.4, 0.8, 1.0]: these are related to thresholds in the effective sampling size (ESS) via ESS = N/(c + 1). With regard to Algorithm 1 we choose M 1 and N' N, that the full batch version; in the CD Algorithm 4 we choose P 4 for the number of ULA steps between restarts. Every run is performed five times and the average between them is shown figures.
Mode collapse in GMM for PCD. Results are shown in Figure 8 In the same setup as above we
select as initial conditions aj = -1 bj 10-1, aß 10-2A (0. 1) and bg 10-2A((0 1) for any E [2, d]: for we fix 20 0. The learning rate is chosen to be 7= = 5 for and Yo.6 = 0.2?> for the means. The time step of ULA PCD algorithm,
0.2. Since the objective to show mode collapse in the just Algorithm for K 10ah iterations and N 10 walkers.
The need for resampling. absence of resampling, we observe dramatic deterioration of ESS (Figure 9). To investigate how this issue is solved by resampling, we compare the three routines presented in Appendix A.4 using three pre-specified lags between the resampling steps. We use the same hyperparameters (learning rate, target distribution, etc.) in Figure The results
21

---

Mass of the first mode
Figure 8: Mode collapse and oscillations in PCD. Evolution of the probability 1/(1 + =)
Mass of the first mode No resampling
0.50
10009
0.2
0.24
0.22
0.18
Standard ESS
0.35
0.15
8000
0.30
0.29
6200
P
100
35p
200
25c
50
100
150
200
259
Rescaled Iterations
Rescaled Iterations
Figure 9: GMM in 50d without resampling step, full batch experiment with N 10 walkers, average of 30 runs, 4.8 10 iterations, same target density of Subsection 4.1 Left Panel: relative mass of the first mode. Right panel: evolution of ESS
shown in Figure I0 looking the upper panels, we that stratified and systematic resampling more stable than multinomial; moreover, if the resampling step is performed too infrequently (green lines), the method converges to result where the target relative mass off. Looking at the lower panels, we see that the minimum statistical error is obtained with systematic resampling: moreover, the behaviour of uncertainty appears be more stable than stratified.
As a side note, systematic resampling requires just single random number, contrarily to multinomial and stratified which need N, making the former also more efficient from computational point of view. This consideration, plus the experimental results we just discussed, strongly motivate the adoption of such method for the resampling step in our proposed algorithm.
Discussion. GMM in high dimension are challenging for the standard CD and PCD algorithms given in 3 and The experimental results of this section also confirm the theoretical amalysis prsented in Appendix beloo:: CD is performing GD but Fisher divergence rather than cross entropy, and as a result it incorrectly estimates the mass of the modes since Fisher divergence is insensitive to this quantity. On the other hand, PCD causes cycles of mode collapse (see Figure and the left panel of Figure (6) and the KL divergence does not decreases monotonically. since the protocol is not ensured to be gradient descent (right panel in Figure
In this example, our Algorithm outperforms these standard methods an implementation GD cross-entropy. In particular, the estimation of the relative mass with Algorithm | is more
22

---

Mass the firs mode Nukromaal resampling
Kass of the first mede- Strabifie resampling
Mass the first mode Systematk resamplig
20400
040
541
0.23
0.2
H4
0.21
106
199
300
10000
33300
20923
0.75
Rescaled Rerstions
Rescales heeratios
Rescaled Berations
Standard crear Mutinomial resampling
Star dand error Stratffied resampling
Standard arror: Systomatic resompling
2003 10691
10003
2.00
0.8
-3
0.004
0.00
000
0.0
313
30000
Rescaled Karations
Rescaled Rerations
Rescaled Rarations
Figure 10: GMM in 50d with three resampling routines, full batch experiment with N 10 walkers, 50 runs per each resampling interval, 4.8 10 iterations. Upper panels: relative mass of the first mode for each resampling method. Lower panels: standard error computed using the empirical standard deviation via Temp/ /50.
accurate than with the CD and PCD algorithms (see Figure 6). Moreover, the computation of KL divergence via Jarzynski weights fairly precise; in Figure 6 the estimated KL and the exact one overlaps beyond the minimum values reached in with the CD and PCD algorithms. With regard to Figure the choice of the threshold for resampling does not appear to be decisive in this regime of hyperparameters; in fact, looking at the evolution of 0 and of KL divergence, the overall behavior of Algorithm is not dramatically influenced by the choice of c.
B.2 MNIST data
In this section, we provide additional details about our numerical experiments on the MNIST dataset using Algorithm 2
confirm that the Jarzynski weights of the generated samples are directly related to the image quality: this is shown in the left panel of Figure where we display images along with their Jarzynski weights. Moreover, the right panel of Figure indicates that resampling at the end of training using these weights helps improve sample quality. Examples of generated images are shown Figure 12
First, as discussed in Section 4.2
B.3 CIFAR-10
A key component of the success of our method is the resampling step. In Figure 13 we plot the Effective Sample Size (ESS): each peak in Figure 13 indicates step of resampling of all the samples. We note that training EBMs with mini-batches has the side effect of a rapid loss of the ESS, which measures the sample quality. So, doing resampling with a proper criterion and a reliable resampler necessary.
C Theoretical analysis of mode collapse in simple 1d-GMM
In this appendix, we explain why and how mode collapse arises when learning multimodal distribu- tions if we do not include the weights prescribed by Jarzynski equality, and why collapse does not arise with these weights included. To show this in the simplest setting, we work with Gaussian mixture model: mixture of two one-dimensional Gaussian densities, with the same variance equal
23

---

3
3
2
6
3
6
2
2
2
2
6
6
3
2
13,50
10.25
10:11
10.09
2
3
3
2
3
3
9
2
2
6
2
9
2
2
9.31
8.73
8.04
7.92
7.36
3
2_
7
2
2
2
3
6
2
2
6
2
2
2
2
2.91
2.73
2.74
2.68
6
2
7
2
2
6
3
2 3
2
2
-
2
2
1.43
1.42
1.41
1.4
1.39
1.38
1.37
1.36
2
2
2
a
3
6
2
9
2
2
0.75
0.74
0.73
0.73
0.73
0.72
0.71
0.71
Generated samples during training (w/o resampling)
Generated samples during training (w/ resampling)
Left panel: Images randomly chosen during the training from the entire set of generated samples with their associate Jarzynski weights. From left to right, top to bottom, the higher the Jarzynski weight is, the better the image quality is. Right panel: Images obtained under the same training conditions, after resampling and continued training for 120 epochs.
Figure 11: MNIST dataset:
2
6
2
3
2
2
6
2
&
2
2
3
2
2
3
2
3
3
2 2
2 2
2
2
2
2
3
5
2
2
2
0
6
6
2
2
)
2
2
2
&
3
2
3
2
3
3
3
6
7
3
8
3
8
2-
8 7 4 1
2
2
3
3
8
6
3
3
2
" 2
3
6
2
3
2
3
7
3
2
6
3
2
2 L
7
1
4
2
2
2
6
6
3
2
3 2
2
2
4
2
2
3
7
3
0
2
3 7 2
3
2
6
0 3
5
2
6
3
6
2
2
2
3 2
2
7
0
2
3
3
6
8
2 2
3
4
}
6
2
4
23
3
Figure 12: MNIST dataset: Left panel: images generated by the mini-batched version of the Right panel: Images generated from the PCD with mini-batches.
Algorithm 2
to 1, with means a, b IR of the modes: our aim is to learn the probability masses of these modes. The general picture is given in Appendix C with details presented afterwards Appendix C.2 where we assume that the learning dynamics is that in (47). with walkers and direct access to the distributions Pa(t) and p., and Appendix CK where we quantify the stochastic fluctuations induced by empirical estimation.

---

Effective Sample Size
1.0
Effective Sample Size
0.6
0.4
0.2
Effective Sample Size
100
200
300
400
500
600
Epochs
Figure 13: CIFAR-IO dataset: The Effective Sample Size (ESS) during the training with Algorithm Each peak in the plot implies step of resampling. Notice that the number of resampling in CIFAR-10 experiments significantly larger than the one in MNIST experiments which suggests the necessity of resampling in the scalability of our method.
C.1
Mode collapse in Gaussian Mixtures
C.1.1 Target distribution and model parameterization
As target density take
px(x) e al- + -0- 2 1/1
(43)
V2T(11 e
Here 2u parameterizes the mass of the second mode relative to the first and the sole parameter of interest. The proportion of samples in both modes is indeed
= 1= 1
(44)
I
Our point is that when both modes are separated by very low-density regions, learning z, without weight correction leads an incorrect estimation of the mode probabilities ('no-learning') to mode collapse depending on the initialization of the learning procedure, whereas using the Jarzynski correction does not. From now on, we will suppose that the modes are separated in the following sense:
10,
(A1)
which will be enough for needs. The more separated they are, the stronger our quantitative bounds will be.
The parameterization for our model potential U2 is consistent with (43
(=(2) = - log (e l0 af? + - 12=
(45)
and the associated partition function and free energy are
6 = V 2T= (1+E=*), F- = - log Z, = log - - - .
(46)
The normalized probability density associated with Uz(xr is thus p=(x) -
C.1.2 Learning procedures
Gradient descent on the cross entropy leads to the following continuous-time dynamics:
=(t) = Ez [2o) E. [0,Uu(()))
(47)
For simplicity, will start this ODE (0) It corresponds proportion of for both modes.
The gradient descent (47) an ideal situation where the expectations EE E, can be exactly analyzed. In practice however, the two terms of (47) are estimated; the second term using a finite
25

---

number of training data { and the first one using a finite number of walkers {X;), with associated weights {e Ii=1. For simplicity we N and the empirical GD dynamics is thus
E
e a. (AU)
(x))
Th 1i=1
(1Us(1))
(t))
(48)
Σi=1 e
and the walkers evolve under the Langevin dynamic
dxi -VV (X))dtt 22dW? for some fixed > 0. Now the nature of the algorithm varies depending on how the walkers are initialized and the Jarzynski weights are evolved.
(49)
1. The standard PCD algorithm sets X = x², that is, the walkers are initialized at the data points, and the weights are not evolved, that A = 0 at any time.
2. Alternatively, the walkers could be initialized at samples of the initial model: X² P=(0), with the weights not evolved. We refer to this algorithm as the unweighted procedure.
3. Our algorithm 2 corresponds to initializing the walkers at samples of the initial model, X3 P=(0). and uses the Jarzynski rule (17) for the weights updates.
For simplicity we analyze the outcome of these algorithms in the continuous-tin set-up, i.e. us- ing (48). This is idealization of the actual algorithms, but it makes the analysis more transparent.
C.1.3 Perturbative analysis
Using simple approximations (see Appendix C for details), we show that in the three cases above, the dynamics 448) can be seen as perturbation of a simpler differential system whose qualitative behaviour fits with our numerical simulations. These systems depend on the initialization of the walkers and thus prone to small stochastic fluctuations. We introduce:
qx the proportion of training data {-i} that are close to mode b (a more precise definition is given Appendix C). and we define 2, as satisfying q. = 2.. /(1 + e e=):);
@(0) the proportion of walkers at initialization that are close to b, and ((0) 1 4(0).
Practically, d. is random variable centered q. and with fluctuations of order n n- and (()) is centered at e e= /(10//1111 (=1+e=(())) 2 with fluctuations of the same order. In the limit where is large, they can be neglected, but in more realistic training settings, the use of mini-batches leads to small but non-negligible fluctuations, as will be clear in Equation (53)
The arguments Appendix [ lead to the following approximations:
In the model-initialized algorithm without Jarzynski correction ('unweighted'), tion of
48 is perturba-
Bunw (t) q(()) - q.
(50)
This system has no stable fixed point since the RHS no longer depends on Zuaw (l), leading to a linear drift Zunw (t) (ğ( & and thus to divergence of zunw (t). Consequently, the mass of the second mode, q(t)
Sumw(()) (t) (1++= Funw((()) (t)), converges time scales, this drift leads to a sudden transfer of all walkers in one the modes, then to complete reversal of the drift of Zuuw which then diverges in the opposite direction, leading a succession of alternating mode-collapses: see also Figure g and Remark C.4 below.
1. However, on longer
the continuous-tin time version of the standard PCD algorithm, (48) is a perturbation of the same ODE as 550) However, in this context, since the initial data {X²} and the training data {xt} are identical, we have (()) = G., and
ápea (4)
(51)
The parameters do not evolve and the system is stuck at apea (0) ('no-learning') Note however that in this version of PCD, the walkers are initialized at the full training data. In practice, the number of walkers is smaller than the number of training data so that one often uses small batch of training data initialize them; in this case we can still have ((0) # q., falling back to the first case above.
26

---

410
(a) Unweighted dynamics: 150
(b) PCD: STT
(c) lalrzynski-coorecteed: (53
Figure 14: Behaviors of the three learning algorithms leading to the ODEs 51
50 and 53 The
on bottom. The blue lines show z(t) (top) and the corresponding g(t) = = + (()) (bottom). Figure 14b; PCD algorithm, which leads to no-learning, with the weights not evolving at all in agreement with the behavior predicted by (51). Figure 14a Dynamics without weights, that leads to mode collapse, with =(t) diverging towards +oo in agreement with (50) when the parameter 7 ((0) d. positive. Figure 14c Jarzynskiicorreeeeed dynamics with weights included, (53). The green line represents (54) and Teatures the small
solid black lines represent the target parameters, 2, on top and 9.- = e /
correction due to the stochastic fluctuations at initialization between the walkers and the model. In all cases, we used 200 walkers and data points. The continuous-time dynamics were discretized using the Euler-Me Maller-mauyyaaa scheme with step-size 0.01 up to time T'= 10*
The continuous-version of Contrastive-Divergenn is equivalent to the well-known Score- Matching technique [24] In this context, (48) is perturbation of
(1) O,
(52)
leading the same 'no-learning' phenomenon.
In the model-initialized algorithm using the Jarzynski correction, 48) is perturbation of
((0)e-))r(1)
Ejar (t) p(0) ((0) (1) + A
(53)
(:00)e=.//000 This system has a unique stable point ب, satisfying e - = =p(0)e 2. /a(0), hence
8(())
= + log
(54)
((0)
Note that the second term at the right hand side is small correction of order O(n -1/2) since ((0)/((0) 1/221/).
These predictions are in good agreement with the results of the simulations shown in Figure 14c
C.2 Analysis of 47
Using the Jarzynski correction
In continuous time, the Jarzynski-correction described in Proposition exactly realizes (47 However, even for simple mixtures of Gaussian densities, the expectations in (43p do not have a simple closed-form that would allow for exact solution. That being said, when and are sufficiently well-separated, the system can be seen perturbation of a simpler system whose solution can be analyzed.
First, we note that
42
0.U2(x) e =
(55)
U2(x)
27

---

The main idea of the approximations to come is that (,U,(x) is almost zero when is far from (and in particular, close a). and is almost if is close to b. The next lemma quantifies this; from now on we will adopt the notation
La = 4
and
1. = 4, 4].
(56)
Lemma 1. Under Assumption A1 for any v E IR,
i c Ias then (_U,(z)) 1 e-
i Ib, then (z) 1| e"~10
Proof. From (55).
that if La then |a - b|
and consequently e- e-11, But
the denominator of (55) itself greater than -= which itself greater than e- 4²/2 since al 4. This gives the first bound and the second is proved similarly.
recall that f N0, 1), then P([5) 1) S -//12/1. hence
P(E) > 4) S - 12 /4 0.0001.
(57)
Lemma 2. Let u. R. Under Assumption A
we have
e 16
E, E.
S
(58)
1
where |e| < 0.0002 + 2c 2e-10 10clol
Proof. The integral is exactly given by
1
vー (E. 6)9/2
e
v-vー (Ev-N)9/2
E
E
(59)
1
(.()
where Sar&o denote two Gaussian random variables with respective means a, By (57). Ea and & respectively contained in Ia [a - 4. and I, = [b - 4, b 4] with probability greater than 0.999. Let us examine the first term of (59). The fraction inside the expectation is always smaller than 1, hence by [emma[]
0_U 0,U,(C))lezo S 11,41 IeEL 0,U,(Ea)) S 1g,41, e 11-10
Consequently, by (57). the expectation is smaller than 0.0001 == The second expectation in (59) is equal to
(X. /2
(Es)
and by the same kind analysis, the expectation here smaller than 0.0001 + v-10 Gathering the bounds yields the result.
In particular, the right hand side of (47) can be approximated by I+e- 14 to an error
=(0
term smaller than 0.0004 2e (c-10(c + el+-)). If =t),, 2, are contained in small interval [-C, with, say, C' < 5, this error term is uniformly small in time, and one might see (7) perturbation of the following system:
-
=(1)
(60)
1 + e (tt)
1+e
system with only one fixed point at 2(t) Z+, the ground-truth solution.
28

---

Mode collapse in absence of Jarzynski corection
Now let us analyze in similar fashion the dynamics without reweighting. Here, 47 replaced by
(()) = (0 (X.) E. [02Us()))
(61)
where the process X, solves
+ dX, -a VVV Coaidd++2allll VadW,, Xo Pz(0)
The probability density function p(t, of X, satisfies a Fokker-Planck equation
a,p= = aV (VUz()) Vp), p(t= = P=(0)
which, in full generality. is hard to solve exactly, and thus exact expressions for the first term of (6T) are intractable. However, depending on whether Xo is close to a or b, the process X₁ can be well approximated by Orsttein-Uhlenheck process, hence p(t, can itself be approximated by Gaussian mixture.
Proposition 3. Suppose that (Al) holds and that
31, C E R+ such that for all t E [0, a 'T1,
=(t) E [-C,
(A2)
Then one has DKL (p())()p(t)) s 8t where 5 0.000025 100c 100e-20e22 In other words, p(t) is approximately constant, up to reasonnable time scales O(188).
Proof. Xo is drawn from P=(0). Gaussian mixture; the probability of it being sampled from Gaussian with mean e==(0)/ /(1 e-=(0)) = 1/2. We will work conditionally on this event En and will note (t) the density of Xo conditional on E"; thus, (0) (a, 1). We set (xx) = 2 al2 so that (z) = (x - a) and we consider the following Ornstein-Uhlenbeck
process:
dY, -a VV + V2ddW, Xo=
whose density will be denoted (1). We use classical bounds on the divergence between (²( and (t). For example, the bounds in [74 Lemma 2.20] directly apply and yield
DRi (1) (1) 1 1 [[VUs(()()))) (V.))) ds.
Since Y is nothing but an inside the integral is a Gaussian expectation and will be shown to be small:
at equilibrium, Y, ~ Naa, 1) for all t 0. The term
E (VV)) VV(Y))
0.0001
400e 2x(1) 20
(62)
Consequently,
0.0001
100e 20 / è (==(())
(t) <
4
Under A2 the overall bound remains smaller than times 0.000025 + 10 as requested, thus proving that (t) = aaa, 1) and (t) close with the same quantitative bound
Similarly, (t), the density of X, conditional on X₂ being sampled from a Gaussian with mean b, is close to (b, 1) with the same quantitative bounds.
Overall, using the chain rule for KL divergences.
(o())()))( S
(1)lo))) + ((E St.
In other words, p(t) is close to mixture of two Gaussians with modes centered at b. and the probability belonging the first mode is the probability of Xg belonging to the first mode, that is. -x(0) /(1 = 1/2.
Proof of
62
We have
(I - a)e la lt=a/2/2 + (x - b)e |z 6/2/1/222
VU+(T)
(63)
(2(x)
29

---

Using Lemma] ! and the fact that if z E La then |x - b) < 16 and xx al < 4, we get (V)(z))
(x a) < 20 with E I n-10, Consequently.
E
-VV)))))))VV))))))) - P(Y & Ia) + (20c 120)-11102
0.0001 + 400e 20-20 20
As a consequence the Proposition 3 the first term of (1 can approximated by
[0,0)|0_U D..Uii((3 = Eg (0.)) which in turn can approximated by e e-=(0)/ /(1 + e=
thanks to (58) Overall, (61 is therefore perturbation of the system
==
a(t)
9- =.
+ = =00)
1e
Since the right hand side no longer depends on =t),, this system leads a constant drift of z(t). that (t) yt. leading to mode collapse since e --=(t)) (1 t goes to either 0 or
Mode collapse for Contrastive-Divereene and Score-Matching
The continuous-timg limit of Contrastive Divergence (Algorithm B is equivalent to Score-Matching minimization (28) The objective function becomes the Stein score,
SM(=) |. log (X V log Pz. (X) EE - vU,, (X))),
which is in theory intractable due to the presence of the unknown parameter 2,; a well-known computation from [24] shows that the gradient dzSM(z) can be estimated using the training samples without resorting to 2.. Note that in this context, there are no "walkers'.
Now, the dynamics
47 is replaced by
VU.. (t) 0, VU=(0) (U,,
(64)
= p.
VE.. (E.) q. 0. - VU.. (S6)
(65)
where here again & ~ Axx, 1). From 62) and the triangle inequality, we have
(E(|VU>() (E.) vU., () 1/2 <
(E[[VU=2()EE.)) - (E. - a) ()) 1/2
(EI (Ea 16)
< 0.0002 + 800e 22xt))
and the same approximation holds for the second part in (65). Overall, we get that for any reasonnable 2, SM(2) that is, every 2 minimizes the score. similar analysis leads to 8_SM(2) ss Consequently, {t) = 0: Score Matching and Contrastive Divergence leads to 'no-learning'.
C.3 Empirical gradient descent analysis of (48
The gradient descent 47 represented an ideal situation where the expectations Ex(r) E, can be exactly analyzed. In practice, the two terms of (47) estimated; the second term using finite number of training data } and the first one using finite number of walkers {Xi}, with associated weights Ai which either evolved using the Jarzynski rule, simply set to 1 in the PCD algorithm. Our goal this section is to explain how these finite-size approximations do not substantially modify the previous analysis and lead to the behaviour presented in C.1 For simplicity we keep the time continuous.
We recall 48
N e CAiisUs( *(t)
En
a,U (xt)
(66)
N
30

---

The dynamics with Jarzynski correction leads to the correct estimation of the empirical weigths The continuous-time dynamics of the walkers and weigths in our method is given by
dx; a VU ( + 22ddW,
(67)
A 0zV=(e) (x)=(1((X())(())
(68)
Let r be the number of training data in In and p. ng/nn their proportion, and similarly q. the proportion in la, and 1- -q. By elementary concentration results, the remainder p- q. can be neglected: with high probability, it is smaller than, eg, 0.0001. We will note ž, the parameter satisfying ς. = I Using Lemma the second term in (66) is approximated by q. Now let us turn to the first term in (66). Still using Lemma we see that the first term in (66) is well approximated by
e-
aien edi Ei=1 e
(69)
Now let us use Lemma The second equation in (67 entails 4i = exp (- OuU
1 if X E Ja. then (_U=(s) = O. Conversely, if x; E Io, then (,Ux(s) = 1. Moreover, Proposition 3 and its proof essentially show that if X3 belongs to the first well (close to a), then with high probability so does X for every 8, and in particular X E la with high probability for every 8. Consequently.
(-/h" Ods =1
e ~ exp
if X Ia:
(70)
ei exp (-)) ;(s)ds) exp -xxp(=(())+000))))0)))))))))) + exx())(()
if if Ib.
(71)
As already explained, the dynamics (67 leaves approximately constant the number of walkers in both modes; consequently, the proportion ((t) of walkers X2 in 4, remains well approximated by the initial proportion, which is g(0), and we obtain that (69) is well approximated by
((0)e==(1)
(72)
p(0) + =e00((0(00))))
where we noted p(0) = 1 - ((0). Note that since =0) 0, with high probability p(0) and ((0) are close to 1/2. The random variable p(0)/q(0) is thus close to
Overall, we obtain that the system (66 a perturbation the following system:
((0) =(1)
(tt)
(73)
((0) ((0)e-8(0) B(())
1 1
This system has unique stable point equal
2 = E, + log log(((0)//(0))
(74)
Remark. If the algorithm had been started at =0)) * 0, could check that the stable point would
become é. + log(((()) + 20).
Freezing the weights leads to mode collapse
If the walkers still evolve under the Langevin dynamics in (67) but the weigths are frozen at e
e ('unweighted" algorithm), then (66) becomes
Σ O.U (XF) N
ΣR
att)
(75)
Keeping the same notation as in the last subsection, the second term is still approximated by q., but this time the first term is instead approximated by
Σ Σxxte 1 qQ() Fz Q(0).
n
31

---

Consequently, 75 is perturbation of the system
e
(tt) = ((0)
1te =: 2,
(76)
which no longer depends on t and thus leads to z(t) = 7t and to mode collapse.
The PCD algorithm leads to no-learning
In the preceding paragraph, initialization, the walkers X are distributed according to the initial model P=(0). In the PCD algorithm 4 the walkers are instead initialized directly at the training data {ær in N==1- However, in this case, the analysis the preceding paragraph remains essentially the same, with single difference: the initial proportion of walkers that close a, noted ((0), now exactly q.- Thus, (66) becomes a perturbation of
(t)) = q. - q. = 0.
(77)
The parameters remain constantly equal to its initial value =(0), i.e. there is no learning.
The CD algorithm leads to no-learning
The Continuous-time Contrastive- Diverrgee--ee algorithm is minimizing the Stein score and is equiva- lent to Sctre-Matching as mentioned above: the direction of the gradient of the log- likelihood is that of the gradient of the Stein Score, leading to no-learing. With the estimation given by the training samples, the analysis exactly the same as above:
2(t) = Σ ( (x!))
i=1
20, VU, (z.) x (VUs(x!) - VU2, (x.)).
If If E La. then explained in the proof of Lemma 62 (v,( Ro (xf a) for every hence
0, V.!) x (VU)(x)) - (v.().)) )) R 0. The same holds for xt, E In, leading to (66) being
perturbation of the system
a(t) = 0.
(78)
C.4 On mode-collapse oscillations
Most of the approximations performed earlier rely on (A2) that is, the learned parameter z(t) remains compact set.
However, in the unweighted algorithm, this is no longer the case at large time scales, since z(I) diverges from (6); in particular, the approximations from Lemma T become meaningless. In fact, Proposition 3 is no longer relevant. The core of Proposition 3 rests upon the fact that if walker X; is close a, then its dynamics (67) close an Crnstein-Uhlenbet
process since
VU=(e) s0 (X a). This fails when x(t) has a large absolute value. Let us suppose for instance
that zt) is very small (negative), that e ==()) is very large, |a(t)] >> |x b/2. In (63 the first term of the numerator is dominated by the second term. Overall, we get
(x²()(X)) R (X: b).
and this is valid for all Xj. Consequently, all the walkers now undergo an Ornteein-Uhlenbee Uhlenbeck process centered b and in particular, the walkers that are close to a are exponentially fast transferred the region close to b. At this point, the first term in (75) becomes close to 1, leading to the approximated system a(t) 1- - q.: (t) oscillates back to +oc, until the same phenomenon happens again and all the walkers transfer to the region close to
This leads to an oscillating behavior that can be observed on longer time scales (see Figure & for example). We do not think that this phenomenon is relevant to real-world situations since most learning algorithms trained for limited time period.
32


Reviews:
[{"review_id": "xrEonykp3F", "reviewer": "Reviewer_1pft", "summary": "This paper examines EBM learning using tools from the Jarzynski and Sequential Monte Carlo. The method is meant to address limitations of current EBM methods, which often rely on simulated Markov Chains to calculate approximate gradient of the log normalizer. Since the chains in practice do not converge to the model distribution, the theoretical validity of standard procedures is questionable and practical issues with correct multi-modal modeling can arise. The central contribution of this work is to introduce a new loss function for EBM learning. Using the Jarzynski equality, the gradient of the log normalizer is shown to be equivalent to the ratio of the expectation of weighted gradients of the EBM evaluated at ULA samples and the expectation of the weight terms. Unlike standard EBM learning which requires an unjustified assumption of MCMC convergence for validity, this proposed loss function is exact even for ULA samples which have not reached the model equilibrium distribution. To improve sample efficiency, weights are periodically resampled during practical implementation. Experiments on Gaussian mixture models and MNIST digits show the proposed method is better able to learn correct relative weightings of modes compared to standard EBM models.", "questions": "* Can the performance of the proposed method on larger scale datasets (e.g. CIFAR-10) be investigated?\n* Can the authors discuss the validity of the weight sampling step in Section 3.3?", "limitations": "Authors did not include explicit sections discussing limitations and broader impacts.", "rating": 5, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "* The theoretical development of the proposed loss function is the central strength of the paper. In particular, developing a loss function that is valid even for non-convergent ULA samples is a major contribution that has the potential to lead to more principled and effective EBM learning.\n* The paper is clearly written and easy to understand.\n* The proposed method does not significantly increase computation costs beyond standard EBM methods.\n* A thorough appendix with full proofs of theoretical results and a significant amount of supporting information and additional experimental results is included.\n* Experimental results give convincing evidence that the proposed method is able to model relative modes weightings more effectively than standard EBM learning.", "weaknesses": "* The primary weakness is the scope of the experimental results. The scale of datasets studied is smaller than typical EBM works, which usually examine at least CIFAR-10 as a baseline. In my experience, many methods which are effective for GMMs and MNIST might have poor performance on CIFAR-10. Since the goal of the proposed method is to improve multi-modal modeling, it is difficult to gauge the potential of the proposed method when experiments only examine datasets with very few modes. FID scores on CIFAR-10 would provide a clearer comparison of the generative abilities of the proposed method compared to related methods.\n* The theoretical validity of weight resampling technique in Section 3.3 is not discussed. My guess is that this step is not theoretically consistent with the loss derivation, although I am not sure about this. The paper should discuss the validity of this step. If it is theoretically consistent with the loss function, then proof should be given. If not, this should be acknowledged as a limitation of the proposed method.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We would like to thank the reviewer for their careful reading, positive view on the manuscript, and constructive suggestions. Below, we address the main points raised.\n\n**Scope of the experimental results.** As stated in the common reply, we did not have the computing resources necessary to perform experiments on large datasets such as CIFAR-10 or ImageNet, unfortunately. However, we plan to perform these experiments in collaboration with other groups with more computational resources in a follow-up paper.\n\nWe also stress that our experiments on MNIST are nontrivial: indeed, by pruning the dataset to 3 digits, we make the overlap between the classes smaller and the energy barriers between them higher, which is a challenge for EBM training especially if we want to recover the relative weights of these classes (as we do).\n\nFinally, let us mention that in the revised version we will include a more quantitative comparison between our proposed method and CD (PCD) for MNIST experiments (including test using the FID as metric, see Fig. 2  in the figure file the ***Common Reply***). Though we do see that our method surpasses PCD in terms of the FID score, we want to mention that the FID score is not a good metric for MNIST images. In order to use the pre-trained Inception V3 network for the computation of the FID score, we copy the grayscale channel three times and resize the images from $28\\times 28$ to $299\\times 299$ pixels. Since there are limited results in the literature on benchmarking the training performance on MNIST images with the FID score, we do not report and compare results from other papers here. Since the Inception v3 network is trained and used for real images (e.g. CIFAR-10, ImageNet) and MNIST hand-written digits are not natural objects, we do not seek to get reasonably low FID scores for MNIST images. In addition to quantitative results on the MNIST dataset, we conducted additional tests on other synthetic examples involving high-dimensional discrete datasets (see the description given in the ***Common Reply*** and Fig. 3  in the figure file).\n\n**Validity of the resampling step.** *The resampling step is theoretically consistent with the computation of the gradient of the KL divergence used as the loss.* This can be seen from Eq. (10): the last term in this equation can be interpreted as a birth/death term whose effect is statistically equivalent (in the infinite number of walkers limit) to keeping the weights and can be implemented in practice via resampling. This duality between tracking the weights of the walkers, or resampling them according to these weights, is valid both in the discrete- and continuous-time setups, and is at the core of sequential Monte-Carlo methods where resampling is used to maintain a good effective sample size (ESS). These arguments are standard, but for the reader's convenience, we will include proof in the final version of our paper. \n\n**Limitations and broader impacts.** We will include a more thorough discussion of the main limitation of our method, namely that the number of walkers must be large enough to ensure statistical accuracy -- the need of using a large number of walkers is common to all EBM training methods, and the impact on efficiency can be mitigated by using mini-batches (as discussed in the answer to Reviewers ripa and d8UC). We will also discuss the broader impact of our approach more thoroughly, in particular emphasizing the widespread usage of EBM and the need to use methods that ensure fairness in this context (as ours does)."}, {"review_id": "dfG9c03o1P", "reviewer": "Reviewer_ripa", "summary": "One of the main challenges in the training of Energy-Based Models (EBMs) is the design of a tractable objective that allows for low-variance gradient estimation. Indeed, the maximum likelihood objective requires sampling from the current density model, which usually is computationally expensive and biased due to poor mixing properties. Many alternatives have been proposed to alleviate the issues of maximum likelihood objective: contrastive divergence, persistent contrastive divergence, non-convergent MCMC training, etc. \n\nThe current paper approaches the optimization problem of the maximum likelihood objective by proposing a new sampling procedure. Once again, the main challenge is to estimate the term\n\n$$\\mathbb{E}_{x\\sim q_\\theta}\\partial_\\theta U_\\theta(x),$$\n\nwhere $q_\\theta(x) \\propto \\exp(-U_\\theta(x))$ is the current density model. After each parameter update, instead of relying on the MCMC sampling from the density, the authors propose to update the current samples using unadjusted Langevin Dynamics and reweight the generated samples to get an estimator of $\\mathbb{E}_{x\\sim q_\\theta}\\partial_\\theta U_\\theta(x)$. Thus, the main development of the paper is Algorithm 1 and Proposition 2 demonstrating that the proposed estimator is consistent (converges to the truth value for the infinite number of samples).\n\nThe authors demonstrate the benefits of the proposed algorithm empirically on two target distributions: the mixture of 2 Gaussians, and 3 (out of 10) MNIST digits. For both experiments, the authors demonstrate that their algorithm allows for a better estimate of the relative mass of the modes.", "questions": "I have no questions for the authors.", "limitations": "Besides poor empirical evaluation, which does not allow for defining the limits of the proposed method, two following points should be discussed.\n\n- The choice of batch size for experiments is worrisome. According to the authors, they used a batch size of $10^5$ for the mixture of two Gaussians and $1024$ for the reduced-MNIST experiment. At the same time [2] uses $128$ negative and positive examples at every iteration.\n- The derived estimate is biased. It is a minor point but it is never mentioned in the paper.", "rating": 6, "confidence": 5, "soundness": 2, "presentation": 3, "contribution": 2, "strengths": "- The proposed algorithm and theoretical developments that justify the algorithm are novel and are definitely of interest to the generative modeling community.\n- The paper is clearly written, and all the results are well-presented.\n\n**EDIT:** During the discussion, the authors provided experimental results for CIFAR-10. However, I still have the following concerns:\n- The baseline in the empirical study for CIFAR-10 demonstrates much worse performance than in the prior work [2].\n- The statements I mentioned in my review are still missing references or justifications, e.g., \"Among these methods, EBMs are perhaps the simplest and most interpretable\" has to be either justified by references (since it's not properly studied in the paper) or removed.\n- The biasedness of the estimator has to be mentioned explicitly next to its introduction since it's an important property (downside) for optimization.\n\nDespite these concerns, I recommend the paper for acceptance.\n\n[2] Du, Y., Li, S., Tenenbaum, J., & Mordatch, I. (2020). Improved contrastive divergence training of energy-based models. arXiv preprint arXiv:2012.01316.", "weaknesses": "- The empirical evaluation of the proposed algorithm is insufficient to demonstrate the practical benefits of the proposed algorithm. Indeed, the considered experiments are rather toy examples compared to the relevant literature [1,2,3]. Since Algorithm 1 is the main development of the paper, this is a major flaw of the current submission.\n- The role of continuous-time analysis is not clear. Namely, Proposition 1 is never used after the introduction in Section 3.1. Thus, it has no immediate practical or theoretical implications. Even though I appreciate the result, I think the paper will benefit from moving this discussion to the Appendix.\n- The following statements in the introduction require references to the corresponding works with justifications.\n    - “Among these methods, EBMs are perhaps the simplest and most interpretable”. It’s also worth mentioning which properties of EBMs make them the most interpretable models.\n    - “methods in this class are provably unable to learn the relative weights of modes when they are separated by low-density regions”\n    - “Unlike the Fisher divergence, the cross-entropy is sensitive to the relative weights in multimodal distributions, but it unfortunately leads to training algorithms that are less justified theoretically and more delicate to use in practice”\n    - “Unfortunately, the approximations made in these algorithms are uncontrolled and they are known to induce biases similar to those observed with score-based methods”. Here it is not clear which biases of the score-based methods are mentioned.\n    - “Score-matching techniques and variants originate from [9, 21, 22]; their shortcoming in the context of EBM training is investigated in [23] and their blindness to the presence of multiple, imbalanced modes in the target density has been known for long and we refer to [24, 25]) for discussions.” In this context, it is necessary to mention [4] (and numerous papers based on it), which is a score-based technique that led to a number of state-of-the-art methods in generative modeling, at the same time, sharing many common properties with EBMs (e.g. density evaluation).\n\n[1] Nijkamp, E., Hill, M., Zhu, S. C., & Wu, Y. N. (2019). Learning non-convergent non-persistent short-run MCMC toward energy-based model. *Advances in Neural Information Processing Systems*, *32*.  \n[2] Du, Y., & Mordatch, I. (2019). Implicit generation and modeling with energy-based models. *Advances in Neural Information Processing Systems*, *32*.  \n[3] Du, Y., Li, S., Tenenbaum, J., & Mordatch, I. (2020). Improved contrastive divergence training of energy-based models. *arXiv preprint arXiv:2012.01316*.  \n[4] Song, Y., Sohl-Dickstein, J., Kingma, D. P., Kumar, A., Ermon, S., & Poole, B. (2020). Score-based generative modeling through stochastic differential equations. *arXiv preprint arXiv:2011.13456*.\n\n\n\n**EDIT:** During the discussion period, the authors provided experimental results for CIFAR-10. However, I still have the following concerns:\n- The baseline in the empirical study for CIFAR-10 demonstrates much worse performance than in the prior work [3].\n- The statements I mentioned in my review are still missing references or justifications, e.g., \"Among these methods, EBMs are perhaps the simplest and most interpretable\" has to be either justified by references (since it's not properly studied in the paper) or removed.\n- The biasedness of the estimator has to be mentioned explicitly next to its introduction since it's an important property (downside) for optimization.\n\nDespite these concerns, I recommend the paper for acceptance.\n\n[3] Du, Y., Li, S., Tenenbaum, J., & Mordatch, I. (2020). Improved contrastive divergence training of energy-based models. arXiv preprint arXiv:2012.01316.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We would like to thank the reviewer for their careful reading, positive view on the manuscript, and constructive suggestions. Below, we address the main points raised.\n\n**Scope of the experimental results.** Unfortunately we did not have the computing resources necessary to perform experiments on large datasets such as CIFAR-10 or ImageNet. However, we plan to perform these experiments in collaboration with other groups with more computational resources in a follow-up paper.\n\nWe also stress that our experiments on MNIST are nontrivial: indeed, by pruning the data set to 3 digits, we make the overlap between the classes smaller and the energy barriers between them higher, which is a challenge for EBM training especially if we want to recover the relative weights of these classes (as we do).\n\nFinally, let us mention that in the revised version we will include a more quantitative comparison between our proposed method and CD (PCD) for the MNIST experiment (including the test using the FID as metric, see Fig. 2  in the figure file the ***Common Reply***). In addition, we conducted additional tests on other synthetic examples involving high-dimensional discrete datasets (see the description given in the ***Common Reply*** and Fig. 3  in the figure file).\n\n**Role of continuous-time analysis.** We believe that the continuous-time analysis will be transparent to the reader familiar with the statistical physics literature and will help them build the intuition for the discrete-time results. For example, the proof of Proposition 1 immediately shows why the resampling step is theoretically consistent (see also our reply to Reviewer 1pft). \n\n**Statements requiring references and justifications.** \n* *Among these methods, EBMs are perhaps the simplest and most interpretable.* We believe that this stems from the fact that these models benefit from a century of intuitions from computational physics to guide the design and interpretation of the energy $U$ and help the sampling of its associated Gibbs distribution. This is stated in the introduction already.\n* *Methods in this class are provably unable to learn the relative weights of modes when they are separated by low-density regions.* This is a known issue, discussed e.g. in Sec. 4.4 of the review paper by Song & Kingma, *How to Train Your Energy-Based Models*, arXiv:2101.03288.  Mathematically, it can be shown that the metric induced by the score is weaker than the one induced by KL since the former is related to the Hessian of the latter. \n* *Unlike the Fisher divergence, the cross-entropy is sensitive to the relative weights in multimodal distributions, but it unfortunately leads to training algorithms that are less justified theoretically and more delicate to use in practice.* The main issue is the computation of the KL and its gradient, which is the problem that we address in our paper. \n* *Unfortunately, the approximations made in these algorithms are uncontrolled and they are known to induce biases similar to those observed with the score-based methods.*  Because these methods (unlike ours) do not perform stochastic gradient descent (SGD) on the KL divergence, it is difficult to analyze their convergence properties and the nature of their fixed point.\n* *\"Score-matching techniques and variants originate from [9, 21, 22]; their shortcoming in the context of EBM training is investigated in [23] and their blindness to the presence of multiple, imbalanced modes in the target density has been known for long and we refer to [24, 25]) for discussions.\" In this context, it is necessary to mention [4] (and numerous papers based on it), which is a score-based technique that led to a number of state-of-the-art methods in generative modeling, at the same time, sharing many common properties with EBMs (e.g. density evaluation).* We thank the reviewer for pointing out these references to us: we will cite them in the new version of our paper. As far as we can tell, these papers propose variants of EBM training that combines it with other techniques, but they do not guarantee that these methods perform training by effectively minimizing the KL divergence. We believe that our method could be used in concert with some of these approaches to remedy this deficiency.\n\n**Limitations:**\n* *Batch size:* In our new numerical experiments we show that the batch size can be reduced, and a gain in efficiency can also be achieved by using mini-batching (see also the reply to Reviewer d8UC).  In particular, we performed experiments with GMM using Algorithm 4 in the paper with mini-batches of various sizes over a total population of $10^4$, showing that the method performs well even with small mini-batches see Fig. 1 in the common reply. Similarly, we did experiments with MNIST using Algorithm 4 with a mini-batch size of $128$ samples over a sample size of $1024$. The performance is given in Fig. 2 of our common reply. The training time of different methods is provided in the caption of Fig. 2. These results show that our method is scalable with mini-batching, which accelerates training without much sacrifice of the performance. For more details on the values of hyperparameters that we used for MNIST experiments, please see our reply to Reviewer d8UC. \n \n* *Biased estimator:* Indeed our estimator is biased. This is a common issue with estimators used in the sequential Monte Carlo method. We did not observe any noticeable effect of this bias for the number of walkers we used."}, {"review_id": "GlhmLS0Sfj", "reviewer": "Reviewer_2q1e", "summary": "The paper reports an enhcements in the eficiency of the training of EBMs using Jarzinsky equality. ", "questions": "Can you please explain what would be the thermodynamics systems, States, temperature, thermal reservoir in the context of the training of a EBM?\nIn a thermodynamics context, the free energy is connected to the work done on the system by means of an inequality. (Delta_F<W). The equality is when the system is slowly varied from a state to another state. In irreversible thermodynamics, according to JE, no matter how fast the systems is transferred from one state to another, the equality remains valid. In JE, the temperature of the systems in equilibrium, T, has a physical meaning and that is the temperature of the heat reservoir with which the system had been in contact. What would be the interpretation of the T in the training steps of the energy based models using JE? \nIn JE, an average is taken over all possible realizations that drives the thermodynamics system from state A to state B. can you please explain how this averaging is performed during the training of the EBM? \nCan you please explain in the context of the thermodynamics, for example using free energy and external work, on why such efficiency is obtained?\n\n", "limitations": "I think a more physical understanding of why the performed modification enhances the training will have be great for gaining an understanding of the physics of the training. ", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "The paper has the potential to help understand the training of EBM better through a thermodynamics based approach. It is well written and elaborated. ", "weaknesses": "The paper shows a more efficient training of EBM. However, it does not explain why such efficient is obtained by performing the modification using the JE. I think the authors should dedicate a section to elaborate from thermodynamics perspective that why such improvements is achieved. ", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We would like to thank the reviewer for their careful reading, and their suggestion to better explain underpinning from the statistical physics of our approach based on the Jarzynski equality (JE).\n\n**Thermodynamic interpretation.** This interpretation is actually the same as for physical systems: since the EBM energy $E_\\theta$ evolves during training, work is being performed on the walkers and we need to account for it if we want to sample the associated Gibbs distribution consistently. This is precisely what our method achieves by associating a scalar weight to every walker that is used to correct the bias introduced by the work from evolving the energy. The gain in efficiency of our method comes from this consistency: in contrast, when one uses e.g. the PCD algorithm or persistent MCMC, we would have to evolve the parameters $\\theta$ infinitely slowly (i.e. adiabatically) compared to the walkers to guarantee that the latter exactly samples the distribution associated with the former (which is akin to using thermodynamic integration in the physical systems). Since an infinite time-scale separation cannot be enforced in practice, work is effectively being performed on the system when one uses the PCD algorithm or persistent MCMC, and neglecting the weights results in a bias.  Our approach based on the Jarzynski equality allows one to remove this bias. \n\nWe also note that in the setup of the paper (as in usually done with EBM), the temperature of the system is absorbed in the energy and therefore set to one. As for the average over the realizations, they are done in practice by averaging over the number of walkers used in the procedure, possibly using mini batches: this is akin to using a parallel implementation of JE. Our approach can also be generalized by making multiple passes on successive time-lags, which would make it more similar to what is usually done to apply JE, but we found that proceeding as we do with walkers was sufficient for our needs.\n\nFor the reader's convenience, we will follow the reviewer's suggestion and add a section in the Appendix discussing this physical interpretation in more detail. "}, {"review_id": "fL5huqDVs0", "reviewer": "Reviewer_YzoY", "summary": "The authors consider EBMs family of algorithms optimized with a cross-entropy objective. They proposed a modification of the unadjusted Langevin algorithm (ULA) to remove bias from the optimization procedure, basically introducing correction coefficients for UCL gradient estimation. Their method is based on Jarzynski equality and is theoretically justified.  Moreover, the paper confirms empirically the method is less biased as it can recover all \"modes\" and proportions between them from data correctly in contrast to other baselines optimized with the cross-entropy loss. ", "questions": "1) The method can be used for density estimation. Is the proposed method beneficial compared to other density estimation methods? Is it worth evaluating it against other methods on the density estimation task? \n2) I am interested in a side-by-side plot with true density and density approximation fitted with a neural network. One can use the toy paper example with two Gaussian modes or more complex distributions. Does this experiment make sense? \n", "limitations": "--", "rating": 7, "confidence": 2, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1) The method is theoretically justified and novel to the best of my knowledge.\n2) The method can be used to find a normalizing constant and estimate density.\n3) The author demonstrates that the method covers \"modes\" correctly on two toy examples:  a Gaussian distribution with two modes, and a label-unbalanced MNIST dataset. The proposed method successfully recovers all the modes and their respective proportions, which is in contrast with the baseline ULA method. The toy experiment designs are great.\n4) The method is easy to implement and extra (compared to ULA) computational cost is low. \n5) The paper is well-written and easy to follow.\n", "weaknesses": "1) The empirical evaluation is weak. The paper only considers two toy datasets. It would be great to explore the proposed method on more challenging datasets. However, the authors noticed it themself in the conclusion. \n2) The paper did not use standard metrics to measure performance. Along with more complex datasets, it would be great to measure and compare performance with standard metrics such as FID,  Inception score, recall, and precision. I am especially interested in the recall and IS scores as they are most related to how the authors measure correct proportion and mode-covering in their toy experiments.  \n3) It is not clear how the new method compares in terms of sample quality to baselines (there are no metrics). \n", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We would like to thank the reviewer for their careful reading, positive view on the manuscript, and constructive suggestions. Below, we address the main points raised.\n\n**Scope of the numerical experiments.** As stated in the common reply, we did not have the computing resources necessary to perform experiments on large datasets such as CIFAR-10 or ImageNet, unfortunately. However, we plan to perform these experiments in collaboration with other groups with more computational resources in a follow-up paper.\n\n**Usage of standard metrics.** In the new version, we will report the FID score as the metric for evaluating the performance of our model on MNIST, see Figure 2 of our Common reply and see also our reply to reviewer 1pft on why we do not think FID score is not good for measuring the performance of EBM training methods on MNIST. As for the Inception Score, we have also tried that but it did not produce reasonable results because the pre-trained Inception V3 network used for computing the IS is trained on natural images, which are fundamentally different than MNIST images.\nThis comparison can be also performed in the GMM experiment, as well as in the new experiments. \n\n\n**Questions about density estimation and direct visualization.** It is not clear how to perform this comparison in general. We did however perform new experiments on a high-dimensional discrete-data set that allows for direct visualization of the learned energies (i.e. the density in logscale) in 2D, see the description and Fig. 3 in the Common Reply. "}, {"review_id": "4XLhbAD79c", "reviewer": "Reviewer_d8UC", "summary": "The paper under consideration proposes a modification of the standard EBM training procedure inspired by Jarzynski equality from thermodynamics. The idea is as follows. It is known that EBM training requires sampling from the distribution for which we know only its energy function given by, e.g., neural network $E_{\\theta}$. Typically, this is done by ULA algorithm or its analogues. Given current NN parameters $\\theta$, we try to sample from $p_{\\theta} \\sim \\exp( - E_{\\theta})$ using the corresponding MCMC algorithm and only after that we update the parameters $\\theta$. The authors take a step further and propose to change the paradigm. They consider a specific discrete procedure which resembles Annealed Importance Sampling MCMC algorithm. This procedure evolves both particles and model (NN) parameters so that the parameters turn out to follow gradient flow with respect to KL divergence. The authors validate and show the advantages of ther idea on moderate-dimensional Gaussian Mixture setup and MNIST dataset.\n\n**EDIT 2:** During the discussion, the authors came up with additional CIFAR10 experiment, which validate the applicability of their approach. On the one hand, the provided results are far from being SOTA (for example, [1] has much better metrics). On the other hand, the authors utilize the simplest CNNs and, from my point of view, such experiments are already valuable because they show that the method works without tricky architectural and (I guess) hyperparameter engineering. So, it is a good proof-of-concept setup. Based on this, I consider rising my score by one (to **7**) and recommend accepting the paper. At the same time, I would suggest authors to devote some time improving the architectures and tuning hyperparameters in order to achieve the performance close to [1] or better. For now, the paper is more methodological and theoretically-biased. From my point of view, a good NeurIPS paper should not underestimate practical aspects.\n\n[1] Du et. al., Improved Contrastive Divergence Training of Energy-Based Models, http://proceedings.mlr.press/v139/du21b/du21b.pdf", "questions": "- Line 163: Are there theoretical guarantees that the integral (which defines $\\hat{\\rho}(t, x)$) converges? \n- A natural question from me as a practitioner of EBMs: how many walkers $N$  we need to support in order to Algorithm 1 (or Algorithm 4) stable converge? Some bounds or at least kind of ablation study for your considered problems (MoG and MNIST) would be quite interesting. \n- Which parameters of Algorithm 1 were used in the MNIST experiment?\n- What is the Assumption (A1) you use in Lemma 1 and Lemma 2 in the appendix? I didn’t manage to find such an assumption in the paper. \n\nMisprints\n- Line 170: In the second equation of the inline formula, one should add the integration w.r.t. $x$.\n- Line 180: “that that” -> “that”.\n- Algorithm 1, line 2: it seems $A_0^i$ should be initialized as zeros.\n- Line 433: “given and” -> “given”.\n- Formula (51), denominator: $U_z(x)$ -> $\\exp(- U_z(x))$.\n- Figure 10 (about MNIST dataset) label refers to formula (25) (which is about MoG experiment). There should be another formula.", "limitations": "The potential negative societal impact of the work are not addressed. However, it seems to be the save as for thousands of papers which deals with generative modelling. So I have no special concerns.", "rating": 7, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "The paper seems to present a novel and theoretically sound idea. To be more precise, I would like to stress out the following points:\n- In spite of the Annealed Importance Sampling (AIS) is known for the decades, and it is also inspired by Jarzynski equality, the paper under consideration is not just utilizes AIS for sampling from an unnormalized distribution, but proposes a theoretically-grounded and assumptions-free procedure which prescribes how to progress particles along with model parameters to consistently converge to the distribution of interest.\n- The proposed method eliminates the problem of slow mixing faced by conventional EBMs methods. It seems that, given large enough particles (walkers), we can just launch the proposed algorithm, and it will succeed in recovering the desired energy function. We don’t need to think about how much steps of Langevin samples should be done for each parameter's update step. \n- The authors demonstrate several more-or-less convincing setups that demonstrate the superiority of their method compared to standard CD and PSD approaches. Additionally, for simple 1D setup they provide the theoretical arguments which supports their thesis regarding this supremacy (to be honest, I didn’t get into details of the analysis in Appendix C, but it seems correct).", "weaknesses": "From my point of view, the main weaknesses of the paper under consideration are as follows:\n- As I understand, the Algorithm 1 can not be fairly adapted for mini-batch setup and requires evolving simultaneously numerous particles (of training dataset order). In other words, for each gradient step w.r.t. model parameters, we need to evolve the whole dataset of points (Algorithm 4 also requires updates for each point from the dataset). From my point of view, it is the really great drawback which limits method applicability. \n- Related to the previous point, it is not clear if the proposed approach scales well to more sophisticated setups like CIFAR-10, Celeba etc. Will all these technical resampling procedures (described in Section 3.3)  work well for these tasks? So, at the current stage, I treat the paper as having limited practical validation in terms of plausibly looking images.\n- There are no quantitative comparison between proposed method and CD (PCD) for MNIST experiment. Showing just images is not very convincing. Perhaps, the better classes balance is achieved by cost of quality degradation.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We would like to thank the reviewer for their careful reading, positive view on the manuscript, and constructive suggestions. Below, we address the main points raised.\n\n**Using mini-batches.** *Algorithm 1 can be easily generalized to use mini-batches*, see Algorithm 4 in the paper, and it does not require evolving all the walkers at every SGD step. The procedure is very simple:\n1. Store the position and the scalar weights of all the walkers;\n2. At each step of SGD, evolve only the position and the weights of a subset of walkers, while keeping the others untouched in a replay buffer;\n3. When a walker is picked from the replay buffer, first adjust its weight to the current value of $U_\\theta$ (which only requires calculating this energy and not its gradient), then evolve its position and weight. \n\nLet us also note that using mini-batching is better justified in our method than in the PCD algorithm. Indeed, not evolving all the walkers in the latter increases the bias (that is, the mismatch between the empirical distribution of these walkers and the Gibbs distribution associated with the current $U_\\theta$); in contrast, our approach removes this bias using the weights *whether the walkers evolve (i.e. are part of the mini-batches) or not*. Of course, the reweighting cannot do miracles if the number of walkers is not large enough -- all sampling methods have the same issue. However, our approach has the advantage that its sample efficiency can be estimated via the calculation of the effective sample size (which involves the weights) and controlled via resampling.\n\n**Scalability of the resampling procedure.** The resampling procedure only involves selecting walkers using their weights. Since these weights are scalar quantities, the cost of this step is independent of the dimensionality of the problem. \n\n**Specific questions** \n* *Line 163: Are there theoretical guarantees that the integral (which defines $\\hat\\rho_\\theta(t,x)$) converges?*\nSince $\\int_0^\\infty \\int_{\\mathbb{R}^d} f(t,a,x) dx da = Z_{\\theta(t)}/Z_{\\theta_0} $ (Eq. 12), this integral can only diverge if $Z_{\\theta(t)}$ does. This cannot happen under our assumption that the target density is of the form $e^{-U_*(x)}/Z_*$ with $Z_*<\\infty$ if $\\theta(t)$ evolves by GD on the KL divergence. We will add a remark to this effect. \n* *Number of walkers needed*  Our method is an instance of a sequential Monte Carlo algorithm: these algorithms are widely used in practice, but they offer few theoretical guarantees, unfortunately. We have however performed more ablation studies on the GMM and MNIST experiments.\n* *Parameters of Algorithm 1 used in the MNIST experiment*\nWe use $n = 9800$, $N = 1024$, $h = \\texttt{5e-3}$. As for the total number of training steps, we run $800$ epochs with batch size $128$.  \n* *What is the Assumption (A1) you use in Lemma 1 and Lemma 2 in the appendix?* Assumption (A1) is defined above line 456. It is defined inline but we will highlight it in a different way to make it clearer.  \n\n**Misprints.**\n     Thank you for pointing these out! We will correct them. And, yes, for Fig. 10, it should be equation (21).  "}, {"review_id": "fZUQ2UbwYO", "reviewer": "Reviewer_nHa5", "summary": "This paper uses Jarzinsky Equality to analyze the training of Energy Based Model (EBM) using cross entropy loss. The authors highlight the problem of directly approximating the gradient using the unweighted mean of nonequilibrium samples and propose a modified algorithm based on reweighting.", "questions": "1. For the resampling step, how to assign new value to each walker $X_{kr}^{i}$ that needs to be resampled? Do you just resample $X_{kr}^{i}$ from $\\rho_{\\theta0}$? Also why can you reset $A{kr}^{i}=0$? For my understanding, according to eq5 and eq14, you should keep track of the whole chain of $A_t$ and $X_t$ through the training and sampling process, or the results might be incorrect/\n\n2. How can you generate new samples after the model is trained? I ask this because there can be a propblem that an EBM model trined with persistent MCMC have difficulty in generating new samples after the model is trained.  \n\n3. The review of related work on EBM training seems to be a bit rough and misses many recent works for EBM training and scaling up. For example:   \n1. A theory of generative convnet.\n2. Flow contrastive estimation of energy-based model.\n3. Your classifier is secretly an energy based model and you should treat it like one\n4. Vaebm: A symbiosis between variational autoencoders and energy-based model.\n5. Generalized energy based mode. \n6. A tale of two flows: Cooperative learning of langevin flow and normalizing flow toward energy-based model.\n7. Guiding energy-based models via contrastive latent variable.", "limitations": "The authors have adequately addressed the limitations.", "rating": 4, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "I like this idea of using Jarzinsky Equality to analyze the current training problem existing in EBM trained with cross entropy loss. I think the analysis is reasonable and hit the core problem of EBM training.", "weaknesses": "However, I remain unconvinced by the proposed algorithm, particularly regarding its potential for practical scalability. I have the following concerns:\n\n1. Sample inefficiency: The proposed method utilizes a reweighted version of gradient updates. While this approach may be theoretically sound for calculating expectations, I am apprehensive about its sample efficiency. Specifically, the reweighting term takes the form of an exponential, which can rapidly become either very large or very small. Consequently, a large batch of samples may need to be maintained and updated at every step. Additionally, with more complex datasets, an increased number of individual chains/walkers may be required, posing scalability challenges for the model.\n\n2. Moreover, the implementation necessitates the tracking of additional information such as $A_{t}^{i}$ and $X_{t}^{i}$. This could impose significant computational and memory burdens.\n\n3. Insufficient experiments: The only complex dataset showcased in this work is a subset of MNIST, containing only three digits. Furthermore, no qualitative measures of overall sample quality, such as FID or Inception Score, are presented. Consequently, with the current results, it is difficult for me to believe that the proposed methods can outperform the original CE and yield satisfactory results. It will be more convincing if the authors can justify their algorithm on more complex dataset like cifar10 or cifar100.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "**Sampling efficiency.** As with our method, the sample size (i.e. the number of wakers) must be large enough to guarantee statistical accuracy. There are few theoretical guarantees about this issue, unfortunately, and it is known that the Jarzynski weight can indeed become large if the sample size is too small, which impacts the accuracy. We stress however that this problem can be mitigated because these weights allow us to estimate the effective sample size (ESS) (which other methods cannot easily do), and the resampling step can be used to control these weights along with the ESS.\n\nWe also stress that our method also allows for mini-batching (see Algorithm 4 in the paper), as explained in our reply to Reviewer d8UC. As stated in this reply, the mini-batching procedure  is very simple:\n1. Store the position and the scalar weights of all the walkers;\n2. At each step of SGD, evolve only the position and the weights of a subset of walkers, while keeping the others untouched in a replay buffer;\n3. When a walker is picked from the replay buffer, first adjust its weight to the current value of $U_\\theta$ (which only requires calculating this energy and not its gradient), then evolve its position and weight. \n\nThis makes the cost of our method comparable to that of the CD or PCD algorithm, though the mini-batching is better justified in our method than in PCD: Indeed, not evolving all the walkers in the latter increases the bias (that is, the mismatch between the empirical distribution of these walkers and the Gibbs distribution associated with the current $U_\\theta$); in contrast, our approach removes this bias using the weights *whether the walkers evolve (i.e. are part of the mini-batches) or not*. \n\n**Variables to track.** Tracking the position $X^i_k$ of the walkers in our method is not different than tracking their positions in the CD and the PCD algorithms. Compared to these methods, we must additionally track the values of the weights $A^i_k$: we stress however that these weights are *scalar* quantities, so there is no noticeable extra cost in storing them. There is also no extra cost from evolving the weights in the mini-batch since this operation uses the gradient $\\nabla U_\\theta(X_k)$ which is already needed to evolve the positions of walkers in the mini-batch (note that this gradient is also needed in the CD and the PCD algorithms).\n\n**Insufficient experiments.** As stated in the common reply, we did not have the  computing resources necessary to perform experiments on large datasets such as CIFAR-10 or ImageNet, unfortunately. However, we plan to perform these experiments in collaboration with other groups with more computational resources in a follow-up paper.\n\nAs also stated in the Common Reply, we did perform additional standard tests on the experiments we have, and now report FID for the MNIST experiments along with ablation studies involving varying the number of walkers and the size of the mini-batches, see Fig. 2 in the Common Reply (this is also done in the GMM experiments, see Fig. 1). In addition, we performed new experiments on a high-dimensional discrete data set, see the description in the Common Reply and Fig. 3.\n\n**Questions:**\n* *Resampling step.* The resampling step involves splitting walkers with large weights, and killing those with low weights. This operation is done within the current population of the walkers (i.e. it does not involve resampling from an external distribution such as $\\rho_{\\theta_0}$ or from the data set, but rather from the empirical distribution of the walkers). This step is performed to control the ESS, and it is standard in sequential Monte Carlo methods: after resampling, all the values of all the log-weights $A^i_k$ are reset to zero, i.e. the ESS is again the total number of walkers (though other strategies evolving partial resampling are also possible). We also stress that the resampling step only involves selecting walkers using their weights. Since these weights are scalar quantities, the cost of this step is independent of the dimensionality of the problem.  \n\n* *Generation of new samples.* In contrast with other methods (like the CD or PCD algorithm, or their variant involving persistent MCMC), at all iterations during training the weighted walkers in our approach sample the EBM by construction. As a result, they can directly be used as new samples, and we can keep evolving them after the training stops to get more new samples. In the procedure, these samples are weighted, but all weights become identical after resampling: these \"resampled samples\" are the ones we can use e.g. as generated images (as we did in the MNIST experiments, see Fig. 2 in the paper). In this sense, our approach truly is an implementation of an EBM that can be used as a generative model.\n\n* *Review of the literature.* We thank the reviewer for pointing out these references to us: we will cite them in the new version of our paper. As far as we can tell, these papers propose variants of EBM training that combine it with other techniques, but they do not guarantee that these methods perform training by effectively minimizing the KL divergence. We believe that our method could be used in concert with some of these approaches to remedy this deficiency."}]

Author rebuttal:
We would like to thank all of the reviewers for their careful reading, constructive suggestions, and overall interest in our work. In this reply, we answer to the main questions common to all reviewers; below, we address their concerns individually and give more details on each point.

***Regarding the theory part:*** Our approach permits the calculation of the KL divergence and its gradient by adding scalar weights to the walkers used to do the sampling. This calculation is exact in the limit of infinitely many walkers and can be simulated using standard techniques from sequential Monte-Carlo. We stress that our method is no more costly than the CD or PCD algorithms. In particular (all these points are elaborated upon in our individual replies to the reviewers):
1. Our resampling step can be justified theoretically and is statistically equivalent to evolving the weights (assuming that the number of walkers is infinite). Since these weights are scalar, their storage and evolution involve no noticeable extra cost.
2. When the number of walkers is finite, the resampling step helps to improve the sample efficiency of our approach since this step allows us to control the effective sample size (which, with methods like CD or PCD, is harder to measure and control since its calculation involves the walkers weights that our method introduces). 
3. The resampling step only involves selecting walkers using their weights. Since these weights are scalar quantities, the cost of this step is independent of the dimensionality of the problem. 
4. Our method can easily be used with mini-batches, in which only a subset of the walkers are evolving at any given time, and the rest are kept in a replay buffer, see Algorithm 4 in the paper. In fact, as we explain in our reply to Reviewer d8UC, using mini-batches with our approach is better justified than doing so in the PCD algorithm. To emphasize this point better, we will report experiments with mini-batches both in the GMM example (see Fig. 1 in the figure file) as well as MNIST (see Fig. 2).


***Regarding the numerical experiments:*** We agree with the reviewers that our method should be tested on larger and more realistic datasets such as CIFAR-10 or ImageNet but unfortunately we do not have the computing resources necessary to perform these experiments. We plan to perform these simulations in collaboration with other groups with more computational resources in a follow-up paper. We also believe that the strong theoretical underpinning of our approach is sufficient to justify its interest for the NeurIPS community. 

We also stress that our experiments on MNIST are nontrivial: indeed, by pruning the dataset to 3 digits, we make the overlap between the data classes smaller, which is a challenge for EBM training especially if we want to recover the relative weights of these classes (as we do). To emphasize this point,  in the revised version we will include a more quantitative comparison between our proposed method and CD (PCD) for our MNIST experiment (see Fig. 2 in the figure file).

Finally, we conducted additional tests on other synthetic examples involving high-dimensional discrete datasets (see the description given below and Fig. 3 in the figure file). These new tests show again that our method performs better than the CD and PCD algorithms at a similar numerical cost. (They also show that our approach can be straightforwardly generalized to EBM defined on discrete datasets, for example graph-structured datasets, on which we plan to test our method on subsequent works)

**Experimental setup of the synthetic experiment involving discrete datasets:** 
We used the experimental setup of Refs [1,2]. We first generate synthetic 2D data $x \in \mathbb{R}^2$ from classical datasets (eg, Moons, 8Gaussians, etc). Then the floating-point number representation of $x$ (with precision up to $\texttt{1e-4}$ ) is turned into a $32$-bit code constructed using concatenated Gray codes (https://en.wikipedia.org/wiki/Gray\_code), which results in a new distribution on the discrete space $\\{0,1\\}^{32}$. The transformation from $\mathbb{R}^2$ to $\\{0,1\\}^{32}$ is challenging since Gray codes are not linear. 
The discrete space has size $2^{32}$  (extremely large). We also introduce variants of the datasets where different modes are assigned different weights (Figure 4 in the figure file). Our codebase directly iterates on the ones found in Refs. [1,2]. 

This experimental setup allows nice visualizations (in the 2D space): visual analysis of the trained models is performed using the heatmaps of the learned energy functions (see Fig. 4 in the figure file). As a metric, we use Maximum Mean Discrepancy with a linear kernel over $\\{0,1\\}^{32}$ (corresponding to the Hamming distance) to compare ${4000}$ samples from the data distribution and $4000$ samples from our model. 

***Hyperparameters:*** We parametrize our EBMs with a 4-layer MLP, ELU activations, and hidden layer size $256$. We did not optimize hyperparameters and directly used the ones from Refs. [1,2]: we use the Adam optimizer with a learning rate of $\texttt{3e-4}$, weight decay of $\texttt{1e-4}$, gradient clipping at 1, batch size of $128$ and $100$ batches per training epoch, with $1000$ training epochs. 

The results in Figure 3 (in the figure file) show the competitiveness of our method. Figure 4 additionally shows how our method is able to retrieve the correct mode masses on imbalanced, multimodal datasets. 

**References**
[1] Hanjun Dai, Rishabh Singh, Bo Dai, Charles Sutton, and Dale Schuurmans. Learning discrete energy-based models via auxiliary-variable local exploration. Advances in Neural Information Processing Systems, 33:10443–10455, 2020.

[2] Meng Liu, Haoran Liu, and Shuiwang Ji. Gradient-guided importance sampling for learning binary energy-based models. arXiv preprint arXiv:2210.05782, 2022.

Meta-review: The authors proposed a modification of the standard procedure to train an energy-based model. Their method is based on Jarzynski equality and is theoretically justified. 

Initially, to demonstrate efficiency of the approach the authors consider GMMs and MNIST dataset, which is not sufficient from my point of view. However, during the rebuttal, the authors demonstrated efficiency also using CIFAR10 dataset. I think this is enough to justify that the method actually works.

Overall, taking into account other positive comments from reviewers I vote to accept the paper.
At the same time, it would be very important to take carefully all comments of the reviewers when finalising the paper.
Decision: Accept (poster)