Exploring Bidirectional Bounds for Minimax-Training of Energy-Based Models

Abstract

Energy-based models (EBMs) estimate unnormalized densities in an elegant framework, but they are generally difficult to train. Recent work has linked EBMs to generative adversarial networks, by noting that they can be trained through a minimax game using a variational lower bound. To avoid the instabilities caused by minimizing a lower bound, we propose to instead work with bidirectional bounds, meaning that we maximize a lower bound and minimize an upper bound when training the EBM. We investigate four different bounds on the log-likelihood derived from different perspectives. We derive lower bounds based on the singular values of the generator Jacobian and on mutual information. To upper bound the negative log-likelihood, we consider a gradient penalty-like bound, as well as one based on diffusion processes. In all cases, we provide algorithms for evaluating the bounds. We compare the different bounds to investigate, the pros and cons of the different approaches. Finally, we demonstrate that the use of bidirectional bounds stabilizes EBM training and yields high-quality density estimation and sample generation.

Appendices

Appendix A Model architecture

Our key motivation for design of $p_g$ is to ensure that the model has support over the entire data space. This is needed for the reason that the KL divergence requires the two distributions to be defined on the same measurable space. Since our energy distribution is defined for the entire data space, we need the generated distribution to have the same support. The simplest way to achieve this is with our choice of $x=G(z)+\epsilon$, where $\epsilon$ is a small-variance Gaussian variable. This concentrates the distribution around the manifold spanned by our generator while ensuring that the generated distribution has full support over the data space. Moreover, if we do not add noise to our generator, our mutual information lower bound is infinite. The added noise makes the bound finite.

Table 9 Selection of most important hyper-parameters and their setting

We consistently assume that $G: \mathbb {R}^d \rightarrow \mathbb {R}^D$ spans an immersed d-dimensional manifold in $\mathbb {R}^D$ as this allows us to apply the change-of-variables formula to get a density for the generator. In order to ensure that this assumption is valid, we minimally restrict the architecture of the generator neural network.

An immersion implies that the Jacobian of G exists and has full rank. We can then ensure existence by using activation functions with at least one derivative almost everywhere. Any smooth activation satisfies this assumption, but we emphasize that, for example, ReLU activations lack a derivative at a single point. The non-smooth region has measure zero if the linear map inside the activation is not degenerate, in which case the change-of-variables technique still applies.

We are unaware of architectural tricks to ensure that the Jacobian has full rank. Still, a minimum requirement is no hidden layer may have dimensionality less than the d dimensions of the latent space. This is a natural requirement that is practically always satisfied. Our model maximizes the generator entropy, which practically ensures full-rank Jacobians (noting that degeneracy in the Jacobian reduces the entropy). We have in practice generally not experienced degenerate Jacobians during training.

Appendix B Training details

Table 9 gives the hyper-parameters used during training. In Eq. 33 we set $p=2$. We normalize the data to $[-1, 1]$ and do not use dequantization. We augment only using random horizontal flips during training. We used the official train-test split from PyTorch for cifar10 and an 85:15 train-test split for animeface. We set $\frac{M}{s_1^2}=\frac{0.1}{z_{dim}}$ for our gradient penalty upper bound, where $s_1$ is the smallest singular value of Jacobian and $z_{dim}$ is the latent dimension. We find this set of parameters generally works well across datasets. Furthermore, we choose $\frac{M}{s_1^2}$ with dynamic decay, i.e.@ $\frac{M}{s_1^2}=\frac{a(b+i)^{-\gamma }}{2}$ where i denotes the iteration, $\gamma = .55$ and a and b are chosen such that $\frac{2M}{s_1^2}$ decreases from 0.01 to 0.0001 during training. These dynamics are consistent with the approximation of $\textrm{KL}\left( p_g \Vert p_\theta \right)$ from the Taylor series point of view. We find that practically for $\textrm{EBM}_{\textrm{SV}+\textrm{GP}}$, $\frac{M}{s_1^2}$ must be adjusted with entropy regularizer $\lambda$ together. We also observe the implications of the network design on the final performance. For example, not using batch normalization in the energy function can improve the stability of training on ANIMEFACE dataset.

Appendix C Selection of bidirectional bound pairings

Given the two available options for both the lower and the upper bound, we provide recommendations for selecting bidirectional bound pairings based on our training experience. For real-world applications, as elaborated in the main text, we generally recommend using combinations $\textrm{EBM}_{\textrm{SV}+\textrm{GP}}$ or $\textrm{EBM}_{\textrm{MI}+\textrm{diff}}$ first due to their respective advantages in memory and time efficiency. If the data and latent space dimensions are relatively small and the network is shallow, we recommend using $\textrm{EBM}_{\textrm{SV}+\textrm{GP}}$, as it tends to achieve better performance in such scenarios. Conversely, for larger dimensions or deeper networks, we suggest using $\textrm{EBM}_{\textrm{MI}+\textrm{diff}}$ to reduce computational cost. Furthermore, $\textrm{GP}$ and $\textrm{SV}$ may suffer from training instability due to their reliance on iterative optimization. Therefore, in cases of unstable training, it is advisable to consider replacing $\textrm{GP}$ or $\textrm{SV}$ with $\textrm{diff}$ or $\textrm{MI}$.

Table 10 Comparison with diffusion-based models on $32\times 32$ CIFAR-10 dataset

Fig. 10

Visual results of CIFAR-10 using UNET

Appendix D Additional results

1.1 D.1 CIFAR-10

Although our bidirectional EBM obtains superior performance among EBMs, it still lags far behind diffusion-based models in terms of generation metrics. To demonstrate the scalability and potential of our model, we also scale up our network structure to UNet for CIFAR-10 dataset. We adopt the same architecture in UOTM (Choi et al., 2024), where its discriminator is utilized as our energy function. Our generator has the same scale as the network in NCSN and DDPM. We borrow some training strategies from UOTM to stabilize training and improve generative performance, such as the form of energy discrepancy and cost function. For the choice of bidirectional bounds, we only use our $\textrm{EBM}_{\textrm{MI}+\textrm{diff}}$ since UNet network is complex and the input has the same dimension as the output, causing a costly Jacobian computation. We present our quantitative results in Table 10 and qualitative results in Fig. 10. All baseline results used for comparison are sourced from Xiao et al. (2022) and Song et al. (2023). For our method, we run the experiments with different seeds and report the standard deviation. For FID evaluation, we use the training set as a reference, aligning common practice in diffusion-based models (Xiao et al., 2022). From Table 10 we can observe our $\textrm{EBM}_{\textrm{MI}+\textrm{diff}}$ achieves results comparable to diffusion-based models while using only one-step generation, offering significantly greater efficiency. Although recent techniques like distillation and consistency models can also enable one-step generation with minimal performance decline, our model demonstrates competitive results compared with them. This experiment suggests that the superiority of diffusion-based models stems more from the large network architecture and meticulous parameter fine-tuning than from the methodology itself. Our bidirectional EBM can achieve performance comparable to recent advanced diffusion-based models when equipped with a similar-scale UNet and an appropriate training strategy. Moreover, as shown in Geng et al. (2024), EBM can be further improved by integrating it with a diffusion process, which presents an opportunity for future work.

Fig. 11

Visual results of $128\times 128$ LSUN Church with two different settings

Table 11 FID scores on $32\times 32$ ImageNet dataset

1.2 D.2 Imagenet

To further evaluate the efficacy of our model on complex datasets that better approximate real-world scenarios, we assess its generative performance on $32\times 32$ ImageNet dataset, which contains 1,281,167 training images and 50,000 test images. We perform experiments using both DCGAN and UNet architectures, as applied to the CIFAR-10 dataset. For $\textrm{EBM}_{\textrm{SV}+\textrm{GP}}$, we set the number of inner loop iterations in $\textrm{SV}$ to zero and find that they contribute to stabilizing the training process; a detailed exploration of this phenomenon is left for future work. Table 11 and Fig. 12 show our quantitative and qualitative results respectively. For quantitative evaluation, we report only the FID score, as it is the most widely used metric for assessing generative models on this dataset. From Table 11 we can see our $\textrm{EBM}_{\textrm{MI}+\textrm{diff}}$-Large gets the best performance compared to other advanced generative models including diffusion-based model DDPM++ (Kim et al., 2021) and the latest energy-based model CDRL (Zhu et al., 2024). Even our $\textrm{EBM}_{\textrm{SV}+\textrm{GP}}$ and $\textrm{EBM}_{\textrm{MI}+\textrm{diff}}$ only use simple DCGAN architecture, they still outperform competitive EBM framework CLEL-Large which employs a deep ResNet network. This experiment further verifies the effectiveness of our bidirectional EBM on large complex datasets.

Fig. 12

Visual results of $32\times 32$ ImageNet dataset with different settings

1.3 D.3 LSUN Church

We also test our model on a large-scale LSUN dataset. For this dataset, we choose the Church outdoor category with a size of $128\times 128$. We use 126,227 training images and 300 testing images following the official PyTorch split. The latent dimension is set to 256. We employ the same ResNet architecture as that used for the ANIMEFACE and CelebA datasets. We trained $\textrm{EBM}_{\textrm{SV}+\textrm{GP}}$ for 100,000 iterations and $\textrm{EBM}_{\textrm{MI}+\textrm{diff}}$ for 200,000 iterations. Figure 11 illustrates the scalability of our models. While these samples are not state-of-the-art, they do showcase the ability of our model to scale to high-dimensional problems. Further improving these results most likely requires work on picking the right network architecture, fine-tuning parameters, and specialized optimization strategies. We leave this for future work.