Learning Probabilistic Models from Generator Latent Spaces with Hat EBM

We briefly review the main components of EBM learning following the standard method derived from works such as [17, 44, 40]. A deep EBM has the form

where $U(x;\theta)$ is a deep neural network with weights $\theta$ and $\mathcal{Z}(\theta)$ is the intractable normalizing constant. Given a true but unknown data density $q(x)$, Maximum Likelihood learning uses the objective $\operatorname*{arg\,min}_{\theta}D_{KL}(q(x)\,\|\,p(x;\theta))$, which can be minimized using the stochastic gradient

where the positive samples $\{X_{i}^{+}\}_{i=1}^{n}$ are a set of data samples and the negative samples $\{X_{i}^{-}\}_{i=1}^{n}$ are samples from the current model $p(x;\theta)$. To obtain the negative samples for a deep EBM, it is common to use MCMC sampling with the $K$ steps of Langevin equation

where $\varepsilon$ is the step size and $V_{k}\sim N(0,I)$. The Langevin trajectories are initialized from a set of states $\{X^{-}_{i,0}\}_{i=1}^{n}$ obtained from a certain initialization strategy.

This section presents our method for adapting a fixed generator network $G(z)$ to be part of an EBM, which we call the Hat EBM. The Hat EBM defines the joint distribution of the random variable $Z\in\mathbb{R}^{m}$ in the $m$-dimensional latent space of the generator network and a random variable $Y\in\mathbb{R}^{d}$ in the $d$-dimensional image space. The joint energy has the form

where $H(x;\theta)$ is a neural network that takes an image $x\in\mathbb{R}^{d}$ as input and returns a scalar output. The weights of $H$ are given by $\theta$. We call $H$ the hat network because it sits atop the generator $G$ to incorporate the generator latent space directly into the probabilistic model.

The random variable $Y$ is meant to accommodate the gap between the output of $G(Z)$ and the image manifold, since in general we expect that $G(Z)$ contains an approximate but imperfect representation of the image distribution which can be refined by the residual state $Y$. In practice, we find that the majority of the appearance of a sampled image $X=G(Z)+Y$ comes from the generator output $G(Z)$ and not the residual image $Y$, indicating that the majority of the sampling behaviors of our model are determined by the latent space of $G$ (see Figure 3).

An appealing aspect of the Hat EBM formulation is that we can learn the model without either calculating the log determinant of the Jacobian of $G(Z)$ as would be required for an energy of the form $U(G(Z);\theta)$, or inferring the $Z$ latent vectors associated with observed images $X$ as done in existing work on learning latent prior EBMs [30]. This is possible because we define the distribution of observed images $X$ by $X=G(Z)+Y$ where the pair $(Y,Z)$ is drawn from a true but unknown density $q(y,z)$. We can then use the Maximum Likelihood framework in Section 3.1 to learn the weights $\theta$ of the hat network $H(x;\theta)$ by minimizing the objective $\operatorname*{arg\,min}_{\theta}D_{KL}(q(y,z)||p(y,z;\theta))$ where

One can obtain negative samples using alternating Langevin updates

which switch off between updates with respect to $z$ and updates with respect to $y$. This sampling algorithm is essentially Metropolis-within-Gibbs since the Langevin update can be written as a Metropolis-Hastings step with Gaussian proposal, in which case (6) and (7) define a valid sampler for $p(y,z;\theta)$. Finally, updating $\theta$ can be accomplished using

where $X_{i}^{+}$ are observed samples and the pairs $(Y_{i}^{-},Z_{i}^{-})$ are obtained via MCMC. In our formulation the observed data $X_{i}^{+}$ are sufficient statistics for $H(x;\theta)$ and there is no need to infer the $(Y_{i}^{+},Z_{i}^{+})$ pairs for the positive samples when learning the weights of the hat network.

Next we present a conditional variant of the Hat EBM. While the previous version of the Hat EBM is applicable to any generator network $G(z)$, including generators from deterministic autoencoders which cannot typically be sampled from, the conditional version of the Hat EBM is tailored towards generator networks which map a trivial latent distribution to complex signals like images. For these kinds of generators, one can use the known latent distribution as an ancestral distribution and learn a conditional distribution of the residual image given a latent sample. We emphasize that the conditional Hat EBM can be used to learn an unconditional distribution of observed images $X$ and that the term conditional refers to the relationship between the latent variables $Y$ and $Z$. Our experiments focus on modeling only observed images $X$ without conditional information such as labels or captions.

Suppose we use a trivial marginal distribution $p_{0}(z)$ for $Z$ and a generator $G$ trained to produce realistic images from this latent distribution. In our experiments, $p_{0}$ is always $N(0,I)$. We can now define a conditional Hat EBM density $p(y|z;\theta)=\frac{1}{\mathcal{Z}_{z}(\theta)}\exp\{-H(G(z)+y;\theta)\}$ and a joint density

In this case, we posit that observed images $X$ are generated according to $X=G(Z)+Y$ for some distribution $q(y,z)=p_{0}(z)q(y|z)$. Obtaining the negative samples $(Z_{i}^{-},Y_{i}^{-})$ is done by first drawing $Z_{i}^{-}$ from $p_{0}(z)$ and then obtaining $Y_{i}^{-}|Z_{i}^{-}$ by using Langevin updates on the conditional probability $p(y|z;\theta)$. Note that $p(y|z;\theta)=p(y|G(z);\theta)$ because of the form of (9). We update $\theta$ using the same equation (8) as the joint Hat EBM because $X_{i}^{+}$ is still a sufficient statistic for learning $H(x;\theta)$ and because we do not need to infer the $Z_{i}^{+}$ for $X_{i}^{+}$ since the prior $p_{0}(z)$ does not contain any model parameters. In practice we initialize Langevin sampling from $Y_{0}=0$ and $Z\sim N(0,I)$, and perform $K$ steps of (6) to draw a residual sample $Y_{K}$ while keeping $Z$ fixed.

Both formulations of the Hat EBM above assume that a pretrained generator network $G(z)$ is available as the basis for learning the Hat EBM. In order to use the Hat EBM as a self-contained learning process for image generation, we now propose a method to learn the weights $\phi$ of a generator network $G(z;\phi)$ and the weights $\theta$ of a hat network $H(x;\theta)$ simultaneously. Our method is based on the cooperative learning [39] strategy. We first briefly review the cooperative learning formulation, and then present the learning formulation for the Hat EBM. A key difference between the derivations is that the original cooperative learning method requires MCMC inference of the latent variable $\hat{Z}$ associated with an MCMC sample $X$ to train $G(z;\phi)$, while in our formulation the latent variable $Z$ is explicitly defined and does not need to be inferred. This enables major computational savings because we can bypass the expensive MCMC inference of $\hat{Z}$.

In the original cooperative learning [39], the generator output $G(z;\phi)$ is trained to match the appearance of a Langevin chain $X_{K}$ sampled from the potential $U(x;\theta)$ and initialized from the state $X_{0}=G(Z_{0};\phi)$ where $Z_{0}\sim N(0,I)$. The model defines the conditional density of an images $X$ given latents $Z\sim N(0,I)$ as $X|Z\sim N(G(Z;\phi),\tau^{2}I)$ for some sufficiently small $\tau>0$. Given a sampled state $X_{K}$, updating $G(z;\phi)$ requires inferring $Z|X_{K}$ using the latent Langevin equation

before updating $\phi$ using the Maximum Likelihood stochastic gradient

where the $i$ index denotes member $i$ of a batch with size $n$. We note that the code released with the cooperative learning method does not infer the latent variable of observed images and $Z_{0}\sim N(0,I)$ is used in place of $Z_{K}$ in the objective (11). In accordance with this approach, we find that inferring $Z_{K}$ often hurts model performance and leads to additional complication. The difficulty of inferring $Z|X_{K}$ and the omission of this step in practice leave an unresolved gap in the cooperative learning formulation. The Hat EBM generator update allows us to bypass the Langevin update for $Z$ without theoretical complications.

To update the weights of a Hat EBM generator, we propose to train the generator to match the image samples produced by the current Hat EBM. Given the current generator weights $\phi_{t}$ and hat network weights $\theta_{t}$ at step $t$, we define our learnable model as $Z\sim N(0,I)$ and $X|Z\sim N(G(Z;\phi),\tau^{2}I)$ and we define the true distribution of $(X,Z)$ as $Z\sim N(0,I)$ and $X=Y_{K}+G(Z;\phi_{t})$ where $Y_{K}|Z$ is drawn from $K$ steps of Langevin sampling with Hat EBM density $p(y|z;\theta_{t},\phi_{t})\propto\exp\{-H(G(z;\phi_{t})+y;\theta_{t})\}$. Langevin updating is only used to obtain $Y_{K}$ while $Z$ remains fixed, as in the method from Section 3.3. Then $\phi$ can be updated using the Maximum Likelihood objective

Conceptually, this loss function encourages $G(Z;\phi)$ to closely match the appearance of samples $X=Y+G(Z;\phi_{t})$ created from a fixed generator $G(Z;\phi_{t})$ and fixed hat network $H(x;\theta_{t})$. In other words, the generator update should achieve $G(Z;\phi_{t+1})\approx Y+G(Z;\phi_{t})$ so that the updated generator absorbs the residual image $Y|Z$ from the current Hat EBM. Once the generator absorbs the current Hat EBM residual, the hat network update should learn to synthesize residual images that refine the the generator output to be more similar to the observed data. Like the original cooperative learning method, our generator is trained using only synthetic images and no observed data images are used to update $\phi$.

Ideal training of $H$ and $G$ would alternate between using the gradient of (12) until generator convergence is reached and using the gradient (8) to update the hat network. In practice we implement the minimization in (12) using only one gradient update initialized from $\phi=\phi_{t}$ to obtain $\phi_{t+1}$ rather than training until full convergence. This is done to increase training efficiency and to avoid maintaining a separate copy of generator weights for the fixed network $G(z;\phi_{t})$.

In our experiments we observe that using a single gradient update with the objective (12) has limited success because the generator output can become too closely tethered to biases of the current hat network. We find the same problems with the original cooperative learning objective (11) (see Appendix F). To overcome these problems, we choose to train $G(z;\phi)$ at time $t+1$ to match the historical distribution of hat networks $H(x;\theta_{t_{\ell}})$ and generators $G(z;\phi_{t_{\ell}})$ for a selection of past epochs $t_{1},t_{2},\dots t_{L}\leq t$ instead of training the generator to match the distribution of the current hat network $H(x;\theta_{t})$ and generator $G(z;\phi_{t})$. This simply involves redefining the true distribution $(X,Z)$ by first sampling $t_{\ell}$ from $\{t_{1},\dots,t_{L}\}$ and then generating $Z\sim N(0,I)$ and $X_{K}=Y_{K}+G(Z;\phi_{t_{\ell}})$ where $Y|Z$ follows the energy $H(G(Z;\phi_{t_{\ell}})+Y;\theta_{t_{\ell}})$. In this case the loss (12), after replacing $\phi_{t}$ with $\phi_{t_{\ell}}$, is a stochastic approximation of the Maximum Likelihood gradient defined by the joint distribution $(t_{\ell},Z,X_{K})$ (see Appendix F). In practice, we implement this procedure by keeping a persistent bank of 10,000 pairs $(X,Z)$ created from past hat network updates. When updating the generator, we draw $n=128$ pairs from the bank and replace it with a newly generated batch of $n=128$ pairs from the current hat EBM. This ensures that the selection $\{t_{1},\dots,t_{L}\}$ of past epochs remains within a close range of the current epoch $t$ with high probability. Saving the generated images $X$ at each EBM update allows us to learn from past generator weights $\phi_{t_{\ell}}$ without maintaining a copy of the weights. See Figure 2 for an illustration and Appendix G for a code sketch.

In this section, we examine the problem of refining samples from a pretrained generator using a joint Hat EBM trained according to the method in Section 3.2. A pretrained SN-GAN with generator is used as a baseline generator $G$. We learn a Hat EBM that incorporates the generator network to refine the initial generator samples. The hat network has the exact same structure as the SN-GAN discriminator except that we remove spectral norm layers. We use fixed batch norm statistics for the generator and our energy is deterministic. The experiment is performed for the CIFAR-10 $32\times 32$ and CelebA $64\times 64$ datasets. To evaluate our method, we compare with Discriminator Driven Latent Sampling (DDLS) [5], which takes the pretrained discriminator $D$ learned with $G$ and samples from the potential $U(z)=D(G(z))+\frac{1}{2}\|z\|_{2}^{2}$. We train $D$ and $G$ using the Mimicry repository for reproducible GAN experiments [22] to obtain stronger baselines for $G(z)$ than those used in [5].

Table 1 shows our results. The Hat EBM can learn a hat network capable of refining a image appearance more successfully than DDLS on both CIFAR-10 and CelebA. We use the joint version of the Hat EBM where both $Y$ and $Z$ are updated during sampling. Sampling is initialized from $Y_{0}=0$ and $Z_{0}\sim N(0,I)$. The hat networks learns to tilt $Z$ away from its initial normal distribution to find a nearby latent vector with a more realistic appearance.

We observe that the majority of the refinement is occurring in the latent space, and that the residual image $Y$ is essentially imperceptible (see Figure 3). We also notice that training quickly becomes unstable when $Y$ is removed from the Hat EBM (see Appendix H). In our experience, it is essential to incorporate the residual for stable learning. A possible reason for this phenomenon is that the hat network can learn to discriminate between generator images and images not from the generator, whether they are realistic or not. If so, the hat network can assign increasingly high energy to generator samples in the absence of the residual $Y$. Even an imperceptible $Y$ is sufficient to prevent the hat network from easily distinguishing positive and negative samples so that learning is stable.

In this section, we incorporate a non-probabilistic generator network $G(z)$ into a probabilistic Hat EBM model. This essentially allows us to sample from the latent space of $G(z)$ to find latent samples whose mapping corresponds to a realistic image. Like the results in Section 4.1, the residual image has small norm in the image space and most of the appearance of the sampled images comes from $G(z)$. This happens naturally without the need to coerce $Y$ to be close to 0 by including a prior term such as $p_{0}(y)\propto\exp\{-\frac{1}{2\sigma^{2}}\|y\|^{2}_{2}\}$, although including a prior term can further limit growth of $Y$.

The autoencoder generator $G(z)$ is pretrained as the second half of a standard inference network and generator network pairing. An image $X$ is fed into the inference network $I$ and converted into a latent state $Z=I(X)$, which is then decoded to obtain a reconstruction $\hat{X}=G(Z)$. The inference network and generator are learned jointly using the MSE loss $\|\hat{X}-X\|_{2}^{2}$. To keep the latent space mapping of $I(X)$ numerically stable, we project the $m$-dimensional raw output of the inference network to the sphere around the origin with radius $\sqrt{m}$ so that $\|I(X)\|_{2}=\sqrt{m}$. More sophisticated methods such as perceptual and adversarial loss could have been used to train the autoencoder, but we use MSE loss to keep our implementation minimal. We observe that when $Z$ is a vector-shaped latent state, it can be difficult to learn reconstructions $\hat{X}$ with sharp appearance even for simple datasets like CIFAR-10. To obtain better reconstructions and therefore a latent space with more realistic mappings to the image space, we use image-shaped latent states $Z$. The details of our autoencoder networks can be found in Appendix E. When $Z$ is an image shaped latent, we treat it exactly the same as a vector latent in the learning and sampling algorithms.

We experiment with assimilating an autoencoder generator into a Hat EBM potential for the CIFAR-10 dataset. Our results are visualized in Figure 3 with additional results in the Appendix K. We train the Hat EBM using shortrun learning in the latent and image space by initializing $Y_{0}=0$ and $Z_{0}\sim N(0,I)$ and using $K=100$ MCMC steps of (6) and (7) from initialization during both training and testing evaluation to generate samples. During the Langevin dynamics, image appearance is refined mostly in the latent space. Our best model achieves a solid FID score of 26.01 $\pm$ 0.09. This demonstrates that Hat EBM can learn a probabilistic model over a non-probabilistic latent space.

In this section, we use the conditional Hat EBM formulation from Section 3.4 to learn a hat network and generator network from scratch for self-contained synthesis. We explore synthesis for unconditional CIFAR-10 at resolution $32\times 32$, CelebA at resolution 64$\times$64, and unconditional ImageNet at resolution 128$\times$128. While recent generative models show promising results for class conditional sampling, unconditional sampling with high quality synthesis remains a significant challenge. We find especially strong results for unconditional ImageNet synthesis using the Hat EBM. This demonstrates strong potential of our synthesis method for learning with highly diverse unstructured datasets.

Our Hat EBM models use SN-GAN architectures for sizes $32\times 32$, $64\times 64$, and $128\times 128$, where the discriminator architecture is used for the hat network. During learning, we keep the generator batch norm parameters fixed to mean 0 and variance 1. We remove all spectral norm layers from the hat network. Training parameters can be found in Appendix J. For ImageNet models, we found that annealing the generator and hat network learning rate by a factor of 10 after 250K weight updates for each network improved the FID score significantly. See Appendix K for uncurated samples from our Hat EBM models. FID results for our model and a representative selction of generative models are shown in Table 2. We were unable to reproduce the GEBM FID scores reported in [3] and we report the score obtained from a reimplementation with the official training code. For the SN-GAN FID score, we report the stronger baseline from our reimplementation in Section 4.1.

Results show strong performance of the Hat EBM compared to competing generative models across all datasets, with an especially strong performance for ImageNet. Our method significantly outperforms other EBM learning methods on CIFAR-10. The Hat EBM synthesis results are on par with the SN-GAN baseline for CIFAR-10 and CelebA, and the Hat EBM significantly outperforms SN-GAN for ImageNet. With a budget of 8 GPUs and about 2.5 days of computing, our Hat EBM achieves an ImageNet FID score of 40.0, outperforming the small-scale SS-GAN.

To our knowledge, the current state-of-the-art model for unconditional ImageNet synthesis at resolution $128\times 128$ is the large-scale SS-GAN [6], which achieves an FID score of 23.4. This model was trained using a BigGAN architecture and 128-core TPUv3 pods. To scale up our Hat EBM, we doubled the number of channel dimensions for both the hat network and generator network from the original SN-GAN architecture and trained on 32-core TPU-v3 pods. Our best FID score for unconditional ImageNet 128$\times$128 was 29.2, which comes within a competitive range of state-of-the-art. We believe that further scaling in future work could enable Hat EBM to match or surpass state-of-the-art. Our results decisively demonstrate the potential of EBM learning for high-quality synthesis well beyond the scale investigated in any prior EBM work.

Experiments in this section assess Hat EBM performance on Out-Of-Distribution (OOD) detection. We use the conditional Hat EBM model trained on CIFAR-10 from Section 4.3 and calculate the energy $H(X;\theta)$ on unseen in-distribution images from the CIFAR-10 test set and images from OOD datasets which include CIFAR-100, CelebA, and SVHN. We follow standard OOD evaluation from works such as [27] which compute the AUROC metric on the energy scores of the in-distribution and OOD samples. This metric measures the ability of the Hat EBM to distinguish between in-distribution samples not seen during training and OOD samples. Following [13, 38], we expect that the energy of the OOD datasets will be higher than the energy of in-distribution test images since higher energy samples should appear with lower frequency in the learn density.

Our results are shown in Table 3. The Hat EBM shows strong performance as an OOD detection method. Among methods that are fully unsupervised, our model has the top performance across all three OOD datasets. Our method approaches the results of methods that are trained with labeled data such as HDGE  [24] and the fine-tuned OOD EBM [25], although we do not yet match these scores. Overall, there is strong evidence that the Hat EBM is naturally an effective method for OOD detection, especially when supervised label information is unavailable.