EBM Life Cycle: MCMC Strategies for Synthesis, Defense, and Density Modeling

1 INTRODUCTION . Generative modeling of complex signals is one of the fundamental challenges of computational cognition . The importance of studying generative models is often justified 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 flexible model p ( x ; θ ) , where θ is a model parameter . This approach is taken by Energy-Based Models ( EBMs ) ( Xie et al. , 2016 ) , normalizing flows ( 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 classifiers , and mature samples represent grounded knowledge of the data density ( or lack thereof in the widespread case of a misaligned steady-state ) . Figure 1 illustrates 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 difficulty of density estimation . Learning midrun sampling trajectories sufficient 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 first 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 difficulty 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 . ( 2018 ) ; 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 significantly 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 purification-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 . 2 HYBRID PERSISTENT COOPERATIVE LEARNING FOR IMAGE SYNTHESIS . 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 . ( 2020 ) ; Alayrac et al . ( 2019 ) . Our method involves hybridizing persistent ( Tieleman , 2008 ) and cooperative ( Xie et al. , 2018 ) initialization . We first briefly cover the strengths and weaknesses of these methods , then discuss our hybridization , and finally present experimental results . 2.1 MOTIVATION : LIMITATIONS OF PERSISTENT AND COOPERATIVE INITIALIZATION . 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 can not efficiently represent the diversity of the data . Initialization from a cooperative generator network is an appealing alternative because it could enable efficient 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 difficulty 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 efficient way of generating new samples from scratch after training that match the samples used during training . In contrast , models trained with persistent initialization often can not 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.2 HYBRID PERSISTENT COOPERATIVE INITIALIZATION . 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 fixed 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 { Zi } ni=1 and { X − i,0 } ni=1 drawn from the latent and image bank respectively and a generator network g ( z ; φ ) , we learn φ using gradient descent on the reconstruction loss L ( φ ) ∝ n∑ i=1 ‖g ( Zi ; φ ) − T ( X−i,0 ; θ ) ‖ 2 2 ( 1 ) Figure 2 : Cooperative-persistent initialization uses paired latent and image states that a drawn from persistent banks to learn the EBM and generator . 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 Zi , but we use a straightthrough estimator ( Bengio et al. , 2013 ) and approximate latent sampling with the identity function as originally done in the official MATLAB implementation from Xie et al . ( 2018 ) . 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 Xi,0 = g ( Zi ; φ ) , 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 fixed φ , thereby perpetuating the instability and preventing further learning . By drawing samples from the image bank , we are effectively choosing Xi,0 = T ′ ( g ( z ; φ′ ) ) where φ′ represents a past generator parameter and T ′ ( 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 .

Summary: The paper describes techniques to train EBMs according to the length of MCMC trajectories required. The paper explores three applications of EBMs: image synthesis, adversarial defense, and density estimation. Each of these applications uses a different regime of MCMC sampling length. An approach is described for modifying the initialization of MCMC both during training and at test-time to improve the performance in each of these applications.  For image synthesis, the paper proposes using a combination of persistent chains and cooperative learning. Persistent chains with rejuvenation can cause instability since negative samples can vary between freshly rejuvenated samples and longer run samples. However, no rejuvenation often leads to lack of sample diversity. Persistent banks don’t scale well to large datasets since they cannot capture all the variability of the larger dataset. Instead, chains can be rejuvenated from a fast generator model, which is trained jointly with the EBM, known as cooperative learning. The authors identify and fix a key issue with cooperative learning. The issue is that generators early in training are unable to generate a diverse set of samples that the EBM can refine. One fix is to use batch normalization in the generator. Instead, the authors propose using a persistent bank over the joint distribution of latent noise and generated images (where the latent noise is input to the generator), which the sampling process is rejuvenated from with the usual process.