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: The paper considers the task of obtaining additional control over a pre-trained generative model without re-training it. Specifically, it aims to generate conditional samples from an unconditional StyleGAN by only training a classifier with respect to the conditioning information. To achieve this, it considers the conditional distribution of StyleGAN latent codes. By Bayes Theorem, this conditional distribution is proportional to the unconditional distribution of the latent codes (which has a known standard normal distribution) times the distribution of conditioning information given the corresponding sample. The latter is modeled as a Gibbs distribution with an energy derived from the classifier. Two approaches to sample from this unnormalized distribution are evaluated: (i) Langevin sampling and (ii) Performing gradient descent on the classifier-derived energy. Experiments demonstrate good performance of this approach with significant improvements over StyleFlow for combinations of conditionings which are not in the training data.