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