Structured Uncertainty in the Observation Space of Variational Autoencoders

1 INTRODUCTION Generative modelling is one of the cornerstones of modern machine learning . One of the most used and widespread classes of generative models is the Variational Autoencoder ( VAE ) ( Kingma & Welling , 2014 ; 2019 ) . VAEs explicitly model the distribution of observations by assuming a latent variable model with low-dimensional latent space and using a simple parametric distribution in observation space . Using a neural network , VAEs decode the latent space into arbitrarily complex observational distributions . Despite many improvements on the VAE model , one oftenoverlooked aspect is the choice of observational distribution . As an explicit likelihood model , the VAE assumes a distribution in observation space – using a delta distribution would not allow gradient based optimisation . Most current implementations , however , employ only simple models , such as pixelwise independent normal distributions , which eases optimisation but limits expressivity . Else , the likelihood term is often replaced by a reconstruction loss – which , in the case of an L2 loss , implicitly assumes an independent normal distribution . Following this implicit assumption , samples are then generated by only predicting the mean , rather than sampling in observation space . An application where this disconnect becomes apparent is image synthesis . The common choices for observational distributions are pixel-wise independent categorical or normal distributions . For pixel-wise independent distributions , regardless of other model choices , sampling from the joint distribution over pixels will result in spatially-incoherent samples due to independent pixel noise ( cf . Figure 1 ) . To address this problem , researchers use the predicted distributions to calculate the log-likelihood in the objective but then discard them in favour of the mean when generating samples or reconstructing inputs . However , this solution is akin to ignoring the issue rather than attempting to solve it . In this work , we explore what happens when we strictly follow VAE theory and sample from the predicted observational distributions . We illustrate the problem of spatial incoherence that arises from using pixel-wise independent distributions . We propose using a spatially dependent joint distribution over the observation space and compare it to the previous scenario . We further compare the samples to the mean of the predicted observational distribution , which is typically used when synthesising images . We note that , in this work , we are not focusing on absolute image quality . Instead , we aim to point to an issue often overlooked in VAE theory and application , which affects most state-of-the-art methods . Thus we analyse the relative difference between using and not using a joint pixel dependent observational distribution for a basic VAE . Yet , our findings are of broad relevance and our proposed model can be used in more advanced VAE variants . 2 RELATED WORK . Modern generative models can be divided into two classes , implicit likelihood models , such as Generative Adversarial Networks ( GANs ) ( Goodfellow et al. , 2014 ) and diffusion models ( Song & Durkan , 2021 ) , and explicit likelihood models , such as VAEs ( Kingma & Welling , 2014 ; Rezende et al. , 2014 ; Kingma & Welling , 2019 ) , flow models ( Dinh et al. , 2017 ; Kingma & Dhariwal , 2018 ; Dinh et al. , 2015 ) and auto-regressive models ( Van Den Oord et al. , 2016a ; b ; Salimans et al. , 2017 ) . Despite implicit likelihood models having achieved impressive results in terms of sample quality without the need for explicitly modelling the observation space ( Karras et al. , 2020 ; Brock et al. , 2019 ) , interest in explicit likelihood models have prevailed due to their appealing properties , such as ease of likelihood estimation . One of the most popular and successful explicit likelihood models is the VAE ( Kingma & Welling , 2014 ; Rezende et al. , 2014 ; Kingma & Welling , 2019 ) . Since its introduction , there have been numerous extensions . For example , Van Den Oord et al . ( 2017 ) ; Razavi et al . ( 2019 ) quantize the latent space to achieve better image quality ; Higgins et al . ( 2017 ) ; Chen et al . ( 2017 ) modify the latent posterior to obtain disentangled and interpretable representations ; and Vahdat & Kautz ( 2020 ) ; Sønderby et al . ( 2016 ) use hierarchical architectures to improve sample quality . Like most other explicit likelihood models , the VAE requires the choice of a parametric observational distribution . This choice is often pixel-wise independent . As a result , practitioners use the distribution to calculate the likelihood but use its expected value when sampling , as the samples themselves , are noisy and of limited use in most applications . However , according to the theory and as previously pointed out by Stirn & Knowles ( 2020 ) and Detlefsen et al . ( 2019 ) , a latent space sample should entail a distribution over observations and not a single point . Despite attempts to enforce spatial dependencies in the decoder architecture ( Miladinovic et al. , 2021 ) , without a pixeldependent joint likelihood , the observation samples will remain noisy . Notable exceptions are autoregressive models and auto-regressive VAE decoders ( Van Den Oord et al. , 2017 ; Razavi et al. , 2019 ; Gulrajani et al. , 2017 ; Nash et al. , 2021 ) . Unlike other explicit likelihood models , auto-regressive models jointly model the observational distribution by sequentially decoding pixels while conditioning on previously decoded values . While this sampling procedure results in spatially coherent samples , it is computationally expensive and uncertainty estimation is not trivial . In the context of non-auto-regressive VAE decoders , work focusing on modelling a joint observational distribution that accounts for pixel dependencies is limited . Monteiro et al . ( 2020 ) use a low-rank multivariate normal distribution to produce spatially consistent samples in a segmentation setting . However , they focus on discriminative models only . For generative models , a notable exception is a work by Dorta et al . ( 2018 ) which , similarly to our proposed method , employs a non-diagonal multivariate normal distribution over observation space . The key difference is the choice of parameterisation used for the covariance matrix . Dorta et al . ( 2018 ) predict the Choleskydecomposed precision matrix , which grows quadratically with the size of the image . To address this computational constraint , the authors use a sparse decomposition . This decomposition considers only a local neighbourhood of pixels , which limits its ability to capture long-range spatial dependencies . In contrast , our approach uses a global parameterisation . 3 METHODS . 3.1 VARIATIONAL AUTOENCODERS . We briefly revisit the theory of standard VAEs as proposed by Kingma & Welling ( 2014 ) ; Rezende et al . ( 2014 ) . We assume a prior distribution over the latent variables , p ( z ) , a probabilistic encoder of the posterior , pθ ( z|x ) , and a probabilistic decoder of the likelihood , pθ ( x|z ) . In practice , the probabilistic encoder describes an intractable posterior . Using variational inference , we approximate this posterior with a distribution , qφ ( z|x ) given parameters φ . Here , VAEs provide the algorithm to jointly learn the parameters θ and φ ( Kingma & Welling , 2014 ) . Considering some dataset X = { x ( i ) } Ni=1 , the VAE objective is given by maximising the evidence lower bound with respect to the parameters θ and φ ( Bank et al. , 2020 ; Kingma & Welling , 2014 ) : L ( θ , φ ; x ( i ) ) = −DKL [ qφ ( z|x ( i ) ) ||pθ ( z ) ] + Eqφ ( z|x ( i ) ) [ log pθ ( x ( i ) |z ) ] ( 1 ) Since the derivative of the lower bound w.r.t φ is problematic due to the stochastic expectation operator , the re-parameterisation trick is used to yield the Stochastic Gradient Variational Bayes ( SGVB ) estimator ( Kingma & Welling , 2014 ) . Note , from the second term in equation 1 , that the optimisation objective requires the choice of an observational distribution to calculate the likelihood that the data comes from the predicted distribution : pθ ( x|z ) . Hence , it follows that a latent sample entails a distribution over observations . The predicted distribution should be as close as possible to the observed distribution . Its samples should look like real observations . However , in the case of highly structured data such as images , this will not be the case when using the commonly-employed pixel-independent joint distribution . 3.2 STRUCTURED OBSERVATION SPACE VARIATIONAL AUTOENCODERS . With the commonly assumed model for the observation space , where the predicted joint distribution is pixel-wise independent , samples can only add noise to the predicted mean , as shown in the examples in section 4.1 . By incorporating spatial dependencies in the predicted distribution , we aim to overcome this limitation and generate more realistic samples under the observed distribution . Our solution is to replace the joint independent distribution predicted by the decoder with one that explicitly models dependencies between outputs . Specifically , we modify the final layer of the decoder to predict a low-rank parameterisation of a fully populated covariance matrix for use in a multivariate normal distribution , pθ ( x|z ) ∼ N ( µ , Σ ) . This small modification can be applied to most existing VAE architectures . Inspired by a recent discriminative model ( Monteiro et al. , 2020 ) , we use an efficient parameterisation , Σ = PPT+D , to model covariance globally , albeit at a low-rank . This yields a compact model with a covariance factor , P ∈ R ( S×C ) ×R and a covariance diagonal , D = diag ( d ) ∈ R ( S×C ) 2 , with diagonal elements , d. Here , S = H ×W is the number of pixels , C is the number of channels , and R is the rank of the parameterisation . Since we have only modified the distribution of the likelihood , the SGVB estimator ( equation 1 ) is applicable without modification . However , we found optimising the SGVB estimator with a non-diagonal covariance to be unstable regarding the variance . The instability comes from the fact that two routes can optimise the likelihood : to find the correct mean and appropriate variance around it or to keep increasing the variance ( uncertainty about the mean ) . The second direction is obviously undesirable and results in implausible samples ( with overly bright colours and high contrast ; cf . Appendix A ) . Monteiro et al . ( 2020 ) observed a similar phenomenon and pre-trained on the mean to avoid it . This problem is not unique to our implementation ; the stability of variance networks has been discussed before ( Stirn & Knowles , 2020 ; Detlefsen et al. , 2019 ) . In our case , we found pre-training the mean to be beneficial but insufficient . To further mitigate the problem , our solution includes weight initialisation , fixing the covariance diagonal to a small positive scalar : D = I 1 , and constraining the entropy of the predicted distribution . Constraining the entropy constrains the variance indirectly , thus giving preference to low-variance solutions . We compute 1We found 10−5 to yield good results . the entropy of the normal distribution in closed form and add it to the objective function . We employ soft-constraints for the entropy and KL divergence using the modified differential method of multipliers ( Platt & Barr , 1988 ) , for its agreeable convergence and stability properties , which results in the following Lagrangian formulation : L̃ ( θ , φ ; x ( i ) ) = 1 M M∑ m=1 log pθ ( x ( i ) |z ( i , m ) ) ( 2 ) − β [ DKL [ ( qφ ( z|x ( i ) ) ||pθ ( z ) ) ] − ξKL ] − λH [ H ( x ( i ) |z ) − ξH ] where β and ξKL are the Lagrangian multiplier and the slack variable , respectively , for the β-VAE constraint ; H ( x ( i ) |z ) is the entropy of the predicted distribution and λH and ξH are the Lagrangian multiplier and the slack variable , respectively , for the new constraint .

Summary: The authors aim at improving the canonical VAE model by replacing the standard iid Gaussian likelihood with a multivariate Gaussian with (low-rank + diagonal) covariance.  In applications to CelebA and a brain MRI dataset from UK Biobank, the authors compare the proposed structured-observation-space VAE with a canonical VAE, showing that the samples from their model have lower Fréchet Inception Distance (FID) scores than both means and samples from the canonical VAE. The authors then evaluate the expressiveness of the representations learned in the observation space by visually evaluating interpolations in the observation space, and images obtained by rescaling the principal components of the observational covariance matrix. Finally, the authors show how their model can be used for interactive editing—i.e., editing a small number of pixels and using the conditional distribution to infer coherent edits in the remainder of the image.