Structured Dropout Variational Inference for Bayesian Neural Networks

Abstract:
Approximate inference in Bayesian deep networks exhibits a dilemma of how to yield high fidelity posterior approximations while maintaining computational efficiency and scalability. We tackle this challenge by introducing a novel variational structured approximation inspired by the Bayesian interpretation of Dropout regularization. Concretely, we focus on the inflexibility of the factorized structure in Dropout posterior and then propose an improved method called Variational Structured Dropout (VSD). VSD employs an orthogonal transformation to learn a structured representation on the variational Gaussian noise with plausible complexity, and consequently induces statistical dependencies in the approximate posterior. Theoretically, VSD successfully addresses the pathologies of previous Variational Dropout methods and thus offers a standard Bayesian justification. We further show that VSD induces an adaptive regularization term with several desirable properties which contribute to better generalization. Finally, we conduct extensive experiments on standard benchmarks to demonstrate the effectiveness of VSD over state-of-the-art variational methods on predictive accuracy, uncertainty estimation, and out-of-distribution detection.

heading: 1 Introduction
text: 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〉.

heading: 2 Background
text: 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.

heading: 3 Variational Structured Dropout
text: 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.

heading: 3.1 The orthogonal approximation for variational structured noise
text: 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.

heading: 3.2 Derivation of tractable variational objective
text: We consider employing an isotropic Gaussian as the prior distribution, namely p(W) = ∏Q j=1 p(W:j) with p(W:j) = N (0, diag(β−1:j )) and β is a hyper-parameter matrix of the same size with W. With the previous analysis, our approximate posterior qt(W) would be absolutely continuous w.r.t the prior p(W), and thus the KL term DKL(qt(W)||p(W)) is defined. Furthermore, the Gaussian prior helps our proposal avoid the pathologies of improper true posterior and ill-posed inference in VD.
Note that, the prior p(W) is a fully factorized Gaussian which usually facilitates simple analysis and efficient computation. This factorized structure is chosen also because we have no reason for the correlation between non-identical weights at first. Moreover, several arguments indicate that a simple prior over parameter p(W), when interacts with neural nets architecture f(D; W), induces a sophisticated prior over function p(f(D; W)), with desirable properties and useful inductive
biases [86]. However, when we are interested in structured approximations in parameter space, the factorized prior may raise some contradictions. By relative entropy decomposition, we have:
DKL(qt(W)||p(W)) = Q∑ j=1 DKL(qt(W:j)||p(W:j)) + I(W:1,W:2, ...,W:Q), (3)
where I(.) is the mutual information measured by the distribution qt(.), and this term is validly defined in VSD. Maximizing the variational lower bound tends to encourage smaller KL term, and hence constrains the components in RHS of equation (3). Intuitively, a relatively small mutual information can break the strong correlations between the columns of W. Several studies have focused on this limitation and suggested using richer priors such as matrix variate Gaussian [75, 91], doubly semi-implicit distribution [58]. Our solution for this scenario is derived naturally from equation (3), in which we leverage the mutual information as an additional regularization term. Concretely, we maximize an alternative variational objective as follows:
LMI(φ) := Eqα(ξ) log p(D|Θ, ξ (t))− DKL(qt(W)||p(W)) + I(W:1,W:2, ...,W:Q)
= Eqα(ξ) log p(D|Θ, ξ (t))− DKL(q?t (W)||p(W)), (4) where q?t (W) := ∏Q j=1 qt(W:j) is the product of marginal column distributions. From the information-theoretic perspective, augmenting the mutual information is a standard principle for structure learning in Bayesian networks [41]. Particularly in our derivation, maximizing the alternative objective LMI(φ) would help to sustain the dependence structure between columns of the network weights, and thus fixes appropriately the broken ELBO as mentioned above. Interestingly, this technique is utilized reasonably in our method. This is because our dependence structure allows to specify explicitly the marginal distribution on each column of W, leading to a tractable objective in equation (4) whose KL divergence between two multivariate Gaussian can be calculated analytically in closed-form. We note that a similar application to other structured approximations, such as low-rank, Kronecker-factored or matrix variate Gaussian, can be non-trivial. We also remark that our alternative variational objective might be not necessarily a valid lower bound of the original model evidence, but would be the lower bound of new model evidence defined on an alternative prior p̂(W), which satisfies DKL(qt(W)||p̂(W)) = DKL(q?t (W)||p(W)). Indeed, the correlated prior determined by p̂(W) ∝ p(W) ∗ qt(W)/q?t (W) meets this constraint, and thus we could reinterpret the use of mutual information as adopting this prior at each iteration of the training procedure. To clarify, our idea was partly motivated by the similar technique that has been extensively adopted in deep latent variable models, in which a mutual information maximization is also added to the variational lower bound to mitigate the degenerate issue of amortized inference in these models [1, 92].
The KL condition in VSD. With the new variational objective in equation (4), to offer VSD complementary advantages of structured Dropout and Bayesian inference, we need to ensure DKL(q?t (W)||p(W)) satisfies the KL condition. We solve this prerequisite by specifying the precision parameter β of the prior p(W) via the Empirical Bayes (EB) approach. As a result, we obtain an optimal value for this hyperparameter in an analytical form β∗. The optimal β∗ is then substituted back into the prior, and thereby we get the optimal KL term with a form independent of the deterministic weight Θ as follows:
DEBKL(q?t (W)||p(W)) = Q
2 K∑ i=1 log 1 +
∑K j=1 αjU 2 ij
αi . (5)
A formal proof of equation (5) is in Appendix B.
The number of transformations steps T in VSD. An appropriate choice of T is essential in our method. For deep learning architectures, the degree T ∗ of the orthogonal matrix P might be relatively large, so it is quite challenging to adjust the empirical value of T in a principled way to meet the basis-kernel representation theorem. We can employ an efficient parameterization introduced in [90], in which we only need the Householder vectors {vt, t ≥ 1} with sizes much smaller than the order of the matrix P . This will facilitate tuning T with a larger range in an applicable computation time. In our method, we instead only adopt a small T ∈ {1, 2, 3} as a form of "low-degree" approxiamtion, and to make the reflections more expressive, we use a fully connected layer between successive Householder vectors, i.e. vt = FC(vt−1). The low dimension of variational noise in VSD leads to a good adaptation, but with a little trade-off in computational complexity when increasing the number of transformation steps T . We show the performance of VSD when using this neural parameterization with T ∈ {1, 2, 3} in Figure 1, where a larger T could potentially improve the results on predictive measures. However, to maintain computational efficiency, we recommend using T = 1 or 2 in large-scale experiments. Indeed, this neural parameterization has been successfully implemented in the context of learning the latent space in deep latent variable models [79, 5].

heading: 3.3 Joint inference with hierarchical prior
text: We further promote our proposal by introducing a prior hierarchy in VSD framework and then obtain a joint approximation for the Dropout posterior. This will facilitate the flexibility of Bayesian inference in our method in terms of the expressiveness of both prior distribution and approximate posterior. We design a two-level hierarchical prior given by: p(W, z) = p(W|z, β)p(z), where p(W|z, β) = ∏Q j=1 p(W:j |z, β:j) and p(W:j |z, β:j) = N (0, diag(z β −1 :j )); the hyperprior p(z) is a distribution with positive support such as Gamma or half-Cauchy distribution; the latent z has the size of the number of rows and is shared across columns of the weight matrix W; the hyper-parameter matrix β is treated as a scaling factor. This prior family has a centered parameterization and induces several compelling properties such as facilitating feature sparsity [46, 12], model selection [21] or improving robustness, uncertainty calibration [15].
We implement variational inference with a joint approximate posterior, also referred to as the joint Dropout posterior, which is parameterized as follows:
qφ(W, z) = qψ(z)qφ(W|z) with qφ(W:j |z) = N (z Θ:j , diag(z2 α Θ2:j)),
where qφ(W|z) is the conditional Dropout posterior, qψ(z) is chosen depending on the family of prior p(z) so that the reparametrization trick can be utilized. Sampling the random weight W from the joint variational posterior qφ(W, z) includes two steps: z∗ ∼ qψ(z) and W∗ ∼ qφ(W|z∗), in which the second one can be reparameterized as: W∗ = diag(z∗)diag(ξ)Θ = diag(z∗ ξ)Θ, with the noise ξ ∼ N (1K , diag(α)). This new representation adapts to the vanilla Dropout procedure but allows our method to regularize each unit layer with different levels of stochasticity. We derive some insights about the role of hierarchical prior in our framework in Appendix A.2. We apply the Householder transformation to the variational noise ξ and obtain a new joint approximate posterior:
qt(W, z) = qψ(z)qt(W|z) with qt(W:j |z) = N (z Θ:j ,VjUdiag(α)(VjU)T ), where Vj = diag(z Θ:j). Similarly, we define q?t (W|z) = ∏Q j=1 qt(W:j |z) and then optimize an alternative variational objective given by:
L(φ, ψ) := Eqψ(z)qt(W|z) log p(D|W)− Eqψ(z)(DKL(q ? t (W|z, φ)||p(W|z, β))− DKL(qψ(z)||p(z)).
We have chosen the (inverse) Gamma and log-Normal distribution for p(z) and q(z) respectively. These distributions have positive supports, can be reparametrized and the KL-divergence between them also has a closed-form due to the conjugacy. A full derivation of L(φ, ψ) including the KL condition is given in Appendix C. The parameterization of the prior hierarchy in our method is flexible without any simplifying assumptions about hyperprior p(z). We can directly apply it for ARDVariational Dropout framework (ARD-VD) [36] (a derivation is in Appendix F). As the experimental results are shown in Figure 2, the hierarchical prior significantly improves the performance of both ARD-VD and VSD on predictive metrics. Therefore, we aim to introduce VSD with hierarchical prior as an unified framework and a general-purpose approach for Bayesian inference, particularly for BNNs.

heading: 3.4 Scalability of Variational Structured Dropout
text: Approximating a structured posterior directly on the random weights of deep convolutional models is challenging. Besides expensive computation, it is difficult to employ the local reparameterization trick [39], leading to the high variance issue in training. We apply VSD to convolutional layer by learning a structured noise with the size of the number of kernels and imposing it to convolutional weights: ξ ∼ N (1K ,Udiag(α)UT ) and Wijk = ξkΘijk, with i, j, k are the indexes representing height, width, and kernel respectively. This simple solution greatly reduces computational complexity while being able to captures the dependencies among kernels of the convolutional layer.
We present the complexity of MAP and VI methods in terms of computational cost and memory storage in Table 1, with the detailed analysis given in Appendix D. VSD adopts the advantage of Dropout training and maintains an efficiency on both criteria. We also give more results in Table 2 about the empirical computation time of VSD and some other methods. Based on these tables, VSD shows more effective running time even than the mean-field BBNs. Although there is a trade-off when using a larger number of T , VSD does not incur much extra computation time compared to VD.

heading: 3.5 On explicit regularization of Variational Structured Dropout
text: There are several compelling theories to explain the tremendous success of Dropout technique, in which regularization-based is one of the most active approaches [81, 26, 55, 53, 83, 10]. We would follow this direction to investigate inductive biases induced by the structured Dropout in VSD, and from which to consolidate our claim of complementary advantages unified in the proposed method. To characterize the regularization of VSD, we consider a deep linear neural net with L layers parameterized by {Θ(i)}Li=1, and define some notations as: x is an input data, B is the data batch; hi is the i-th hidden layer; Ji(x) denotes the Jacobian of network output w.r.t hi(x); Hi(x) and Hout(x) denotes the Hessian of the loss w.r.t hi(x) and the network output, respectively. Then we have Ji = ( ∏L l=i Θ
(l))T , Θ[i:L] the transposition of linear multiplication of weight matrices from i-th layer to the last one. From a detailed derivation presented in Appendix E, VSD induces an explicit regularization given by:
RV SD = E(x∼B) L∑ i=1 〈 Hi, diag(hi)Udiag(α)UT diag(hi) 〉 .
Note that, Hi can be approximated by JTi HoutJi after ignoring the non-PSD term which is less important empirically [73]. This regularizer offers some intriguing but popular interpretations related
to the curvature of loss landscape [83, 10] (see a detailed explanation in Appendix E). We now show novel properties of VSD induced by the orthogonal matrix U.
VSD imposes a Tikhonov-like regularization and reshapes the gradient of network weights: We rewrite our regularization corresponding to layer i-th by: R(i)V SD = E(x∼B)‖H 1/2 i diag(hi)Udiag(α
1/2)‖2F . This form can be interpreted as the Tikhonov-like regularization imposed on the square root of Hessian matrix Hi, in which the Tikhonov matrix Γ := diag(hi)Udiag(α1/2) is automatically learned during training. This principle can improve the conditioning of the estimation problem. Furthermore, when considering the case of regression problem, we have:
R (i) V SD = E(x∼B) [ Θ[i:L]diag(hi)Udiag(α)UT diag(hi)Θ[i:L].T ] . (6)
This is a data-dependent regularization with adaptive structure determined by the matrix ΓΓT = diag(hi)Udiag(α)UT diag(hi). From the algorithmic perspective, this regularizer allows VSD to reshape the gradient of network weights according to the geometry of the data based on both scale and direction information [14, 25]. Meanwhile ΓΓT = diag(hi)diag(α)diag(hi) only plays as a scaling factor in VD.
VSD penalizes implicitly the spectral norm of weight matrices: Let Ωi := diag(hi)JTi HoutJidiag(hi), then our regularizer can be rewritten as: R(i)V SD = E(x∼B) ∑K k=1 α 2 kU T :kΩiU:k. Since the trainable matrix U satisfies UT:kU:k = 1 for any k, a penalty on U T :kΩiU:k implies that VSD likely prefers a solution with smaller spectral norms of the matrix H1/2out Jidiag(hi) and thus of the network weights. This implication points us to the well-studied theories about generalization bound based on the spectral norm [4, 65]. Concretely, Neyshabur et al. [65] suggests that smaller spectral norm and stable rank can lead to better generalization. This expectation can be observed empirically in VSD through
Table 9 in Appendix E. A more solid investigation about the generalization of VSD is of interest.

heading: 4 Experiments
text: In this section, we provide experimental evaluations to show the effectiveness of our proposed methods compared with the existing methods in terms of both predictability and scalability. We focus mainly on variational inference methods of the following two directions: the first one is direct approximations of the posterior on the random weights of Bayesian nets, including Bayes by Backprop (BBB) [8], Variational Matrix Gaussian (VMG) [47], low-rank approximations (SLANG, ELRG) [56, 80]; and the other one is the Bayesian Dropout methods with MC Dropout (MCD) [19, 20], Variational Dropout (VD) [39], and our method-Variational Structured Dropout (VSD). In addition, we evaluate the performance of point estimate framework MAP and two standard non-variational baselines Deep Ensemble (D.E) [44] and SWAG [50]. Details about data descriptions, network architectures, hyper-parameter tuning are presented in Appendix I.

heading: 4.1 Image classification
text: We now compare the predictive performance of the aforementioned methods for classification tasks on three standard image datasets: MNIST [45], CIFAR10 [43], and SVHN [64]. We evaluate the
predictive probabilities using negative log-likelihood (NLL), error rate, and expected calibration error (ECE) [61, 23]. Details on experimental settings are available in Appendix I.4.
The synthesis results of this experiment are in Table 3. On MNIST, VMG achieves err. rates of 1.17% and 1.27% with FC 400x2 and FC 750x3 respectively, while SLANG reports 1.72% err. rate with FC 400x2. In general, VSD outperforms consistently other variational methods in most settings. For D.E and SWAG, VSD exhibits competitive results on all three metrics. Especially, the figures on NLL and ECE indicate well-calibrated probabilities in our model. This is also a common but noteworthy behavior in structured approximations. On the other hand, for the remaining methods such as MAP and BBB, the error rates are worse by a large margin compared to VSD (respectively about 9% and 5% on CIFAR10, 4% and 2% on SVHN). On CIFAR10 and SVHN, these two methods and VD all show poor results on both NLL and ECE measures, implying that it will be difficult for them to reason properly about the model uncertainty especially in the out-of-distribution context. For MC Dropout, we observe a pretty good performance with the second-best result in variational methods that is similar to those reported of other works [56, 68, 80]. These results of the Bayesian Dropout methods are competitive with structured methods such as VMG, SLANG, ELRG. This further reinforces our motivation about the potential of Dropout methods for improving the predictive performance.

heading: 4.2 Scaling up Bayesian deep convolutional networks
text: We conduct additional experiments to integrate VSD into large-scale convolutional networks. We reproduce the experiments proposed in [80] (ELRG), in which we trained AlexNet and ResNet18 on 4 datasets CIFAR10, SVHN, CIFAR100 [43], and STIL10 [11] to evaluate the predictive performance of our proposal compared to other methods.
The final results are given in Table 4 and Table 5, in which the top two results will be highlighted in bold. For AlexNet, the performance of VSD is more consistent and higher than other variational methods. Modern deeper architectures facilitate deterministic estimates like MAP to better learn discriminative information extracted from training data but also makes its predictions more confident when picking excessively on unique optima. Therefore, although MAP has comparable predictive accuracy on some settings, it comes at a trade-off with the worst results on both NLL and ECE. Meanwhile, ELRG with a low-rank structure on the variational posterior gains desirable properties on uncertainty metrics. Its performance on NLL and ECE are competitive to that of VSD, however, our method obtains significant improvements on accuracy metric. For the remaining methods, BBB performs poorly on CIFAR10 and CIFAR100 in predictive accuracy, but it still has better performance than MAP in terms of NLL and ECE.
For ResNet18 architecture, while MAP and BBB still exhibit the same behavior as mentioned above, VSD continues to achieve the convincing results, namely, it has the best performance on CIFAR100 and SVHN over all three metrics when compared to variational-based methods. On CIFAR10 and STL10, VSD also gets competitive statistics on NLL and accuracy. Overall, compared with MCD, VD, and ELRG, our method maintains good performance with better stability.

heading: 4.3 Predictive entropy performance
text: We now evaluate the predictive uncertainty of each model on out-of-distribution settings that have been implemented in previous works [44, 48, 68, 80]. We evaluate the entropy of predictive distribution p(y∗|x∗,D) and use the density of this entropy to assess the quality of uncertainty estimates. Basically, an accurate and well-calibrated model is expected to represent entropy values being concentrated mostly around 0 (i.e. high confidence) when the test data comes from the same underlying distribution as the training data, and in the opposite case, the predictive entropies should be evenly distributed (i.e. higher uncertainty). In fact, the deep learning models do not achieve simultaneously on both expectations at the most ideal, but instead, accurate and well-calibrated ones tend to exhibit a moderate level of confidence on in-distribution data, and then provide a reasonable representation for uncertainty estimates on out-of-distribution data.
For LeNet, we train the model on SVHN dataset and then consider out-of-distribution data from CIFAR10 and CIFAR100. The results are shown in Figure 3. All methods work well on in-distribution data SVHN with the entropy value being distributed mostly around zero. However, the entropy densities of MAP and VD are concentrated excessively. This indicates that these methods would tend to make overconfident predictions on out-of-distribution data. This claim is consolidated by the qualitative results on CIFAR10 and CIFAR100 datasets. In contrast, MCD, BBB, and VSD are well-calibrated with a moderate level of confidence for in-distribution data. On CIFAR10 and CIFAR100 datasets, VSD gains better results with entropy values being distributed over a larger support, meaning that the predictions of VSD are closer to uniform on unseen classes.
We run a similar experiment, in which we train AlexNet on CIFAR10 and use SVHN, CIFAR100 as out-of-distribution data. The results are shown in the top row of Figure 4 in Appendix G.1. While MAP and VD still exhibit the same overconfident phenomenon as on LeNet, we observe the underconfident predictions of BBB and MC Dropout even on in-distribution data, which possibly leads to a high uncertainty on out-of-distribution data. We hypothesize that this is because the models trained with these methods are likely underfit with a low accuracy on the in-distribution training data. In contrast, VSD estimates reasonably the predictive entropy in both settings. The remaining scenario with ResNet18 trained on CIFAR100 is shown in the bottom row of Figure 4 with the same behaviors.

heading: 5 Conclusions
text: We proposed a novel approximate inference framework for Bayesian deep nets, named Variational Structured Dropout (VSD). In VSD, we learn a structured approximate posterior via the Dropout principle. VSD is able to yield a flexible inference while maintaining computational efficiency and scalability for deep convolutional models. The extensive experiments have evidenced the advantages of VSD such as well-calibrated prediction, better generalization, good uncertainty estimation. Given a consistent performance of VSD as presented throughout the paper, an extension of that method to other problems, such as Bayesian active learning or reinforcement learning, is of interest.

heading: Acknowledgements
text: We would like to thank the anonymous reviewers for their insightful comments and valuable suggestions. We also thank Ngo Trung Nghia (VinAI Research) for helpful discussions throughout this work. Khoat Than was funded by Gia Lam Urban Development and Investment Company Limited, Vingroup, and supported by Vingroup Innovation Foundation (VINIF) under project code VINIF.2019.DA18.

Review summaries:
This paper introduces Variational Structured Dropout (VSD), a scalable method building on previous Variational Dropout methods, and with a structured representation of the variational noise. The paper derives the method, introduces a hierarchical prior, and provides results on many benchmarks while comparing to many baselines. VSD performs reasonably well in all benchmarks on uncertainty metrics and accuracy. VSD also performs well on out-of-distribution uncertainty experiments.

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).