Title: ON THE PARAMETERIZATION OF GAUSSIAN MEAN FIELD POSTERIORS IN BAYESIAN NEURAL NETWORKS

Abstract: Variational Bayesian Inference is a popular methodology for approximating posterior distributions in Bayesian neural networks. Recent work developing this class of methods has explored ever richer parameterizations of the approximate posterior in the hope of improving performance. In contrast, here we share a curious experimental finding that suggests instead restricting the variational distribution to a more compact parameterization. For a variety of deep Bayesian neural networks trained using Gaussian mean-field variational inference, we find that the posterior standard deviations consistently exhibit strong low-rank structure after convergence. This means that by decomposing these variational parameters into a low-rank factorization, we can make our variational approximation more compact without decreasing the models' performance. Furthermore, we find that such factorized parameterizations improve the signal-to-noise ratio of stochastic gradient estimates of the variational lower bound, resulting in faster convergence.

1. INTRODUCTION:

Bayesian Neural Networks explicitly represent their parameter-uncertainty by forming a posterior distribution over model parameters, instead of relying on a single point estimate for making predictions, as is done in traditional deep learning. For neural network weights w, features x and labels y, the posterior distribution p(w|x, y) is computed using Bayes' rule, which multiplies the prior distribution p(w) and data likelihood p(y|w, x) and renormalizes. When predicting with Bayesian neural networks, we form an average over model predictions where each prediction is generated using a set of parameters that is randomly sampled from the posterior distribution. Bayesian neural networks are thus a type of ensembling, of which various types have proven highly effective in deep learning (see e.g. Goodfellow et al., 2016, sec 7.11) .

Besides offering improved predictive performance over single models, Bayesian ensembles are also more robust because ensemble members will tend to make different predictions on hard examples (Raftery et al., 2005) . In addition, the diversity of the ensemble represents predictive uncertainty and can be used for out-of-domain detection or other risk-sensitive applications (Ovadia et al., 2019) .

Variational inference is a popular class of methods for approximating the posterior distribution p(w|x, y), since the exact Bayes' rule is often intractable to compute for models of practical interest. This class of methods specifies a distribution q θ (w) of given parametric or functional form as the posterior approximation, and optimizes the approximation by solving an optimization problem. In particular, we minimize the Kullback-Leibler (KL) divergence D KL between the variational distribution q θ (w) and the true posterior distribution p(w|x, y), which is given by D KL [q θ (w)||p(w|x, y)] = E q log q θ (w) p(w|x, y) = E q log q θ (w) p(w)p(y|w, x)/p (y|x) .

Here, we do not know the normalizing constant of the exact posterior p(y|x), but since this term does not depend on w, we may ignore it for the purpose of optimizing our approximation q. We are then left with what is called the negative Evidence Lower Bound (negative ELBO):

L q = D KL [q θ (w)||p(w)] -E q [p(y|w, x)].

In practice, the expectation of the likelihood p(y|w, x) with respect to q is usually not analytically tractable and instead is estimated using Monte Carlo sampling:

E q [log p(y|w, x)] ≈ 1 S S s=1

log p(y|w (s) , x), w (s) ∼ q θ (w),

where the ELBO can then be optimized by differentiating this stochastic approximation with respect to the variational parameters θ (Salimans et al., 2013; Kingma & Welling, 2013) . : Approaches to variational Bayes on Bayesian neural networks, ordered by i) whether they factorize the variational distribution q, and ii) whether they tie the variational parameters.

In Gaussian Mean Field Variational Inference (GMFVI), we choose the variational approximation to be a fully factorized Gaussian distribution q = N (µ q , Σ q ) with w lij ∼ N (µ lij , σ 2 lij ), where l is a layer number, and i and j are the row and column indices in the layer's weight matrix. Although this type of approximation is considered to be one of the simplest types of variational approximations, it already doubles the parameter count compared to deterministic neural networks. In addition, Bayesian neural networks with mean-field Gaussian posterior approximations often are harder to train than deterministic neural networks because they suffer from increased noise in stochastic gradient estimates.

Beyond mean-field variational inference, recent work on approximate Bayesian inference has explored ever richer parameterizations of the approximate posterior in the hope of improving the performance of Bayesian neural networks (see Figure 1 ). In contrast, here we study a simpler, more compactly parameterized variational approximation, which we show can also work well for a variety of models. In particular we find that:

• Converged posterior standard deviations under GMFVI consistently display strong lowrank structure. This means that by decomposing these variational parameters into a lowrank factorization, we can make our variational approximation more compact without decreasing our model's performance. • Factorized parameterizations of posterior standard deviations are easier to train since they improve the signal-to-noise ratio of stochastic gradient estimates of the variational lower bound. Using a more limited parameterization of the approximate posterior can thus lead to faster convergence and a reduction in number of parameters compared to standard GMFVI.

2. MEAN FIELD POSTERIOR STANDARD DEVIATIONS NATURALLY HAVE LOW-RANK STRUCTURE:

In this section we show that the converged posterior standard deviations of Bayesian neural networks trained using standard GMFVI consistently display strong low-rank structure. We also show that it is possible to compress the learned posterior standard deviation matrix using a low-rank approximation without decreasing the network's performance. We first briefly introduce the mathematical notation for our GMFVI setting and the low-rank approximation that we explore. We then provide experimental results that support the two main claims of this section.

To avoid any confusion among the readers, we would like to clarify that we use the terminology "low-rank" in a particular context. While variational inference typically makes use of low-rank decompositions to compactly represent the dense covariance of a Gaussian variational distribution (see numerous references in Section 4), we investigate instead underlying low-rank structures within the already diagonal covariance of a Gaussian fully-factorized variational distribution. We will make this explanation more formal in the next section.

2.1. METHODOLOGY:

To introduce the notation we consider layers that consist of a linear transformation followed by a non-linearity f ,

a l = h l W l + b l , h l+1 = f (a l ),

(1) where W l ∈ R m×n , h l ∈ R 1×m and b l , a l , h l+1 ∈ R 1×n . To simplify the notation in the following, we drop the subscript l such that W = W l , µ q = µ ql , Σ q = Σ ql and we focus on the kernel matrix W for a single layer.

In GMFVI, (Blundell et al., 2015) , we model the variational posterior as

EQUATION

where µ q ∈ R mn×1 is the posterior mean vector, Σ q ∈ R mn×mn + is the diagonal posterior covariance matrix. The weights are then usually sampled using a reparametrization trick (Kingma & Welling, 2013), i.e, for the s-th sample, we have s) , ∼ N (0, 1).

EQUATION

(3) In practice, we often represent the posterior standard deviation parameters σ ij in the form of a matrix A ∈ R m×n + . Note that we have the relationship Σ q = diag(vec(A 2 )) where the elementwise-squared A is vectorized by stacking its columns, and then expanded as a diagonal matrix into R mn×mn + .

In the sequel, we start by empirically studying the properties of the spectrum of matrices A post training, while using standard Gaussian mean-field variational distributions. Interestingly, we observe that those matrices naturally exhibit a low-rank structure (see Section 2.3 for the corresponding experiments), i.e, A ≈ UV T (4) for some U ∈ R m×k , V ∈ R n×k and k a small value (e.g., 2 or 3). This observation motivates the introduction of the following variational family, which we name k-tied Normal:

EQUATION

)

where the squaring of the matrix UV T is applied elementwise. Due to the tied parametrization of the diagonal covariance matrix, we emphasize that this variational family is smaller-i.e., included in-the standard Gaussian mean-field variational distribution family.

Variational family Parameters (total) Multivariate Normal mn + mn (mn+1) (2017) , is related to our k-tied Normal distribution with k = 1 when MN uses diagonal row and column covariances. Interestingly, we prove that for k ≥ 2, our k-tied Normal distribution cannot be represented by any MN distribution. This illustrates the main difference of our approach from the most closely related previous work of Louizos & Welling (2016) (see also Figure 1 ). Furthermore, notice that diagonal covariance Σ q repeatedly reuses the same elements of U and V, which results in parameter sharing across different weights.

2 diagonal Normal mn + mn MN mn + m(m+1) 2 + n(n+1) 2 MN (diagonal) mn + m + n k-tied Normal mn + k(m + n)

The total number of the standard deviation parameters in our method is k(m + n) from U and V, compared to mn in the standard GMFVI parametrization from A. Given that in our experiments the k is very low (e.g. k = 2) this reduces the number of parameters from quadratic to linear in the dimensions of the layer, see Table 1 . More importantly, such parameter sharing across the weights leads to higher signal-to-noise ratio during training and thus faster convergence. We demonstrate this phenomena in the next section. In the rest of this section, we will first demonstrate that the standard GMFVI methods already learn a low-rank structure in the posterior standard deviation matrix A. Furthermore, we will show that replacing the full matrix A with its low-rank approximation does not reduce the predictive performance.

2.2. EXPERIMENTAL SETTING:

Before providing the experimental results we briefly explain the key properties of the experimental setting. We analyse three types of GMFVI Bayesian neural network models:

• Multilayer Perceptron (MLP): a network of 3 dense layers and ReLu activations that we train on the MNIST dataset (LeCun & Cortes, 2010) . We use the last 10,000 examples of the training set as a validation set.

• Convolutional Neural Network (CNN): a LeNet architecture (LeCun et al., 1998) with 2 convolutional layers and 2 dense layers that we train on the CIFAR-100 dataset (Krizhevsky et al., 2009b) . We use the last 10,000 examples of the training set as a validation set.

• Long Short-Term Memory (LSTM): a model that consists of an embedding and an LSTM cell (Hochreiter & Schmidhuber, 1997) , followed by a single unit dense layer. We train it on an IMBD dataset (Maas et al., 2011) , in which we use the last 5,000 examples of the training set as a validation set.

• Residual Convolutional Neural Network (ResNet): a ResNet-18 1 architecture (He et al., 2016) trained on all 50,000 training examples of the CIFAR-10 dataset (Krizhevsky et al., 2009a) .

In each of the four models we use a mean-field Normal posterior and a Normal prior with a single scalar standard deviation hyper-parameter for all the layers. We optimize the variational parameters using an Adam optimizer (Kingma & Ba ( 2014)). For a more comprehensive explanation of the experimental setup used in this section please refer to Appendix A.1. Finally, we highlight that our experiments focus primarily on the comparison across a broad range of model types rather than competing with the state-of-the-art results over the specifically used datasets. Therefore, we use mainly small to medium models (MLP, CNN, LSTM) that are known to train well using the standard GMFVI approach explored in this paper. Encouraged by the reviewers, we also briefly show that our results can extend to larger models such as the ResNet-18 model. However, scaling GMFVI to such larger model sizes is still a challenging research problem (Osawa et al., 2019) .

2.3. MAIN EXPERIMENTAL OBSERVATION:

Our main experimental observation is that the standard GMFVI learns posterior standard deviation matrices that have a low-rank structure across different model types (MLP, CNN, LSTM, ResNet) and layer types (dense and convolutional). To show this, we investigate the results of the SVD decomposition of posterior standard deviation matrices for the four types of models trained until ELBO convergence using GMFVI. While we evaluate the low-rank structure only for the dense layers of the first three models (MLP, CNN and LSTM), we investigate also the low-rank structure of the convolutional layers of the ResNet model.

Figure 2 shows the percentage of explained variance per singular value k of the SVD decomposition of dense layers in the first three models. The percent of explained variance for the singular value k is calculated as γ 2 k / i γ i , where γ i are singular values. We observe that most of the variance in the posterior standard deviation parameters is captured in the rank-1 approximation. However, a more fine-grained analysis shows that a rank-2 approximation can encompass nearly all of the remaining variance. Finally, we note that we do not observe the same behaviour for the posterior mean parameters as we do for the posterior standard deviation parameters. Figure 9 in Appendix D further supports this claim visually by comparing the heat maps of the full-rank posterior standard deviations matrix with its rank-1 and rank-2 approximations. In particular, we observe that the rank-2 approximation results in the heat-map looking visually very similar to the full-rank matrix. Finally, Figure 3 illustrates that the low-rank structure is also visible in both the dense and convolutional layers of the ResNet model. See Appendix C for more details. In the analysis of the experiments, we use the shorthand SEM to refer to the standard error of the mean. Motivated by the above observation, we show that it is possible to replace the full-rank posterior standard deviation matrix with its low-rank approximation without a decrease in performance. Table 2 shows the comparison of performance of models with different ranks of approximation to their posterior standard deviation matrix for the MLP, CNN and LSTM models. Figure 3 contains analogous results for the ResNet model. The results show that the post-training approximation with ranks higher than 1 achieves predictive performance close to that of the fullrank matrix for all the analyzed model and layer types. This observation itself could be used as a form of posttraining network compression. Moreover, it gives rise to further interesting exploration directions such as formulating posteriors that exploit such a low rank structure.

2.4. LOW-RANK APPROXIMATION OF VARIANCE MATRICES:

In the next section we explore this particular direction while focusing on the first three model types (MLP, CNN, LSTM).

3. THE k-TIED NORMAL DISTRIBUTION: EXPLOITING LOW-RANK PARAMETER-STRUCTURE IN MEAN FIELD POSTERIORS:

In the previous section we have shown that it is possible to replace a full-rank matrix of posterior standard deviations trained with GMFVI with this matrix's low-rank approximation without decreasing the predictive performance. In this section we show that it is also possible to exploit this observation during training time. We achieve this by exploiting our novel variational family, the k-tied Normal distribution (see Section 2.1).

We show that using this distribution in the context of GMFVI in Bayesian neural networks allows to reduce the number of network parameters and in some cases speed up model convergences while maintaining the predictive performance of the standard parametrization of the GMFVI.

We start by recalling the definition of the k-tied Normal distribution:

k-tied -N (W; µ q , U, V) = N µ q , diag vec (UV T ) 2

where the variational parameters are comprised of {µ q , U, V}.

We now introduce the experimental setting in which we evaluate the GMFVI variational posterior parametrized by the k-tied Normal distribution. We assess the impact of the described posterior in terms of predictive performance and reduction in the number of parameters for the same first three model types (MLP, CNN, LSTM) and respective datasets (MNIST, CIFAR-100, IMDB) as we used in the previous section. Additionally, we also analyse the impact of tying in the posterior on the signal-to-noise ratio of stochastic gradient estimates of the variational lower bound for the CNN model as a representative example. Overall the experimental setup is very similar to the one introduced in the previous section. Therefore, we highlight here only the key differences.

We apply the k-tied Normal variational posterior distribution to the layers which we analysed in the previous section. More concretely, we use the k-tied Normal variational posterior for all the three layers of the MLP model, the two dense layers of the CNN model and the LSTM cell's kernel and recurrent kernel. We initialize the parameters u ik and v jk of the k-tied Normal distribution so that after the outer-product operation the respective standard deviations σ ij have the same mean values as we obtain in the standard GMFVI posterior parametrization. In the experiments for this section we use KL annealing where we linearly scale-up the contribution of the KL term from a fraction of its full value to its full contribution over the course of training. Finally, similarly to the previous section, we picked the key hyper-parameters such as the KL annealing rate, a learning rate, a batch size and a prior standard deviation based on the performance on respective validation sets. For more information about the experimental details for this section please refer to Appendix A.2.

3.2. EXPERIMENTAL RESULTS:

We first investigate the comparison of predictive performance of GMFVI Bayesian neural network models trained using the k-tied Normal posterior distribution with the models trained using the standard parametrization of the GMFVI. Figure 4 shows the results for the three model types on the test splits of their respective datasets. We analyse variants of the k-tied Normal posterior distribution with different levels of tying k. We observe that for k ≥ 2 the k-tied Normal posterior is able to achieve the performance competitive with the standard GMFVI posterior parametrization, while reducing the total number of model parameters. The benefits of using the k-tied Normal posterior are the most visible for models where the dense layers with the k-tied Normal posterior constitute for a significant portion of the total number of the model parameters (e.g. MLPs and CNNs with dense layers for classification).

We further investigate the impact of the k-tied Normal posterior distribution on the signal-to-noise ratio2 (SNR) of stochastic gradient estimates of the variational lower bound (ELBO). In particular, we focus on the gradient SNR of the GMFVI posterior standard deviation parameters for which we perform the tying. These parameters are either u ik or v jk for the k-tied Normal posterior or σ ij for the standard GMFVI parametrization, all optimized in their log forms for numerical stability. The SNR results in the Figure 4 show a significant increase in the gradient SNR when using the k-tied Normal posterior.

Consequently, we observe that the increase in the gradient SNR translates into faster convergence of the negative ELBO objective in some of the analyzed models. Figure 5 shows the convergence plots of validation negative ELBO for all the three model types. We observe that the impact of the k-tied Normal posterior on convergence depends on the model type. As shown in Figure 4 , the impact on the MLP model is strong and consistent with the k-tied Normal posterior increasing convergence speed compared to the standard GMFVI parameterization. For the LSTM model we also observe a similar speed-up. However, for the CNN model the impact of the k-Normal posterior on the ELBO convergence is much smaller. We hypothesize that this is due to the fact that we use the k-tied Normal posterior for all the layers trained using GMFVI in the MLP and the LSTM models, while in the CNN model we use the k-tied Normal posterior only for some of the GMFVI trained layers. More precisely, in the CNN model we use the k-tied Normal posterior only for the two dense layers, while the two convolutional layers are trained using the standard parameterization of the GMFVI. Finally, we point out that the k-tied Normal posterior does not increase the training step time compared to the standard parameterization of the GMFVI. See 

4. RELATED WORK:

The application of variational inference to neural networks dates back at least to Peterson (1987) ; Hinton & Van Camp (1993) . Many developments3 have followed those seminal research efforts, in particular regarding (1) the expressiveness of the variational posterior distribution and (2) the way the variational parameters themselves can be structured to lead to compact, easier-to-learn and scalable formulations. We organize the discussion of this section around those two aspects, with a specific focus on the Gaussian case.

Full Gaussian posterior. Because of their substantial memory and computational cost, Gaussian variational distributions with full covariance matrices have been primarily applied to (generalized) linear models and shallow neural networks (Jaakkola & Jordan, 1997; Barber & Bishop, 1998; Marlin et al., 2011; Titsias & Lázaro-Gredilla, 2014; Miller et al., 2017; Ong et al., 2018) .

To represent the dense covariance matrix efficiently in terms of variational parameters, several schemes have been proposed, including the sum of low-rank plus diagonal matrices (Barber & Bishop, 1998; Seeger, 2000; Miller et al., 2017; Zhang et al., 2017; Ong et al., 2018) , the Cholesky decomposition (Challis & Barber, 2011) or by operating instead on the precision matrix (Tan & Nott, 2018; Mishkin et al., 2018) .

Gaussian posterior with block-structured covariances. In the context of Bayesian neural networks, the layers represent a natural structure to be exploited by the covariance matrix. When assuming independence across layers, the resulting covariance matrix exhibits a block-diagonal structure that has been shown to be a well-performing simplification of the dense setting (Sun et al., 2017; Zhang et al., 2017) , with both memory and computational benefits.

Within each layer, the corresponding diagonal block of the covariance matrix can be represented by a Kronecker product of two smaller matrices (Louizos & Welling, 2016; Sun et al., 2017) , possibly with a parametrization based on rotation matrices (Sun et al., 2017) . Finally, using similar techniques, Zhang et al. (2017) proposed to use a block tridiagonal structure that better approximates the behavior of a dense covariance.

Fully factorized mean-field Gaussian posterior. A fully factorized Gaussian variational distribution constitutes the simplest option for variational inference. The resulting covariance matrix is diagonal and all underlying parameters are assumed to be independent. While the mean-field assumption is known to have some limitations-e.g., underestimated variance of the posterior distribution (Turner & Sahani, 2011) and robustness issues (Giordano et al., 2018) -it leads to scalable formulations, with already competitive performance, as for instance illustrated by the recent uncertainty quantification benchmark of Ovadia et al. (2019) .

Because of its simplicity and scalability, the fully-factorized Gaussian variational distribution has been widely used for Bayesian neural networks (Graves, 2011; Ranganath et al., 2014; Blundell et al., 2015; Hernández-Lobato & Adams, 2015; Zhang et al., 2017; Khan et al., 2018) .

Our approach can be seen as an attempt to further reduce the number of parameters of the (already) diagonal covariance matrix. Closest to our approach is the work of Louizos & Welling (2016) .

Their matrix variate Gaussian distribution instantiated with the Kronecker product of the diagonal row-and column-covariance matrices leads to a rank-1 tying of the posterior variances. In contrast, we explore tying strategies beyond the rank-1 case, which we show to lead to better performance (both in terms of ELBO and predictive metrics). Importantly, we further prove that tying strategies with a rank greater than one cannot be represented in a matrix variate Gaussian distribution, thus clearly departing from (Louizos & Welling, 2016) (see Appendix B for details).

Our approach can be also interpreted as a particular case of hierarchical variational inference (Ranganath et al., 2016) where the prior on the variational parameters corresponds to a Dirac distribution, non-zero only when a pre-specified low-rank tying relationship holds.

We close this related work section by mentioning the existence of other strategies to produce more flexible approximate posteriors, e.g., normalizing flows (Rezende & Mohamed, 2015) and extensions thereof (Louizos & Welling, 2017) .

5. CONCLUSION:

In this work we have shown that Bayesian Neural Networks trained with standard Gaussian Mean-Field Variational Inference learn posterior standard deviation matrices that can be approximated with little information loss by low-rank SVD decompositions. This suggests that richer parameterizations of the variational posterior may not always be better, and that compact parameterizations can also work well. We used this insight to propose a simple, yet effective variational posterior parametrization, which speeds up training and reduces the number of variational parameters without degrading predictive performance on 3 different model types.

In future work, we hope to scale up variational inference with compactly parameterized approximate posteriors to much larger models and more complex problems. For mean-field variational inference to work well in that setting several challenges will likely need to be addressed (Osawa et al., 2019) ; improving the signal-to-noise ratio of ELBO gradients using our compact variational parameterizations may provide a piece of the puzzle.

• It is possible that the SVD decomposition of posterior standard deviation matrix will result in a low-rank approximation that contains negative values. In such cases, we threshold the minimum values of the resulting approximation at very low positive constant.

A.2 EXPERIMENTAL DETAILS FOR SECTION 3

We provide here some of the experimental details that we skipped over in the main text of the Section 3 of the paper:

• For the k-tied Normal posteriors to work reliably we used KL annealing with a linearly increasing scaling of the KL term during training. We select the best linear coefficient of the KL annealing from 5 × 10 -5 and 5 × 10 -6 per batch and update it's value every 100 batches.

• We initalize the parameters u ik and v jk so that after the outer-product operation the respective σ ij standard deviations have means at 0.01 before transforming to log-domain. That means that in the log domain the parameters u ik and v jk are initialized as 0.5(log(0.01)log(k)). We also add white noise N (0, 0.1) to the values of u ik and v jk in the log domain to break symmetry.

B PROOF OF THE MATRIX VARIATE NORMAL PARAMETERIZATION:

In this section of the appendix, we formally explain the connections between the k-tied Normal distribution and the matrix variate Gaussian distribution (Gupta & Nagar, 2018), referred to as MN .

Consider positive definite matrices Q ∈ R r×r and P ∈ R c×c and some arbitrary matrix M ∈ R r×c . We have by definition that W ∈ R r×c ∼ MN (M, Q, P) if and only if vec(W) ∼ N (vec(M), P⊗ Q), where vec(•) stacks the columns of a matrix and ⊗ is the Kronecker product

The MN has already been used for variational inference by Louizos & Welling (2016) and Sun et al. (2017) . In particular, Louizos & Welling (2016) consider the case where both P and Q are restricted to be diagonal matrices. In that case, the resulting distribution corresponds to our k-tied Normal distribution with k = 1 since

P ⊗ Q = diag(p) ⊗ diag(q) = diag(vec(qp )).

Importantly, we prove below that, in the case where k ≥ 2, the k-tied Normal distribution cannot be represented as a matrix variate Gauussian distribution. Lemma B.1 (Rank-2 matrix and Kronecker product). Let B be a rank-2 matrix in R r×c + . There do not exist matrices Q ∈ R r×r and P ∈ R c×c such that diag(vec(B)) = P ⊗ Q.

Proof. Let us introduce the shorthand D = diag(vec(B)). By construction, D is diagonal and has its diagonal terms strictly positive (it is assumed that B ∈ R r×c + , i.e., b ij > 0 for all i, j). We proceed by contradiction. Assume there exist Q ∈ R r×r and P ∈ R c×c such that D = P ⊗ Q.

This implies that all diagonal blocks of P⊗Q are themselves diagonal with strictly positive diagonal terms. Thus, p jj Q is diagonal for all j ∈ {1, . . . , c}, which implies in turn that Q is diagonal, with non-zero diagonal terms and p jj = 0. Moreover, since the off-diagonal blocks p ij Q for i = j must be zero and Q = 0, we have p ij = 0 and P is also diagonal.

To summarize, if there exist Q ∈ R r×r and P ∈ R c×c such that D = P ⊗ Q, then it holds that D = diag(p) ⊗ diag(q) with p ∈ R c and q ∈ R r . This last equality can be rewritten as b ij = p j q i for all i ∈ {1, . . . , r} and j ∈ {1, . . . , c}, or equivalently

B = qp .

This leads to a contradiction since qp has rank one while B is assumed to have rank two. Encouraged by the reviewers, we analyze the applicability of our method to different layer types and larger models. In particular, we investigate the low-rank structure of the posterior standard deviations of both the dense layers and the convolutional layers in a ResNet-18 (He et al., 2016) GMFVI Bayesian neural network trained on the CIFAR-10 dataset (Krizhevsky et al., 2009a) . More precisely, we use a publicly available model4 from the example repository of Tensorflow Probability (Dillon et al., 2017) . To perform the low-rank decomposition of the convolutional layers, we flatten all the dimensions of the convolutional layers except the last dimension (e.g., a layer with shape [3, 3, 512, 512 ] is reshaped to [3 • 3 • 512, 512]). See Figure 7 for the examples of the resulting 2-dimensional matrices. We then analyze the SVD spectrum of the posterior means and standard deviations of such flattened convolutional layers. Interestingly, we observe a similar low-rank structure in the flattened convolutional layers as observed in the dense layers, see Figure 8 . The low-rank structure is the most visible for the last convolutional layer, which also contain the highest number of parameters.

Importantly, note that after performing the low-rank approximation in the 2-dimensional space, we can reshape the materialized 2-dimensional matrices back into the 4-dimensional form of a convolutional layer. Table 3 shows that such a low-rank approximation in the convolutional layers of the analyzed ResNet-18 model can be performed without a loss in the model's predictive performance, while halving the total number of parameters. These results suggest that our method can be applicable not only to the different model types (MLP, CNN and LSTM), but also to a range of layer types (dense and convolutional) in large and practically applicable models such as the ResNet-18 model. 

Rank:



:

SNR for each gradient value is calculated asE[g 2 b ]/Var[g 2 b ], where g b is the gradient value for a single parameter. The expectation E and variance V ar of the gradient values g b are calculated over a window of last 10 batches.

We refer the interested readers to Zhang et al. (2018) for a recent review of variational inference.

https://github.com/tensorflow/probability/blob/master/tensorflow_ probability/examples/cifar10_bnn.py

Figures and Tables:

Figure1: Approaches to variational Bayes on Bayesian neural networks, ordered by i) whether they factorize the variational distribution q, and ii) whether they tie the variational parameters.

Figure2: Explained variance per singular value from SVD of matrices of posterior means and posterior standard deviations for different layers of three types of models trained using standard GMFVI: MLP (left), CNN (center), LSTM (right). Posterior standard deviations clearly display strong low-rank structure, with most of the variance contained in the top few singular values, while this is not the case for posterior means.

Figure3: Unlike posterior means, the posterior standard deviations of both dense and convolutional layers in the ResNet model trained using standard GMFVI display strong low-rank structure and can be approximated without loss in predictive metrics. Top: Explained variance per singular value of the matrices of converged posterior means and standard deviations. Bottom: Impact of post training low-rank approximation of the posterior standard deviation matrices on model's performance. We report mean and SEM of each metric across 100 models samples.

Figure 5: Convergence of negative ELBO (lower is better) reported for validation dataset when training with tied variational posterior standard deviations for MLP (left), CNN (center), and LSTM (right) with different low-rank factorizations of the posterior standard deviation matrix. Full-rank is the standard parametrization of the GMFVI.

Figure6provides an illustration of the difference between the k-tied Normal and the MN distribution.

Number of variational parameters for a variational family for a matrix W ∈ R m×n . MN (diagonal) is fromLouizos & Welling (2016).

Table content: <table><tr><td>As formally discussed in Appendix B, the</td></tr><tr><td>matrix variate Gaussian distribution (Gupta</td></tr><tr><td>&amp; Nagar, 2018), referred to as MN</td></tr><tr><td>and already used for variational inference</td></tr><tr><td>by Louizos &amp; Welling (2016) and Sun et al.</td></tr></table>

Impact of low-rank approximation of the GMFVI-trained posterior standard deviation matrix on ELBO and predictive performance, for three types of models. We report mean and SEM of each metric across 100 models samples.

Table content: <table><tr><td/><td/><td>MLP</td><td/><td/><td>CNN</td><td/><td/><td/><td>LSTM</td></tr><tr><td>Rank</td><td>-ELBO ↓</td><td>NLL ↓</td><td colspan="2">Accuracy ↑ -ELBO ↓</td><td>NLL ↓</td><td colspan="2">Accuracy ↑</td><td>-ELBO ↓</td><td>NLL ↓</td><td>Accuracy ↑</td></tr><tr><td>Full</td><td colspan="2">0.431±0.0057 0.100±0.0034</td><td>97.6±0.15</td><td colspan="2">3.83±0.020 2.23±0.017</td><td>42.1±0.49</td><td/><td colspan="2">0.536±0.0058 0.493±0.0057</td><td>80.1±0.25</td></tr><tr><td>1</td><td>3.41±0.019</td><td>0.677±0.0040</td><td>93.6±0.25</td><td colspan="2">4.33±0.021 2.30±0.016</td><td>41.7±0.49</td><td/><td colspan="2">0.687±0.0058 0.491±0.0056</td><td>80.0±0.25</td></tr><tr><td>2</td><td colspan="2">0.456±0.0059 0.107±0.0033</td><td>97.6±0.15</td><td colspan="2">3.88±0.020 2.24±0.017</td><td>42.2±0.49</td><td/><td colspan="2">0.621±0.0058 0.494±0.0057</td><td>80.1±0.25</td></tr><tr><td>3</td><td colspan="2">0.450±0.0059 0.106±0.0033</td><td>97.6±0.15</td><td colspan="2">3.86±0.020 2.24±0.017</td><td>42.1±0.49</td><td/><td colspan="2">0.595±0.0058 0.493±0.0056</td><td>80.1±0.25</td></tr><tr><td/><td/><td/><td/><td/><td/><td>Percent of explained variance</td><td>10 -5 10 -4 10 -3 10 -2 10 -1 10 0</td><td/><td>conv2d stddevs conv2d means conv2d_7 stddevs conv2d_7 means dense stddevs dense means</td></tr><tr><td/><td/><td/><td/><td/><td/><td/><td>10 -6</td><td colspan="2">Rank 1 2 3 4 5 6 7 8 9 10</td></tr><tr><td/><td/><td/><td/><td/><td/><td>Rank</td><td/><td>-ELBO ↓</td><td>NLL ↓</td><td>Accuracy ↑</td></tr><tr><td/><td/><td/><td/><td/><td/><td>Full</td><td colspan="3">122.61±0.012 0.495±0.0080</td><td>83.5±0.37</td></tr><tr><td/><td/><td/><td/><td/><td/><td>1</td><td colspan="3">122.57±0.012 0.658±0.0069</td><td>81.7±0.39</td></tr><tr><td/><td/><td/><td/><td/><td/><td>2</td><td colspan="3">122.77±0.012 0.503±0.0080</td><td>83.2±0.37</td></tr><tr><td/><td/><td/><td/><td/><td/><td>3</td><td colspan="3">122.67±0.012 0.501±0.0079</td><td>83.2±0.37</td></tr></table>

Table 4 in Appendix D for the support of this claim. Left: impact of the k-tied Normal posterior on test ELBO, test predictive performance and number of model parameters. Test performance is reported as a mean and SEM across 100 weights samples after training each model for ≈300 epochs. Right top: mean gradient SNR in the Dense 2 layer of the MNIST MLP model at increasing training steps for different ranks of tying k.We observe a similar increase in the SNR from tying for the CNN and the LSTM models as for the MLP model shown here. We report mean and SEM across 3 training runs with different random seeds. Right bottom: Negative ELBO on the MNIST validation data set at increasing training steps for different ranks of tying k. See also Figure5, which shows negative ELBO convergence plots for the all three models types.

Table content: <table><tr><td colspan="2">Model &amp; Dataset Rank k MNIST, MLP full MNIST, MLP 1 MNIST, MLP 2 MNIST, MLP 3 CIFAR100, CNN full</td><td colspan="2">-ELBO ↓ 0.501±0.0061 0.133±0.0040 NLL ↓ 0.539±0.0063 0.155±0.0043 0.520±0.0063 0.129±0.0039 0.497±0.0060 0.120±0.0038 3.72±0.018 2.16±0.016</td><td colspan="2">Accuracy ↑ #Par. [k] ↓ 96.8±0.18 957 96.1±0.19 482 96.8±0.18 484 96.9±0.18 486 43.9±0.50 4,405</td><td>Rank k full 1 2 3</td><td>MNIST, MLP Dense 2, SNR at step 1000 5000 9000 4.13±0.027 4.45±0.091 3.21±0.035 5840±190 158±3.8 5.3±0.20 7500±240 140±11 4.3±0.26 7000±270 117±1.7 4.1±0.20</td></tr><tr><td>CIFAR100, CNN</td><td>1</td><td>3.65±0.017</td><td>2.12±0.015</td><td>45.5±0.50</td><td>2,262</td><td/></tr><tr><td>CIFAR100, CNN CIFAR100, CNN IMDB, LSTM IMDB, LSTM IMDB, LSTM IMDB, LSTM</td><td>2 3 full 1 2 3</td><td colspan="2">3.76±0.019 3.73±0.018 0.538±0.0054 0.478±0.0052 2.15±0.016 2.13±0.016 0.592±0.0041 0.512±0.0040 0.560±0.0042 0.484±0.0041 0.550±0.0051 0.491±0.0050</td><td>44.3±0.50 44.3±0.50 79.5±0.26 77.6±0.26 78.2±0.26 78.8±0.26</td><td>2,268 2,273 2,823 2,693 2,694 2,695</td><td>Rank k full 1 2 3</td><td>MNIST, MLP, -ELBO at step 1000 5000 9000 42.16±0.070 26.52±0.016 15.39±0.016 43.11±0.039 14.85±0.017 2.06±0.027 42.74±0.090 13.97±0.023 1.82±0.017 42.63±0.068 13.61±0.020 1.80±0.031</td></tr><tr><td>Figure 4:</td><td/><td/><td/><td/><td/><td/></tr></table>

Impact of the low-rank approximation of the GMFVI-trained posterior standard deviations of a ResNet-18 model on the model's ELBO and predictive performance. We report mean and SEM of each metric across 100 models samples.

Table content: <table><tr><td/><td>-ELBO ↓</td><td>NLL ↓</td><td colspan="3">Accuracy ↑ #Params ↓ %Params ↓</td></tr><tr><td>Full</td><td colspan="2">122.61±0.012 0.495±0.0080</td><td>83.5±0.37</td><td>9,814,026</td><td>100.0</td></tr><tr><td>1</td><td colspan="2">122.57±0.012 0.658±0.0069</td><td>81.7±0.39</td><td>4,929,711</td><td>50.2</td></tr><tr><td>2</td><td colspan="2">122.77±0.012 0.503±0.0080</td><td>83.2±0.37</td><td>4,946,964</td><td>50.4</td></tr><tr><td>3</td><td colspan="2">122.67±0.012 0.501±0.0079</td><td>83.2±0.37</td><td>4,964,217</td><td>50.6</td></tr></table>

Peer review (iclr2020):
This paper showed the (diagonal) variance parameters in mean-field VI for BNNs exhibit a low-rank structure, and that training from scratch using such a low-rank parameterization lead to comparable performance as well as increased SNR of the gradient. While the observation is somewhat interesting, currently it is only verified in a narrow range of network architectures, and it*s unclear if the observation and the proposed method will still be useful on network architectures used in real-world applications. As such, I believe this work would be more suitable as a workshop presentation. More specifically, the models considered are MLP on MNIST, LeNet on CIFAR-100 and LSTM on IMDB. These choices are not practical, as the reported performance indicates (e.g. 45% accuracy on CIFAR-100); as such these results cannot support the claim that the proposed low-rank parameterization could be useful in practice: while MFVI can be useful on some model architectures, it lead to pathologies on others, especially on smaller networks. See Fig.1 in [1] for an example and [2] for a possible explanation. Also note that the reported accuracy on MNIST is ~2% worse than the typical values in BNN papers using a comparable setting, e.g. [3]. These facts unfortunately lead to the doubt that the proposed low-rank parameterization could only match the performance of MFVI when MFVI is not that useful. Another major concern is that I*m not sure if the proposed low-rank variational would actually save parameters in practice, since the variance parameter in MFVI could already be stored as the preconditioners in Adam-like optimizers [4-5]. Suggestions for future improvement: * Re-do the experiments using more complex network architectures, and optionally, on larger datasets / more complex tasks (e.g. image segmentation as in [6]). * Also, consider setups more commonly used in previous BNN papers, e.g. VGG/ResNet on CIFAR-10 has been used in [4,7,8]. CIFAR-100 could be sufficiently complex as a BNN benchmark, but few papers reported results on it. * Report the quality of the learned uncertainty, either directly as in [9] or using performance of downstream tasks, e.g. RL and adversarial robustness as in [4,8]. # References [1] Structured and Efficient Variational Deep Learning with Matrix Gaussian Posteriors [2] Overpruning in Variational Bayesian Neural Networks [3] A Unified Particle-Optimization Framework for Scalable Bayesian Sampling [4] Noisy Natural Gradient as Variational Inference [5] Fast and Scalable Bayesian Deep Learning by Weight-Perturbation in Adam [6] Bayesian Uncertainty Estimation for Batch Normalized Deep Networks [7] Learning Weight Uncertainty with Stochastic Gradient MCMC for Shape Classification [8] Function Space Particle Optimization for Bayesian Neural Networks [9] Can You Trust Your Model’s Uncertainty? Evaluating Predictive Uncertainty Under Dataset Shift