Learning Energy-based Model via Dual-MCMC Teaching

Summary

This paper proposes a novel method for training and inference of energy-based model (EBM). The author claim that although the use of generator as an initializer model may improve MCMC sampling, training unbiased generator is an open problem . Thus they propose a joint learning framework in which generator is learned by MLE via posterior sampling. While EBMs have a number of appealing properties they are difficult to inference. Therefore learning an unbiased generator for faster MCMC sampling from the EBM distribution is a hot topic.

Strengths

The paper is clearly written and well presented. The potential is impact of the proposed dual-MCMC teaching is considerable. The problem arised in the paper is actual and mostly uncovered by the previous research. The paper provides a number of novel ideas which could encorage the future research in this direction. The experiments are very clear and solid including fair comparison with concurent works. Remarkably, the numerical results for the proposed method are spectacular.

Weaknesses

1. The section 2 is clear-written and self-contained but it is large and has little connection to the methodology proposed in the paper. As the authors reforlulate each component described in the section 2 it might be helpful to better explain the benefits of the choice for proposed method. 2. The objective proposed for learning EBM includes maximization of Kulback-Leibler divergence term (equation 7) which is stated to perform as a self-critic for EBM. The grounding beneath this “​​self-adversarial learning” is not very confident. Either literature referenced by the authors lack the verification or some formal reasoning behind this objective. Compared to the objectives for generator and inference model this one doesn’t pose an upper bound for MLE objective. 3. In the Figure 2 in the main text you provide an energy profile after performing $30$ steps (as during training) of MCMC revision on $x$. FID improves from $5.94$ on $x_0$ to $5.15$ on $x_{30}$, still the energy profile is not convergent. It would be more convincing if you provide a result for long-run MCMC (until energy profile convergence) revision for learned EBM to verify your EBM. 4. Please include the training details, including optimization hyperparameters, step sizes for Langevin dynamics, etc. Typos: 1. In the line 156 there is a typo for the marginal distribution of x, but inside the brackets you use a proper notation. 2. In the line 214 “second” is used instead of “first”.

Questions

1. It was noticed [1] that for EBM distribution with density $\pi(x)$ and generator $p_{\theta} = p_{\theta}(x|z)p(z)$ sampling latent variable from $p(z)\pi(g_{\theta}(z))$ and pushing it by $g_{\theta}$ to the ambient space is equivalent to sampling from $\pi(x)$. Why do you propose to perform EBM inference in the ambient space instead of latent space? As traversing high-dimensional space with MCMC sampling is a hard problem. Please elaborate on why don’t you use an advantage of sampling from EBM in the latent space of the generator model. 2. The paper massively adopts the method proposed in [2], nevertheless considering results on CIFAR-10 dataset FID drops from $30.10$ for [2] to $9.26$ for the current paper. What is a major reason for such a dramatic improvement? In the Appendix you provide a quality for a "simple network" which seems to share architecture with paper [2]. It would be helpful for understanding if you provide a comparison of model complexity and/or time, GPU consumption in addition to image quality comparison. [1] Che, Tong, et al. "Your gan is secretly an energy-based model and you should use discriminator driven latent sampling." Advances in Neural Information Processing Systems 33 (2020): 12275-12287. [2] Han, Tian, et al. "Divergence triangle for joint training of generator model, energy-based model, and inferential model." Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition. 2019.

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

3 good

Presentation

4 excellent

Contribution

3 good

Limitations

The limitations of the method are obscured. The authors only mention that “this work may share the limitation with other MCMC-based methods in terms of the computational cost” which is a bit vague.

Summary

This paper investigates the problem of Maximum Likelihood (ML) estimation for Energy-Based Models (EBM). Building on previous research, the authors propose a novel approach that involves learning a surrogate generative model to initialize the costly Langevin steps. The authors chose a latent-based generative model thus requiring additional Markov Chain Monte Carlo (MCMC) sampling for estimation. This new MCMC procedure is amortized using a simpler third generative model, referred to as the inference model. This work provides a procedure to learn the three models jointly. Through extensive experimental evaluations, the authors demonstrate the effectiveness of their proposed framework. They showcase the benefits of incorporating the additional models and highlight the superiority of their combined algorithm over other existing methods. The results substantiate the value of joint learning, showcasing improved accuracy.

Strengths

- Experiments are well detailled. - Experiments on Out-of-Distribution Detection (OOD) in App 1.2 allow to evaluate the EBM beyond generating samples.

Weaknesses

- The motivation to learn the generator with observed samples is unclear. - The improved accuracy of this method is not balanced with the computational and memory costs generated. - The quality of the density estimation task is barely considered.

Questions

- One of the main justifications of your learning procedure is to learn the generator using training data. Can you explain why not doing this (like in cooperative learning [1,2]) would introduce a bias ? This claim was made multiple times in the main paper. - As a justification of your revisited densities (or in L129), you say that Langevin steps could transform something unimodal into something multimodal which is not the case (see for instance [3]) as shown in Fig 2 of Sec 5.2. Is there another justification for those Langevin steps ? - The ablation studies (Sec 5.2 and 5.3) highlight the usefulness of the inference model with surprisingly low computational overhead. What about the memory overhead ? [1] 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. [2] Jianwen Xie, Zilong Zheng, and Ping Li. Learning energy-based model with variational auto-encoder as amortized sampler. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 35, pages 10441–10451, 2021 [3] Véronique Gayrard, Anton Bovier, Markus Klein, Metastability in reversible diffusion processes II: precise asymptotics for small eigenvalues. J. Eur. Math. Soc. 7 (2005), no. 1, pp. 69–99

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

3 good

Limitations

Not applicable.

Summary

This work presents a new approach for training energy-based models (EBMs). These are commonly trained by approximating the maximum-likelihood estimate of the model parameters. This is done by minimising a certain Kullback--Leibler divergence via stochastic-gradient descent. To estimate the necessarily gradient, one must approximate an expectation w.r.t. the law of the data under the energy-based model. This is typically done via an MCMC algorithm. The present work attempts to speed up this procedure by training an additional surrogate model (the "generator" model) from which the MCMC chain can be initialised (so that fewer MCMC iterations are needed). This generator model and the EBM are trained as part of a larger scheme which additionally trains an inference model (i.e. the posterior distribution of the latent variables given the observed data). ----------------------- EDIT (2023-09-04): I have now read all the reviews and rebuttals. All my questions have been addressed. I continue to be happy for this paper to be published.

Strengths

1. The method appears to be sound and novel. 2. The paper is mostly well written.

Weaknesses

Main comments: My expertise with training EBMs in this fashion is too limited to offer deeper insights. However: 1. Some of the objectives used for training the three models appear a little ad-hoc and their motivation/justification sometimes leans on convergence of the MCMC chains (which seems unlikely to occur in practice given the small number of MCMC iterations used). At the very least, some added justification/derivation of these objectives would help the reader gain more intuition for the method. 2. I'd ask the authors to confirm that the numerical comparisons are fair in the sense that all algorithms have roughly the same computational cost (see question below). 3. I think it would help to have some (brief/high-level) pseudo code in the main paper. Minor comments/typos: L114: grammar in "such Langevin process" L180: are the KL divergences equivalent to the marginal version or is the /minimisation/ of these KL divergences equivalent to minimising the marginal divergences? L152--160: I understand what the authors mean with the "transition-kernel" notation. But I think the notation/explanation in this paragraph needs to be improved. For example, in Line 156, the stated expression is referred to as "marginal distribution on $x$" but it clearly represents distribution on joint the space of $(x, z)$. I would state the formal definitions (i.e. those found in Lines 158 and 160) right from the start to enhance clarity. It might also be worth reminding the reader that $p_\theta(x)$ is the marginal of $p(z)p_\theta(x|z)$. L256--257: the MCMC sampling -> MCMC sampling - there are sometimes missing spaces between "Eq" and the equation number and between "Sec." and the section number. - punctuation is needed even if a sentence ends in a (displayed) equation - the font size in most of the figures is too small (especially in Figure 3) and axis labels are sometimes missing - it may just be my lack of familiarity with the literature in this area but is it clear to the reader what "explaining away inference" means in this specific context? If not, maybe it'd be worth adding a brief explanation. - table captions normally go /above/ tables (unless the NeurIPS style guide says otherwise)

Questions

In Table 1 (and also in the numerical results from Sections 5.2 and 5.3), is the computational cost of all algorithms comparable?

Rating

6: Weak Accept: Technically solid, moderate-to-high impact paper, with no major concerns with respect to evaluation, resources, reproducibility, 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

2 fair

Limitations

NA

Summary

This work investigates learning EBMs for image data using auxiliary generator and inference models. The model generates samples by drawing a latent normal vector, passing it through the generator, and refining the generator sample using MCMC with the EBM. An inference network, which predicts latent vectors from images, is used to assist with training the EBM and generator. The key innovation is to update both generator samples using image space MCMC and inferred latent vectors for data using latent space MCMC. This allows the generator to be efficiently trained using both revised MCMC samples, as in Cooperative Learning, and revised latent codes for data, which is unique to this work. Training the generator to match both EBM and data samples is crucial for improving performance beyond the results achieved by related methods. Training algorithms for the EBM, generator, and inference network are presented using several KL terms between different joint distributions of images and latent codes. Experiments show strong EBM synthesis results for CIFAR-10, CelebA, CelebA HQ, and LSUN-Church-64.

Strengths

* The work identifies and addresses a key limitation of the Cooperative Learning framework, which is the fact that the generator is trained to reconstruct only refined EBM samples. This can limit the potential of the generator since EBM samples are less realistic than true data. While other works such as Divergence Triangle train the generator to reconstruct data images using inferred codes for observed data, these works lack an MCMC refinement stage which is crucial for improving sample quality. This work bridges the gap between Cooperative Learning style methods and Divergence Triangle style methods to incorporate both MCMC sampling and generator data reconstruction into a single learning framework. * Experimental results show significant benefits from the proposed method compared to a variety of other EBM learning methods.

Weaknesses

* The loss functions appear to be significantly inspired by the Divergence Triangle, but the paper lacks a thorough comparison between the proposed method and Divergence Triangle. Such a discussion would help elucidate the differences between the proposed method and prior work. To my understanding, the primary differences are 1) replace all LHS of KL terms with distributions defined by short-run MCMC distributions with frozen model parameters, 2) add KL($\tilde{P} || P$) and KL($\tilde{Q} || Q$) terms to the generator and inference network updates respectively. * The motivation of each loss term is not entirely clear, which makes the proposed loss gradients appear somewhat ad-hoc. An intuitive explanation of the role of each KL term would be very helpful for the reader. A clearer discussion of connections to the Divergence Triangle could also help with understanding the KL terms. * The proposed method is somewhat complex, and the loss functions are defined sequentially given the current fixed model weights. Such an approach is common to Cooperative Learning and related methods, while amortized methods like Divergence Triangle have a loss function which can be written in closed form without reference to fixed model weights. While this is not a major issue in my view, the reliance on loss defined by a sequence of fixed model weights as opposed to a single loss function should be discussed as a limitation of the proposed method. Minor Problems * The related work [12] uses MCMC sampling initialized from the generator to update the EBM similar to [36, 37], instead direct generator samples like [11, 9, 19].

Questions

Could the authors provide a detailed comparison between their proposed method and the Divergence Triangle?

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

3 good

Presentation

3 good

Contribution

3 good

Limitations

A broader impact section was not included. Computational limitations are very briefly discussed, but more thorough details about runtime and compute costs compared to related methods would strengthen the paper.

Thanks for the responses

Thanks to the authors for their responses. I have read the responses to myself and other reviewers and find that they provide satisfactory answers. I will keep my score.

Thank you

I want to thank the authors for their detailed response. I think I'd be happy to see this published. That said, this is somewhat outside my field of expertise (as my low confidence score indicates) e.g., I was unaware of the divergence-triangle paper.