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: This work introduces a new method to embed used-designed inductive-biases into VAE architectures. the basic idea is straightforward: a fixed user-defined deterministic transformation is used to map the latent z into a "structured latent" y. The transformation is used to induces topological and/or sparsity constraints. Technically, this is identical to using the transformation as the first layer of the decoder architecture. Besides of the general idea, the authors introduce a variety of simple but clever transformations used to enforce latent constraints. The approach is shown to offer an interpretable latent space and attractive performance in a large class of density estimation tasks.