O N THE P ARAMETERIZATION OF G AUSSIAN M EAN F IELD P OSTERIORS IN B AYESIAN N EURAL N ETWORKS
Anonymous authors
Paper under double-blind review
A BSTRACT
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 ﬁnding that suggests instead restricting the variational distribution to a more compact parameterization.
For a variety of deep Bayesian neural networks trained using Gaussian mean-ﬁeld variational inference, we ﬁnd 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 ﬁnd that such factorized parameterizations improve the signal-to-noise ratio of stochastic gradient estimates of the variational lower bound, resulting in faster convergence.
I NTRODUCTION
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 speciﬁes 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
[Equation]
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):
[Equation]
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:
[Equation]
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).
[Figure]
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-ﬁeld Gaussian posterior approximations often are harder to train than deterministic neural networks because they suffer from increased noise in stochastic gradient estimates.
Beyond mean-ﬁeld 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 ﬁnd 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.
M EAN 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 ﬁrst brieﬂy 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. M ETHODOLOGY
To introduce the notation we consider layers that consist of a linear transformation followed by a non-linearity f ,
[Equation]
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
[Equation]
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-ﬁeld variational distributions.
Interestingly, we observe that those matrices naturally exhibit a low-rank structure (see Section 2.3 for the corresponding experiments), i.e,
[Equation]
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-ﬁeld variational distribution family.
[Table]
[Equation]
As formally discussed in Appendix B, the matrix variate Gaussian distribution (Gupta & Nagar, 2018), referred to as MN and already used for variational inference by Louizos & Welling (2016) and Sun et al.
(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.
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 ﬁrst 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. E XPERIMENTAL SETTING
Before providing the experimental results we brieﬂy 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) withconvolutional 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-ﬁeld 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 speciﬁcally 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 brieﬂy 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. M AIN 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 ﬁrst 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 ﬁrst three models.
The percent of explained variance for the singular value k is calculated as γ 2 k / (cid:80) i (cid:48) γ i (cid:48) , where γ i (cid:48) 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 ﬁne-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 rank2 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
L OW - RANK APPROXIMATION OF VARIANCE MATRICES
[Figure] [Figure] [Figure] 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 thanachieves 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.
In the next section we explore this particular direction while focusing on the ﬁrst three model types (MLP, CNN,
LSTM).
T HE k - TIED N ORMAL D ISTRIBUTION : E XPLOITING L OW -R ANK P ARAMETER -S TRUCTURE IN M EAN F IELD P OSTERIORS
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 deﬁnition of the k -tied Normal distribution:
[Equation]
where the variational parameters are comprised of { µ q , U , V } .
3. E XPERIMENTAL 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 ﬁrst 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. E XPERIMENTAL RESULTS
We ﬁrst 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 beneﬁts 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 signiﬁcant portion of the total number of the model parameters (e.g. MLPs and CNNs with dense layers for classiﬁcation).
We further investigate the impact of the k -tied Normal posterior distribution on the signal-to-noise ratio 2 (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 signiﬁcant 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 Table 4 in Appendix D for the support of this claim
[Table] [Figure] [Figure] l d
[Figure]
o u t n e g
[Figure]
L B
O
[Figure]
n [Figure]  R ELATED WORK The application of variational inference to neural networks dates back at least to Peterson (1987);
Hinton & Van Camp (1993).
Many developments 3 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 speciﬁc 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 efﬁciently 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 simpliﬁcation of the dense setting (Sun et al., 2017;
Zhang et al., 2017), with both memory and computational beneﬁts.
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-ﬁeld 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-ﬁeld 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 quantiﬁcation benchmark of Ovadia et al.
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-speciﬁed low-rank tying relationship holds.
We close this related work section by mentioning the existence of other strategies to produce more ﬂexible approximate posteriors, e.g., normalizing ﬂows (Rezende & Mohamed, 2015) and extensions thereof (Louizos & Welling, 2017).
C ONCLUSION
In this work we have shown that Bayesian Neural Networks trained with standard Gaussian MeanField 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-ﬁeld 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.
A A DDITIONAL D ETAILS A. E XPERIMENTAL DETAILS FOR SECTION
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 ﬁlters of sizes 32 and 64 and 2 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 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.
• 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.
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 N (0 . 01 , 0 . 001) .
• We use a Normal prior N (0 , σ p ) for all the weights and select the σ p for each of the models separately from a set of [0 . 2 , 0 . 3] based on validation data perforfmance.
• For optimization we used an Adam optimizer (Kingma & Ba (2014)), for which we picked the optimal learning rate for each model from the set of [0 . 0001 , 0 . 0003 , 0 . 001 , 0 . 003] based on validation set performance.
Additional details:
• 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. E XPERIMENTAL DETAILS FOR SECTION
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 coefﬁcient of the KL annealing from 5 × 10 − 5 and 5 × 10 − 6 per batch and update it’s value everybatches.
• 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 P ROOF OF THE M ATRIX V ARIATE N ORMAL P ARAMETERIZATION
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 deﬁnite matrices Q ∈ R r × r and P ∈ R c × c and some arbitrary matrix M ∈ R r × c .
We have by deﬁnition 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.
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
[Equation]
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
[Equation]
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 (cid:54) = 0 .
Moreover, since the off-diagonal blocks p ij Q for i (cid:54) =
j must be zero and Q (cid:54) = 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 (cid:62) .
This leads to a contradiction since qp (cid:62) has rank one while B is assumed to have rank two. Figure 6 provides an illustration of the difference between the k -tied Normal and the MN distribution. tied parameters non-factorized distribution fully factorized distribution free parameters m + n mn m(m+1)/2 + n(n+1)/ d e cr e a s i n g n u m b e r o f p a r a m e t e r
s
k
=
M a t r i x - v a r i a t e N o r m a l (f u ll )
M a t r i x - v a r i a t e N o r m a l ( d i a g ) i n cr e a s i n g k k - t i e d
N o r m a l
G a u ss i a n m e a n f i e l d mn(mn+1)/
C L OW - 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, we investigate the low-rank structure of the posterior standard deviations of both the dense layers and the convolutional layers in a ResNet-
+}
(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 model 4 from the example repository of Tensorﬂow Probability (Dillon et al., 2017).
To perform the low-rank decomposition of the convolutional layers, we ﬂatten 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 ﬂattened convolutional layers.
Interestingly, we observe a similar low-rank structure in the ﬂattened 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.
[Table]
[Figure]
[Figure]
P e r c e n t
o f e x p l a i n e d v a r i a n c e
[Figure]
D A DDITIONAL EXPERIMENTAL RESULTS
[Table]
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 to hold more generally because the biggest additional operation per step when using the k - tied Normal posterior is the UV T multiplication to materialize the matrix of posterior standard deviations A , where U ∈ R m × k , V ∈ R n × k and k is a small value (e.g., 2 or 3).
The time complexity of this operations is O ( kmn ) , which is usually negligible compared to the time complexity of data-weight matrix multiplication O ( bmn ) , where b is the batch size.
Our GMFVI implementation of this model is available under