Generalized Energy Based Models

1 INTRODUCTION . Energy-based models ( EBMs ) have a long history in physics , statistics and machine learning ( LeCun et al. , 2006 ) . They belong to the class of explicit models , and can be described by a family of energies E which define probability distributions with density proportional to exp ( −E ) . Those models are often known up to a normalizing constantZ ( E ) , also called the partition function . The learning task consists of finding an optimal function that best describes a given system or target distribution P. This can be achieved using maximum likelihood estimation ( MLE ) , however the intractability of the normalizing partition function makes this learning task challenging . Thus , various methods have been proposed to address this ( Hinton , 2002 ; Hyvärinen , 2005 ; Gutmann and Hyvärinen , 2012 ; Dai et al. , 2019a ; b ) . All these methods estimate EBMs that are supported over the whole space . In many applications , however , P is believed to be supported on an unknown lower dimensional manifold . This happens in particular when there are strong dependencies between variables in the data ( Thiry et al. , 2021 ) , and suggests incorporating a low-dimensionality hypothesis in the model . Generative Adversarial Networks ( GANs ) ( Goodfellow et al. , 2014 ) are a particular way to enforce low dimensional structure in a model . They rely on an implicit model , the generator , to produce samples supported on a low-dimensional manifold by mapping a pre-defined latent noise to the sample space using a trained function . GANs have been very successful in generating high-quality samples on various tasks , especially for unsupervised image generation ( Brock et al. , 2018 ) . The generator is trained adversarially against a discriminator network whose goal is to distinguish samples produced by the generator from the target data . This has inspired further research to extend the training procedure to more general losses ( Nowozin et al. , 2016 ; Arjovsky et al. , 2017 ; Li et al. , 2017 ; Bińkowski et al. , 2018 ; Arbel et al. , 2018 ) and to improve its stability ( Miyato et al. , 2018 ; Gulrajani et al. , 2017 ; Nagarajan and Kolter , 2017 ; Kodali et al. , 2017 ) . While the generator of a GAN has effectively a low-dimensional support , it remains challenging to refine the distribution of mass on that support using pre-defined latent noise . For instance , as shown by Cornish et al . ( 2020 ) for normalizing flows , when the latent distribution is unimodal and the target distribution possesses multiple disconnected low-dimensional components , the generator , as a continuous map , compensates for this mismatch using steeper slopes . In practice , this implies the need for more complicated generators . ∗Correspondence : michael.n.arbel @ gmail.com . In the present work , we propose a new class of models , called Generalized Energy Based Models ( GEBMs ) , which can represent distributions supported on low-dimensional manifolds , while offering more flexibility in refining the mass on those manifolds . GEBMs combine the strength of both implicit and explicit models in two separate components : a base distribution ( often chosen to be an implicit model ) which learns the low-dimensional support of the data , and an energy function that can refine the probability mass on that learned support . We propose to train the GEBM by alternating between learning the energy and the base , analogous to f -GAN training ( Goodfellow et al. , 2014 ; Nowozin et al. , 2016 ) . The energy is learned by maximizing a generalized notion of likelihood which we relate to the Donsker-Varadhan lower-bound ( Donsker and Varadhan , 1975 ) and Fenchel duality , as in ( Nguyen et al. , 2010 ; Nowozin et al. , 2016 ) . Although the partition function is intractable in general , we propose a method to learn it in an amortized fashion without introducing additional surrogate models , as done in variational inference ( Kingma and Welling , 2014 ; Rezende et al. , 2014 ) or by Dai et al . ( 2019a ; b ) . The resulting maximum likelihood estimate , the KL Approximate Lower-bound Estimate ( KALE ) , is then used as a loss for training the base . When the class of energies is rich and smooth enough , we show that KALE leads to a meaningful criterion for measuring weak convergence of probabilities . Following recent work by Chu et al . ( 2020 ) ; Sanjabi et al . ( 2018 ) , we show that KALE possesses well defined gradients w.r.t . the parameters of the base , ensuring well-behaved training . We also provide convergence rates for the empirical estimator of KALE when the variational family is sufficiently well behaved , which may be of independent interest . The main advantage of GEBMs becomes clear when sampling from these models : the posterior over the latents of the base distribution incorporates the learned energy , putting greater mass on regions in this latent space that lead to better quality samples . Sampling from the GEBM can thus be achieved by first sampling from the posterior distribution of the latents via MCMC in the low-dimensional latent space , then mapping those latents to the input space using the implicit map of the base . This is in contrast to standard GANs , where the latents of the base have a fixed distribution . We focus on a class of samplers that exploit gradient information , and show that these samplers enjoy fast convergence properties by leveraging the recent work of Eberle et al . ( 2017 ) . While there has been recent interest in using the discriminator to improve the quality of the generator during sampling ( Azadi et al. , 2019 ; Turner et al. , 2019 ; Neklyudov et al. , 2019 ; Grover et al. , 2019 ; Tanaka , 2019 ; Wu et al. , 2019b ) , our approach emerges naturally from the model we consider . We begin in Section 2 by introducing the GEBM model . In Section 3 , we describe the learning procedure using KALE , then derive a method for sampling from the learned model in Section 4 . In Section 5 we discuss related work . Finally , experimental results are presented in Section 6 with code available at https : //github.com/MichaelArbel/GeneralizedEBM . 2 GENERALIZED ENERGY-BASED MODELS . In this section , we introduce generalized energy based models ( GEBM ) , that combine the strengths of both energy-based models and implicit generative models , and admit the first of these as a special case . An energy-based model ( EBM ) is defined by a set E of real valued functions called energies , where eachE∈E specifies a probability density over the data spaceX ⊂Rd up to a normalizing constant , Q ( dx ) =exp ( −E ( x ) −A ) dx , A=log Å∫ exp ( −E ( x ) ) dx ã . ( 1 ) While EBMs have been shown recently to be powerful models for representing complex high dimensional data distributions , they still unavoidably lead to a blurred model whenever data are concentrated on a lower-dimensional manifold . This is the case in Figure 1 ( a ) , where the ground truth distribution is supported on a 1-D line and embedded in a 2-D space . The EBM in Figure 1 ( d ) learns to give higher density to a halo surrounding the data , and thus provides a blurred representation . That is a consequence of EBM having a density defined over the whole space , and can result in blurred samples for image models . An implicit generative model ( IGM ) is a family of probability distributions Gθ parametrized by a learnable generator function G : Z 7→X that maps latent samples z from a fixed latent distribution η to the data space X . The latent distribution η is required to have a density over the latent space Z and is often easy to sample from . Thus , Sampling from G is simply achieved by first sampling z from η then applyingG , x∼G ⇐⇒ x=G ( z ) , z∼η . ( 2 ) GANs are popular instances of these models , and are trained adversarially ( Goodfellow et al. , 2014 ) . When the latent spaceZ has a smaller dimension than the input spaceX , the IGM will be supported on a lower dimensional manifold of X , and thus will not possess a Lebesgue density on X ( Bottou et al. , 2017 ) . IGMs are therefore good candidates for modelling low dimensional distributions . While GANs can accurately learn the low-dimensional support of the data , they can have limited power for representing the distribution of mass on the support . This is illustrated in Figure 1 ( b ) . A generalized energy-based model ( GEBM ) Q is defined by a combination of a base G and an energyE defined over a subsetX of Rd . The base component can typically be chosen to be an IGM as in ( 2 ) . The generalized energy component can refine the mass on the support defined by the base . It belongs to a class E of real valued functions defined on the input spaceX , and represents the negative log-density of a sample from the GEBM with respect to the base G , Q ( dx ) =exp ( −E ( x ) −AG , E ) G ( dx ) , AG , E=log Å∫ exp ( −E ( x ) ) G ( dx ) ã , ( 3 ) where AG , E is the logarithm of the normalizing constant of the model w.r.t . G. Thus , a GEBM Q re-weights samples from the base according to the un-normalized importance weights exp ( −E ( x ) ) . Using the latent structure of the base G , this importance weight can be pulled-back to the latent space to define a posterior latent distribution ν , ν ( z ) : =η ( z ) exp ( −E ( G ( z ) ) −AG , E ) . ( 4 ) Hence , the posterior latent ν can be used instead of the latent noise η for sampling from Q , as summarized by Proposition 1 : Proposition 1 . Sampling from Q requires sampling a latent z from ν ( 4 ) then applying the mapG , x∼Q ⇐⇒ x=G ( z ) , z∼ν . ( 5 ) In order to hold , Proposition 1 does not need the generatorG to be invertible . We provide a proof in Appendix C.1 which relies on a characterization of probability distribution using generalized moments . We will see later in Section 4 how equation ( 5 ) can be used to provide practical sampling algorithms from the GEBM . Next we discuss the advantages of GEBMs . Advantages of Generalized Energy Based Models . The GEBM defined by ( 3 ) can be related to exponential tilting ( re-weighting ) ( Siegmund , 1976 ; Xie et al. , 2016 ) of the base G. The important difference over classical EBMs is that the base G is allowed to change its support and shape in space . By learning the base G , GEBMs can accurately learn the low-dimensional support of data , just like IGMs do . They also benefit from the flexibility of EBMs for representing densities using an energyE to refine distribution of mass on the support defined by G , as seen in Figure 1 ( c ) . Compared to EBMs , that put mass on the whole space by construction ( positive density ) , GEBMs have the additional flexibility to concentrate the probability mass on a low-dimensional support learned by the base G , provided that the dimension of the latent spaceZ is smaller than the dimension of the ambient space X : see Figure 1 ( c ) vs Figure 1 ( d ) . In the particular case when the dimension of Z is equal to the ambient dimension andG is invertible , the baseG becomes supported over the whole space X , and GEBM recover usual EBMs . The next proposition further shows that any EBM can be viewed as a particular cases of GEBMs , as proved in Appendix C.1 . Proposition 2 . Any EBM with energyE ( as in ( 1 ) ) can be expressed as a GEBM with baseG given as a normalizing flow with density exp ( −r ( x ) ) and a generalized energy Ẽ ( x ) =E ( x ) −r ( x ) . In this particular case , the dimension of the latent is necessarily equal to the data dimension , i.e . dim ( Z ) =dim ( X ) . Compared to IGMs , that rely on a fixed pre-determined latent noise distribution η , GEBMs offer the additional flexibility of learning a richer latent noise distribution . This is particularly useful when the data is multimodal . In IGMs , such a GANs , the latent noise η is usually unimodal thus requiring a more sophisticated generator to distort a unimodal noise distribution into a distribution with multiple modes , as shown by Cornish et al . ( 2020 ) . Instead , GEBMs allow to sample from a posterior ν over the latent noise defined in ( 4 ) . This posterior noise can be multimodal in latent space ( by incorporating information from the energy ) and thus can put more or less mass in specific regions of the manifold defined by the base G. This allows GEBMs to capture multimodality in data , provided the support of the base is broad enough to subsume the data support Figure 1 ( c ) . The base can be simpler , compared to GANs , as it doesn ’ t need to distort the input noise too much to produce multimodal samples ( see Figure 8 in Appendix G.4 ) . This additional flexibility comes at no additional training cost compared to GANs . Indeed , GANs still require another model during training , the discriminator network , but do not use it for sampling . Instead , GEBMs avoid this waist since the base and energy can be trained jointly , with no other additional model , and then both are used for sampling .

Summary: This paper proposes a framework called GEBM that combine an implicit generator and an EBM to define a probabilistic model on low dimensional manifold. Specifically, the implicit generator defines the base distribution, and the EBM refines the base. Finally, this method is equivalent to define an EBM on the latent space of the implicit generator together with a mapping from the latent space to the data space. The authors propose to use the KALE to train the generator, and provide a theoretical guarantee about the validness of the KALE.