Learnable Uncertainty under Laplace Approximations

Laplace approximations are classic , computationally lightweight means for constructing Bayesian neural networks ( BNNs ) . As in other approximate BNNs , one can not necessarily expect the induced predictive uncertainty to be calibrated . Here we develop a formalism to explicitly “ train ” the uncertainty in a decoupled way to the prediction itself . To this end we introduce uncertainty units for Laplaceapproximated networks : Hidden units with zero weights that can be added to any pre-trained , point-estimated network . Since these units are inactive , they do not affect the predictions . But their presence changes the geometry ( in particular the Hessian ) of the loss landscape around the point estimate , thereby affecting the network ’ s uncertainty estimates under a Laplace approximation . We show that such units can be trained via an uncertainty-aware objective , making the Laplace approximation competitive with more expensive alternative uncertaintyquantification frameworks . 1 INTRODUCTION . The point estimates of neural networks ( NNs ) —constructed as maximum a posteriori ( MAP ) estimates via ( regularized ) empirical risk minimization—empirically achieve high predictive performance . However , they tend to underestimate the uncertainty of their predictions , leading to an overconfidence problem ( Hein et al. , 2019 ) , which could be disastrous in safety-critical applications such as autonomous driving . Bayesian inference offers a principled path to overcome this issue . The goal is to turn a “ vanilla ” NN into a Bayesian neural network ( BNN ) , where the posterior distribution over the network ’ s weights are inferred via Bayes ’ rule and subsequently taken into account when making predictions . Since the cost of exact posterior inference in a BNN is often prohibitive , approximate Bayesian methods are employed instead . Laplace approximations ( LAs ) are classic methods for such a purpose ( MacKay , 1992b ) . The key idea is to obtain an approximate posterior by “ surrounding ” a MAP estimate of a network with a Gaussian , based on the loss landscape ’ s geometry around it . A standard practice in LAs is to tune a single hyperparameter—the prior precision—which is inflexible ( Ritter et al. , 2018b ; Kristiadi et al. , 2020 ) . Here , we aim at improving the flexibility of uncertainty tuning in LAs . To this end , we introduce Learnable Uncertainty under Laplace Approximations ( LULA ) units , which are hidden units associated with a zeroed weight . They can be added to the hidden layers of any MAP-trained network . Because they are inactive , such units do not affect the prediction of the underlying network . However , they can still contribute to the Hessian of the loss with respect to the parameters , and hence induce additional structures to the posterior covariance under a LA . LULA units can be trained via an uncertainty-aware objective ( Hendrycks et al. , 2019 ; Hein et al. , 2019 , etc . ) , such that they improve the predictive uncertainty-quantification ( UQ ) performance of the Laplace-approximated BNN . Figure 1 demonstrates trained LULA units in action : They improve the UQ performance of a standard LA , while keeping the MAP predictions in both regression and classification tasks . In summary , we ( i ) introduce LULA units : inactive hidden units for uncertainty tuning of a LA , ( ii ) bring a robust training technique from non-Bayesian literature for training these units , and ( iii ) show empirically that LULA-augmented Laplace-approximated BNNs can yield better UQ performance compared to both previous tuning techniques and contemporary , more expensive baselines . 2 BACKGROUND . 2.1 BAYESIAN NEURAL NETWORKS . Let f : Rn × Rd → Rk defined by ( x , θ ) 7→ f ( x ; θ ) be an L-layer neural network . Here , θ is the concatenation of all the parameters of f . Suppose that the size of each layer of f is given by the sequence of ( nl ∈ Z > 0 ) Ll=1 . Then , for each l = 1 , . . . , L , the l-th layer of f is defined by a ( l ) : = W ( l ) h ( l−1 ) + b ( l ) , with h ( l ) : = { ϕ ( a ( l ) ) if l < L a ( l ) if l = L , ( 1 ) where W ( l ) ∈ Rnl×nl−1 and b ( l ) ∈ Rnl are the weight matrix and bias vector of the layer , and ϕ is a component-wise activation function . We call the vector h ( l ) ∈ Rnl the l-th hidden units of f . Note that by convention , we consider n0 : = n and nL : = k , while h ( 0 ) : = x and h ( L ) : = f ( x ; θ ) . From the Bayesian perspective , the ubiquitous training formalism of neural networks amounts to MAP estimation : The empirical risk and the regularizer are interpretable as the negative loglikelihood under an i.i.d . dataset D : = { xi , yi } mi=1 and the negative log-prior , respectively . That is , the loss function is interpreted as L ( θ ) : = − m∑ i=1 log p ( yi | f ( xi ; θ ) ) − log p ( θ ) = − log p ( θ | D ) . ( 2 ) In this view , the de facto weight decay regularizer amounts to a zero-mean isotropic Gaussian prior p ( θ ) = N ( θ | 0 , λ−1I ) with a scalar precision hyperparameter λ . Meanwhile , the usual softmax and quadratic output losses correspond to the Categorical and Gaussian distributions over yi in the case of classification and regression , respectively . MAP-trained neural networks have been shown to be overconfident ( Hein et al. , 2019 ) and BNNs can mitigate this issue ( Kristiadi et al. , 2020 ) . They quantify epistemic uncertainty by inferring the full posterior distribution of the parameters θ ( instead of just a single point estimate in MAP training ) . Given that p ( θ | D ) is the posterior , then the prediction for any test point x ∈ Rn is obtained via marginalization p ( y | x , D ) = ∫ p ( y | f ( x ; θ ) ) p ( θ | D ) dθ , ( 3 ) which captures the uncertainty encoded in the posterior . 2.2 LAPLACE APPROXIMATIONS . In deep learning , since the exact Bayesian posterior is intractable , approximate Bayesian inference methods are used . An important family of such methods is formed by LAs . Let θMAP be the minimizer of ( 2 ) , which corresponds to a mode of the posterior distribution . A LA locally approximates the posterior using a Gaussian p ( θ | D ) ≈ N ( θ | θMAP , Σ ) : = N ( θ | θMAP , ( ∇2L|θMAP ) −1 ) . Thus , LAs construct an approximate Gaussian posterior around θMAP , whose precision equals the Hessian of the loss at θMAP—the “ curvature ” of the loss landscape at θMAP . While the covariance of a LA is tied to the weight decay of the loss , a common practice in LAs is to tune the prior precision under some objective , in a post-hoc manner . In other words , the MAP estimation and the covariance inference are thought as separate , independent processes . For example , given a fixed MAP estimate , one can maximize the log-likelihood of a LA w.r.t . the prior precision to obtain the covariance . This hyperparameter tuning can thus be thought as an uncertainty tuning . A recent example of LAs is the Kronecker-factored Laplace ( KFL , Ritter et al. , 2018b ) . The key idea is to approximate the Hessian matrix with the layer-wise Kronecker factorization scheme proposed by Heskes ( 2000 ) ; Martens & Grosse ( 2015 ) . That is , for each layer l = 1 , . . . , L , KFL assumes that the Hessian corresponding to the l-th weight matrix W ( l ) ∈ Rnl×nl−1 can be written as the Kronecker product G ( l ) ⊗ A ( l ) for some G ( l ) ∈ Rnl×nl and A ( l ) ∈ Rnl−1×nl−1 . This assumption brings the inversion cost of the Hessian down to Θ ( n3l +n 3 l−1 ) , instead of the usual Θ ( n 3 l n 3 l−1 ) cost . The approximate Hessian can easily be computed via tools such as BackPACK ( Dangel et al. , 2020 ) . Even with a closed-form Laplace-approximated posterior , the predictive distribution ( 3 ) in general does not have an analytic solution since f is nonlinear . Instead , one can employ Monte-Carlo ( MC ) integration by sampling from the Gaussian : p ( y | x , D ) ≈ 1 S S∑ s=1 p ( y | f ( x ; θs ) ) ; θs ∼ N ( θ | θMAP , Σ ) , for S number of samples . In the case of binary classification with f : Rn × Rd → R , one can use the following well-known approximation , due to MacKay ( 1992a ) : p ( y = 1 | x , D ) ≈ σ ( f ( x ; θMAP ) √ 1 + π/8 v ( x ) ) , ( 4 ) where σ is the logistic-sigmoid function and v ( x ) is the marginal variance of the network output f ( x ) , which is often approximated via a linearization of the network around the MAP estimate : v ( x ) ≈ ( ∇θf ( x ; θ ) |θMAP ) > Σ ( ∇θf ( x ; θ ) |θMAP ) . ( 5 ) ( This approximation has also been generalized to multi-class classifications by Gibbs ( 1997 ) . ) In particular , as v ( x ) increases , the predictive probability of y = 1 goes to 0.5 and therefore the uncertainty increases . This relationship has also been shown empirically in multi-class classifications with MC-integration ( Kristiadi et al. , 2020 ) . 3 LULA UNITS . The problem with the standard uncertainty tuning in LAs is that the only degree-of-freedom available for performing the optimization is the scalar prior precision and therefore inflexible.1 We shall address this by introducing “ uncertainty units ” , which can be added on top of the hidden units of any MAP-trained network ( Section 3.1 ) and can be trained via an uncertainty-aware loss ( Section 3.2 ) . 3.1 CONSTRUCTION . Let f : Rn × Rd → Rk be a MAP-trained L-layer neural network with parameters θMAP = { W ( l ) MAP , b ( l ) MAP } Ll=1 . The premise of our method is simple : At each hidden layer l = 1 , . . . , L − 1 , 1While one can also use a non-scalar prior precision , it appears to be uncommon in deep learning . In any case , such a element-wise weight-cost would interact with the training procedure . MAP . The additional units are represented by the additional block at the bottom of each layer . Dashed lines correspond to the free parameters Ŵ ( 1 ) , . . . , Ŵ ( L−1 ) , while dotted lines to the zero weights . suppose we add ml ∈ Z≥0 additional hidden units , under the original activation function , to h ( l ) . As a consequence , we need to augment each of the weight matrices to accommodate them . Consider the following construction : for each layer l = 1 , . . . , L − 1 of the network , we expand W ( l ) and b ( l ) to obtain the block matrix and vector W̃ ( l ) : = ( W ( l ) MAP 0 Ŵ ( l ) 1 Ŵ ( l ) 2 ) ∈ R ( nl+ml ) × ( nl−1+ml−1 ) ; b̃ ( l ) : = ( b ( l ) MAP b̂ ( l ) ) ∈ Rnl+ml , ( 6 ) respectively , with m0 = 0 since we do not add additional units to the input . For l = L , we define W̃ ( L ) : = ( W ( L ) MAP , 0 ) ∈ R k× ( nL−1+mL−1 ) ; b̃ ( L ) : = b ( L ) MAP ∈ R k , so that the output dimensionality is unchanged . For brevity , we denote Ŵ ( l ) : = ( Ŵ ( l ) 1 , Ŵ ( l ) 2 ) . Refer to Figure 2 for an illustration and Algorithm 2 in Appendix B for a step-by-step summary . Taken together , we denote the resulting augmented network as f̃ and the resulting parameter vector as θ̃MAP ∈ Rd̃ , where d̃ it the resulting number of parameters . Note that we can easily extend this construction to convolutional nets by expanding the “ channel ” of a hidden layer.2 Let us inspect the implication of this construction . Here for each l = 1 , . . . , L − 1 , since they are zero , the upper-right quadrant of W̃ ( l ) deactivates the ml−1 additional hidden units in the previous layer , thus they do not contribute to the original hidden units in the l-th layer . Meanwhile , the submatrix Ŵ ( l ) and the sub-vector b̂ ( l ) contain parameters for the additional ml hidden units in the l-th layer . We are free to choose the the values of these parameters since the following proposition guarantees that they will not change the output of the network ( the proof is in Appendix A ) . Proposition 1 . Let f : Rn × Rd → Rk be a MAP-trained L-layer network parametrized by θMAP . Suppose f̃ : Rn × Rd̃ → R and θ̃MAP ∈ Rd̃ are obtained via the previous construction . For any input x ∈ Rn , we have f̃ ( x ; θ̃MAP ) = f ( x ; θMAP ) . So far , it looks like all our changes to the network are inconsequential . However , they do affect the curvature of the landscape of L,3 and thus the uncertainty arising in a LA . Let θ̃ be a random variable in Rd̃ and θ̃MAP be an instance of it . Suppose we have a Laplace-approximated posterior p ( θ̃ | D ) ≈ N ( θ̃ | θ̃MAP , Σ̃ ) over θ̃ , where the covariance Σ̃ is the inverse Hessian of the negative log-posterior w.r.t . the augmented parameters at θ̃MAP . Then , Σ̃ contains additional dimensions ( and thus in general , additional structured , non-zero uncertainty ) absent in the original network , which depend on the values of the free parameters { Ŵ ( l ) , b̂ ( l ) } L−1l=1 . 2E.g . if the hidden units are a 3D array of ( channel × height × width ) , then we expand the first dimension . 3More formally : The principal curvatures of the graph of L , seen as a d-dimensional submanifold of Rd+1 . The implication of the previous finding can be seen clearly in real-valued networks with diagonal LA posteriors . The following proposition shows that , under such a network and posterior , the construction above will affect the output uncertainty of the original network f ( the proof is in Appendix A ) . Proposition 2 . Suppose f : Rn × Rd → R is a real-valued network and f̃ is as constructed above . Suppose further that diagonal Laplace-approximated posteriors N ( θ | θMAP , diag ( σ ) ) , N ( θ̃ | θ̃ , diag ( σ̃ ) ) are employed . Using the linearization ( 5 ) , for any input x ∈ Rn , the variance over the output f̃ ( x ; θ̃ ) is at least that of f ( x ; θ ) . In summary , the construction along with Propositions 1 and 2 imply that the additional hidden units we have added to the original network are uncertainty units under LAs , i.e . hidden units that only contribute to the Laplace-approximated uncertainty and not the predictions . This property gives rise to the name Learnable Uncertainty under Laplace Approximations ( LULA ) units .

Summary: The paper proposes a post-hoc uncertainty tuning pipeline for Bayesian neural networks. After getting the point estimate, it adds extra dimensions to the weight matrices and hidden layers, which has no effect on the network output, with the hope that it would influence the variance of the original network weights under the Laplacian approximation. More specifically, it tunes the extra weights by optimizing another objective borrowed from the non-Bayesian robust learning literature, which encourages low uncertainty over real (extra, validation) data, and high uncertainty over manually constructed, out-of-distribution data.