Conditional Image Generation by Conditioning Variational Auto-Encoders

1 INTRODUCTION . A major challenge with applying variational auto-encoders ( VAEs ) to high-dimensional data is the typically slow training times . For example , training a state-of-the-art VAE ( Vahdat & Kautz , 2020 ; Child , 2020 ) on the 256× 256 FFHQ dataset ( Karras et al. , 2019 ) takes on the order of 1 GPU-year , but a state-of-the-art generative adversarial network ( GAN ) ( Lin et al. , 2021 ; Karras et al. , 2020 ) can be trained on the same dataset in a matter of GPU-weeks . One hypothesis for the cause of this disparity is that , whereas the “ mass-covering ” training objective for a VAE forces it to assign probability mass over the entirety of the data distribution , a GAN can “ cut corners ” by dropping modes ( Arora & Zhang , 2017 ; Arora et al. , 2017 ) . We focus on the problem of conditional generative modelling : given an input ( e.g . a partially blanked-out image ) , we wish to map to a distribution over outputs ( e.g . plausible completions of the image ) . Both conditional GANs ( Zheng et al. , 2019 ; Zhao et al. , 2021 ) and conditional VAEs ( Sohn et al. , 2015 ; Ivanov et al. , 2018 ) are applicable to this problem , with the same disparity in training times that we described for their unconditional counterparts . We present an approach based on the conditional VAE framework but , to mitigate the associated slow training times , we design the architecture so that we can incorporate pretrained unconditional VAEs . We show that re-using publicly available pretrained models in this way can lead to training times competitive with GANs . While requiring an existing pretrained model is a limitation , we note that : ( I ) The unconditional VAE need not have been ( pre- ) trained on the same dataset as the conditional model ; we show unconditional models trained on ImageNet are suitable for later use with various photographic image datasets . ( II ) A single unconditional VAE can be used for later training of conditional VAEs on any desired conditional generation tasks ( e.g . the same image model may be later used for image completion or image colourisation ) . ( III ) There is an increasing trend in the machine learning community towards sharing large , expensively trained models ( Wolf et al. , 2020 ) , sometimes referred to as foundation models ( Bommasani et al. , 2021 ) . Most of the unconditional VAEs in our experiments use publicly-available pretrained weights released by Child ( 2020 ) . By presenting a use case for foundation models in image modelling , we hope to encourage even more sharing of pretrained weights in this domain . We demonstrate our approach on several conditional generation tasks in the image domain but focus in particular on stochastic image completion : the problem of inferring the posterior distribution over images given the observation of a subset of pixel values . The visual quality of image completions produced by our method ( see Fig . 1 ) is close to the state-of-the-art ( Zhao et al. , 2021 ) , and we show results indicating that our coverage of the “ true ” posterior over image completions is superior to that of any of our baselines . For some applications such as photo-editing the implicit distribution defined by GANs is good enough . We argue that our approach has substantial advantages when image completion is used as part of a larger pipeline . We demonstrate one such situation in Section 5 : Bayesian optimal experimental design ( BOED ) for guiding a sensor or hard attention mechanism ( Ma et al. , 2018 ; Harvey et al. , 2019 ; Rangrej & Clark , 2021 ) . In this case , missing modes of the posterior over images is likely to lead to bad decisions . We show that our objective corresponds to the mass-covering KL divergence and so covers the posterior well , outperforming GANs which typically learn a distribution with a fraction of the data distribution ’ s support ( Arora & Zhang , 2017 ; Arora et al. , 2017 ) . Contributions We develop a method to cheaply convert pretrained unconditional VAEs into conditional VAEs . We demonstrate that this allows conditional VAE training times competitive with GANs , while avoiding the mode-dropping behaviour associated with GANs . Finally , we showcase an application in Bayesian optimal experimental design that benefits from these capabilities . 2 VARIATIONAL AUTO-ENCODERS . We describe VAEs in terms of three components . ( I ) A decoder with parameters θ ∈ Θ maps from latent variables z to a distribution over data x , which we call pmodel ( x|z ; θ ) . ( II ) There is a prior over latent variables , pmodel ( z ; θ ) . This may have learnable parameters , which we consider to be part of θ . Together , the prior and decoder define a joint distribution , pmodel ( z , x ; θ ) . Finally , ( III ) an encoder with parameters φ ∈ Φ maps from data to an approximate posterior distribution over latent variables , q ( z|x ; φ ) ≈ pmodel ( z|x ; θ ) . Ideally , θ would be learned to maximise the log likelihood log pmodel ( x ; θ ) = log ∫ pmodel ( z , x ; θ ) dz , averaged over training examples . Since this is intractable , θ and φ are instead trained jointly to maximise an average of the evidence lower-bound ( ELBO ) over each training example x ∼ pdata ( · ) : Epdata ( x ) [ ELBO ( θ , φ , x ) ] = Epdata ( x ) Eq ( z|x ; φ ) [ log pmodel ( z ; θ ) pmodel ( x|z ; θ ) q ( z|x ; φ ) ] ( 1 ) = −H [ pdata ( x ) ] − KL ( pdata ( x ) q ( z|x ; φ ) ∥∥ pmodel ( z , x ; θ ) ) . ( 2 ) The data distribution ’ s entropy , H [ pdata ( x ) ] , is typically a finite constant , and this is guaranteed in our experiments where x is an image with discrete pixel values . Maximising the above objective will therefore drive pmodel ( z , x ; θ ) towards pdata ( x ) q ( z|x ; φ ) , and so the marginal pmodel ( x ; θ ) towards pdata ( x ) . The KL divergence shown leads to mass-covering behaviour from pmodel ( z , x ; θ ) ( Bishop , 2006 ) so pmodel ( x ; θ ) should assign probability broadly over the data distribution pdata ( x ) . For notational simplicity in the rest of the paper , parameters θ and φ are not written when clear from the context . Hierarchical VAEs ( Gregor et al. , 2015 ; Kingma et al. , 2016 ; Sønderby et al. , 2016 ; Klushyn et al. , 2019 ) partition the latent variables z in a way which has been found to improve the fidelity of the learned pmodel ( x ) , especially for the image domain ( Vahdat & Kautz , 2020 ; Child , 2020 ) . In particular , they define z to consist of L disjoint groups , z1 , . . . , zL . The prior for each zl can depend on the previous groups through the factorisation pmodel ( z ) = L∏ l=1 pmodel ( zl|z < l ) . ( 3 ) where z < l is the null set for l = 1 and { z1 , . . . , zl−1 } otherwise . Fig . 2a shows the hierarchical VAE architecture we base this work on , in which the dependency of the prior for each zl on z < l is maintained via the decoder ’ s ( shown in black ) hidden state hl . The distribution produced by the encoder ( shown in orange ) for each zl also depends on the previous hidden state hl−1 and so factorises as q ( z|x ) = ∏L l=1 q ( zl|z < l , x ) . We will parameterise pmodel ( zl|z < l ) and q ( zl|z < l , x ) as diagonal Gaussian distributions , as is common for hierarchical VAEs ( Sønderby et al. , 2016 ; Vahdat & Kautz , 2020 ; Child , 2020 ) . 3 AMORTIZED INFERENCE IN A PRETRAINED ARTIFACT . To convert an unconditional VAE architecture to a conditional architecture , we introduce a partial encoder with parameters φ̂ ∈ Φ̂ . This is fed a conditioning input y . For example , in the case of image completion , y could be an image with some pixels masked out . The partial encoder then defines an approximate posterior over the latent variables , q̂ ( z|y ; φ̂ ) . The conditional generation task is to approximate pdata ( x|y ) . Using the partial encoder , we define pcond ( x|y ; θ , φ̂ ) : = ∫ pmodel ( x|z ; θ ) q̂ ( z|y ; φ̂ ) dz ( 4 ) with learnable parameters θ and φ̂ . We can sample from pcond ( x|y ) by sampling z ∼ q̂ ( ·|y ) and then x ∼ pmodel ( ·|z ) as shown in Fig . 2c . This defines a conditional VAE architecture which , unique amongst related work with high-dimensional x and y ( Sohn et al. , 2015 ; Zheng et al. , 2019 ; Ivanov et al. , 2018 ; Wan et al. , 2021 ) , has a decoder pmodel ( x|z ; θ ) with no dependence on y . This decoder can therefore use an architecture identical to that of an unconditional VAE and also , as we will show later , re-use unconditional VAE weights . Before describing our method in further detail , we briefly introduce some notation . Let the distribution of paired data be pdata ( x , y ) . Then recall that training an unconditional VAE matches two joint distributions : the distribution of samples from the generator , pmodel ( z , x ) ; and the distribution resulting from sampling data x and encoding it , pdata ( x ) q ( z|x ) . For notational convenience , we define the following extensions of these joint distributions to include y : pmodel ( z , x , y ; θ ) = pmodel ( z ; θ ) pmodel ( x|z ; θ ) pdata ( y|x ) , ( 5 ) r ( z , x , y ; φ ) = pdata ( x , y ) q ( z|x ; φ ) , ( 6 ) where pdata ( y|x ) is a ( potentially intractable ) conditional distribution under pdata ( x , y ) . Note that pmodel ( z , x , y ; θ ) and r ( z , x , y ; φ ) are exactly the two distributions matched by the unconditional VAE objective in Eq . ( 2 ) with an additional factor of pdata ( y|x ) . Therefore , if the unconditional VAE represented by θ and φ is well trained , pmodel and r will be close . From now on , we will use pmodel and r to refer to any marginals and conditionals of the above joint distributions , with the specific marginal or conditional clear from context .

Summary: The paper proposes a method of leveraging state-of-the-art VAE decoders (foundational models) to create new conditional VAEs by training a new encoder for prior of the latent space.  The authors demonstrate quite convincingly that the new conditional VAE can generate images with more variation than a state-of-the-art GAN while achieving a comparable quality of image.  The advantage of the approach is that it substantially reduces training times (a drawback of traditional VAEs).  They then propose a task in the domain of X-ray imaging where the improved variation leads to improved performance in targetted X-ray measurements.