Model Selection for Bayesian Autoencoders

1 Introduction . The problem of learning useful representations of data that facilitate the solution of downstream tasks such as clustering , generative modeling and classification , is at the crux of the success of many machine learning applications [ see , e.g. , 5 , and references therein ] . From a plethora of potential solutions to this problem , unsupervised approaches based on autoencoders [ 10 ] are particularly appealing as , by definition , they do not require label information and have proved effective in tasks such as dimensionality reduction and information retrieval [ 20 ] . Autoencoders are neural network models composed of two parts , usually referred to as the encoder and the decoder . The encoder maps input data to a set of lower-dimensional latent variables . The decoder maps the latent variables back to the observations . The bottleneck introduced by the low-dimensional latent space is what characterizes the compression and representation learning capabilities of autoencoders . It is not surprising that these models have connections with principal component analysis [ 4 ] , factor analysis and density networks [ 33 ] , and latent variable models [ 28 ] . In applications where quantification of uncertainty is a primary requirement or where data is scarce , it is important to carry out a Bayesian treatment of these models by specifying a prior distribution over their parameters , i.e. , the weights of the encoder/decoder . However , estimating the posterior 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . distribution over the parameters of these models , which we refer to as Bayesian autoencoders ( BAEs ) , is generally intractable and requires approximations . Furthermore , the need to specify priors for a large number of parameters , coupled with the fact that autoencoders are not generative models , has motivated the development of Variational Autoencoders ( VAEs ) as an alternative that can overcome these limitations [ 25 ] . Indeed , VAEs have found tremendous success and have become one of the preferred methods in modern machine-learning applications [ see , e.g. , 26 , and references therein ] . To recap , three potential limitations of BAEs hinder their widespread applicability in order to achieve a similar or superior adoption to their variational counterpart : ( i ) lack of generative modeling capabilities ; ( ii ) intractability of inference and ( iii ) difficulty of setting sensible priors over their parameters . In this work we revisit BAEs and deal with these limitations in a principled way . In particular , we address the first limitation in ( i ) by employing density estimation in the latent space . Furthermore , we deal with the second limitation in ( ii ) by exploiting recent advances in Markov chain Monte Carlo ( MCMC ) and , in particular , stochastic gradient Hamiltonian Monte Carlo ( SGHMC ) [ 9 ] . Finally , we believe that the third limitation ( iii ) , which we refer to as the difficulty of carrying out model selection , requires a more detailed treatment because choosing sensible priors for Bayesian neural networks is an extremely difficult problem , and this is the main focus of this work . Contributions . Specifically , in this paper we provide a novel , practical , and elegant way of performing model selection for BAEs , which allows us to revisit these models for applications where VAEs are currently the primary choice . We start by considering the common practice of estimating prior ( hyper- ) parameters via type-II maximum likelihood , which is equivalent to minimizing the Kullback-Leibler divergence ( KL ) between the distribution induced by the BAE and the data generating distribution . Because of the intractability of this objective and the difficulty to estimate it through samples , we resort to an alternative formulation where we replace the KL with the distributional sliced-Wasserstein distance ( DSWD ) between these two distributions . The advantages of this formulation are that we can estimate the DSWD based on samples and , thanks to the slicing , we can handle large dimensional problems . Once BAE hyper-parameters are optimized , we estimate the posterior distribution over the BAE parameters via SGHMC [ 9 ] , which is a powerful sampler that operates on mini-batches and has proven effective for Bayesian deep/convolutional networks [ 45 , 52 , 23 ] . Furthermore , we turn our BAE into a generative model by fitting a flexible mixture model in the latent space , namely the Dirichlet Process Mixture Model ( DPMM ) . We evaluate our approach qualitatively and quantitatively using a vast experimental campaign on a number of unsupervised learning tasks , with particular emphasis on the challenging task of generative modeling when the number of observations is small . 2 Preliminaries on Bayesian Autoencoders . An autoencoder ( AE ) is a neural network parameterized by a set of parameters w , which transforms an unlabelled dataset , x def= { xn } Nn=1 , into a set of reconstructions f def = { fn } Nn=1 , with xn , fn ∈ RD . An AE is composed of two parts : ( 1 ) an encoder fenc which maps an input sample xn to a latent code zn ∈ RK , K D ; and ( 2 ) a decoder fdec which maps the latent code to a reconstructed datapoint fn . In short , f = f ( x ; w ) = ( fdec ◦ fenc ) ( x ) , where we denote w : = { wenc , wdec } the union of parameters of the encoder and decoder . The Bayesian treatment of AEs dictates that a prior distribution p ( w ) is placed over all parameters of fenc and fdec and the posterior is inferred using Bayes ’ rule . For the sake of presentation , we can consider a BAE as a classic Bayesian neural network ( BNN ) for a supervised learning task [ 34 , 37 ] , where we have labels y def= { yn } Nn=1 associated with each input point xn . From this perspective , we can write the posterior on w as follows : p ( w |y , x ) = p ( y |w , x ) p ( w ) p ( y |x ) , ( 1 ) where p ( y |w , x ) is the likelihood defined by the network architecture and the denominator— p ( y |x ) —constitutes the marginal likelihood . In order to keep the notation uncluttered , we can drop the explicit dependency on x as input , which leads to : p ( w |y ) = p ( y |w ) p ( w ) p ( y ) , ( 2 ) We observe that for an AE the labels y are the same as the datapoints x , meaning that the likelihood is computed in x . Likelihood model . Before giving an in-depth treatment on priors for BAEs in the next section , we briefly discuss the likelihood , which can be chosen according to the type of data . Firstly , we assume factorization of the likelihood on the datapoints , i.e . p ( y |w ) = ∏N n=1 p ( yn |w ) . Secondly , given that our experiments mainly focuses on image datasets where pixel values are normalized in the [ 0 , 1 ] range , we will use the continuous Bernoulli distribution [ 31 ] : p ( yn |w ) = D∏ i=1 K ( λi ) λ yn , i i ( 1− λi ) 1−yn , i : = p ( yn | fn ) , ( 3 ) where K ( λi ) is a properly defined normalization constant [ 31 ] and λi = fi ( xn ; w ) = fn , i ∈ [ 0 , 1 ] is the i-th output from the BAE given the input xn . We note that , as fn depends deterministically on w , we will use the above expression to refer to both p ( yn |w ) and p ( yn | fn ) , where the latter term will be of crucial importance when we define the functional prior induced over the reconstruction f . Finally , we remark that in the Bayesian scheme both the prior and likelihood are modeling choices . In fact , in our experiments we explore an additional likelihood model , namely the truncated Gaussian , and we show how the problem of selecting good priors is orthogonal to the choice of the likelihood . Inference . Although the posterior over the BAE parameters is analytically intractable , it can be approximated by variational methods or using MCMC sampling . Within the large family of approximate Bayesian inference schemes , SGHMC [ 9 ] allows us to sample from the true posterior by efficiently simulating a Hamiltonian system [ 38 ] . Unlike more traditional methods , SGHMC can scale up to large datasets by relying on noisy but unbiased estimates of the potential energy function U ( w ) = − ∑N n=1 log p ( yn |w ) − log p ( w ) . These can be computed by considering a mini-batch of size M of the data and approximating ∑N n=1 log p ( yn |w ) ≈ N M ∑ j∈IM log p ( yj |w ) , where IM is a set of M random indices . More details on SGHMC can be found in the Appendix . Pathologies of standard priors . The choice of the prior is important for the Bayesian treatment of any model as it characterizes the hypothesis space [ 34 , 36 ] . Specifically for BAEs , one should note that placing a prior on the parameters of the encoder and decoder has an implicit effect on the prior over the network output ( i.e . the reconstruction ) . In addition , the highly nonlinear nature of these models implies that interpreting the effect of the architecture is theoretically intractable and practically challenging . Several works argue that a vague prior such as N ( 0 , 1 ) is good enough for some tasks and models , like classification with convolutional neural networks ( CNNs ) [ 49 ] . However , for BAEs this is not enough , as illustrated in Fig . 1 . The realizations obtained by sampling weights/biases from aN ( 0 , 1 ) prior indicate that this choice provides poor inductive bias . Meanwhile , by encoding better beliefs via an optimized prior , which is the focus of the next section , the samples can capture main characteristics intrinsic to the data , even when the model is fed with out-ofdistribution inputs . 3 Model Selection for Bayesian Autoencoders via Prior Optimization . One of the main advantages of the Bayesian paradigm is that we can incorporate prior knowledge into the model in a principled way . Let us assume a prior distribution pψ ( w ) on the parameters of the AE network , where now we are explicit on the set of ( hyper- ) parameters that determine the prior , i.e. , ψ. Specifying the prior is easy , e.g. , a Gaussian . Determining the effective functional prior , i.e. , the prior over the network output f is not trivial due to the complex nonlinear forms of fenc and fdec , which induce a non-trivial effect on the output ( functional ) prior : pψ ( f ) = ∫ f ( x ; w ) pψ ( w ) dw , ( 4 ) where , as before , f = f ( x ; w ) is the functional output of the BAE . Although pψ ( f ) can not be evaluated analytically , it is possible to draw samples from it . Prior parameterization . The only two requirements needed to design a parameterization for the prior are : to be able to ( 1 ) draw samples from it and ( 2 ) to compute its log-density at any point . The latter is required by many inference algorithms such as SGHMC . We consider a fully-factorized Gaussian prior over weights and biases at layer l : p ( wl ) = N ( wl ; µlw , σ2lw ) , p ( bl ) = N ( bl ; µlb , σ 2 lb ) , ( 5 ) Notice that , as we shall see in § 3.2 and § 3.3 , in order to estimate our prior hyper-parameters , we will require gradient back-propagation through the stochastic variables wl and bl . Thus , we treat these parameters in a deterministic manner by means of the reparameterization trick [ 42 , 25 ] .

Summary: The goal of this paper is to train a Bayesian autoencoder (BAE), where the AE parameters are specified as random variables with trainable hyperparameters. Three contributions are claimed: (1) the distributional sliced-Wasserstein (DSW) distance [25] is leveraged to enable sample-based training of BAE hyperparameters (which model the distribution of AE parameters); (2) to instantiate the AE parameters, the SGHMC [9] is used; (3) to sample from a BAE, the authors propose to learn an ex-post density to model the distribution of the projected latent codes, mimicking [12,8,16]. Experiments on MNIST, FREY-YALE, and CELEBA are conducted.