Efficient Training of Energy-Based Models Using Jarzynski Equality

Summary

This paper uses Jarzinsky Equality to analyze the training of Energy Based Model (EBM) using cross entropy loss. The authors highlight the problem of directly approximating the gradient using the unweighted mean of nonequilibrium samples and propose a modified algorithm based on reweighting.

Strengths

I like this idea of using Jarzinsky Equality to analyze the current training problem existing in EBM trained with cross entropy loss. I think the analysis is reasonable and hit the core problem of EBM training.

Weaknesses

However, I remain unconvinced by the proposed algorithm, particularly regarding its potential for practical scalability. I have the following concerns: 1. 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. 2. 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. 3. 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.

Questions

1. For the resampling step, how to assign new value to each walker $X_{kr}^{i}$ that needs to be resampled? Do you just resample $X_{kr}^{i}$ from $\rho_{\theta0}$? Also why can you reset $A{kr}^{i}=0$? For my understanding, according to eq5 and eq14, you should keep track of the whole chain of $A_t$ and $X_t$ through the training and sampling process, or the results might be incorrect/ 2. 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. 3. 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: 1. A theory of generative convnet. 2. Flow contrastive estimation of energy-based model. 3. Your classifier is secretly an energy based model and you should treat it like one 4. Vaebm: A symbiosis between variational autoencoders and energy-based model. 5. Generalized energy based mode. 6. A tale of two flows: Cooperative learning of langevin flow and normalizing flow toward energy-based model. 7. Guiding energy-based models via contrastive latent variable.

Rating

4: Borderline reject: Technically solid paper where reasons to reject, e.g., limited evaluation, outweigh reasons to accept, e.g., good evaluation. Please use sparingly.

Confidence

4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.

Soundness

3 good

Presentation

3 good

Contribution

3 good

Limitations

The authors have adequately addressed the limitations.

Summary

The paper under consideration proposes a modification of the standard EBM training procedure inspired by Jarzynski equality from thermodynamics. The idea is as follows. It is known that EBM training requires sampling from the distribution for which we know only its energy function given by, e.g., neural network $E_{\theta}$. Typically, this is done by ULA algorithm or its analogues. Given current NN parameters $\theta$, we try to sample from $p_{\theta} \sim \exp( - E_{\theta})$ using the corresponding MCMC algorithm and only after that we update the parameters $\theta$. The authors take a step further and propose to change the paradigm. They consider a specific discrete procedure which resembles Annealed Importance Sampling MCMC algorithm. This procedure evolves both particles and model (NN) parameters so that the parameters turn out to follow gradient flow with respect to KL divergence. The authors validate and show the advantages of ther idea on moderate-dimensional Gaussian Mixture setup and MNIST dataset. **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. [1] Du et. al., Improved Contrastive Divergence Training of Energy-Based Models, http://proceedings.mlr.press/v139/du21b/du21b.pdf

Strengths

The paper seems to present a novel and theoretically sound idea. To be more precise, I would like to stress out the following points: - 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. - 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. - The authors demonstrate several more-or-less convincing setups that demonstrate the superiority of their method compared to standard CD and PSD approaches. Additionally, for simple 1D setup they provide the theoretical arguments which supports their thesis regarding this supremacy (to be honest, I didn’t get into details of the analysis in Appendix C, but it seems correct).

Weaknesses

From my point of view, the main weaknesses of the paper under consideration are as follows: - 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. - 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. - 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.

Questions

- Line 163: Are there theoretical guarantees that the integral (which defines $\hat{\rho}(t, x)$) converges? - 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. - Which parameters of Algorithm 1 were used in the MNIST experiment? - 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. Misprints - Line 170: In the second equation of the inline formula, one should add the integration w.r.t. $x$. - Line 180: “that that” -> “that”. - Algorithm 1, line 2: it seems $A_0^i$ should be initialized as zeros. - Line 433: “given and” -> “given”. - Formula (51), denominator: $U_z(x)$ -> $\exp(- U_z(x))$. - Figure 10 (about MNIST dataset) label refers to formula (25) (which is about MoG experiment). There should be another formula.

Rating

7: Accept: Technically solid paper, with high impact on at least one sub-area, or moderate-to-high impact on more than one areas, with good-to-excellent evaluation, resources, reproducibility, and no unaddressed ethical considerations.

Confidence

4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.

Soundness

3 good

Presentation

3 good

Contribution

3 good

Limitations

The potential negative societal impact of the work are not addressed. However, it seems to be the save as for thousands of papers which deals with generative modelling. So I have no special concerns.

Summary

The authors consider EBMs family of algorithms optimized with a cross-entropy objective. They proposed a modification of the unadjusted Langevin algorithm (ULA) to remove bias from the optimization procedure, basically introducing correction coefficients for UCL gradient estimation. Their method is based on Jarzynski equality and is theoretically justified. Moreover, the paper confirms empirically the method is less biased as it can recover all "modes" and proportions between them from data correctly in contrast to other baselines optimized with the cross-entropy loss.

Strengths

1) The method is theoretically justified and novel to the best of my knowledge. 2) The method can be used to find a normalizing constant and estimate density. 3) 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. 4) The method is easy to implement and extra (compared to ULA) computational cost is low. 5) The paper is well-written and easy to follow.

Weaknesses

1) The empirical evaluation is weak. The paper only considers two toy datasets. It would be great to explore the proposed method on more challenging datasets. However, the authors noticed it themself in the conclusion. 2) 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. 3) It is not clear how the new method compares in terms of sample quality to baselines (there are no metrics).

Questions

1) The method can be used for density estimation. Is the proposed method beneficial compared to other density estimation methods? Is it worth evaluating it against other methods on the density estimation task? 2) 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?

Rating

7: Accept: Technically solid paper, with high impact on at least one sub-area, or moderate-to-high impact on more than one areas, with good-to-excellent evaluation, resources, reproducibility, and no unaddressed ethical considerations.

Confidence

2: You are willing to defend your assessment, but it is quite likely that you did not understand the central parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.

Soundness

3 good

Presentation

3 good

Contribution

3 good

Limitations

--

Summary

The paper reports an enhcements in the eficiency of the training of EBMs using Jarzinsky equality.

Strengths

The paper has the potential to help understand the training of EBM better through a thermodynamics based approach. It is well written and elaborated.

Weaknesses

The paper shows a more efficient training of EBM. However, it does not explain why such efficient is obtained by performing the modification using the JE. I think the authors should dedicate a section to elaborate from thermodynamics perspective that why such improvements is achieved.

Questions

Can you please explain what would be the thermodynamics systems, States, temperature, thermal reservoir in the context of the training of a EBM? In 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? In 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? Can you please explain in the context of the thermodynamics, for example using free energy and external work, on why such efficiency is obtained?

Rating

6: Weak Accept: Technically solid, moderate-to-high impact paper, with no major concerns with respect to evaluation, resources, reproducibility, ethical considerations.

Confidence

3: You are fairly confident in your assessment. It is possible that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.

Soundness

3 good

Presentation

3 good

Contribution

3 good

Limitations

I think a more physical understanding of why the performed modification enhances the training will have be great for gaining an understanding of the physics of the training.

Summary

One of the main challenges in the training of Energy-Based Models (EBMs) is the design of a tractable objective that allows for low-variance gradient estimation. Indeed, the maximum likelihood objective requires sampling from the current density model, which usually is computationally expensive and biased due to poor mixing properties. Many alternatives have been proposed to alleviate the issues of maximum likelihood objective: contrastive divergence, persistent contrastive divergence, non-convergent MCMC training, etc. The 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 $$\mathbb{E}_{x\sim q_\theta}\partial_\theta U_\theta(x),$$ where $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). The 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.

Strengths

- The proposed algorithm and theoretical developments that justify the algorithm are novel and are definitely of interest to the generative modeling community. - The paper is clearly written, and all the results are well-presented. **EDIT:** During the discussion, the authors provided experimental results for CIFAR-10. However, I still have the following concerns: - The baseline in the empirical study for CIFAR-10 demonstrates much worse performance than in the prior work [2]. - 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. - The biasedness of the estimator has to be mentioned explicitly next to its introduction since it's an important property (downside) for optimization. Despite these concerns, I recommend the paper for acceptance. [2] Du, Y., Li, S., Tenenbaum, J., & Mordatch, I. (2020). Improved contrastive divergence training of energy-based models. arXiv preprint arXiv:2012.01316.

Weaknesses

- The empirical evaluation of the proposed algorithm is insufficient to demonstrate the practical benefits of the proposed algorithm. Indeed, the considered experiments are rather toy examples compared to the relevant literature [1,2,3]. Since Algorithm 1 is the main development of the paper, this is a major flaw of the current submission. - 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. - The following statements in the introduction require references to the corresponding works with justifications. - “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. - “methods in this class are provably unable to learn the relative weights of modes when they are separated by low-density regions” - “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” - “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. - “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). [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*. [2] Du, Y., & Mordatch, I. (2019). Implicit generation and modeling with energy-based models. *Advances in Neural Information Processing Systems*, *32*. [3] Du, Y., Li, S., Tenenbaum, J., & Mordatch, I. (2020). Improved contrastive divergence training of energy-based models. *arXiv preprint arXiv:2012.01316*. [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*. **EDIT:** During the discussion period, the authors provided experimental results for CIFAR-10. However, I still have the following concerns: - The baseline in the empirical study for CIFAR-10 demonstrates much worse performance than in the prior work [3]. - 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. - The biasedness of the estimator has to be mentioned explicitly next to its introduction since it's an important property (downside) for optimization. Despite these concerns, I recommend the paper for acceptance. [3] Du, Y., Li, S., Tenenbaum, J., & Mordatch, I. (2020). Improved contrastive divergence training of energy-based models. arXiv preprint arXiv:2012.01316.

Questions

I have no questions for the authors.

Rating

6: Weak Accept: Technically solid, moderate-to-high impact paper, with no major concerns with respect to evaluation, resources, reproducibility, ethical considerations.

Confidence

5: You are absolutely certain about your assessment. You are very familiar with the related work and checked the math/other details carefully.

Soundness

2 fair

Presentation

3 good

Contribution

2 fair

Limitations

Besides poor empirical evaluation, which does not allow for defining the limits of the proposed method, two following points should be discussed. - 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. - The derived estimate is biased. It is a minor point but it is never mentioned in the paper.

Summary

This paper examines EBM learning using tools from the Jarzynski and Sequential Monte Carlo. The method is meant to address limitations of current EBM methods, which often rely on simulated Markov Chains to calculate approximate gradient of the log normalizer. Since the chains in practice do not converge to the model distribution, the theoretical validity of standard procedures is questionable and practical issues with correct multi-modal modeling can arise. The central contribution of this work is to introduce a new loss function for EBM learning. Using the Jarzynski equality, the gradient of the log normalizer is shown to be equivalent to the ratio of the expectation of weighted gradients of the EBM evaluated at ULA samples and the expectation of the weight terms. Unlike standard EBM learning which requires an unjustified assumption of MCMC convergence for validity, this proposed loss function is exact even for ULA samples which have not reached the model equilibrium distribution. To improve sample efficiency, weights are periodically resampled during practical implementation. Experiments on Gaussian mixture models and MNIST digits show the proposed method is better able to learn correct relative weightings of modes compared to standard EBM models.

Strengths

* The theoretical development of the proposed loss function is the central strength of the paper. In particular, developing a loss function that is valid even for non-convergent ULA samples is a major contribution that has the potential to lead to more principled and effective EBM learning. * The paper is clearly written and easy to understand. * The proposed method does not significantly increase computation costs beyond standard EBM methods. * A thorough appendix with full proofs of theoretical results and a significant amount of supporting information and additional experimental results is included. * Experimental results give convincing evidence that the proposed method is able to model relative modes weightings more effectively than standard EBM learning.

Weaknesses

* The primary weakness is the scope of the experimental results. The scale of datasets studied is smaller than typical EBM works, which usually examine at least CIFAR-10 as a baseline. In my experience, many methods which are effective for GMMs and MNIST might have poor performance on CIFAR-10. Since the goal of the proposed method is to improve multi-modal modeling, it is difficult to gauge the potential of the proposed method when experiments only examine datasets with very few modes. FID scores on CIFAR-10 would provide a clearer comparison of the generative abilities of the proposed method compared to related methods. * 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.

Questions

* Can the performance of the proposed method on larger scale datasets (e.g. CIFAR-10) be investigated? * Can the authors discuss the validity of the weight sampling step in Section 3.3?

Rating

5: Borderline accept: Technically solid paper where reasons to accept outweigh reasons to reject, e.g., limited evaluation. Please use sparingly.

Confidence

4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.

Soundness

3 good

Presentation

3 good

Contribution

2 fair

Limitations

Authors did not include explicit sections discussing limitations and broader impacts.

Answer

I thank the authors for their answers and clarifications. Indeed, Algorithm 4 does not require updating weights for the whole dataset at each minibatch step. At the same time, I still think that for granting better score, more convincing (large-scale) experiments which demonstrates the generative ability of the proposed method are required. So I consider leaving my score unchanged.

Thanks for the author responses. I will keep my score.

Thanks to the authors for their responses to myself and other reviewers. I greatly appreciate the efforts of this work to develop an EBM learning framework that is less reliant on assumptions which are demonstrably not true in practice. This work has ideas which could be a step towards learning proper unnormalized densities. As the authors acknowledge, this method relies on approximations as well, namely the use of a finite number of walkers to approximate certain expectations, and the accuracy of these approximations in the finite sample case remain opaque. While one can speculate about the relative accuracy of the approximations for current EBM methods and the proposed method for high-dimensional data, rigorous comparisons cannot be made. Empirical results are the usual method comparison in this case, but this paper does not perform experiments at a scale that is sufficient for a fair comparison. I agree with Reviewer nHa5 that there are public resources for EBM learning which are very accessible with relatively small compute cost (especially given free/low-cost cloud GPU platforms). I do not view it as necessary to achieve state-of-the-art results on CIFAR-10 with the current method, but at least viable FID scores are needed to validate the claims in this paper. Overall, I choose to keep my score.

Thanks to the Authors for addressing my comments. I keep my score.

Thanks to the Authors for addressing my comments and questions. I keep my score.

New experiment on the CIFAR-10 image dataset

**CIFAR-10 Experiments** Following the Reviewers suggestion, we performed an empirical evaluation of our method on the full CIFAR-10 (32 $\times$ 32) image dataset. We used the setup of Ref. [1], with the same neural architecture with $n_f = 64$ features, and compared the results obtained with our approach with mini-batching (Algorithm 4) to those obtained with PCD and PCD with data augumentation of Ref. [2] (which consists of a combination of color distortion, horizontal flip, rescaling, and Gaussian blur augumentations, to help the mixing of the MCMC sampling and stabilize training). The code used to performed these new experiments is available (along with some generated images) 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 $\eta = 10^{-4}$ and inject a Gaussian noise of standard deviation $\sigma = 3\times 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 22 hours with the PCD algorithm (w/ and w/o data augumentation) and about 23 hours with our method. The results in terms of FID and IS are: | Method | FID | IS | ------- |------ |------ |PCD |45.63 | 5.62 |PCD with data augmentation |42.27| 6.34 |Our method (Algorithm 4)| 38.49 | 6.76 We will add these new results to the revised version of our paper. They indicate that our method can achieve slightly better performance than the PCD algorithms w/ and w/o data augumentation 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. ***References:*** [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. [2] Du, Y., Li, S., Tenenbaum, J., & Mordatch, I. (2020). Improved contrastive divergence training of energy-based models. arXiv preprint arXiv:2012.01316.