Fixing Asymptotic Uncertainty of Bayesian Neural Networks with Infinite ReLU Features

1 INTRODUCTION . Calibrated uncertainty is crucial for safety-critical decision making by neural networks ( NNs ) ( Amodei et al. , 2016 ) . Standard training methods of NNs yield point estimates that , even if they are highly accurate , can still be severely overconfident ( Guo et al. , 2017 ) . Approximate Bayesian methods , which turn NNs into Bayesian neural networks ( BNNs ) , can be used to address this issue . Kristiadi et al . ( 2020 ) recently showed that for binary ReLU classification networks , far away from the training data ( more precisely : when scaling any input x with a scalar α > 0 and taking the limit α → ∞ ) , the uncertainty of BNNs can be bounded away from zero . This is an encouraging result when put in contrast to the standard point-estimated networks , for which Hein et al . ( 2019 ) showed earlier that the same asymptotic limit always yields arbitrarily high ( over- ) confidence . Nevertheless , BNNs can still be asymptotically overconfident ( albeit less so than the standard NNs ) since the aforementioned uncertainty bound can be loose . This issue is our principal interest in this paper . An intuitive interpretation is that ReLU NNs “ miss out on some uncertainty ” even in their Bayesian formulation , because they fit a finite number of ReLU features to the training data , by “ moving around ” these features within the coverage of the data . This process has no means to encode a desideratum that the model should be increasingly uncertain away from the data . In this work , we “ add in ” additional uncertainty by considering an infinite number of additional ReLU features spaced at regular intervals away from the data in the input and hidden spaces of the network . Since these features have negligible values in the data region , they do not contribute to the training process . Hence , we can consider a prior for their weights , chosen to be an independent Gaussian , and arrive at a specific Gaussian process ( GP ) which covariance function is a generalization of the classic cubic-spline kernel ( Wahba , 1990 ) . This GP prior can be added to any pre-trained ReLU BNN as a simple augmentation to its output . Considering the additive combination of a parametric BNN and GP prior together , we arrive at another view of the method : It approximates the “ full GP posterior ” that models the residual of a point-estimated NN ( Blight & Ott , 1975 ; Qiu et al. , 2020 ) . In our factorization , the BNN models uncertainty around the training data , while the GP prior models uncertainty far away from them . By factorizing these two parts from each other , our formulation requires no ( costly ) GP posterior inference , and thus offers lightweight , modular uncertainty calibration . See Fig . 1 for illustration . Theoretical analysis is a core contribution of this work . We show that the proposed method ( i ) preserves the predictive performance of the base ReLU BNN . Furthermore , it ( ii ) ensures that the surrounding output variance asymptotically grows cubically in the distance to the training data , and thus ( iii ) yields uniform asymptotic confidence in the multi-class classification setting . These results extend those of Kristiadi et al . ( 2020 ) in so far as their analysis is limited to the binary classification case and their bound can be loose . Furthermore , our approach is complementary to the method of Meinke & Hein ( 2020 ) which attains maximum uncertainty far from the data for non-Bayesian point-estimate NNs . Finally , our empirical evaluation confirms our analysis and shows that the proposed method also improves uncertainty estimates in the non-asymptotic regime . 2 BACKGROUND . 2.1 BAYESIAN NEURAL NETWORKS . Let f : RN ×RD → RC defined by ( x , θ ) 7→ f ( x ; θ ) = : fθ ( x ) be a neural network . Here , θ is the collection of all parameters of f . Given an i.i.d . dataset D : = ( xm , ym ) Mm=1 , the standard training procedure amounts to finding a point estimate θ∗ of the parameters θ , which can be identified in the Bayesian framework with maximum a posteriori ( MAP ) estimation1 θ∗ = arg max θ log p ( θ | D ) = arg max θ M∑ m=1 log p ( ym | fθ ( xm ) ) + log p ( θ ) . While this point estimate may yield highly accurate predictions , it does not encode uncertainty over θ , causing an overconfidence problem ( Hein et al. , 2019 ) . Bayesian methods can mitigate this issue , specifically , by treating the parameter of f as a random variable and applying Bayes rule . The resulting network is called a Bayesian neural network ( BNN ) . The common way to approximate the posterior p ( θ | D ) of a BNN is by a Gaussian q ( θ | D ) = N ( θ | µ , Σ ) , which can be constructed for example by a Laplace approximation ( MacKay , 1992b ) or variational Bayes ( Hinton & Van Camp , 1993 ) . Given such an approximate posterior q ( θ | D ) and a test point x∗ ∈ RN , one then needs to marginalize the parameters to make predictions , i.e . we compute the integral y∗ = ∫ h ( f ( x∗ ; θ ) ) q ( θ | D ) dθ , where h is an inverse link function , such as the identity function for regression or the logistic-sigmoid and softmax functions for binary and multi-class classifications , respectively . Since the network f is a non-linear function of θ , this integral does not have an analytic solution . However , one can obtain a useful approximation via the following network linearization : Let x∗ ∈ RN be a test point and q ( θ | D ) = N ( θ | µ , Σ ) be a Gaussian approximate posterior . Linearizing f around µ yields the following marginal distribution over the function output f ( x∗ ) :2 p ( f ( x∗ ) | x∗ , D ) ≈ N ( f ( x∗ ) | f ( x∗ ; µ ) ︸ ︷︷ ︸ = : m∗ , J > ∗ ΣJ∗︸ ︷︷ ︸ = : V∗ ) , ( 1 ) where J∗ is the Jacobian of f ( x∗ ; θ ) w.r.t . θ at µ . ( In the case of a real-valued network f , we use the gradient g∗ : = ∇θf ( x∗ ; θ ) |µ instead of J∗ . ) This distribution can then be used as the predictive distribution p ( y∗ | x∗ , D ) in the regression case . For classifications , we need another approximation since h is not the identity function . One such approximation is the generalized probit approximation 1In the statistical learning view , log p ( ym | fθ ( xm ) ) is identified with the empirical risk , log p ( θ ) with the regularizer . The two views are equivalent in this regard . 2See Bishop ( 2006 , Sec . 5.7.1 ) for more details . ( Gibbs , 1997 ; Spiegelhalter & Lauritzen , 1990 ; MacKay , 1992a ) : p ( y∗ = c | x∗ , D ) ≈ exp ( m∗c κ∗c ) ∑C i=1 exp ( m∗i κ∗i ) , for all c = 1 , . . . , C , ( 2 ) where for each i = 1 , . . . , C , the real numbers m∗i is the i-th component of the vector m∗ , and κ∗i : = ( 1 + π/8 v∗ii ) −1/2 where v∗ii is the i-th diagonal term of the matrix V∗ . These approximations are analytically useful , but can be expensive due to the computation of the Jacobian matrix J∗ . Thus , Monte Carlo ( MC ) integration is commonly used as an alternative , i.e . we approximate y∗ ≈ 1S ∑S s=1 h ( f ( x∗ ; θs ) ) with θs ∼ q ( θ | D ) . Finally , given a classification predictive distribution p ( y∗ | x∗ , D ) , we define the predictive confidence of x∗ as the maximum probability conf ( x∗ ) : = maxc∈ { 1 , ... , C } p ( y∗ = c | x∗ , D ) over class labels . 2.2 RELU AND GAUSSIAN PROCESSES . The ReLU activation function ReLU ( z ) : = max ( 0 , z ) ( Nair & Hinton , 2010 ) has become the defacto choice of non-linearity in deep learning . Given arbitrary real numbers c , it can be generalized as ReLU ( z ; c ) : = max ( 0 , z − c ) , with the “ kink ” at location c. An alternative formulation , useful below , is in terms of the Heaviside function H as ReLU ( z ; c ) = H ( z − c ) ( z − c ) . We may define a collection of d such ReLU functions evaluated at some point in R as the function φ : R → RK with z 7→ ( ReLU ( z ; c1 ) , . . . , ReLU ( z ; cK ) ) > . We call this function the ReLU feature map ; it can be interpreted as “ placing ” ReLU functions at different locations in R. Consider a linear model g : R × RK → R defined by g ( x ; w ) : = w > φ ( x ) . Suppose φ regularly places the K generalized ReLU functions centered at ( ci ) Ki=1 over [ cmin , cmax ] ⊂ R , where cmin < cmax . If we consider a Gaussian prior p ( w ) : = N ( w ∣∣0 , σ2K−1 ( cmax − cmin ) I ) over the weights w then , as K goes to infinity , the distribution over g ( x ) is a Gaussian process with mean 0 and covariance ( using the shorthand gx : = g ( x ) and x̄ : = min ( x , x′ ) ; full derivation in Appendix A ) : lim K→∞ cov ( gx , gx′ ) = σ 2H ( x̄− cmin ) ( 1 3 ( x̄3 − c3min ) − 1 2 ( x̄2 − c2min ) ( x+ x′ ) + ( x̄− cmin ) xx′ ) = : k1 ( x , x′ ; cmin , σ 2 ) , for x̄ ≤ cmax . Since this expression does not depend on cmax , we consider the limit cmax → ∞ . The resulting covariance function is the cubic spline kernel ( Wahba , 1990 ) . 3 METHOD . Hein et al . ( 2019 ) showed that the confidence of point-estimated ReLU networks ( i.e . feed-forward nets which use piecewise-affine activation functions and are linear in the output layer ) approaches 1 with increasing distance from the training data . For binary classification , Kristiadi et al . ( 2020 ) showed that Gaussian-approximated ReLU BNNs f instead approach a constant confidence bounded away from 1 , but not necessarily close to the maximum uncertainty value of 1/2 . Thus , just being Bayesian as such does not fix overconfidence entirely . A close look at their proof suggests that the issue is a structural limitation of the deep model itself : for any input x∗ and a sufficiently large scalar α , both the mean and standard deviation of the output f ( αx∗ ) are linear functions of x∗ . Intuitively , this issue arises because the net only has finitely many ReLU features available to “ explain ” the data , and thus it “ lacks ” ReLU features for modeling uncertainty away from the data . In this section , we will utilize the cubic spline kernel to construct a new kernel and method that , intuitively speaking , adds an infinite number ReLU features away from the data to pre-trained BNNs . This construction adds the “ missing ” ReLU features and endows BNNs with super-quadratic output variance growth , without affecting predictions . All proofs are in Appendix B .

Summary: The authors consider the issue of overconfidence in ReLU NN and BNNs, particularly for data that are far (in Euclidean distance) from the training data. They address this by modeling the residual (to the NN) in the latent space with a GP. The kernel for this GP is derived as the limit of infinitely many ReLU-based random features. Specifically, this kernel has the property that it scales cubically with the norm of the input, and so causes large uncertainty away from the origin. Crucially, the GP term changes little from its prior distribution after conditioning on the data, so no expensive inference is required under the approximation made.