Sparse Uncertainty Representation in Deep Learning with Inducing Weights

1 Introduction . 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 2 Inducing variables for variational inference . 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 .

Summary: 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.