EBM L IFE C YCLE : MCMC S TRATEGIES FOR S YNTHESIS , D EFENSE , AND D ENSITY M ODELING
Anonymous authors
Paper under double-blind review
A BSTRACT
This work presents strategies to learn an Energy-Based Model (EBM) according to the desired length of its MCMC sampling trajectories.
MCMC trajectories of different lengths correspond to models with different purposes.
Our experiments cover three different trajectory magnitudes and learning outcomes: 1) shortrun sampling for image generation; 2) midrun sampling for classiﬁer-agnostic adversarial defense; and 3) longrun sampling for principled modeling of image probability densities.
To achieve these outcomes, we introduce three novel methods of MCMC initialization for negative samples used in Maximum Likelihood (ML) learning.
With standard network architectures and an unaltered ML objective, our MCMC initialization methods alone enable signiﬁcant performance gains across the three applications that we investigate.
Our results include state-of-the-art FID scores for unnormalized image densities on the CIFAR-10 and ImageNet datasets; state-ofthe-art adversarial defense on CIFAR-10 among puriﬁcation methods and the ﬁrst
EBM defense on ImageNet; and scalable techniques for learning valid probability densities.
I NTRODUCTION
Generative modeling of complex signals is one of the fundamental challenges of computational cognition.
The importance of studying generative models is often justiﬁed by the Analysis-bySynthesis hypothesis, which posits that human cognition is an interplay of bottom-up information from world states and top-down synthesis of imagined states.
Powerful generative models could provide the foundation of the imaginative capabilities of future machine intelligence.
One approach to generative modeling is to posit the existence of a density q ( x ) of signals x that generates data samples, and to learn q ( x ) using a ﬂexible model p ( x ; θ ) , where θ is a model parameter.
This approach is taken by Energy-Based Models (EBMs) (Xie et al., 2016), normalizing ﬂows (Kingma & Dhariwal, 2018), score-based models (Song & Ermon, 2019; 2020), and auto-regressive models (Oord et al., 2016), as well as by Variational Auto-encoders (VAEs) (Kingma & Welling, 2013) using a joint model p ( x, z ; θ ) .
In this work we consider modeling q ( x ) using an EBM density p ( x ; θ ) for image signals x .
See Appendix A for a brief review of Maximum Likelihood learning with an EBM.
The primary goal of most deep generative modeling is to generate realistic images.
It is well-known that shortrun sampling with an EBM is an effective method for image generation, but synthesis results still lag behind GANs (Goodfellow et al., 2014), score models, and diffusion models (Ho et al., 2020).
Other directions of EBM research include learning valid densities (Nijkamp et al., 2020), combining discriminative and generative learning via the Joint Energy Model (Grathwohl et al., 2020) and relatives, and using an EBM for defense against adversarial attacks (Hill et al., 2021).
In this work, we focus on unconditional learning where EBMs are trained exclusively on unlabeled images.
We explore the tasks of image synthesis, adversarial defense, and density modeling to examine a breadth of capabilities for the unconditional EBM.
Image synthesis is a shared goal across all generative models, while defense and density modeling are tasks that are especially suitable for EBMs.
See Appendices B and I for a thorough comparison of EBM and other generative models from the perspective of synthesis, defense, and density estimation.
Each task is naturally associated with a certain length of MCMC trajectory.
Image synthesis is most effective with shortrun trajectories (about 20 to 200 steps) that can rapidly generate new images.
Adversarial defense requires midrun trajectories (about 200 to 2000 steps) that can preserve the class features while sampling removes adversarial signals.
Density modeling requires longrun trajectories (50K steps or more) to ensure proper calibration of probability mass for the model steady-state.
It is known that EBMs have a near universal tendency to learn a misaligned steady-state focusing on unrealistic images (Nijkamp et al., 2020), and there are very few existing solutions to correct this.
We intuitively refer to the spectrum of trajectories as the life cycle of an MCMC sample, from youth through middle age to maturity.
Young samples have the highest quality visual appearance, middle age samples are useful for securing classiﬁers, and mature samples represent grounded knowledge of the data density (or lack thereof in the widespread case of a misaligned steady-state).
Figureillustrates the sampling paths that we study in this work.
Ideally, a valid density estimator would also be a good synthesizer and a good defender.
However, synthesis quality and density modeling are goals that tend to be at odds with each other.
High-quality synthesis is easier to achieve using shortrun sampling with a defective density rather than long-run sampling with a stable density (Nijkamp et al., 2020).
Unfortunately, models with high-quality shortrun synthesis lack the stability needed for defense.
The EBM defense from Hill et al.
(2021) advocates for the use of a valid density approximation to stabilize trajectories, but this requirement might be too strict given the difﬁculty of density estimation.
Learning midrun sampling trajectories sufﬁcient for defense is much more feasible for complex datasets such as ImageNet than full density estimation, even if longrun samples from the defensive model are not realistic.
Our work is the ﬁrst to explore this possibility.
The overarching purpose of this paper is to discuss techniques for building sample paths from shortrun to longrun, with the hope that these techniques will eventually enable both high quality synthesis and defense to be accomplished with a model that is also a valid density.
For now, the difﬁculty of valid density modeling leads us to restrict our focus to separate time scales of the sampling regime necessary for each task.
Interestingly, EBM training naturally accommodates learning at different trajectory lengths.
Strategies for improving EBM learning beyond the standard framework (e.g. Xie et al.
; Du & Mordatch (2019); Nijkamp et al.
(2020)) can broadly be divided into methods that focus on the initialization of MCMC samples (Xie et al., 2018; Gao et al., 2018; Nijkamp et al., 2019) and methods that focus on the ML learning objective (Yu et al., 2020b).
Some works explore both (Gao et al., 2020a;
Du et al., 2020).
Our novel learning methods focus on MCMC initialization, and we retain the standard ML objective and use conventional network architectures.
We introduce three new MCMC initialization strategies which are tailored to the three different trajectories lengths we explore.
During training we exclusively use shortrun MCMC to ensure computational feasibility.
Learning models with midrun and longrun trajectories is accomplished by simulating longer trajectories via well- chosen initialization and optimizer annealing.
Our initialization strategies are able to signiﬁcantly improve the state-of-the-art across the tasks we investigate.
We summarize our contributions below.
• We propose a hybridization of persistent (Tieleman, 2008) and cooperative (Xie et al., 2018) initialization to overcome limitations of each.
The proposed method yields state-of-the-art FID scores for unconditional unnormalized image densities for the CIFAR-10 and ImageNet datasets.
The ImagetNet results surpass GANs trained with similar resources.
See Section 2.
• We show that persistent initialization with appropriately tuned rejuvenation from in-distribution states can be used to train EBMs with stable trajectories of several thousand MCMC steps.
This allows us to extend the method of Hill et al.
(2021) to obtain state-of-the-art puriﬁcation-based defense for CIFAR-10 and to scale the EBM defense to ImageNet.
See Section 3.
• We propose a method for principled density estimation that allows incorporation of a rejuvenation step for persistent states.
Incorporating rejuvenation allows us to learn well-formed EBM densities at a greater scale than previously possible.
See Section 4.
Our initialization methods will primarily build upon persistent (Tieleman, 2008; Du & Mordatch, 2019) and cooperative (Xie et al., 2018) initialization.
All of our methods will use a generator network as the source of rejuvenation for persistent states.
Our shortrun experiments will learn the generator in tandem with the EBM so the synthesis process is self-contained, while our midrun and longrun experiments will use pretrained generators since we will apply sampling paths from in-distribution initial images rather than synthesizing from scratch.
H YBRID P ERSISTENT C OOPERATIVE L EARNING FOR I MAGE S YNTHESIS
In this section, we focus on on the conventional task of learning an EBM for high quality synthesis with shortrun sampling.
We restrict our attention to achieving high quality synthesis without use of pretrained models, in contrast with works such as Che et al.
Alayrac et al.
Our method involves hybridizing persistent (Tieleman, 2008) and cooperative (Xie et al., 2018) initialization.
We ﬁrst brieﬂy cover the strengths and weaknesses of these methods, then discuss our hybridization, and ﬁnally present experimental results.
2. M OTIVATION : L IMITATIONS OF P ERSISTENT AND C OOPERATIVE I NITIALIZATION
A common framework for learning synthesis with an EBM is to use persistent initialization for
MCMC samples with a certain rate of rejuvenation (e.g. 5% chance) from a noise distribution (Du & Mordatch, 2019).
This approach can cause instability because shortrun samples used to update the model include a mix of higher-energy burn-in samples and lower-energy realistic images, which can destabilize training by increasing the variance of the gradient of the negative samples in (4).
On the other hand, removing rejuvenation can decrease the quality of the learned images because persistent images often become stuck in local modes and develop defects that linger for many updates.
Persistent banks also scale poorly to large datasets such as ImageNet because the bank cannot efﬁciently represent the diversity of the data.
Initialization from a cooperative generator network is an appealing alternative because it could enable efﬁcient in-distribution generation of highly diverse appearances.
Using a generator for rejuvenation allows samples to begin much closer to the correct energy spectrum, thereby avoid the instability of noise rejuvenation.
However, we observe a major limitation of cooperative learning that affects the generator ability to produce diverse initial states.
In particular, a generator trained with the cooperative learning objective has difﬁculty breaking the symmetry of its activations.
This happens because shortrun EBM samples are unable to provide novel diversity if the generator initialization already lacks diversity.
This is illustrated in Figure 5.
Learning quickly becomes unstable without auxiliary techniques such as batch normalization to break generator symmetry.
The lack of diversity of shortrun EBM samples limits the results of cooperative learning even when training succeeds.
The strengths and limitations of these methods are complementary.
Persistent initialization can reliably provide a diverse set of initial images that represent prior samples of several different model snapshots.
Even if a single EBM update is biased to a certain image defect (e.g. MCMC samples are blue-tinged or too bright), samples across previous iterations have the correct diversity on average.
On the other hand, the generator can quickly propose in-distribution states for rejuvenation instead of relying on out-of-distribution noise samples or aggressively augmented persistent samples.
Furthermore, training the generator in tandem with the EBM provides an efﬁcient way of generating new samples from scratch after training that match the samples used during training.
In contrast, models trained with persistent initialization often cannot produce good samples after training without an in-distribution initialization (Nijkamp et al., 2020).
Learning a generator and EBM jointly enables self-contained and reproducible FID implementation to evaluate the learned model.
2. H YBRID P ERSISTENT C OOPERATIVE I NITIALIZATION
We visualize our proposed initialization technique in Figure 2.
The method features two banks of persistent states.
One bank is for persistent images and the other bank is for persistent latent vectors.
There will always be a one-to-one pairing between persistent images and persistent latents, meaning that paired states will be drawn from the banks, return to the banks, and experience rejuvenation at the same time.
All samples are rejuvenated from the current generator.
Other than the pairing and generator rejuvenation source, sampling from banks and rejuvenation with a ﬁxed probability is conducted in the same way as persistent initialization.
This pairing is necessary for the cooperative learning loss used to train the generator.
Given a batch of paired samples { Z i } ni =1 and { X − i, 0 } ni = drawn from the latent and image bank respectively and a generator network g ( z ; φ ) , we learn φ using gradient descent on the reconstruction loss
[Equation]
where T ( X − i, 0 ; θ )
=
X − i represents a shortrun Langevin trajectory with the EBM p ( x ; θ ) .
This loss can be derived in the ML framework as a way to teach the output of g ( z ; φ ) to match the distribution of p ( x ; θ ) .
A full implementation of cooperative learning requires an additional sampling process on Z i , but we use a straightthrough estimator (Bengio et al., 2013) and approximate latent sampling with the identity function as originally done in the ofﬁcial MATLAB implementation from Xie et al.
We review further details of the cooperative learning loss in Appendix C.
Intuitively, the loss (1) encourages the output of the generator to match the outcome of the Langevin sampling process (5).
Cooperative learning always uses
X i, 0 = g
( Z i ; φ ) , which corresponds to rejuvenating with probability 1 in our method.
This is the root of the limitation of cooperative learning.
If g ( z ; φ ) has little diversity, T ( g ( z ; φ ); θ ) will also have little diversity.
This causes extreme oscillation in both the EBM and generator output as p ( x ; θ ) attempts to cover the modes of q ( x ) with shortrun samples from low diversity initialization.
Despite oscillation in appearance across φ , generator samples tend to remain nearly identical for any ﬁxed φ , thereby perpetuating the instability and preventing further learning.
By drawing samples from the image bank, we are effectively choosing X i,
= T (cid:48)
( g ( z ; φ (cid:48) )) where φ (cid:48) represents a past generator parameter and T (cid:48) ( x ) represents the composition of Langevin sampling with past EBMs.
This allows us to learn φ using appropriately diverse initializations spanning samples of several past models.
We observe that this simple adjustment has a dramatic effect for improving
EBM learning with generator initialization.
2. E XPERIMENTS : I MAGE S YNTHESIS WITH S HORTRUN MCMC
We present the results of our new learning process applied to the CIFAR-10, Celeb-A, and ImageNet benchmark datasets, using the image sizes 32 × 32 , 64 × 64 , and 128 × 128 respectively.
A sketch of the hybrid learning algorithm is given in Appendix D. Besides standard deep learning techniques for stable optimization, our framework is fully described by the ML objective and MCMC initialization
[Table]
above.
We use the SNGAN (Miyato et al., 2018) architectures for all models, where EBMs use the SNGAN discriminator with no normalization.
The generator has batch normalization for the
CIFAR-10 and Celeb-A experiments only.
Table 1 displays the FID scores achieved by our model in comparison with prior methods
[Table]
When calculating FID scores, samples are ﬁrst initialized from the generator then updated with the
EBM using a number of Langevin updates tuned to provide optimal synthesis quality.
As mentioned before, an appealing aspect of EBM learning with generator initialization is the ease of generating new images from scratch, in contrast with persistent initialization.
We use 50K samples with the ofﬁcial
FID code from Heusel et al.
(2017) to calculate all FID scores.
In particular, our scores and framework are consistent with the GAN replication library from Lee & Town (2020).
We publicly release the checkpoints, learning code, and FID code for each model.
The checkpoints are representative of what is achievable within an ordinary run of the code provided.
Our CIFAR-10 results show a signiﬁcant improvement over prior EBM synthesis and contribute to closing the gap between EBMs and other generative models.
Surprisingly, our ImageNet results surpass the results of GANs such as SNGAN (Miyato et al., 2018) and SSGAN (Chen et al., 2019) on unconditional synthesis using a similar magnitude of computational resources.
M IDRUN S AMPLERS FOR A DVERSARIAL D EFENSE
This section presents a method for learning EBMs that are capable of preserving the appearance of an in-distribution initial state across several thousand MCMC steps.
Such models are useful for the purpose of adversarial defense.
Our defense framework is based on the approach in Hill et al.
(2021), which uses an EBM to defend an independent naturally trained classiﬁer.
Appendix E brieﬂy reviews the EBM defense and compares this approach with other defense methods.
We then present our proposed method for learning a defensive EBM (3), which is based on persistent initialization using a ﬁxed pretrained generator as a source of rejuvenation.
Finally, we apply our defense to achieve state-of-the-art performance for puriﬁcation-based defense on CIFAR-10 and ImageNet.
One limitation of prior EBM defense is the reliance on persistent initialization with no rejuvenation to learn the defensive model.
This is done to ensure that defensive sampling trajectories remain stable for arbitrary numbers of steps.
However, removing rejuvenation from persistent learning has drawbacks discussed in Section 2.1.
In particular, it becomes very difﬁcult to learn meaningful EBMs for large and complex datasets such as ImageNet in this framework because the persistent bank cannot represent the diversity of the dataset and the quality of persistent images without rejuvenation quickly degrades.
Methods for efﬁcient learning of defensive EBMs at a greater scale are needed to extend the EBM defense to more realistic situations.
We propose to overcome this obstacle by building on the observation that fully stable sampling paths are not required for successful defense.
While the defense from Hill et al.
(2021) uses EBMs with stable sampling for 100K steps or more, the defense results require less than 2000 steps.
A natural question is whether it is possible to learn stable MCMC trajectories for only a predeﬁned midrun range to achieve the defensive beneﬁts without fully stabilizing samples over longrun trajectories.
Deﬁning a learning procedure to obtain such models is the goal of this section.
Efﬁciently learning
EBMs with stable midrun trajectories allows us to scale up the EBM defense to signiﬁcantly more challenging domains.
3. I N -D ISTRIBUTION R EJUVENATION FOR L EARNING A D EFENSIVE EBM
The initialization method for our midrun sampler is similar to standard persistent initialization with the adjustment that persistent states are rejuvenated from a frozen pretrained generator rather than noise.
We use a generator in our experiments so that the EBM and generator could be used after training to sample from scratch.
This choice is not necessary and rejuvenation from data samples, or another efﬁcient in-distribution initialization, is also effective for learning defensive EBMs.
We note that MCMC initialization from a trained generator or from data samples is not explored in recent work because the goal of most current EBM learning is image synthesis, which becomes trivial if EBM trajectories are always initialized from high-quality samples.
From the perspective of synthesis our learning process is nearly invisible but from the perspective of defense its utility becomes concrete.
and negative samples are
MCMC samples initialized from a batch from an image bank
Learning an EBM for defense involves tuning the length of the sampling trajectory via the number of shortrun training steps and the rejuvenation rate, and tuning the annealing schedule of the EBM optimizer.
Given a desired number of MCMC steps K def for a defensive update and a shortrun trajectory K ≈ 100 , we simply set the rejuvenation rate to p rejuv =
K/K def to ensure that on average samples will travel K def steps before rejuvenation.
While in practice we use K def = 2000 and p rejuv = 0 .
05 , we have found that this method can yield stable paths for at least K def =
50 K MCMC steps when p rejuv is low.
Initialization alone is insufﬁcient to stabilize MCMC pathways when model weights are changing quickly.
Using a low learning rate late in training is a key aspect of stabilizing sampling paths (Nijkamp et al., 2020; Hill et al., 2021).
Intuitively, if the EBM optimizer has a sufﬁciently low learning rate then MCMC trajectories in the persistent image bank can function as approximate trajectories from the current model, since weights change very little as the persistent states are updated.
By annealing in tandem with our initialization, we are effectively using midrun trajectories of length K def initialized from the generator to update the EBM while we are actually using shortrun trajectories of length K from the persistent bank.
Annealing is a crucial component for stabilizing both midrun and longrun trajectories.
Without annealing, sampling paths are not able to maintain realism for large K def .
3. E XPERIMENTS : D EFENDING N ATURAL C LASSIFIERS WITH AN EBM
We train our EBMs using the persistent initialization described above in tandem with a pretrained generator on both CIFAR-10 and ImageNet.
We use the same SNGAN models as before for our CIFAR-10 experiments, with the exception that the generator is pretrained instead of learned.
For our ImageNet experiments, we use the BigGAN (Brock et al., 2019) discriminator architecture for our EBM modiﬁed for input size 224 × 224 and a pretrained BigGAN Generator.
Our naturally trained classiﬁer f ( x ) is a pretrained WideResNet 28-10 (Zagoruyko & Komodakis, 2016) for CIFARand a pretrained EfﬁcientNetB-7 architecture (Tan & Le, 2020) for ImageNet.
We evaluate our models using the attack gradient (8) and Algorithm 2.
For CIFAR-10 we use
K def = [-
-]500 Langevin steps with l ∞ adversarial parameters ε = 8255 and α = 2255 , where ε is the size of the l ∞ ball and α is the gradient step size.
For ImageNet we used K def = 200 Langevin steps for defense with l ∞ adversarial parameters ε = 2 255 and α = 1 255 .
We attack ImageNet for 50 attacks
[Table]
steps across 10K val samples.
For CIFAR-10 we perform the same number of attacks across 4K validation samples.
[Table]
On CIFAR-
[Table]
we surpass the robustness of the existing EBM defense using a much more reliable learning framework.
The importance of midrun learning is clearly demonstrated by our successful application of EBM defense to ImageNet at the resolution 224 × 224 .
The robustness of a naturally trained classiﬁer secured by our EBM is comparable with adversarial training.
While the ImageNet results for EBM defense are not yet on par with state-of-the-art adversarially trained models, our experiments are an important proof of concept that the method can be scaled.
See Appendix F for diagnostics and further discussion.
L ONGRUN S AMPLING FOR D ENSITY E STIMATION
Our ﬁnal objective is to introduce scalable tools for learning a valid image density with an EBM.
Nijkamp et al.
(2020) revealed that, in the absence of careful implementation, EBM learning always results in an unexpected outcome where steady-state samples from the learned density p ( x ; θ ) have an oversaturated and unrealistic appearance that differs drastically from shortrun samples used during learning.
This outcome affects all EBMs that are not speciﬁcally trained to overcome this defect, as well as related models such as Score Matching (Song & Ermon, 2019; 2020) as shown in Figure 8.
To our knowledge, across different generative models, the only way to learn a well-formed potential energy surface approximating a complex density q ( x ) that is compatible with efﬁcient MCMC sampling is using an EBM.
Despite the theoretical formulation of the EBM as a potential surface, the problem of principled density estimation has received relatively little attention.
Successful image synthesis is sometimes used as misleading evidence of successful density estimation, but we emphasize these outcomes not equivalent.
In this section, we address the lack scalable methods to learn valid densities of complex signals by introducing an MCMC initialization can incorporate rejuvenation while still simulating extremely long trajectories.
Learning a valid density is a fundamental computational problem that is important in its own right, and we further hope that our learning method leads to EBM clustering techniques based on the potential energy basins of a well-formed density (Nijkamp et al., 2020).
4. I NCORPORATING R EJUVENATION IN D ENSITY E STIMATION
Prior work has suggested that persistent learning is the most effective method for learning a valid
EBM density.
Furthermore, works that learn a valid density have avoided rejuvenation because the incorporation of newly rejuvenated samples into the persistent bank ensures that EBM updates will always include samples that are not at the steady-state.
However, persistent learning without rejuvenation has shortcomings mentioned in Section 3.
We present hypothesized conditions for learning a valid density that motivate the design of our MCMC initialization:
• After a certain point in training, all samples used to update the EBM must be approximate steadystate samples of the current model p ( x ; θ ) .
• Persistent samples that are newly rejuvenated (up to about 50K Langevin steps since rejuvenation, and possibly many more) cannot be approximate steady-state samples for any known rejuvenation sources, including data, generators, and noise.
• Persistent samples that have undergone sufﬁciently many lifetime Langevin updates for a model whose weights are changing very slowly can be approximate steady-state samples.
Both the second and third point are corroborated by prior work (Nijkamp et al., 2020; Hill et al., 2021) as well as our own observations.
The third point means that persistent states updated with shortrun Langevin can eventually act as longrun Langevin samples if the optimizer learning rate is small, because the EBM samples in previous timesteps are essentially samples from the current EBM.
Learning a valid density that includes rejuvenation while satisfying the conditions above requires separating the newly rejuvenated samples from samples that are used to update the
EBM.
This leads us to introduce two persistent image banks: one for newly rejuvenated samples, and one for samples that will be used to update the EBM.
Samples in the newly rejuvenated bank that have been updated sufﬁciently many times will eventually replace samples from the bank used to update the EBM, at which point newly rejuvenated states will be added to the burn-in bank.
Figure 4 shows a visualization of the MCMC initialization method.
As in Section 3, we will use a pre-trained generator to efﬁciently obtain high-quality rejuvenated samples so that the generation process is fully synthetic, although data samples could be used as well.
Our goal is to preserve the sample quality for an arbitrary number of MCMC steps.
We note that this is a sufﬁcient but not necessary condition for learning a valid EBM density.
Nonetheless, stable sampling is an important step towards rigorous probabilistic EBMs. after 100K Langevin and 1M
Langevin steps.
FID remains stable over long trajectories.
[Table]
Even with our improved initialization, we ﬁnd that extremely longrun trajectories of 1 million or more MCMC steps still tend to oversaturate, although to a much lesser degree.
To further stabilize the appearance of extremely longrun Langevin samples, we include prior energy terms in the model.
Our longrun EBMs have the form U
[Equation] [Equation]
where U (cid:48)
( x ; θ ) is the model whose weights are updated, U 0 ( x ; θ 0 ) is a prior EBM with ﬁxed weights θ 0 and σ is a parameter controlling the strength of a Gaussian prior.
We used a prior EBM in a shortrun manner.
The role of the prior EBM is to provide some stability but also to provide a tendency to oversaturate at longer trajectories so that the current EBM learns to correct oversaturation.
The Gaussian prior is meant to discourage unbounded activations outside of the image hypercube.
Further discussion is in Appendix H. We ﬁnd that including both of these terms signiﬁcantly improves the ability to learn quality synthesis over long trajectories.
4. E XPERIMENTS : EBM D ENSITY E STIMATION
In these experiments, we learn EBM models that have stable MCMC trajectories for an arbitrary number of steps.
We apply the longrun learning method in Algorithm 4 to CIFAR-10, Celeb-A, and
ImageNet at the resolutions 32 × 32 , 64 × 64 , and 64 × 64 respectively.
Pre-trained SNGANs are used to rejuvenate images for CIFAR-10 and Celeb-A, while a resized images from a pre-trained
BigGAN are used to rejuvenate ImageNet samples.
We require that samples have been updated for 50K to 75K Langevin steps since the last rejuvenation before states are moved from the burn-in bank to the update bank.
Our EBMs all use the discriminator architecture of the SNGAN model without spectral normalization.
We use pretrained midrun samples for our prior EBM terms in the learned energy (2).
See Appendix K for an algorithm sketch.
Figure 1 shows the evolution of FID score over
Langevin steps starting from the generator network for our CIFAR-10 model.
Table 4 shows the
FID scores of 5K samples for each dataset.
Visualizations of longrun samples with 100K steps and extremely longrun samples using 1 million Langevin steps are shown in Appendix M. See Appendix J for a discussion of our evaluation in comparison to log likelihood.
R ELATED W ORK
Energy-Based Models.
Early forms of EBMs include the exponential family distribution, the
FRAME model (Zhu et al., 1998) and Restricted Boltzmann Machines (Hinton, 2002).
Recent work has introduced the EBM with a ConvNet potential (Xie et al., 2016; Du & Mordatch, 2019).
This dramatically increased the learning capacity of the model which led to many follow-up works on image synthesis (Gao et al., 2018; Lee et al., 2018; Nijkamp et al., 2019), adversarial robustness (Hill et al., 2021), and joint learning of discriminative and generative models (Grathwohl et al., 2020).
Several works investigate training and EBM in tandem with an auxilary model.
Kim & Bengio (2016) train an EBM and generator and tandem without MCMC by using samples from the generator as direct approximations of the EBM density and training the generator using a variational objective.
A similar approach is explored by Grathwohl et al.
Cooperative learning (Xie et al., 2018) trains the EBM and generator by using the generator to initialize samples needed to train the EBM and uses reconstruction loss between generator and EBM samples to learn the generator.
Gao et al.
(2020a) learn an EBM using Noise Contrastive Estimation with an auxiliary ﬂow model.
Xiao et al.
(2021) use a pretrained VAE to facilitate EBM learning.
Our work builds on cooperative learning by identifying and resolving symmetry breaking problems in early training, leading to state-of-the-art
EBM synthesis for unconditional ImageNet.
Despite the formulation of the EBM as an unnormalized density, it has been shown that most EBMs have strong misaligned steady-state distributions (Nijkamp et al., 2020).
Our work introduces new methods to learn a model with correct steady-state alignment.
Adversarial Robustness.
Adversarial Training (AT) (Madry et al., 2018), which trains a classiﬁer using PGD-generated adversaries, is the most popular and studied adversarial defense.
Many variations and improvements have been introduced, including optimizing the training loop by recycling gradients of past adversaries (Shafahi et al., 2019), combining single step FGSM with random initialization to achieve similar robustness Wong et al.
(2020), learning with auxiliary unlabeled data
Carmon et al.
(2019), local linearization (Qin et al., 2019), and the use of smooth activation functions
Xie et al.
An alternative approach to adversarial training involves the use of preprocessing transformations.
Randomized smoothing (Cohen et al., 2019) and related methods (Salman et al., 2020) add noise to the input signal to remove adversarial signals.
Many other preprocessing defenses have been proposed (Guo et al., 2018; Song et al., 2018;
Yang et al., 2019), but nearly all of these methods can be broken by adaptive attacks that are aware of the preprocessing method (Athalye et al., 2018).
A notable exception is the EBM defense (Hill et al., 2021), which uses midrun MCMC trajectories to purify images.
We ease the restriction of learning EBMs with stability for arbitrary
MCMC runs in the EBM defense by introducing a midrun sampler that enables faster learning of defensive EBMs and allows the EBM defense to scale to more complex datasets.
C ONCLUSION AND F UTURE W ORK
We have described three unique MCMC initializations for EBM using different sampling trajectories: shortrun for synthesis, midrun for defense, and long-run for density estimation.
Furthermore, we have elaborated on different MCMC initialization strategies used to stabilize these models for different sampling lengths.
We have demonstrated the ﬂexibility of these mechanisms by using similar architectures, data, and training platforms to create different EBMs for different applications.
We hope that future research incorporates these new training initialization schemes to improve their generative models for a wide variety of tasks.
R EPRODUCIBILITY
We have made our results reproducible by releasing the anonymized github repository https: //anonymous.4open.science/r/lifecycle-3346/ .
All hyperparameters can be found in the conﬁguration ﬁles, all checkpoints for the various EBMs are also under the release section and there is a README.MD describing how to use this repository to reproduce the results described herein.
We will release a non-anonymous repository at a later date.
All our results were trained on the standard CIFAR-10, ImageNet and Celeb-A datasets which are publicly available.
Our computational resources were primarily 5 TPUv2-8 machines and 5 TPUv3-8 machines made available by the
Tensorﬂow Research Cloud (TRC) program sponsored by Google.
These resources are widely available to researchers who apply to the program from accredited institutions.
We are very grateful to Google for granting us access to the resources that made this work possible.
A R EVIEW OF EBM L EARNING AND MCMC I NITIALIZATION
We brieﬂy review the main components of EBM learning following the standard method derived from works such as Hinton (2002); Zhu et al.
Xie et al. (2016).
Throughout our experiments, we will use this framework and focus on exploring different possibilities for MCMC initialization.
A deep EBM has the form
[Equation]
where U ( x ; θ ) is a ConvNet with weights θ and Z ( θ ) is the intractable normalizing constant.
Maximum Likelihood (ML) learning uses the objective arg min θ D KL ( q ( x ) || p ( x ; θ )) , which can be minimized using the stochastic gradient
[Equation]
where the positive samples { X + i } ni =1 are a set of data samples and the negative samples { X − i } ni = are samples from the current model p ( x ; θ ) .
To obtain the negative samples for a deep EBM, it is common to use MCMC sampling with the K steps of Langevin equation
[Equation]
where η is the step size and Z k ∼ N (0 , I ) .
The Langevin trajectories are initialized from a set of states { X − i, 0 } ni =1 obtained from a certain initialization strategy.
The method of MCMC initialization determines many of the properties of the learned model.
There are a number of different methods of MCMC initialization for EBM learning.
Early Restricted
Boltzmann Machines use data samples as the initial states in every iteration (Hinton, 2002).
Persistent initialization (Zhu et al., 1998; Tieleman, 2008) maintains a bank of MCMC samples from previous model updates as the source of MCMC states in the current iteration.
Persistent states can be rejuvenated with a certain probability instead of returning to the image bank.
Potential sources for rejuvenation include noise samples (Du & Mordatch, 2019), data augmentation applied to persistent states (Du et al., 2020), data samples, or samples from a generator network.
Cooperative learning initializes MCMC states from a generator learned in tandem with an EBM (Xie et al., 2018).
The Multi-grid model (Gao et al., 2018) initializes samples from lower resolution images.
Noise initialization is explored by Nijkamp et al.
The VAEBM (Xiao et al., 2021) initializes samples from a pretrained VAE generator.
EBMs can learn stable shortrun synthesis from virtually any initialization distribution as long as the initialization distribution remains fairly consistent across training updates.
EBM sampling can produce realistic images after training when the testing initialization closely matches the training initialization, while synthesis can fail for test initializations unseen during training.
(Nijkamp et al., 2020).
B C OMPARISON OF EBM AND O THER G ENERATIVE M ODELS
In this section, we brieﬂy compare the EBM with other kinds of generative models.
Besides the probabilistic model families mentioned in the introduction, we consider two other popular approaches.
One approach is to design an objective that promotes realistic synthesis by producing fake samples and distinguishing these from real samples, as done by the Generative Adversarial Network (GAN) model (Goodfellow et al., 2014) and its many variations.
Another approach comes from recent diffusion models that deﬁne a sequence of intermediate distributions to transport a signal between a reference density p 0 ( x ) and q ( x ) (Ho et al., 2020; Song et al., 2020).
Among probabilistic models, the synthesis quality is typically the highest for score matching models that only try to learn distribution gradients and typically lower for models like normalizing ﬂows and auto-regressive objectives that model the full normalized density.
EBM synthesis usually lies somewhere in between.
Score-matching does not require MCMC sampling during learning and as a consequence learning tends to be more straightforward than learning for EBMs.
Yet we believe that
MCMC sampling is not yet fully optimized for EBM learning and that EBMs can match or surpass score-matching models.
Our shortrun initialization method is an effort in this direction.
In fact, we ﬁnd cooperative learning is not possible at all without generator batch normalizatin.
Our hybrid initialization increases the diversity of shortrun samples by including persistent samples from previous EBM updates.
This accelerates learning and dramatically improves synthesis results.
GANs and diffusion models differ fundamentally from probabilistic models of the form p ( x ; θ ) in the sense that the former are explicitly formulated from the perspective of synthesis while the latter are formulated for density modeling and synthesis is a byproduct.
In the literature, the former types of models are sometimes referred to as generative models (in a more speciﬁc sense of the term than typical) while the latter are called descriptive models (Guo et al., 2003; Zhu, 2003).
Generative models map one set to another (e.g. a generator mapping a latent vector to an image), while descriptive models deﬁne a concept in terms of the set of states that have a certain descriptive property (e.g. steady-state samples of an EBM).
A descriptive model can be used for generative purposes (e.g.
Langevin sampling with an EBM)
but this is not the direct goal of descriptive learning.
From this perspective, it is perhaps unsurprising that generative methods such as GANs and diffusion models surpass EBMs in terms of synthesis, since that the model formulations are geared directly towards generation.
Nonetheless, we believe that EBMs have unique properties that distinguish them from generative models and other descriptive models.
In particular, EBM learning offers a control of the MCMC sampling paths from a learned model that p ( x ; θ ) is, to our knowledge, is unobtainable under other current frameworks.
Unlike GAN models which map a trivial distribution to a complex signal distribution without actually describing the signal distribution, EBMs directly model the complex signal distribution via a potential energy surface.
Normalizing ﬂow models are too large for efﬁcient MCMC sampling.
VAE models cannot be efﬁciently marginalized.
Score-based models exhibit the same oversaturation behavior as miscalibrated EBMs when sampling at a constant noise level for many steps (see Figure 8), and one cannot adjust MCMC stability during learning as we do for EBMs, because sampling is not used when training a score model.
Diffusion-based methods do not learn an approximation p ( x ; θ ) , but rather can be formulated as learning a conditional probability p ( x | x (cid:48) ; θ ) that describes a state x given a state x (cid:48) a rung above or below x in the ladder of distributions between p 0 ( x ) and q ( x ) (Gao et al., 2020b).
Auto-regressive models depend on a certain sequence of components for generation, and are not compatible with joint updates of all pixels.
On the other hand, MCMC sampling with an EBM is a natural and relatively efﬁcient process.
Uniquely among all methods, the desired stability and length of MCMC sampling trajectories for an EBM can easily be controlled by adjusting the sampling phase of Maximum Likelihood (ML) learning.
This can yield model capabilities not obtainable by other kinds of unsupervised models.
C C OOPERATIVE L EARNING FOR G ENERATOR N ETWORK
In this section we brieﬂy review the Maximum Likelihood framework for cooperatively learning the generator in tandem with the EBM (Xie et al., 2018).
The generator model assumes a joint distribution ( X, Z ) of images X and latent signals Z given by
Z ∼ N (0 , I ) and X | Z ∼ N ( g ( Z ; φ ) , τ 2 I ) for a generator network g with weights φ and a Gaussian parameter τ controls the spread of X around g ( Z ; φ ) .
Given i.i.d.
pairs { X i , Z i } ni =1 , the observations have a joint negative log probability − log p ( { X i ,
Z i } ni
=1 ; φ ) =
[Equation]
(cid:88) [Equation] In practice, the X i are shortrun MCMC samples from an EBM and the Z i are unobserved.
The full formulation of cooperative learning involves latent MCMC using the conditional distribution p ( Z
| X i ; φ ) to infer each Z i .
In practice, since the X i are shortrun samples that are paired with Z i through the process of rejuvenation, one can roughly approximate the latent MCMC process with an identity function and obtain a gradient (1).
We ﬁnd that this works as well or better than including a latent sampling step on Z , and it is consistent with the original implementation from Xie et al.
Figure 5 illustrates the importance of including a persistent image bank in the cooperative learning framework.
Shortrun samples from the generator in standard cooperative learning have low diversity and it is difﬁculty for the generator to break the symmetry of its activations early in training.
This can lead to catastrophic instability without auxiliary techniques like batch norm to help break the symmetry.
Incorporating persistent states from prior EBM updates when updating the generator helps the shortrun distribution achieve a health diversity that enables stable and effective learning.
D C OOPERATIVE -P ERSISTENT A LGORITHM FOR S HORTRUN L EARNING
[Table]
[Table]
E B ACKGROUND ON EBM D EFENSE
The most popular method for adversarial defense is adversarial training (AT) (Madry et al., 2018;
Wong et al., 2020; Shafahi et al., 2019) which aims to train a classiﬁer to correctly predict adversarial samples within a small ball around a natural input.
Another popular defense method is randomized smoothing (Cohen et al., 2019), which adds Gaussian noise to images to remove adversarial signals before classiﬁcation.
Both of these approaches modify classiﬁer training.
Although many methods have been proposed to secure a naturally trained classiﬁer, most have been broken by stronger attacks.
A recent method that has been shown to secure a natural classiﬁer is Langevin sampling with an EBM (Hill et al., 2021).
The EBM defense uses a classiﬁer trained with labelled natural images and an EBM trained with unlabelled natural images.
The two networks are trained independently, which is a key advantage of
EBM defense over adversarial training and randomized smoothing.
Since the EBM is independent of the classiﬁer, EBM defense has the potential to secure many classiﬁers across diverse tasks with a single defensive model, while existing methods are typically tailored to a single method.
Starting with a naturally classiﬁer trained on natural images f ( x ) , we deﬁne its robust counterpart as
[Equation]
) where T ( x ) is a random variable representing K steps of the Langevin transformation (5) initialized from a state x .
We cannot evaluate F ( x ) directly so we approximate it using
[Equation]
where
[Equation]
where f ( . ) is a forward pass of our classiﬁer to return logits and where the accuracy of approximation is driven by the number of replicates H .
Meaningful evaluation of adversarial defenses must be based on adaptive attack methods which are aware of both f ( x ) and T ( x ) .
For our attack we use the
BPDA+EOT formulation from (Athalye et al., 2018; Tramer et al., 2020) to obtain the attack gradient [Equation] BPDA
[Equation]
[Equation] y [Equation] =
[Equation]
which we use in the standard PGD framework to generate adversarial examples
Algorithm 2 gives a sketch of the defense evaluation. Algorithm
Require: Natural images { x + m } Mm =1 , EBM U ( x ; θ ) , classiﬁer f , Langevin noise η , attack replicates H adv , defense replicates H def , l ∞ radius ε , attack step size, α , Langevin steps K Ensure: Defense record { D m } Mm =1 for each image initialized as ones.
for 1 ≤ i ≤ M do select ( X i , y i ) from batch
Randomly initialize adversary ˆ X 0 inside L ∞ ball around X
i for 1 ≤ j ≤ N do c j = arg max (cid:96) [
ˆ F H adv ( ˆ X j − 1 )] (cid:96) ∆
j = ∆ BPDA+EOT ( ˆ X j − 1 , y i )
if c j (cid:54) =
y
i then c (cid:48)
j = arg max (cid:96) [ ˆ F H def ( ˆ X j − 1 )] (cid:96) if c (cid:48)
j (cid:54) =
y i then D i =end if end if ˆ X j = PGD ( ˆ X j − 1 , ∆ j , ε, α ) end for end for
F EBM D EFENSE E XPERIMENT D ETAILS AND D IAGNOSTICS
[Table]
To verify the integrity of our results we ran an attack with heavily increased resources for ImageNet and CIFAR-10 compared to our standard evaluation.
While using these resources for all attacks is infeasible in practice, we want to ensure our defense maintains robustness as attacker resources increase.
As shown in Table 5, our benchmark accuracy remains consistent when we increase the number of attack steps (from 50 to 200 ) and EOT attack replicates ( H adv ).
We also demonstrate results over varying numbers of langevin steps K .
We can see in Fig 6 that our sampling trajectory for defending imagenet at K = 200 is reasonable to achieve high natural image classiﬁcation as well as robustness.
On CIFAR-10 we are able to reduce the number of langevin steps from 1500 in Hill et al.
(2021) to K = 500 , which greatly reduces our compute overhead and demonstrates that a non-convergent model can outperform a convergent model ( +13% robustness) while using less steps.
[Figure]
G I MPORTANCE OF L EARNING R ATE A NNEALING
This section demonstrates the importance of learning rate annealing for learning a robust energy landscape.
We repeat the midrun and longrun learning experiments for CIFAR-10 except that we never anneal the learning rate.
We then sample with the models for 1500 steps for the model trained with the midrun method and 100K steps for the model trained with the longrun method.
The results in Figure 7 show that learning rate annealing is essential for stabilizing both midrun and longrun trajectories.
The importance of annealing can be understood as follows. If the EBM is being updated with a very low learning rate, then samples from recent EBM snapshots can function as samples from the current EBM. In the case of midrun trajectory, annealing allows the model to robustify trajectories that are approximately as long as the lifetime of a persistent sample between rejuvenation.
In the case of longrun learning, annealing allows the burnin samples to approximately reach the model steady-state before they are included in the update bank.
This allows the persistent samples in the update bank to function as approximate steady-state samples from the current EBM, leading to proper modeling of probability mass.
H P RIOR EBM FOR L ONGRUN L EARNING
This section discusses the choice of training longrun EBMs by using a ﬁxed shortrun EBM as a prior distribution.
Intuitively, the prior EBM learns a reasonable but imperfect approximation of the energy landscape that we know will have energy basins which leak to low-energy states.
Using a ﬁxed misaligned prior EBM allows the EBM which is being actively updated to focus on learning landscape features which were not correctly learned by the prior EBM.
In particular, our longrun allows the actively updated EBM to focus on sealing the leaky energy basins of the prior EBM so that the full model learn the correct distribution of probability mass.
By separating the burn-in and update banks and longrun learning, we guarantee that the prior EBM alone would lead to extreme oversaturation by the time a state reaches the update bank unless the actively updated EBM is counteracting the oversaturation tendency to preserve realism.
This additional effect of directly driving states towards the oversaturated regions in the absence of a correcting force leads to improved stability for longrun models learned with a prior EBM compared to models learned without a prior EBM.
The technique has precedent in ResNets.
In both cases, it is easier to learn an approximate direct model ﬁrst and then to approximate the residual that remains.
I C OMPARISON TO S CORE -B ASED D ENOISERS , N ORMALIZING F LOWS , AND D IFFUSION M ODELS
The score-based model from (Song & Ermon, 2020) and its annealed Langevin dynamics process has recently been used for purifying adversarial signals (Lee & Oh, 2021; Yoon et al., 2021).
One approach is add noise and the using the score model to denoise (Lee & Oh, 2021).
The robustness of this method is upper-bounded by standard randomized smoothing (Cohen et al., 2019).
Another approach is to initiate the Langevin process of the score model from a natural image as done in the
EBM defense.
A score model can be used in a langevin process T ( x ) that is a direct analogue to the
EBM langevin process in (5) and (8).
Given a score model S θ ( x ) for a low noise value σ , one can deﬁne the Langevin equation
[Equation]
where S θ ( x ) ≈ ∇ x log q ( x ) since σ is low.
In both cases, the score-based model is trained with the loss:
[Equation]
where q σ i (˜ x | x ) represents adding Gaussian noise of standard deviation σ i to an image x sampled from q .
We experimented with this process as a defense mechanism by selecting the smallest trained noise value σ = 0 .
01 and using the Langevin process as a method to purify adversaries.
We evaluated this method using our BPDA+EOT attack framework over different numbers of Langevin steps during puriﬁcation.
In contrast to reports from Yoon et al. (2021), we were unable to obtain any signiﬁcant defense using a pretrained score model when initializing sampling directly from adversarial or natural images.
In Figure 8, one can see that the score-based model drives natural images toward saturation quickly, leading to a sharp decrease in natural classiﬁcation that undermines the possibility of robustness from sampling.
Since the score-based model does not perform sampling during training, and one cannot adjust the stability of its sampling process as we do in this work.
While it is not immediately clear how to overcome this problem, we believe that defense with a score model is possible and we hope that our observation lead to efforts to stabilize the sampling paths of score models as we do for EBMs in this work.
To underscore our claims about the difﬁculty of calibrating the probability mass of a density model, we further investigate longrun MCMC samples from a normalizing ﬂow and diffusion model. We ﬁnd that the normalizing ﬂow from the GLOW model (Kingma & Dhariwal, 2018) and the recovery likelihood diffusion model (Gao et al., 2020b) have misaligned steady-states as well. This shows that the problem of improper density estimation extends well beyond the EBM. Tractable density modeling with a normalizing ﬂow does not prevent steady-state misalignment.
These experiments corroborate our claim that log likelihood experiments in previous works are not able to detect the misaligned probability mass of many prior models.
We believe that the calibration of the model steady-state is currently best diagnosed with longrun MCMC sampling because the distribution of longrun MCMC samples represents the probability mass of the model.
Figure 9 ( left ) displays initial and ﬁnal states from a GLOW model density after 100K sampling steps.
Despite the fact that the GLOW model has a fully tractable density, it is unable to learn a valid distribution of probability mass.
Figure 9 ( right ) shows initial and ﬁnal samples from the Recovery
Likelihood T6 model after 100K steps of the conditional model at the lowest noise value.
We observe the same oversaturation for the conditional density as for the unconditional density of a standard
EBM.
Code for the T1K model that the authors evaluate in their longrun experiments is not released so we have not yet been able to directly test their results.
The T1K model is equivalent to an EBM version of the score model in Figure 8 which we have shown has a misaligned steady-state.
Further, we note the longrun experiments with the T1K model are very misleading because the experiments use 100 steps with 1000 distinct conditional models and claim this is a longrun evaluation of 100K steps.
The correct evaluation is to use 100K steps on a single conditional model.
We strongly believe that the Recovery Likelihood model as originally presented has a misaligned steady-state like many other methods.
We hope that the observations in our work can lead to efforts to stabilize the sampling trajectories of many existing models.
J R EGARDING D ENSITY E STIMATION AND L OG L IKELIHOOD
In this work, we are primarily interested in learning density models that assign the majority of probability mass in realistic regions of the image space.
This goal is, surprisingly, distinct from the goal of likelihood maximization.
In particular, one can achieve high likelihood with mixture models where only an inﬁnitesimal portion of mass is assigned to a mixture component that approximates the true density, while the majority of probability mass is assigned to a degenerate distribution (see
Theis et al.
(2016), “Great Log Likelihood but Poor Samples”).
The steady-state of MCMC sampling with such a mixture distribution would concentrate on the degenerate distribution even though the structures of the true density exist in high-energy regions.
Therefore, log likelihood cannot detect if a model has assigned probability mass in realistic region of the image space.
Recent observations have shown that this situation is not just hypothetical (Nijkamp et al., 2020).
Non-convergent EBMs are practical examples, since these models consist of a mixture of partially formed high-energy energy basins enabling effective shortrun synthesis in realistic regions of the image space and much lowerenergy basins in unrealistic image regions that dominate the probability mass.
The same misaligned landscape structure extends beyond EBM to RBMs (Decelle et al., 2021), score models, normalizing ﬂows, and diffusion models (see Appendix I).
Log likelihood can only indicate the presence of energy landscape features that are similar to the ground truth landscape, but it cannot detect whether these features will leak into lower-energy basins that represent the true mass distribution.
Like shortrun sample quality, high log likelihood is often misleading false evidence that a model density concentrates on realistic images.
Overall, we believe that FID using true samples from the EBM and data samples is an appropriate, if rough, measure of successful density modeling.
The subtlety involves generating true samples from the EBM.
We believe that longrun MCMC sampling is the most practical and principled tool for obtaining approximate samples from the EBM steady-state, and our investigations show that density modeling with properly calibrated probability mass is currently most effectively accomplished using
MCMC-based EBM training.
Though we use the FID metric, our density modeling experiments are best described as realism preservation experiments rather than image generation experiments.
Require: Natural images { x + m } Mm =1 , EBM U ( x ; θ ) , frozen pre-trained generator g ( z ) Langevin noise η , Langevin steps K , EBM optimizer h U , initial weights θ 0 , update threshold D , number of training iterations
T .
Ensure:
Weights θ T for EBM stable longrun samples.
Initialize burn bank
Initialize
T do
Select
Get do
Randomly
Rejuvenate end end
[Figure]
S HORTRUN S YNTHESIS R ESULTS
M L ONGRUN S YNTHESIS R ESULTS
[Figure]
N T ABLES OF E XPERIMENT P ARAMETERS
This section gives a full list of experiment parameters.
Some notable minor additions beyond what are discussed in the text are a ﬁxed MCMC temperature parameter to regulate gradient strength at the beginning of training, a gradient clipping parameter to restrict the magnitude of EBM and generator weight updates.
We only use gradient clipping for shortrun experiments.
We use the annealing schedule γ anneal =
[(10 − 4 , 0) , (10 − 5 , 50000) , (10 − 6 , 75000) , (10 − 7 , 100000) , (10 − 8 , 125000)] in all experiments that involve annealing, where the pairs denote a learning rate and model update step at which that learning rate is ﬁrst used.
We denote this schedule as γ anneal in the tables.
Shortrun Training
Dataset
Celeb-A
CIFARImageNet
Training Steps
Data Epsilon 1e1e- 1eEBM LR 1e1e- 1eEBM Optimizer
Adam
Adam
Adam
EBM Gradient Clip
Langevin Epsilon 3e5e- 3eMCMC Steps
Rejuvenation Probability 0.
MCMC Temperature 1e1e- 1eMax Update Rounds
Persistent Bank Size
Generator LR 1e1e- 5eGenerator Optimizer
Adam
Adam
Adam
Generator Gradient Clip
Generator Batch Norm
No
Yes
No
Shortrun Evaluation
Dataset
Celeb-A
CIFARImageNet
Langevin Epsilon 3e5e- 3eMCMC Steps
MCMC Temperature 1e1e- 1e-
Defense Training
Dataset
CIFARImageNet
Training Steps
Data Epsilon 2.5e3e-
EBM LR γ anneal γ anneal EBM Optimizer
Adam
Adam
Langevin Epsilon 1.25e2e-
MCMC Steps
Rejuvenation Probability 0.
MCMC Temperature 1e1e-
Persistent Bank Size
Defense Evaluation
Dataset
CIFARImageNet
Adversarial Steps
Adversarial Epsilon
Adversarial Eta
EOT Attack Reps
EOT Defense Reps
Langevin Steps
Langevin Epsilon 1.25e2e-
MCMC Temperature 1e1e-
Longrun Training
Dataset
Celeb-A
CIFARImageNet
Training Steps
Data Epsilon 2e2e- 2eEBM
LR γ anneal γ anneal γ anneal EBM Optimizer
Adam
Adam
Adam
Langevin Epsilon 1e1e- 1eMCMC Steps
Burn-in Update Rounds
MCMC Steps Burn-in
MCMC Temperature 1e1e- 1eTau 1.5e1.5e1.5ePersistent Bank Size
Burn-in Bank Size
Prior EBM
Dataset
Celeb-A
CIFARImageNet
Training Steps
Data Epsilon 2e2e- 1.5eEBM LR 1e1e- 1e-
[Table]
1e
[Table]
1e
[Table]
1e
[Table]
[Table]
Steps
[Table]
[Table]
[Table]
[Table]
Temperature 1e
[Table]
1e
[Table]
1e
[Table]
Tau 1.5e
[Table]
1.5e
[Table]
1.5e
[Table]
[Table]