# JEM++: Improved Techniques for Training JEM

Xiulong Yang, Shihao Ji
Department of Computer Science
Georgia State University
<EMAIL_ADDRESS>
###### Abstract

Joint Energy-based Model (JEM) [17] is a recently proposed hybrid model that
retains strong discriminative power of modern CNN classifiers, while
generating samples rivaling the quality of GAN-based approaches. In this
paper, we propose a variety of new training procedures and architecture
features to improve JEM’s accuracy, training stability, and speed altogether.
1) We propose a proximal SGLD to generate samples in the proximity of samples
from previous step, which improves the stability. 2) We further treat the
approximate maximum likelihood learning of EBM as a multi-step differential
game, and extend the YOPO framework [59] to cut out redundant calculations
during backpropagation, which accelerates the training substantially. 3)
Rather than initializing SGLD chain from random noise, we introduce a new
informative initialization that samples from a distribution estimated from
training data. 4) This informative initialization allows us to enable batch
normalization in JEM, which further releases the power of modern CNN
architectures for hybrid modeling.111Code: https://github.com/sndnyang/JEMPP

##  1 Introduction

Deep neural networks (DNNs) have made significant breakthroughs in various
discriminative tasks and generative tasks, including image classification,
object detection, and high-quality image and text generation [34, 25, 4, 6].
However, prior works on discriminative models and generative models are
largely separated. Even though a few researches (e.g., [7, 10]) have shown
that generative training is beneficial to discriminative models, most recent
works on generative models focus primarily on qualitative sample quality [5,
48, 52], and the discriminative performances of state-of-the-art generative
models are still far behind discriminative ones [3, 8, 11].

Among different discriminative and generative models, energy-based models
(EBMs) [35] are an appealing class of probabilistic models, which can be
viewed as hybrid models with both discriminative and generative powers [17].
Compared to the popular generative models, such as VAE [33] and GAN [14],
which train explicit functions to generate samples, EBMs only need to train a
single network with a set of shared features for discriminative tasks and
generative tasks, and exploit implicit sampling for generation. Since an EBM
is the only object that needs to be trained, it generally achieves a higher
simplicity and stability than approaches that use multiple networks. Hence,
there is a great interest recently in encompassing the generative capabilities
into discriminative models without sacrificing their discriminative powers.
Specifically, a series of recent works propose to train a CNN as an EBM for
image classification and generation [57, 23, 11, 17]. Among them, JEM [17] is
one of the most representative ones, which reinterprets the modern CNN
classifier (e.g., Wide-ResNet [58]) as an EBM for image generation and
achieves impressive performances in image classification and image generation
simultaneously. JEM demonstrates the potential of EBMs in hybrid modeling and
ignites a series of follow-up works [60, 18, 13, 19].

However, training EBMs is still a challenging task. As shown in Table 1,
existing methods demonstrate a great deal of tradeoffs among different
algorithmic features in the quest of improved training algorithms. Most of the
works [41, 11, 17] adopt the SGLD sampling [54] to train EBMs, where $K$
sweeps of forward and backward propagations are required in each sampling
step. These training methods can be prolonged with a large $K$, preventing
them from long training procedures required by large-scale datasets. In
addition, SGLD can be precarious and easily diverged, which further hinders
the prevalence of EBMs. To avoid the long sampling process of SGLD, recent
works introduce auxiliary models [24, 56, 19] or use special architectures
[18, 53] to amortize the SGLD sampling or improve its stability. Given the
architectural simplicity of the SGLD-based methods, especially JEM [17], we
ask the following question: Is it possible to develop new training methods of
JEM to reduce the number of sampling steps required by SGLD while improving
its training stability?

Table 1: Characteristics of different EBM training methods. Training Method | Fast | Stable | High dim | No aux. model | Unrestricted arch | Approx. likelihood
---|---|---|---|---|---|---
SGLD-based [41, 11, 17] | ✗ | ✗ | ✓ | ✓ | ✓ | ✓
Score Matching [53, 29] | ✓ | ✗ | ✓ | ✓ | ✗ | ✗
Noise Contrastive [13, 22] | ✓ | ✓ | ✗ | ✗ | ✓ | ✗
Regularized Generator [24, 19] | ✓ | ✓ | ✓ | ✗ | ✓ | ✓
JEM++ (ours) | $\uparrow$ | $\uparrow$ | ✓ | ✓ | ✓ | ✓

In this paper, we introduce a variety of training procedures and architecture
features to improve JEM’s accuracy, training stability, and speed altogether.
After a thorough investigation on JEM, we find that JEM sometimes generates
abnormal images containing pixels with extreme values beyond a reasonable
range. This motivates us to constrain the SGLD sampling by projecting samples
to an $L_{p}$-norm ball of previous samples. Secondly, JEM does not support
modern architecture features such as batch norm [30]222Although the authors
stated they have been able to train JEM with batch norm successfully, no
details are disclosed in their paper or code.. We find that a huge statistic
gap between the initial noisy samples of SGLD and real data incurs the
training difficulty of JEM when batch norm is enabled. Hence, we introduce a
new informative initialization that closes the gap between initial samples and
real data. Moreover, we find that batch-norm-enabled JEM supports a larger
learning rate, which further increases the convergence rate of JEM. Finally,
we extend YOPO [59], a general framework for PGD [36] acceleration, to the
maximum likelihood learning of EBM and speed up the training of JEM even
further. Our main contributions are summarized as follows:

  1. 1.

We propose a proximal SGLD to generate samples in the proximity of samples
from previous step, which improves the stability of JEM.

  2. 2.

We further treat the approximate maximum likelihood learning of EBM as a
multi-step differential game, which can be accelerated by cutting out
redundant calculations during backpropagation, while retaining the overall
predictive performance.

  3. 3.

We introduce a new informative initialization to initialize the SGLD chain,
which stabilizes the training further and accelerates the convergence rate of
SGLD sampling.

  4. 4.

This new informative initialization also enables batch norm to train JEM
successfully and release the power of modern CNN architectures. What’s more,
with the informative initialization and batch norm, JEM++ can be optimized
with a large learning rate, while JEM fails to.

  5. 5.

JEM++ matches or outperforms prior state-of-the-art hybrid models on
discriminative and generative tasks, while enjoying improved stability and
training speed over the original JEM.

##  2 Energy-Based Models

Energy-based models (EBMs) [35] define an energy function that assigns low
energy values to samples drawn from data distribution and high values
otherwise, such that any probability density $p_{\bm{\theta}}(\bm{x})$ can be
expressed via a Boltzmann distribution as

$p_{\bm{\theta}}(\bm{x})=\frac{\exp\left(-E_{\bm{\theta}}(\bm{x})\right)}{Z(\bm{\theta})},$
(1)

where $E_{\bm{\theta}}(\bm{x})$ is an energy function that maps each input
$\bm{x}\in\mathcal{X}$ to a scalar, and $Z(\bm{\theta})$ is the normalizing
constant (also known as the partition function) such that
$p_{\bm{\theta}}(\bm{x})$ is a valid density function.

The key challenge of training EBMs lies in estimating the partition function
$Z(\bm{\theta})$, which is notoriously intractable. The standard maximum
likelihood estimation of parameters $\bm{\theta}$ is not straightforward
either, and a number of sampling-based approaches have been proposed to
approximate it effectively. Specifically, the derivative of the log-likelihood
of a single sample $\bm{x}\in\mathcal{X}$ w.r.t. $\bm{\theta}$ can be
expressed as

$\displaystyle\frac{\partial\log
p_{\bm{\theta}}(\bm{x})}{\partial\bm{\theta}}=\mathbb{E}_{p_{\bm{\theta}}(\bm{x}^{\prime})}\frac{\partial
E_{\bm{\theta}}(\bm{x}^{\prime})}{\partial\bm{\theta}}-\frac{\partial
E_{\bm{\theta}}(\bm{x})}{\partial\bm{\theta}},$ (2)

where the expectation is over the density function
$p_{\bm{\theta}}(\bm{x}^{\prime})$, sampling from which is challenging due to
the intractable $Z(\bm{\theta})$. Therefore, MCMC and Gibbs sampling [28] have
been proposed previously to estimate the expectation efficiently. To speed up
the mixing for effective sampling, recently Stochastic Gradient Langevin
Dynamics (SGLD) [54] has been employed to train EBMs by using the gradient
information [41, 11, 17]. Specifically, to sample from
$p_{\bm{\theta}}(\bm{x})$, SGLD follows

$\displaystyle\bm{x}^{0}\sim p_{0}(\bm{x}),$
$\displaystyle\bm{x}^{t+1}=\bm{x}^{t}-\frac{\alpha}{2}\frac{\partial
E_{\bm{\theta}}(\bm{x}^{t})}{\partial\bm{x}^{t}}+\alpha\epsilon^{t},\;\;\epsilon^{t}\sim\mathcal{N}(0,1),$
(3)

where $p_{0}(\bm{x})$ is typically a uniform distribution over $[-1,1]$, whose
samples are refined via a noisy gradient decent with step-size $\alpha$ over a
SGLD chain.

Prior works [11, 17, 40] have investigated the effect of hyper-parameters in
SGLD sampling in terms of stability and speed, and showed that the SGLD-based
approaches suffer from poor stability and computational challenges from
sequential sampling at every iteration. Specifically, Nijkamp et al. [40] find
that the noise term in SGLD is not important, and including a noise of low
variance appears to improve synthesis quality. What’s more, for unnormalized
densities, it’s desirable to generate samples from SGLD chain after it
converges. This requires the step-size $\alpha$ to decay with a polynomial
schedule and an infinite number of sampling steps, which is not realistic in
practical applications. Instead, JEM [17] uses a constant step-size $\alpha$
during sampling and approximates the samples with a sampler that runs only for
a finite number of steps. To improve the sampling stability, the model would
require to quadruple the number of SGLD steps, which greatly increases the
run-time.

##  3 JEM++: The Improved Training of JEM

We first give a brief introduction of JEM [17] and then discuss a variety of
new training procedures to improve its accuracy, stability and speed.

Joint Energy-based Models (JEM) [17] reinterprets modern CNN classifiers as
EBMs. Considering a CNN classifier of parameters $\bm{\theta}$, given an input
$\bm{x}$ the classifier first maps the input to a vector of $C$ real-valued
numbers (or logits): $f_{\bm{\theta}}(\bm{x})[y],\forall y\in[1,\cdots,C]$,
where $C$ is the number of classes; the logits are then normalized via the
softmax function to yield a probability vector:
$p_{\bm{\theta}}(y|\bm{x})=e^{f_{\bm{\theta}}(\bm{x})[y]}/\sum_{y^{\prime}}e^{f_{\bm{\theta}}(\bm{x})\left[y^{\prime}\right]}$.
Interestingly, the same vector of logits $f_{\bm{\theta}}(\bm{x})[y]$ can also
be used to define an EBM for the joint density:
$p_{\bm{\theta}}(\bm{x},y)=e^{f_{\bm{\theta}}(\bm{x})[y]}/Z(\bm{\theta})$,
where $Z(\bm{\theta})$ is an unknown normalizing constant (regardless of
$\bm{x}$ or $y$). Then a marginal density of $\bm{x}$ can be achieved by
marginalizing the joint density as:
$p_{\bm{\theta}}(\bm{x})=\sum_{y}p_{\bm{\theta}}(\bm{x},y)=\sum_{y}e^{f_{\bm{\theta}}(\bm{x})\left[y\right]}/Z(\bm{\theta})$.
Comparing this density with Eq. 1, it is readily to show that the
corresponding energy function of $\bm{x}$ is defined as

$E_{\bm{\theta}}(\bm{x})\\!=\\!-\log\\!\sum_{y}\\!e^{f_{\bm{\theta}}(\bm{x})\left[y\right]}\\!=\\!-\text{LSE}(f_{\bm{\theta}}(\bm{x})),$
(4)

where $\text{LSE}(\cdot)$ denotes the Log-Sum-Exp function.

To optimize the model parameter $\bm{\theta}$, JEM proposes to maximize the
joint density function $p_{\bm{\theta}}(\bm{x},y)$, which can be factorized
as:

$\log p_{\bm{\theta}}(\bm{x},y)=\log p_{\bm{\theta}}(y|\bm{x})+\log
p_{\bm{\theta}}(\bm{x}),$ (5)

where the first term is the conventional cross-entropy objective for
classification, and the second term can be optimized by the maximum likelihood
learning of EBM as shown in Eq. 2 with the SGLD sampling defined in (2). In
the paper, we follow the same objective function of JEM and focus on how to
improve the stability of SGLD sampling as well as accelerate the maximum
likelihood learning of EBM.

###  3.1 Training EBM as a Minimax Optimization

In practice, when we employ the maximum likelihood estimate of model
parameters $\bm{\theta}$ with Eq. 2, a minibatch of $B$ samples
$\\{\bm{x}_{1},\bm{x}_{2},\cdots,\bm{x}_{B}\\}\sim p_{\bm{\theta}}(\bm{x})$
and a minibatch of $B$ real data samples
$\\{\bm{x}^{r}_{1},\bm{x}^{r}_{2},\cdots,\bm{x}^{r}_{B}\\}\sim\mathcal{X}$ are
used. To avoid notational clutter, we assume $B=1$ in the rest of the paper,
but the results are readily extended to $B>1$.

Similar to Nijkamp et al. [40] who have found the insignificance of noise term
in the SGLD sampling (2), our empirical study also confirms this observation.
Thus, we ignore the noise term in Eq. 2 and treat it as an artifact that
generates some stochasticity in the sampling process to facilitate the
optimization. Under this assumption, the SGLD sampling (2) can be
reinterpreted approximately as an SGD iteration, with a learning rate of
$\alpha/2$, initialized from a random sample of $p_{0}(\bm{x})$. Assume the
convergence can be achieved, the objective of the SGLD sampling (2) is to
solve the following optimization problem approximately 333The entire pipeline
is still a stochastic sampler since samples are generated by running finite-
length stochastic gradient decent with random initialization.

$\displaystyle\bm{x}^{*}=\operatorname*{argmin}_{\bm{x}}E_{\bm{\theta}}(\bm{x}).$
(6)

Therefore, the maximum likelihood learning of EBM with Eq. 2 is to
approximately solve the following minimax game

$\displaystyle\max_{\bm{\theta}}\left[\min_{\bm{x}}E_{\bm{\theta}}(\bm{x})-E_{\bm{\theta}}(\bm{x}^{r})\right].$
(7)

To have a robust convergence behavior, we can solve the inner minimization
problem of (7) by using the Proximal Point Method [42]. We can further treat
the minimax optimization problem (7) as a multi-step differential game and
extend YOPO [59], a general framework of accelerating PGD, to speed up the
training of EBM. Next, we describe these new training procedures in details.

###  3.2 Proximal SGLD

Prior works on EBMs reveal the tradeoff between training stability and
computational time of SGLD-based approaches [41, 11, 17]. However, the cause
of instability of SGLD-based EBMs is still under investigation. Empirically,
we observe that upon the divergence of EBM, SGLD generates abnormal samples
with extreme values that have a severe negative impact on model parameter
update. Hence, we introduce our first improvement to stabilize the inner
minimization problem with a proximal SGLD.

Proximal point methods are widely used in optimization [46, 42]. To solve the
inner minimization problem of (7), the algorithm generates a sequence
$\\{\bm{x}^{t}\\}_{t=1,2,\cdots}$ by the following proximal point iteration:

$\displaystyle\bm{x}^{t+1}=\operatorname*{argmin}_{\bm{x}}E_{\bm{\theta}}(\bm{x})\;s.t.\;||\bm{x}-\bm{x}^{t}||_{p}<\varepsilon,$
(8)

which solves a constrained minimization problem at each iteration $t$, i.e.,
the current solution should be in the proximity of previous one, measured by
an $L_{p}$ norm. Compared with the standard SGD iteration, the proximal point
iteration has a robust convergence behavior. Moreover, even if the proximal
operator defined in Eq. 8 is not exactly minimized in each iteration, it still
has a stronger convergence guarantee than standard SGD, giving rise to the
inexact proximal point method [46]. Thus, if we solve each minimization
problem (8) inexactly with one step of SGD, we obtain an inexact proximal
point iteration

$\displaystyle\bm{x}^{t+1}=\bm{x}^{t}-\frac{\alpha}{2}L_{p}(\nabla_{\bm{x}}E_{\bm{\theta}}(\bm{x}^{t}),\varepsilon),$
(9)

where $L_{p}(\cdot,\varepsilon)$ projects the gradient to an $L_{p}$-norm ball
of a radius $\varepsilon$. Empirically, we find the $L_{\infty}$-norm works
well across different architectures and datasets. Hence, we only consider the
$L_{\infty}$-norm in the rest of the paper. With an $L_{\infty}$-norm, Eq. 9
can be rewritten as

$\displaystyle\bm{x}^{t+1}=\bm{x}^{t}-\frac{\alpha}{2}\text{clamp}(\nabla_{\bm{x}}E_{\bm{\theta}}(\bm{x}^{t}),\varepsilon)+\alpha\epsilon^{t},$
(10)

where the $\text{clamp}(\cdot,\varepsilon)$ operator clamps the gradient in
the range of $[-\varepsilon,\varepsilon]$. Note that to incorporate
stochasticity into the inexact proximal point iteration, we add the noise term
back to Eq. 10, which resembles the original SGLD sampling (2) but with a
gradient clamping operator used to enforce the proximity constraint.

###  3.3 Training EBM as a Differential Game

As discussed in Section 3.1, the maximum likelihood learning of EBM (7) solves
a minimax game approximately. This objective has a close relationship to
adversarial training with the PGD attack [36]. Hence, we can extend methods
for accelerating adversarial training to EBM and reduce the computational
complexity of multi-step SGLD.

Inspired by Pontryagin’s Maximum Principle [43], a general framework in
optimal control, Zhang et al. [59] propose an optimization method called YOPO
(You-Propogate Only Once) to accelerate multi-step adversarial training such
as PGD. The key factor in YOPO is that the adversarial perturbation is only
coupled with the first layer’s weights in a neural network. Then YOPO can
decouple the adversary update from training of network parameters, and reduce
the total number of full forward and backward propagations to only one in each
group of adversary updates.

Similarly, we can extend YOPO to the maximum likelihood learning of EBM
because the objective (7) can also be treated as a multi-step differential
game and the sampled image $\bm{x}$ from proximal SGLD (10) is only coupled
with the first layer’s weights. By inserting the energy function (4) into (7),
we can rewrite the minimax objective as:

$\displaystyle\max_{\bm{\theta}}$
$\displaystyle\left[\min_{\bm{x}}-\text{LSE}\left(g_{\tilde{\bm{\theta}}}\left(f_{0}\left(\bm{x},\bm{\theta}_{0}\right)\right)\right)-E_{\bm{\theta}}(x^{r})\right]$
(11)

where $f_{0}$ denotes the first layer of a CNN-based EBM,
$g_{\tilde{\bm{\theta}}}=f_{T-1}^{\bm{\theta}_{T-1}}\circ
f_{T-2}^{\bm{\theta}_{T-2}}\circ\cdots f_{1}^{\bm{\theta}_{1}}$ denotes the
network without the first layer, such that
$f_{\bm{\theta}}(\bm{x})=g_{\tilde{\bm{\theta}}}\left(f_{0}\left(\bm{x},\bm{\theta}_{0}\right)\right)$.
Given a sample $\bm{x}$, the gradient of energy function (4) can be calculated
by chain rule as:

$\displaystyle\frac{\partial E_{\bm{\theta}}(\bm{x})}{\partial\bm{x}}=$
$\displaystyle-\nabla_{g_{\tilde{\bm{\theta}}}}\text{LSE}\left(g_{\tilde{\bm{\theta}}}\left(f_{0}\left(\bm{x},\bm{\theta}_{0}\right)\right)\right)$
$\displaystyle\cdot\nabla_{f_{0}}g_{\tilde{\bm{\theta}}}(f_{0}(\bm{x},\bm{\theta}_{0}))\cdot\nabla_{\bm{x}}f_{0}(\bm{x},\bm{\theta}_{0}).$
(12)

Figure 1: Comparison between SGLD-$K$ sampling and PYLD-$M$-$N$ sampling.

Proximal SGLD (10) conducts $K$ sweeps of full forward and backward
propagations for each update of $\bm{\theta}$. To stabilize the training of
EBM, it requires a large $K$, which greatly increases the run-time. To reduce
the total number of thorough forward and backward propagations, we follow YOPO
and introduce a slack variable:

$\displaystyle
p\\!=\\!-\nabla_{g_{\tilde{\bm{\theta}}}}\text{LSE}\\!\left(g_{\tilde{\bm{\theta}}}\left(f_{0}\\!\left(\bm{x},\bm{\theta}_{0}\right)\right)\right)\cdot\nabla_{f_{0}}g_{\tilde{\bm{\theta}}}\left(f_{0}\\!\left(\bm{x},\\!\bm{\theta}_{0}\right)\right),$
(13)

and freeze it as a constant in the inner loop of the sample update. We call
our accelerated Proximial SGLD algorithm PYLD-$M$-$N$ (Proximal-YOPO-SGLD)
with $M$ outer loops and $N$ inner loops. Figure 1 demonstrates a conceptual
comparison between SGLD-$K$ and PYLD-$M$-$N$. SGLD-$K$ accesses the data $K$
times requiring $K$ full forward and backward propagations. On the contrary,
PYLD-$M$-$N$ accesses the data $M\times N$ times, while only requiring $M$
full forward and backward propagations and a inner loop of $M\times N$ cheap
sample updates. Similar to YOPO [59], when $M\times N\approx K$, PYLD can
achieve a similar sample quality as SGLD. But PYLD-$M$-$N$ has the flexibility
of increasing $N$ and reducing $M$ to achieve approximately the same level of
movement with much less computation cost. We will demonstrate this when we
present results.

The pseudo-code of our PYLD is described in Algorithm 1. For more details of
YOPO, we refer the readers to [59].

Algorithm 1 PYLD-$M$-$N$ sampling: Given network $g_{\tilde{\bm{\theta}}}$ and
$f_{0}$ with $\bm{\theta}_{0}$, step-size $\alpha$, number of steps $M$ and
$N$

1: $\bm{x}^{0}\sim p_{0}(\bm{x})$

2: for $t\in[0,1,\cdots,M-1]$ do

3: % calculate the slack variable

4:
$p\\!=\\!-\nabla_{g_{\tilde{\bm{\theta}}}}\text{LSE}\\!\left(g_{\tilde{\bm{\theta}}}\\!\left(f_{0}\\!\left(\bm{x}^{t},\\!\bm{\theta}_{0}\right)\right)\right)\cdot\nabla_{f_{0}}g_{\tilde{\bm{\theta}}}\\!\left(f_{0}\\!\left(\bm{x}^{t},\\!\bm{\theta}_{0}\right)\right)$

5: $\bm{x}^{t,0}=\bm{x}^{t}$

6: for $s\in[0,1,\cdots,N-1]$ do

7:
$\gamma=\text{clamp}(p\cdot\nabla_{\bm{x}^{t,s}}f_{0}\\!\left(\bm{x}^{t,s},\bm{\theta}_{0}\right),\varepsilon)$

8: $\bm{x}^{t,s+1}=\bm{x}^{t,s}-\alpha/2\cdot\gamma$

9: end for

10: $\bm{x}^{t+1}=\bm{x}^{t,N}+\alpha\epsilon^{t}$

11: end for

12: return $\bm{x}^{M}$

###  3.4 Informative Initialization

The initial sampling distribution $p_{0}(\bm{x})$ also plays an important role
in the training of EBM. Nijkamp et al. [40] summarize two main types of SGLD
initializations for $\bm{x}^{0}$: non-informative initialization and
informative initialization. The former initializes the sample $\bm{x}^{0}$
from a noise distribution independent to the training data, such as a uniform
or Gaussian distribution, while the latter samples from an approximate
distribution close to the data distribution. One typical informative
initialization is to use samples from training data directly, as proposed in
Contrastive Divergence (CD) [28]. Based on this, Tieleman [51] proposes
Persistent Contrastive Divergence (PCD) and uses samples from previous
learning iteration as the initial samples for the current iteration. In
contrast to common wisdom, Nijkamp et al. [41] propose a short-run MCMC
sampler which always starts from the random noise distribution such as a
uniform distribution. Moreover, to train EBMs, Xie et al. [57] propose another
persistent initialization, which combines non-informative and informative
initialization and samples short SGLD chains from data samples of previous
iterations and occasionally (with a small probability $\rho$) reinitializes
the chains from random noise. This is also the sampling approach adopted by
IGEBM [11] and JEM [17], which maintain a replay buffer of samples from
previous iterations and replace a small percentage of samples in the buffer
with random noise to train EBMs.

In this paper, we explore informative initialization to initialize the SGLD
chain, and use the PCD with a replay buffer. The main difference is that we
substitute the random noise samples with samples from a Gaussian mixture
distribution estimated from the training dataset. That is, we define the
initial sampling distribution as

$\displaystyle p_{0}(\bm{x})$
$\displaystyle=\sum\nolimits_{y}\mathcal{\pi}_{y}{N}(\bm{\mu}_{y},\bm{\Sigma}_{y})$
(14) $\displaystyle\text{with}\quad\pi_{y}$
$\displaystyle=|\mathcal{D}_{y}|/\sum\nolimits_{y^{\prime}}|\mathcal{D}_{y^{\prime}}|,\quad\bm{\mu}_{y}=\mathbb{E}_{\bm{x}\sim\mathcal{D}_{y}}[\bm{x}],$
$\displaystyle\bm{\Sigma}_{y}$
$\displaystyle=\mathbb{E}_{\bm{x}\sim\mathcal{D}_{y}}\left[\left(\bm{x}-\bm{\mu}_{y}\right)\left(\bm{x}-\bm{\mu}_{y}\right)^{\top}\right],$

where $\mathcal{D}_{y}$ denotes the set of training samples with label $y$. As
an example, Figure 2 visualizes the
$\\{\bm{\mu}_{1},\bm{\mu}_{2},\cdots,\bm{\mu}_{10}\\}$ (categorical centers)
estimated from the CIFAR10 training dataset. Similar visualizations on
CIFAR100 and SVHN as well as example samples from the informative
initialization can be found in the supplementary material.

Figure 2: The categorical centers of CIFAR10.

The informative initialization brings sufficient information into $\bm{x}^{0}$
to guide the SGLD chain to converge faster than from a random noise since the
initial sample $\bm{x}^{0}$ is now much closer to the real data manifold.
Empirically, we also observe the improved training stability. What’s more, the
informative initialization allows us to enable batch norm [30], a modern
architecture feature of DNNs, that is excluded by IGEBM and JEM due to the
training difficulty introduced by batch norm.

###  3.5 Batch Normalization and Learning Rate

Batch norm [30] is an essential component in many state-of-the-art CNN
architectures. Batch norm normalizes input features by the mean and variance
computed within each mini-batch, which mitigates the vanishing gradient issue
of training very deep networks and dramatically improves the convergence rate
of gradient-based methods. Moreover, batch norm allows a much larger learning
rate and mitigates the need of tedious finetuning.

However, state-of-the-art EBMs, such as IGEBM [11] and JEM [17], do not
support batch norm. If batch norm is enabled in JEM, the model can neither
achieve a high classification accuracy nor generate realistic images. This is
because one intrinsic assumption of batch norm is that the input features
should come from a single or similar distributions. This normalization
behavior could be problematic if the mini-batch contains data from different
distributions, therefore resulting in inaccurate statistics estimation.
Unfortunately, this might be the case for the original IGEBM and JEM.
Apparently, if the initial samples $\bm{x}^{0}$ are sampled from a uniform or
Gaussian distribution as in IGEBM and JEM, $\bm{x}^{0}$ and real data samples
have different underlying distributions, violating the assumption of batch
norm.

Similar phenomenon has also been observed by Xie et al. [55] who demonstrates
the different statistics between clean data and adversarial examples. They
show that both clean data accuracy and adversarial robustness can be improved
by using two branches of batch norm: one main branch for clean data and one
auxiliary branch for adversarial examples. Instead of using two batch norms,
we mitigate the training difficulty of batch norm from a different
perspective. Since we have the choice of designing sampling distribution
$p_{0}(\bm{x})$, we can use the informative initialization discussed above to
enable batch norm in the EBM training. Since the Gaussian mixture distribution
(15) is actually estimated from real training examples, we can close the
statistic gap between initial samples of SGLD and real data and enable batch
norm in JEM++ successfully. What’s more, with the informative initialization
and batch norm, JEM++ can also use a much larger learning rate to improve
convergence rate even further.

In summary, Algorithm 2 provides the pseudo-code for JEM++ training, which
follows a similar design of JEM [17] and IGEBM [11] with a replay buffer. For
brevity, only one real sample $(\bm{x}^{r},y^{r})\sim\mathcal{D}$ and one
generated sample $\bm{x}^{M}\sim p_{\bm{\theta}}(\bm{x})$ are used to optimize
the parameter $\bm{\theta}$. It is straightforward to generalize the pseudo-
code above to a mini-batch setting, which we use in the experiments.

Algorithm 2 Training JEM++: Given network $f_{\bm{\theta}}$, step-size
$\alpha$, replay buffer $\mathbb{B}$, number of steps $M$ and $N$,
reinitialization frequency $\rho$, and number of classes $C$

1: while not converged do

2: Sample $(\bm{x}^{r},y^{r})\sim\mathcal{D}$

3: Sample $\bm{x}^{0}\sim\mathbb{B}$ with probability $1-\rho$, else
$\bm{x}^{0}\sim\mathcal{N}(\bm{\mu}_{y},\bm{\Sigma}_{y}),y\sim p(y)=\bm{\pi}$

4: Apply PYLD in Algo. 1 to sample $\bm{x}^{M}$ from $\bm{x}^{0}$

5: Calculate gradient with Eq. 2 from $\bm{x}^{r}$ and $\bm{x}^{M}$, and
gradient of CE loss from $(\bm{x}^{r},y^{r})$, and update model parameter
$\bm{\theta}$

6: Add / replace sample $\bm{x}^{M}$ back to $\mathbb{B}$

7: end while

##  4 Experiments

We evaluate the performance of JEM++ on multiple discriminative and generative
tasks, including image classification, image generation, adversarial
robustness, calibration of uncertainty, and out-of-distribution (OOD)
detection. Since our main goal is to improve JEM’s accuracy, training
stability and speed, we present these results in the main text and relegate
its downstream applications, such as adversarial robustness, calibration and
OOD detection, to the supplementary material. For a fair comparison with JEM
[17], our experiments closely follow the settings provided in the source code
of JEM444https://github.com/wgrathwohl/JEM. All our experiments are performed
with PyTorch on Nvidia RTX GPUs.

Table 2: Hybrid Modeling Results on CIFAR10. We report JEM++’s performance
with different $M$s when $N=5$ is fixed. We also report the per epoch speedup
between JEM and JEM++.

Class | Model | Acc % $\uparrow$ | IS∗ $\uparrow$ | FID∗ $\downarrow$
---|---|---|---|---
| Residual Flow [8] | 70.3 | 3.60 | 46.4
| Glow [32] | 67.6 | 3.92 | 48.9
Single | IGEBM [11] | 49.1 | 8.30 | 37.9
Hybrid | JEM (K=20) [17] $1\times$ | 92.9 | 8.76 | 38.4
Model | JEM++ (M=5) $2.4\times$ | 91.1 | 7.81 | 37.9
| JEM++ (M=10) $1.5\times$ | 93.5 | 8.29 | 37.1
| JEM++ (M=20) $.92\times$ | 94.1 | 8.11 | 38.0
Reg | VERA† ($\alpha$=100) $2.8\times$ | 93.2 | 8.11 | 30.5
Gen. | VERA [19] ($\alpha$=1) $2.8\times$ | 76.1 | 8.00 | 27.5
Disc. | WRN w/ BN | 95.8 | N/A | N/A
Gen. | SNGAN [37] | N/A | 8.59 | 25.5
| NCSN [49] | N/A | 8.91 | 25.3

  * •

†VERA uses an auxiliary generator to amortize the SGLD sampling and reports a
$2.8\times$ speedup without much details on how the evaluation is performed.

  * •

∗A fair evaluation of IS and FID is challenging as different methods use
different ways to measure the image quality. JEM uses an ensemble of models to
evaluate its IS and FID, while JEM++ only uses a single model for evaluation.
No more details are provided in JEM. Thus, it is difficult to have a fair
comparison.

Table 3: Test Accuracy (%) on SVHN and CIFAR100. Model | SVHN | CIFAR100
---|---|---
Softmax (w/ BN) | 97.0 | 78.9
VERA [19] | 96.8 | 72.2
JEM (K=20) | 96.7 | 72.2
JEM++ (M=5) | 96.7 | 72.0
JEM++ (M=10) | 96.9 | 74.5

###  4.1 Hybrid Modeling

We train JEM++ on three benchmark datasets: CIFAR10, CIFAR100 [34] and SVHN
[39], and compare it to the state-of-the-art hybrid models, as well as
standalone generative and discriminative models. Following the settings of JEM
[17], all our experiments are based on the Wide-ResNet architecture [58], with
the details of hyper-parameter settings of JEM++ provided in the supplementary
material. It’s worth mentioning that applying the SGD optimizer with $lr=0.1$
to JEM++ achieves better accuracy than the default setting of JEM using Adam
with $lr=0.0001$555JEM cannot use a learning rate larger than 0.0001.
Otherwise, it is extremely unstable and diverges easily at early epochs.. To
evaluate the quality of generated images, we adopt the Inception Score (IS)
[47] and Fréchet Inception Distance (FID) [27].

The results on CIFAR10, CIFAR100 and SVHN are reported in Table 2 and 3,
respectively. It can be observed that JEM++ ($M\\!=\\!10$) outperforms JEM and
other single-network hybrid models in terms of accuracy (93.5%), FID score
(37.1) and per epoch speedup (1.5$\times$), while being slightly worse in IS
score. Since no IS and FID scores are commonly reported on SVHN and CIFAR100,
we present the classification accuracy and generated samples on these two
benchmarks. Our JEM++ ($M\\!\\!=\\!\\!10$) model achieves an accuracy of 96.9%
and 74.5% on SVHN and CIFAR100, respectively, outperforming JEM by notable
margins. Example images generated by JEM++ for CIFAR10, SVHN and CIFAR100 are
shown in Figure 3 and 4, respectively. Additional JEM++ generated images can
be found in the supplementary material.

We also investigated JEM++’s performances on several downstream applications,
including adversarial robustness, calibration of uncertainty, and OOD
detection, where JEM++ achieves improved performances over the original JEM in
most of the cases. Due to page limit, the details are relegated to the
supplementary material.

(a) Unconditional Samples

(b) Class-conditional Samples

Figure 3: JEM++ generated CIFAR10 samples.

(a) SVHN

(b) CIFAR100

Figure 4: JEM++ generated class-conditional samples of SVHN and CIFAR100. Each
row corresponds to one class.

###  4.2 Training Stability and Speed

The main limitation of the SGLD-based training is the tradeoff between
training time and stability. The more SGLD sampling steps are used, the more
stable and better performance EBMs can achieve. In this section, we evaluate
JEM and JEM++ in terms of training stability and speed.

We first compare the training stability of JEM and JEM++. From our empirical
study, the official JEM ($K$-step SGLD with $K=20$) suffers from training
instability, i.e., it regularly diverges before 60 epochs. Prior works [11,
17], including JEM, fail to find a reasonably small $K$ to completely
stabilize the training of EBMs, and thus rely on checkpoints to resume the
training when divergences occur. Figure 5 shows the learning curves of JEM++
trained on CIFAR10 with different configurations. As can be seen, JEM++ is
much more stable and does not diverge when $M\\!=\\!20$. What’s more, JEM++
with $M\\!=\\!10$ can achieve high stability; even JEM++ with $M\\!=\\!5$ is
more stable than JEM with $K\\!=\\!20$. As discussed in Section 3, the
informative initialization improves JEM’s stability because the initial
samples $\bm{x}^{0}$ of SGLD are now close to the real data manifold. Hence,
the sampling process requires fewer steps to reach the low energy region of
the energy function, which we conjure should be much smoother than other
regions. In addition, the proximity constraint also improves the stability of
JEM++ as demonstrated in Figure 5.

Figure 5: The learning curves of JEM++ trained on CIFAR10 with different
configurations: (1) Number of steps $M$ with $N=5$, and (2) the proximity
constraint. The official JEM uses $K=20$, but it regularly diverges before 60
epochs.

We further compare the training speed between JEM and JEM++ in terms of run-
time per epoch. The results are reported in Table 4, where we compare JEM and
JEM++ trained on CIFAR10 with different configurations of $M$ and $N$. It can
be observed that $M$ specifies the total number of forward and backward
propagations of PYLD, consuming most of the run-time, while $N$ has a minor
impact on the run-time as it specifies the number of inner loops for sample
update, which is relatively inexpensive. Therefore, we can increase $N$ and
reduce $M$ to achieve approximately the same level of sample quality with much
less computation cost. Considering the training stability (Figure 5) and
training speed (Table 4), $M\\!=\\!10$ and $N\\!=\\!5$ achieves a good balance
between the two criteria and therefore is our default configuration of JEM++.

Table 4: Run-time comparison of JEM and JEM++ on CIFAR10. Model | Minutes per epoch | Speedup
---|---|---
JEM | 30.1 | $1\times$
JEM++, $M=5$ |  |
$N=5$ | 12.5 | $2.41\times$
$N=10$ | 12.6 | $2.39\times$
$N=20$ | 13.0 | $2.31\times$
JEM++, $M=10$ |  |
$N=5$ | 20.1 | $1.49\times$
$N=10$ | 20.3 | $1.48\times$
$N=20$ | 20.4 | $1.47\times$
JEM++, $M=20$ |  |
$N=5$ | 32.5 | $.93\times$
$N=10$ | 32.7 | $.92\times$
$N=20$ | 32.9 | $.91\times$

###  4.3 Ablation Study

JEM++ introduces a variety of new training procedures and architecture
features to improve JEM’s accuracy, training stability and speed. In this
section, we study the effect of different components of JEM++ on the
performance of image classification and image generation. Specifically, we
conduct the ablation study on CIFAR10 with an exhaustive comparison of
different components. We measure the effects of 1) w/o proximity constraint,
2) with Adam optimizer, 3) random initialization with batch norm enabled, and
4) two different types of initialization w/o batch norm.

The results are reported in Table 5. It can be observed that each component
contributes to JEM++’s performance positively. The proximity constraint in
Proximal SGLD improves both stability and accuracy. Our experiments show that
when a smaller $M$ enlarges the instability, the proximity constraint not only
helps to stabilize the training, but also improves the accuracy of the trained
models. The informative initialization also takes a significant role in JEM++,
which enables both batch norm and the use of SGD with larger learning rates.
When batch norm is enabled in JEM, we find that it can neither achieve a high
classification accuracy nor generate realistic images. On the other hand, JEM
without batch norm can achieve decent classification accuracy and generate
quality images, but it’s precarious and easily diverged at early epochs. The
informative initialization itself w/o batch norm is still beneficial to
stabilize the training, as manifested by the improved classification accuracy
and image quality. It’s worth mentioning that when batch norm is disabled,
only Adam [31] with a small learning rate no greater than 0.0001 yields a
stable training. However, when batch norm is enabled, the SGD optimizer with a
much larger learning rate can be applied to train JEM++ successfully,
outperforming the default Adam optimizer (with a very small learning rate)
used in JEM.

Table 5: Ablation study of different components of JEM++. All the models are
trained on CIFAR10 with $M=10$ and $N=5$.

Ablation | Acc % $\uparrow$ | IS $\uparrow$ | FID $\downarrow$
---|---|---|---
JEM++ | 93.5 | 8.29 | 37.1
w/o Proximity | 92.9 | 7.92 | 36.0
w/ Adam | 92.5 | 7.65 | 42.7
random init (w/ BN)1 | - | - | -
random init (w/o BN)2 | 88.6 | 7.64 | 35.1
informative init (w/o BN)3 | 91.1 | 7.92 | 39.8

  * •

1 It fails to achieve a high accuracy and generate realistic images.

  * •

2 It diverges early at epoch 28.

  * •

3 Without batch norm, only ADAM with $lr\\!=\\!0.0001$ can be used.

###  4.4 Classification Accuracy vs. Image Quality

One interesting phenomenon we observed from our experiments is the tradeoff
between classification accuracy and image quality. Figure 6 shows the
evolution of classification accuracy, IS and FID scores as a function of the
training epochs. At the early stage of training (before epoch 100), both
classification accuracy and image quality can be improved jointly. After that,
there is a clear competition between accuracy and image quality, where
improving accuracy hurts image quality. This probably can be explained by our
minimax objective (7), in which the classifier and the implicit generator
compete with each other to achieve an equilibrium. Compared to the standard
GANs [14], the difference is that we have only one network that serves both as
classifier and generator. How to balance the discriminative and generative
powers within one model is unclear. It would be interesting to investigate
this further in the future.

Figure 6: The evolution of JEM++’s classification accuracy, IS and FID scores
as a function of training epochs on CIFAR10. The spike around epoch 125 is due
to training instability and thanks to the proximity constraint, JEM++
stabilizes the training eventually.

##  5 Conclusion

In this paper, we propose JEM++ which improves JEM’s accuracy, training
stability and speed altogether with a number of new training procedures and
architecture features. We demonstrate the effectiveness of these improvements
on multiple benchmark datasets with state-of-the-art results in most of the
tasks of image classification, image generation, adversarial robustness,
uncertainty calibration and OOD detection. Most importantly, JEM++ enjoys
stable and accelerated training over the original JEM.

As for future work, we plan to investigate the tradeoff between the
classification accuracy and image quality as shown in Figure 6. We are
interested in what the optimal tradeoff is and how we can achieve the optimum
with architecture design and/or new training methodologies (e.g., [19, 2]). We
also plan to explore JEM++ to large-scale benchmarks, such as ImageNet, and
its application to other domains, such as NLP.

##  6 Acknowledgment

We would like to thank the anonymous reviewers for their comments and
suggestions, which helped improve the quality of this paper. We would also
gratefully acknowledge the support of VMware Inc. for its university research
fund to this research.

## References

  * [1] Naveed Akhtar, Jian Liu, and Ajmal Mian.  Defense against universal adversarial perturbations.  In Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition, 2018.
  * [2] Lynton Ardizzone, Radek Mackowiak, Carsten Rother, and Ullrich Köthe.  Training normalizing flows with the information bottleneck for competitive generative classification.  In Neural Information Processing Systems (NeurIPS), 2020.
  * [3] Jens Behrmann, Will Grathwohl, Ricky TQ Chen, David Duvenaud, and Jörn-Henrik Jacobsen.  Invertible residual networks.  In International Conference on Machine Learning (ICML), 2018.
  * [4] Andrew Brock, Jeff Donahue, and Karen Simonyan.  Large scale GAN training for high fidelity natural image synthesis.  In International Conference on Learning Representations(ICLR), 2019\.
  * [5] Andrew Brock, Jeff Donahue, and Karen Simonyan.  Large scale GAN training for high fidelity natural image synthesis.  In International Conference on Learning Representations (ICLR), 2019\.
  * [6] Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared D Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, Sandhini Agarwal, Ariel Herbert-Voss, Gretchen Krueger, Tom Henighan, Rewon Child, Aditya Ramesh, Daniel Ziegler, Jeffrey Wu, Clemens Winter, Chris Hesse, Mark Chen, Eric Sigler, Mateusz Litwin, Scott Gray, Benjamin Chess, Jack Clark, Christopher Berner, Sam McCandlish, Alec Radford, Ilya Sutskever, and Dario Amodei.  Language models are few-shot learners.  In Advances in Neural Information Processing Systems (NeurIPS), 2020\.
  * [7] Olivier Chapelle, Bernhard Scholkopf, and Alexander Zien.  Semi-supervised learning.  IEEE Transactions on Neural Networks, 2009.
  * [8] Ricky TQ Chen, Jens Behrmann, David Duvenaud, and Jörn-Henrik Jacobsen.  Residual flows for invertible generative modeling.  arXiv preprint arXiv:1906.02735, 2019.
  * [9] Ping-yeh Chiang, Renkun Ni, Ahmed Abdelkader, Chen Zhu, Christoph Studer, and Tom Goldstein.  Certified defenses for adversarial patches.  In ICLR 2020, 2020.
  * [10] Arthur P Dempster, Nan M Laird, and Donald B Rubin.  Maximum likelihood from incomplete data via the em algorithm.  Journal of the Royal Statistical Society: Series B (Methodological), 1977.
  * [11] Yilun Du and Igor Mordatch.  Implicit generation and generalization in energy-based models.  In Advances in Neural Information Processing Systems (NeurIPS), 2019\.
  * [12] Gintare Karolina Dziugaite, Zoubin Ghahramani, and Daniel M Roy.  A study of the effect of jpg compression on adversarial images.  arXiv preprint arXiv:1608.00853, 2016.
  * [13] Ruiqi Gao, Erik Nijkamp, Diederik P. Kingma, Zhen Xu, Andrew M. Dai, and Ying Nian Wu.  Flow contrastive estimation of energy-based models.  In IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR), 2020.
  * [14] Ian Goodfellow, Jean Pouget-Abadie, Mehdi Mirza, Bing Xu, David Warde-Farley, Sherjil Ozair, Aaron Courville, and Yoshua Bengio.  Generative adversarial nets.  In Advances in Neural Information Processing Systems (NeurIPS), 2014\.
  * [15] Ian Goodfellow, Jonathon Shlens, and Christian Szegedy.  Explaining and harnessing adversarial examples.  In International Conference on Learning Representations (ICLR), 2015\.
  * [16] Ian J Goodfellow, Jonathon Shlens, and Christian Szegedy.  Explaining and harnessing adversarial examples.  In International Conference on Learning Representations, 2015.
  * [17] Will Grathwohl, Kuan-Chieh Wang, Joern-Henrik Jacobsen, David Duvenaud, Mohammad Norouzi, and Kevin Swersky.  Your classifier is secretly an energy based model and you should treat it like one.  In International Conference on Learning Representations (ICLR), 2020\.
  * [18] Will Grathwohl, Kuan-Chieh Wang, Joern-Henrik Jacobsen, David Duvenaud, and Richard Zemel.  Learning the stein discrepancy for training and evaluating energy-based models without sampling.  In Proceedings of the 37th International Conference on Machine Learning (ICML), 2020.
  * [19] Will Sussman Grathwohl, Jacob Jin Kelly, Milad Hashemi, Mohammad Norouzi, Kevin Swersky, and David Duvenaud.  No mcmc for me: Amortized sampling for fast and stable training of energy-based models.  In International Conference on Learning Representations (ICLR), 2021\.
  * [20] Chuan Guo, Geoff Pleiss, Yu Sun, and Kilian Q Weinberger.  On calibration of modern neural networks.  In International Conference on Machine Learning (ICML), 2017.
  * [21] Chuan Guo, Mayank Rana, Moustapha Cisse, and Laurens Van Der Maaten.  Countering adversarial images using input transformations.  arXiv preprint arXiv:1711.00117, 2017.
  * [22] Michael Gutmann and Aapo Hyvärinen.  Noise-contrastive estimation: A new estimation principle for unnormalized statistical models.  In Proceedings of the Thirteenth International Conference on Artificial Intelligence and Statistics, 2010.
  * [23] Tian Han, Erik Nijkamp, Xiaolin Fang, Mitch Hill, Song-Chun Zhu, and Ying Nian Wu.  Divergence triangle for joint training of generator model, energy-based model, and inference model.  In Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition, 2019.
  * [24] Tian Han, Erik Nijkamp, Linqi Zhou, Bo Pang, Song-Chun Zhu, and Ying Nian Wu.  Joint training of variational auto-encoder and latent energy-based model.  In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR), 2020.
  * [25] Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun.  Deep residual learning for image recognition.  In IEEE Conference on Computer Vision and Pattern Recognition (CVPR), 2016.
  * [26] Dan Hendrycks and Kevin Gimpel.  A baseline for detecting misclassified and out-of-distribution examples in neural networks.  In International Conference on Learning Representations(ICLR), 2016\.
  * [27] Martin Heusel, Hubert Ramsauer, Thomas Unterthiner, Bernhard Nessler, and Sepp Hochreiter.  Gans trained by a two time-scale update rule converge to a local nash equilibrium.  In Advances in Neural Information Processing Systems (NeurIPS), 2017\.
  * [28] Geoffrey E Hinton.  Training products of experts by minimizing contrastive divergence.  Neural computation, 2002.
  * [29] Jonathan Ho, Ajay Jain, and Pieter Abbeel.  Denoising diffusion probabilistic models.  In Advances in Neural Information Processing Systems (NeurIPS), 2020\.
  * [30] Sergey Ioffe and Christian Szegedy.  Batch normalization: Accelerating deep network training by reducing internal covariate shift.  In International Conference on Machine Learning (ICML), 2015.
  * [31] Diederik Kingma and Jimmy Ba.  Adam: A method for stochastic optimization.  In International Conference on Learning Representations (ICLR), 2015\.
  * [32] Durk P Kingma and Prafulla Dhariwal.  Glow: Generative flow with invertible 1x1 convolutions.  In Advances in Neural Information Processing Systems, 2018.
  * [33] Diederik P Kingma and Max Welling.  Auto-encoding variational bayes.  In International Conference on Learning Representations (ICLR), 2014\.
  * [34] Alex Krizhevsky and Geoffrey Hinton.  Learning multiple layers of features from tiny images.  Technical report, Citeseer, 2009.
  * [35] Yann LeCun, Sumit Chopra, Raia Hadsell, M Ranzato, and F Huang.  A tutorial on energy-based learning.  Predicting structured data, 2006.
  * [36] Aleksander Madry, Aleksandar Makelov, Ludwig Schmidt, Dimitris Tsipras, and Adrian Vladu.  Towards deep learning models resistant to adversarial attacks.  In International Conference on Learning Representations (ICLR), 2018\.
  * [37] Takeru Miyato, Toshiki Kataoka, Masanori Koyama, and Yuichi Yoshida.  Spectral normalization for generative adversarial networks.  In International Conference on Learning Representations (ICLR), 2018\.
  * [38] Eric Nalisnick, Akihiro Matsukawa, Yee Whye Teh, Dilan Gorur, and Balaji Lakshminarayanan.  Do deep generative models know what they don’t know?  arXiv preprint arXiv:1810.09136, 2018.
  * [39] Yuval Netzer, Tao Wang, Adam Coates, Alessandro Bissacco, Bo Wu, and Ng andrew Y.  Reading digits in natural images with unsupervised feature learning.  2011\.
  * [40] Erik Nijkamp, Mitch Hill, Tian Han, Song-Chun Zhu, and Ying Nian Wu.  On the anatomy of mcmc-based maximum likelihood learning of energy-based models.  In AAAI Conference on Artificial Intelligence (AAAI), 2019.
  * [41] Erik Nijkamp, Song-Chun Zhu, and Ying Nian Wu.  On learning non-convergent short-run mcmc toward energy-based model.  In Advances in Neural Information Processing Systems (NeurIPS), 2019\.
  * [42] Neal Parikh and Stephen Boyd.  Proximal algorithms.  Foundations and Trends in Optimization, 1(3):127–239, 2014.
  * [43] Lev Semenovich Pontryagin.  Mathematical theory of optimal processes.  CRC, 1987.
  * [44] Jonas Rauber, Wieland Brendel, and Matthias Bethge.  Foolbox: A python toolbox to benchmark the robustness of machine learning models.  arXiv preprint arXiv:1707.04131, 2017.
  * [45] Herbert Robbins and Sutton Monro.  A stochastic approximation method.  The Annals of Mathematical Statistics, 1951.
  * [46] R. Rockafellar.  Augmented lagrangians and applications of the proximal point algorithm in convex programming.  Mathematics of Operations Research, 1(2):97–116, 1976.
  * [47] Tim Salimans, Ian Goodfellow, Wojciech Zaremba, Vicki Cheung, Alec Radford, and Xi Chen.  Improved techniques for training gans.  In Advances in neural information processing systems (NeurIPS), 2016\.
  * [48] Shibani Santurkar, Andrew Ilyas, Dimitris Tsipras, Logan Engstrom, Brandon Tran, and Aleksander Madry.  Image synthesis with a single (robust) classifier.  In Advances in Neural Information Processing Systems (NeurIPS), 2019\.
  * [49] Yang Song and Stefano Ermon.  Generative modeling by estimating gradients of the data distribution.  arXiv preprint arXiv:1907.05600, 2019.
  * [50] Christian Szegedy, Wojciech Zaremba, Ilya Sutskever, Joan Bruna, Dumitru Erhan, Ian Goodfellow, and Rob Fergus.  Intriguing properties of neural networks.  In International Conference on Learning Representations (ICLR), 2014\.
  * [51] Tijmen Tieleman.  Training restricted boltzmann machines using approximations to the likelihood gradient.  In International conference on Machine learning (ICML), 2008.
  * [52] Arash Vahdat and Jan Kautz.  Nvae: A deep hierarchical variational autoencoder.  In Advances in Neural Information Processing Systems, 2020.
  * [53] Pascal Vincent.  A connection between score matching and denoising autoencoders.  Neural computation, 2011.
  * [54] Max Welling and Yee W Teh.  Bayesian learning via stochastic gradient langevin dynamics.  In Proceedings of the 28th international conference on machine learning (ICML), 2011.
  * [55] Cihang Xie, Mingxing Tan, Boqing Gong, Jiang Wang, Alan L. Yuille, and Quoc V. Le.  Adversarial examples improve image recognition.  In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR), 2020.
  * [56] 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:27–45, Jan. 2020.
  * [57] Jianwen Xie, Yang Lu, Song-Chun Zhu, and Yingnian Wu.  A theory of generative convnet.  In International Conference on Machine Learning, 2016.
  * [58] Sergey Zagoruyko and Nikos Komodakis.  Wide residual networks.  In The British Machine Vision Conference (BMVC), 2016.
  * [59] Dinghuai Zhang, Tianyuan Zhang, Yiping Lu, Zhanxing Zhu, and Bin Dong.  You only propagate once: Accelerating adversarial training via maximal principle.  In Advances in Neural Information Processing Systems (NeurIPS), 2019\.
  * [60] Stephen Zhao, Jorn-Henrik Jacobsen, and Will Grathwohl.  Joint energy-based models for semi-supervised classification.  In ICML 2020 Workshop on Uncertainty and Robustness in Deep Learning, 2020.

##  Appendix A Experimental Details

To have a fair comparison with JEM [17], all our experiments are based on the
Wide-ResNet architecture [58] and follow JEM’s settings whenever possible. As
we discussed in the main text, JEM++ enables batch norm [30] and the SGD
optimizer [45] with a large learning rate, which we find works better than
Adam [31] with a very small learning rate of $1e\\!-\\!4$ that is used by JEM.
Specifically, we use SGD with an initial learning rate of 0.1 and a decay rate
of 0.2, and train all our models for 150 epochs. We reduce the learning rate
at epoch [50, 100, 125]. Table 6 lists the hyperparameters of JEM++. Note that
JEM++ is still highly stable even with $M=5$. More experimental details can be
found in our code, which is publicly available at
https://github.com/sndnyang/JEMPP.

Table 6: Hyperparameters of JEM++ for CIFAR10 Variable | Value
---|---
Number of outer steps $M$ | 5, 10
Number of inner steps $N$ | 5
Proximity constraint $\varepsilon$ | 1
Buffer size $|\mathbb{B}|$ | 10,000
Reinitialization freq. $\rho$ | 5%
PYLD step-size $\alpha$ | 0.2

##  Appendix B Informative Initialization

(a) SVHN

(b) CIFAR100

Figure 7: The categorical centers of SVHN and CIFAR100.

(a) Categorical centers of CIFAR10

(b) Samples from each category

Figure 8: The categorical centers of and corresponding samples of CIFAR10.

In this paper, we introduce a novel informative initialization to start the
SGLD chain. Specifically, instead of using a uniform distribution, we sample
from a Gaussian mixture distribution estimated from the training data as

$\displaystyle p_{0}(\bm{x})$
$\displaystyle=\sum\nolimits_{y}\mathcal{\pi}_{y}{N}(\bm{\mu}_{y},\bm{\Sigma}_{y})$
(15) $\displaystyle\text{with}\quad\pi_{y}$
$\displaystyle=|\mathcal{D}_{y}|/\sum\nolimits_{y^{\prime}}|\mathcal{D}_{y^{\prime}}|,\quad\bm{\mu}_{y}=\mathbb{E}_{\bm{x}\sim\mathcal{D}_{y}}[\bm{x}],$
$\displaystyle\bm{\Sigma}_{y}$
$\displaystyle=\mathbb{E}_{\bm{x}\sim\mathcal{D}_{y}}\left[\left(\bm{x}-\bm{\mu}_{y}\right)\left(\bm{x}-\bm{\mu}_{y}\right)^{\top}\right],$

where $\mathcal{D}_{y}$ denotes the set of training samples with label $y$.
Figure 7 visualizes the categorical centers ($\bm{\mu}$’s) estimated from the
SVHN and CIFAR100 training datasets. Figure 8 visualizes the categorical
centers and the corresponding samples $\bm{x}_{0}\sim p_{0}(\bm{x})$ for
CIFAR10. Note that no extra information is used to train JEM++ over JEM.

##  Appendix C Applications

In the main text, we compared JEM and JEM++ in terms of classification
accuacy, image quality, training stability and speed. Here we compare JEM and
JEM++ in other downstream applications, such as adversarial robustness,
calibration and out of distribution (OOD) detection.

###  C.1 Robustness

It’s well known that DNNs are particularly vulnerable to adversarial examples
[50, 16] in the form of small perturbations to inputs that lead DNNs to
predict incorrect outputs. Specially, the widely explored adversarial examples
are defined as perturbed inputs $\tilde{\bm{x}}=\bm{x}+\bm{\delta}$ under an
$L_{p}$-norm constraint $\|\bm{\delta}\|_{p}<\varepsilon$. To overcome the
security threat posed by adversarial examples, a variety of defense algorithms
have been proposed in the past few years to improve the robustness of models
[15, 12, 21, 1, 36, 9]. Among them, adversarial training [15, 36] has been
proved to be the most effective one to defend adversarial examples.

As we discussed in Section 3.1, there is a close relationship between the
maximum likelihood learning of EBM (7) and adversarial training with PGD [36]
as both solve a similar minimax objective. Therefore, the maximum likelihood
trained EBMs should be more robust to adversarial examples than the standard
trained softmax classifiers, and this has been empirically verified by recent
works (e.g., [11, 17]). Since JEM++ improves JEM’s accuracy, training
stability and speed, it’s interesting to check if JEM++ can improve model
robustness as well.

(a) $L_{\infty}$ Robustness

(b) $L_{2}$ Robustness

Figure 9: Adversarial robustness under the PGD attacks.

To evaluate the robustness of a given model, we run a white-box PGD attack
[36] under an $L_{\infty}$ or $L_{2}$ constraint using foolbox [44], with the
results reported in Figure 9. It can be observed that JEM++ achieves a similar
robustness with JEM under the $L_{\infty}$ and $L_{2}$ PGD attacks, while both
are more robust than the standard softmax classifiers. The adversarial
training with PGD [36, 48] achieves the highest robustness since it is trained
and test under the same PGD attacks, while JEM/JEM++ are trained on real and
generated samples from the energy function, without the access to the PGD
samples for training.

###  C.2 Calibration

Recent researches have shown that the predictions from modern DNNs could be
over-confident [20], i.e., they often output incorrect but confident
predictions, which could have catastrophic consequences. Hence, calibration of
uncertainty for DNNs is a critical task with an enormous practical impact
nowadays. Here, the confidence is defined as $\max_{y}p(y|\bm{x})$ which is
used to decide when to output a prediction. In this section, we compare the
calibration qualities of models trained by JEM and JEM++ as well as the
standard softmax classifiers on the CIFAR10/100 dataset.

#### Expected Calibration Error

(ECE) is a standard metric to evaluate the calibration quality of a classifier
[20]. It firstly computes the confidence of the model,
$\max_{y}p(y|\bm{x}_{i})$, for each $\bm{x}_{i}$ in the dataset. Then it
groups the predictions into equally spaced buckets
$\\{B_{1},B_{2},\cdots,B_{M}\\}$ based on the confidence scores. For example,
if $M$ = 20, then $B_{1}$ would represent all examples for which the model’s
confidence scores were between 0 and 0.05. Then ECE is calculated as

$\mathrm{ECE}=\sum_{m=1}^{M}\frac{\left|B_{m}\right|}{n}\left|\operatorname{acc}\left(B_{m}\right)-\operatorname{conf}\left(B_{m}\right)\right|,$
(16)

where $n$ is the number of data in the dataset, acc($B_{m}$) is the average
accuracy of the model on all the examples in $B_{m}$ and conf($B_{m}$) is the
average confidence on all the examples in $B_{m}$. In our experiments, we set
$M$ = 20. For a perfectly calibrated model, the ECE will be 0 for any $M$.

Figures 10 and 11 report the results on CIFAR10 and CIFAR100, respectively. As
we can see, the models trained by JEM and JEM++ are better calibrated than the
standard softmax classifiers, while JEM++ achieves better calibration
qualities than JEM on CIFAR10 (2.35% vs. 4.2%) and CIFAR100 (3.3% vs. 4.87%)
with notable margins.

(a) Standard Softmax

(b) JEM (K=20)

(c) JEM++ (M=5)

(d) JEM++ (M=10)

Figure 10: Calibration results on CIFAR10. The smaller ECE is, the better.

(a) Standard Softmax

(b) JEM (K=20)

(c) JEM++ (M=5)

(d) JEM++ (M=10)

Figure 11: Calibration results on CIFAR100. The smaller ECE is, the better.

###  C.3 Out-Of-Distribution Detection

The OOD detection is a binary classification problem, which outputs a score
$s_{\bm{\theta}}(\bm{x})\in\mathbb{R}$ for a given query $\bm{x}$. The model
should be able to assign lower scores to OOD examples than to in-distribution
examples, such that it can be used to distinguish two sets of examples.
Following the settings of JEM [17], we use the Area Under the Receiver-
Operating Curve (AUROC) [26] to evaluate the performance of OOD detection. In
our experiments, two standard score functions are considered: the input
density $p_{\bm{\theta}}(\bm{x})$ [38] and the predictive distribution
$p_{\bm{\theta}}(y|\bm{x})$ [26].

#### Input Density

A natural choice of $s_{\bm{\theta}}(\bm{x})$ is the input density
$p_{\bm{\theta}}(\bm{x})$. For OOD detection, intuitively we consider examples
with low $p(\bm{x})$ to be OOD. Quantitative results can be found in Table 7
(top row), where CIFAR10 is the in-distribution data and SVHN, an interpolated
CIFAR10, CIFAR100 and CelebA are treated as out-of-distribution data,
respectively. Moreover, the corresponding distributions of scores are
visualized in Table 8. As can be seen, the JEM++ model assigns higher
likelihoods to in-distribution data than to the OOD data, outperforming JEM
and all the other models by significant margins.

#### Predictive Distribution

Another useful OOD score is the maximum probability from a classifier’s
predictive distribution:
$s_{\bm{\theta}}(\bm{x})=\max_{y}p_{\bm{\theta}}(y|\bm{x})$. Hence, OOD
performance using this score is highly correlated with a model’s
classification accuracy. The results can be found in Table 7 (bottom row).
Again, JEM++ outperforms JEM and all the other models by notable margins.

Table 7: OOD detection results. Models are trained on CIFAR10. Values are AUROC. $s_{\bm{\theta}}(\bm{x})$ | Model | SVHN | CIFAR10 Interp | CIFAR100 | CelebA
---|---|---|---|---|---
$\log p_{\bm{\theta}}(\bm{x})$ | Uncond Glow | .05 | .51 | .55 | .57
IGEBM | .63 | .70 | .50 | .70
JEM (K=20) | .67 | .65 | .67 | .75
JEM++ (M=5) | .89 | .73 | .81 | .74
JEM++ (M=10) | .63 | .68 | .64 | .59
| JEM++ (M=20) | .85 | .57 | .68 | .89
$\max_{y}p_{\bm{\theta}}(y|\bm{x})$ | WideResNet | .93 | .77 | .85 | .62
IGEBM | .43 | .69 | .54 | .69
JEM (K=20) | .89 | .75 | .87 | .79
JEM++ (M=5) | .88 | .78 | .86 | .78
JEM++ (M=10) | .91 | .78 | .88 | .82
| JEM++ (M=20) | .94 | .77 | .88 | .90
JEM |  |  |
---|---|---|---
JEM++(M=5) |  |  |
JEM++(M=10) |  |  |
JEM++(M=20) |  |  |
Table 8: Histograms of $\log_{\bm{\theta}}p(\bm{x})$ for OOD detection. Green
corresponds to in-distribution dataset, while red corresponds to OOD dataset.

##  Appendix D Additional Generated Samples

Additional JEM++ generated samples of SVHN and CIFAR100 are provided in Figure
12. Additional JEM++ generated class-conditional (best and worst) samples of
CIFAR10 are provided in Figures 13-22. It is worth noting that the worst
images (the lowest $p(\bm{x})$ or $p(y|\bm{x})$) generated by JEM++ are more
visually appealing than JEM generated (see examples in the Appendix of JEM
[17]).

(a) SVHN (Conditional)

(b) CIFAR100 (Conditional)

Figure 12: JEM++ generated class-conditional samples of SVHN and CIFAR100.
Each row corresponds to one class.

(a) Samples with highest $p(\bm{x}$)

(b) Samples with lowest $p(\bm{x})$

(c) Samples with highest $p(y|\bm{x}$)

(d) Samples with lowest $p(y|\bm{x})$

Figure 13: JEM++ generated class-conditional samples of Plane

(a) Samples with highest $p(\bm{x}$)

(b) Samples with lowest $p(\bm{x})$

(c) Samples with highest $p(y|\bm{x}$)

(d) Samples with lowest $p(y|\bm{x})$

Figure 14: JEM++ generated class-conditional samples of Car

(a) Samples with highest $p(\bm{x}$)

(b) Samples with lowest $p(\bm{x})$

(c) Samples with highest $p(y|\bm{x}$)

(d) Samples with lowest $p(y|\bm{x})$

Figure 15: JEM++ generated class-conditional samples of Bird

(a) Samples with highest $p(\bm{x}$)

(b) Samples with lowest $p(\bm{x})$

(c) Samples with highest $p(y|\bm{x}$)

(d) Samples with lowest $p(y|\bm{x})$

Figure 16: JEM++ generated class-conditional samples of Cat

(a) Samples with highest $p(\bm{x}$)

(b) Samples with lowest $p(\bm{x})$

(c) Samples with highest $p(y|\bm{x}$)

(d) Samples with lowest $p(y|\bm{x})$

Figure 17: JEM++ generated class-conditional samples of Deer

(a) Samples with highest $p(\bm{x}$)

(b) Samples with lowest $p(\bm{x})$

(c) Samples with highest $p(y|\bm{x}$)

(d) Samples with lowest $p(y|\bm{x})$

Figure 18: JEM++ generated class-conditional samples of Dog

(a) Samples with highest $p(\bm{x}$)

(b) Samples with lowest $p(\bm{x})$

(c) Samples with highest $p(y|\bm{x}$)

(d) Samples with lowest $p(y|\bm{x})$

Figure 19: JEM++ generated class-conditional samples of Frog

(a) Samples with highest $p(\bm{x}$)

(b) Samples with lowest $p(\bm{x})$

(c) Samples with highest $p(y|\bm{x}$)

(d) Samples with lowest $p(y|\bm{x})$

Figure 20: JEM++ generated class-conditional samples of Horse

(a) Samples with highest $p(\bm{x}$)

(b) Samples with lowest $p(\bm{x})$

(c) Samples with highest $p(y|\bm{x}$)

(d) Samples with lowest $p(y|\bm{x})$

Figure 21: JEM++ generated class-conditional samples of Ship

(a) Samples with highest $p(\bm{x}$)

(b) Samples with lowest $p(\bm{x})$

(c) Samples with highest $p(y|\bm{x}$)

(d) Samples with lowest $p(y|\bm{x})$

Figure 22: JEM++ generated class-conditional samples of Truck