Activation-level uncertainty in deep neural networks

1 INTRODUCTION . Deep Neural Networks ( DNNs ) have achieved state-of-the-art performance in many different tasks , such as speech recognition ( Hinton et al. , 2012 ) , natural language processing ( Mikolov et al. , 2013 ) or computer vision ( Krizhevsky et al. , 2012 ) . In spite of their predictive power , DNNs are limited in terms of uncertainty estimation . This has been a classical concern in the field ( MacKay , 1992 ; Hinton & Van Camp , 1993 ; Barber & Bishop , 1998 ) , which has attracted a lot of attention in the last years ( Lakshminarayanan et al. , 2017 ; Guo et al. , 2017 ; Sun et al. , 2019 ; Wenzel et al. , 2020 ) . Indeed , this ability to “ know what is not known ” is essential for critical applications such as medical diagnosis ( Esteva et al. , 2017 ; Mobiny et al. , 2019 ) or autonomous driving ( Kendall & Gal , 2017 ; Gal , 2016 ) . Bayesian Neural Networks ( BNNs ) address this problem through a Bayesian treatment of the network weights1 ( MacKay , 1992 ; Neal , 1995 ) . This will be refered to as weight-space stochasticity . However , dealing with uncertainty in weight space is challenging , since it contains many symmetries and is highly dimensional ( Wenzel et al. , 2020 ; Sun et al. , 2019 ; Snoek et al. , 2019 ; Fort et al. , 2019 ) . Here we focus on two specific limitations . First , it has been recently shown that BNNs with well-established inference methods such as Bayes by Backprop ( BBP ) ( Blundell et al. , 2015 ) and MC-Dropout ( Gal & Ghahramani , 2016 ) underestimate the predictive uncertainty for instances located in-between two clusters of training points ( Foong et al. , 2020 ; 2019 ; Yao et al. , 2019 ) . Second , the weight-space prior does not allow BNNs to guide extrapolation to out-of-distribution ( OOD ) data ( Sun et al. , 2019 ; Nguyen et al. , 2015 ; Ren et al. , 2019 ) . Both aspects are illustrated graphically in Figure 3 , more details in Section 3.1 . ∗Work developed mostly while visiting Cambridge University , UK . 1The bias term will be absorbed within the weights throughout the work . As an alternative to standard BNNs , Functional Bayesian Neural Nets ( fBNN ) specify the prior and perform inference directly in function space ( Sun et al. , 2019 ) . This provides a mechanism to guide the extrapolation in OOD data , e.g . predictions can be encouraged to revert to the prior in regions of no observed data . However , the posterior stochastic process is still defined by a factorized Gaussian on the network weights ( i.e . as in BBP ) , see ( Sun et al. , 2019 , Sect . 3.1 ) . We will show that this makes fBNN inherit the problem of underestimating the predictive uncertainty for in-between data . In this work , we adopt a different approach by moving stochasticity from the weights to the activation function , see Figure 1 . This will be referred to as auNN ( activation-level uncertainty for Neural Networks ) . The activation functions are modelled with ( one-dimensional ) GP priors , for which a triangular kernel inspired by the ReLu non-linearity ( Nair & Hinton , 2010 ; Glorot et al. , 2011 ) is used . Since non-linearities are typically simple functions ( e.g . ReLu , sigmoid , tanh ) , our GPs are sparsified with few inducing points . The network weights are deterministic parameters which are estimated to maximize the marginal likelihood of the model . The motivation behind auNN is to avoid inference in the complex space of weights . We hypothesise that it could be enough to introduce stochasticity in the activation functions that follow the linear projections to provide sensible uncertainty estimations . We show that auNN obtains well-calibrated estimations for in-between data , and its prior allows to guide the extrapolation to OOD data by reverting to the empirical mean . This will be visualized in a simple 1D example ( Figure 3 and Table 1 ) . Moreover , auNN obtains competitive performance in standard benchmarks , is scalable ( datasets of up to ten millions training points are used ) , and can be readily used for classification . The use of GPs for the activations establishes an interesting connection with deep GPs ( DGPs ) ( Damianou & Lawrence , 2013 ; Salimbeni & Deisenroth , 2017 ) . The main difference is the linear projection before the GP , recall Figure 1 ( c-d ) . This allows auNN units to model simpler mappings between layers , which are defined along one direction of the input space , similarly to neural networks . However , DGP units model more complex mappings defined on the whole input space , see also Figure 2a . We will show that auNN units require fewer inducing points and are better suited for deep architectures , achieving superior performance . Also , a thorough discussion on additional related work will be provided in Section 4 . In summary , the main contributions of this paper are : ( 1 ) a new approach to model uncertainty in DNNs , based on deterministic weights and simple stochastic non-linearities ( in principle , not necessarily modelled by GPs ) ; ( 2 ) the specific use of non-parametric GPs as a prior , including the triangular kernel inspired by the ReLu ; ( 3 ) auNN addresses a well-known limitation of BNNs and fBNNs ( uncertainty underestimation for in-between data ) , can guide the extrapolation to OOD data by reverting to the empirical mean , and is competitive in standard prediction tasks ; ( 4 ) auNN units require fewer inducing points and are better suited for deep architectures than DGP ones , achieving superior performance . 2 PROBABILISTIC MODEL AND INFERENCE . Model specification . We focus on a supervised task ( e.g . regression or classification ) with training data2 { xn , : , yn , : } Nn=1 . The graphical model in Figure 2b will be useful throughout this section . We 2The output is represented as a vector since all the derivations apply for the multi-output case . assume a model of L layers , each one with Dl units as in Figure 1c . Each activation is modelled with a ( 1D ) GP prior , i.e . f ld ( a l d ) ∼ GP ( µld , kld ) , with µld : R → R and kld : R × R → R. The GP hyperparameters θld will be omitted for clarity ( for the kernels used here , θ l d includes the amplitude and the lengthscale ) . Assuming independence between units , each layer depends on the previous one as : p ( Fl|Fl−1 , Wl ) = p ( Fl|Al ) = ∏Dl d=1 p ( f l d|ald ) , ( 1 ) where Fl is the N ×Dl matrix of outputs of the l-th layer for N inputs , Wl is the Dl−1 ×Dl matrix of weights in that layer , and Al is the N ×Dl matrix of pre-activations , i.e . Al = Fl−1 ·Wl . As usual , the columns and rows of Fl are denoted as f ld and f l n , : , respectively ( and analogously for the other matrices ) . Since the activation is defined by a GP , we have p ( f ld|ald ) = N ( f ld|µld , Kld ) , with µld ( resp . Kld ) the result of evaluating µ l d ( resp . k l d ) on a l d ( that is , µ l d is a N -dimensional vector and K l d is a N ×N matrix ) . To fully specify the model , the output Y is defined from the last layer with a distribution that factorizes across data points , i.e . p ( Y|FL ) = ∏N n=1 p ( yn , :|fLn , : ) . This formulation resembles that of DGPs ( Damianou & Lawrence , 2013 ; Salimbeni & Deisenroth , 2017 ) . The main difference is that we model Fl|Fl−1 through Dl 1D GPs evaluated on the pre-activations Al ( i.e . the projections of Fl−1 through Wl ) , whereas DGPs use Dl GPs of dimension Dl−1 evaluated directly on Fl−1 , recall Figure 1 ( c-d ) . Variational Inference . Inference in the proposed model is intractable . To address this , we follow standard sparse variational GP approaches ( Titsias , 2009 ; Hensman et al. , 2013 ; 2015 ) , similarly to the Doubly Stochastic Variational Inference ( DSVI ) for DGPs ( Salimbeni & Deisenroth , 2017 ) . Specifically , in each unit of each layer we introduce M l inducing values uld , which are the result of evaluating the GP on the one-dimensional inducing points zld . We naturally write U l and Zl for the corresponding M l × Dl matrices associated to the l-th layer , respectively . Following eq . ( 1 ) , the augmented model for one layer is p ( Fl , Ul|Fl−1 , Wl , Zl ) = p ( Fl|Ul , Al , Zl ) p ( Ul|Zl ) = ∏Dl d=1 p ( f l d|uld , ald , zld ) p ( uld|zld ) . ( 2 ) Variational inference ( VI ) involves the approximation of the true posterior p ( { Fl , Ul } l|Y ) . Following ( Hensman et al. , 2013 ; Salimbeni & Deisenroth , 2017 ) , we propose a posterior given by p ( F|U ) and a parametric Gaussian on U : q ( { Fl , Ul } l ) = ∏L l=1 p ( F l|Ul , Al , Zl ) q ( Ul ) = ∏L l=1 ∏Dl d=1 p ( f l d|uld , ald , zld ) q ( uld ) , ( 3 ) where q ( uld ) = N ( uld|mld , Sld ) , with mld ∈ RM l and Sld ∈ RM l×M l variational parameters to be estimated . Minimizing the KL divergence between q ( { Fl , Ul } l ) and the true posterior is equivalent to maximizing the following evidence lower bound ( ELBO ) : log p ( Y| { Wl , Zl } l ) ≥ ELBO = N∑ n=1 Eq ( fLn , : ) [ log p ( yn , :|fLn , : ) ] − L∑ l=1 Dl∑ d=1 KL ( q ( uld ) ||p ( uld ) ) . ( 4 ) In the ELBO , the KL term can be computed in closed-form , as both q ( uld ) and p ( u l d ) are Gaussians . The log likelihood term can be approximated by sampling from the marginal posterior q ( fLn , : ) , which can be done efficiently through univariate Gaussians as in ( Salimbeni & Deisenroth , 2017 ) . Specifically , Ul can be analytically marginalized in eq . ( 3 ) , which yields q ( { Fl } l ) = ∏ l q ( F l|Fl−1 , Wl ) = ∏ l , dN ( f ld|µ̃ l d , Σ̃ l d ) , with : [ µ̃ld ] i = µ l d ( a l id ) +α l d ( a l id ) ᵀ ( mld − µld ( zld ) ) , ( 5 ) [ Σ̃ l d ] ij = k l d ( a l id , a l jd ) −αld ( alid ) ᵀ ( kld ( zld ) − Sld ) αld ( aljd ) , ( 6 ) where αld ( x ) = k l d ( x , z l d ) [ k l d ( z l d ) ] −1 and aln , : = W lf l−1n , : . Importantly , the marginal posterior q ( f l n , : ) is a Gaussian that depends only on aln , : , which in turn only depends on q ( f l−1 n , : ) . Therefore , sampling from f ln , : is straightforward using the reparametrization trick ( Kingma & Welling , 2013 ) : f lnd = [ µ̃ l d ] n + ε · [ Σ̃ l d ] 1/2 nn , with ε ∼ N ( 0 , 1 ) , and f0n , : = xn , : . ( 7 ) Training consists in maximizing the ELBO , eq . ( 4 ) , w.r.t . variational parameters { mld , Sld } , inducing points { zld } , and model parameters ( i.e . weights { wld } and kernel parameters { θ l d } ) . This can be done in batches , allowing for scalability to very large datasets . The complexity to evaluate the ELBO is O ( NM2 ( D1 + · · ·+DL ) ) , the same as DGPs with DSVI ( Salimbeni & Deisenroth , 2017 ) .3 Predictions . Given a new x∗ , : , we want to compute4 p ( fL∗ , :|X , Y ) ≈ Eq ( { Ul } ) [ p ( fL∗ , :| { Ul } ) ] . As in ( Salimbeni & Deisenroth , 2017 ) , this can be approximated by sampling S values up to the ( L− 1 ) -th layer with the same eq . ( 7 ) , but starting with x∗ , : . Then , p ( fL∗ , :|X , Y ) is given by the mixture of the S Gaussians distributions obtained from eqs . ( 5 ) - ( 6 ) . Triangular kernel . One of the most popular kernels in GPs is the RBF ( Williams & Rasmussen , 2006 ) , which produces very smooth functions . However , the ReLu non-linearity led to a general boost in performance in DNNs ( Nair & Hinton , 2010 ; Glorot et al. , 2011 ) , and we aim to model similar activations . Therefore , we introduce the use of the triangular ( TRI ) kernel . Just like RBF , TRI is an isotropic kernel , i.e . it depends on the distance between the inputs , k ( x , y ) = γ · g ( |x− y|/ ` ) , with γ and ` the amplitude and lengthscale . For RBF , g ( t ) = e−t 2/2 . For TRI , g ( t ) = max ( 1− t , 0 ) . This is a valid kernel ( Williams & Rasmussen , 2006 , Section 4.2.1 ) . Similarly to the ReLu , the functions modelled by TRI are piecewise linear , see Figure 6a in the main text and Figure 8 in Appendix C. Comparison with DGP . The difference between auNN and DGP units is graphically illustrated in Figure 2a . Whereas DGP mappings from one layer to the next are complex functions defined on Dl−1 dimensions ( Dl−1 = 2 in the figure ) , auNN mappings are defined just along one direction via the weight projection . This is closer in spirit to NNs , whose mappings are also simpler and better suited for feature extraction and learning more abstract concepts . Moreover , since the GP is defined on a 1D space , auNN requires fewer inducing points than DGP ( which , intuitively , can be interpreted as inducing ( hyper ) planes in the Dl−1-dimensional space before the projection ) .

Summary: Either putting the uncertainty on the weights (e.g., Bayes by BP) or on the activation (e.g., fast dropout or variants of natural-parameter networks [2,3] or Bayesian dark knowledge [4]) or both [1] have been investigated before. The idea of moving the uncertainty from the weight to the activation function is not new. One could argue that VAE-style parameterization or local reparameterization trick is also a kind of methods that put uncertainty in the activation function. In fact the proposed method does involve the reprarameterization trick in each layer as shown in Eq. 7.