Imagining the Latent Space of a Variational Auto-Encoders

1 INTRODUCTION . Variational auto-encoders ( VAEs ) have made a significant impact since their introduction by Kingma and Welling ( 2014 ) . However , one of their perceived problems is their reconstruction performance . This has spawned a wave of research into trying to improve the reconstruction performance ( Zhao et al. , 2017 ; Dai and Wipf , 2019 ; Larsen et al. , 2016 ; Gao et al. , 2017 ; Brock et al. , 2017 ) . We argue that such attempts are misguided . The whole point of VAEs is to capture only compressible information and discard information specific to any particular image . This is a consequence of the well known evidence lower bound or ELBO objective function consisting of a negative log-probability of generating the original image from the latent representation ( this is often implemented as a mean squared error between the image and the reconstruction , although as we argue in Appendix A this term should be proportional to the logarithm of the mean squared error ) and a KL-divergence between the probability distribution representing a latent code and a ‘ prior distribution ’ ( usually taken as a multivariate normal with mean zero and unit variance ) . These two terms have a nice interpretation in terms of the minimum description length ( Rissanen , 1978 ) —this has been described elsewhere , for example , Chen et al . ( 2016 ) . The KL-term can be viewed as a measure of the amount of information in the latent code while the log-probability of the image measures the amount of information required to change the image produced by the decoder into the input image ( see Section 3 for details ) . That is , the latent space of a VAE can be viewed as a model of the dataset—capturing compressible information while not encoding any image specific information ( which is cheaper to communicate using the reconstruction loss ) . The great strength of a VAE is that it builds a model of the dataset that does not over-fit ( i.e . code for in-compressible features found in specific images ) . However , because of this it typically will not do a good job of reconstructing images as the latent code does not contain enough information to do the reconstruction ( for very restrictive dataset such as MNIST and Celeb-A a lot of information can be captured in the latent space , but for more complex datasets like ImageNet or CIFAR the reconstructions are poor ) . Of course , if you want good reconstructions on the training set then the simplest solution is to remove the KL-divergence term and just use an autoencoder . However , having a model that does not over-fit the dataset can be useful , but in this case the decoder of a standard VAE should not be regarded as a generative model—that is not its purpose . If we wish to generate realistic looking images we need to imagine the information discarded by the encoder . As a rather simplified analogy , consider a verbal description of an image “ a five year old girl in a blue dress standing on a beach ” . If we asked different artists to depict such scene there is clearly not enough information to provide pixel-wise or feature-wise similarity between their interpretation although each artist could render a convincing image that satisfies the description . In a similar manner if we want a VAE to act as a generative model we need to build a renderer that will imagine an image consistent with the latent variable representation . A simple way to achieve this is using a modified Generative Adversarial Network ( GAN ) . We call such a model a latent space renderer-GAN ( or LSR-GAN ) . To generate an image we choose a latent vector z from the prior distribution for the VAE . This is passed to a generator network that generates an image , x̂ , with the same dimensions as that of the dataset used to train the VAE . The generated image has both to convince a discriminator network that it is a real image—as is usual for a GAN ( Goodfellow et al. , 2014 ) —at the same time the VAE encoder should map x̂ close to z . To accomplish this we add an additional cost to the normal GAN loss function for the generator ( LGEN ) LGEN − λ log ( qφ ( z|x̂ ) ) ( 1 ) where qφ ( ·|x̂ ) is the probability distribution generated by the VAE encoder given an image x̂ and z is the latent vector that was put into the GAN generator . Note that when training the LSR-GAN we freeze the weights of the VAE encoder . The constant λ is an adjustable hyperparameter providing a trade-off between how realistic the image should look and how closely it captures the information in the latent space . This modification of the objective function can clearly be applied to any GAN or used with any VAE . Although the idea is simple , it provides a powerful method for visualising ( imagining ) the information stored in a latent space . Interestingly , it also appears to provide a powerful regularisation mechanism to stabilize the training for GANs . Combinations of VAEs and GANs are , of course , not new ( Makhzani et al. , 2016 ; Larsen et al. , 2016 ; Brock et al. , 2017 ; Huang et al. , 2018 ; Srivastava et al. , 2017 ) . In all cases we are aware of GANs have been combined with VAEs to “ correct ” for the poor reconstruction performance of the VAE ( see Appendix B for a more detailed discussion of the literature on VAE-GAN hybrids ) . As we have argued ( and expound on in more detail in Section 3 ) , we believe that the decoder of a VAE does the job it is designed to do . They can not reconstruct images accurately , because the latent space of a VAE loses information about the image , by design . All we can do is imagine the type of image that a point in the latent space represents . In the next section , we show examples of images generated by the LSR-GAN for both normal VAEs and β-VAEs ( we also spend time describing VAEs , β-VAEs and the LSR-GAN in more detail ) . In addition , in this section we present a number of systematic experiments showing the performance of a VAE and LSR-GAN . In Section 3 , we revisit the minimum description length formalism to explain why we believe a VAE is doomed to fail as a generative model . We conclude in Section 4 . We cover more technical aspects in the appendices . In Appendix A we show that the correct loss function for a VAE requires minimising a term proportional to the logarithm of the mean squared error . In Appendix B we draw out the similarities and differences between our approach to hybridising VAEs with GANs and other work in this area . We present some additional experimental results in Appendix C. A detailed description of the architecture of LSR-GAN is given in Appendix D. We end the paper with Appendix E by showing some samples generated by randomly drawing latent variables and feeding them to the LSR-GAN . 2 IMAGINING LATENT SPACES . A natural question to ask is what information about an image gets represented in the latent space of a VAE . To answer this we can use the VAE encoder to generate a distribution qφ ( z|x ) representing that image in the latent space ( see Sections 2.1 for details on VAEs ) . From this distribution we can sample points in the latent space and feed this to the LSR-GAN generator . We show examples of this for both CIFAR-10 and ImageNet ( down-sampled to 64 × 64 ) in Figure 1 . In all cases in this paper the input images are taken from a test set that is independent of the training set . Note that both CIFAR-10 and ImageNet are “ hard ” for VAEs in the sense that they represent extremely diverse sets of images . As a consequence , the VAE latent space will struggle to store detailed information about the images and the VAE reconstructions will be poor . We have repeated this for a β-VAE ( see section 2.3 for a full description of β-VAEs ) . We note that there is very little variation between the different samples drawn from qφ ( z|x ) , particularly for the standard VAE ( β = 1 ) , showing that the latent space of the VAE is relatively smooth ( there is more variation when β = 20 ) . To get a sense of the variation in the information stored in latent spaces we show in Figure 2 inputoutput pairs , where the left image is the input and right image is the output generated by the LSRGAN generator seeded with a latent vector encoding of the input image . The reconstructions capture the shape and background , but clearly loses a lot of detail . In some cases it appears that the type of object is being captured , although in the case of the boat with the β-VAE ( with β = 20 ) the wrong object is being rendered . 2.1 VARIATIONAL AUTOENCODERS . The structure of a VAE is represented schematically below . x ∼ D Encoder ( parameters φ ) ( µφ , σ 2 φ ) z ∼ q ( z|x , φ ) Decoder ( parameters θ ) x̂ = Dθ ( z ) We sample an input x from some dataset , D. To be concrete we will consider the case where the inputs are images , although clearly a VAE can be used to represent many different types of data . For each input x the encoder outputs a mean vector , µ , and standard deviation vector , σ , that describes an axis aligned normal distribution , qφ ( z|x ) = N ( z ∣∣µφ ( x ) , diag ( σ2φ ( x ) ) ) . A latent variable z is sampled from this distribution and then fed to a decoder . For simple black and white datasets such as MNIST the decoder outputs a scalar at each pixel location that can be interpreted as the probability that the pixel is black . For more complex datasets the decoder ususal generates a “ reconstruction ” x̂ = Dθ ( z ) . The probability of generating a pixel value xi is then usually taken as a normal distribution with mean x̂i ( i.e . pθ ( xi|z ) = N ( xi|x̂i , σ2 ) ) and variance σ2 that measures the expected size of the errors between the input images , x , and the reconstructions , x̂ . The loss function for a VAE is equal to the negative evidence lower bound ( ELBO ) L = −Ex∼D [ Ez∼qφ ( z|x ) [ log ( pθ ( x|z ) ) ] + KL ( qφ ( z|x ) ∥∥N ( 0 , I ) ) ] . ( 2 ) As explained in Appendix A , log ( pθ ( x|z ) ) is chosen to be proportional to the logarithm of the reconstruction error between x̂ and the input image x—in our experiments this produced better reconstructions than replacing log ( pθ ( x|z ) ) with the mean squared error . 2.2 LSR-GAN . LSR-GAN is a novel hybridization of VAE and GAN model . The most distinct difference of LSRGAN from previous work is that it is a two-stage model . In the first stage we train the VAE model . Having done this we freeze the weights of the VAE and train the GAN . We train the discriminator , D , of LSR-GAN in the same way as a normal GAN . That is , we minimise a loss function LD = −Ex [ log ( D ( x ) ) ] − Ez [ log ( 1−D ( G ( z ) ) ) ] ( 3 ) where G is the generator or the decoder of LSR-GAN . The job of the discriminator , D is , to decide whether its import is a real image or not . Thus , to optimise the loss function we neet to maximize the log-probability of passing the real data , x , while minimising the log-probability of accepting a random sampling G ( z ) generated by a generator G seeded with a random latent vector z . The architecture of the generator is the same as that of a normal GAN but the loss function is slightly different . We add an additional term giving LG = Ez [ log ( D ( G ( z ) ) ) ] + λ log ( qφ ( z|G ( z ) ) ) . ( 4 ) The parameters of the discriminator and generator are trained in the usual tick-tock fashion using gradient descent . We built the VAE and the generator of GAN using a ResNet ( He et al. , 2016 ) as it gave slightly better performance than using a standard CNN . The architecture of the discriminator is the same as DCGAN ( Radford et al. , 2016 ) . The architecture is described in Appendix D. To test the LSR-GAN we use the VAE to generate a latent representation z for an image drawn from an independent test set . The latent vector is then used as a seed value for the generator in the LSRGAN . The LSR-GAN can get sharper reconstruction images than VAE ( see Figure 3 ) . Although not visually so obvious , we have used a quantitative measure of sharpness computed as luminancenormalised Laplacian ( San Pedro and Siersdorfer , 2009 , Section 3.1.2 ) . For the reconstructed images from the VAE we obtained a measure of 0.17± 0.03 while for the LSR-GAN we obtain 0.28± 0.08 ( i.e . an improvement of a factor of two ) . We have also computed the FID measure of image quality for CIFAR-10 ( Heusel et al. , 2017 ) . For images seeded from a testing example the VAE achieved a score of 89.8 while LSR-GAN achieved a score of 44.1 , while for images seeded with random latent variable ( i.e . z ∼ N ( 0 , I ) ) the FID score for the VAE is 138.6 while for the LSR-GAN it is 47.4 . This should not be surprising . The decoder of the VAE is training only where there are training images . Despite the fact that the KL-divergence tries to ensure that as much latent space as possible is used , the constraint of minimising the reconstruction loss means that most of the latent space is far from a training example . Although the VAE does not do too badly generating testing examples , these tend to be substantially closer in the latent space to the training examples than random samples . In contrast the LSR-GAN is trained on random samples so that the generator will have to produce “ realistic ” images over the whole latent space . Of course , whether these generated images represents anything recognisable is open to question . For diverse training sets such as CIFAR-19 and ImageNet this may be very difficult . What image should we expect from a latent vector halfway between a truck and a bird ? In Appendix E we show images generated by seeding LSR-GAN with random latent variables for CIFAR-10 , ImageNet , MNIST and Celeb-A .

Summary: This paper proposes a hybrid VAE-GAN model, called the latent space renderer-GAN (LSR-GAN), with the goal to “imagine” the latent space of a VAE, and to improve the decoding and sampling quality of a VAE. First, a VAE-like model is trained, after which the encoder weights are frozen, and the decoder is trained as the generator of a GAN (together with an auxiliary discriminator). The generator loss also contains a reconstruction-like term in the latent space, described by the negative log density of the encoding distribution of the original latent conditioned on the output of the generator: -log q(z|g(z)).