Controllable and Compositional Generation with Latent-Space Energy-Based Models

1 Introduction . Deep generative learning has made tremendous progress in image synthesis . High-quality and photorealistic images can now be generated with generative adversarial networks ( GANs ) [ 17 , 28 , 29 ] , score-based models [ 23 , 51 ] , and variational autoencoders ( VAEs ) [ 32 , 54 ] . However , a key requirement for the success of these generative models in many real-world applications is the controllability of the generation process . Controllable generation , in particular the compositional ability to generate novel concept combinations , still remains a great challenge to current methods . A common approach for achieving controllable generation is to train a conditional model [ 36 , 31 , 10 , 41 , 38 ] , where the conditional information , such as semantic attributes , is specified during training to guide the generation . There are two major drawbacks with this approach : ( i ) Since the set of attributes is fixed and pre-defined , it is difficult to introduce new attributes to an existing model . This , in turn , limits the model ’ s compositional ability for unseen attributes and their novel combinations [ 12 ] . ( ii ) Training conditional generative models from scratch for new attributes is computationally expensive , and it can be challenging to reach to the same image quality of the uncontrollable counterparts [ 10 , 41 ] . Recent works have attempted to overcome these issues by first training an unconditional generator , and then converting it to a conditional model with a small cost [ 26 , 20 , 1 , 42 ] . This is often achieved by discovering semantically meaningful directions in the latent space of the unconditional model . This way , one would pay the most computational cost only once for training the unconditional model . However , these approaches often struggle with compositional generation , in particular with rare 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . combination of the attributes and the introduction of new attributes . An appealing solution to the compositionality problem is to use energy-based models ( EBMs ) for controllable generation [ 12 , 18 , 13 ] . This is due to the fact that energy functions representing different semantics can be combined together to form compositional image generators . However , existing approaches train EBMs directly in the pixel space , making them slow to sample and difficult to generate high-resolution images [ 13 ] . In this paper , we leverage the compositionality of EBMs and the generative power of state-of-the-art pre-trained models such as StyleGANs [ 28 , 29 ] to achieve high-quality controllable and compositional generation . Particularly , we propose an EBM that addresses the problem of controlling an existing generative model instead of generating images directly with EBMs or improving the sampling quality . Our main contributions are summarized as follows : • We propose a novel formulation of a joint EBM in the latent space of pre-trained generative models for controllable generation . In particular , the EBM training reduces to training a classifier only , and it can be applied to any latent-variable generative model . • Based on our EBM formulation , we propose a new sampling method that relies on ordinary different equation ( ODE ) solvers . The sample quality is robust to hyperparameters . • We show that our method is fast to train and efficient to sample , and it outperforms previous methods in both conditional sampling and sequential editing . • We show strong compositionality with zero-shot generation on unseen attribute combinaions , and with logical compositions of energy functions in generating high-quality images . Specifically , we build an EBM in the joint space of data and attributes where the marginal data distribution is denoted by an implicit distribution ( e.g. , a pre-trained GAN generator ) , and the conditional distribution of attributes , given data , is represented by an attribute classifier . Using the reparameterization trick , we show that our EBM formulation becomes an energy function in the latent space where the latent distribution is known ( i.e. , a standard Gaussian ) . This way , we only need to train the classifier in the data space and do the sampling in the latent space . Because our method only requires training a classifier to add controllability , it is conceptually simple and fast-to-train . For sampling , most existing EBMs rely on the Langevin dynamics ( LD ) sampling which is often computationally expensive and sensitive to the choice of sampling parameters . Instead , we build on the recent score-based model [ 51 ] , and define the sampling process in the latent space using the corresponding probability flow ODEs , induced by the reverse diffusion process . Thus , we rely on an ODE solver for sampling from our EBM formulation in the latent space . With its adaptive step size , the ODE sampling in the latent space is efficient and robust to the sampling parameters . In experiments , we show our method achieves higher image quality and better controllability in both the conditional sampling and sequential editing , compared to various baselines , including StyleFlow [ 1 ] and JEM [ 18 ] . For the training time , we show our method is 25⇥ faster than StyleFlow on the FFHQ data . For the inference time , our sampling is at least 49⇥ and 876⇥ faster than EBMs [ 18 , 12 ] and score-based models [ 51 ] in the pixel space , respectively , on CIFAR-10 . More importantly , our method excels on zero-shot generation with unseen attribute combinations where StyleFlow almost completely fails . By composing energy functions with logical operators [ 12 ] , our method is the first to show such compositionality in generating photo-realistic images of resolution 1024⇥1024 . 2 Method . In this section , we will first give an overview of EBMs and then propose our new EBM formulation in the latent space , based on which , we introduce a new sampling method through the ODE solver . 2.1 Energy-based models . EBMs [ 34 ] represent data x 2 Rd by learning an unnormalized probability distribution pθ ( x ) / e Eθ ( x ) , where Eθ ( x ) : R d ! R is the energy function , parameterized by a neural network . To train an EBM on a data distribution pdata with maximum likelihood estimation ( MLE ) , we can estimate the gradient of the data log-likelihood L ( ✓ ) = Ex⇠pdata [ log pθ ( x ) ] as [ 22 ] : rL ( ✓ ) = Ex⇠pdata [ rθEθ ( x ) ] Ex⇠pθ [ rθEθ ( x ) ] ( 1 ) To sample x from pθ for training and inference , Langevin dynamics ( LD ) [ 57 ] is applied as follows , x0 ⇠ p0 ( x ) , xt+1 = xt ⌘ 2 rxEθ ( xt ) + ✏t , ✏t ⇠ N ( 0 , ⌘I ) ( 2 ) where p0 is the initial distribution , ⌘ is the step size . When ⌘ ! 0 and t ! 1 , xt is guaranteed to follow the probability pθ ( x ) under some regularity conditions [ 50 ] . 2.2 Modelling joint EBM in the latent space . To formulate our EBM for controllable generation , we consider the following setting : We have the data x 2 X ✓ Rd and its attribute code c = { c1 , · · · , cn } 2 C ✓ R n , which is a n-dimensional vector with each entry ci representing either a discrete mi-class attribute where ci 2 { 0 , · · · , mi 1 } , or a continuous attribute where ci 2 R. The goal is to learn a generative model and semantically control its generated samples by manipulating its conditioning attribute code c. To begin with , we define the joint generative model on both the data and attribute as : pθ ( x , c ) : = pg ( x ) pθ ( c|x ) / pg ( x ) e Eθ ( c|x ) ( 3 ) where pg ( x ) is an implicit distribution defined by a pre-trained generator g ( such as GANs ) in the form x = g ( z ) with the latent variable z 2 Z . And , pθ ( c|x ) is conditional distribution on attributes given x , modeled by a conditional energy function Eθ ( c|x ) . In this paper , we assume that the generator g is fixed and we only train the energy function Eθ ( c|x ) . We also assume the attributes are conditionally independent , i.e. , Eθ ( c|x ) = Pn i=1 Eθ ( ci|x ) . Since , our goal is to preserve the image generation quality , we design the joint in Eq . ( 3 ) such that the marginal data distribution pθ ( x ) from our joint satisfies pθ ( x ) = pg ( x ) . This is obtained with an energy function that is normalized up to a constant . Thus , we define : Eθ ( ci|x ) = ⇢ fi ( x ; ✓ ) [ ci ] + log P ci exp ( fi ( x ; ✓ ) [ ci ] ) if ci is discrete , 1 2σ2 ( ci fi ( x ; ✓ ) ) 2 if ci is continuous ( 4 ) where fi ( x ; ✓ ) is the output of a multi-class classifier mapping from X to R mi if the i-th attribute is discrete or a regression network mapping from X to R if it is continuous . Without loss of generality , we will always call fi ( x ; ✓ ) a classifier . Here , 2 is a hyperparameter to adjust the continuous energy . In practice , we set 2 = 0.01 after normalizing all the continuous attributes to [ 0 , 1 ] . We can not use Eq . ( 3 ) for sampling , since pg ( x ) is defined implicitly in GAN generators . However , in Appendix A.1 , we show that using the reparameterization trick [ 6 , 58 ] , sampling x from pθ ( x , c ) is reparameterized to sampling z from pθ ( z , c ) / e Eθ ( c|g ( z ) ) +log p ( z ) , and then transferring the z samples to the data space using generator x=g ( z ) . In most generative models , the prior distribution p ( z ) is a standard Gaussian . Thus , the joint becomes pθ ( z , c ) /e Eθ ( z , c ) with the energy function : Eθ ( z , c ) = X i Eθ ( ci|g ( z ) ) + 1 2 kzk22 ( 5 ) where each conditional energy function Eθ ( ci|g ( z ) ) is given by Eq . ( 4 ) . For sampling , we can run LD of Eq . ( 2 ) in the z-space using the energy function above . At the end of a chain , we pass the final z samples to the generator g to get the x samples . Since our latent-space EBM is formulated for controllable generation , we term it as LAtent Controllable EBM ( LACE ) throughout this work . Inspecting the joint energy Eθ ( z , c ) in Eq . ( 5 ) shows that when the unconditional generator g is fixed , the only trainable component is the classifier fi ( x ; ✓ ) that represents Eθ ( ci|x ) . Therefore , unlike the previous joint EBMs [ 18 ] that train both pθ ( x ) and pθ ( c|x ) , our method only needs to train a classifier for the conditional pθ ( c|x ) , and does not require sampling with LD in the x-space during training . That is , for controllable generation , we can train the classifier only in the x-space and do the LD sampling in the z-space . This fact makes our method conceptually simple and easy to train . Note that the considered x-space for training the classifier does not have to be the pixel space . Instead , we can choose any intermediate layer of the pre-trained generator as the x-space . Take StyleGAN2 [ 29 ] as an example , if we train the classifier in the w-space , then g ( z ) in Eq . ( 5 ) actually corresponds to the mapping network of StyleGAN2 , and once we obtain the w samples via the LD sampler , we can pass them to the synthesis network of StyleGAN2 to get the final images . In other words , since data x is a deterministic function of z , the distribution of attributes is uniquely defined by z . We have the freedom to define pθ ( ci|z ) by applying the classifier on top of z or any representation extracted from z including the generator g ( z ) or its intermediate representations such as w-space in StyleGAN2 .

Summary: This paper studies conditional learning by training an EBM in the latent space of a pre-trained top-down generator such as StyleGANs. The EBM is a joint distribution of data and attributes together, and sampling from it is formulated as solving an ordinary differential equation. Experimental results show that the method outperforms the state-of-the-art in both conditional sampling and sequential editing. The contribution of the paper lies on obtaining state-of-the-art performance of conditional learning by combining existing technologies, such as StyleGAN, latent space energy-based learning, and score-based generative model via stochastic differential equations.