Fast Predictive Uncertainty for Classification with Bayesian Deep Networks

In Bayesian Deep Learning , distributions over the output of classification neural networks are approximated by first constructing a Gaussian distribution over the weights , then sampling from it to receive a distribution over the categorical output distribution . This is costly . We extend existing work to construct a Dirichlet approximation of this output distribution , yielding an analytic map between Gaussian distributions in logit space and Dirichlet distributions in the output space . We argue that the resulting Dirichlet distribution has theoretical and practical advantages , in particular more efficient computation of the uncertainty estimate , scaling to large datasets and networks like ImageNet and DenseNet . We demonstrate the use of this Dirichlet approximation by using it to construct a lightweight uncertainty-aware output ranking for the ImageNet setup . 1 INTRODUCTION . Quantifying the uncertainty of Neural Networks ’ ( NNs ) predictions is important in safety-critical applications such as medical-diagnosis ( Begoli et al. , 2019 ) and self-driving vehicles ( McAllister et al. , 2017 ; Michelmore et al. , 2018 ) . Architectures for classification tasks produce a probability distribution as their output , constructed by applying the softmax to the point-estimate output of the penultimate layer . However , it has been shown that this distribution is overconfident ( Nguyen et al. , 2015 ; Hein et al. , 2019 ) and thus can not be used for predictive uncertainty quantification . Approximate Bayesian methods provide quantified uncertainty over the network ’ s parameters and thus the outputs in a tractable fashion . The commonly used Gaussian approximate posterior ( MacKay , 1992a ; Graves , 2011 ; Blundell et al. , 2015 ; Ritter et al. , 2018 ) approximately induces a Gaussian distribution over the logits of a NN ( Mackay , 1995 ) . But the associated predictive distribution does not have an analytic form . It is thus generally approximated by Monte Carlo ( MC ) integration requiring multiple samples . Predictions in Bayesian Neural Networks ( BNNs ) are thus generally expensive operations . In this paper , we re-consider an old but largely overlooked idea originally proposed by David JC MacKay ( 1998 ) in a different setting ( arguably the inverse of the Deep Learning setting ) which transforms a Dirichlet distribution into a Gaussian . Dirichlet distributions are generally defined on the simplex . But when its variable is defined on the inverse softmax ’ s domain , its shape effectively approximates a Gaussian . The inverse of this approximation , which will be called the Laplace Bridge here ( Hennig et al. , 2012 ) , analytically maps a Gaussian distribution onto a Dirichlet distribution . Given a Gaussian distribution over the logits of a NN , one can thus efficiently obtain an approximate Dirichlet distribution over the softmax outputs ( Figure 1 ) . Our contributions are : We re-visit MacKay ’ s derivation with particular attention to a symmetry constraint that becomes necessary in our “ inverted ” use of the argument from the Gaussian to the Dirichlet family . We then validate the quality of this approximation both theoretically and empirically , and demonstrate significant speed-up over MC-integration . Finally , we show a use-case , leveraging the analytic properties of Dirichlets to improve the popular top-k metric through uncertainties . Section 2 provides the mathematical derivation . Section 3 discusses the Laplace Bridge in the context of NNs . We compare it to the recent approximations of the predictive distributions of NNs in Section 4 . Experiments are presented in Section 5 . 2 THE LAPLACE BRIDGE . Laplace approximations1 are a popular and light-weight method to approximate a general probability distribution q ( x ) with a Gaussian N ( x|µ , Σ ) when q ( x ) is twice differentiable and the Hessian at the mode is positive definite . It sets µ to a mode of q , and Σ = − ( ∇2 log q ( x ) |µ ) −1 , the inverse Hessian of log q at that mode . This scheme can work well if the true distribution is unimodal and defined on the real vector space . The Dirichlet distribution , which has the density function Dir ( π|α ) : = Γ ( ∑K k=1 αk ) ∏K k=1 Γ ( αk ) K∏ k=1 παk−1k , ( 1 ) is defined on the probability simplex and can be “ multimodal ” in the sense that the distribution diverges in the k-corner of the simplex when αk < 1 . Both issues preclude a Laplace approximation , at least in the naı̈ve form described above . However , MacKay ( 1998 ) noted that both can be fixed , elegantly , by a change of variable ( see Figure 2 ) . Details of the following argument can be found in the supplements . Consider the K-dimensional variable π ∼ Dir ( π|α ) defined as the softmax of z ∈ RK : πk ( z ) : = exp ( zk ) ∑K l=1 exp ( zl ) , ( 2 ) for all k = 1 , . . . , K. We will call z the logit of π . When expressed as a function of z , the density of the Dirichlet in π has to be multiplied by the absolute value of the determinan of the Jacobian det ∂π ∂z = ∏ k πk ( z ) , ( 3 ) thus removing the −1 terms in the exponent : Dirz ( π ( z ) |α ) : = Γ ( ∑K k=1 αk ) ∏K k=1 Γ ( αk ) K∏ k=1 πk ( z ) αk , ( 4 ) This density of z , the Dirichlet distribution in the softmax basis , can now be accurately approximated by a Gaussian through a Laplace approximation , yielding an analytic map from the parameter space α ∈ RK+ to the parameter space of the Gaussian ( µ ∈ RK and symmetric positive definite Σ ∈ 1For clarity : Laplace approximations are also one out of several possible ways to construct a Gaussian approximation to the weight posterior of a NN , by constructing a second-order Taylor approximation of the empirical risk at the trained weights . This is not the way they are used in this section . The Laplace Bridge is agnostic to how the input Gaussian distribution is constructed . It could , e.g. , also be constructed as a variational approximation , or the moments of Monte Carlo samples . See also Section 5 . RK×K ) , given by µk = logαk − 1 K K∑ l=1 logαl , ( 5 ) Σk ` = δk ` 1 αk − 1 K [ 1 αk + 1 α ` − 1 K K∑ u=1 1 αu ] . ( 6 ) The corresponding derivations require care because the Gaussian parameter space is evidently larger than that of the Dirichlet and not fully identified by the transformation . A pseudo-inverse of this map was provided by Hennig et al . ( 2012 ) . It maps the Gaussian parameters to those of the Dirichlet as αk = 1 Σkk ( 1− 2 K + eµk K2 K∑ l=1 e−µl ) ( 7 ) ( this equation ignores off-diagonal elements of Σ , more discussion below ) . Together , Eqs . 5 , 6 and 7 will here be used for Bayesian Deep Learning , and jointly called the Laplace Bridge . A full derivation of the Laplace Bridge can be found in B and C. Even though the Laplace Bridge implies a reduction of the expressiveness of the distribution , we show in Section 3 that this map is still sufficiently accurate . Figure 1 shows the quality of the resulting approximation . We consider multiple different µ , Σ in three dimensions . We exhaustively sample from the Gaussian and apply the softmax . The resulting histogram is compared to the PDF of the corresponding Dirichlet . The first part of the figure emphasizes that a point estimate is insufficient . Since the mean for the Dirichlet is the normalized parameter vector α , the parameters ( α1 = [ 2 , 2 , 6 ] > and α2 = [ 11 , 11 , 51 ] > ) yield the same point estimate even though their distributions are clearly different . The second part shows how the Laplace Bridge maps w.r.t decreasing uncertainty . 3 THE LAPLACE BRIDGE FOR BNNS . Let fθ : RN → RK be an L-layer Neural Network parametrized by θ ∈ RP , with a Gaussian approximate posterior N ( θ|µθ , Σθ ) . For any input x ∈ RN , one way to obtain an approximate Gaussian distribution on the pre-softmax output ( logit vector ) fθ ( x ) = : z is as q ( z|x ) ≈ N ( z|fmµθ ( x ) , J ( x ) > ΣθJ ( x ) ) , ( 8 ) where J ( x ) is the P × K Jacobian matrix representing the derivative ∂z∂θ ( Mackay , 1995 ) . This is a linearized approximation of the network which doesn ’ t account for non-linearities and is often computationally infeasible . For the experiments with larger networks we therefore use a last-layer Laplace approximation of the network as successfully used , e.g . by Snoek et al . ( 2015 ) ; Wilson et al . ( 2016 ) ; Brosse et al . ( 2020 ) ; Kristiadi et al . ( 2020 ) . It is given by q ( z|x ) ≈ N ( z|µW ( l ) φ ( x ) , φ ( x ) TΣW ( l ) φ ( x ) ) , ( 9 ) where φ ( x ) denotes the output of a forward pass through the first L−1 layers , µW ( l ) is the maximum likelihood estimate for the weights of the last layer , and ΣW ( l ) is the inverse of the negative Hessian w.r.t . the loss ΣW ( l ) = − { ∂ 2L ∂2W ( l ) } −1 . Approximating the density of the softmax of this Gaussian random variable as a Dirichlet , using the Laplace Bridge , analytically approximates the predictive distribution in a single step , as opposed to many samples . From Eq . ( 7 ) , this requires O ( K ) computations to construct the K parameters αk of the Dirichlet . In contrast , MC-integration has computational costs of O ( MJ ) , where M is the number of samples and J is the cost of sampling from q ( z|x ) ( typically J is of order K2 after an initial O ( K3 ) operation for a matrix decomposition of the covariance ) . The Monte Carlo approximation has the usual sampling error of O ( 1/ √ M ) , while the Laplace Bridge has a fixed but small error ( empirical comparison in Section 5.2 ) . We now discuss several qualitative properties of the Laplace Bridge relevant for the uncertainty quantification use case in Deep Learning . For output classes of “ comparably high ” probability , as defined below , the variance Var ( πk|α ) under the Laplace Bridge increases with the variance of the underlying Gaussian . In this sense , the Laplace Bridge approximates the uncertainty information encoded in the output of a BNN . Proposition 1 ( proof in supplements ) . Let Dir ( π|α ) be obtained via the Laplace Bridge from a Gaussian O N ( z|µ , Σ ) over RK . Then , for each k = 1 , . . . , K , letting α6=k : = ∑ l 6=k αl , if αk > 1 4 ( √ 9α26=k + 10α 6=k + 1− α6=k − 1 ) , then the variance Var ( πk|α ) of the k-th component of π is increasing in Σkk . Intuitively , this result describes the condition that needs to be fulfilled such that the variance of the resulting Dirichlet scales with the variance of the k-th component of the Gaussian . It can be seen as a proxy for a high-quality approximation . An empirical evaluation showing that this condition is fulfilled in most cases can be found in the supplements . Further benefits of this approximation arise from the convenient analytical properties of the Dirichlet exponential family . For example , a point estimate of the posterior predictive distribution is directly given by the Dirichlet ’ s mean , Eπ = ( α1∑K l=1 αl , . . . , αK∑K l=1 αl ) > , ( 10 ) Additionally , Dirichlets have Dirichlet marginals : If p ( π ) = Dir ( π|α ) , then p π1 , π2 , . . . , πj , ∑ k > j πk > = Dir α1 , α2 , . . . , αj , ∑ k > j αk . ( 11 ) An additional benefit of the Laplace Bridge for BNNs is that it is more flexible than an MC-integral . If we let p ( π ) be the distribution over π : = softmax ( z ) : = [ ez1/ ∑ l e zl , . . . , ezK/ ∑ l e zl ] > , then the MC-integral can be seen as a “ point-estimate ” of this distribution since it approximates Eπ . In contrast , the Dirichlet distribution Dir ( π|α ) approximates the distribution p ( π ) . Thus , the Laplace Bridge enables tasks that can be done only with a distribution but not a point estimate . For instance , one could ask “ what is the distribution of the first L classes ? ” when one is dealing with K-class ( L < K ) classification . Since the marginal distribution can be computed analytically ( 11 ) , the Laplace Bridge provides a convenient yet cheap way of answering this question .

Summary: This article improves the efficiency of Bayesian neural networks (BNNs). It follows the sampling-free solutions within Bayesian deep learning (Wu et al., 2018; Haussmann et al., 2019, etc.). The difference is the proposed Laplace Bridge that approximates the  full distribution over the softmax outputs of a neural network. The Laplace Bridge tackles the problem of approximating the distribution over the softmax outputs of the ubiquitous Gaussian-approximated BNNs without any additional training  procedure. This allows the Laplace Bridge to be used with pre-trained networks and emphasizes its non-invasive nature.