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: In this paper, the authors focus on training conditional variational autoencoders. They propose an architecture and training objective which leverages pretraining an initial unconditional VAE. The approach effectively infers the latent variables of the original unconditional VAE given the new conditioning input. They demonstrate the method on a number of datasets, specifically focusing on image completion. They show competitive performance with adversarial approaches. They also demonstrate an application in the realm of medical imaging.