Sparse Uncertainty Representation in Deep Learning with Inducing Weights
Year: 2022

Abstract:
Bayesian Neural Networks and deep ensembles represent two modern paradigms of uncertainty quantification in deep learning. Yet these approaches struggle to scale mainly due to memory inefficiency, requiring parameter storage several times that of their deterministic counterparts. To address this, we augment each weight matrix with a small inducing weight matrix, projecting the uncertainty quantification into a lower dimensional space. We further extend Matheron’s conditional Gaussian sampling rule to enable fast weight sampling, which enables our inference method to maintain reasonable run-time as compared with ensembles. Importantly, our approach achieves competitive performance to the state-of-the-art in prediction and uncertainty estimation tasks with fully connected neural networks and ResNets, while reducing the parameter size to ≤ 24.3% of that of a single neural network.

heading: 1 Introduction
text: Deep learning models are becoming deeper and wider than ever before. From image recognition models such as ResNet-101 (He et al., 2016a) and DenseNet (Huang et al., 2017) to BERT (Xu et al., 2019) and GPT-3 (Brown et al., 2020) for language modelling, deep neural networks have found consistent success in fitting large-scale data. As these models are increasingly deployed in real-world applications, calibrated uncertainty estimates for their predictions become crucial, especially in safety-critical areas such as healthcare. In this regard, Bayesian Neural Networks (BNNs) (MacKay, 1995; Blundell et al., 2015; Gal & Ghahramani, 2016; Zhang et al., 2020) and deep ensembles (Lakshminarayanan et al., 2017) represent two popular paradigms for estimating uncertainty, which have shown promising results in applications such as (medical) image processing (Kendall & Gal, 2017; Tanno et al., 2017) and out-of-distribution detection (Ovadia et al., 2019).
Though progress has been made, one major obstacle to scaling up BNNs and deep ensembles is their high storage cost. Both approaches require the parameter counts to be several times higher than their deterministic counterparts. Although recent efforts have improved memory efficiency (Louizos & Welling, 2017; Świątkowski et al., 2020; Wen et al., 2020; Dusenberry et al., 2020), these still use more parameters than a deterministic neural network. This is particularly problematic in hardware-constrained edge devices, when on-device storage is required due to privacy regulations.
Meanwhile, an infinitely wide BNN becomes a Gaussian process (GP) that is known for good uncertainty estimates (Neal, 1995; Matthews et al., 2018; Lee et al., 2018). But perhaps surprisingly, this infinitely wide BNN is “parameter efficient”, as its “parameters” are effectively the datapoints, which have a considerably smaller memory footprint than explicitly storing the network weights. In addition, sparse posterior approximations store a smaller number of inducing points instead (Snelson & Ghahramani, 2006; Titsias, 2009), making sparse GPs even more memory efficient.
∗Work done at Microsoft Research Cambridge.
35th Conference on Neural Information Processing Systems (NeurIPS 2021).
Can we bring the advantages of sparse approximations in GPs — which are infinitely-wide neural networks — to finite width deep learning models? We provide an affirmative answer regarding memory efficiency, by proposing an uncertainty quantification framework based on sparse uncertainty representations. We present our approach in BNN context, but the proposed approach is also applicable to deep ensembles. In detail, our contributions are as follows:
• We introduce inducing weights — an auxiliary variable method with lower dimensional counterparts to the actual weight matrices — for variational inference in BNNs, as well as a memory efficient parameterisation and an extension to ensemble methods (Section 3.1).
• We extend Matheron’s rule to facilitate efficient posterior sampling (Section 3.2). • We provide an in-depth computation complexity analysis (Section 3.3), showing the signifi-
cant advantage in terms of parameter efficiency. • We show the connection to sparse (deep) GPs, in that inducing weights can be viewed as
projected noisy inducing outputs in pre-activation output space (Section 5.1). • We apply the proposed approach to BNNs and deep ensembles. Experiments in classification,
model robustness and out-of-distribution detection tasks show that our inducing weight approaches achieve competitive performance to their counterparts in the original weight space on modern deep architectures for image classification, while reducing the parameter count to ≤ 24.3% of that of a single network.
We open-source our proposed inducing weight approach, together with baseline methods reported in the experiments, as a PyTorch (Paszke et al., 2019) wrapper named bayesianize: https: //github.com/microsoft/bayesianize. As demonstrated in Appendix I, our software makes the conversion of a deterministic neural network to a Bayesian one with a few lines of code: import bnn # our pytorch wrapper package net = torchvision.models.resnet18() # construct a deterministic ResNet18 bnn.bayesianize_(net, inference="inducing") # convert it into a Bayesian one

heading: 2 Inducing variables for variational inference
text: Our work is built on variational inference and inducing variables for posterior approximations. Given observations D = {X,Y} with X = [x1, ...,xN ], Y = [y1, ...,yN ], we would like to fit a neural network p(y|x,W1:L) with weights W1:L to the data. BNNs posit a prior distribution p(W1:L) over the weights, and construct an approximate posterior q(W1:L) to the exact posterior p(W1:L|D) ∝ p(D|W1:L)p(W1:L), where p(D|W1:L) = p(Y|X,W1:L) = ∏N n=1 p(yn|xn,W1:L).
Variational inference Variational inference (Hinton & Van Camp, 1993; Jordan et al., 1999; Zhang et al., 2018a) constructs an approximation q(θ) to the posterior p(θ|D) ∝ p(θ)p(D|θ) by maximising a variational lower-bound:
log p(D) ≥ L(q(θ)) := Eq(θ) [log p(D|θ)]−KL [q(θ)||p(θ)] . (1)
For BNNs, θ = {W1:L}, and a simple choice of q is a Fully-factorized Gaussian (FFG): q(W1:L) = ∏L l=1 ∏dlout i=1 ∏dlin j=1N (m (i,j) l , v (i,j) l ), withm (i,j) l , v (i,j) l the mean and variance of W (i,j) l and dlin, d l out the respective number of inputs and outputs to layer l. The variational parameters are then φ = {m(i,j)l , v (i,j) l }Ll=1. Gradients of L w.r.t. φ can be estimated with mini-batches of data (Hoffman et al., 2013) and with Monte Carlo sampling from the q distribution (Titsias & LázaroGredilla, 2014; Kingma & Welling, 2014). By setting q to an BNN, a variational BNN can be trained with similar computational requirements as a deterministic network (Blundell et al., 2015).
Improved posterior approximation with inducing variables Auxiliary variable approaches (Agakov & Barber, 2004; Salimans et al., 2015; Ranganath et al., 2016) construct the q(θ) distribution with an auxiliary variable a: q(θ) = ∫ q(θ|a)q(a)da, with the hope that a potentially richer mixture distribution q(θ) can achieve better approximations. As then q(θ) becomes intractable, an auxiliary variational lower-bound is used to optimise q(θ,a) (see Appendix B):
log p(D) ≥ L(q(θ,a)) = Eq(θ,a)[log p(D|θ)] + Eq(θ,a) [ log
p(θ)r(a|θ) q(θ|a)q(a)
] . (2)
Here r(a|θ) is an auxiliary distribution that needs to be specified, where existing approaches often use a “reverse model” for r(a|θ). Instead, we define r(a|θ) in a generative manner: r(a|θ) is the “posterior” of the following “generative model”, whose “evidence” is exactly the prior of θ:
r(a|θ) = p̃(a|θ) ∝ p̃(a)p̃(θ|a), such that p̃(θ) := ∫ p̃(a)p̃(θ|a)da = p(θ). (3)
Plugging Eq. (3) into Eq. (2):
L(q(θ,a)) = Eq(θ)[log p(D|θ)]− Eq(a) [KL[q(θ|a)||p̃(θ|a)]]−KL[q(a)||p̃(a)]. (4) This approach yields an efficient approximate inference algorithm, translating the complexity of inference in θ to a. If dim(a) < dim(θ) and q(θ,a) = q(θ|a)q(a) has the following properties:
1. A “pseudo prior” p̃(a)p̃(θ|a) is defined such that ∫ p̃(a)p̃(θ|a)da = p(θ);
2. The conditionals q(θ|a) and p̃(θ|a) are in the same parametric family, so can share parameters; 3. Both sampling θ ∼ q(θ) and computing KL[q(θ|a)||p̃(θ|a)] can be done efficiently; 4. The designs of q(a) and p̃(a) can potentially provide extra advantages (in time and space
complexities and/or optimisation easiness).
We call a the inducing variable of θ, which is inspired by variationally sparse GP (SVGP) with inducing points (Snelson & Ghahramani, 2006; Titsias, 2009). Indeed SVGP is a special case (see Appendix C): θ = f , a = u, the GP prior is p(f |X) = GP(0,KXX), p(u) = GP(0,KZZ), p̃(f ,u) = p(u)p(f |X,u), q(f |u) = p(f |X,u), q(f ,u) = p(f |X,u)q(u), and Z are the optimisable inducing inputs. The variational lower-bound is L(q(f ,u)) = Eq(f)[log p(Y|f)]−KL[q(u)||p(u)], and the variational parameters are φ = {Z, distribution parameters of q(u)}. SVGP satisfies the marginalisation constraint Eq. (3) by definition, and it has KL[q(f |u)||p̃(f |u)] = 0. Also by using small M = dim(u) and exploiting the q distribution design, SVGP reduces run-time from O(N3) to O(NM2 + M3) where N is the number of inputs in X, meanwhile it also makes storing a full Gaussian q(u) affordable. Lastly, u can be whitened, leading to the “pseudo prior” p̃(f ,v) = p(f |X,u = K1/2ZZ v)p̃(v), p̃(v) = N (v; 0, I) which could bring potential benefits in optimisation. We emphasise that the introduction of “pseudo prior” does not change the probabilistic model as long as the marginalisation constraint Eq. (3) is satisfied. In the rest of the paper we assume the constraint Eq. (3) holds and write p(θ,a) := p̃(θ,a). It might seem unclear how to design such p̃(θ,a) for an arbitrary probabilistic model, however, for a Gaussian prior on θ the rules for computing conditional Gaussian distributions can be used to construct p̃. In Section 3 we exploit these rules to develop an efficient approximate inference method for Bayesian neural networks with inducing weights.

heading: 3 Sparse uncertainty representation with inducing weights

heading: 3.1 Inducing weights for neural network parameters
text: Following the above design principles, we introduce to each network layer l a smaller inducing weight matrix Ul to assist approximate posterior inference in Wl. Therefore in our context, θ = W1:L and a = U1:L. In the rest of the paper, we assume a factorised prior across layers p(W1:L) = ∏ l p(Wl), and drop the l indices when the context is clear to ease notation.
Augmenting network layers with inducing weights Suppose the weight W ∈ Rdout×din has a Gaussian prior p(W) = p(vec(W)) = N (0, σ2I) where vec(W) concatenates the columns of the weight matrix into a vector. A first attempt to augment p(vec(W)) with an inducing weight variable U ∈ RMout×Min may be to construct a multivariate Gaussian p(vec(W), vec(U)), such that ∫ p(vec(W), vec(U))dU = N (0, σ2I). This means for the joint covariance matrix of (vec(W), vec(U)), it requires the block corresponding to the covariance of vec(W) to match the prior covariance σ2I . We are then free to parameterise the rest of the entries in the joint covariance matrix, as long as this full matrix remains positive definite. Now the conditional distribution p(W|U) is a function of these parameters, and the conditional sampling from p(W|U) is further discussed in Appendix D.1. Unfortunately, as dim(vec(W)) is typically large (e.g. of the order of 107), using a full covariance Gaussian for p(vec(W), vec(U)) becomes computationally intractable.
We address this issue with matrix normal distributions (Gupta & Nagar, 2018). The prior p(vec(W)) = N (0, σ2I) has an equivalent matrix normal distribution form as p(W) =
MN (0, σ2rI, σ2cI), with σr, σc > 0 the row and column standard deviations satisfying σ = σrσc. Now we introduce the inducing variable U in matrix space, as well as two auxiliary variables Ur ∈ RMout×din , Uc ∈ Rdout×Min , so that the full augmented prior is:(
W Uc Ur U
) ∼ p(W,Uc,Ur,U) :=MN (0,Σr,Σc), (5)
with Lr = ( σrI 0 Zr Dr ) s.t. Σr = LrL>r = ( σ2rI σrZ > r σrZr ZrZ > r +D 2 r ) and Lc = ( σcI 0 Zc Dc ) s.t. Σc = LcL>c = ( σ2cI σcZ > c σcZc ZcZ > c +D 2 c ) .
See Fig. 1(a) for a visualisation of the augmentation. Matrix normal distributions have similar marginalisation and conditioning rules as multivariate Gaussian distributions, for which we provide further examples in Appendix D.2. Therefore the marginalisation constraint Eq. (3) is satisfied for any Zc ∈ RMin×din ,Zr ∈ RMout×dout and diagonal matrices Dc,Dr. For the inducing weight U we have p(U) = MN (0,Ψr,Ψc) with Ψr = ZrZ>r + D2r and Ψc = ZcZ>c + D2c . In the experiments we use whitened inducing weights which transforms U so that p(U) =MN (0, I, I) (Appendix H), but for clarity we continue with the above formulas in the main text.
The matrix normal parameterisation introduces two additional variables Ur,Uc without providing additional expressiveness. Hence it is desirable to integrate them out, leading to a joint multivariate normal with Khatri-Rao product structure for the covariance:
p(vec(W), vec(U)) = N ( 0, ( σ2cI ⊗ σ2rI σcZ>c ⊗ σrZ>r σcZc ⊗ σrZr Ψc ⊗Ψr )) . (6)
As the dominating memory complexity here isO(doutMout+dinMin) which comes from storingZr and Zc, we see that the matrix normal parameterisation of the augmented prior is memory efficient.
Posterior approximation in the joint space We construct a factorised posterior approximation across the layers: q(W1:L,U1:L) = ∏ l q(Wl|Ul)q(Ul). Below we discuss options for q(W|U).
The simplest option is q(W|U) = p(vec(W)| vec(U)) = N (µW|U,ΣW|U), similar to sparse GPs. A slightly more flexible variant adds a rescaling term λ2 to the covariance matrix, which allows efficient KL computation (Appendix E):
q(W|U) = q(vec(W)| vec(U)) = N (µW|U, λ2ΣW|U), (7)
R(λ) := KL [q(W|U)||p(W|U)] = dindout(0.5λ2 − log λ− 0.5). (8)
Plugging θ = W1:L, a = U1:L and Eq. (8) into Eq. (4) returns the following variational lower-bound L(q(W1:L,U1:L)) = Eq(W1:L)[log p(D|W1:L)] − ∑L
l=1 (R(λl) + KL[q(Ul)||p(Ul)]), (9)
with λl the associated scaling parameter for q(Wl|Ul). Again as the choices of Zc,Zr,Dc,Dr do not change the marginal prior p(W), we are safe to optimise them as well. Therefore the variational parameters are now φ = {Zc,Zr,Dc,Dr, λ, dist. params. of q(U)} for each layer.
Two choices of q(U) A simple choice is FFG q(vec(U)) = N (mu, diag(vu)), which performs mean-field inference in U space (c.f. Blundell et al., 2015), and here KL[q(U)||p(U)] has a closedform solution. Another choice is a “mixture of delta measures” q(U) = 1K ∑K k=1 δ(U = U
(k)), i.e. we keep K distinct sets of parameters {U (k)1:L}Kk=1 in inducing space that are projected back into the original parameter space via the shared conditionals q(Wl|Ul) to obtain the weights. This approach can be viewed as constructing “deep ensembles” in U space, and we follow ensemble methods (e.g. Lakshminarayanan et al., 2017) to drop KL[q(U)||p(U)] in Eq. (9). Often U is chosen to have significantly lower dimensions than W, i.e. Min << din and Mout << dout. As q(W|U) and p(W|U) only differ in the covariance scaling constant, U can be regarded as a sparse representation of uncertainty for the network layer, as the major updates in (approximate) posterior belief is quantified by q(U).

heading: 3.2 Efficient sampling with extended Matheron’s rule
text: Computing the variational lower-bound Eq. (9) requires samples from q(W), which requires an efficient sampling procedure for q(W|U). Unfortunately, q(W|U) derived from Eq. (6) & Eq. (7) is not a matrix normal, so direct sampling is prohibitive. To address this challenge, we extend Matheron’s rule (Journel & Huijbregts, 1978; Hoffman & Ribak, 1991; Doucet, 2010) to efficiently sample from q(W|U), with derivations provided in Appendix F. The original Matheron’s rule applies to multivariate Gaussian distributions. As a running example, consider two vector-valued random variablesw, u with joint distribution p(w,u) = N (0,Σ). Then the conditional distribution p(w|u) = N (µw|u,Σw|u) is also Gaussian, and direct sampling from it requires decomposing the conditional covariance matrix Σw|u which can be costly. The main idea of Matheron’s rule is that we can transform a sample from the joint Gaussian to obtain a sample from the conditional distribution p(w|u) as follows:
w = w̄ + ΣwuΣ −1 uu(u− ū), w̄, ū ∼ N (0,Σ), Σ = ( Σww Σwu Σuw Σuu ) . (10)
One can check the validity of Matheron’s rule by computing the mean and variance of w above:
Ew̄,ū[w] = ΣwuΣ−1uuu = µw|u, Vw̄,ū[w] = Σww −ΣwuΣ−1uuΣuw = Σw|u. It might seem counter-intuitive at first sight in that this rule requires samples from a higher dimensional space. However, in the case where decomposition/inversion of Σ and Σuu can be done efficiently, sampling from the joint Gaussian p(w,u) can be significantly cheaper than directly sampling from the conditional Gaussian p(w|u). This happens e.g. when Σ is directly parameterised by its Cholesky decomposition and dim(u) << dim(w), so that sampling w̄, ū ∼ N (0,Σ) is straight-forward, and computing Σ−1uu is significantly cheaper than decomposing Σw|u.
Unfortunately, the original Matheron’s rule cannot be applied directly to sample from q(W|U). This is because q(W|U) = q(vec(W)| vec(U)) differs from p(vec(W)| vec(U)) only in the variance scaling λ, and for p(vec(W)| vec(U)), its joint distribution counter-part Eq. (6) does not have an efficient representation for the covariance matrix. Therefore a naive application of Matheron’s rule requires decomposing the covariance matrix of p(vec(W), vec(U)) which is even more expensive than direct conditional sampling. However, notice that for the joint distribution p(W,Uc,Ur,U) in an even higher dimensional space, the row and column covariance matrices Σr and Σc are parameterised by their Cholesky decompositions, so that sampling from this joint distribution can be done efficiently. This inspire us to extend the original Matheron’s rule for efficient sampling from q(W|U) (details in Appendix F, when λ = 1 it also applies to sampling from p(W|U)):
W = λW̄ + σZ>r Ψ −1 r (U − λŪ)Ψ−1c Zc; W̄, Ū ∼ p(W̄, Ūc, Ūr, Ū) =MN (0,Σr,Σc).
(11) Here W̄, Ū ∼ p(W̄, Ūc, Ūr, Ū) means we first sample W̄, Ūc, Ūr, Ū from the joint then drop Ūc, Ūr; in fact Ūc, Ūr are never computed, and the other samples W̄, Ū can be obtained by:
W̄ = σE1, Ū = ZrE1Z > c + L̂rẼ2Dc +DrẼ3L̂ > c +DrE4Dc,
E1 ∼MN (0, Idout , Idin); Ẽ2, Ẽ3,E4 ∼MN (0, IMout , IMin), (12) L̂r = chol(ZrZ>r ), L̂c = chol(ZcZ > c ).
The run-time cost isO(2M3out+2M3in+doutMoutMin+Mindoutdin) required by inverting Ψr,Ψc, computing L̂r, L̂c, and the matrix products. The extended Matheron’s rule is visualised in Fig. 1
with a comparison to the original Matheron’s rule for sampling from q(vec(W)| vec(U)). Note that the original rule requires joint sampling from Eq. (6) (i.e. sampling the white blocks in Fig. 1(b)) which has O((doutdin +MoutMin)3) cost. Therefore our recipe avoids inverting and multiplying big matrices, resulting in a significant speed-up for conditional sampling.

heading: 3.3 Computational complexities
text: In Table 1 we report the complexity figures for two types of inducing weight approaches: FFG q(U) (FFG-U) and Delta mixture q(U) (Ensemble-U). Baseline approaches include: DeterministicW, variational inference with FFG q(W) (FFG-W, Blundell et al., 2015), deep ensemble in W (Ensemble-W, Lakshminarayanan et al., 2017), as well as parameter efficient approaches such as matrix-normal q(W) (Matrix-normal-W, Louizos & Welling (2017)), variational inference with k-tied FFG q(W) (k-tied FFG-W, Świątkowski et al. (2020)), and rank-1 BNN (Dusenberry et al., 2020). The gain in memory is significant for the inducing weight approaches, in fact with Min < din and Mout < dout the parameter storage requirement is smaller than a single deterministic neural network. The major overhead in run-time comes from the extended Matheron’s rule for sampling q(W|U). Some of the computations there are performed only once, and in our experiments we show that by using a relatively low-dimensional U and large batch-sizes, the overhead is acceptable.

heading: 4 Experiments
text: We evaluate the inducing weight approaches on regression, classification and related uncertainty estimation tasks. The goal is to demonstrate competitive performance to popular W-space uncertainty estimation methods while using significantly fewer parameters. We acknowledge that existing parameter efficient approaches for uncertainty estimation (e.g. k-tied or rank-1 BNNs) have achieved close performance to deep ensembles. However, none of them reduces the parameter count to be smaller than that of a single network. Therefore we decide not to include these baselines and instead focus on comparing: (1) variational inference with FFG q(W) (FFG-W, Blundell et al., 2015) v.s. FFG q(U) (FFG-U, ours); (2) deep ensemble in W space (Ensemble-W, Lakshminarayanan et al., 2017) v.s. ensemble in U space (Ensemble-U, ours). Another baseline is training a deterministic neural network with maximum likelihood. Details and additional results are in Appendices J and K.

heading: 4.1 Synthetic 1-D regression
text: The regression task follows Foong et al. (2019), which has two input clusters x1 ∼ U [−1,−0.7], x2 ∼ U [0.5, 1], and targets y ∼ N (cos(4x+ 0.8), 0.01). For reference we show the exact posterior results using the NUTS sampler (Hoffman & Gelman, 2014). The results are visualised in Fig. 2 with predictive mean in blue, and up to three standard deviations as shaded area. Similar to historical results, FFG-W fails to represent the increased uncertainty away from the data and in between clusters. While underestimating predictive uncertainty overall, FFG-U shows a small increase in predictive uncertainty away from the data. In contrast, a per-layer Full-covariance Gaussian (FCG) in both weight (FCG-W) and inducing space (FCG-U) as well as Ensemble-U better capture the increased predictive variance, although the mean function is more similar to that of FFG-W.
Table 2: CIFAR in-distribution metrics (in %).
CIFAR10 CIFAR100 Method Acc. ↑ ECE ↓ Acc. ↑ ECE ↓ Deterministic 94.72 4.46 75.73 19.69 Ensemble-W 95.90 1.08 79.33 6.51 FFG-W 94.13 0.50 74.44 4.24 FFG-U 94.40 0.64 75.37 2.29 Ensemble-U 94.94 0.45 75.97 1.12
12 4 8 16 32 Number of samples
0
5
10
re la
tiv e
ru n-
tim e
Speed FFG-U (M=128, R18) FFG-U (M=128, R50) FFG-U (M=64, R18) FFG-U (M=64, R50)
FFG-W (R18) FFG-W (R50)
3264 128 256 M
0
0.5
1
2
re la
tiv e
pa ra
m . c
ou nt
deterministic
FFG-W
acc=94.4% 94.7%
FFG-U
Model size (R50)
Figure 3: Resnet run-times & model sizes.

heading: 4.2 Classification and in-distribution calibration
text: As the core empirical evaluation, we train Resnet-50 models (He et al., 2016b) on CIFAR-10 and CIFAR-100 (Krizhevsky et al., 2009). To avoid underfitting issues with FFG-W, a useful trick is to set an upper limit σ2max on the variance of q(W) (Louizos & Welling, 2017). This trick is similarly applied to the U-space methods, where we cap λ ≤ λmax for q(W|U), and for FFG-U we also set σ2max for the variance of q(U). In convolution layers, we treat the 4D weight tensor W of shape (cout, cin, h, w) as a cout × cinhw matrix. We use U matrices of shape 64× 64 for all layers (i.e. M = Min = Mout = 64), except that for CIFAR-10 we set Mout = 10 for the last layer.
In Table 2 we report test accuracy and test expected calibration error (ECE) (Guo et al., 2017) as a first evaluation of the uncertainty estimates. Overall, Ensemble-W achieves the highest accuracy, but is not as well-calibrated as variational methods. For the inducing weight approaches, Ensemble-U outperforms FFG-U on both datasets; overall it performs the best on the more challenging CIFAR-100 dataset (close-to-Ensemble-W accuracy and lowest ECE). Tables 5 and 6 in Appendix K show that increasing the U dimensions to M = 128 improves accuracy but leads to slightly worse calibration.
In Fig. 3 we show prediction run-times for batch-size = 500 on an NVIDIA Tesla V100 GPU, relative to those of an ensemble of deterministic nets, as well as relative parameter sizes to a single ResNet-50. The extra run-times for the inducing methods come from computing the extended Matheron’s rule. However, as they can be calculated once and cached for drawing multiple samples, the overhead reduces to a small factor when using larger number of samples K, especially for the bigger Resnet-50. More importantly, when compared to a deterministic ResNet-50, the inducing weight models reduce the parameter count by over 75% (5, 710, 902 vs. 23, 520, 842) for M = 64.
Hyper-parameter choices We visualise in Fig. 4 the accuracy and ECE results for computationally lighter inducing weight ResNet-18 models with different hyper-parameters (see Appendix J).
Performance in both metrics improves as the U matrix size M is increased (right-most panels), and the results for M = 64 and M = 128 are fairly similar. Also setting proper values for λmax, σmax is key to the improved results. The left-most panels show that with fixed σmax values (or Ensemble-U), the preferred conditional variance cap values λmax are fairly small (but still larger than 0 which corresponds to a point estimate for W given U). For σmax which controls variance in U space, we see from the top middle panel that the accuracy metric is fairly robust to σmax as long as λmax is not too large. But for ECE, a careful selection of σmax is required (bottom middle panel).

heading: 4.3 Robustness, out-of-distribution detection and pruning
text: To investigate the models’ robustness to distribution shift, we compute predictions on corrupted CIFAR datasets (Hendrycks & Dietterich, 2019) after training on clean data. Fig. 5 shows accuracy and ECE results for the ResNet-50 models. Ensemble-W is the most accurate model across skew intensities, while FFG-W, though performing well on clean data, returns the worst accuracy under perturbation. The inducing weight methods perform competitively to Ensemble-W with Ensemble-U being slightly more accurate than FFG-U as on the clean data. For ECE, FFG-U outperforms Ensemble-U and Ensemble-W, which are similarly calibrated. Interestingly, while the accuracy of FFG-W decays quickly as the data is perturbed more strongly, its ECE remains roughly constant.
Table 3 further presents the utility of the maximum predicted probability for out-of-distribution (OOD) detection. The metrics are the area under the receiver operator characteristic (AUROC) and the precision-recall curve (AUPR). The inducing-weight methods perform similarly to Ensemble-W; all three outperform FFG-W and deterministic networks across the board.
Parameter pruning We further investigate pruning as a pragmatic alternative for more parameter-efficient inference. For FFG-U, we prune entries of the Z matrices, which contribute the largest number of parameters to the inducing methods, with the smallest magnitude. For FFG-W we follow Graves (2011) in setting different fractions of W to 0 depending on their variational mean-to-variance ratio and repeat the previous experiments after fine-tuning the distributions on the remaining variables. We stress that, unlike FFG-U, the FFG-W pruning corresponds to a post-hoc change of the probabilistic model and no longer performs inference in the original weight-space.
For FFG-W, pruning 90% of the parameters (leaving 20% of parameters as compared to its deterministic counterpart) worsens the ECE, in particular on CIFAR100, see Fig. 6. Further pruning to 1%
worsens the accuracy and the OOD detection results as well. On the other hand, pruning 50% of the Z matrices for FFG-U reduces the parameter count to 13.2% of a deterministic net, at the cost of only slightly worse calibration. See Appendix K for the full results.

heading: 5 Discussions

heading: 5.1 A function-space perspective on inducing weights
text: Although the inducing weight approach performs approximate inference in weight space, we present in Appendix G a function-space inference perspective of the proposed method, showing its deep connections to sparse GPs. Our analysis considers the function-space behaviour of each network layer’s output and discusses the corresponding interpretations of the U variables and Z parameters.
The interpretations are visualised in Fig. 7. Similar to sparse GPs, in each layer, the Zc parameters can be viewed as the (transposed) inducing input locations which lie in the same space as the layer’s input. The Uc variables can also be viewed as the corresponding (noisy) inducing outputs that lie in the pre-activation space. Given that the output dimension dout can still be high (e.g. > 1000 in a fully connected layer), our approach performs further dimension reduction in a similar spirit as probabilistic PCA (Tipping & Bishop, 1999), which projects the column vectors of Uc to a lower-dimensional space. This returns the inducing weight variables U, and the projection parameters are {Zr,Dr}. Combining the two steps, it means the column vectors of U can be viewed as collecting the “noisy projected inducing outputs” whose corresponding “inducing inputs” are row vectors of Zc (see the red bars in Fig. 7).
In Appendix G we further derive the resulting variational objective from the function-space view, which is almost identical to Eq. (9), except for scaling coefficients on the R(λl) terms to account for the change in dimensionality from weight space to function space. This result nicely connects posterior inference in weight- and function-space.

heading: 5.2 Related work
text: Parameter-efficient uncertainty quantification methods Recent research has proposed Gaussian posterior approximations for BNNs with efficient covariance structure (Ritter et al., 2018; Zhang et al., 2018b; Mishkin et al., 2018). The inducing weight approach differs from these in introducing structure via a hierarchical posterior with low-dimensional auxiliary variables. Another line of work reduces the memory overhead via efficient parameter sharing (Louizos & Welling, 2017; Wen et al., 2020; Świątkowski et al., 2020; Dusenberry et al., 2020). The third category of work considers a hybrid approach, where only a selective part of the neural network receives Bayesian treatments, and the other weights remain deterministic (Bradshaw et al., 2017; Daxberger et al., 2021). However, both types of approaches maintain a “mean parameter” for the weights, making the memory footprint at least that of storing a deterministic neural network. Instead, our approach shares parameters via the augmented prior with efficient low-rank structure, reducing the memory use compared to a deterministic network. In a similar spirit, Izmailov et al. (2019) perform inference in a d-dimensional sub-space obtained from PCA on weights collected from an SGD trajectory. But this approach does not leverage the layer-structure of neural networks and requires d× memory of a single network.
Network pruning in uncertainty estimation context There is a large amount of existing research advocating network pruning approaches for parameter-efficient deep learning, e.g. see Han et al. (2016); Frankle & Carbin (2018); Lee et al. (2019). In this regard, mean-field VI approaches have also shown success in network pruning, but only in terms of maintaining a minimum accuracy level (Graves, 2011; Louizos et al., 2017; Havasi et al., 2019). To the best of our knowledge, our empirical study presents the first evaluation for VI-based pruning methods in maintaining uncertainty estimation quality. Deng et al. (2019) considers pruning BNNs with stochastic gradient Langevin dynamics (Welling & Teh, 2011) as the inference engine. The inducing weight approach is orthogonal to these BNN pruning approaches, as it leaves the prior on the network parameters intact, while the pruning
approaches correspond to a post-hoc change of the probabilistic model to using a sparse weight prior. Indeed our parameter pruning experiments showed that our approach can be combined with network pruning to achieve further parameter efficiency improvements.
Sparse GP and function-space inference As BNNs and GPs are closely related (Neal, 1995; Matthews et al., 2018; Lee et al., 2018), recent efforts have introduced GP-inspired techniques to BNNs (Ma et al., 2019; Sun et al., 2019; Khan et al., 2019; Ober & Aitchison, 2021). Compared to weight-space inference, function-space inference is appealing as its uncertainty is more directly relevant for predictive uncertainty estimation. While the inducing weight approach performs computations in weight-space, Section 5.1 establishes the connection to function-space posteriors. Our approach is related to sparse deep GP methods with Uc having similar interpretations as inducing outputs in e.g. Salimbeni & Deisenroth (2017). The major difference is that U lies in a low-dimensional space, projected from the pre-activation output space of a network layer.
The original Matheron’s rule (Journel & Huijbregts, 1978; Hoffman & Ribak, 1991; Doucet, 2010) for sampling from conditional multivariate Gaussian distributions has recently been applied to speed-up sparse GP inference (Wilson et al., 2020, 2021). As explained in Section 3.2, direct application of the original rule to sampling W conditioned on U still incurs prohibitive cost as p(vec(W), vec(U)) does not have a convenient factorisation form. Our extended Matheron’s rule addresses this issue by exploiting the efficient factorisation structure of the joint matrix normal distribution p(W,Uc,Ur,U), reducing the dominating factor of computation cost from cubic (O(d3outd3in)) to linear (O(doutdin)). We expect this new rule to be useful for a wide range of models/applications beyond BNNs, such as matrix-variate Gaussian processes (Stegle et al., 2011).
Priors on neural network weights Hierarchical priors for weights has also been explored (Louizos et al., 2017; Krueger et al., 2017; Atanov et al., 2019; Ghosh et al., 2019; Karaletsos & Bui, 2020). However, we emphasise that p̃(W,U) is a pseudo prior that is constructed to assist posterior inference rather than to improve model design. Indeed, parameters associated with the inducing weights are optimisable for improving posterior approximations. Our approach can be adapted to other priors, e.g. for a Horseshoe prior p(θ, ν) = p(θ|ν)p(ν) = N (θ; 0, ν2)C+(ν; 0, 1), the pseudo prior can be defined as p̃(θ, ν, a) = p̃(θ|ν, a)p̃(a)p(ν) such that ∫ p̃(θ|ν, a)p̃(a)da = p(θ|ν). In general, pseudo priors have found broader success in Bayesian computation (Carlin & Chib, 1995).

heading: 6 Conclusion
text: We have proposed a parameter-efficient uncertainty quantification framework for neural networks. It augments each of the network layer weights with a small matrix of inducing weights, and by extending Matheron’s rule to matrix-normal related distributions, maintains a relatively small run-time overhead compared with ensemble methods. Critically, experiments on prediction and uncertainty estimation tasks show the competence of the inducing weight methods to the state-of-the-art, while reducing the parameter count to under a quarter of a deterministic ResNet-50 before pruning. This represents a significant improvement over prior Bayesian and deep ensemble techniques, which so far have not managed to go below this threshold despite various attempts of matching it closely.
Several directions are to be explored in the future. First, modelling correlations across layers might further improve the inference quality. We outline an initial approach leveraging inducing variables in Appendix H. Second, based on the function-space interpretation of inducing weights, better initialisation techniques can be inspired from the sparse GP and dimension reduction literature. Similarly, this interpretation might suggest other innovative pruning approaches for the inducing weight method, thereby achieving further memory savings. Lastly, the run-time overhead of our approach can be mitigated by a better design of the inducing weight structure as well as vectorisation techniques amenable to parallelised computation. Designing hardware-specific implementations of the inducing weight approach is also a viable alternative for such purposes.

Review summaries:
This work proposes a specific parametrisation for the Gaussian prior and approximate posterior distribution in variational Bayesian neural networks in terms of inducing weights. The general idea is an instance of the sparse variational inference scheme for GPs proposed by Titsias back in 2009; for a given model with a prior p(W) perform variational inference on an extended model with a hierarchical prior p(U) p(W | U), that has the same marginal p(W) = \int p(U)p(W | U)dU as the original model. The authors then consider “U” to be auxiliary weights that are jointly Gaussian with the actual weights “W” and then use the decomposition p(W|U)p(U), q(W|U)q(U) for the prior and approximate posterior (which can easily be computed via the conditional Gaussian rules). Furthermore, they “tie” (almost) all of the parameters between q(W|U) and p(W|U) (similarly to Titsias, 2009). The main benefit from these two things is that since the mean and covariance of the Gaussian distribution over W conditioned on U can be efficiently represented as functions of U, whenever dim(U) << dim(W) we get reductions in memory for storing the distributions over the parameters in the network. The authors furthermore, discuss how to efficiently parametrize the joint distribution over W, U, discuss different choices for q(U) (that can lead to either traditional VI or something like deep ensembles). In addition, they also discuss how more efficient sampling from q(W|U) can be realised via an extension of the Matheron’s rule to the case of matrix random variables. Finally, they evaluate their method against traditional mean field variational Bayesian neural networks and deep ensembles on several tasks that include regression, classification, calibration and OOD performance. 

This paper proposes a method for reducing the storage cost of BNNs. The idea is to apply the inducing point method and fast Matheron's sampling rule, commonly used in GPs, to standard variational BNNs. Experiments show that the resulting BNNs stay competitive to standard BNNs while enjoying up to 75% reduction in parameter size.

The paper proposes to represent uncertainty in neural networks by augmenting neural network layers with inducing weights and then assuming that the resulting vector of weights is drawn from a matrix normal distribution which is structured so that it has fewer parameters than the original network. The parameters of the distribution are learned via variational inference and the authors present a procedure to sample efficiently from this distribution for optimizing the model parameters and for inference at test time. Experiments on classification and out of distribution detection are presented to validate the approach.

The paper proposes a parameter-efficient method for uncertanity quantification in deep neural network using per layer inducing weights. The paper develops a network weight sampling scheme based on inducing weights. The proposed method employs Matheron's rule for efficient sampling. The method requires fewer parameters than even deterministic networks, while keeping good calibration and performance, at the cost of longer compute times, when few MC weight samples are used. 

This paper proposes a layer-wise variational inference scheme for Bayesian neural networks that uses "inducing weights" for each layer before up-sampling to the full (layer-wise) weight distribution. The scheme is equivalent to a hierarchical prior formulation with an analogous hierarchical variational posterior. Experiments are performed on the standard BDL benchmarks of synthetic one dimensional regression, ResNets on CIFAR-100, and out of distribution detection.