On Incorporating Inductive Biases into VAEs

1 INTRODUCTION . VAEs provide a rich class of deep generative models ( DGMs ) with many variants ( Kingma & Welling , 2014 ; Rezende & Mohamed , 2015 ; Burda et al. , 2016 ; Gulrajani et al. , 2016 ; Vahdat & Kautz , 2020 ) . Based on an encoder-decoder structure , VAEs encode datapoints into latent embeddings before decoding them back to data space . By parameterizing the encoder and decoder using expressive neural networks , VAEs provide a powerful basis for learning both generative models and representations . The standard VAE framework assumes an isotropic Gaussian prior . However , this can cause issues , such as when one desires the learned representations to exhibit some properties of interest , for example sparsity ( Tonolini et al. , 2020 ) or clustering ( Dilokthanakul et al. , 2016 ) , or when the data distribution has very different topological properties from a Gaussian , for example multi-modality ( Shi et al. , 2020 ) or group structure ( Falorsi et al. , 2018 ) . Therefore , a variety of recent works have looked to use non-Gaussian priors ( van den Oord et al. , 2017 ; Tomczak & Welling , 2018 ; Casale et al. , 2018 ; Razavi et al. , 2019 ; Bauer & Mnih , 2019 ) , often with the motivation of adding inductive biases into the model ( Davidson et al. , 2018b ; Mathieu et al. , 2019b ; Nagano et al. , 2019 ; Skopek et al. , 2019 ) . In this work , we argue that this approach of using non-Gaussian priors can be a problematic , and even ineffective , mechanism for adding inductive biases into VAEs . Firstly , non-Gaussian priors will often necessitate complex encoder models to maintain consistency with the prior ’ s shape and dependency structure ( Webb et al. , 2018 ) , which typically no longer permit simple parameterization . Secondly , the latent encodings are still not guaranteed to follow the desired structure because the ‘ prior ’ only appears in the training objective as a regularizer on the encoder . Indeed , Mathieu et al . ( 2019b ) find that changing the prior is typically insufficient in practice to learn the desired representations at a population level , with mismatches occurring between the data distribution and learned model . To provide an alternative , more effective , approach that does not suffer from these pathologies , we introduce Intermediary Latent Space VAEs ( InteL-VAEs ) , an extension to the standard VAE framework that allows a wide range of powerful inductive biases to be incorporated while maintaining an isotropic Gaussian prior . This is achieved by introducing an intermediary set of latent variables that deal with the stochasticity of the encoding process before incorporating the desired inductive biases via a parametric function that maps these intermediary latents to the latent representation itself , with the decoder taking this final representation as input . See Fig . 1 for an example . 1Department of Statistics , University of Oxford , 2University of Edinburgh * Correspondence to : Ning Miao < ning.miao @ stats.ox.ac.uk > , Tom Rainforth < rainforth @ stats.ox.ac.uk > The InteL-VAE framework provides a variety of advantages over directly replacing the prior . Firstly , it directly enforces our inductive biases on the representations , rather than relying on the regularizing effect of the prior to encourage this implicitly . Secondly , it provides a natural congruence between the generative and representational models via sharing of the mapping function , side-stepping the issues that non-Gaussian priors can cause for the inference model . Finally , it allows for more general and more flexible inductive biases to be incorporated , by removing the need to express them with an explicit density function and allowing for parts of the mapping to be learned during training . To further introduce a number of novel specific realizations of the InteL-VAE framework , showing how they can be used to incorporate various inductive biases , enforcing latent representations that are , for example , multiply connected , multi-modal , sparse , or hierarchical . Experimental results show their superiority compared with baseline methods in both generation and feature quality , most notably providing state-of-the-art performance for learning sparse representations in the VAE framework . To summarize , we a ) highlight the need for inductive biases in VAEs and explain why directly changing the prior is a suboptimal means for incorporating them ; b ) propose InteL-VAEs as a simple but effective general framework to introduce inductive biases ; and c ) introduce specific InteL-VAE variants which can learn improved generative models and representations over existing baselines on a number of tasks . Accompanying code is provided at https : //github.com/NingMiao/InteL-VAE . 2 THE NEED FOR INDUCTIVE BIASES IN VAES . Variational auto-encoders ( VAEs ) are deep stochastic auto-encoders that can be used for learning both deep generative models and low-dimensional representations of complex data . Their key components are an encoder , qϕ ( z|x ) , which probabilistically maps from data x ∈ X to latents z ∈ Z ; a decoder , pθ ( x|z ) , which probabilistically maps from latents to data ; and a prior , p ( z ) , that completes the generative model , p ( z ) pθ ( x|z ) , and regularizes the encoder during training . The encoder and decoder are parameterized by deep neural networks and are simultaneously trained using a dataset { x1 , x2 , ... , xN } and a variational lower bound on the log-likelihood , most commonly , L ( x , θ , ϕ ) : = Ez∼qϕ ( z|x ) [ log pθ ( x|z ) ] −DKL ( qϕ ( z|x ) ∥ p ( z ) ) . ( 1 ) Namely , we optimize L ( θ , ϕ ) : = Ex∼pdata ( x ) [ L ( x , θ , ϕ ) ] , where pdata ( x ) represents the empirical data distribution . Here the prior is typically fixed to a standard Gaussian , i.e . p ( z ) = N ( z ; 0 , I ) . While it is well documented that this standard VAE setup with a ‘ Gaussian ’ latent space can be suboptimal ( Davidson et al. , 2018a ; Mathieu et al. , 2019b ; Tomczak & Welling , 2018 ; Bauer & Mnih , 2019 ; Tonolini et al. , 2020 ) , there is perhaps less of a unified high-level view on exactly when , why , and how one should change it to incorporate inductive biases . Note here that the prior does not play the same role as in a Bayesian model : because the latents themselves are somewhat arbitrary and the model is learned from data , it does not encapsulate our initial beliefs in the way one might expect . We argue that there are two core reasons why inductive biases can be important for VAEs : ( a ) standard VAEs can fail to encourage , and even prohibit , desired structure in the representations we learn ; and ( b ) standard VAEs do not allow one to impart prior information or desired topological characteristic into the generative model . Considering the former , one often has some a priori desired characteristics , or constraints , on the representations learned ( Bengio et al. , 2013 ) . For example , sparse features can be desirable because they can improve data efficiency ( Yip & Sussman , 1997 ) , and provide robustness to noise ( Wright et al. , 2009 ; Ahmad & Scheinkman , 2019 ) and attacks ( Gopalakrishnan et al. , 2018 ) . In other settings one might desire clustered ( Jiang et al. , 2017 ) , disentangled ( Ansari & Soh , 2019 ; Kim & Mnih , 2018 ; Higgins et al. , 2018 ) or hierarchical representations ( Song & Li , 2013 ; Sønderby et al. , 2016 ; Zhao et al. , 2017 ) . The KL-divergence term in Eq . ( 1 ) regularizes the encoding distribution towards the prior and , as a standard Gaussian distribution typically does not exhibit our desired characteristics , this regularization can significantly hinder our ability to learn representations with the desired properties . Not only can this be problematic at an individual sample level , it can cause even more pronounced issues at the population level : desired structural characteristics of our representations often relate to the pushforward distribution of the data in the latent space , qϕ ( z ) : = Epdata ( x ) [ qϕ ( z|x ) ] , which is both difficult to control and only implicitly regularized to the prior ( Hoffman & Johnson , 2016 ) . Inductive biases can also be essential to the generation quality of VAEs : because the generation process of standard VAEs is essentially pushing-forward the Gaussian prior on Z to data space X by a ‘ smooth ’ decoder , there is an underlying inductive bias that standard VAEs prefer sample distributions with similar topology structures to Gaussians . As a result , VAEs can perform poorly when the data manifold exhibits certain different topological properties ( Caterini et al. , 2020 ) . For example , they can struggle when data is clustered into unconnected com- ponents as shown in Fig . 2 , or when data is not simply-connected . This renders learning effective mappings using finite datasets and conventional architectures ( potentially prohibitively ) difficult . In particular , it can necessitate large Lipschitz constants in the decoder , causing knock-on issues like unstable training and brittle models ( Scaman & Virmaux , 2018 ) , as well as posterior collapse ( van den Oord et al. , 2017 ; Alemi et al. , 2018 ) . In short , the Gaussian prior of a standard VAE can induce fundamental topological differences to the true data distribution ( Falorsi et al. , 2018 ; Shi et al. , 2020 ) . 3 SHORTFALLS OF VAES WITH NON-GAUSSIAN PRIORS . Though directly replacing the Gaussian prior with a different prior sounds like a simple solution , effectively introducing inductive biases can , unfortunately , be more complicated . Firstly , the only influence of the prior during training is as a regularizer on the encoder through the DKL ( qϕ ( z|x ) ∥ p ( z ) ) term . This regularization is always competing with the need for effective reconstructions and only has an indirect influence on qϕ ( z ) . As such , simply replacing the prior can be an ineffective way of inducing desired structure at the population level ( Mathieu et al. , 2019b ) , particularly if p ( z ) is a complex distribution that it is difficult to fit ( see , e.g. , Fig . 3a ) . Mismatches between qϕ ( z ) and p ( z ) can also have further deleterious effects on the learned generative model : the former represents the distribution of the data in latent space during training , while the latter is what is used by the learned generative model , leading to unrepresentative generations if there is mismatch . Secondly , it can be extremely difficult to construct appropriate encoder mappings and distributions for non-Gaussian priors . While the typical choice of a mean-field Gaussian for the encoder distribution is simple , easy to train , and often effective for Gaussian priors , it is often inappropriate for other choices of prior . For example , in Fig . 3 , we consider replacement with a sparse prior . A VAE with a Gaussian encoder struggles to encode points in a manner that even remotely matches the prior . One might suggest replacing the encoder distribution as well , but this has its own issues , most notably that other distributions can be hard to effectively parameterize or train . In particular , the form of the required encoding noise might become heavily spatially variant ; in our sparse example , the noise must be elongated in a particular direction depending on where the mean embedding is . If the prior has constraints or topological properties distinct from the data , it can even be difficult to learn a mean encoder mapping that respects these , due to the continuous nature of neural networks .

Summary: The paper presents a novel method for introducing inductive biases into variational auto-encoders (VAEs). The inductive biases are specified through a deterministic transformation in the latent space, which is shared between the encoder and the decoder. Through experiments on toy data, the authors demonstrate how to design appropriate inductive biases by matching the topological properties of the transformed representations to those of the real data distributions. They further quantify the performance of the model's ability to learn informative representations through experiments on MNIST, Fashion-MNIST, and CelebA, and show that it outperforms competitors both in terms of the generation quality with and without sparsity constraints as well as terms of downstream classification based on the representations.