Learning energy-based models by self-normalising the likelihood

We want to thank reviewer qYkm for their careful analysis of our paper and suggestions for improvement.

1 - This paper has a well-motivated idea and contains comprehensive theoretical derivation for understanding the key idea. However, as mentioned by the author, the NCE method is related, it would be nice to have a deeper theoretical connection and comparison with the NCE method. For now, the major comparison is shown by empirical experiments.

Thank you for suggesting an investigation of the relationship between the NCE and SNL objectives. We found that the SNL loss resembles one of the generalisations of NCE proposed in a theoretical paper on NCE [2]. We will discuss this paper in the revised version.

2 - Many other prior works can be applied to some more challenging real data, such as CIFAR-10, CelebA-64, or even the high-resolution (CelebA-HQ-256), so what limited this learning algorithm for such dataset?

The goal of our paper is to introduce a new method for training Energy-Based Models. As it was done for other methods in the field, scaling such methods requires tuning and tricks that are a completely new avenue of research. For image modelling, we will require more complicated tuning and proposal than just a simple Gaussian maybe using a flow as a proposal. For instance [1] focuses on MCMC-based methods but scales it using a replay buffer.

3 - As a novel learning method, it would be nice to have a practical learning algorithm to simplify and illustrate the main idea.

We added an algorithm description for our method in appendix E. Thank you for the suggestions.

[1] Du, Y. and Mordatch, I. "Implicit Generation and Modeling with Energy Based Models "Advances in Neural Information Processing Systems 32 (2019)

[2] Pihlaja et al. (2010), A family of computationally efficient and simple estimators for unnormalized statistical models, Uncertainty in Artificial Intelligence

We thank reviewer 3VY8 for their study of our paper and the recommendation for future iterations of the paper.

1 - Can you provide more real-world experiments? For instance in generative modeling (without the VAE component) or out-of-distribution detection. 2 - As mentioned in the weaknesses, I would expect your method to be very sensitive to the design of $q$. (a) Did you run a sensitivity comparison with [1], [2] (which develop similar ideas) or NCE in higher dimensional settings (compared to Sec 4.2)? (b) Did you try to learn as done in [3] for NCE ?

Answer for 1 and 2: We agree that the method will be sensitive to the proposal, especially when scaling it for image modelling. However, we feel that the work required for finding and tuning a correct proposal is a new avenue of work in itself. Indeed, all these papers are mostly focused on scaling existing methods by finding and tuning a proposal (noisy neural proposal in [1] using $l_{IS}$ in our paper or flow proposal in [3] for NCE). Note that we did compare $l_{IS}$ and $l_{SNL}$ with a Gaussian for the 2D toy experiment and the image regression experiments. In such cases, training the EBM was less stable than with our method or NCE.

3 - In [4], the authors give a very similar result as your theorem 3.1 but for NCE. Are there more theoretical comparisons to be drawn against NCE ?

The paper you are mentioning also provides more theoretical insights into the loss landscape for NCE. Conducting such a study for SNL would require substantial work but is an interesting new avenue of work but we will integrate it to a revised version of the paper. Thank you for suggesting an investigation of the relationship between the NCE and SNL objectives. We do not know if there is a direct relationship between our loss and the NCE one. However, we found that the SNL loss resembles one of the generalisations of NCE proposed in a theoretical paper on NCE [5]. We will discuss this paper in the revised version.

4 - As mentioned in the weaknesses, I think it would be nice to compare SNL against MCMC-based methods (at least Langevin based) with apple-to-apple computational budgets.

Due to the lack of time required to implement and tune MCMC-based methods, we will provide a comparison of SNL with MCMC based method using similar computational budgets (ie number of call of a neural networks and time) upon acceptance or for a further iteration of the paper.

[1] Will Grathwohl, Jacob Kelly, Milad Hashemi, Mohammad Norouzi, Kevin Swersky, & David Duvenaud. (2021). No MCMC for me: Amortized sampling for fast and stable training of energy-based models.

[2] Hanjun Dai, Rishabh Singh, Bo Dai, Charles Sutton, & Dale Schuurmans. (2020). Learning Discrete Energy-based Models via Auxiliary-variable Local Exploration.

[3] Ruiqi Gao, Erik Nijkamp, Diederik P. Kingma, Zhen Xu, Andrew M. Dai, & Ying Nian Wu. (2020). Flow Contrastive Estimation of Energy-Based Models.

[4] Bingbin Liu, Elan Rosenfeld, Pradeep Ravikumar, & Andrej Risteski. (2021). Analyzing and Improving the Optimization Landscape of Noise-Contrastive Estimation.

[5] Pihlaja et al. (2010), A family of computationally efficient and simple estimators for unnormalized statistical models, Uncertainty in Artificial Intelligence

We thank the reviewer CuDM for their careful review of our paper and address their concern below.

1 - From my own experience, the most challenging part when training the EBM is to get valid samples from the current fitted distribution to estimate the (gradient of) normalizing constant. [...]

Our experience matches the one of reviewer CuDM, where the difficulty is to get good samples from the current fitted distribution. In MCMC-based methods training the EBM requires long chains or tricks (for instance, keeping a buffer in [1]) to avoid biased training. However, our method does not require sampling from the model to estimate gradients and that is why we believe this method is more flexible and avoids biases training loss.

2 - To do this Monte Carlo estimation, the work employs important sampling using a base distribution and simple distributions like Gaussian. I suspect this algorithm works because the target distributions tested in this work are very simple,[...] .

Though we agree that having a simple distribution like a Gaussian might not scale to more complicated high-dimensional data like images, it does not necessarily mean the need to resort to MCMC-based methods again. Indeed, having a more complex proposal distribution like a flow as a proposal or base distribution could be a way to handle this issue.

3 - The proposed algorithm introduces a variational parameter b, which requires updating b and the energy function iteratively. Then similar to the VAE case, whether there can be a mismatch between the estimate of b and the energy function. (Not sure whether the term will make the training more unstable if b is not well optimized.) Or in other words, how difficult is it to design the schedule of updating b and energy function to make this algorithm work?

It is true that the quality of the gradient estimate depends on how close $b$ is to the normalization constant. This is clearly seen in equation (17) of the paper where we rewrite SNL gradients as the likelihood gradients with some negligible terms if $b$ verifies the aforementioned condition :

\nabla_{\theta} \ell_{\text{snl}}(\theta,b) &= - \frac{1}{n} \sum_{i=1}^{n} \nabla_{\theta} E_{\theta}(x) - e^{-b+\log{Z_{\theta}}} \nabla \log{Z_{\theta}} \\ & = \nabla_{\theta} \ell_{\theta} + \nabla_{\theta} \log{Z_{\theta}} (1-e^{-b+\log{Z_{\theta}}}). \end{aligned}

In practice, we train both $\theta$ and $b$ within the same iteration and very quickly during training, $b$ gets very close to $\log{Z_{\theta}}$ thus this did not require more tuning from our side. Note that we treat the extra parameter $b$ as the bias of the last layer of neural network and as such as any other parameter weight of the neural network, optimized using Adam.

4 - As also mentioned in 2, the modelled distributions in the experiments are too simple to be convincing to me. The modelled experiments are either unconditional distribution on toy data or with image input but only model the conditional distribution on some low dimensional label. The VAE experiment in 5.3 models binary MNIST (which is also not very complex).[...]

Though we agree that the experiments are low-dimensional, we disagree that they are not convincing enough. Indeed, to the best of our knowledge, this is the first time an EBM is used for density estimation (and actually provides an upper and lower bound of the likelihood) on the UCI datasets. These datasets usually require complicated networks (autoregressive flows for instance), whereas we reach similar results with much simpler networks and a simple Gaussian proposal. Similarly, we show state-of-the-art performance on EBMs for image regression which is in itself useful [3] for visual tracking or pose-estimation in computer vision. Finally, the point of the paper is to introduce a new method for training general EBMs, without a specific focus on high-dimensional data. As it was done for other methods in the field, scaling such methods requires tuning and tricks that are a completely new avenue of research.

4 bis - I think in order to make the proposed algorithm more convincing, the authors need to demonstrate better results than pure MLE loss on more complex distributions like real images (face or cifar or SVHN).

Could you specify what you call pure MLE loss? Is that just the standard VAE?

5- We added the recommended paper to the references, we want to thank the reviewer for their valuable suggestions.

[1] Du, Y. and Mordatch, I. "Implicit Generation and Modeling with Energy Based Models "Advances in Neural Information Processing Systems 32 (2019)

[2] Will Grathwohl, et al. "No MCMC for me: Amortized samplers for fast and stable training of energy-based models." (2021).

[3] Gustafsson, F. K., Danelljan, M., Timofte, R., & Schön, T. B. (2020). How to train your energy-based model for regression. arXiv preprint arXiv:2005.01698.

Thank you for clarifying your question.

Also, I further found that the performance is measured by SNELBO. Will this metric favor the model trained with ? What if pure ELBO is reported?

Thank you for your question. "Pure" ELBO is actually reported when it is possible to report it (ie for Vanilla VAE or Mixture of Gaussian Prior VAE where we have access to the exact prior value). When there is an EBM in the prior distribution of the VAE, it is not possible to report "pure" ELBO because one still needs to estimate the normalization constant :

$ELBO(\theta, \phi, \gamma) = E_{q_{\gamma}(z|x)}\left[\log p_{\phi}(x|z)\right] + E_{q_{\gamma}(z|x)} \left[\log{\frac{d(z)}{q_{\gamma}(z|x)}}\right] + E_{q_{\gamma}(z|x)}\left[-E_{\theta}(z)\right] - \log(Z)$

+ E_{q_{\gamma}(z|x)}\left[-E_{\theta}(z)\right] - \log \int_z e^{-E_{\theta}(z)} .$$ It is then possible, to estimate an upper bound of the ELBO, using yet another proposal distribution $q$ and Jensen's inequality : $$ELBO(\theta, \phi, \gamma) \leq E_{q_{\gamma}(z|x)}\left[\log p_{\phi}(x|z)\right] + E_{q_{\gamma}(z|x)} \left[\log{\frac{d(z)}{q_{\gamma}(z|x)}}\right] + E_{q_{\gamma}(z|x)}\left[-E_{\theta}(z)\right] - E_q \left[-E_{\theta}(z) \right] .$$ This is not a good estimation because there is no guarantee that this is a lower bound of the likelihood $\ell_{\theta, \phi, \gamma}$. This is why we resort to SNELBO which is a lower bound of the ELBO and thus a lower bound of the likelihood. Note that this was not possible to get such an estimate of the likelihood for VAE-EBM models before SNL, maybe this is the reason why there is no ELBO or likelihood evaluation in [1] > However, this doesn't necessarily imply that the EBM trained with the techniques proposed in this paper (loss) is superior to one trained with the standard maximum likelihood loss, as practiced in previous studies[1]. We want to point out to the reviewer that [1] is not doing maximum likelihood but moment matching. Indeed, [1] extend short term non persistent non convergent mcmc methods [2] from EBM to VAE-EBM. In [2], it is explicitely stated that this is not pure maximum likelihood. To do maximum likelihood, we will do longer mcmc chains than [1]. > Instead, we use the plain maximum likelihood loss as shown in eq 3 and 4. (As done in [1].) We provide here results on Binary MNIST with two extra results. In both cases, we learn a similar model to [1] using either short term MCMC (ie 50 steps of Langevin Dynamics with the same step-size as [1]) and long term MCMC (ie 1000 steps of Langevin Dynamics using the same step size as [1]) akin to "plain maximum likelihood". We train the Short-Term MCMC VAE-EBM for 100000 training steps as [1] and the long term MCMC VAE-EBM for 25000 training steps (training is still on-going and may perform better than VAE-EBM SNELBO but is much slower). | | (SN)ELBO | FID | | --------------------------- |:----------:|:-----:| | VAE | -89.10 | 17.23 | | VAE-MoG | -88.73 | 14.3 | | {VAE-EBM SNELBO Post-Hoc} | -88.11 | 14.24 | | VAE-EBM SNELBO | **-87.09** | 6.42 | | Short Term MCMC VAE-EBM [1] | -850.1 | 21.68 | | Long Term MCMC VAE-EBM | -980.3 | 8.06 | Estimating SNELBO with MCMC is not straightforward. First, there is a lot of variance in the normalization constant estimate required calculating the SNELBO. More over, SNELBO requires an approximate proposal. Here we train a posteriori (i.e., after training) an approximate posterior with a Gaussian parameterized with the same Neural network. This is the same setting as the other VAE, but the bad results in ELBO may show that this is not a correct parameterization for the posterior. We also report the FID and show that VAE-EBM (trained with SNELBO) performed better. [1] Pang, B., Han, T., Nijkamp, E., Zhu, S. C., & Wu, Y. N. (2020). Learning latent space energy-based prior model. Advances in Neural Information Processing Systems, 33, 21994-22008. [2] 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.