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: This paper presents a number of methods to scale up training and sampling of EBMs on image data. The main contribution consists of an approach for progressively growing the model by increasing the image resolution as training progresses. This approach echos similar approaches used for scaling up GAN training. The approach involves slowly annealing in new blocks to the model during training which processes the image at increasing resolutions. This allows training of image EBMs at notably larger resolutions than published in prior work.