On the Parameterization of Gaussian Mean Field Posteriors in Bayesian Neural Networks
Authors: Jakub Świątkowski
Kevin Roth
Bastiaan S. Veeling
Linh Tran
Joshua V. Dillon
Jasper Snoek
Stephan Mandt
Tim Salimans
Rodolphe Jenatton
Sebastian Nowozin
Venue:  (2020)
Keywords: variational Bayes
Bayesian neural networks
mean field

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 exhibits 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. What's more, we find that such factorized parameterizations are easier to train since they improve the signal-to-noise ratio of stochastic gradient estimates of the variational lower bound, resulting in faster convergence.

Under review as a conference paper at ICLR 2020
ON THE PARAMETERIZATION OF GAUSSIAN MEAN FIELD POSTERIORS IN BAYESIAN NEURAL NETWORKS
Anonymous authors Paper under double-blind review
ABSTRACT
Variational Bayesian Inference is popular methodology for approximating poste- rior distributions in Bayesian neural networks. Recent work developing this class of methods has explored ever richer parameterizations of the approximate poste- rior in the hope of improving performance. In contrast, here we share curious experimental finding that suggests instead restricting the variational distribution to more compact parameterization. For variety of deep Bayesian neural net- works 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 fac- torized parameterizations improve the signal-to-noise ratio of stochastic gradient estimates of the variational lower bound, resulting in faster convergence.
INTRODUCTION
Bayesian Neural Networks explicitly represent their parameter-uncerrainty by forming a posterior distribution over model parameters, instead of relying on single point estimate for making predic- tions, as is done in traditional deep learning. For neural network weights w, features x and labels y. the posterior distribution p(w/x,y y) is computed using Bayes' rule, which multiplies the prior distribution p(w) (w) and data likelihood p(y w x) and renormalizes. When predicting with Bayesian form average over model predictions where each prediction is generated
neural networks, using a set of parameters that is randomly sampled from the posterior distribution. Bayesian neural networks are thus type of ensembling, of which various types have proven highly effective in deep learning (see c.g. Goodfellow et al., 2016 sec 7.1
Besides offering improved predictive performance over single models, Bayesian ensembles 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 of-mman-omain detection or other tisk-sensitive applications (Ovadia et al.) 2019).
Variational inference is a popular class of methods for approximating the posterior distribution p(wjx,y), since the exact Bayes' rule is often intractable to compute for models of practical in- terest. This class of methods specifies a distribution qo (w) of given parametric or functional form as the posterior approximation, and optimizes the approximation by solving an optimization prob- lem. In particular, we minimize the Kullback-Leible (KL) divergence DKL. between the variational distribution qo (w) and the true posterior distribution p(w xx, y), which is given by
(w(w) p(w E. log
(o(w) p(w)p(y W. wx)p/(x'xx)
DEL (2) (pw)owwwwwpll x, y)] = Ey log
Here, 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 We are then left with what called the negative Evidence Lower Bound (negative ELBO):
Tag Dxalgo(w) (w)]lp(w)) Eg(p(y) w. x)).

---

Under review as a conference paper at ICLR 2020
In practice, the expectation of the likelihood ply w, with respect is usually not analytically tractable and instead is estimated using Monte Carlo sampling:
$ Σ
E [log p(y/w, x)]
log phy) p(y/p(o), xx), w()) (w(),
where the ELBO can then be optimized by differentiating this stochastic approximation with respect to the variational parameters 0 (Salimans al., 2013; Kingma & Welling 2013).
In Gaussian Mean Field Variational Infer- ence (GMFVT), we choose the variational ap- proximation to be fully factorized Gaussian distribution = (pp 2,,) with Wyy ~ N (uun, Ryy), where a layer number, and i and 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 variational ap- proximations, it already doubles the param- eter count compared to deterministic neu- ral networks. In addition, Bayesian neural networks with mean-field Gaussian posterior approximations often are harder train than deterministic neural networks because they suffer from increased noise in stochastic gra- dient estimates.
paramcters
Noive Matrin-wuriate wariate
Eseracctical
VD
joint distribution
factorized distribution
sotmalining
Dowa
Gaussion zocon Sold
stractured
faat
parameter:
Figure Approaches to variational Bayes on Bayesian neural networks, ordered by i) whether they factorize the variational distribution q, and ii) whether they tie the vari- ational parameters.
Beyond mean-field variational inference, recent work on approximate Bayesian inference has ex- plored ever richer parameterizations of the approximate posterior in the hope of improving the per- formance of Bayesian neural networks (see Figure [ In contrast, here we study 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 low- rank structure. This means that by decomposing these variational parameters into a low- make our variational approximation
rank factorization, creasing our model's performance.
compact without de-
Factorized parameterizations of posterior standard deviations are easier train since they improve the signal-to-none 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 possible compress the learned posterior standard deviation matrix using a low-rank approxi- mation 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 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 S)c we investigate instead underlying low-rank structures within the already diagonal covariance of a Gaussian make this explanation more formal in the next section.
variational distribution. We will
2

---

Under review as a conference paper at ICLR 2020
2.1 METHODOLOGY
To introduce the notation we consider layers that consist of a linear transformation followed by non-linearity S.
= W + bj, , h, E RRxm and by, ar, h1-1 E RRinn To simplify the notation in the following, we drop the subscript such that W Wt, My = Mut. Eg Egt and focus on the kernel matrix W for single layer.
h441 = s(ar),
(1)
where W E RRxx,
In GMFVI, (Blundell et al. 2015). we model the variational posterior as 1
(W) = (p Mgg)
IIII (wi)), with ((wrg) Nuue oty),
where H4g E Rmax1 is the posterior mean vector, En E R is the diagonal posterior co- variance matrix. The weights are then usually sampled using reparametrization trick (Kingma & Welling 2013), i.e, for the s-th sample, we have
(+) Pij + dzje(") c~ c~, 1).
Wij
In practice, we often represent the posterior standard deviation parameters oij in the form of matrix A B=x= Note that have the relationship Σ, diag (lag(vec(A2)) where the elementwise-squared A is vectorized by stacking its columns, and then expanded as a diagonal matrix into mm
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 ob- serve that those matrices naturally exhibit a low-rank structure (see Section 2.3 for the corresponding experiments), i.e.
A UVT
(4)
for some U Rmxk V Baak and i small value (e.g.. or 3). This observation motivates the introduction of the following variational family, which we name k-tied Normal:
k-tied-N pa: U, U)V)= = diag ) ((UV))?)))
(5)
where the squaring of the matrix UVT 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.
As formally discussed in Appendix B the matrix variate Gaussian distribution (Gupta & Nagar 2018). referred to as MM and already used for variational inference Louizos & Welling (2016) Snn et al.) (2017), is related to our k-tied Normal dis- tribution with & = 1 when MN uses diag- onal row and column covariances. Interest- ingly, we prove that for k 2 2, 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 (0066) (see also Figure TT. Further- more, notice that diagonal covariance Eg repeatedly reuses the same elements of U and V, which results in parameter sharing across different weights.
Variational family
Parameters (total)
Multivariate Normal
mn mn (m=+1)
diagonal Normal
mit + mrn mn m(m+1] 2 + a(n+1) 2
MN (M(diagonal) k-tied Normal
mn + m +
mn + k(m + n)
Table 1: Number of variational parameters for a variational family for a matrix W E Rmx (dd(diogona)) is from
Louizos & Welling 2016).
The total number of the standard deviation parameters in our method is k(m + n2) from U and V, compared to mn in the standard GMFVI parametrization from A. Given that in our experiments the & is very low (e.g. h- = 2) this reduces the number of parameters from quadratic to linear in the dimensions the layer, see Table I 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 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.
3

---

Under review as a conference paper at ICLR 2020
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 (MIP): network of 3 dense layers and ReLu activations that 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 convolutional layers and 2 dense layers that we train on the CIFAR- 100 dataset (Krizhevsky We use the last 10,000 examples of the training set as a validation set.
et al. 2009b).
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 validation set.
Residual Convolutional Neural Network (ResNet):: ResNet-118 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 mean-field Normal posterior and Normal prior with single scalar standard deviation hyper-parameter for all the layers. We optimize the variational parameters using an Adam optimizer (Kingma & Ba (2014). For more comprehensive explanation of the experimental setup used in this section please refer to Appendix A. 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 (MI.,, 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 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 > Σx "V, where YYu are singular values. We observe that most of the variance in the posterior standard deviation parameters is captured in the rank-1 approximation. However, more fine- ane-grained analysis shows that rank-2 approximation remaining variance. Finally, we note that we do not observe the same behaviour for the posterior mean parameters
encompass nearly all of the
do for the posterior standard deviation parameters. Figure 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.
2.4 LOW-RANK APPROXIMATION OF VARIANCE MATRICES
probabi1icty/exxaplee///fff
4

---

Under review as a conference paper at ICLR 2020
Ceroc marr
Denel mesty
Rans more
Deroe mars
Bese2 moats
Racererikene mese
Ceroe more
Barrel Aedees
200
Deroe sedees
Cerno sidees
300
Figure 2: Explained variance per singular value from SVD of matrices of posterior means and posterior stan- dard 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.
MLP
CNN
LSTM
Rank Full
-ELBO
NLC 1
Accuracy
-ELBO І
NLL J
Accuracy
-ELBO 1
NLL
Accuracy
2.03
0.100.2000
97.6.a15
3.83
2.23
42.1 g2.1
0.536_nors
g.493
80. g
0.677
93.6
4.33
1.30
41.7-0m
1.687
0.491
quaee
80.0
2
0.456
0.107
97.6+uu
3.88
2.24
42.2=04
0.621
0.494
80.1
3
0.450
0.106=coee
97.6
2.24 s00r7
42.1
0.595
coozm
0.493
80.1 1a2r
Table 2: Impact of low-rank approximation of the GMFVI-traine 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.
come 26 JEEROO
10-2
TTATTS
Motivated by the above observation, we show that it is possible to replace the full-rank posterior standard devi- ation matrix with its low-rank approximation without decrease in performance. Table 2 shows the comparison of performance of models with different ranks of approxi- mation to their posterior standard deviation matrix for the MLP, CNN and LSTM models. Figure contains analo- gous results for the ResNet model. The results show that the post-training approximation with ranks higher than achieves predictive performance close to that of the full- rank matrix for all the analyzed model and layer types. This observation itself could be used as form of post- training network compression. Moreover, gives rise to further interesting exploration directions such formu- lating posteriors that exploit such a low rank structure. In the next section we explore this partieular direction while focusing on the first three model types (MLP, CNN, LSTM).
dance ctities dameerrass
10'
6
9
10
Rank
Rank
-ELBO
NLL 4
Accurucy
Full
122.61 =co02
0,495
83.5
1
122.57
0.658
-00000
81.7=00
2
122.77=00m
0,503
83 23.2
122.67 Lomm
0.501 Liaar
83.2
Figure 3: Unlike posterior means, the pos- terior standard deviations of both dense and convolutional layers in the ResNet model trained using standard GMFVI dis- play strong low-rank structure and can be approximated without loss in predictive met- rics. Top: Explained variance per singular value of the matrices of converged posterior means and standard deviations. Bottom: Im- pact of post training low-rank approxima- tion of the posterior standard deviation ma- trices on model's performance. We report mean and SEM of each metric across 100 models samples.
THE k-TIED
NORMAL DISTRIBUTION: EXPLOITING LOW-RANK CRRAMETER-STRUCTUSE 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 de- creasing 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 the context GMFVI 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-lied-N phg U, V) = N dia (vV(()))))
5

---

Under review as a conference paper at ICLR 2020
where the variational parameters are comprised of (Mg U, V).
3.1 EXPERIMENTAL SETTING
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 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.
Wc 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 Usk and Fisk of the k-tied Normal distribution so that after the outer-product operation the respective standard deviations oij have the same mean values as we obtain in the standard GMFVI posterior parametrization. In the experiments for this section KL annealing where linearly scale-up the contribution of the KL term from fraction of its full value to its full contribution over the course of training. Finally, similarly to the previous section, we picked the key pprrmeprrsettrrs such as the KL annealing rate, learning rate, 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. Tigure 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 & 2 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- signal-to-noiss to-so ratid" (SNR) of stochastic gradient estimates of the variational lower bound (ELBO). In particular, focus the gradient SNR of the GMFVI posterior standard deviation parameters for which perform the tying. These parameters are either W4k or Øyk for the k-tied Normal posterior or diy for the standard GMFVI parametrization, all optimized in their log forms for numerical stability. The SNR results in the Figure 4 show 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 [ 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
ASNR for each gradient value is calculated as E 2 /arr(gg)) where is the gradient value for a single parameter. The expectation F and variance Var of the gradient values go are calculated over window of last 10 batches.

---

and SEM across Under review as a conference paper at ICLR 2020
trained layers. More precisely, in the CNN model 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, point out that the k-tied Normal posterior does not increase the training step time compared to the standard parameterization of the GMFVI. See Table in Appendix D for the support of this claim.
MNIST,
Model Dataset
Rank
0.501
-ELBO
NLL
Accuracy
#aar.
Rank
MNIST. MLP Dense 2. SNR step
1000
5000
9000
MNIST, MLP
0.539
0.155
96. Lac
482
Full
4.T3gom
4,45
3,21
MNIST, MLP
o.520
a.129
96.8
484
5840 Pm0
158 135
MNIST, MLP
0.497
0.120
96.9
486
7500
140
4.3
CIFARTOO, CNN
tall
3.72
2,16g
23.9
4,405
7000
117.10
CIFARIOO, CNN
2.12
45.5
2.262
CIFAR100, CNN
3.76
2.15
soute
44.3
2,268
MNIST, MLP ELBO at step
CIFARIO0, CNN
3.73.o0m
2.13 coune
44.3 sc:s
2,273
Rank A
1000
5000
9000
IMDB, LSTM
Tull
0.538 ponest
1.4781
79.5
2.821
Full
42.16
26.52 roate
15.39
IMDB, LSTM
IMDB, LSTM
0.592 FOUME
0.500
0.484
0.512
77.6
78.2
2,694
2,693
2
42.7
42.74
13.97
amem
2000
2.06-300
188
IMDB, LSIM
0.550
0.491
78.8 -0.2
2,695
3
42.63-004
13.61
1.00
Figure 4: Left: impact of the k-tied Normal posterior on test ELBO, test predictive performance and number of model parameters. Test performance reported as mean and SEM across 100 weights samples after training each model for sx300 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 training runs with different random seeds. Right bottom: Negative ELBO on the MNIST validation data set increasing training steps for different ranks tying k. See also Figure 5 which shows negative EI. EI,BO convergence plots for the all three models types.
Tull-mnnc
10
2 Rarii-3
20089
00003 stop
12232
30000
19000
11000
3050
Transdip
Stanstop
Figure 5: Convergence of negative BLBBO (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 the standard parametrization of the GMFVI.
4 RELATED WORK
The application of variational inference to neural networks dates back at least to Peterson (1987);
Hinton & Van Camp 1993). Many developments have followed those seminal research efforts, in particular regarding (I) the expressiveness of the variational posterior distribution and (2) the way the variational parameters themselves can be structured to lead to compact, easier-to-leam and scalable formulations. We organize the discussion of this section around those two aspects, with a specific focus 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 Mar-
lin et al. 201T Titsias & Lázaro-Gredilla 2014 Miller et al. 1017.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 Sceger 2000; Miller et al. 2017 Zhang et al. 2017; Ong al. 2018), the Cholesky
decomposition (Challis & Barber 2011 or by operating instead on the precision matrix (Tan &
Nott 2018 Mishkin al. (018).
'We refer the interested readers Zhang et al. 2018 for recent review of variational inference.

---

Under review as a conference paper at ICLR 2020
Gaussian posterior with block-structured covariances. In the context of Bayesian neural net- works, the layers represent a natural structure to be exploited by the covariance matrix. When assum- ing independence across layers, the resulting covariance matrix exhibits a that has been shown to be a well-performing simplification of the dense setting (Sun et al. 2017 Zhang al. 2017). with both memory and computational benefits.
structure
Within each layer, the corresponding diagonal block of the covariance matrix can be represented by Kronecker product of two smaller matrices (Louizos & Welling 2016|Sum al. 2017), possibly with a parametrization based on rotation matrices (Sun et al. 2017). Finally, using similar tech- niques, Zhang et al. (2017) proposed to use block tridiagonal structure that better approximates the behavior of dense covariance.
Fully factorized mean-field Gaussian posterior. A fully factorized Gaussian variational distri- bution constitutes the simplest option for variational inference. The resulting covariance matrix diagonal and all underlying parameters are assumed to be independent. While the mean-fiell field assumption is known to have some limitations- e.g.. underestimated variance of the posterior dis- tribution Turnerr& & Sahanil 2011) and robustness issues (Giordano et al.] 2018)-it leads to scalable formulations, with already competitive performance, for instance illustrated by the recent uncer- tainty 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
al. 2015 Hernández-Lobato & Adams, 2015 Zhang et al.[2017 Khan 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 covvrrince matrices leads 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 metries). Importantly, we further prove that tying strategies with a rank greater than one cannot be represented in matrix variate Gaussian distribution, thus clearly departing from (Louizos & Welling 2016) (see Appendix B for details).
Our approach can be also interpreted as particular case of hierarchical variational inference (Ran- ganath et al. 2016) where the prior on the variational parameters corresponds to a Dirac distribution, non-zero only when pre-specified low-rank tying relationship holds.
We close this related work section by mentioning the existence other strategies to produce more flexible approximate posteriors, e.g., normalizing flows (Rezende & Mohamed 2015) and exten- sions thereof (Louizos & Welling 2017).
5 CONCLUSION
In this work we have shown that Bayesian Neural Networks trained with standard Gaussian Mcan- 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 parametriza- tion, 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 approx- imate 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.
8

---

Under review as a conference paper at ICLR 2020
REFERENCES
David Barber and Christopher M Bishop. Ensemble learning for multi-layer networks. In Advances in neural information processing systems, pp. 395 395-401, 1998.
Charles Blundell, Julien Cornebise, Koray Kavukeuoglu, and Daan Wierstra. Weight uncertainty in neural networks. arXiv preprint arXiv:1505.05444. 2015.
Edward Challis and David Barber. Concave gaussian variational approximations for inference in large-scale bayesian linear models. In Proceedings of the Fourteenth International Conference on Artificial Intelligence and Statisties, pp. 199-207, 2011.
Joshua Dillon, Ian Langmore, Dustin Tran, Eugene Brevdo, Srinivas Vasudevan, Dave Moore, Brian Patton, Alex Alemi, Matt Hoffman, and Rif A Saurous. Tensorflow distributions. arXiv preprint arXiv: 10604, 2017.
Ryan Giordano, Tamara Broderick, and Michael I Jordan. Covariances, robustness and variational bayes. The Journal of Machine Learning Reseurch, 19()):18112002) 2018.
Ian Goodfellow, Yoshua Bengio, and Aaron Courville. Deep learning. MIT press, 2016.
Alex Graves. Practical variational inference for neural networks. In Advances in neural information processing systems, pp. 2348 -236,, 2011.
Arjun K Gupta and Daya K Nagar. Matrix variate distributions. Chapman and Hall/CRC, 2018.
Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Delving deep into rectifiers: Surpassing human-level performance on imagenet classification. In Proceedings of the IEEE international conference on computer vision, pp. 1026-1034, 2015.
Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recog- nition. In Proceedings of the IEEE conference on computer vision and pailern recognition, pp- 770-778, 2016.
José Miguel Herbandez-Lobato and Ryan Adams. Probabilistic backpropagation for scalable learn- ing of bayesian neural networks. In International Conference on Machine Learning, pp. 1861- 1869, 2015.
Geoffrey Hinton and Drew Van Camp. Keeping neural networks simple by minimizing the descrip- tion length of the weights. In in Proc. of the 6th Ann. ACM Conf. Computational Learning Theory. Citeseer, 1993.
Sepp Hochreiter and Jürgen Schmidhuber. Long short-term memory. Neural compulation, 9(8): 1735-1780, 1997.
Tommi Jaakkola and Michael Jordan. A variational approach to bayesian logistic regression models and their extensions. Sixth International Workshop on Artificial Intelligence and Statisties, volume 82, pp. 4, 1997.
Mohammad Emtiyaz Khan, Didrik Nielsen, Voot Tangkaratt, Wu Lin, Yarin Gal, and Akash Srivas- tava. Fast and scalable bayesian deep learning by 1rXiv:1806.04854 2018.
in adam. arXiv preprint
Diederik Kingma and Jimmy Ba. Adam: method for stochastic optimization. arXiv preprint arXiv:1412.6980 2014.
Diederik Kingma and Max Welling. Auto-encoding variational bayes. arXiv preprint 13/261 2013.
Alex Krizhevsky, Geoffrey Hinton, et al. Learning multiple layers of features from tiny images. Technical report, Citeseer, 2009a.
Alex Krizhevsky, Geoffrey Hinton, et al. Learning multiple layers of features from tiny images. Technical report, Citeseer, 2009b.

---

Under review as a conference paper at ICLR 2020
Yann LeCun and Corinna Cortes. MNIST handwritten digit database. 2010. URL htRp://t:///an
lecun. com/ eeeeexd//mee http///ddmmmni///////m//
Yann LeCun, Léon Bottou, Yoshua Bengio, Patrick Haffner, et al. Gradient-based learning applied to document recognition. Proceedings of the IEEE, 86(1 1):2218-2222444 1998.
Christos Louizos and Max Welling. Structured and efficient variational deep learning with matrix gaussian posteriors. In International Conference Machine Learning, pp. 1708-1716, 2016.
Christos Louizos and Max Welling. Multiplicative normalizing flows for variational bayesian neural networks. In Proceedings of the 34th International Conference on Machine Learning- -oonne 70, pp. 2218-2221. JMLR. org, 2017.
Andrew L Maas, Raymond E Daly, Peter T Pham, Dan Huang, Andrew Y Ng, and Christopher Potts. Learning word vectors for sentiment analysis. In Proceedings of the 49th annual meeting of the association for computational linguistics: Human language technologies-voluee I, pp. 1 142-150. Association for Computational Linguistics, 2011.
Benjamin M Marlin, Mohammad Emtiyaz Khan, and Kevin P Murphy. Piecewise bounds for esti- mating bernoulli-logistic latent gaussian models. In Proceedings of the International Conference on Machine Learning, pp. 633-640, 2011.
Andrew C Miller, Nicholas Foti, and Ryan P Adams. Variational boosting: Iteratively refining pos- terior approximations. Proceedings of the 34th International Conference on Muchine Learning. pp. 2420-2429. JMLR. org, 2017.
Aaron Mishkin, Frederik Kunstner, Didrik Nielsen, Mark Schmidt, and Mohammad Emtiyaz Khan. Slang: Fast structured covariance approximations for bayesian deep learning with natural gradi- ent. In Advances Neural Information Pro Processing Systems, pp. 6245-6255, 2018.
Victor M-H Ong, David Nott, and Michael S Smith. Gaussian variational approximation with factor covariance structure. Journal of Computational and Graphical Statistics, 27(3):465-44) 478,
2018.
Kazuki Osawa, Siddharth Swaroop, Anirudh Jain, Runa Eschenhagen, Richard Turner, Rio Yokota, and Mohammad Emtiyaz Khan. Practical deep learning with bayesian principles. arXiv preprint arXiv: 1206.0206, 2019.
Yaniv Ovadia, Emily Fertig, Jie Ren, Zachary Nado, D Sculley, Sehastian Nowozin, Joshua V Dil- lon, Balaji Lakshminarayanan, and Jasper Snoek. Can you trust your model's uncertainty? eval- uating predictive uncertainty under dataset shift. arXiv preprint D0XXvv:99660053333 2019.
Carsten Peterson. A mean field theory learning algorithm for neural networks. Complex systems, 1: 995-1019, 1987.
Adrian E Raftery, Tilmann Gneiting, Fadoua Balabdaoui, and Michael Polakowski. Using bayesian model averaging to calibrate forecast ensembles. Monthly weather review, (33(5):1155-111)) 2005.
Rajesh Ranganath, Sean Gerrish, and David Blei. Black box variational inference. Artificial
Intelligence and Statistics, pp. 814-8224 2014.
Rajesh Ranganath, Dustin Tran, and David Blei. Hierarchical variational models. In International Conference on Machine Learning, pp. 324-333, 2016.
Danilo Rezende and Shakir Mohamed. Variational inference with normalizing flows. In Interna- tional Conference on Machine Learning, pp. 1530-1538, 2015.
Tim Salimans, David Knowles, et al. Fixed-form variational posterior approximation through stochastic linear regression. Bayesian Analysis. 8(4))837-88)) 2013.
Matthias Seeger. Bayesian model selection for support vector machines, gaussian processes and other kernel classifiers. In Advances in neural information processing systems, pp. 603-609, 2000.
10

---

Under review as a conference paper at ICLR 2020
Shengyang Sun, Changyou Chen, and Lawrence Carin. Learning structured weight uncertainty in bayesian neural networks. In Artificial Intelligence and Statistics, pp. 1283-1292, 2017.
Linda SL Tan and David Nott. Gaussian variational approximation with sparse precision matrices. Statistics and Computing, 28(2):25-2275, 2018.
Michalis Titsias and Miguel
Gredilla. Doubly stochastic variational bayes for non-conjugate inference. In International conference on machine learning, pp. 1971-1979, 2014.
Richard Turner and Maneesh Sahani. Two problems with variational expectation maximisation for time-series models, pp. 109-130. Cambridge University Press, 2011.
Cheng Zhang, Judith Butepage. Hedvig Kjellstrom, and Stephan Mandt. Advances in variational inference. IEEE Iransactions on patlern analysis and machine intelligence, 2018.
Guodong Zhang, Shengyang Sun, David Duvenaud, and Roger Grosse. Noisy natural gradient as variational inference. arXiv preprint arXiv: 2017.
A ADDITIONAL DETAILS
A.1 EXPERIMENTAL DETAILS FOR SECTION 2
We provide here some of the experimental details that we skipped over in the main text of the Section
2 of the paper.
More architecture details:
The 3 layers of the MLP have number of units respectivelly 400, 400 and 10.
The 2 convolutional layers of the CNN have filters of sizes 32 and 64 and dense layers with the number of units 512 and 100.
The embedding and LSTM cell in the LSTM model are both of size 128.
In the MI.s and the CNN models we train both the kernel and bias weights using GMFVI. In the LSTT model we train only the kernel weights using GMFVI and we trained the bias weights using MAP estimation.
It is possible that the SVD decomposition of posterior standard deviation matrix will result in a low-rank approximation that contains negative values. such cases, we threshold the minimum values of the resulting approximation at very low positive constant.
Details of the GMFVI optimisation:
In the MLPs and the CNN models we train both the kernel and bias weights using GMFVI. In the LSTM model we train only the kernel weights using GMFVI and we trained the bias weights using MAP estimation.
For each of the three models the GMFVI uses the standard reparametrization trick (Kingma & Welling 2013).
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.
We initialize the variational posterior means using the standard He initialization (He et al. 2015) and the posterior standard deviations using samples from (0.001, 0.001).
We use Normal prior (00,, ap)) 2 for all the weights and select the o for each of the models separately from set of (0.2, 0.3) based on validation data perforfmance.
For optimization we used an Adam optimizer (Kingma Ba (204)), for which we picked the optimal learning rate for each model from the set of 00.0011, 0. 0.0003, 0.001, 0.003] based on validation set performance.
Additional details:
11

---

Under review as a conference paper at ICLR 2020
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 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 linearly increasing scaling of the KL term during training. We select the best linear coefficient of the KL annealing from x 10 and x 10 10 per batch and update it's value every 100 batches.
We initalize the parameters u,a and Vij that after the outer-product operation the respec- tive dij standard deviations have means at 0.01 before transforming to log-domain. That means that in the log domain the parameters Uzk and Uyk initialized 0.5(log(0.01) log(k)). ). We also add white noise (0, 0.1) 1) to the values of uik and Uyk in the log domain 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 MM.
Consider positive definite matrices Q R=x* and P E Rexe and some arbitrary matrix M Rexe We have by definition that W E IR'x ~ (M, Q, P) and only if vec W) (VeccMM) P& Q), where vee(.) stacks the columns of matrix and 69 is the Kronecker product
The MN has already been used for variational inference by Louizos & Welling (016) and Sun al. (2077). In particular, Louizos & Welling (2016) consider the case where both P and Q restricted to be diagonal matrices. In that case, the resulting distribution corresponds to our k-tied
Normal distribution with k since
P Q - diag(p) & diag(1) ()) = diag (iag@vee(<p))) T)).
Importantly, we prove below that, in the case where k: 2 2, the k-tied Normal distribution cannot be represented matrix variate Gauussian distribution.
Lemma B.1 (Rank-2 matrix and Kronecker product). Ler B be rank-2 matrix in R. There do not exist matrices Q E R and P E cR such that
diag(vec(B)) = Q.
Proof. Let us introduce the shorthand D diag(vec(B)) (B)). By construction, D diagonal and has its diagonal terms strictly positive (it assumed that B R7 i.c., bij > for all i,
We proceed by contradiction. Assume there exist Q Re and P E Rexx such that
Q
This implies that all diagonal blocks of P Q are themselves diagonal with strictly positive diagonal terms. Thus, Pjş Q diagonal for all j (1....,,c) ,J, which implies turn that Q is diagonal, with non-zero diagonal terms and Pjj # 0. Moreover, since the off-diagonal blocks P.ş Q for i + must be zero and Q 전 O, we have Pij = 0 and P also diagonal.
To summarize, if there exist Q E grxr and P E Rexx such that D P Q then it holds that D diag(p) & diag(q) with p E R and q E Rr. This last equality can rewritten as b, = p2 for all E [1...., ,7) and {1.....c) c]. equivalently
B = Bpqp.
This leads to contradiction since qp" has rank one while B is assumed to have rank two.
Figure provides an illustration of the difference between the k-tied Normal and the MA distribu- tion.
12

---

Under review as a conference paper at ICLR 2020
tied parameters
vatria-variate Normal (full)
Normal (ding)
k=1
m(m+1) n(n+ 12/2
m
decresoing nambler of paremeeter
increosing
mn(mn+112
mn
free
parameters
Gaussian
non-factorized distribution
fully factorized distribution
moon field
Figure 6: Illustration of the difference between the k-tied Normal (green), the MN distribution (red), the Gaussian mean field (blue) and the full Gaussian covariance (black) for a matrix of posterior standard devia- tions for layer of size m X k-tied Normal with k is equivalent to MMN with diagonal row and column covariance matrices (half-red, half-green circle). Our experiments show that the formance the mean field. On the other hand, while the full/non-diagonal MM increases the expressiveness of the posterior, it also increases the number of parameters. In contrast, ls-ied Normal with & 2 2 not only decreases the number of parameters, but also matches the predictive performance of the mean field.
1 fails capture the per-
C LOW-RANK STRUCTURE IN POSTERIOR STANDARD DEVIATIONS OF CONVOLUTIONAL LAYERS
Encouraged by the reviewers, we analyze the applicability of our method to different layer types and larger models. In particular, investigate the low-rank structure of the posterior standard deviations of both the dense layers and the convolutional layers in a ResNet-18 18 (He et al. 2016)
GMFVI Bayesian neural network trained the CIFAR-10 10 dataset (Krizhevsky et al. 2009a). More
precisely, we use a publicly available mode from the example repository of Tensorflow Probabil- ity (Dillon al., 2017). To perform the low-rank decomposition of the convolutional layers, flatten all the dimensions of the convolutional layers except the last dimension (e.g.. a layer with shape 3, 512, 512) is reshaped [3 3 512, 512).. See Figure for the examples of the resulting 2 2-dimensional matrices. We then analyze the SVD spectrum of the posterior means and standard deviations of such flattened convolutional layers. Interestingly, we observe similar low-rank struc- ture in the flattened convolutional layers as observed in the dense layers, see Figure 8 The low-rank structure 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 the 2-dimensional space, we can reshape the materialized 2-dimensional matrices back into the 4-dimensional form of a convo- lutional 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 applica- ble not only to the different model types (MLP, CNN and LSTM), but also range of layer types (dense and convolutional) in large and practically applicable models such as the ResNet- 18 model.
Rank
-ELBO 1
NLL 1
Accuracy 1
#Params
&Prras
Full
122.61
10012
0.495
=0.0080
83.5
9,814,026
100.0
122.57 20012
0.658=n.000
200000
81.74039
4,929,711
50.2
2
122.77
20012
0.503
=0,000
83.2
4,946,964
50.4
3
122.67 1:00122
0.501 2000079
83.2 10:37
4,964,217
50.6
Table Impact low-rank approximation of GMFVi-trained posterior standard deviations ResNet-18 model on the model's ELBO and predictive performance. We report mean and SEM of each metric across 100 models samples.
proba
13

---

Under review as a conference paper at ICLR 2020
00R/23 10DIA
conv2d Nopoot
-0.39
D11
-238
135
-114
064
412
SANX 8M
Figure 7: Heat maps of the partially flattened posterior standard deviation tensors for the selected convolutional layers of the ResNet-18 18 GMFVI BNN trained on CIFAR-10. The partially flattened posterior standard deviation tensors of the convolutional layers display similar low-rank patterns that observe for the dense layers.
corvec_ J1 por
carvol Beok
109
10
More
Moae
Figure 8: Explained variance per singular value from SVD partially flattened tensors of posterior means and posterior standard deviations for different convolutional layers of the ResNet-18 18 GMFVI BNN trained on CIFAR-10. 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.
D ADDITIONAL EXPERIMENTAL RESULTS
011
222
30701
201
Figure 9: Heat map of the posterior standard deviation matrix for the weights the first dense layer of LeNet CNN trained using GMFVI on the CIFAR-100 dataset (left), as well as its rank-1 approximation (middle) and rank-2 approximation (right). The rank-2 approximation looks visually similar to the full-rank matrix, confirming our numerical results from Figure
Training method
Train step time [ms] J.
Point estimate
2.00
Standard GMFVI
7.17
K-tied GMFVI
1.1410011
Table 4: Training step evaluation times for simple model architecture with two dense layers for different training methods. We report mean and SEM across a single training run. The k-tied Normal posterior does not increase the train step evaluation times compared to the standard parameterization of the GMFVI posterior. We expect this hold more generally because the biggest additional operation per step when using the k- tied Normal posterior is the UV" multiplication to materialize the matrix of posterior standard deviations A. where U E 2TTh xk V E & n and & is a small value (e.g.. 2 or 3). The time complexity of this operations is Okmm)), which is usually negligible compared to the time complexity of data-weight matrix multiplication O Obmm), where b is the batch size.
EEg.. as used in en ttntuutoiiass////ooooo kwwas/ model
Our GMFVI implementation of this
is available under
https://co ab. resea www.rcc..oooo e com/ wwt.o//drie///// vedgge
VG5s49x cxx cxB- CXw--UFFS9ooryyyvae
14


Reviews:
[{"review_id": "S1xrRMiJcr", "reviewer": "AnonReviewer1", "summary": "", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "The paper proposes a low-rank approximation to the diagonal of Gaussian mean field posterior which reduces the number of parameters to fit. They show that the predictive performance doesn't drop much compared with the full covariance but the number of parameters is significantly reduced. \n\n1. Why Matrix normal distribution is related to k-tied Normal distribution when k=1. When k=1, the rank of UV^\\top is 1. The covariance is matrix normal is U\\otimes V, whose rank is rank(U)rank(V). Also MN has a full covariance for the Gaussian approximation, not mean field. MN is only equal to 1-tied when only diagonal row and column covariances are considered. If that's what the paper means, k-tied is only compared to MN with diagonal row and column covariances. Better make this point clear. \n\n2. Figure 4 should show running time instead of training step. I don't think low-rank approximation is gonna influence the convergence that much. It should affect the evaluation speed of each step. \n\nI think the trick the paper uses is a practical one but not significantly novel enough for the ICLR community. It feels like a standard trick people would do when fitting parameters for large matrices, i.e. exploring the low-rank structure and fitting the factorized matrices. Matrix normal is more significant since it reduces the number of parameters as well as maintaining a full covariance matrix with structures. If just focusing on the diagonal covariance, it already throws away the full covariance. Low-rank won't help much with improving the posterior distribution. ", "extra_scores": {"experience_assessment": "I have published in this field for several years.", "review_assessment:_thoroughness_in_paper_reading": "I read the paper at least twice and used my best judgement in assessing the paper.", "review_assessment:_checking_correctness_of_experiments": "I assessed the sensibility of the experiments.", "review_assessment:_checking_correctness_of_derivations_and_theory": "I carefully checked the derivations and theory."}, "extra_text": {}}, {"review_id": "HJxXIP82YB", "reviewer": "AnonReviewer2", "summary": "", "questions": "", "limitations": "", "rating": 1, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "The paper considers variational Bayesian inference (learning) for neural networks assuming that both the prior distributions and the posterior distributions of the network weights are factorising Gaussians. It is well known that the respective optimisation task for the parameters of the posterior weight distributions (ELBO) is tractable by stochastic gradient ascent.  The authors propose to simplify (restrict) the model even further, by assuming that the posterior variances for fully connected layers, (when seen as a matrix) have low rank. It is pretty obvious that the ELBO optimisation task remains tractable in this case.\n\nThe authors perform two types of experiments to show that the proposed simpler model does not decrease the performance of the network, when compared to the full rank factorising model. They first show that the learned variances of dense layers indeed exhibit a low rank structure for three models learned on the corresponding data (MLP on MNIST, LeNet on CIFAR and LSTM on IMDB). Then, in a second step, they consider the proposed rank constraint during learning and show that with rank k >= 2 the models are able to achieve performance competitive with the full rank model and, moreover, exhibit better signal to noise ratio for the gradients during learning.\n\nThe paper is well written, clearly structured and technically correct. However, in my opinion, its novelty and new insights are restricted, which is why I suggest to reject the paper. The reasons for this are the following.\n\n- The authors restrict their analysis to dense layers only. Moreover, it remains conceptually unclear, why and when the proposed model should be useful and represent a good approximation of the full rank model. \n\n- The experimental study is restricted to small models, e.g. in the case of CIFAR a quite shallow version of LeNet which achieves only ~45% validation accuracy. It remains unclear, whether the observed low rank structure of the variance matrices (of the dense layers) will scale to deeper models that could achieve competitive validation accuracies.\n\n- When measuring the impact of the low rank tying of the posterior variances, the authors compare to the full rank model only. I am missing the \"point estimate\" baseline for these models. For if the the positive impact of the Bayesian inference approach with the full rank model is small when compared to the point estimate, then, as a consequence, the impact of the low rank tying must be small when compared to the full rank model.\n\n- The numbers reported in Table 3 (second group of experiments) raise some questions not addressed by the authors. It remains unclear, why the low rank model with ranks 1,2,3 gives better accuracies on CIFAR than the full rank model. The same holds for the LSTM experiment. This may indicate some kind of \"overfitting\" and should have been analysed by the authors in order to \"disentangle\" possible overfitting issues from the question of validity of the proposed low rank model.", "extra_scores": {"experience_assessment": "I have published one or two papers in this area.", "review_assessment:_checking_correctness_of_derivations_and_theory": "I carefully checked the derivations and theory.", "review_assessment:_checking_correctness_of_experiments": "I carefully checked the experiments.", "review_assessment:_thoroughness_in_paper_reading": "I read the paper thoroughly."}, "extra_text": {}}, {"review_id": "B1xP-y9tFr", "reviewer": "AnonReviewer3", "summary": "", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "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. \n\nWhile 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. \n\nMore 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.\n\nAnother 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]. \n\nSuggestions for future improvement:\n* Re-do the experiments using more complex network architectures, and optionally, on larger datasets / more complex tasks (e.g. image segmentation as in [6]).\n* 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.\n* 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].\n\n# References\n[1] Structured and Efficient Variational Deep Learning with Matrix Gaussian Posteriors\n[2] Overpruning in Variational Bayesian Neural Networks\n[3] A Unified Particle-Optimization Framework for Scalable Bayesian Sampling\n[4] Noisy Natural Gradient as Variational Inference\n[5] Fast and Scalable Bayesian Deep Learning by Weight-Perturbation in Adam\n[6] Bayesian Uncertainty Estimation for Batch Normalized Deep Networks\n[7] Learning Weight Uncertainty with Stochastic Gradient MCMC for Shape Classification\n[8] Function Space Particle Optimization for Bayesian Neural Networks\n[9] Can You Trust Your Model’s Uncertainty? Evaluating Predictive Uncertainty Under Dataset Shift", "extra_scores": {"experience_assessment": "I have published one or two papers in this area.", "review_assessment:_checking_correctness_of_derivations_and_theory": "I assessed the sensibility of the derivations and theory.", "review_assessment:_checking_correctness_of_experiments": "I carefully checked the experiments.", "review_assessment:_thoroughness_in_paper_reading": "I read the paper thoroughly."}, "extra_text": {}}]

Author rebuttal:

Meta-review: This paper proposes to reduce the number of variational parameters for mean-field VI. A low-rank approximation is used for this purpose. Results on a few small problems are reported.

As R3 has pointed out, the main reason to reject this paper is the lack of comparison of uncertainty estimates. I also agree that, recent Adam-like optimizers do use preconditioning that can be interpreted as variances, so it is not clear why reducing this will give better results.

I agree with R2's comments about missing the "point estimate" baseline. Also the reason for rank 1,2,3 giving better accuracies is unclear and I think the reasons provided by the authors is speculative.

I do believe that reducing the parameterization is a reasonable idea and could be useful. But it is not clear if the proposal of this paper is the right one. Due to this reason, I recommend to reject this paper. However, I highly encourage the authors to improve their paper taking these points into account.
Decision: Reject