Structured Dropout Variational Inference for Bayesian Neural Networks

1 Introduction . Bayesian Neural Networks ( BNNs ) [ 49 , 63 ] offer a probabilistic interpretation for deep learning models by imposing a prior distribution on the weight parameters and aiming to infer a posterior distribution instead of only point estimates . By marginalizing over this posterior for prediction , BNNs perform a procedure of ensemble learning . These principles improve the model generalization , robustness and allow for uncertainty quantification . However , exactly computing the posterior of non-linear BNNs is infeasible and approximate inference has been devised . The core challenge is how to construct an expressive approximation for the true posterior while maintaining computational efficiency and scalability , especially for modern deep learning architectures . Variational inference is a popular deterministic approximation approach to deal with this challenge . The first practical methods were proposed in [ 22 , 8 , 39 ] , in which the approximate posteriors are assumed to be fully factorized distributions , also called mean-field variational inference . In general , the mean-field approximation family encourages several advantages in inference including computational tractability and effective optimization with the stochastic gradient-based methods . However , it ignores the strong statistical dependencies among random weights of neural nets , leading to the inability to capture the complicated structure of the true posterior and to estimate the true model uncertainty . ∗These two authors contributed equally . †Correspondence to Son Nguyen : < v.sonnv27 @ vinai.io > 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . To overcome this limitation , many extensive studies proposed to provide posterior approximations with richer expressiveness . For instance , [ 47 ] treats the weight matrix as a whole via a matrix variate Gaussian [ 24 ] and approximates the posterior based on this parametrization . Several later works have exploited this distribution to investigate different structured representations for the variational Gaussian posterior , such as Kronecker-factored [ 89 , 71 , 72 ] , k-tied distribution [ 77 ] , noncentered or rank-1 parameterization [ 21 , 15 ] . Another original idea to represent the true covariance matrix of Gaussian posterior is by employing the low-rank approximation [ 67 , 35 , 80 ] . For robust approximation with multimodality , [ 48 ] adopted hierarchical variational model framework [ 69 ] for inferring an implicit marginal distribution in high dimensional Bayesian setting . Despite significant improvements in both predictive accuracy and uncertainty calibration , some of these methods incur a large computational complexity and are difficult to integrate into deep convolutional networks . Motivations . In this paper , we approach the structured posterior approximation in Bayesian neural nets from a different perspective which has been inspired by the Bayesian interpretation of Dropout training [ 74 , 51 ] . More specifically , the methods proposed in [ 39 , 20 ] reinterpret Dropout regularization as approximate inference in Bayesian deep models and base on this connection to learn a variational Dropout posterior over the weight parameters . From the literature , inference approaches based on Bayesian Dropout have shown competitive performances in terms of predictive accuracy on various tasks , even compared to the structured Bayesian methods aforementioned , but with much cheaper computational complexity . Moreover , with the solid and intriguing theories on effective regularization [ 81 , 26 , 83 ] , generalization bound [ 52 , 59 ] , convergence rate and robust optimization [ 55 , 54 , 7 ] , Dropout principle offers several potentials to further improve approximate inference in Bayesian deep networks . However , since these Bayesian Dropout methods also employed simple structures of the mean-field family , their approximations often fail to obtain satisfactory uncertainty estimates [ 17 ] . In addition , Variational Dropout methods based on multiplicative Gaussian noise also suffer from theoretical pathologies , including improper prior leading to ill-posed true posterior , and singularity of the approximate posterior making the variational objective undefined [ 31 ] . Contributions . With the above insights , we propose a novel structured variational inference framework , which rationally acquires complementary benefits of the flexible Bayesian inference and Dropout inductive bias . Our method adopts an orthogonal approximation called Householder transformation to learn a structured representation for multiplicative Gaussian noise in Variational Dropout method [ 39 , 57 ] . As a consequence of the Bayesian interpretation , we go beyond the mean-field family and obtain a variational Dropout posterior with structured covariance . Furthermore , to make our framework more expressive , we deploy a hierarchical Dropout procedure , which is equivalent to inferring a joint posterior in a hierarchical Bayesian neural nets . We name the proposed method as Variational Structured Dropout ( VSD ) and summarize its advantages as follows : 1 . Our structured approximation is implemented on low dimensional space of variational noise with considerable computational efficiency . VSD can be employed for deep CNNs in a direct way while maintaining the backpropagation in parallel and optimizing efficiently with gradient-based methods . 2 . Especially , VSD has a standard Bayesian justification , in which our method can overcome the critiques from the non-Bayesian perspective of previous Variational Dropout methods . Our inference framework uses a proper prior , non-singular approximate posterior and derives a tractable variational lower bound without further simplified approximation . 3 . Compared with previous Bayesian Dropout methods which are relatively inflexible by some strict conditions , VSD is more efficient on both the criteria of expressive approximation and flexible hierarchical modeling . Therefore , VSD is promising to be a general-purpose approach for Bayesian inference and in particular for BNNs . 4 . To reinforce the complementary advantages unified in our proposal , we also investigate the inductive biases induced by the adaptive regularization of structured Dropout noise . We further provide an interpretation that VSD implicitly facilitates the networks to converge to a local minima with smaller spectral norms and stable rank . This properties suggests better generalization and we present empirical results to support this implication . 5 . Finally , we carry out extensive experiments with standard datasets and different network architectures to validate the effectiveness of our method on many criteria , including scalability , predictive accuracy , uncertainty calibration , and out-of-distribution detection , in comparison to popular variational inference methods . Notation . For a matrix A , ‖A‖F and A > denotes the Frobenius norm and the transpose matrix , Ai : and A : j denote the i-th row and the j-th column . For an interger i , ei is the i-th standard basis , 1i ∈ Ri is the vector of all ones . The diagonal matrix with diagonal entries as the elements of a vector x is denoted by diag ( x ) . The inner product between two matrices A and B is denoted by 〈A , B〉 . 2 Background . Variational inference for Bayesian neural networks : Given a dataset D consisting of input-output pairs ( X , Y ) : = { ( xn , yn ) } Nn=1 . In BNNs , we impose a prior distribution over random weights p ( W ) whose form is in a tractable parametric family and aim to infer an intractable true posterior p ( W|D ) . Variational inference ( VI ) [ 30 , 33 ] can do this by specifying a variational distribution qφ ( W ) with free parameter φ and then minimizing the Kullback-Leibler ( KL ) divergence DKL ( qφ ( W ) ‖p ( W|D ) ) . This optimization is equivalent to maximizing the Evidence Lower Bound ( ELBO ) with respect to variational parameters φ as follows : L ( φ ) = Eqφ ( W ) log p ( D|W ) − DKL ( qφ ( W ) ‖p ( W ) ) . ( 1 ) By leveraging the reparameterization trick [ 40 ] combined with the Monte Carlo integration , we can derive an unbiased differentiable estimation for the variational objective above . Then , this estimation can be effectively optimized using stochastic gradient methods with the variance reduction technique such as the local reparameterization trick [ 39 ] . Variational Bayesian inference with Dropout regularization : Given a deterministic neural net with the weight parameter Θ of size K ×Q , training this model with stochastic regularization techniques such as Dropout [ 29 , 74 ] can be interpreted as approximate inference in Bayesian probabilistic models . This is because injecting a stochastic noise into the input layer is equivalent to multiplying the rows of subsequent deterministic weight by the same random variable , namely with each datapoint ( xn , yn ) and a noise vector ξ , we have : yn = ( xn ξ ) Θ = xndiag ( ξ ) Θ . This induces a BNN with random weight matrix defined by W : = diag ( ξ ) Θ . Applying VI to this Bayesian model , with some specific choices for prior and approximate posterior , the variational lower bound ( 1 ) can resemble the form of Dropout objective in the original deterministic network . This principle is referred to as the KL condition [ 18 ] . Gal et al . [ 20 ] , Kingma et al . [ 39 ] used this principle to propose Bayesian Dropout inference methods such as MC Dropout ( MCD ) and Variational Dropout ( VD ) . Dropout inference is practical approximate framework especially in high dimensional setting . However , the scope of Bayesian inference in these methods is restricted in terms of flexibility of both prior and approximate posterior . Concretely , the Dropout posteriors qφ ( W ) in MCD and VD both have simple structures of mean-field approximation which often underestimate the variance of true posterior , possibly leading to a poor uncertainty representation [ 17 ] . Moreover , in theory , VD employed an improper log-uniform prior which can result in an ill-posed true posterior and generally push the parameters towards the penalized maximum likelihood solution [ 31 ] . In addition , VD also suffers from the singularity issue of approximate posterior that makes the KL divergence term undefined . Our work gains an efficient remedy to these pathologies . 3 Variational Structured Dropout . We focus on Bayesian Dropout methods using multiplicative Gaussian noise with correlated parameterization [ 39 ] . This procedure induces a random weight W = diag ( ξ ) Θ , where the Dropout noise ξ is a multivariate Gaussian with diagonal covariance qα ( ξ ) = N ( 1K , diag ( α ) ) , and α is the droprate vector . The corresponding Dropout posterior then is given by qφ ( W ) = Law ( diag ( ξ ) Θ ) . This distribution on each column exhibits a factorized structure with the form of q ( W : j ) = N ( Θ : j , diag ( α Θ2 : j ) ) , whilst allows a correlation on each row because each Wi : is shared by the same scalar noise ξi respectively . The parameters φ : = ( α , Θ ) are optimized via maximizing a variational lower bound as follows : L ( φ ) : = Eqφ ( W ) log p ( D|W ) −DKL ( qφ ( W ) ‖p ( W ) ) = Eqα ( ξ ) log p ( D|ξ , Θ ) − DKL ( qφ ( W ) ‖p ( W ) , where the later equation is derived from the change of variables formula . 3.1 The orthogonal approximation for variational structured noise . Intuitively , a richer representation for the noise distribution can enrich the expressiveness of Dropout posterior via the Bayesian interpretation . We implement this intuition with an assumption that the Dropout noise could be sampled from a Gaussian distribution with a full covariance matrix instead of a diagonal structure , namely , qΣ ( ξ ) = N ( 1K , Σ ) with Σ is a positive definite matrix of size K ×K . To make this covariance matrix learnable , we first represent Σ in the form of the spectral decomposition : Σ = PΛPT , where P is an orthogonal matrix with its eigenvectors in columns , Λ is a diagonal matrix where diagonal elements are the eigenvalues . By the basis-kernel representation theorem [ 6 , 76 ] , we parameterize the orthogonal matrix P as a product of Householder matrices in the following form : P = HT∗HT∗−1 ... H1 , where Ht = I− 2vtvTt /‖vt‖22 , vt is the Householder vector of size K , and T ∗ is the degree of P. This parameterization relaxes the orthogonal constraint of matrix P , and we can then directly optimize the covariance matrix Σ via gradient-based methods . Notably , this transformation can be interpreted as a sequence of invertible mappings . More explicitly , we extract a zero-mean Gaussian noise η ( 0 ) from the original noise ξ ( 0 ) ∼ N ( 1K , diag ( α ) ) in the form of ξ ( 0 ) = 1 + η ( 0 ) , and by successively transforming η ( 0 ) through a chain of T Householder reflections , we obtain the induced noise and the corresponding density at each step t as follows : ξ ( t ) : = 1 + HtHt−1 ... H1η ( 0 ) = 1 + Uη ( 0 ) , and qt ( ξ ) : = N ( 1K , Udiag ( α ) UT ) . By injecting the structured noise ξ ( t ) into the deterministic weight Θ , we obtain a random weight W ( t ) : = diag ( ξ ( t ) ) Θ and an induced Dropout posterior qt ( . ) = Law ( diag ( ξ ( t ) ) Θ ) with fully correlated representation , in which the marginal column distribution is given by : qt ( W : j ) = N ( Θ : j , diag ( Θ : j ) Udiag ( α ) UT diag ( Θ : j ) ) . Detailed discussion about the expressiveness of this correlated structure is presented in Appendix A.1 . With the above derivations , we use variational inference with approximate posterior qt ( W ) and optimize the variational lower bound as follows : L ( φ ) : = Eqt ( W ) log p ( D|W ) − DKL ( qt ( W ) ||p ( W ) ) = Eqα ( ξ ) log p ( D|Θ , ξ ( t ) ) − DKL ( qt ( W ) ||p ( W ) ) . ( 2 ) Overcoming the singularity issue of approximate posterior . In Variational Dropout method with correlated parameterization , there is a mismatch in support between the approximate posterior and the prior , thus making the KL term DKL ( q ( W ) ||p ( W ) ) undefined [ 31 ] . Specifically , the form W = diag ( ξ ) Θ is equivalent to multiplying each row Wi : by the same scalar noise ξi , namely Wi : = ξiΘi : with q ( ξi ) = N ( 1 , αi ) . This means that the approximate distribution always assigns all its mass on subspaces defined by the directions aligned with the rows of Θ . These subspaces have Lebesgue measure zero causing the singularities in approximate posterior , and the KL term will be undefined whenever the prior p ( W ) puts zero mass to these subspaces . However , in VSD the scalar noises are not treated separately due to the structured correlation , namely VSD would injects each ξi into the whole matrix Θ instead of some individual directions . Indeed , we have : W ( V D ) : = diag ( ξ ( 0 ) ) Θ = Θ + diag ( η ( 0 ) ) Θ = Θ + ∑K i=1 η ( 0 ) i ( diag ( ei ) Θ ) = Θ + ∑K i=1 η ( 0 ) i Θ ( i ) , and W ( V SD ) : = diag ( ξ ( t ) ) Θ = Θ + diag ( Uη ( 0 ) ) Θ = Θ + ∑K i=1 η ( 0 ) i ( diag ( Ui : ) Θ ) , where Θ ( i ) is the matrix Θ with only the i-th row retained . While VD causes singular components represented by { Θ ( i ) } Ki=1 , VSD maintains a trainable orthogonal matrix U which prevents the approximate posterior from having degenerate supports with measure zero , thereby avoiding the singularity issue . In the following section , we will present an appropriate choice of the prior distribution p ( W ) such that the KL term is well-defined , and then derive a tractable objective function satisfying the KL condition in Bayesian Dropout frameworks .

Summary: The paper proposes Variational Structured Dropout, an extension to variational dropout that uses a full covariance matrix to model the posterior at the pre-activations. This covariance matrix, after spectral decomposition, is parameterized as the product of T Householder matrices.  The new parameterization of the posterior does not suffer from the singularity issue of Variational Dropout. When using a diagonal posterior approximation, the singularity issue is that the approximate posterior has zero measure, which results in the KL-term being undefined. VSD avoids this by using a full covariance matrix thus ensuring that the approximate posterior has non-zero measure.  To train the model using the new parameterization, the authors optimize a new variational lower bound. This lower bound is a lower bound to the ELBO.  Regarding the prior, the paper advocates for using a hierarchical prior distribution where the prior is a zero mean Gaussian distribution over the weights, with the hyperprior determining the standard deviation.  The experiments first examine the impact of T, the number of matrices used to parameterize the approximate posterior. Next, large scale experiments on deep neural networks demonstrate the performance of VSD against a number of baseline Bayesian inference methods measuring both predictive performance and predictive entropy (for OOD detection).