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 author(s) face the design of Bayesian autoencoders (BAE) and, in particular, the model selection problem. A likelihood model is proposed, then to optimize it the prior of the weights is also needed, a normal distribution is proposed where mean and variances are to be learned. Marginal likelihood estimation is brought here to learn these hyperparameters, but being intractable is transformed into an approximated expression using distributional sliced-Wassersteing distance (DSWD). This expression is estimated with Monte Carlo integration and optimized via gradient descent. In the end, the prior is found by matching the true data distribution and the parametrized marginal distribution of the joint probability of the data and the weights.  The contribution of the manuscript is to develop the mathematical framework to optimize the hyperparameters of a prior for the weights in an autoencoder using DSWD. The solution is designed for a normal prior with to-be-learned mean and variance and continuous Bernoulli likelihood for the probability of the data conditioned to the weights.