\section{An Improved Contrastive Divergence Framework for Energy-Based Models}
Energy-Based Models (EBMs) represent the likelihood of a probability distribution $p_D({\bm{x}})$ for ${\bm{x}} \in \mathbb{R}^D$ as $p_{\theta}({\bm{x}}) = \frac{\exp(-E_{\theta}({\bm{x}}))}{Z(\theta)}$ where the function $E_{\theta}({\bm{x}}): \mathbb{R}^D \rightarrow \mathbb{R}$, is known as the \textit{energy function}, and $Z(\theta) = \int_{{\bm{x}}} \exp{-E_{\theta}({\bm{x}})}$ is known as the partition function. Thus, an EBM can be represented by an neural network that takes ${\bm{x}}$ as input and outputs a scalar.

Training an EBM through maximum likelihood (ML) is not straightforward, as $Z(\theta)$ cannot be reliably computed, since this involves  integration over the entire input domain of ${\bm{x}}$. However, the gradient of log-likelihood with respect to a data sample ${\bm{x}}$ can be represented as
\begin{equation}
    \label{eq:aml}
    \resizebox{0.85\hsize}{!}{$
    \frac{\partial \log p_{\theta} ({\bm{x}})}{\partial \theta} = -\left( \frac{\partial E_\theta({\bm{x}})}{\partial \theta} - \mathbb{E}_{p_{\theta}({\bm{x}}')} \left [\frac{ \partial E_\theta({\bm{x}}')}{\partial \theta} \right] \right) $}.
\end{equation}
Note that \eqn{eq:aml} is still not tractable, as it requires using Markov Chain Monte Carlo (MCMC) to draw samples from the model distribution $p_{\theta}({\bm{x}})$, which often takes exponentially long to mix. As a practical approximation to the above objective, \citep{Hinton2002Training} proposes the contrastive divergence objective
\begin{equation}
    \kldiv{p_D({\bm{x}})}{p_{\theta} ({\bm{x}})} - \kldiv{\Pi_\theta^t(p_D({\bm{x}}))}{p_{\theta}({\bm{x}})},
    \label{eq:cd}
\end{equation}
where $\Pi_\theta$ represents a MCMC transition kernel for $p_{\theta}$, and $\Pi_\theta^t(p_D({\bm{x}}))$ represents $t$ sequential MCMC transitions starting from $p({\bm{x}})$.  In this objective, if we can guarantee that%
\begin{equation}
    \kldiv{p_D({\bm{x}})}{p_{\theta} ({\bm{x}})} \ge \kldiv{\Pi^t_\theta(p_D({\bm{x}}))}{p_{\theta}({\bm{x}})},
    \label{eqn:ge}
\end{equation}%
then the objective guarantees that $p_{\theta}({\bm{x}})$ converges to the data distribution $p_D({\bm{x}})$, since the objective is only zero (at its fixed point) when $p_{\theta}({\bm{x}}) = p_D({\bm{x}})$.

If $\Pi$ represents a MCMC transition kernel, this property is guaranteed \citep{lyu2011unifying}. Note that $\Pi$ \textit{does not} need to converge to the underlying probability distribution, and only a finite number of steps of MCMC sampling may be used. In fact, this objective may be utilized to maximize likelihood even if $\Pi$ is \textit{not} a MCMC transition kernel, but instead a model such as an amortized generator, as long as we ensure that \eqn{eqn:ge} holds. In the appendix, we show that our approach applies even when MCMC chains are not initialized from the data distribution.

\input{figText/pipeline}

\subsection{A Missing Term in Contrastive Divergence}
\label{sect:missing}
When taking the negative gradient of the contrastive divergence objective (\eqn{eq:cd}), we obtain the expression%
\begin{equation}
    \label{eq:kl_grad}
    \resizebox{0.8\hsize}{!}{$
    \begin{split}
    - ( \mathbb{E}_{p_D({\bm{x}})} \left [\frac{\partial E_{\theta}({\bm{x}})}{\partial \theta} \right] &- \mathbb{E}_{q_{\theta}({\bm{x}}')} [\frac{\partial E_{\theta}({\bm{x}}')}{\partial \theta}] \\
    &+ {\color{red} \frac{\partial q({\bm{x}}')}{\partial \theta} \frac{\partial \kldiv{q_{\theta}({\bm{x}}')}{p_{\theta}({\bm{x}}')}}{\partial q_{\theta}({\bm{x}}')} } )
    \end{split} $},
\end{equation}
where for brevity, we summarize $\Pi_\theta^t(p({\bm{x}})) = q_{\theta}({\bm{x}})$. The first two terms are identical to those of \eqn{eq:aml} and the third gradient term (which we refer to as the {\color{red}KL divergence term}) corresponds to minimizing the divergence between $q_{\theta}({\bm{x}})$ and $p_{\theta}({\bm{x}})$. In practice, past contrastive divergence approaches have ignored the third gradient term, which was difficult to estimate and claimed to be empirically negligible \citep{hinton1999products} (which in \fig{fig:kl_cd} we show to be non-negligible), leading
to the incorrect optimization of \eqn{eq:cd}. To correctly optimize \eqn{eq:cd}, we construct a new joint loss expression $\mathcal{L}_{\text{Full}}$,  consisting of traditional contrastive loss $\mathcal{L}_{\text{CD}}$ and a new loss expression $\mathcal{L}_{\text{KL}}$, to accurately exhibit all three gradient terms. Specifically, we have
$\mathcal{L}_{\text{Full}} = \mathcal{L}_{\text{CD}} + \mathcal{L}_{\text{KL}}$ where $\mathcal{L}_{\text{CD}}$ is
\begin{equation}
    \mathcal{L}_{\text{CD}} =  \mathbb{E}_{p_D({\bm{x}})} [E_{\theta}({\bm{x}})] - \mathbb{E}_{\text{stop\_grad}(q_{\theta}({\bm{x}}'))} [E_{\theta}({\bm{x}}')],
    \label{eq:cd_loss}
\end{equation}
and the ignored KL divergence term corresponding to the following KL loss:
\begin{equation}
    \resizebox{0.87\hsize}{!}{$
    \mathcal{L}_{\text{KL}} =  \mathbb{E}_{q_{\theta}({\bm{x}})} [E_{\text{stop\_grad}(\theta)}({\bm{x}})] + \mathbb{E}_{q_{\theta}({\bm{x}})}[\log (q_{\theta}({\bm{x}}))] $}.
    \label{eq:kl_loss}
\end{equation}
Despite being difficult to estimate, we show that $\mathcal{L}_{\text{KL}}$ is a useful tool for both speeding up and stabilizing training of EBMs.  We provide derivations showing the equivalence of gradients of $\mathcal{L}_{\text{Full}}$ and that of \eqn{eq:cd} in the appendix, where stop gradient operators are necessary to ensure correct gradients. \fig{fig:pipeline_img} illustrates  the overall effects of both losses. \eqn{eq:cd_loss} encourage the energy function to assign low energy to real samples and high energy for generated samples. However, only optimizing \eqn{eq:cd_loss} often leads to an adversarial mode where the energy function learns to simply generate an energy landscape that makes sampling difficult. The KL divergence term counteracts this effect and encourages sampling to closely approximate the underlying distribution $p_\theta({\bm{x}})$, by encouraging samples to be both low energy under the energy function as well as diverse. Empirically, we find that including for KL term significantly improves both the stability, generation quality, and robustness to different model architectures (\fig{fig:kl_loss}). Next, we will discuss our approach towards estimating this KL divergence.

\subsection{Estimating the missing gradient term}

Estimating $\mathcal{L}_{\text{KL}}$ can further be decomposed into two separate objectives, minimizing the energy of samples from $q_{\theta}({\bm{x}})$, which we refer to as $\mathcal{L}_{\text{opt}}$ (\eqn{eq:opt_loss}) and maximizing the entropy of samples from $q_{\theta}({\bm{x}})$ which we refer to as $\mathcal{L}_{\text{ent}}$ (\eqn{eq:ent_loss}).

\textbf{Minimizing Sampler Energy.} To minimize the energy of samples from $q_{\theta}({\bm{x}})$ we can directly differentiate through both the energy function and MCMC sampling. We follow recent work in EBMs and utilize Langevin dynamics \citep{du2019implicit, nijkamp2019learning, grathwohl2019your} for our MCMC transition kernel, and note that each step of Langevin sampling is fully differentiable with respect to underlying energy function parameters. Precisely, gradient of  $\mathcal{L}_{\text{opt}}$,  $\frac{ \partial \mathcal{L}_{\text{opt}}}{ \partial \theta}$ becomes
\begin{equation}
    \label{eq:opt_loss}
    \resizebox{0.9\hsize}{!}{$
     \mathbb{E}_{q_{\theta}({\bm{x}}_0', \ldots, {\bm{x}}_t')} \left [\frac{\partial E_{\text{stop\_grad}(\theta)} ({\bm{x}}_{t-1}' - \nabla_{{\bm{x}}_{t-1}'} E_{\theta} ({\bm{x}}_{t-1}') + \omega) }{\partial \theta} \right],
    $}
\end{equation}

where $\omega \sim \mathcal{N}(0,\lambda)$ and ${\bm{x}}_i'$ represents the $\text{i}^{\text{th}}$ step of Langevin sampling. To reduce the memory overhead of this differentiation procedure, we only differentiate through the last step of Langevin sampling as also done in \citep{pmlr-v119-vahdat20a}. In the appendix we show that this leads to the same effect as differentiation through Langevin sampling.

\textbf{Entropy Estimation.} To maximize the entropy of samples from $q_{\theta}({\bm{x}})$, we use a non-parametric nearest neighbor entropy estimator \citep{beirlant1997nonparametric}, which is shown to be mean square consistent \citep{kozachenko1987sample} with root-n convergence rate \citep{tsybakov1996root}.  The entropy $H$ of a distribution $p({\bm{x}})$ can be estimated through a set $X = {x_1, x_2, \ldots, x_n}$ of $n$ different points sampled from $p({\bm{x}})$ as $H(p_{\theta}({\bm{x}})) = \frac{1}{n} \sum_{i=1}^n \ln (n \cdot \operatorname{NN}(x_i, X)) + O(1)$ where the function $\operatorname{NN}(x_i, X)$ denotes the nearest neighbor distance of $x_i$ to any other data point in $X$. Based off the above entropy estimator, we write $\mathcal{L}_{\text{ent}}$ as the entropy loss:

\begin{equation}
    \mathcal{L}_{\text{ent}} = \mathbb{E}_{q({\bm{x}})}[-\log ( \operatorname{NN}({\bm{x}}, B))]
      \label{eq:ent_loss}
\end{equation}
where we measure the nearest neighbor with respect to a set $B$ of 100 past samples from MCMC chains. We utilize L2 distance as the metric for computing nearest neighbors. This type of nearest entropy estimator is known to scale poorly to high dimensions (requiring an exponential number of samples to yield an accurate entropy estimate). However, in our setting, we do not need an accurate estimate of entropy. Instead, our computation of entropy is utilized as a fast regularizer to prevent sampling from collapsing.

\subsection{Data Augmentation Transitions}
Langevin sampling, our MCMC transition kernel, is prone to falling into local probability modes \citep{neal2011mcmc}. In the image domain, this manifests with sampling chains always converging to a fixed image \citep{du2019implicit}.  A core difficulty is that distances between two qualitatively similar images can be significantly far away from each other in the input domain, on which sampling is applied. While $\mathcal{L}_{\text{KL}}$ serves as a regularizer to prevent sampling collapse in Langevin dynamics,   Langevin dynamics alone is not enough to encourage large jumps in a finite number of steps. It is further beneficial to have an individual sampling chain have the ability to mix between probability modes.

 \input{tables/training_alg}
\input{figText/multi-scale}
To encourage greater exploration between similar inputs in our model, we propose to augment chains of MCMC sampling with periodic data augmentation transitions that encourages movement between ``similar'' inputs. In particular, we utilize a combination of color, horizontal flip, rescaling, and Gaussian blur augmentations. Such combinations of augmentation has recently seen success applied in unsupervised learning \citep{chen2020simple}. Specifically, during training time,  we initialize MCMC sampling from a data augmentation applied to an input sampled from the buffer of past samples. At test time, during the generation, we apply a random augmentation to the input after every 20 steps of Langevin sampling. We illustrate this process in the bottom of \fig{fig:pipeline_img}. Data augmentation transitions are always taken.

\subsection{Compositional Multi-scale Generation}

To encourage energy functions to focus on features in both low and high resolutions, we define our energy function as the composition (sum) of a set of energy functions operating on different scales of an image, illustrated in \fig{fig:multi-scale}. Since the downsampling operation is fully differentiable, Langevin based sampling can be directly applied to the energy function. In our experiments, we utilize full, half, and quarter resolution image as input and show that this improves the generation performance.

\input{tables/sampling_alg}

\subsection{Training Algorithm and Sampling}
We provide an overview of our overall proposed training algorithm in Algorithm \ref{alg:energy}. Our overall approach is  similar to the algorithm presented in \citep{du2019implicit}, with two notable differences. First, we apply data augmentation to samples drawn for the replay buffer. Second, we propagate gradients through sampling to  efficiently compute $\mathcal{L}_{\text{KL}}$.  We further present the sample generation algorithm for a trained EBMs in Algorithm \ref{alg:sample}. We iteratively apply $N$ steps of data augmentation and Langevin sampling to mimic the replay buffer utilized during training.

\section{Conclusion}
\vspace{-5pt}
We propose a simple and general framework for improving generation and ease of training of EBMs. We show that the framework enables high resolution compositional image generation and out-of-distribution robustness. In the future, we are interested in further computational scaling of our framework and its applications to additional domains such as text, video, and reasoning.

\section{Acknowledgements}

We would like to thank Jasha Sohl Dickstein, Bo Dai, Rif Sauros, Simon Osindero, Alex Alemi and anonymous reviewers for there helpful feedback on initial versions of the manuscript. Yilun Du is supported by a NSF GFRP fellowship. This work is in part supported by ONR MURI N00014-18-1-2846 and IBM Thomas J. Watson Research Center CW3031624.
\section{Related Work}
Our work is related to a large, growing body of work on different approaches for training EBMs. Our approach is based on contrastive divergence \citep{Hinton2002Training}, where an energy function is trained to contrast negative samples drawn from a model distribution and from real data. In recent years, such approaches have been applied to the image domain \citep{Salakhutdinov2009Deep, xie2016theory, gao2018learning, du2019implicit}. \citep{gao2018learning} also proposes a multiscale approach towards generating images from EBMs, but different from our work, uses each sub-scale EBM to initialize the generation of the next EBM, while we jointly sample across each resolution. Our work builds on existing works, and aims to provide improvements in generation and stability.

A difficulty with contrastive divergence training is negative sample generation. To sidestep this issue, a separate line of work utilizes an auxiliary network to amortize the negative portions of the sampling procedure \citep{kim2016deep, kumar2019maximum, han2019divergence, xie2018cooperative, song2018learning, dai2019exponential, che2020gan,  grathwohl2020vera, arbel2020ebm}. One line of work  \citep{kim2016deep, kumar2019maximum, song2018learning}, utilizes a separate generator network for negative image sample generations.  In contrast,  \citep{xie2018cooperative}, utilizes a generator to warm start generations for negative samples and \citep{han2019divergence} minimizes a divergence triangle between three models.  While such approaches enable better qualitative generation, they also lose some of the flexibility of the EBM formulation. For example, separate energy models can no longer be composed together for generation.

In addition, other approaches towards training EBMs seek instead to investigate separate objectives to train the EBM.  One such approach is score matching, where the gradients of an energy function are trained to match the gradients of real data \citep{hyvarinen2005estimation, song2019generative}, or with a related denoising  \citep{sohl2015deep, saremi2018deep, ho2020denoising} approach. Additional objectives include noise contrastive estimation \citep{gao2020flow}, learned Steins discrepancy \citep{grathwohl2020cutting}, and learned F divergences \citep{yu2020training}.

Most prior work in contrastive divergence has ignored the KL term \citep{hinton1999products, Salakhutdinov2009Deep}. A notable exception is \citep{ruiz2019contrastive}, which obtains a similar KL divergence term to ours. \citet{ruiz2019contrastive} use a high variance REINFORCE estimator to estimate the gradient of the KL term, while our approach relies on auto-differentiation and nearest neighbor entropy estimators. Differentiation through model generation procedures has previously been explored in other models \citep{Finn2017Meta, metz2016unrolled}. Other related entropy estimators include those based on Stein's identity \citep{liu2017stein} and MINE \citep{belghazi2018mine}. In contrast to these approaches, our entropy estimator relies only on nearest neighbor calculation, and does not require the training of an independent neural network.

\section{Introduction}

Energy-Based models (EBMs) have received an influx of interest recently and have been applied to realistic image generation \citep{han2019divergence, du2019implicit}, 3D shapes synthesis \citep{xie2018learning} , out of distribution and adversarial robustness \citep{lee2018wasserstein, du2019implicit, grathwohl2019your}, compositional generation \citep{hinton1999products,du2020compositional}, memory modeling \citep{bartunov2019meta}, text generation \citep{deng2020residual}, video generation \citep{xie2018video},   reinforcement learning \citep{haarnoja2017reinforcement,du2019model}, continual learning \citep{li2020energy}, protein design and folding \citep{ingraham2019learning,du2020energy} and biologically-plausible training \citep{scellier2017equilibrium}.
Contrastive divergence is a popular and elegant procedure for training EBMs proposed by \citep{Hinton2002Training}  which lowers the energy of the training data and raises the energy of the sampled confabulations generated by the model. The model confabulations are generated via an MCMC process (commonly Gibbs sampling or Langevin dynamics), leveraging the extensive body of research on sampling and stochastic optimization. The appeal of contrastive divergence is its simplicity and extensibility. It does not require training additional auxiliary networks \citep{kim2016deep, dai2019exponential} (which introduce additional tuning and balancing demands), and can be used to compose models zero-shot.

Despite these advantages, training EBMs with contrastive divergence has been challenging due to training instabilities. Ensuring training stability required either combinations of spectral normalization and Langevin dynamics gradient clipping \citep{du2019implicit}, parameter tuning \citep{grathwohl2019your}, early stopping of MCMC chains \citep{nijkamp2019learning}, or avoiding the use of modern deep learning components, such as self-attention or layer normalization \citep{du2019implicit}. These requirements limit modeling power, prevent the compatibility with modern deep learning architectures, and prevent long-running training procedures required for scaling to larger datasets. With this work, we aim to maintain the simplicity and advantages of contrastive divergence training, while resolving stability issues and incorporating complementary deep learning advances.

\input{figText/teaser.tex}
An often overlooked detail of contrastive divergence formulation is that changes to the energy function change the MCMC samples, which introduces an additional gradient term in the objective function (see \sect{sect:missing} for details). This term was claimed to be empirically negligible in the original formulation and is typically ignored \citep{Hinton2002Training,liu2017learning} or estimated via high-variance likelihood ratio approaches \citep{ruiz2019contrastive}. We show that this term can be efficiently estimated for continuous data via a combination of auto-differentiation and nearest-neighbor entropy estimators. We also empirically show that this term contributes significantly to the overall training gradient and has the effect of stabilizing training. It enables inclusion of self-attention blocks into network architectures, removes the need for capacity-limiting spectral normalization, and allows us to train the networks for longer periods. We do not introduce any new objectives or complexity - our procedure is simply a more complete form of the original formulation.

We further present techniques to improve mixing and mode exploration of MCMC transitions in contrastive divergence. We propose data augmentation as a useful tool to encourage mixing in MCMC by directly perturbing input images to related images. By incorporating data augmentation as semantically meaningful perturbations, we are able to greatly improve mixing and diversity of MCMC chains. We also leverage compositionality of EBMs to evaluate an image sample at multiple image resolutions when computing energies. Such evaluation and coarse and fine scales leads to samples with greater spatial coherence, but leaves MCMC generation process unchanged. We note that such hierarchy does not require specialized mechanisms such as progressive refinement \citep{Karras2017Progressive}

Our contributions are as follows: firstly, we show that a gradient term neglected in the popular contrastive divergence formulation is both tractable to estimate and is important in avoiding training instabilities that previously limited applicability and scalability of energy-based models. Secondly, we highlight how data augmentation and multi-scale processing can be used to improve model robustness and generation quality. Thirdly, we empirically evaluate stability of model architectures and show improved performance on a host of benchmarks and use cases, such as image generation, OOD detection, and compositional generation\footnote{Project page and code: https://energy-based-model.github.io/improved-contrastive-divergence/}.

\section{Experiments}
\label{sec:experiments}
We perform empirical experiments to validate the following set of questions: (1) What are the effects of each proposed component towards training EBMs? (2) Are our trained EBMs able to perform well on downstream applications of EBMs, such as image generation, out-of-distribution detection, and concept compositionality?

\subsection{Experimental Setup}
We investigate the efficacy of our proposed approach. Models are trained using the Adam Optimizer \citep{Kingma2015Adam}, on a single 32GB Volta GPU for CIFAR-10 for 1 day, and for 3 days on 8 32GB Volta GPUs for CelebaHQ, LSUN and ImageNet 32x32 datasets. We provide detailed training configuration details in the appendix.

Our improvements are largely built on top of the EBMs training framework proposed in \citep{du2019implicit}. We use a buffer size of 10000, with a resampling rate of 0.1\%. Our approach is significantly more stable than IGEBM, allowing us to remove aspects of regularization in \citep{du2019implicit}. We remove the clipping of gradients in Langevin sampling as well as spectral normalization on the weights of the network. In addition, we add self-attention blocks and layer normalization blocks in residual networks of our trained models. In multi-scale architectures, we utilize 3 different resolutions of an image, the original image resolution, half the image resolution and a quarter the image resolution. We report detailed architectures in the appendix. When evaluating models, we utilize the EMA model with EMA weight of 0.9999.

\input{figText/main_quant.tex}
\input{figText/sample_init}
\subsection{Image Generation}

\input{figText/diverse_sample}
\input{figText/mode_convergence_kl}
We evaluate our approach on  CIFAR-10, Imagenet 32x32 \citep{Deng2009Imagenet}, and CelebA-HQ \citep{Karras2017Progressive} datasets. Additional quantitative comparisons, results, and ablations can be found in the appendix of the paper.

\textbf{Image Quality.} We evaluate our approach on unconditional generation in \tbl{tbl:main_quant}. We utilize Inception \citep{Salimans2016Improved} and FID \citep{heusel2017gans} implementations from \citep{du2019implicit} to evaluate samples. On CIFAR-10, we find that our approach outperforms many past EBM approaches in both FID and Inception scores using a similar number of parameters. We find that our performance is slightly worse than that of SNGAN and SSGAN.  On CelebA-HQ, we find that our approach outperforms our reimplementation of SNGAN using default ImageNet hyperparameters, and is close to the reported numbers of SSGAN. Finally, on Imagenet32x32, we find that our approach outperforms previous EBM models and achieves performance comparable to that of the large PixelIQN model in terms of FID and Inception score. We note, however, that our model is significantly smaller than the PixelIQN model and has one tenth the number of parameters.  We present example qualitative images from CelebA-HQ in \fig{fig:celeba} and present qualitative images on other datasets in the appendix of the paper.
While our overall generative performance are not the best reported, we emphasize that it improves existing generative performance of EBMs, which have unique benefits such as compositionality (\sect{sect:compositionality}).

\vspace{-3pt}
\textbf{Effect of Data Augmentation.} We evaluate the effect of data augmentation on sampling in EBMs. In \fig{fig:sample_init} we show that by combining Langevin sampling with data augmentation transitions, we are able to enable chains to mix across different images, whereas prior works have shown Langevin converging to fixed images. In \fig{fig:diverse_sample} we show that given a fixed random noise initialization, data augmentation transitions enable to reach a diverse number of different samples, while  sampling without data augmentation transitions leads all chains to converge to the same face.

\vspace{-3pt}
\textbf{Mode Convergence.} We further investigate high likelihood modes of our model. In \fig{fig:mode_convergence}, we compare very low energy samples (obtained after running gradient descent 1000 steps on an energy function) for both our model with data augmentation and KL loss and the IGEBM model. Due to improved mode exploration, we find that low temperature samples under our model with data augmentation/KL loss reflect typical high likelihood "modes" in the training dataset, while our baseline models converges to odd shapes,  also noted in \citep{nijkamp2019anatomy}.

\vspace{-3pt}
\textbf{Stability/KL Loss.} EBMs are difficult to train and are sensitive to both the exact architecture and to various hyper-parameters. We found that the addition of a KL term into our training objective significantly improved the stability of training, by encouraging the sampling distribution to match the model distribution. In \fig{fig:kl_loss}, we measure training stability by measuring the energy differences between real and generated images. Stable training occurs when the energy difference is close to zero. Without $\mathcal{L}_{\text{KL}}$, we found that training our model with or without self-attention were both unstable, with differences spiking. Adding spectral normalization stabilizes training, but the addition of self-attention once again destabilizes training. In contrast with $\mathcal{L}_{\text{KL}}$, the addition of self-attention is also stable. We further compare Inception scores in \fig{fig:kl_loss} over training and find that while spectral normalization stabilizes training, it does at the expense of decreased improvement of Inception score.

The addition of the KL term itself is not too expensive, simply requiring an additional nearest neighbor computation during training, a relatively insignificant cost compared to the number of negative sampling steps used during training. With a intermediate number of negative sampling steps (60 steps) during training, adding the KL term incurs a roughly 20\% computational cost. This difference is further decreased with a larger number of sampling steps.

\vspace{-3pt}
\textbf{Ablations.} We ablate each portion of our proposed approach in \tbl{tbl:ablation}. We find that each our proposed components have significant gains in generation performance. In particular, we find a large gain in overall generative performance when adding the KL loss. This is in part due to a large boost in training stability (\fig{fig:kl_loss}), enabling significantly longer training times with both multiscale sampling and data augmentation.
\input{figText/ablations}
\textbf{KL Gradient.} We plot the overall gradient magnitudes of $\mathcal{L}_{\text{CD}}$ and $\mathcal{L}_{\text{KL}}$ when training an EBM on CIFAR-10 in \fig{fig:kl_cd}. We find that relative magnitude of gradients of both training objectives remains constant across training, and that the gradient of the KL objective is non-negligible.

\input{figText/kl_cd}

\input{figText/robustness}

\vspace{-3pt}
\subsection{Out of Distribution Robustness}
\vspace{-3pt}
Energy-Based Models (EBMs) have also been shown to exhibit robustness to both out-of-distribution and adversarial samples \citep{du2019implicit, grathwohl2019your, grathwohl2020vera}.
We evaluate out-of-distribution detection of our trained energy function through log-likelihood using the AUROC evaluation metrics proposed in \citet{hendrycks2016baseline}.
We similarly evaluate out-of-distribution detection of an unconditional CIFAR-10 model.

\textbf{Results.}
We present out-of-distribution results in \tbl{tbl:robustness}, comparing with both likelihood models and EBMs using log-likelihood to detect outliers. We find that our approach significantly outperforms other baselines, with the exception of CIFAR-10 interpolations. We note the JEM \citep{grathwohl2019your} further requires supervised labels to train the energy function, which has to shown to improve out-of-distribution performance. We posit that by more efficiently exploring modes of the energy distribution at training time, we are able to reduce the spurious modes of the energy function and thus improve out-of-distribution performance.

\input{figText/comp_table}
\input{figText/igebm_comp.tex}
\input{figText/compositionality}

\subsection{Compositionality}
\label{sect:compositionality}
Energy-Based Models (EBMs) have the ability to \textit{compose} with other models at generation time  \citep{hinton1999products, haarnoja2017reinforcement, du2020compositional}. We investigate to what extent EBMs trained under our new framework can also exhibit compositionality. See \citep{du2020compositional} for a discussion of various compositional operators and applications in EBMs.
In particular, we train \textit{independent} EBMs $E({\bm{x}}|c_1)$, $E({\bm{x}}|c_2)$, $E({\bm{x}}|c_3)$, that learn conditional generative distribution over concept factors $c$ such as facial expression. We test whether we can compose independent energy functions together to generate images with each concept factor simultaneously. We test compositions on the CelebA-HQ dataset, where we train separate EBMs on face attributes of age, gender, smiling, and wavy hair and on a rendered Blender dataset where we train separate EBMs on object attributes of size, position, rotation, and identity.

\vspace{-3pt}
\textbf{Qualitative Results.} We present qualitative results on  compositions of energy functions on CelebA-HQ in \fig{fig:comp_qual}. We consider compositional generation on the factors of young, female, smiling, and wavy hair. Compared to baselines, our approach is able to successfully generate images with each of conditioned factor, with faces being significantly higher resolution than baselines. In \fig{fig:qual_cube}, we further consider compositions of energy functions over object attributes. We consider compositional generation on factors of size, type, position, and rotation. Again, we find that each associated image generation exhibits the corresponding conditioned attribute. Images  are further visually consistent in terms of lighting, shadows and reflections. We note that different from most past work, generations of these combinations of different factors are \textit{only} specified at generation time, with models being trained \textit{independently}. Our results indicate that our framework for training EBMs is a promising direction for high resolution compositional visual generation.

\vspace{-3pt}
\textbf{Quantitative Comparison.} We quantitatively compare compositional generations of our model with IGEBM and JVAE \citep{Vedantam2018Generative} models  on the CelebA-HQ dataset. We assess the compositional generation accuracy of different models by measuring the accuracy with which a pretrained ResNet18 classifier can recover the underlying conditioned attributes. In \tbl{tbl:comp_table}, we find that our model has significantly higher attribute recovery compared to baselines across all compositions of attributes.
\section{More Image Results}
\label{sect:image}
\subsection{Nearest Neighbor Generations}

We present L2 nearest neighbors in CelebA-HQ training dataset of unconditional image samples from our trained EBM in \fig{fig:celeba_nn}. We find that our approach generates images distinct from the training set.

\subsection{Additional Quantitative Results}

We further quantitatively compare our generations with those of SNGAN on LSUN 128x128 bedroom scenes. We find that an SNGAN model trained on LSUN 128x128 bedroom scenes obtains an FID of 64.05 compared to our approach, which obtains an FID of 33.46. To report SNGAN scores, we re-implemented the SNGAN model using the default hyper parameters to train models on ImageNet 128x128.

\subsection{Additional Qualitative Images}

We present qualitative visualizations of unconditional samples generated from an EBM. \fig{fig:lsun_uncond} shows unconditional image generations from LSUN bedroom scenes. \fig{fig:cifar_uncond} shows unconditional image generations on the CIFAR-10 dataset. Finally, \fig{fig:imagenet_uncond} shows unconditional image generations on the ImageNet 32x32 dataset. In all three different settings, we find that our generated unconditional images are relatively globally coherent.

\section{Training Details}
\label{sect:method}
\subsection{Model Architectures}

In this part, we provide the model architectures used in  our experiments. When training multiscale energy functions, our final output energy function is the sum of energy functions applied to the full resolution image, the half resolution image, and the quarter resolution image. We use the architecture reported in \tbl{fig:cifar} for the full resolution image on CIFAR-10 and ImageNet 32x32 (used in the main paper Section 3.2 and 3.3). The model architecture used on the CelebA-HQ and LSUN datasets are reported in \tbl{fig:celeba_arch} (used in the main paper Section 3.2 and 3.4). The half-resolution models share the architecture listed in \tbl{fig:cifar}, but with the first down-sampled residual block removed. Similarly, the quarter resolution models share the architectures listed, but with the first two down-sampled residual blocks removed. We utilize group normalization \citep{Wu2018Group} inside each residual block and utilize the Swish nonlinearity \citep{ramachandran2018searching}.

\vspace{-10pt}
\subsection{Experiment Configurations For Different Datasets}

\myparagraph{CIFAR-10/ImageNet 32x32.} For CIFAR-10 and ImageNet 32x32, we use 40 steps of Langevin sampling to generate a negative sample. The Langevin sampling step size is set to be 500, with Gaussian noise of magnitude 0.001 at each iteration. The data augmentation transform consists of color augmentation of strength 1.0 from \citep{chen2020simple}, a random horizontal flip, and a image resize between 0.02 and 1.0. This is used in the main paper Section 3.2 and 3.3.

\myparagraph{CelebA/LSUN Bedroom.} For the CelebA-HQ and LSUN bed datasets, we use 40 steps of Langevin sampling to generate negative samples. The Langevin sampling step size is set to be 1000, with Gaussian noise of magnitude 0.001 applied at each iteration. The data augmentation transform consists of color augmentation of strength 1.0 from \citep{chen2020simple}, a random horizontal flip, and a image resize between 0.08 and 1.0. This is used in the main paper Section 3.2 and 3.4.

\section{Loss Gradient Derivation}
\label{sect:objective}

We show that the gradient of the contrastive divergence objective, $\mathcal{L}_{\text{CD\_Full}}$ is equivalent to that of the $\mathcal{L}_{\text{Full}} = \mathcal{L}_{\text{KL}} + \mathcal{L}_{\text{CD}}$. Recall that the contrastive divergence objective is given by
\begin{equation}
    \mathcal{L}_{\text{CD\_Full}} = \kldiv{p_D({\bm{x}})}{p_{\theta} ({\bm{x}})} - \kldiv{q_{\theta}({\bm{x}})}{p_{\theta}({\bm{x}})}.
    \label{eq:cd_sup}
\end{equation}
The gradient of the first KL term with respect to $\theta$, $\frac{\partial \kldiv{p_D({\bm{x}})}{p_{\theta} ({\bm{x}})}}{\partial \theta}$  is
\begin{equation}
   \resizebox{0.3\hsize}{!}{$
  - \mathbb{E}_{p_D({\bm{x}})}  \left [\frac{\partial E_{\theta}({\bm{x}})}{\partial \theta} \right]
  $}
\end{equation}
while the gradient of the second KL term with respect to $\theta$, $\frac{\kldiv{q_{\theta}({\bm{x}})}{p_{\theta}({\bm{x}})}}{\partial \theta}$
\begin{equation}
   \resizebox{0.8\hsize}{!}{$
  \frac{\partial q({\bm{x}}')}{\partial \theta} \frac{\partial \kldiv{q_{\theta}({\bm{x}}')}{p_{\theta}({\bm{x}}')}}{\partial q_{\theta}({\bm{x}}')}) - \mathbb{E}_{q_{\theta}({\bm{x}}')} [\frac{\partial E_{\theta}({\bm{x}}')}{\partial \theta}]
  $}
\end{equation}
with the overall gradient being
\begin{equation}
    \label{eq:cd_obj_grad_sup}
    \resizebox{0.8\hsize}{!}{$
    \begin{split}
    \frac{\mathcal{L}_{\text{CD\_Full}}}{\partial \theta} =
    - ( \mathbb{E}_{p_D({\bm{x}})} \left [\frac{\partial E_{\theta}({\bm{x}})}{\partial \theta} \right] &- \mathbb{E}_{q_{\theta}({\bm{x}}')} [\frac{\partial E_{\theta}({\bm{x}}')}{\partial \theta}] \\
    &+ {\color{red} \frac{\partial q({\bm{x}}')}{\partial \theta} \frac{\partial \kldiv{q_{\theta}({\bm{x}}')}{p_{\theta}({\bm{x}}')}}{\partial q_{\theta}({\bm{x}}')})}
    \end{split}$}.
\end{equation}
We have that
\begin{equation}
    \mathcal{L}_{\text{CD}} =  \mathbb{E}_{p_D({\bm{x}})} [E_{\theta}({\bm{x}})] - \mathbb{E}_{\text{stop\_grad}(q_{\theta}({\bm{x}}'))} [E_{\theta}({\bm{x}}')],
\end{equation}
with  corresponding gradients
\begin{equation}
    \resizebox{0.8\hsize}{!}{$
     \frac{\partial \mathcal{L}_{\text{CD}}}{\partial \theta} = \mathbb{E}_{p_D({\bm{x}})} \left [\frac{\partial E_{\theta}({\bm{x}})}{\partial \theta} \right] - \mathbb{E}_{q_{\theta}({\bm{x}}')} [\frac{\partial E_{\theta}({\bm{x}}')}{\partial \theta}]$}.
    \label{eq:cd_loss_grad_sup}
\end{equation}
Furthermore, we have that
\begin{equation}
    \mathcal{L}_{\text{KL}} =  \mathbb{E}_{q_{\theta}({\bm{x}})} [E_{\text{stop\_grad}(\theta)}({\bm{x}})] + \mathbb{E}_{q_{\theta}({\bm{x}})}[\log (q_{\theta}({\bm{x}}))] ,
\end{equation}
can be rewritten as
\begin{align}
    \mathcal{L}_{\text{KL}} &=  \mathbb{E}_{q_{\theta}({\bm{x}})} [-\log(p_{\theta}(x))] + \mathbb{E}_{q_{\theta}({\bm{x}})}[\log (q_{\theta}({\bm{x}}))]  \\
    &= \kldiv{q_\theta({\bm{x}})}{p_{\text{stop\_gradient}(\theta)}({\bm{x}})}.
\end{align}

The corresponding gradient of the objective is
\begin{equation}
     \frac{\partial \mathcal{L}_{\text{KL}}}{\partial \theta} =\frac{\partial q({\bm{x}})}{\partial \theta} \frac{\partial \kldiv{q_{\theta}({\bm{x}})}{p_{\theta}({\bm{x}})}}{\partial q_{\theta}({\bm{x}})}.
    \label{eq:kl_loss_grad_sup}
\end{equation}

Thus the sum of the gradients in $\frac{\partial \mathcal{L}_{\text{CD}}}{\partial \theta}$ (\eqn{eq:cd_loss_grad_sup}) and $\frac{\partial \mathcal{L}_{\text{KL}}}{\partial \theta}$ (\eqn{eq:kl_loss_grad_sup}) is equal to the full contrastive divergence gradient $\frac{\mathcal{L}_{\text{CD\_Full}}}{\partial \theta}$ (  \eqn{eq:cd_obj_grad_sup}).

\section{Additional Analysis}
\label{sect:analysis}

\subsection{Alternative Sampling Distributions} Instead of utilizing $q_\theta({\bm{x}})$ as $\Pi_\theta^t(p_D({\bm{x}}))$, as noted in the method section, our approach can further maximize likelihood as long as $\kldiv{p_D({\bm{x}})}{p_\theta({\bm{x}})}$ is greater $\kldiv{q_\theta({\bm{x}})}{p_\theta({\bm{x}})}$. We test an alternative sampler $q_\theta({\bm{x}})$ consisting of initializing Langevin dynamics from random noise in \fig{fig:noise}. We find again that our approach improves the training stability.

\begin{figure}
    \centering
    \includegraphics[width=\linewidth]{fig/rebuttal_fig.pdf}
    \vspace{-10pt}
    \caption{\small Inception Score and energy difference plots when $\mathcal{L}_{\text{KL}}$ is applied to MCMC initialized from random noise.}
    \label{fig:noise}
    \vspace{-15pt}
\end{figure}

\subsection{Analysis of Truncated Langevin Backpropagation}
To better understand the training effect of $\mathcal{L}_{\text{KL}}$, we analyze the effect of truncating backpropogation through Langevin sampling.  We train two separate models on MNIST, one with backpropogation through all Langevin steps, and one with backpropogation through only the last Langevin step. We obtain an FID of 90.54 with backpropogation through only 1 step of Langevin sampling and an FID of 94.85 with backpropogation through all steps of Langevin sampling. We present illustrations of samples generated with one step in \fig{fig:mnist_one} and with all steps in \fig{fig:mnist_all}. Overall, we find little degradation in performance with the truncation of backpropogation, but note that backpropogation through all steps of sampling is over 3 times slower to train.

\subsection{Analysis of Effect of KL Loss on Mode Sampling}

We illustrate the effect of $\mathcal{L}_{\text{KL}}$ as a regularizer to prevent EBM sampling collapse. When training an EBM, $\mathcal{L}_{\text{KL}}$ serves as a repelling term encouraging MCMC samples from an EBM to both have low energy and exhibit diversity. In the absence of $\mathcal{L}_{\text{KL}}$, we find that EBM sampling always collapses and eventually always generates samples illustrated in \fig{fig:collapse}. These samples are significantly less diverse than those generated when training with $\mathcal{L}_{\text{KL}}$ ( \fig{fig:cifar}), which never suffers from sampling collapse.

\begin{figure}[h]
    \centering
    \includegraphics[width=1\linewidth]{fig/celeba_nn2.pdf}
    \vspace{-5pt}
    \caption{\small Nearest neighbors in the L2 space of generated images in CelebA-HQ 128x128.}
    \label{fig:celeba_nn}
\end{figure}

\input{figText/lsun_uncond}
\input{figText/cifar10_uncond}
\input{figText/imagenet32x32_uncond}

\begin{figure}

\begin{minipage}{0.45\linewidth}
\captionof{table}{The model architecture used for CIFAR-10 and ImageNet-32x32 experiments.}
\centering
\begin{tabular}{c}
    \toprule
    \toprule
    3x3 conv2d, 64 \\
    \midrule
    ResBlock  64 \\
    \midrule
    ResBlock Down 64 \\
    \midrule
    ResBlock  64 \\
    \midrule
    ResBlock  Down 64 \\
    \midrule
    Self Attention 64 \\
    \midrule
    ResBlock  128 \\
    \midrule
    ResBlock  Down 128 \\
    \midrule
    ResBlock  256 \\
    \midrule
    ResBlock  Down 256 \\
    \midrule
    Global Mean Pooling \\
    \midrule
    Dense $\rightarrow$ 1 \\
    \bottomrule
\end{tabular}
\label{fig:cifar}
\end{minipage}%
\hfill
\begin{minipage}{0.45\linewidth}
\vspace{-19pt}
\captionof{table}{The model architecture used for CelebA-HQ/LSUN room experiments.}
\centering
\begin{tabular}{c}
    \toprule
    \toprule
    3x3 conv2d, 64 \\
    \midrule
    ResBlock Down 64 \\
    \midrule
    ResBlock Down 128 \\
    \midrule
    ResBlock Down 128 \\
    \midrule
    ResBlock 256 \\
    \midrule
    ResBlock Down 256 \\
    \midrule
    Self Attention 512 \\
    \midrule
    ResBlock  512 \\
    \midrule
    ResBlock  Down 512 \\
    \midrule
    Global mean Pooling \\
    \midrule
    Dense $\rightarrow$ 1 \\
    \bottomrule
\end{tabular}
\label{fig:celeba_arch}
\end{minipage}

\end{figure}

\begin{figure}
    \centering
    \includegraphics[width=1\linewidth]{fig/mnist_one.png}
    \caption{\small Generations on MNIST with backpropogation through 1 step of Langevin sampling. }
    \label{fig:mnist_one}
\end{figure}

\begin{figure}
    \centering
    \includegraphics[width=1\linewidth]{fig/mnist_all.png}
    \caption{\small Generations on MNIST with backpropogation through all steps of Langevin sampling.}
    \label{fig:mnist_all}
\end{figure}

\begin{figure}
    \centering
    \includegraphics[width=\linewidth]{fig/collapse.png}
    \vspace{-10pt}
    \caption{\small Illustration of collapsed sampling from an EBM. Sampling does not collapse with the addition of the KL loss.}
    \label{fig:collapse}
\end{figure}