Learning Energy-Based Generative Models via Coarse-to-Fine Expanding and Sampling

1 INTRODUCTION Recently , energy-based models ( EBMs ) ( Zhu et al. , 1998 ; LeCun et al. , 2006 ) parameterized by modern neural networks have drawn much attention from the deep learning communities . Successful applications with EBMs include generations of images ( Xie et al. , 2016 ; 2018b ; Du & Mordatch , 2019 ) , videos ( Xie et al. , 2017 ; 2019 ) , 3D volumetric shapes ( Xie et al. , 2018c ; 2020 ) , unordered point clouds ( Xie et al. , 2021a ) , texts ( Deng et al. , 2020 ) , molecules ( Ingraham et al. , 2018 ; Du et al. , 2019 ) , etc. , as well as image-to-image translation ( Xie et al. , 2021b ; c ) , out-of-distribution detection ( Liu et al. , 2020 ) and inverse optimal control ( Xu et al. , 2019 ) . EBMs are characterized by ( i ) Simplicity : The maximum likelihood learning of EBMs unifies representation and generation in a single model , and ( ii ) Explicitness : EBMs provide an explicit density distribution of data by training an energy function that assigns lower values to observed data and higher values to unobserved ones . However , it is still difficult to train an EBM to synthesize diverse and high-fidelity images . The maximum likelihood estimation ( MLE ) of EBMs requires the Markov chain Monte Carlo ( MCMC ) ( Liu , 2008 ; Barbu & Zhu , 2020 ) to sample from the model and then updates the model parameters according to the difference between those samples and the observed data . Such an “ analysis by synthesis ” ( Grenander et al. , 2007 ) learning scheme is challenging because the sampling step is neither efficient nor stable . In particular , when the energy function is multimodal due to the highly varied or high resolution training data , it is not easy for the MCMC chains to traverse the modes of the learned model . Fortunately , it is common knowledge that the manifold residing in a downsampled low-dimensional image space is smoother than that in the original high-dimensional counterpart . Thus , learning an EBM from low-dimensional data is much stabler and faster than learning from high-dimensional data in terms of convergence ( Odena et al. , 2017 ; Gao et al. , 2018 ) . Inspired by the above knowledge , we propose to train EBMs via a multistage coarse-to-fine expanding and sampling strategy ( CF-EBM ) . As shown in Figure 1 ( a ) , the approach starts with learning a coarselevel EBM on low resolution images and then smoothly transits to learn the finer-level EBM by adding new layers that take into account the higher resolution information as the learning progresses . The gradient-based short-run MCMC ( Nijkamp et al. , 2019 ) , e.g. , Langevin dynamics ( Neal et al. , 2011 ) , is used for sampling . From the modeling aspect , the coarse-level training can be useful for exploring the global structure of image , while the fine-level training will then gradually refine the image details . Recent works have demonstrated the advantages of this incremental learning ( Karras et al. , 2018 ; Wang et al. , 2018 ) . However , there have been no works focusing on the incremental learning of EBMs that incorporates bottom-up representation and top-down sampling in a single net . Besides , as shown in Figure 1 ( a ) , the top-down gradient information for synthesis flows from coarse-level layers towards fine-level layers . Thus , during the coarse-to-fine expanding , we can use the coarse-level synthesis to help the fine-level synthesis to stabilize the sampling . Such a coarse-to-fine expanding and sampling scheme is useful for high-fidelity synthesis in several vision tasks . See Figure 1 ( b ) . Furthermore , we propose a one-sided energy-based unsupervised image-to-image translation method and scale it up to high resolution . The approach is immediately available with the FC-EBM by using its iterative Langevin dynamics without the need of the cycle consistency ( Zhu et al. , 2017 ) or geometry constraints ( Fu et al. , 2019 ) . Specifically , we learn an EBM of target domain with Langevin dynamics initialized by the examples from source domain . The resulting translator is the short-run MCMC . Compared with those prior works ( Zhu et al. , 2017 ; Huang et al. , 2018 ; Park et al. , 2020 ) that learn black-box encoder-decoder networks between domains , our method is much more interpretable in the sense that ours can be explained by a visualization method ( Simonyan et al. , 2014 ; Adebayo et al. , 2018 ) that uses gradients to visualize the most essential regions , i.e. , the generative saliency , when translating an image from the source domain to the target domain . See Figure 1 ( c ) . The contributions of our paper can be summarized as : • To the best of our knowledge , this is the first work that trains EBMs under the “ analysis by synthesis ” scheme via a multistage coarse-to-fine expanding and sampling strategy . Besides , we propose several essential techniques for improving EBM , e.g. , smooth activations . Particularly , our work is the first to train a pure EBM for synthesizing 512× 512 images . • We propose a novel energy-based unsupervised image-to-image translation approach , which is essentially different from all other existing GAN-based approaches . We demonstrate noticeable results in terms of both translation quality and efficiency of time and memory . • We conduct extensive experiments to validate our approach , including image generation , denoising , inpainting , out-of-distribution detection and unsupervised image translation . Strong results show that our method outperforms or is competitive with the prior art . The rest of the paper is organized as follows . Section 2 summarizes related works and how the proposed method is different from the prior art . Section 3 introduces the proposed methodology in detail . Section 4 presents extensive experiments to test our method . Section 5 concludes our paper and discuss some future research directions . 2 RELATED WORK . 2.1 ENERGY-BASED GENERATIVE MODELS . The main challenge to train EBMs via MLE lies in drawing fair samples from the model , especially when the energy function is parameterized by a highly non-linear ConvNet . The contrastive divergence ( CD ) ( Hinton , 2002 ; Tieleman , 2008 ) , with MCMC chains initialized from data distribution , is an efficient but biased way to train EBMs . Another direction is to adopt the idea of energy-based correction of a more tractable model to train EBMs . Noise contrastive estimation ( NCE ) ( Gutmann & Hyvärinen , 2010 ; Gao et al. , 2020 ) and introspective neural networks ( INNs ) ( Lazarow et al. , 2017 ; Jin et al. , 2017 ; Lee et al. , 2018b ) belong to this theme . Generative cooperative networks ( CoopNets ) ( Xie et al. , 2018b ; 2021d ; b ) train an EBM with a generator or a variational auto-encoder ( VAE ) ( Kingma & Welling , 2014 ) as amortized sampler by MCMC teaching ( Xie et al. , 2018a ) . Triangle divergence ( Han et al. , 2019 ) trains an EBM without MCMC by amortizing the MCMC via a VAE . However , these frameworks still struggle to scale up and model multimodal data . There have been several strategies to improve the EBM training . Gao et al . ( 2018 ) adopts a multi-grid method that trains multiple EBMs at different grids simultaneously , where the EBM at coarser grid is used to initialize the image generation by EBM at finer grid . However , optimizing and sampling from multiple EBMs will result in low efficiency of both time and memory . To stabilize the training , Nijkamp et al . ( 2019 ) ; Grathwohl et al . ( 2020 ) add the Gaussian white noise to the observed data , resulting in noisy synthesized images . In contrast , our paper proposes to train a single EBM via a coarse-to-fine growing strategy , along with some improved techniques . With smooth parameter training and image sampling , our model can preserve EBM ’ s compatibility and synthesize high-fidelity images . 2.2 UNSUPERVISED IMAGE-TO-IMAGE TRANSLATION . Unsupervised image-to-image translation aims at learning two directions of mappings between two unpaired domains . Recent successes are all based on adversarial learning , e.g. , CycleGAN ( Zhu et al. , 2017 ) , UNIT ( Liu et al. , 2017 ) , MUNIT ( Huang et al. , 2018 ) , DRIT ( Lee et al. , 2018a ) and U-GATIT ( Kim et al. , 2020 ) . These methods typically train two GANs with two levels of learning objectives : ( i ) Distribution level : Two adversarial losses are used to capture style discrepancy between source and target domain ; ( ii ) Instance level : To tackle the difficulty of unpaired setting , they adopt a cycle consistency loss for content preservation . This loss enables an instance-level supervision to regularize the training of two mappings by enforcing them to be a bijective function between two domains . Except for works about two-sided translation , efforts on research about one-sided unsupervised image translation have also been made , e.g. , DistanceGAN ( Benaim & Wolf , 2017 ) , GcGAN ( Fu et al. , 2019 ) and CUT ( Park et al. , 2020 ) , which apply geometric or contrastive constraints . We solve this problem from the prospective of EBM , which is different from GAN-based methods . The proposed concise EBM solution only relies on its built-in objective , which is a distribution-level statistics matching , to accomplish the one-sided image translation . It transfers the style and preserves the source content by MCMC without using the cycle-consistency loss . The model demonstrates better performances with less time and memory . Another distinction between our method and GAN-based methods is the natural interpretability of Langevin dynamics . It provides a gradient-based saliency map to visualize those key regions that make the two domains distinct , as illustrated in Figure 1 ( c ) . 3 METHOD . In this section , we first present the EBM learning framework , and then the proposed CF-EBM approach . After that , we generalize our model to the unsupervised image-to-image translation . 3.1 MCMC-BASED MAXIMUM LIKELIHOOD LEARNING OF ENERGY-BASED MODEL . Let x ∈ RD be the observed example , e.g. , an image . An energy-based model is defined as follows : pθ ( x ) = 1 Z ( θ ) exp ( −Eθ ( x ) ) , ( 1 ) where Eθ ( x ) : RD −→ R is the energy function defined by a bottom-up ConvNet parameterized by θ . Z ( θ ) = ∫ exp ( −Eθ ( x ) ) dx is the intractable normalizing constant or the partition function . Given N observed examples { xi } Ni=1 ∼ pdata ( x ) , where pdata ( x ) denotes the unknown data distribution , the model can be trained by maximizing the log-likelihood L ( θ ) = 1N ∑N i=1 log pθ ( xi ) ≈ Ex∼pdata ( x ) log ( pθ ( x ) ) . The derivative of the negative log-likelihood is given by −∇θL ( θ ) = Ex∼pdata ( x ) [ ∇θEθ ( x ) ] − Ex̃∼pθ ( x ) [ ∇θEθ ( x̃ ) ] , ( 2 ) where the second expectation term under pθ ( x ) is intractable and can be approximated via MCMC . Given that , the EBM is updated by gradient descent . To sample x̃ ∼ pθ ( x ) via MCMC , we rely on gradient-based Langevin dynamics that recursively computes the following step x̃t+1 = x̃t − ηt 2 ∇x̃Eθ ( x̃t ) + √ ηt t , t ∼ N ( 0 , I ) , ( 3 ) where ηt is the step size of Langevin step and also the variance of Gaussian noise t. Theoretically , to ensure convergence , the MCMC is typically performed with infinite steps and an infinitesimal stepsize ( Welling & Teh , 2011 ) . However , it is impractical for training EBMs . In this paper , we follow Nijkamp et al . ( 2019 ) to use short-run MCMC , which always starts from a fixed noise distribution and runs a fixed number T of Langevin steps in both training and testing stages . The training with a short-run MCMC might result in a biased estimation of EBM but the learned short-run MCMC is still a valid generator , which enables us to synthesize realistic images and efficiently train the model , as seen in most well-established EBM works ( Nijkamp et al. , 2019 ; Grathwohl et al. , 2020 ; Pang et al. , 2020 ) . In this paper , we keep the step size constant and linearly decay the noise variance till 0 .

Summary: Much like progressive growing of GANs two years ago, this paper adopts a similar coarse-to-fine procedure for scaling EBMs to higher resolutions. In particular, the approach starts from learning EBMs on low-resolution images and then smoothly transitions to higher resolution by carefully designing an expand layer and a smooth sampling procedure. Authors were able to obtain competitive FID scores on CIFAR-10, and demonstrate the first set of 256x256 image samples from EBMs. In addition, authors demonstrate successful application of EBMs to unpaired image-to-image translation.