Title: Efficient Training of Energy-Based Models Using Jarzynski Equality
Venue: NeurIPS 2023
Accepted: yes
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

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.

1 Introduction

Probabilistic models have become a 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 [1, 2, 3] generative adversarial networks [4, 5], normalizing flows [6, 7, 8, 9, 10], diffusion-based models [11, 12, 13], restricted Boltzmann machines [14, 15, 16], and energy-based models (EBMs) [17, 18, 19].

Inspired by physics, EBMs are unnormalized probability models, specified via an energy function U , with the underlying probability density function (PDF) defined as ρ = exp( -U ) /Z , where Z is a normalization constant: this constant is often intractable but crucially it is not needed for data generation via Monte-Carlo sampling. In statistical physics, such PDFs are called Boltzmann-Gibbs 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 a 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 a 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 applications: for example, under-representation of elements from a particular population. It is important to ensure that these biases are either avoided or compensated, leading to methods that ensure fairness ([23]); however, the training routines for EBMs often have a 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. In 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, ∇ U = -∇ log ρ . 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 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. Indeed, gradient-based optimization procedures on the cross-entropy require sampling from the model distribution using Markov-Chain Monte-Carlo (MCMC) methods or the unadjusted Langevin algorithm (ULA) until mixing, which in a high-dimensional context or without log-concavity 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 crossentropy, 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 to overcome this issue, for example, based on MCMC sampling [29, 30, 31] or on various regularizations of CD [32, 33]. 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 as 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.

We also apply our method on the MNIST and CIFAR-10 datasets. In the first, we intentionally bias the proportion of the different digits, to 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 [39, 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 [43] 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, 26, 27]. 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 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; [33] introduced a way to approximate the missing term in the gradient of the CD objective; and [28, 55] 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 in the method we propose: it is always the cross-entropy.

Jarzynski equality (JE) was introduced in [34] and gives an exact expression relating the normalizing constants (or equivalently free energy [22]) of two distributions linked by out-of-equilibrium 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 methods have been used in the context of generative models based on variational autoencoders [57, 58], normalizing flows [59, 60], and diffusion-based models [61, 62]. In contrast, here we use sequential Monte-Carlo method to train EBM on the cross-entropy directly.

2 Energy-Based Models

Setup, notations, and assumptions. The problem we consider can be formulated as follows: we assume that we are given n ∈ N data points { x ∗ i } n i =1 in R d drawn from an unknown probability distribution that is absolutely continuous with respect to the Lebesgue measure on R d , with a positive probability density function (PDF) ρ ∗ ( x ) > 0 (also unknown). Our aim is to estimate this PDF via an energy-based model (EBM), i.e. to find a suitable energy function in a parametric class, U θ : R d → [0 , ∞ ) with parameters θ ∈ Θ , such that the associated Boltzmann-Gibbs PDF

is an approximation of the target density ρ ∗ ( x ) . The normalization factor Z θ is known as the partition function in statistical physics [22]. This factor is hard to estimate and one advantage of EBMs is that they provide generative models that do not require the explicit knowledge of Z θ since Markov Chain Monte-Carlo (MCMC) methods can in principle be used to sample ρ θ knowing only U θ - the design of such MCMC methods is an integral part of the problem of building an EBM.

To proceed we will assume that the parametric class of energy we use is such that, for all θ ∈ Θ ,

These assumptions guarantee that Z θ < ∞ (i.e. we can associate a PDF ρ θ to U θ via (1) for any θ ∈ Θ ) and that the Langevin equations as well as their time-discretized versions we will use to sample ρ θ have global solutions and are ergodic [63, 64, 65]. We stress that (2) does not imply that U θ is convex (i.e. that ρ θ is log-concave): in fact, we will be most interested in situations where U θ has multiple local minima so that ρ θ is multimodal. For simplicity we will also assume that ρ ∗ is in the parametric class, i.e. ∃ θ ∗ ∈ Θ : ρ θ ∗ = ρ ∗ . Our aims are primarily to identify this θ ∗ and to sample ρ θ ∗ ; in the process, we will also show how to estimate Z θ ∗ .

Cross-entropy minimization. To measure the quality of the EBM and train its parameters one can use the cross-entropy of the model density ρ θ relative to the target density ρ ∗

where we used the definition of ρ θ in (1) to get the second equality. The cross-entropy is related to the Kullback-Leibler divergence via H ( ρ θ , ρ ∗ ) = H ( ρ ∗ ) + D KL ( ρ ∗ || ρ θ ) , where H ( ρ ∗ ) is the entropy of ρ ∗ , and its gradient with respect to the parameter θ can be calculated using the identity ∂ θ log Z θ = -∫ R d ∂ θ U θ ( x ) ρ θ ( x ) dx , to obtain

The cross-entropy is more 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 ρ ∗ separated by low-density regions [36]. Unfortunately, the cross entropy is also much harder to use in practice since evaluating it requires estimating Z θ , and evaluating its gradient requires calculating the expectation E θ [ ∂ θ U θ ] (in contrast E ∗ [ U θ ] and E ∗ [ ∂ θ U θ ] can be readily estimated on the data). Typical training methods, e.g. based on the CD or the PCD algorithms, give up on estimating Z θ and resort to various approximations to calculate the expectation E θ [ ∂ θ U θ ] -see Appendix A.5 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.

3 Training via sequential Monte-Carlo methods based on Jarzynski equality

In this section we use tools from nonequilibrium statistical mechanics [34, 56] to write exact expressions for both E θ [ ∂ θ U θ ] and Z θ (Sec 3.1) that are amenable to empirical estimation via sequential Monte-Carlo methods [35], thereby enabling gradient descent-type algorithms for the optimization of EBMs (Sec. 3.2).

3.1 Jarzynski equality in discrete-time

Proposition 1. Assume that the parameters θ are evolved by some time-discrete protocol { θ k } k ∈ N 0 and that (2) hold. Given any h ∈ (0 , L ) , let X k ∈ R d and A k ∈ R be given by the iteration rule

where U θ ( x ) is the model energy, { ξ k } k ∈ N 0 are independent N (0 d , I d ) , and we defined

Then, for all k ∈ N 0 ,

where the expectations on the right-hand side are over the law of the joint process ( X k , A k ) .

Algorithm 1 Sequential Monte-Carlo training with Jarzynski correction

1: Inputs: data points { x i ∗ } n i =1 ; energy model U θ ; optimizer step opt ( θ, D ) using θ and the empirical CE gradient D ; initial parameters θ 0 ; number of walkers N ∈ N 0 ; set of walkers { X i 0 } N i =1 sampled from ρ θ 0 ; total duration K ∈ N ; ULA time step h ; set of positive constants { c k } k ∈ N .

2: A i 0 = 0 for i = 1 , . . . , N .

3: for k = 0 , . . . , K -1 do

4: p i k = exp( A i k ) / ∑ N exp( A j

j =1 k ) ▷ normalized weights

5: ˜ D k = ∑ N i =1 p i k ∂ θ U θ k ( X i k ) -n -1 ∑ n j =1 ∂ θ U θ k ( x j ∗ )

6: θ k +1 = opt ( θ k , ˜ D k )

7: for i = 1 , ..., N do

10: end for

empirical CE gradient

▷

optimization step

11:

Resample the walkers and reset the weights if ESS

12: end for

13: Outputs: Optimized energy U θ K ; set of weighted walkers { X i K , A i K } N i =1 sampling ρ θ K ; partition function estimate ˜ Z θ K = Z θ 0 N -1 ∑ N i =1 exp( A i K ) ; CE estimate log ˜ Z θ K + n -1 ∑ n j =1 U θ K ( x j ∗ )

̸

The proof of the proposition is given in Appendix A.2: for completeness we also give a continuoustime version of this proposition in Appendix A.1. We stress that the inclusion of the weights in (7) is key, as E [ ∂ θ U θ k ( X k )] = E θ k [ ∂ θ U θ k ] in general. We also stress that (7) holds exactly despite the fact that for h > 0 the iteration step for X k in (5) is that of the unadjusted Langevin algorithm (ULA) with no Metropolis correction as in MALA [66, 67]. That is, the inclusion of the weights A k exactly corrects for the biases induced by both the slow mixing and the time-discretization errors in ULA.

Proposition 1 shows that we can evolve the parameters by gradient descent over the cross-entropy by solving (5) concurrently with

where γ k > 0 is the learning rate and k ∈ N 0 with θ 0 given. We can also replace the gradient step for θ k in (8) by any update optimization step (via AdaGrad, ADAM, etc.) that uses as input the gradient D k of the cross-entropy evaluated at θ k to get θ k +1 . Assuming that we know Z θ 0 we can track the evolution of the cross-entropy via

3.2 Practical implementation

Empirical estimators and optimization step. We introduce N independent pairs of walkers and weights, { X i k , A i k } N i =1 , which we evolve independently using (5) for each pair. To evolve θ k from some prescribed θ 0 we can then use the empirical version of (8):

where ˜ D k is the estimator for the gradient in θ of the cross-entropy:

These steps are summarized in Algorithm 1, which is a specific instance of a sequential Monte-Carlo algorithm. We can also use mini-batches of { X i k , A i k } N i =1 and the data set { x j ∗ } n j =1 at every iteration (see Algorithm 2 in Appendix A.3), and switch to any optimizer step that uses θ k and D k as input to get the updated θ k +1 . During the calculation, we can monitor the evolution of the partition function and the cross-entropy using as estimators

k

+1

< c

k

+1

, see (13).

▷

resampling step

▷

These steps are summarized in Algorithm 1, which is a specific instance of a sequential Monte-Carlo algorithm. We adapted our routine to mini-batches 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, N < ∞ and n < ∞ . Regarding n , we may need to add a 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 i k 's will spread away from zero during the optimization, implying that the weights exp( A i k ) will become non-uniform, thereby reducing the effective sample size. This is a known issue with sequential Monte-Carlo 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, i.e.

The effective sample size of the N walkers is ESS k N . Initially, since A i 0 = 0 , ESS 0 = 1 , but it decreases with k . At each iteration k r such that ESS k r < c k r , where { c k } k ∈ N is a set of predefined positive constants in (0 , 1) , we then:

Resample the walkers X i k r using p i k r = e A i kr / ∑ N j =1 e A j kr as probability to pick walker i ;

Reset A i k r = 0 ;

Use the update Z θ k = Z θ kr N -1 ∑ N i =1 exp( A i k ) for k ≥ k r until the next resampling step.

This resampling is 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 i k evolve, the algorithm produces weighted samples X i k . At any iteration, however, equal-weight samples can be generated by resampling. Notice that, even if we no longer evolve the model parameters θ k , the algorithm is such that it removes the bias from ULA coming from h > 0 - this bias removal is not perfect, again because N is finite, but this can be controlled by increasing N at the stage when the EBM is used as a generative model.

4 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

where a ∗ , b ∗ ∈ R d specify the means of the two modes and z ∗ ∈ R controls their relative weights p ∗ = 1 / (1 + e -z ∗ ) and q ∗ = 1 -p ∗ = e -z ∗ / (1 + e -z ∗ ) . The values of a ∗ , b ∗ , z ∗ are carefully chosen such that the modes are well separated and the energy barrier between the modes is high 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

where θ = ( a, b, z ) are the parameters to be optimized. We choose this model as it allows us to calculate the partition function of the model at any value of the parameters, Z θ = (2 π ) d/ 2 (1 + e -z ) . We use this information as a benchmark to compare the prediction with those produced by our method.

0

50

100

150

200

250

Rescaled Iterations

0.0

0.1

0.2

0.3

0.4

0.5

p

Mass of the first mode

p

*

Algo 1

PCD

CD

0

50

100

150

200

250

Rescaled Iterations

10

8

6

4

2

0

2

4

6

a

,

b

Means

Algo 1

PCD

CD

a

*

b

*

0

50

100

150

200

250

Rescaled Iterations

10

4

10

3

10

2

10

1

10 0

10 1

KL

KL divergence

Algo 1

Eq. (21)

PCD

CD

In our numerical experiments, we set d = 50 , use N = 10 5 walkers with a mini-batch of N ′ = 10 4 and n = 10 5 data points. We initialize the model at θ 0 = ( a 0 , b 0 , z 0 ) with a 0 and b 0 drawn from an N (0 , ϵ 2 I d ) with ϵ = 0 . 1 and z 0 = 0 , meaning that the initial ρ θ 0 is close to the PDF of an N (0 , I d ) . The training is performed using Algorithm 2 with h = 0 . 1 and fixed learning rates γ k = 0 . 2 for a k and b k and γ k = 1 for z k . We perform the resampling step by monitoring ESS k defined in (13) with constant 1 /c k = 1 . 05 and using the systematic method. We also compare our results to those obtained using ULA with these same parameters (which is 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 h = 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 ( ρ θ k , ρ ∗ ) 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

The results are shown in Figure 1. As can be seen, all three methods learn well the values of a ∗ and b ∗ specifying the positions of the modes. However, only our approach learns the value of z 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 z , consistent with the fact that it effectively uses the Fisher divergence as the objective. The results 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 our 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, we perform empirical experiments on the MNIST dataset to answer the following question: when it comes to high-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 n = 9800 data points), so that in this manufactured dataset the digits are in have respective weights 4 / 7 , 2 / 7 , and 1 / 7 .

6333223

Z2322&22

242222d 2

3333362

3

3622b62

63233323

323232

２2236２3

6&63

222233

6666663

32B60322

23262332

2

2

2323

2223326

Generated samples(rightafterresampling)

ImagesfromMNiSTtestset

Generated samples (without Jarzynski)

Errorintheweightestimation(digit2)

Errorin the weightestimation(digit6)

w/oJarzynski

w/oJarzynski

w/oJarzynski

0.7

w/Jarzynski

2.0

w/Jarzynski

1.75

w/Jarzynski

0.6

1.50

1.5

Relativeerror

0.5

1.25

0.4

1.00

1.0

0.3

0.75

0.2

0.5

0.50

0.1

0.25

0.0

2

0.0

0.00

0

100

200

300

400

500

600

700

800

0

100

200

300

400

500

600

700

800

0

100

200

300

400

500

600

700

008

Epoch

Epoch

Epoch

We train two EBMs to represent this data set, the first using our Algorithm 2, and the second using the PCD Algorithm 4. We represent the energy using a simple six-layer convolutional neural network with the swish activation and about 77 K parameters. We use the ADAM optimizer for the training with a learning rate starting from 10 -4 and linearly decaying to 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.

The results are shown in Figure 2. Both our Algorithm 1 and the PCD Algorithm 4 generate 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 3, we show the relative error on the weight estimation of each mode obtained by using a classifier pre-trained on the data set (see Appendix B.2 for 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.

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 as we do not know the 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 a plot of which is given in the right panel Figure 4 for the mini-batched training algorithm (Algorithm 2).

In these experiments with MNIST, we used the adaptive resampling scheme described after equation 13. In 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 4) a single resampling step was made. We also found empirically that the results are insensitive to the choice of of the parameters c k used for resampling: the rule of

Convergenceofthecross-entropyestimate

EffectiveSampleSize

Estimated cross-entropy

0

cross-entropyestimate

1.0

Size

-1

0.8

2

Sample

-3

0.6

Effective

-4

0.4

Resampling

-5

0.2

EffectiveSampleSize

-6

0

100

200

300

400

500

600

700

800

0

100

200

300

400

500

600

Epochs

Epochs

thumb we found is to not resample at the beginning of training, to avoid possible mode collapse, and use resampling towards the end, to improve the image quality.

4.3 CIFAR-10

We perform an empirical evaluation of our method on the full CIFAR-10 (32 × 32) image dataset. We use the setup of [47], with the same neural architecture with n f = 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 a combination of color distortion, horizontal flip, rescaling, and Gaussian blur augmentations, to help the mixing of the MCMC sampling and stabilize training).

GeneratedCIFAR-10sampleswithourapproachGeneratedCIFAR-10sampleswithPCD

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 η = 10 -4 and inject a Gaussian noise of standard deviation σ = 3 × 10 -2 to the dataset while performing gradient clipping in Langevin sampling for better performance. All the experiments were performed on a 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 1. They indicate that our method can achieve slightly better performance than the PCD algorithms w/ and w/o data augmentation at a similar computational cost. Furthermore, these results on CIFAR-10 suggest that our method scales well to complicated training tasks on more realistic data sets.

Table 1: textitCIFAR-10 dataset: Comparison of FID and Inception Score (IS) for PCD and Algorithm 2. Experiments performed using the neural architecture in [47] to model the energy.

It is worth stressing that a key component of the success of our method is 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 one to minimize the cross-entropy between the model and the target densities to train EBMs. Our approach rests on results from non-equilibrium thermodynamics (in particular, Jarzynski equality) that show how to take exact expectations over a 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 a 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 a 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.

References

Diederik P Kingma and Max Welling. An introduction to variational autoencoders. Foundations and Trends in Machine Learning , 12(4):307-392, 2019.

Carl Doersch. Tutorial on variational autoencoders. arXiv preprint arXiv:1606.05908 , 2016.

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.

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.

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.

Esteban G. Tabak and Eric Vanden-Eijnden. Density estimation by dual ascent of the loglikelihood. Communications in Mathematical Sciences , 8(1):217 - 233, 2010.

E. G. Tabak and Cristina V. Turner. A family of nonparametric density estimation algorithms. Communications on Pure and Applied Mathematics , 66(2):145-164, 2013.

Danilo Rezende and Shakir Mohamed. Variational inference with normalizing flows. In International conference on machine learning , pages 1530-1538. PMLR, 2015.

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.

Emile Mathieu and Maximilian Nickel. Riemannian continuous normalizing flows. Advances in Neural Information Processing Systems , 33:2503-2515, 2020.

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.

Jonathan Ho, Ajay Jain, and Pieter Abbeel. Denoising diffusion probabilistic models. In Advances in Neural Information Processing Systems , volume 33, pages 6840-6851, 2020.

Yang Song, Conor Durkan, Iain Murray, and Stefano Ermon. Maximum likelihood training of score-based diffusion models. Advances in Neural Information Processing Systems , 34: 1415-1428, 2021.

Geoffrey E Hinton. A practical guide to training restricted boltzmann machines. Neural Networks: Tricks of the Trade: Second Edition , pages 599-619, 2012.

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.

Hannes Schulz, Andreas Müller, Sven Behnke, et al. Investigating convergence of restricted boltzmann machine learning. In NIPS 2010 Workshop on Deep Learning and Unsupervised Feature Learning , volume 1, pages 6-1, 2010.

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.

Michael Gutmann and Aapo Hyvärinen. Noise-contrastive estimation: A new estimation principle for unnormalized statistical models. In Proceedings of the thirteenth international conference on artificial intelligence and statistics , pages 297-304. JMLR Workshop and Conference Proceedings, 2010.

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.

Ludwig Boltzmann. Weitere studien über das wärmegleichgewicht unter gasmolekülen. Kinetische Theorie II , pages 115-225, 1970.

Josiah Willard Gibbs. Elementary principles in statistical mechanics: developed with especial reference to the rational foundations of thermodynamics . C. Scribner's sons, 1902.

Evgenii Mikhailovich Lifshitz and Lev Petrovich Pitaevskii. Statistical physics: theory of the condensed state , volume 9. Elsevier, 2013.

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.

Aapo Hyvärinen and Peter Dayan. Estimation of non-normalized statistical models by score matching. Journal of Machine Learning Research , 6(4), 2005.

Geoffrey E Hinton. Training products of experts by minimizing contrastive divergence. Neural computation , 14(8):1771-1800, 2002.

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.

Miguel A Carreira-Perpinan and Geoffrey Hinton. On contrastive divergence learning. In International workshop on artificial intelligence and statistics , pages 33-40. PMLR, 2005.

Aapo Hyvarinen. Connections between score matching, contrastive divergence, and pseudolikelihood for continuous-valued variables. IEEE Transactions on neural networks , 18(5): 1529-1531, 2007.

Tijmen Tieleman. Training restricted boltzmann machines using approximations to the likelihood gradient. In International conference on Machine learning , pages 1064-1071, 2008.

Pierre E Jacob, John O'Leary, and Yves F Atchadé. Unbiased markov chain monte carlo methods with couplings. Journal of the Royal Statistical Society: Series B (Statistical Methodology) , 82 (3), 2020.

Yixuan Qiu, Lingsong Zhang, and Xiao Wang. Unbiased contrastive divergence algorithm for training energy-based latent variable models. In International Conference on Learning Representations , 2020.

Yilun Du and Igor Mordatch. Implicit generation and modeling with energy based models. In Advances in Neural Information Processing Systems , 2019.

Yilun Du, Shuang Li, Joshua Tenenbaum, and Igor Mordatch. Improved contrastive divergence training of energy-based models. In International Conference on Machine Learning , pages 2837-2848. PMLR, 2021.

C Jarzynski. Nonequilibrium equality for free energy differences. Physical Review Letters , 78 (14):2690, 1997.

Arnaud Doucet, Nando De Freitas, Neil James Gordon, et al. Sequential Monte Carlo methods in practice , volume 1. Springer, 2001.

Yang Song and Diederik P Kingma. How to train your energy-based models. arXiv preprint arXiv:2101.03288 , 2021.

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.

Will Grathwohl, Kuan-Chieh Wang, Joern-Henrik Jacobsen, David Duvenaud, Mohammad Norouzi, and Kevin Swersky. Your classifier is secretly an energy based model and you should treat it like one. In International Conference on Learning Representations , 2019.

Steve Brooks, Andrew Gelman, Galin Jones, and Xiao-Li Meng. Handbook of Markov chain Monte Carlo . CRC press, 2011.

Jun S Liu and Jun S Liu. Monte Carlo strategies in scientific computing , volume 75. Springer, 2001.

Pascal Vincent. A connection between score matching and denoising autoencoders. Neural computation , 23(7):1661-1674, 2011.

Kevin Swersky, Marc'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.

Li Kevin Wenliang. On the failure of variational score matching for V AE models. arXiv preprint arXiv:2210.13390 , 2022.

Li Wenliang, Danica J Sutherland, Heiko Strathmann, and Arthur Gretton. Learning deep kernels for exponential family densities. In International Conference on Machine Learning , 2019.

Yang Song and Stefano Ermon. Generative modeling by estimating gradients of the data distribution. In Advances in neural information processing systems , volume 32, 2019.

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 , 42(1):27-45, 2018.

Erik Nijkamp, Mitch Hill, Song-Chun Zhu, and Ying Nian Wu. Learning non-convergent nonpersistent short-run MCMC toward energy-based model. In Advances in Neural Information Processing Systems , volume 32, 2019.

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 IEEE/CVF Conference on Computer Vision and Pattern Recognition , pages 7518-7528, 2020.

Zhisheng Xiao, Karsten Kreis, Jan Kautz, and Arash Vahdat. Vaebm: A symbiosis between variational autoencoders and energy-based models. In International Conference on Learning Representations , 2020.

Jianwen Xie, Zilong Zheng, and Ping Li. Learning energy-based model with variational autoencoder as amortized sampler. In Proceedings of the AAAI Conference on Artificial Intelligence , volume 35, pages 10441-10451, 2021.

Jianwen Xie, Yaxuan Zhu, Jun Li, and Ping Li. A tale of two flows: Cooperative learning of langevin flow and normalizing flow toward energy-based model. In International Conference on Learning Representations , 2021.

Hankook Lee, Jongheon Jeong, Sejun Park, and Jinwoo Shin. Guiding energy-based models via contrastive latent variables. In The Eleventh International Conference on Learning Representations , 2022.

Omer Yair and Tomer Michaeli. Contrastive divergence learning is a time reversal adversarial game. In International Conference on Learning Representations , 2021.

Carles Domingo-Enrich, Alberto Bietti, Eric Vanden-Eijnden, and Joan Bruna. On energybased models with overparametrized shallow neural networks. In International Conference on Machine Learning , pages 2771-2782, 2021.

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:2107.05134 , 2021.

Radford M Neal. Annealed importance sampling. Statistics and computing , 11:125-139, 2001.

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 https://openreview.net/forum?id=BJ8c3f-0b .

Xinqiang Ding and David J. Freedman. Learning deep generative models with annealed importance sampling, 2020.

Laurence Illing Midgley, Vincent Stimper, Gregor N. C. Simm, Bernhard Schölkopf, and José Miguel Hernández-Lobato. Flow annealed importance sampling bootstrap. In The Eleventh International Conference on Learning Representations , 2023. URL https://openreview. net/forum?id=XCTVFJwS9LJ .

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 mcmc. In International Conference on Machine Learning , pages 8489-8510. PMLR, 2023.

Jascha Sohl-Dickstein, Eric Weiss, Niru Maheswaranathan, and Surya Ganguli. Deep unsupervised learning using nonequilibrium thermodynamics. In Francis Bach and David Blei, editors, Proceedings of the 32nd International Conference on Machine Learning , volume 37 of Proceedings of Machine Learning Research , pages 2256-2265, Lille, France, 07-09 Jul 2015. PMLR. URL https://proceedings.mlr.press/v37/sohl-dickstein15.html .

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.

Bernt Oksendal. Stochastic Differential Equations . Springer-Verlag Berlin Heidelberg, 6 edition, 2003.

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-232, October 2002.

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.

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 http://www.jstor.org/stable/3318418 .

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.

Pierre Del Moral, Arnaud Doucet, and Ajay Jasra. On adaptive resampling strategies for sequential Monte Carlo methods. Bernoulli , 18(1):252 - 278, 2012.

Lawrence C Evans. Partial differential equations , volume 19. American Mathematical Society, 2022.

Neil J Gordon, David J Salmond, and Adrian FM Smith. Novel approach to nonlinear/nongaussian bayesian state estimation. In IEE proceedings F (radar and signal processing) , volume 140, pages 107-113. IET, 1993.

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.

James Carpenter, Peter Clifford, and Paul Fearnhead. Improved particle filter for nonlinear problems. IEE Proceedings-Radar, Sonar and Navigation , 146(1):2-7, 1999.

Tiancheng Li, Miodrag Bolic, and Petar M Djuric. Resampling methods for particle filtering: classification, implementation, and strategies. IEEE Signal processing magazine , 32(3):70-86, 2015.

Michael S Albergo, Nicholas M Boffi, and Eric Vanden-Eijnden. Stochastic interpolants: A unifying framework for flows and diffusions. arXiv preprint arXiv:2303.08797 , 2023.

A Additional theoretical results

A.1 Jarzynski equality in continuous-time

Before proving Proposition 1, we give a continuous-time version of this result whose proof helps guide the intuition.

Proposition 2. Assume that the parameters θ are evolved according to some time-differentiable protocol θ ( t ) such that θ (0) = θ 0 . Given any α > 0 , let X t ∈ R d and A t ∈ R be the solutions of

where U θ ( x ) is the model energy and W t ∈ R d is a standard Wiener process. Then, for any t ≥ 0 ,

where the expectations on the right-hand side are over the law of the joint process ( X t , A t ) .

̸

The second equation in (18) can also be written in term of the free energy F θ = -log Z θ as F θ ( t ) = F θ 0 -log E [ e A t ] : this is Jarzynski's equality [34]. We stress that it is key to include the weights in (18) and, in particular, E [ ∂ θ U θ ( t ) ( X t )] = E θ ( t ) [ ∂ θ U θ ( t ) ] . This is because the PDF of X t alone lags behind the model PDF ρ θ ( t ) : the larger α , the smaller this lag, but it is always there if α < ∞ , 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 θ ( t ) by the gradient descent flow over the cross-entropy by solving (17) concurrently with

since by (18) the right hand side of (19) is precisely -∂ θ H ( ρ θ ( t ) , ρ ∗ ) = E θ ( t ) [ ∂ θ U θ ( t ) ] -E ∗ [ ∂ θ U θ ( t ) ] . Assuming that we know Z θ 0 we can also track the evolution of the cross-entropy (3) via

Proof. The joint PDF f ( t, x, a ) of the process ( X t , A t ) satisfying (17) solves the Fokker-Planck equation (FPE)

Let us derive an equation for

To this end, multiply (21) by e a , integrate the result over a ∈ ( -∞ , ∞ ) , and use integration by parts for the last term at the right-hand side to obtain:

By general results for the solutions of parabolic PDEs such as (23) (see [69], Chapter 7), we know that the solution to this equation is unique, and we can check by direct substitution that it is given by

Since by definition E [ e A t ] = ∫ R d ∫ ∞ -∞ e a f ( t, x, a ) dadx = ∫ R d ˆ ρ ( t, x ) dx this establishes the second equation in (18) . To establish the first notice that

This implies that

The need for Jarzynski's correction. Suppose that the walkers satisfy the Langevin equation (first equation in (17)):

where θ ( t ) is evolving according to some protocol. The probability density function ρ ( t, x ) of X t then satisfies the Fokker-Planck equation (compare (23))

̸

The solution to this equation is not available in closed form, and in particular ρ ( t, x ) = ρ θ ( t ) ( x ) -ρ ( t, x ) is only close to ρ θ ( t ) ( x ) if we let α → ∞ , so that the walkers X t 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 U θ ( t ) is complicated). As a result E [ ∂ θ U θ ( t ) ( X t )] = E θ ( t ) [ ∂ θ U θ ( t ) ] , implying the necessity to include the weights in the expectation (18).

̸

A.2 Proof of Proposition 1

The iteration rule for A k in (5) implies that

For k ∈ N 0 , let

be the transition probability density of the ULA update in (5), i.e. β k ( X k , X k +1 ) is the probability density of X k +1 conditionally on X k . By the definition of A k , we have

where in the second line we added and subtracted | X q -X q -1 | 2 / 4 h and used the definition of α k ( x, y ) given in (6). Since the joint probability density function of the path ( X 0 , X 1 , . . . , X k ) at any k ∈ N is

we deduce from (31) and (32) that, given an f : R d → R , we can express the expectation E [ f ( X k ) e A k ] as the integral

Since ∫ R d β k ( x, y ) dy = 1 for all k ∈ N 0 and all x ∈ R d , we can perform the integrals over x 0 , then x 1 , etc. in this expression to be left with

Setting f ( x ) = 1 in this expression gives

which implies the second equation in (7); setting f ( x k ) = ∂ θ U θ k ( x k ) in (34) gives

which can be combined with (35) to arrive at the first equation in (7). □

A.3 Mini-batched version of Algorithm 1.

In Algorithm 2 we present a mini-batched version of Algorithm 1, 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 back-propagation). This mini-batched 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 ∗ i } n i =1 ; energy model U θ ; optimizer step opt ( θ, D ) using θ and the empirical CE gradient D ; initial parameters θ 0 ; number of walkers N ∈ N 0 ; batch size N ′ ∈ N 0 with N ′ < N , set of walkers { X i 0 } N i =1 sampled from ρ θ 0 ; total duration K ∈ N ; ULA time step h ; set of positive constants { c k } k ∈ N .

2: A i 0 = 0 for i = 1 , . . . , N .

3: for k = 1 : K -1 do

4: p i k = exp( A i k ) / ∑ N j =1 exp( A j k ) ▷ normalized weights

5: ˜ D k = ∑ N i =1 p i k ∂ θ U θ k ( X i k ) -n -1 ∑ n j =1 ∂ θ U θ k ( x j ∗ ) ▷ empirical CE gradient

6: θ k +1 = opt ( θ k , ˜ D k ) ▷ optimization step

7: Randomly select a mini-batch { X j k } j ∈ B with # B = N ′ from the set of walkers { X i k } N i =1 8: for j ∈ B do √

11:

12:

end for for

j

̸∈

B

do

15:

end for

16: Resample the walkers and reset the weights if ESS k +1 < c k +1 , see (13). ▷ resampling step 17: end for

18: Outputs: Optimized energy U θ K ; set of weighted walkers { X i K , A i K } N i =1 sampling ρ θ K ; partition function estimate ˜ Z θ K = Z θ 0 N -1 ∑ N i =1 exp( A i K ) ; CE estimate log ˜ Z θ K + n -1 ∑ n j =1 U θ K ( x j ∗ )

A.4 Resampling Routines

Resampling schemes are necessary to tackle the decay of effective sample size (13). 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 to the review [73] for more details.

Given a set of normalized scalar weights { p i } N i =1 ∈ [0 , 1] with ∑ N i =1 p i = 1 associated to the N walkers, we define the cumulative sum

All three methods prescribe a way to choose a set { u n } N n =1 ∈ (0 , 1] used to perform the resampling in the following way: for every n, m ∈ { 1 , . . . , N } , the m -th particle is chosen during the n -th extraction if

Let us now specify how the set of u n is selected in the cases in study. We denote by U ( a, b ] the uniform probability distribution on the interval ( a, b ] .

Multinomial resampling. Sample u mult n ∼ U (0 , 1] , independently for every n ∈ { 1 , . . . , N } . This approach leads to a 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 , 1] into N sub-intervals, or strata, of size 1 /N ; then, sample u str n ∼ U (( n -1) /N,n/N ] independently for each n = { 1 , . . . , N } . This approach picks a single u n 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 sys 1 ∼ U (0 , 1 /N ] , and u sys n = u sys 1 +( n -1) /N for n > 1 . This method also reduces the number of possible resampled configurations.

Note that there are various modifications of these three methods (see the review [73]), all of them meeting the so-called unbiasedness condition, that is, the i -particle is expected to be sampled in average Np i times. However, these extensions do not lead to a critical lowering of the number of possible resampled configurations compared to systematic resampling. For this reason in our numerical experimentals we only tested the three methods above. Note also that, regardless of the resampling routine one uses, a fundamental role is played by the variance of the weights: if 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.

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 1 and 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

Interestingly, we can write down an equation that mimics the evolution of the PDF of the walkers in the CD algorithm, at least in the continuous-time limit: this equation reads

where the parameter ν > 0 controls the rate at which the walkers are reinitialized at the data points: the last term in (39) is a birth-death term that captures the effect of these reinitializations. The solution

Algorithm 4 Persistent contrastive divergence (PCD) algorithm

1: Inputs: data points Ω = { x ∗ i } n i =1 in R d ; energy model U θ ; optimizer step opt ( θ, D ) using θ and the empirical CE gradient D ; initial parameters θ 0 ; number of walkers N ∈ N 0 with N < n ; total duration K ∈ N ; ULA time step h .

2: X i 0 = SampleMultinomial (Ω) for i = 1 , . . . , N .

3: for k = 1 , . . . , K -1 do

4:

˜

D

N

k

=

N

-

1

i

=1

∑

∂

θ

5: θ k +1 = opt ( θ k , ˜ D k )

6: for i = 1 , ..., N do

8: end for

9: end for

10: Outputs: Optimized energy U θ K ; set of walkers { X i K } N i =1 .

̸

to this equation is not available in closed from (and ˇ ρ ( t, x ) = ρ θ ( t ) ( x ) in general), but in the limit of large ν (i.e. with very frequent reinitializations), we can show [55] that

As a result

The leading order term at the right hand side is precisely ν -1 times the gradient with respect to θ of the Fisher divergence

where ∆ denotes the Laplacian and we used ∫ R d ∇ U θ ( x ) · ∇ log ρ ∗ ( x ) ρ ∗ ( x ) dx = ∫ R d ∇ U θ ( x ) · ∇ ρ ∗ ( x ) dx = -∫ R d ∆ U θ ( x ) ρ ∗ ( x ) dx . 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 we give additional details about our numerical results. The codes are available at: https: //github.com/submissionx12/EBMs_Jarzynski .

B.1 Gaussian mixture distributions

Here we give some additional numerical results for the teacher-student model discussed in 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 1 ∗ = -10 and b 1 ∗ = 6 , and a α ∗ = b α ∗ = 0 for any α = 2 , . . . , d ; we also set z ∗ = -log(3) , corresponding to a mass p ∗ = 1 / (1 + e -z ∗ ) = 0 . 25 of the mode centered at a .

All the simulations are performed in d = 50 , with a time step of h = 0 . 1 for the ULA update. The number of data points is n = 10 4 . The setup of the teacher is the same for every simulation we display here and the optimization step is performed with full batch gradient descent with learning rate constant in time.

U

θ

k

(

X

i

k

)

-

-

n

1

n

i

=1

∑

∂

θ

U

θ

k

(

x

i

∗

)

▷

empirical gradient

▷

optimization step

0

50

100

150

200

250

Rescaled Iterations

0.0

0.1

0.2

0.3

0.4

0.5

p

Mass of the first mode

p

*

Algo 1

PCD

CD

0

50

100

150

200

250

Rescaled Iterations

10

8

6

4

2

0

2

4

6

a

,

b

Means

Algo 1

PCD

CD

a

*

b

*

0

50

100

150

200

250

Rescaled Iterations

10

5

10

4

10

3

10

2

10

1

10 0

10 1

KL

KL divergence

Algo 1

Eq. (21)

PCD

CD

10 0

10 1

10 2

Rescaled Iterations

10 4

6×10 3

7×10 3

8×10 3

9×10 3

ESS

9090.9

8333.3

7142.9

5555.6

5000.0

GMMfor different choices of c k . Results are shown in Figures 6 and 7. As initial conditions we select a 1 0 = -10 -1 , b 1 0 = 10 -1 , a α 0 ∼ 10 -2 N (0 , 1) and b α 0 ∼ 10 -2 N (0 , 1) for any α = 2 , . . . , d ; this perturbation around a = b = 0 is prescribed to avoid numerical degeneracy. For z , we fix z 0 = 0 . We run Algorithm 1, as well as the CD and PCD Algorithms 3 and 4 using N = 10 4 walkers for K = 8 · 10 3 iterations. We use a different learning rate for z and a, b , namely γ z = 0 . 125 and γ a,b = 0 . 2 γ z . Moreover, these values are multiplied by a factor 10 in CD. With regard to the resampling step in Algorithm 1, we choose a threshold c k = c which is fixed in time. We display the result for five possible values of this hyperparameter, namely c = [0 . 1 , 0 . 2 , 0 . 4 , 0 . 8 , 1 . 0] ; these are related to thresholds in the effective sampling size (ESS) via ESSthresh = N/ ( c +1) . With regard to Algorithm 1, we choose M = 1 and N ′ = N , that is 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 in figures.

Mode collapse in GMM for PCD. Results are shown in Figure 8. In the same setup as above we select as initial conditions a 1 0 = -10 -1 , b 1 0 = 10 -1 , a α 0 ∼ 10 -2 N (0 , 1) and b α 0 ∼ 10 -2 N (0 , 1) for any α ∈ [2 , d ] ; for z , we fix z 0 = 0 . The learning rate is chosen to be γ z = 5 for z and γ a,b = 0 . 2 γ z for the means. The time step of ULA is h = 0 . 2 . Since the objective is to show mode collapse in the PCD algorithm, we run just Algorithm 4 for K = 10 4 iterations and N = 10 4 walkers.

The need for resampling. In absence of resampling, we observe a 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.) as in Figure 9. The results are

0

2000

4000

6000

8000

10000

Iterations

0.0

0.2

0.4

0.6

0.8

1.0

p

Mass of the first mode

p

*

1

p

*

0

50

100

150

200

250

Rescaled Iterations

0.15

0.20

0.25

0.30

0.35

0.40

0.45

0.50

p

Mass of the first mode - No resampling

160

180

200

220

240

0.16

0.18

0.20

0.22

0.24

0.26

p

*

p

0

50

100

150

200

250

Rescaled Iterations

6000

7000

8000

9000

10000

Standard ESS

shown in Figure 10: looking at the upper panels, we see that stratified and systematic resampling are more stable than multinomial; moreover, if the resampling step is performed too infrequently (green lines), the method converges to a result where the target relative mass is off. Looking at the lower panels, we see that the minimum statistical error is obtained with systematic resampling; moreover, the behaviour of uncertainty appears to be more stable than stratified.

As a side note, systematic resampling requires just a single random number, contrarily to multinomial and stratified which need N , making the former also more efficient from a 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. GMMin high dimension are challenging for the standard CD and PCD algorithms given in 3 and 4. The experimental results of this section also confirm the theoretical analysis prsented in Appendix C below: CD is performing GD but on 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 8 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 6).

In this example, our Algorithm 1 outperforms these standard methods as it is an implementation GD on cross-entropy. In particular, the estimation of the relative mass with Algorithm 1 is more

0

50

100

150

200

250

Rescaled Iterations

0.25

0.30

0.35

0.40

0.45

0.50

p

Mass of the first mode - Multinomial resampling

100

150

200

250

2.4 × 10

1

2.5 × 10

1

2.6 × 10

1

Resampling interval

100

1000

10000

0

50

100

150

200

250

Rescaled Iterations

0.25

0.30

0.35

0.40

0.45

0.50

p

Mass of the first mode - Stratified resampling

100

150

200

250

2.4 × 10

1

2.5 × 10

1

2.6 × 10

1

Resampling interval

100

1000

10000

0

50

100

150

200

250

Rescaled Iterations

0.25

0.30

0.35

0.40

0.45

0.50

p

Mass of the first mode - Systematic resampling

100

150

200

250

2.4 × 10

1

2.5 × 10

1

2.6 × 10

1

Resampling interval

100

1000

10000

0

50

100

150

200

250

Rescaled Iterations

0.000

0.002

0.004

0.006

0.008

0.010

0.012

0.014

p

Standard error - Multinomial resampling

Resampling interval

100

1000

10000

0

50

100

150

200

250

Rescaled Iterations

0.0000

0.0005

0.0010

0.0015

0.0020

p

Standard error - Stratified resampling

Resampling interval

100

1000

10000

0

50

100

150

200

250

Rescaled Iterations

0.0000

0.0002

0.0004

0.0006

0.0008

0.0010

p

Standard error - Systematic resampling

Resampling interval

100

1000

10000

accurate than with the CD and PCD algorithms (see Figure 6). Moreover, the computation of KL divergence via Jarzynski weights is 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 7, 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 θ and of KL divergence, the overall behavior of Algorithm 1 is not dramatically influenced by the choice of c .

B.2 MNIST data set

In this section, we provide additional details about our numerical experiments on the MNIST dataset using Algorithm 2. First, as discussed in Section 4.2, we 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 11, where we display images along with their Jarzynski weights. Moreover, the right panel of Figure 11 indicates that resampling at the end of training using these weights helps improve sample quality. Examples of generated images are shown in Figure 12.

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 a 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 is 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 distributions 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 a Gaussian mixture model: a mixture of two one-dimensional Gaussian densities, with the same variance equal

3

3

3

13.59

11.01

11.01

10.25

10.11

10.09

10.02

9.51

2

2

2

2

9.31

8.77

8.73

8.63

8.04

7.92

7.86

7.8

2.91

2.82

2.82

2.82

2.75

2.74

2.68

2.67

3

2

1.43

1.42

1.41

1.4

1.39

1.38

1.37

1.36

3

2

0.75

0.74

0.73

0.73

0.73

0.72

0.71

0.71

Generatedsamplesduringtraining(w/oresampling)

Generated samples during training (w/resampling)

623226

6

2

2

22

2

h

2

3

332

2

2

2

2

5

)32？22

2622？2

2

2

2

66

2

2

2

223

23

E

3

√

3

3

3

2

3

3

3

2

6

2

3

2

3

2

6

3

2

6

3

2

2

2

2

2

2

2

3

a

2

3

2

G

5

2

6

3

2

2

2

3

3

3

6

∞N2

3

2

3

to 1, with means a, b ∈ R of the modes: our aim is to learn the probability masses of these modes. The general picture is given in Appendix C.1, with details presented afterwards in Appendix C.2, where we assume that the learning dynamics is that in (47), with no walkers and direct access to the distributions ρ z ( t ) and ρ ∗ , and Appendix C.3, where we quantify the stochastic fluctuations induced by empirical estimation.

EffectiveSampleSize

1.0

.2

0.8

e

0.6

S

Effective

e

0.4

0.2

EffectiveSampleSize

0.0

0

100

200

300

400

500

600

Epochs

C.1 Mode collapse in Gaussian Mixtures

C.1.1 Target distribution and model parameterization

As target density we take

Here z ∗ ∈ R parameterizes the mass of the second mode relative to the first and is the sole parameter of interest. The proportion of samples in both modes is indeed

Our point is that when both modes are separated by very low-density regions, learning z ∗ without weight correction leads to an incorrect estimation of the mode probabilities ('no-learning') or 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:

which will be enough for our needs. The more separated they are, the stronger our quantitative bounds will be.

The parameterization for our model potential U z is consistent with (43):

and the associated partition function and free energy are

The normalized probability density associated with U z ( x ) is thus ρ z ( x ) = e -U z ( x )+ F z .

C.1.2 Learning procedures

Gradient descent on the cross entropy leads to the following continuous-time dynamics:

For simplicity, we will start this ODE at z (0) = 0 . It corresponds to a proportion of 1 2 for both modes.

The gradient descent (47) is an ideal situation where the expectations E z ( t ) , E ∗ can be exactly analyzed. In practice however, the two terms of (47) are estimated; the second term using a finite

number of training data { x i ∗ } n i =1 , and the first one using a finite number of walkers { X i t } N i =1 , with associated weights { e A i t } N i =1 . For simplicity we set N = n and the empirical GD dynamics is thus

and the walkers evolve under the Langevin dynamic

for some fixed α > 0 . Now the nature of the algorithm varies depending on how the walkers are initialized and the Jarzynski weights are evolved.

The standard PCD algorithm sets X i 0 = x i ∗ , that is, the walkers are initialized at the data points, and the weights are not evolved, that is A i t = 0 at any time.

Alternatively, the walkers could be initialized at samples of the initial model: X i 0 ∼ ρ z (0) , with the weights not evolved. We refer to this algorithm as the umweighted procedure.

Our algorithm 2 corresponds to initializing the walkers at samples of the initial model, X i 0 ∼ ρ z (0) , and uses the Jarzynski rule (17) for the weights updates.

For simplicity we analyze the outcome of these algorithms in the continuous-time set-up, i.e. using (48). This is an 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 (48) can be seen as a perturbation of a simpler differential system whose qualitative behaviour fits with our numerical simulations. These systems depend on the initialization of the walkers and are thus prone to small stochastic fluctuations. We introduce:

ˆ q ∗ the proportion of training data { x i ∗ } that are close to mode b (a more precise definition is given in Appendix C), and we define ˆ z ∗ as satisfying ˆ q ∗ = e -ˆ z ∗ / (1 + e -ˆ z ∗ ) ;

ˆ q (0) the proportion of walkers at initialization that are close to b , and ˆ p (0) = 1 -ˆ q (0) .

Practically, ˆ q ∗ is a random variable centered at q ∗ and with fluctuations of order n -1 / 2 , and ˆ q (0) is centered at e -z (0) / (1 + e -z (0) ) = 1 2 with fluctuations of the same order. In the limit where n 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 in Appendix C lead to the following approximations:

In the model-initialized algorithm without Jarzynski correction ('unweighted'), (48) is a perturbation of

This system has no stable fixed point since the RHS no longer depends on z unw ( t ) , leading to a linear drift z unw ( t ) = (ˆ q (0) -ˆ q ∗ ) t and thus to a divergence of z unw ( t ) . Consequently, the mass of the second mode, q ( t ) = e -z unw ( t ) / (1 + e -z unw ( t ) ) , converges to 0 or 1. However, on longer time scales, this drift leads to a sudden transfer of all walkers in one the modes, then to a complete reversal of the drift of z unw which then diverges in the opposite direction, leading to a succession of alternating mode-collapses; see also Figure 8 and Remark C.4 below.

In the continuous-time version of the standard PCD algorithm, (48) is a perturbation of the same ODE as (50). However, in this context, since the initial data { X i 0 } and the training data { x i ∗ } are identical, we have ˆ q (0) = ˆ q ∗ , and

̸

The parameters do not evolve and the system is stuck at z pcd (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 a small batch of training data to initialize them; in this case we can still have ˆ q (0) = ˆ q ∗ , falling back to the first case above.

(a) Unweighted dynamics: (50)

(b) PCD: (51)

The continuous-version of Contrastive-Divergence is equivalent to the well-known ScoreMatching technique ([24]). In this context, (48) is a perturbation of

leading to the same 'no-learning' phenomenon.

In the model-initialized algorithm using the Jarzynski correction, (48) is a perturbation of

This system has a unique stable point ˜ z ∗ satisfying e -˜ z ∗ = ˆ p (0) e -ˆ z ∗ / ˆ q (0) , hence

Note that the second term at the right hand side is a small correction of order O ( n -1 / 2 ) since ˆ q (0) / ˆ p (0) = 1 + O ( n -1 / 2 ) .

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 2 exactly realizes (47). However, even for simple mixtures of Gaussian densities, the expectations in (43) do not have a simple closed-form that would allow for an exact solution. That being said, when a and b are sufficiently well-separated, the system can be seen as a perturbation of a simpler system whose solution can be analyzed.

First, we note that

The main idea of the approximations to come is that ∂ z U z ( x ) is almost zero when x is far from b (and in particular, close to a ), and is almost 1 if x is close to b . The next lemma quantifies this; from now on we will adopt the notation

Lemma 1. Under Assumption (A1) , for any v ∈ R ,

if x ∈ I a , then ∂ z U v ( x ) ≤ e -v -10 ;

if x ∈ I b , then | ∂ z U v ( x ) -1 | ≤ e v -10 .

Proof. From (55), we see that if x ∈ I a then | x -b | > 6 and consequently e -( x -b ) 2 / 2 ≤ e -18 . But the denominator of (55) is itself greater than e -( x -a ) 2 / 2 which is itself greater than e -4 2 / 2 = e -8 since | x -a | < 4 . This gives the first bound and the second is proved similarly.

We recall that if ξ ∼ N (0 , 1) , then P ( | ξ | > t ) ≤ e -t 2 / 2 /t , hence

Lemma 2. Let u, v ∈ R . Under Assumption (A1) , we have

where | ε | ≤ 0 . 0002 + 2 e -10 e | v | .

Proof. The integral is exactly given by

where ξ a , ξ b denote two Gaussian random variables with respective means a, b . By (57), ξ a and ξ b are respectively contained in I a = [ a -4 , a +4] and I b = [ 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 Lemma 1,

Consequently, by (57), the expectation is smaller than 0 . 0001 + e -v -10 . The second expectation in (59) is equal to

and by the same kind of analysis, the expectation here is smaller than 0 . 0001 + e v -10 . Gathering the bounds yields the result.

In particular, the right hand side of (47) can be approximated by e -z ( t ) 1+ e -z ( t ) -e -z ∗ 1+ e -z ∗ up to an error term smaller than 0 . 0004+2 e -10 ( e | z ( t ) | + e | z ∗ | ) . If z ( t ) , z ∗ are contained in a small interval [ -C, C ] with, say, C < 5 , this error term is uniformly small in time, and one might see (47) as a perturbation of the following system:

a system with only one fixed point at z ( t ) = z ∗ , the ground-truth solution.

Mode collapse in absence of Jarzynski corection

Now let us analyze in a similar fashion the dynamics without reweighting. Here, (47) is replaced by

where the process X t solves

The probability density function ρ ( t, x ) of X t satisfies a Fokker-Planck equation

which, in full generality, is hard to solve exactly, and thus exact expressions for the first term of (61) are intractable. However, depending on whether X 0 is close to a or b , the process X t can be well approximated by an Ornstein-Uhlenbeck process, hence ρ ( t, x ) can itself be approximated by a Gaussian mixture.

Proposition 3. Suppose that (A1) holds and that

Then one has D KL ( ρ (0) | ρ ( t )) ≤ δt where δ = 0 . 000025 + 100 e -20 e 2 C .

In other words, ρ ( t ) is approximately constant, up to reasonnable time scales t = O (1 /δ ) .

Proof. X 0 is drawn from ρ z (0) , a Gaussian mixture; the probability of it being sampled from a Gaussian with mean a is e -z (0) / (1 + e -z (0) ) = 1 / 2 . We will work conditionally on this event E a and we will note ρ a ( t ) the density of X 0 conditional on E a ; thus, ρ a (0) = N ( a, 1) . We set V ( x ) = 1 2 | x -a | 2 so that ∇ V ( x ) = ( x -a ) and we consider the following Ornstein-Uhlenbeck process: √

whose density will be denoted ˜ ρ a ( t ) . We use classical bounds on the divergence between ρ a ( t ) and ˜ ρ a ( t ) . For example, the bounds in [74, Lemma 2.20] directly apply and yield

Since Y t is nothing but an Ornstein-Uhlenbeck at equilibrium, Y t ∼ N ( a, 1) for all t ≥ 0 . The term inside the integral is a Gaussian expectation and will be shown to be small:

Consequently,

Under (A2), the overall bound remains smaller than t times 0 . 000025 + 100 e -20 e 2 C as requested, thus proving that ˜ ρ a ( t ) = N ( a, 1) and ρ a ( t ) are close with the same quantitative bound

Similarly, ρ b ( t ) , the density of X t conditional on X 0 being sampled from a Gaussian with mean b , is close to N ( b, 1) with the same quantitative bounds.

Overall, using the chain rule for KL divergences,

In other words, ρ ( t ) is close to a mixture of two Gaussians with modes centered at a, b , and the probability of belonging to the first mode is the probability of X 0 belonging to the first mode, that is, e -z (0) / (1 + e -z (0) ) = 1 / 2 .

Proof of (62) . We have

Using Lemma 1 and the fact that if x ∈ I a then | x -b | < 16 and | x -a | < 4 , we get |∇ U v ( x ) -( x -a ) | ≤ 20 ε with ε ≤ e -v -10 . Consequently,

As a consequence of the Proposition 3, the first term of (61) can be approximated by E z (0) [ ∂ z U z ( t ) ( X t )] = E z (0) [ ∂ z U z ( t ) ] , which in turn can be approximated by e -z (0) / (1 + e -z (0) ) thanks to (58). Overall, (61) is therefore a perturbation of the system

Since the right hand side no longer depends on z ( t ) , this system leads to a constant drift of z ( t ) , that is z ( t ) = γt , leading to mode collapse since (1 + e -z ( t ) ) -1 ≈ (1 + e γt ) -1 goes to either 0 or 1.

Mode collapse for Contrastive-Divergence and Score-Matching

The continuous-time limit of Contrastive Divergence (Algorithm 3) is equivalent to Score-Matching minimization ([28]). The objective function becomes the Stein score,

which is in theory intractable due to the presence of the unknown parameter z ∗ ; a well-known computation from [24] shows that the gradient ∂ z SM( z ) can be estimated using the training samples without resorting to z ∗ . Note that in this context, there are no 'walkers'.

Now, the dynamics (47) is replaced by

where here again ξ x ∼ N ( x, 1) . From (62) and the triangle inequality, we have

and the same approximation holds for the second part in (65). Overall, we get that for any reasonnable z , SM( z ) ≈ 0 : that is, every z minimizes the score. A similar analysis leads to ∂ z SM( z ) ≈ 0 . Consequently, ˙ z ( 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 E z ( t ) , E ∗ can be exactly analyzed. In practice, the two terms of (47) are estimated; the second term using a finite number of training data { x i ∗ } , and the first one using a finite number of walkers { X i t } , with associated weights e A i t which are either evolved using the Jarzynski rule, or simply set to 1 in the PCD algorithm. Our goal in 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):

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

Let ˆ n a ∗ be the number of training data in I a and ˆ p ∗ = ˆ n a ∗ /n their proportion, and similarly ˆ q ∗ the proportion in I b , and r = 1 -ˆ p ∗ -ˆ q ∗ . By elementary concentration results, the remainder 1 -ˆ p ∗ -ˆ q ∗ can be neglected: with high probability, it is smaller than, eg, 0 . 0001 . We will note ˆ z ∗ the parameter satisfying ˆ q ∗ = e -ˆ z ∗ 1+ e -ˆ z ∗ . Using Lemma 1, the second term in (66) is approximated by ˆ q ∗ . Now let us turn to the first term in (66). Still using Lemma 1, we see that the first term in (66) is well approximated by

The second equation in (67) entails e A i t = exp ( -∫ t 0 ∂ z U z ( s ) ( X i s ) ˙ z ( s ) ds ) . Now let us use Lemma 1: if X i s ∈ I a , then ∂ z U z ( s ) ( X i s ) ≈ 0 . Conversely, if X i s ∈ I b , then ∂ z U z ( s ) ( X i s ) ≈ 1 . Moreover, Proposition 3 and its proof essentially show that if X i 0 belongs to the first well (close to a ), then with high probability so does X a s for every s , and in particular X a s ∈ I a with high probability for every s . Consequently,

As already explained, the dynamics (67) leaves approximately constant the number of walkers in both modes; consequently, the proportion ˆ q ( t ) of walkers X i t in I b remains well approximated by the initial proportion, which is ˆ q (0) , and we obtain that (69) is well approximated by

where we noted ˆ p (0) = 1 -ˆ q (0) . Note that since z (0) = 0 , with high probability ˆ p (0) and ˆ q (0) are close to 1 / 2 . The random variable ˆ p (0) / ˆ q (0) is thus close to 1.

Overall, we obtain that the system (66) is a perturbation of the following system:

This system has a unique stable point equal to

̸

Remark. If the algorithm had been started at z (0) = 0 , one could check that the stable point would become ˆ z ∗ +log(ˆ q (0) / ˆ p (0)) + z (0) .

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 A i t = e 0 = 1 ('unweighted' algorithm), then (66) becomes

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

Consequently, (75) is a perturbation of the system

which no longer depends on t and thus leads to z ( t ) = ˆ γt and to mode collapse.

The PCD algorithm leads to no-learning

In the preceding paragraph, at initialization, the walkers X i 0 are distributed according to the initial model ρ z (0) . In the PCD algorithm 4, the walkers are instead initialized directly at the training data { x i ∗ } n i =1 . However, in this case, the analysis of the preceding paragraph remains essentially the same, with a single difference: the initial proportion of walkers that are close to a , noted ˆ q (0) , is now exactly ˆ q ∗ . Thus, (66) becomes a perturbation of

The parameters remain constantly equal to its initial value z (0) , i.e. there is no learning.

The CD algorithm leads to no-learning

The Continuous-time Contrastive-Divergence algorithm is minimizing the Stein score and is equivalent to Score-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 is exactly the same as above:

If x i ∗ ∈ I a , then as explained in the proof of Lemma 62, ∇ U s ( x i ∗ ) ≈ ( x i ∗ -a ) for every s , hence ∂ z ∇ U z ( x i ∗ ) × ( ∇ U z ( x i ∗ ) -∇ U z ∗ ( x i ∗ )) ≈ 0 . The same holds for x i ∗ ∈ I b , leading to (66) being a perturbation of the system

C.4 On mode-collapse oscillations

Most of the approximations performed earlier rely on (A2), that is, the learned parameter z ( t ) remains in a compact set.

However, in the unweighted algorithm, this is no longer the case at large time scales, since z ( t ) diverges from (76); in particular, the approximations from Lemma 1 become meaningless. In fact, Proposition 3 is no longer relevant. The core of Proposition 3 rests upon the fact that if a walker X i t is close to a , then its dynamics (67) is close to an Ornstein-Uhlenbeck process since ∇ U z ( t ) ( X i t ) ≈ ( X i t -a ) . This fails when z ( t ) has a large absolute value. Let us suppose for instance that z ( t ) is very small (negative), so that e -z ( t ) is very large, | z ( t ) | ≫| x -b | 2 . In (63), the first term of the numerator is dominated by the second term. Overall, we get

and this is valid for all X i t . Consequently, all the walkers now undergo an Ornstein-Uhlenbeck process centered at b and in particular, the walkers that are close to a are exponentially fast transferred to the region close to b . At this point, the first term in (75) becomes close to 1, leading to the approximated system ˙ z ( t ) = 1 -ˆ q ∗ : z ( t ) oscillates back to + ∞ , until the same phenomenon happens again and all the walkers transfer to the region close to a .

This leads to an oscillating behavior that can be observed on longer time scales (see Figure 8 for example). We do not think that this phenomenon is relevant to real-world situations since most learning algorithms are trained for a limited time period.

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS:\nI 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.\n\nWEAKNESSES:\nHowever, 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.\n\nQUESTIONS:\n1. 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.\n\nLIMITATIONS:\nThe authors have adequately addressed the limitations.", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-02"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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\n\nSTRENGTHS:\nThe 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).\n\nWEAKNESSES:\nFrom 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.\n\nQUESTIONS:\n- 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.\n\nLIMITATIONS:\nThe 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.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-05"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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.\n\nSTRENGTHS:\n1) 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\nWEAKNESSES:\n1) 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\nQUESTIONS:\n1) 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\nLIMITATIONS:\n--", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 2, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-06"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe paper reports an enhcements in the eficiency of the training of EBMs using Jarzinsky equality.\n\nSTRENGTHS:\nThe paper has the potential to help understand the training of EBM better through a thermodynamics based approach. It is well written and elaborated.\n\nWEAKNESSES:\nThe 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.\n\nQUESTIONS:\nCan 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\nLIMITATIONS:\nI 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.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-07"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nOne 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.\n\nSTRENGTHS:\n- 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.\n\nWEAKNESSES:\n- 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.\n\nQUESTIONS:\nI have no questions for the authors.\n\nLIMITATIONS:\nBesides 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.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 5, "SOUNDNESS_CORRECTNESS": 2, "DATE": "2023-07-20"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS:\n* 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.\n\nWEAKNESSES:\n* 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.\n\nQUESTIONS:\n* 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?\n\nLIMITATIONS:\nAuthors did not include explicit sections discussing limitations and broader impacts.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-20"}