Original: Fast Predictive Uncertainty for Classification with Bayesian Deep Networks
Revised: Fast Predictive Uncertainty for Classification with Bayesian Deep Networks

Original: Abstract
Revised: Abstract

Original: 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 softmax outputs.
Revised: In Bayesian Deep Learning, distributions over the output of classification neural networks are often approximated by first constructing a Gaussian distribution over the weights, then sampling from it to receive a distribution over the softmax outputs.

Original: This is costly.
Revised: This is costly.

Original: We reconsider old work (Laplace Bridge) to construct a Dirichlet approximation of this softmax output distribution, which yields an analytic map between Gaussian distributions in logit space and Dirichlet distributions (the conjugate prior to the Categorical distribution) in the output space.
Revised: We reconsider old work (Laplace Bridge) to construct a Dirichlet approximation of this softmax output distribution, which yields an analytic map between Gaussian distributions in logit space and Dirichlet distributions (the conjugate prior to the Categorical distribution) in the output space.

Original: Importantly, the vanilla Laplace Bridge comes with certain limitations.
Revised: Importantly, the vanilla Laplace Bridge comes with certain limitations.

Original: We analyze those and suggest a simple solution that compares favorably to other commonly used estimates of the softmax-Gaussian integral.
Revised: We analyze those and suggest a simple solution that compares favorably to other commonly used estimates of the softmax-Gaussian integral.

Original: We demonstrate that the resulting Dirichlet distribution has multiple advantages, in particular, more efficient computation of the uncertainty estimate and scaling to large datasets and networks like ImageNet and DenseNet.
Revised: We demonstrate that the resulting Dirichlet distribution has multiple advantages, in particular, more efficient computation of the uncertainty estimate and scaling to large datasets and networks like ImageNet and DenseNet.

Original: We further demonstrate the usefulness of this Dirichlet approximation by using it to construct a lightweight uncertainty-aware output ranking for ImageNet.
Revised: We further demonstrate the usefulness of this Dirichlet approximation by using it to construct a lightweight uncertainty-aware output ranking for ImageNet.

Original:  INTRODUCTION
Revised:  INTRODUCTION

Original: 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
Revised: 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

Original: [McAllister et al., 2017, Michelmore et al., 2018], but it is often limited by computational constraints.
Revised: [McAllister et al., 2017, Michelmore et al., 2018], but it is often limited by computational constraints.

Original: 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.
Revised: 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.

Original: However, it has been shown that this distribution is overconfident[Nguyen et al., 2015, Hein et al., 2019] and thus cannot be
Revised: However, it has been shown that this distribution is overconfident

Original: [Figure]
Revised: [Figure]

Original: e
Revised: p ( y | x, D ) denotes the marginalized softmax output, i.e. the mean of the Dirichlet.

Original:  used for predictive uncertainty quantification.
Revised: [Nguyen et al., 2015, Hein et al., 2019] and thus cannot be used for predictive uncertainty quantification.

Original: Approximate
Revised: Approximate

Original: Bayesian methods provide quantified uncertainty over the
Revised: Bayesian methods provide quantified uncertainty over the

Original: NN’s parameters in a tractable fashion.
Revised: NN’s parameters in a tractable fashion.

Original: The commonly used
Revised: The commonly used

Original: 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
Revised: 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

Original: NN [Mackay, 1995], but the associated predictive distribution is not analytic.
Revised: NN [Mackay, 1995], but the associated predictive distribution is not analytic.

Original: It is typically approximated by Monte
Revised: It is typically approximated by Monte

Original: Carlo (MC) integration.
Revised: Carlo (MC) integration.

Original: This requires multiple samples, making prediction in Bayesian Neural Networks (BNNs) a comparably expensive operation.
Revised: This requires multiple samples, making prediction in Bayesian Neural Networks (BNNs) a comparably expensive operation.

Original: Here we reconsider 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.
Revised: Here we reconsider 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.

Original: When Dirichlet distributions are transformed with the inverse-softmax function, its shape effectively approximates a Gaussian.
Revised: When Dirichlet distributions are transformed with the inverse-softmax function, its shape effectively approximates a Gaussian.

Original: The inverse of this approximation, which will be called the Laplace Bridge (LB) here [Hennig et al., 2012], can also in principle analytically map the parameters of a Gaussian distribution onto those of a Dirichlet distribution.
Revised: The inverse of this approximation, which will be called the Laplace Bridge (LB) here [Hennig et al., 2012], can also in principle analytically map the parameters of a Gaussian distribution onto those of a Dirichlet distribution.

Original: Given a Gaussian distribution over the logits of a NN, one can thus efficiently obtain an approximate Dirichlet distribution over the softmax outputs.
Revised: Given a Gaussian distribution over the logits of a NN, one can thus efficiently obtain an approximate Dirichlet distribution over the softmax outputs.

Original: However, the bridge was previously used to map in the Gaussian to Dirichlet direction.
Revised: However, the bridge was previously used to map in the Gaussian to Dirichlet direction.

Original: The inverse direction of the vanilla LB has some limitations, arguably caused by the larger state-space of
Revised: The inverse direction of the vanilla LB has some limitations, arguably caused by the larger state-space of

Original: Gaussian relative to the Dirichlet exponential family.
Revised: Gaussian relative to the Dirichlet exponential family.

Original: π p ( π ) in basis

Original: B N −−   x p ( x )

Original: Laplace in logit basis

Original: B N 0.

Original: π p ( π )

Original: Back to standard basis

Original: B N

Original: Our contributions are a)
Revised: Our contributions are a)

Original: We analyze these limits and suggest a solution that allows for the practical application of the LB.
Revised: We analyze these limits and suggest a solution that allows for the practical application of the LB.

Original: b) We show how the result can be used in the context of
Revised: b) We show how the result can be used in the context of

Original: BNNs (see Figure 1 and 3).
Revised: BNNs (see Figure 1 and 3).

Original: c)
Revised: c)

Original: We empirically evaluate the quality of this approximation, its speed-up, and its performance for out-of-data distribution tasks.
Revised: We empirically evaluate the quality of this approximation, its speed-up, and its performance for out-of-data distribution tasks.

Original: d) Finally, we show a use case on ImageNet, leveraging the analytic properties of Dirichlets to improve the popular top- k metric through uncertainties.
Revised: d) Finally, we show a use case on ImageNet, leveraging the analytic properties of Dirichlets to improve the popular top- k metric through uncertainties.

Original:  THE LAPLACE BRIDGE
Revised:  THE LAPLACE BRIDGE

Original: Laplace approximations 1 [MacKay, 1992a, Daxberger et al., 2021] are a popular and lightweight method to approximate general probability distributions q ( x ) with a Gaussian
Revised: Laplace approximations 1 [MacKay, 1992a, Daxberger et al., 2021] are a popular and lightweight method to approximate general probability distributions q ( x ) with a Gaussian

Original: N ( x | µ , Σ ) when q ( x ) is twice differentiable and the Hessian at the mode is positive definite.
Revised: N ( x | µ , Σ ) when q ( x ) is twice differentiable and the Hessian at the mode is positive definite.

Original: They set µ to a mode of q , and Σ = − ( ∇ 2 log q ( x ) | µ ) − 1 , the inverse Hessian of log q at that mode.
Revised: They set µ to a mode of q , and Σ = − ( ∇ 2 log q ( x ) | µ ) − 1 , the inverse Hessian of log q at that mode.

Original: This scheme can work well if the true distribution is unimodal and defined on the real vector space.
Revised: This scheme can work well if the true distribution is unimodal and defined on the real vector space.

Original: The Dirichlet distribution, which has the density function
Revised: The Dirichlet distribution, which has the density function

Original: [Equation]
Revised: [Equation]

Original: is defined on the probability simplex and can be “multimodal” in the sense that the distribution diverges in the kcorner of the simplex when α k < 1 .
Revised: 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 .

Original: Both issues preclude a
Revised: This precludes a

Original: Laplace approximation, at least in the naïve form described above.
Revised: Laplace approximation, at least in the naïve form described above.

Original: However, MacKay [1998] noted that both can be fixed elegantly by a change of variable (Figure 2).
Revised: However, MacKay [1998] noted that both can be fixed elegantly by a change of variable (Figure 2).

Original: Details of the following argument can be found in Appendices B and C. Consider the K -dimensional variable π ∼ Dir( π | α ) defined as the softmax of z ∈ R K :
Revised: Details of the following argument can be found in Appendices B and C. Consider the K -dimensional variable π ∼ Dir( π | α ) defined as the softmax of z ∈ R K :

Original: [Equation]
Revised: [Equation]

Original: for all k = 1 , . . .
Revised: for all k = 1 , . . .

Original: , K .
Revised: , K .

Original: We will call z the logit of π .
Revised: We will call z the logit of π .

Original: When expressed as a function of z , the density of the Dirichlet in π has to be multiplied by the absolute value of the determinant of the Jacobian
Revised: When expressed as a function of z , the density of the Dirichlet in π has to be multiplied by the absolute value of the determinant of the Jacobian

Original: [Equation]
Revised: [Equation]

Original: thus removing the “ − 1 ” terms in the exponent:
Revised: thus removing the “ − 1 ” terms in the exponent:

Original: [Equation]
Revised: [Equation]

Original: This density of z , the Dirichlet distribution in the softmax basis , can now be accurately approximated by a Gaussian through a Laplace approximation (see Figure 2), yielding an analytic map from the parameter α ∈ R K + to the parameters of the Gaussian ( µ ∈ R K and symmetric positive definite Σ ∈ R K × K ), given by
Revised: This density of z , the Dirichlet distribution in the softmax basis , can now be accurately approximated by a Gaussian through a Laplace approximation (see Figure 2), yielding an analytic map from the parameter α ∈ R K + to the parameters of the Gaussian ( µ ∈ R K and symmetric positive definite Σ ∈ R K × K ), given by

Original: [Equation]
Revised: [Equation]

Original: The corresponding derivations require care because the
Revised: The corresponding derivations require care because the

Original: Gaussian parameter space is evidently larger than that of the Dirichlet and not fully identified by the transformation.
Revised: Gaussian parameter space is evidently larger than that of the Dirichlet and not fully identified by the transformation.

Original: A pseudo-inverse of this map was provided as a side result in Hennig et al.
Revised: A pseudo-inverse of this map was provided as a side result in Hennig et al.

Original: It maps the Gaussian parameters to those of the Dirichlet as
Revised: It maps the Gaussian parameters to those of the Dirichlet as

Original: [Equation]
Revised: [Equation]

Original: (this equation ignores off-diagonal elements of Σ , more discussion below).
Revised: (this equation ignores off-diagonal elements of Σ , more discussion in Appendix C).

Original: Together, Eqs.
Revised: Together, Eqs.

Original: (5), (6) and (7) will be called the Laplace Bridge .
Revised: (5), (6) and (7) will be called the Laplace Bridge .

Original: For Bayesian Deep Learning, we only use Equation (7) which maps from µ, Σ to α .
Revised: For Bayesian Deep Learning, we only use Equation (7) which maps from µ , Σ to α .

Original: Even though the LB implies a reduction of the distribution’s expressiveness, we show in Section 3 that this map is still sufficiently accurate.
Revised: Even though the LB implies a reduction of the distribution’s expressiveness, we show in Section 3 that this map is still sufficiently accurate.

Original:  THE LAPLACE BRIDGE FOR BNNS
Revised:  THE LAPLACE BRIDGE FOR BNNS

Original: The Laplace Bridge can be applied to any NN setup that maps from a Gaussian to probabilities by using the softmax.
Revised: The Laplace Bridge can be applied to any NN setup that maps from a Gaussian to probabilities by using the softmax.

Original: Throughout this paper we use a last-layer Laplace approximation of the network as successfully used e.g. by Snoek et al.
Revised: Throughout this paper, we use a last-layer Laplace approximation of the network as successfully used e.g. by Snoek et al.

Original: [2015], Kristiadi et al.
Revised: [2015], Kristiadi et al.

Original: It is given by
Revised: It is given by

Original: [Equation]
Revised: [Equation]

Original: where ϕ ( x ) denotes the output of the first L − 1 layers, µ W ( l ) is the maximum a posteriori (MAP) estimate for the weights of the last layer, and Σ W ( l ) is the inverse of the negative loss Hessian w.r.t.
Revised: where ϕ ( x ) denotes the output of the first L − 1 layers, µ W ( l ) is the maximum a posteriori (MAP) estimate for the weights of the last layer, and Σ W ( l ) is the inverse of the negative loss Hessian w.r.t.

Original: W ( l ) , Σ W ( l ) = − ( ∇ 2 W ( l ) L ) − 1 .
Revised: W ( l ) , Σ W ( L ) = − ( ∇ 2 W ( L ) L ) −around the MAP estimate W ( L ) .

Original: Even though last-layer Laplace approximations only use uncertainty from the last linear layer, they empirically perform as well as full Laplace approximations
Revised: Even though last-layer Laplace approximations only use uncertainty from the last linear layer, they empirically perform as well as full Laplace approximations

Original: [Kristiadi et al., 2020].
Revised: [Kristiadi et al., 2020].

Original: Furthermore, they allow for very fast inference, thus being a good match for the LB.
Revised: Furthermore, they allow for very fast inference, thus being a good match for π p ( π ) B N −−   z p ( x ) Laplace in logit basis B N 0.

Original:  We use diagonal and Kronecker approximations to the Hessian (see Appendix D).
Revised: π p ( π ) Back basis B N the LB. We use diagonal and Kronecker approximations to the Hessian (see Appendix D).

Original: Using the LB we can analytically approximate the density of the softmax-Gaussian random variable that is the output of the BNN as a Dirichlet rather than using many samples.
Revised: Using the LB we can analytically approximate the density of the softmax-Gaussian random variable that is the output of the BNN as a Dirichlet rather than using many samples.

Original: As shown in Eq.
Revised: As shown in Eq.

Original: (7), it requires O ( K ) computations to construct the K parameters α k of the Dirichlet.
Revised: (7), it requires O ( K ) computations to construct the K parameters α k of the Dirichlet.

Original: 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 K 2 after an initial O ( K 3 ) operation for a matrix decomposition of the covariance).
Revised: 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 K 2 after an initial O ( K 3 ) operation for a matrix decomposition of the covariance).

Original: The MC approximation has the usual sampling error of O (1 / √ M ) , while the LB has a fixed but small error (empirical comparison in Section 6.4).
Revised: The MC approximation has the usual sampling error of O (1 / √ M ) , while the LB has a fixed but small error (empirical comparison in Section 6.4).

Original: This means that computing the LB is faster than drawing a single MC sample while yielding a full distribution.
Revised: This means that computing the LB is faster than drawing a single MC sample while yielding a full distribution.

Original: Further benefits of this approximation arise from the convenient analytical properties of the Dirichlet exponential family.
Revised: Further benefits of this approximation arise from the convenient analytical properties of the Dirichlet exponential family.

Original: For example, a point estimate of the posterior predictive distribution is directly given by the Dirichlet’s mean,
Revised: For example, a point estimate of the posterior predictive distribution is directly given by the Dirichlet’s mean,

Original: [Equation]
Revised: [Equation]

Original: This removes the necessity for MC integration and can be computed analytically.
Revised: This removes the necessity for MC integration and can be computed analytically.

Original: Additionally, Dirichlets have Dirichlet marginals: If p ( π ) =
Revised: Additionally, Dirichlets have Dirichlet marginals: If p ( π ) =

Original: Dir( π | α ) , then 
Revised: Dir( π | α ) , then p 

Original: [Equation]
Revised: [Equation]

Original: Thus marginal distributions  arbitrary subsets of outputs (including binary marginals) can be computed in closed-form.
Revised: Thus marginal distributions of arbitrary subsets of outputs (including binary marginals) can be computed in closedform.

Original: An additional benefit of the LB for BNNs is that it is more flexible than an MC-integral.
Revised: An additional benefit of the LB for BNNs is that it is more flexible than an MC-integral.

Original: If we let p ( π ) be the distribution over π := softmax( z ) :=
Revised: If we let p ( π ) be the distribution over π := softmax( z ) : =

Original: [ e z 1 / (cid:80) l e z l , . . , e z K /
Revised: [Equation]

Original: [Equation]
Revised: [Equation]

Original: l e z l ] ⊤ , then the MC-integral can be seen as a “point-estimate” of this distribution since it approximates E [ π ] .
Revised: l e z l ] ⊤ , then the MC-integral can be seen as a “point-estimate” of this distribution since it approximates E [ π ] .

Original: In contrast, the Dirichlet distribution
Revised: In contrast, the Dirichlet distribution

Original: Dir( π | α ) approximates the distribution p ( π ) .
Revised: Dir( π | α ) approximates the distribution p ( π ) .

Original: Thus, the LB enables tasks that can be done only with a distribution but not a point estimate.
Revised: Thus, the LB enables tasks that can be done only with a distribution but not a point estimate.

Original: For instance, one could ask “what is the distribution of the softmax output of the first L classes?”
Revised: For instance, one could ask “what is the distribution of the softmax output of the first L classes?”

Original: when one is dealing with K -class
Revised: when one is dealing with K -class

Original: ( L < K ) classification.
Revised: ( L < K ) classification.

Original: Since the marginal distribution can be computed analytically with Eq.
Revised: Since the marginal distribution can be computed analytically with Eq.

Original: (10), the LB provides a convenient yet cheap way of answering this question.
Revised: (10), the LB provides a convenient yet cheap way of answering this question.

Original:  LIMITATIONS OF THE LAPLACE BRIDGE
Revised:  LIMITATIONS OF THE LAPLACE BRIDGE

Original: There are two limitations to applying the LB as presented in
Revised: There are two limitations to applying the LB as presented in

Original: Equation (7).
Revised: Equation (7).

Original: First, the LB assumes that the random variable of the Gaussian sums to zero  (see Appendix C).
Revised: First, the LB assumes that the random variable of the Gaussian sums to zero due to the difference in degrees of freedom between Dirichlet and Gaussian (see Appendix C).

Original: Thus, we have to add a correction that projects from any arbitrary Gaussian to one that fulfills this constraint.
Revised: Thus, we have to add a correction that projects from any arbitrary Gaussian to one that fulfills this constraint.

Original: The resulting Gaussian (see Appendix A) is
Revised: The resulting Gaussian (see Appendix A) is

Original: [Equation]
Revised: [Equation]

Original: where 1 is the one-vector of size K .
Revised: where 1 is the one-vector of size K .

Revised: Second, the softmax-Dirichlet distribution is asymmetric

Original: [Figure] [Figure]
Revised: [Figure] [Figure]

Original: provides good fit

Original: However some high

Original: -Dirichlet becomes asymmetric and thus the Gaussian is a suboptimal propose projects a lowervariance region before the LB

Original: This can be understood “pulling the to the dynamic of the logistic function ( ) and thus yields a better approximation. Second, the softmax-Dirichlet distribution is asymmetric for extremely sparse cases (see Figure 4).
Revised: for extremely sparse cases (see Figure 4).

Original: These arise in regions where the logistic transform (the 1D special case of the softmax) is completely flat (as indicated by its derivative in Figure 4).
Revised: These arise in regions where the logistic transform (the 1D special case of the softmax) is nearly flat (as indicated by its derivative in Figure 4).

Original: For these two reasons the LA is suboptimal in these high-variance cases.
Revised: Therefore, the LA is suboptimal in these highvariance cases.

Original: This limitation can also be explained by looking at Equation (7).
Revised: This limitation can also be explained by looking at Equation (7).

Original: We observe that Σ contributes linearly to α with 1Σ kk while µ contributes exponentially with exp( µ k ) .
Revised: We observe that Σ contributes linearly to α with 1Σ kk while µ contributes exponentially with exp( µ k ) .

Original: For settings where Σ is small this doesn’t have a large effect.
Revised: For settings where Σ is small, this doesn’t have a large effect.

Original: However, when Σ kk and µ k grow the LB results differ from softmax Gaussian samples.
Revised: However, when Σ kk and µ k grow the LB results differ from softmax Gaussian samples.

Original: In the LB, the resulting α is dominated by the mean and the linear influence of the variance cannot correct sufficiently.
Revised: In the LB, the resulting α is dominated by the mean and the linear influence of the variance cannot correct sufficiently.

Original: For Monte Carlo sampling, on the other hand, the result is mostly determined by the large variance and then amplified through the softmax.
Revised: For Monte Carlo sampling, on the other hand, the result is mostly determined by the large variance and then amplified through the softmax.

Original: Our proposed normalization to the LB reduces this effect (see
Revised: Our proposed normalization to the LB reduces this effect (see

Original: Figure 3).
Revised: Figure 3).

Original: In BNNs, we often encounter such cases, especially far away from the data (see Figure 3 top).
Revised: In BNNs, we often encounter such cases, especially far away from the data (see Figure 3 top).

Original: Therefore, we propose an additional correction for practical purposes:
Revised: Therefore, we propose an additional correction for practical purposes:

Original: [Equation]
Revised: [Equation]

Original: where v mean (Σ) denotes the mean variance of Σ .
Revised: where v mean (Σ) denotes the mean variance of Σ , v mean (Σ)

Revised: = (cid:80) i Σ ii .

Original: The factor of
Revised: The factor of

Original: [Equation]
Revised: [Equation]

Original: / 2 is added because we found that higher dimensionalities require less correction.
Revised: is added because we found that higher dimensionalities require less correction.

Original: Since our correction is just a rescaling, Equation (11) is still fulfilled.
Revised: Since our correction is just a rescaling, the zero-sum constrained is still fulfilled.

Original: This normalization that can be understood as “pulling back” the distribution into a space where it is symmetric has higher approximation quality.
Revised: This normalization that can be understood as “pulling back” the distribution into a space where it is symmetric has higher approximation quality.

Revised: This correction is applied after the zero-sum constraint correction.

Revised: We want to point out that our correction is motivated by experimentation and the theoretical insights detailed above.

Revised: There is no theoretical derivation from first principles for the correction.

Revised: We provide additional explanations and figures in Appendix A.

Original: Throughout the paper, we will call this normalizing correction LB-norm and explicitly state when we use it.
Revised: Throughout the paper, we will call this normalizing correction LB-norm and explicitly state when we use it.

Original: Otherwise, we will use the vanilla version with zero sum correction.
Revised: Otherwise, we will use the vanilla version with zero-sum correction.

Original:  RELATED WORK
Revised:  RELATED WORK

Original: In BNNs, analytic approximations of posterior predictive distributions have attracted a great deal of research.
Revised: In BNNs, analytic approximations of posterior predictive distributions have attracted a great deal of research.

Original: In the binary classification case, for example, the probit approximation [Gibbs, 1997, Lu et al., 2020] has been proposed already in the 1990s [Spiegelhalter and Lauritzen, 1990,
Revised: In the binary classification case, for example, the probit approximation [Gibbs, 1997, Lu et al., 2020] has been proposed already in the 1990s [Spiegelhalter and Lauritzen, 1990,

Original: MacKay, 1992b].
Revised: MacKay, 1992b].

Original: However, while there exist some bounds [Titsias, 2016] and approximations of the expected logsum-exponent function [Ahmed and Xing, 2007, Braun and
Revised: However, while there exist some bounds [Titsias, 2016] and approximations of the expected logsum-exponent function [Ahmed and Xing, 2007, Braun and

Original: McAuliffe, 2010], in the multi-class case, obtaining a good analytic approximation of the expected softmax function under a Gaussian measure is an open problem.
Revised: McAuliffe, 2010], in the multi-class case, obtaining a good analytic approximation of the expected softmax function under a Gaussian measure is an open problem.

Original: Our LB can be used to produce a close analytical approximation of this integral.
Revised: Our LB can be used to produce a close analytical approximation of this integral.

Original: It thus furthers the trend of sampling-free solutions within Bayesian Deep Learning [Wu et al., 2018, Haussmann et al., 2019, etc.].
Revised: It thus furthers the trend of sampling-free solutions within Bayesian Deep Learning [Wu et al., 2018, Haussmann et al., 2019, etc.].

Original: The crucial difference is that, unlike these methods, the LB approximates the full distribution over the softmax outputs of a deep network.
Revised: The crucial difference is that, unlike these methods, the LB approximates the full distribution over the softmax outputs of a deep network.

Original: Recently, it has been proposed to model the distribution of softmax outputs of a network directly.
Revised: Previous approaches proposed to model the distribution of softmax outputs of a network directly.

Original: Similar to the 
Revised: Similar to the LB,

Original: LB, Malinin and Gales [2018, 2019], Sensoy et al.
Revised: Malinin and Gales [2018, 2019], Sensoy et al.

Original: [2018] proposed to use the Dirichlet distribution to model the posterior predictive for non-Bayesian networks.
Revised: [2018] proposed to use the Dirichlet distribution to model the posterior predictive for non-Bayesian networks.

Original: They further proposed novel training techniques in order to directly learn the Dirichlet.
Revised: They further proposed novel training techniques in order to directly learn the Dirichlet.

Revised: Additionally, different work on Distillation [Malinin et al., 2019, Vadera et al., 2020] takes larger models and distills them into a smaller one.

Revised: The result of some distillation methods is a Dirichlet similar to the LB.

Original:  In contrast, the LB tackles the problem of approximating the distribution over the softmax outputs of the ubiquitous Gaussian-approximated BNNs [Graves, 2011, Blundell et al., 2015, Louizos and Welling, 2016, Sun et al., 2017, etc] without any additional training procedure.
Revised: We compare against prior nets in the experiments. In contrast, the LB tackles the problem of approximating the distribution over the softmax outputs of the ubiquitous Gaussian-approximated BNNs [Graves, 2011, Blundell et al., 2015, Louizos and Welling, 2016, Sun et al., 2017, etc] without any additional training procedure.

Original: Therefore the LB can, for example, be used with pre-trained weights on large datasets.
Revised: Therefore the LB can, for example, be used with pre-trained weights on large datasets while prior networks and distillation usually require training from scratch.

Original:  EXPERIMENTS
Revised:  EXPERIMENTS

Original: We conduct five experiments.
Revised: We conduct multiple experiments.

Original: Firstly, we compare the LB to the MC-integral on a 2D toy example (Section 6.1).
Revised: Firstly, we compare the LB to the MC-integral on a 2D toy example (Section 6.1).

Original: Secondly, we apply the same comparison to out-of-distribution (OOD) detection  (Section 6.2).
Revised: Secondly, we apply the same comparison to out-ofdistribution (OOD) detection in many settings (Section 6.2).

Original: Thirdly, we compare the commonly used probit approximation to the LB in section 6.3 Fourthly, we compare their computational cost and contextualize the speed-up for the prediction process in Section 6.4.
Revised: Thirdly, we compare the commonly used probit approximation to the LB in section 6.3 Fourthly, we compare their computational cost and contextualize the speed-up for the prediction process in Section 6.4.

Original: Finally, in Section 6.5, we present analysis on ImageNet [Russakovsky et al., 2014] to demonstrate the scalability of the LB and the advantage of having a full Dirichlet distribution over softmax outputs.
Revised: Finally, in Section 6.5, we present analysis on ImageNet [Russakovsky et al., 2014] to demonstrate the scalability of the LB and the advantage of having a full Dirichlet distribution over softmax outputs.

Original: We extended Laplace torch [Daxberger et al., 2021] for the experiments.
Revised: We extended Laplace torch [Daxberger et al., 2021] for the experiments.

Original: Code can be found in the accompanying GitHub repository.
Revised: Code can be found in the accompanying GitHub repository.

Original: For all experiments, a last-layer Laplace approximation has been applied.
Revised: For all experiments, a last-layer Laplace approximation has been applied.

Original: This scheme has been successfully used by
Revised: This scheme has been successfully used by

Original: Snoek et al.
Revised: Snoek et al.

Original: [2015], Kristiadi et al.
Revised: [2015], Kristiadi et al.

Original: We use diagonal and Kronecker-factorized (KFAC)[Ritter et al., 2018,
Revised: We use diagonal and Kronecker-factorized (KFAC)[Ritter et al., 2018,

Original: Martens and Grosse, 2015] approximations of the Hessian, since inverting the exact Hessian is too costly.
Revised: Martens and Grosse, 2015] approximations of the Hessian, since inverting the exact Hessian is too costly.

Original: A detailed mathematical explanation and setup of the experiments can be found in appendix D. While the LB could also be applied to different approximations of a Gaussian posterior predictive such as Variational Inference [Graves, 2011, Blundell et al., 2015], we used a Laplace approximation in our experiments to construct such an approximation.
Revised: A detailed mathematical explanation and setup of the experiments can be found in appendix D. While the LB could also be applied to different approximations of a Gaussian posterior predictive such as Variational Inference [Graves, 2011, Blundell et al., 2015], we used a Laplace approximation in our experiments to construct such an approximation.

Original: This is for two reasons: (i) it is one of the fastest ways to get a Gaussian posterior predictive and (ii) it can be applied to pre-trained networks which is especially useful for large problems such as ImageNet.
Revised: This is for two reasons: (i) it is one of the fastest ways to get a Gaussian posterior predictive and (ii) it can be applied to pre-trained networks which is especially useful for large problems such as ImageNet.

Original: Nevertheless, we want to emphasize again that the LB can be applied to any Gaussian over the outputs independent of the way it was generated.
Revised: Nevertheless, we want to emphasize again that the LB can be applied to any Gaussian over the outputs independent of the way it was generated.

Original: 6.1 2D TOY EXAMPLE
Revised: 6.1 2D TOY EXAMPLE

Original: We train a simple ReLU network on the 2D half-moon problems from scikit-learn [Pedregosa et al., 2011].
Revised: We train a simple ReLU network on the 2D half-moon problems from scikit-learn [Pedregosa et al., 2011].

Original: As can be seen in Figure 3 the MAP estimate and vanilla LB are overconfident for the reasons discussed in 4 but the normalized version yields a near-perfect fit.
Revised: As can be seen in Figure 3 the MAP estimate and vanilla LB are overconfident for the reasons discussed in 4 but the normalized version yields a near-perfect fit.

Original: 6.2 OOD DETECTION
Revised: 6.2 OOD DETECTION

Original: We compare the performance of the LB to the MC-integral (Diagonal and KFAC) on a standard OOD detection benchmark suite, to test whether the LB gives similar results to the MC sampling methods.
Revised: We compare the performance of the LB to the MC-integral (Diagonal and KFAC) on a standard OOD detection benchmark suite, to test whether the LB gives similar results to the MC sampling methods.

Original: Following prior literature, we use the standard expected calibration error (ECE) and area under the ROC-curve (AUROC) metrics
Revised: Following prior literature, we use the standard expected calibration error (ECE) and area under the ROC-curve (AUROC) metrics

Original: [Hendrycks and
Revised: [Hendrycks and

Original: Gimpel, 2016].
Revised: Gimpel, 2016].

Original: For the exact setup we refer the reader to Appendix D.
Revised: For the exact setup, we refer the reader to Appendix D.

Original: We use the mean of the Dirichlet to obtain a comparable approximation to the MC-integral.
Revised: We use the mean of the Dirichlet to obtain a comparable approximation to the MC-integral.

Original: The results are presented in Table 1.
Revised: The results are presented in Table 1.

Original: We find that the results of the LB or its normalized version are comparable throughout the entire benchmark suite.
Revised: We find that the results of the LB or its normalized version [Table] [Table] [Table]

Revised: [Table]

Revised: are comparable throughout the entire benchmark suite.

Original: Since the LB is much faster it can be a good replacement for MC in time-sensitive applications.
Revised: Since the LB is much faster it can be a good replacement for MC in time-sensitive applications.

Revised: Furthermore, we compare the LB to prior networks (PNs) in Table 4 since PNs also yield a Dirichlet distribution as an output on classification tasks.

Revised: We find that PNs outperform the LB in most cases.

Revised: However, we don’t think this is a major problem since they have different aims and use cases.

Revised: The LB creates a Dirichlet distribution on top of an already existing Gaussian model while PNs describe a training procedure and have to be trained from scratch.

Revised: Thus, the primary comparison for the LB should be against sampling and other integral approximations like in Table 2.

Original: 6.3 COMPARISON TO THE PROBIT APPROXIMATION
Revised: Lastly, we compare the LB for a full-layer vs. last-layer Laplace approximation of the network in Table 3. We find that, as expected, the full-layer setting yield slightly better results. However, since the primary advantage of the LB is its speed, we think the natural fit for it is a last-layer approximation.

Revised: [Figure]

Revised: i

Revised: [Figure]

Revised: e r g e n ce

Revised: MonteCarlo LaplaceBridge

Revised: [Figure]

Revised: i

Revised: [Figure]

Revised: e r g e n ce

Revised: [Table]

Original:  The multi-class probit approximation [Gibbs, 1997, Lu et al., 2020] is a commonly used approximation for the softmaxGaussian integral.
Revised: Fraction [Table] 6.3 COMPARISON TO THE PROBIT APPROXIMATION The multi-class probit approximation [Gibbs, 1997, Lu et al., 2020] is a commonly used approximation for the softmaxGaussian integral.

Original: We compare it to the diagonal normalized
Revised: We compare it to the diagonal normalized

Original: LB in Table 2.
Revised: LB in Table 2.

Original: We find that the LB norm outperforms the probit approximation in most OOD tasks.
Revised: We find that the LB norm outperforms the probit approximation in most OOD tasks.

Original: When we use a
Revised: When we use a

Original: KFAC approximation of the Hessian, this trend still holds (see Table 4 in Appendix D).[Table] [Table] [Figure]
Revised: KFAC approximation of the Hessian, this trend still holds (see Table 6 in Appendix D).

Original: [Figure]

Original: i v e r g e n

Original: [Figure]

Original: [Table]

Original: + MC

Original: [Table]

Original: + MC(100) + MC(10)

Original: + Laplace Bridge

Original: [Table]

Original: [Table]

Original: Fraction of overall time 0.06/0.38/0.83/0.

Original: [Table]

Original: 6.4 TIME COMPARISON
Revised: 6.4 TIME COMPARISON

Original: We compare the computational cost of the density-estimated p sample distribution via sampling and the Dirichlet obtained from the LB
Revised: We compare the computational cost of the density-estimated p sample distribution via sampling and the Dirichlet obtained from the LB

Original: p LB for approximating the true p true over MCsampling.
Revised: p LB for approximating the true p true over MCsampling.

Original: Different numbers of samples are drawn from the Gaussian, the softmax is applied and the KL-divergence between the histogram of the samples with the true distribution is computed.
Revised: Different numbers of samples are drawn from the Gaussian, the softmax is applied and the KL-divergence between the histogram of the samples with the true distribution is computed.

Original: We use KL-divergences D KL ( p true ∥ p sample ) and D KL ( p true ∥ p LB ) , respectively, to measure similarity between approximations and ground truth while the number of samples for p sample is increased exponentially.
Revised: We use KL-divergences D KL ( p true ∥ p sample ) and D KL ( p true ∥ p LB ) , respectively, to measure similarity between approximations and ground truth while the number of samples for p sample is increased exponentially.

Original: The true distribution p true is constructed via MC with 100 k samples.
Revised: The true distribution p true is constructed via MC with 100 k samples.

Original: The experiment is conducted for three different Gaussian distributions over R 3 .
Revised: The experiment is conducted for three different Gaussian distributions over R 3 .

Original: Since the softmax applied to a Gaussian does not have an analytic form, the algebraic calculation of the approximation error is not possible and an empirical evaluation via sampling is the best option.
Revised: Since the softmax applied to a Gaussian does not have an analytic form, the algebraic calculation of the approximation error is not possible and an empirical evaluation via sampling is the best option.

Original: The fact that there[Figure] is no analytic solution is part of the justification for using the LB in the first place.
Revised: The fact that there is no analytic solution is part of the justification for using the LB in the first place.

Original: Figure 5 suggests that the number of samples required such that the distribution p sample approximates the true distribution p true as good as the Dirichlet distribution obtained via the LB is large, i [Figure] e. somewhere between 750 and 10 k.
Revised: Figure 5 suggests that the number of samples required such that the distribution p sample approximates the true distribution p true as good as the Dirichlet distribution obtained via the LB is large, i.e. somewhere between 750 and 10 k.

Original: This translates to a wall-clock time advantage of at least a factor of 100 before sampling becomes competitive in quality with the LB.
Revised: This translates to a wall-clock time advantage of at least a factor of 100 before sampling becomes competitive in quality with the LB.

Original: To further demonstrate the low compute cost of the LB, we timed different parts of the process for our setup.
Revised: To further demonstrate the low compute cost of the LB, we timed different parts of the process for our setup.

Original: On our hardware and setup, training a ResNet-18 on CIFAR10 over 130 epochs takes 71 minutes and 30 seconds.
Revised: On our hardware and setup, training a ResNet-18 on CIFAR10 over 130 epochs takes 71 minutes and 30 seconds.

Original: Computing a
Revised: Computing a

Original: Hessian for the network from the training data can be done with B ACK PACK [Dangel et al., 2020] at the cost of one backward pass over the training data or around 29 seconds.
Revised: Hessian for the network from the training data can be done with B ACK PACK [Dangel et al., 2020] at the cost of one backward pass over the training data or around 29 seconds.

Original: This one additional backward pass is the only change to the training procedure compared to conventional training.
Revised: This one additional backward pass is the only change to the training procedure compared to conventional training.

Original: Since the LB only applies to the last step of the prediction pipeline, it is important to compare it to a forward pass through the rest of the network.
Revised: Since the LB only applies to the last step of the prediction pipeline, it is important to compare it to a forward pass through the rest of the network.

Original: Re-using the ResNet-18 and CIFAR setup we measure the time in seconds for a forward pass, for the application of the LB, and for the sampling procedure with 10, 100, and 1000 samples.
Revised: Re-using the ResNet-18 and CIFAR setup we measure the time in seconds for a forward pass, for the application of the LB, and for the sampling procedure with 10, 100, and 1000 samples.

Original: The resulting sum total time for the entire test set is given in Table 3.
Revised: The resulting sum total time for the entire test set is given in Table 5.

Original: We find that sampling takes up between 94% (for 1000 samples) and 17% (for 10 samples) of the entire prediction while the LB is only 4%.
Revised: We find that sampling takes up between 94% (for 1000 samples) and 17% (for 10 samples) of the entire prediction while the LB is only 4%.

Original: Thus, the acceleration through the LB is a significant improvement for the prediction process as a whole, not only for a part of the pipeline. 
Revised: Thus, the acceleration through the LB is a significant improvement for the prediction process as a whole, not only for a part of the pipeline. [Figure]

Original: 6.5 UNCERTAINTY-AWARE OUTPUT RANKING ON IMAGENET
Revised: 6.5 UNCERTAINTY-AWARE OUTPUT RANKING ON IMAGENET

Original: Due to the cost of sampling-based inference, classification on large datasets with many classes, like ImageNet, is rarely done in a Bayesian fashion.
Revised: Due to the cost of sampling-based inference [Figure] classification on large datasets with many classes, like ImageNet, is rarely done in a Bayesian fashion.

Original: Instead, models for such tasks are often compared along a top- k metric (e.g. k = 5 ).
Revised: Instead, models for such tasks are often compared along a top- k metric (e [Figure] g. k = 5 ).

Original: Although widely accepted, this metric has some pathologies:
Revised: Although widely accepted [Figure] this metric has some pathologies:

Original: Depending on how close the point predictions are relative to their uncertainty , the total number of likely class labels should be allowed to vary from case to case.
Revised: Depending on how close the point predictions are relative to their uncertainty , the total number of likely class labels should be allowed to vary from case to case.

Original: Figure 6 shows examples:
Revised: Figure 6 shows examples:

Original: In some cases (panel 2) the classifier is quite confident that the image in question belongs to one out of only two classes and all others are highly unlikely.
Revised: In some cases (panel 2) the classifier is quite confident that the image in question belongs to one out of only two classes and all others are highly unlikely.

Original: In others (e.g. panel 1), a larger set of hypotheses are all nearly equally probable.
Revised: In others (e.g. panel 1), a larger set of hypotheses are all nearly equally probable.

Original: The Laplace Bridge, in conjunction with the last-layer
Revised: The Laplace Bridge, in conjunction with the last-layer

Original: Laplace approximations, can be used to address this issue.
Revised: Laplace approximations, can be used to address this issue.

Original: To this end, the analytic properties of its Dirichlet prediction are particularly useful: Recall that the marginal distribution p ( π i , (cid:80) j ̸
Revised: To this end, the analytic properties of its Dirichlet prediction are particularly useful: Recall that the marginal distribution p ( π i , (cid:80) j ̸

Original: = i π j ) over each component of a Dirichlet relative to all other components is Beta( α i , (cid:80)
Revised: = i π j ) over each component of a Dirichlet relative to all other components is Beta( α i , (cid:80)

Original: j ̸
Revised: j ̸

Original: = i α j ) .
Revised: = i α j ) .

Original: We leverage this property to propose a simple uncertaintyaware top-
Revised: We leverage this property to propose a simple uncertaintyaware top-

Original: k decision rule inspired by statistical tests.
Revised: k decision rule inspired by statistical tests.

Original: In-

Original: Top-k

Original: C oun t 0. overlap 0. A cc u r ac y L i s t l e ng t h stead of keeping k fixed, it uses the model’s confidence to adapt k (pseudo-code in Algorithm 1).
Revised: Instead of keeping k fixed, it uses the model’s confidence to adapt k (pseudo-code in Algorithm 1).

Original: We begin by sorting the class predictions in order of their expected probability α i .
Revised: We begin by sorting the class predictions in order of their expected probability α i .

Original: Then we compute the Beta marginal of the most likely class.
Revised: Then we compute the Beta marginal of the most likely class.

Original: Now, we compute the overlap of the next marginal and add that class to the list iff the overlap is more than some threshold (e.g. 0.05).
Revised: Now, we compute the overlap of the next marginal and add that class to the list iff the overlap is more than some threshold (e.g. 0.05).

Original: Continuing in this fashion, the algorithm terminates with a finite value k ≤ K of “non-separated” top classes.
Revised: Continuing in this fashion, the algorithm terminates with a finite value k ≤ K of “non-separated” top classes.

Original: The intuition behind this rule is that, if any Beta density overlaps with the most likely one more than the threshold of, say, 5% , the classifier cannot confidently predict one class over the other.
Revised: The intuition behind this rule is that, if any Beta density overlaps with the most likely one more than the threshold of, say, 5% , the classifier cannot confidently predict one class over the other.

Original: Thus, all classes sufficiently overlapping with the top contender should be returned as the top estimates.
Revised: Thus, all classes sufficiently overlapping with the top contender should be returned as the top estimates.

Original: We evaluate this decision rule on the test set of ImageNet.
Revised: We evaluate this decision rule on the test set of ImageNet.

Original: The overlap is calculated through the inverse CDF 3 of the respective Beta marginals.
Revised: The overlap is calculated through the inverse CDF 3 of the respective Beta marginals.

Original: The original top- 1 accuracy of
Revised: The original top- 1 accuracy of

Original: DenseNet on ImageNet is 0 . 744 .
Revised: DenseNet on ImageNet is 0 . 744 .

Original: In contrast, the uncertaintyaware top-k method yields accuracies of over 0 .
Revised: In contrast, the uncertaintyaware top-k method yields accuracies of over 0 .

Original: 85 while average list lengths stay below 3 (see Figure 7).
Revised: 85 while average list lengths stay below 3 (see Figure 7).

Original: Furthermore, we find that most of the predictions given by the uncertainty-aware metric still yielded a top- 1 prediction.
Revised: Furthermore, we find that most of the predictions given by the uncertainty-aware metric still yielded a top- 1 prediction.

Original: This means that using uncertainty does not imply adding meaningless classes to the prediction.
Revised: This means that using uncertainty does not imply adding meaningless classes to the prediction.

Original: Furthermore, there are non-negligibly many cases where k equals to 2 , 3 , or(all values larger than 10 are in the 10 bin).
Revised: Furthermore, there are non-negligibly many cases where k equals to 2 , 3 , or(all values larger than 10 are in the 10 bin).

Original: Thus, using the uncertainty-aware prediction rule above, the classifier can use its uncertainty to adaptively return a longer or shorter list of predictions.
Revised: Thus, using the uncertainty-aware prediction rule above, the classifier can use its uncertainty to adaptively return a longer or shorter list of predictions.

Original: This not only allows it to improve accuracy over a hard top-1 threshold.
Revised: This not only allows it Top-k C oun t 0. overlap 0.

Original: Arguably, the ability to vary the size of the predicted set of classes is a practically useful functionality in itself.
Revised: A cc u r ac y

Original: As Figure
Revised: L i s t l e ng t h

Original: Algorithm
Revised: Algorithm

Original: -aware
Revised: -aware

Original: Input: A Dirichlet parameter α ∈ R K obtained by applying the LB to the Gaussian over the logit of an input, a percentile threshold T e.g. 0 .
Revised: Input: A Dirichlet parameter α ∈ R K obtained by applying the LB to the Gaussian over the logit of an input, a percentile threshold T e.g. 0 .

Original: 05 , a function class _ of that returns the underlying class of a sorted index.
Revised: 05 , a function class _ of that returns the underlying class of a sorted index.

Original: ˜ α
Revised: ˜ α

Original: = sort _ descending( α ) // start with the highest confidence α 

Revised: = sort _ descending( α ) // start with the highest confidence α

Original: = (cid:80)
Revised: = (cid:80)

Original: i
Revised: i

Original: α
Revised: α

Original: i
Revised: i

Original: C = { class _ of(1) } // initialize top-
Revised: C = { class _ of(1) } // initialize top-

Original: k , must include at least one class
Revised: k , must include at least one class

Original: F 1 =
Revised: F 1 =

Original: Beta(˜ α 1 , α 0 − ˜ α 1 ) //
Revised: Beta(˜ α 1 , α 0 − ˜ α 1 ) //

Original: the first marginal CDF l 1 = F − 1 1 ( T/ 2) // left T 2 percentile of the first marginal for i = 2 , . .
Revised: the first marginal CDF l 1 = F − 1 1 ( T/ 2) // left T 2 percentile of the first marginal for i = 2 , . .

Original: K do
Revised: K do

Original: F i =
Revised: F i =

Original: Beta(˜ α i , α 0 − ˜ α i ) //
Revised: Beta(˜ α i , α 0 − ˜ α i ) //

Original: the current marginal CDF r i = F − 1 i (1 − T/ 2) //
Revised: the current marginal CDF r i = F − 1 i (1 − T/ 2) //

Original: right T 2 percentile of the current marginal if r i > l 1 then
Revised: right T 2 percentile of the current marginal if r i > l 1 then

Original: C = C ∪ { class _ of( i ) } // overlap detected, add the current class else break // No more overlap, end the algorithm end if end for
Revised: C = C ∪ { class _ of( i ) } // overlap detected, add the current class elsebreak // No more overlap, end the algorithm end if end for

Original: Output: C // return the resulting top- 
Revised: Output: C // return the resulting top- k prediction to improve accuracy over a hard top-1 threshold.

Revised: Arguably, the ability to vary the size of the predicted set of classes is a practically useful functionality in itself.

Original: k prediction shows anecdotally, some of the labels (like “notebook” and “laptop”) are semantically so similar to each other that it would seem only natural for the classifier to use them synonymously.
Revised: As Figureshows anecdotally, some of the labels (like “notebook” and “laptop”) are semantically so similar to each other that it would seem only natural for the classifier to use them synonymously.

Original:  CONCLUSION
Revised:  CONCLUSION

Original: We have adapted a previously developed approximation scheme for new use in Bayesian Deep Learning.
Revised: We have adapted a previously developed approximation scheme for new use in Bayesian Deep Learning.

Original: Given a
Revised: Given a

Original: Gaussian approximation to the weight-space posterior of a
Revised: Gaussian approximation to the weight-space posterior of a

Original: NN (which can be constructed by various means, including another Laplace approximation), and an input, the Laplace
Revised: NN (which can be constructed by various means, including another Laplace approximation), and an input, the Laplace

Original: Bridge analytically maps the marginal Gaussian prediction on the logits onto a Dirichlet distribution over the softmax vectors.
Revised: Bridge analytically maps the marginal Gaussian prediction on the logits onto a Dirichlet distribution over the softmax vectors.

Original: The associated computational cost of O ( K ) for Kclass prediction compares favorably to that of MC sampling.
Revised: The associated computational cost of O ( K ) for Kclass prediction compares favorably to that of MC sampling.

Original: The proposed method empirically preserves predictive uncertainty, offering an attractive, low-cost, high-quality alternative to Monte Carlo sampling.
Revised: The proposed method empirically preserves predictive uncertainty, offering an attractive, low-cost, high-quality alternative to Monte Carlo sampling.

Original: In conjunction with a low-cost, last-layer Bayesian approximation, it is useful in real-time applications wherever uncertainty is required— especially because it drastically reduces the cost of predicting a posterior distribution at test time for a minimal increase in cost at training time.
Revised: In conjunction with a low-cost, last-layer Bayesian approximation, it is useful in real-time applications wherever uncertainty is required— especially because it drastically reduces the cost of predicting a posterior distribution at test time for a minimal increase in cost at training time.

Original: The vanilla LB has some limitations, for which we proposed a simple correction that outperforms alternative softmax-integral approximations such as the commonly used multi-class probit.
Revised: The vanilla LB has some limitations, for which we proposed a simple correction that outperforms alternative softmax-integral approximations such as the commonly used multi-class probit.

Original: We demonstrate the utility of the scheme for large-scale Bayesian inference by using it to construct an uncertainty-aware top- k ranking on ImageNet. 
Revised: We demonstrate the utility of the scheme for large-scale Bayesian inference by using it to construct an uncertainty-aware top- k ranking on ImageNet. Author Contributions

Revised: MH wrote the code, created the figures, and wrote most of the paper.

Revised: AK gave guidance and supervision and assisted throughout the entire process and greatly helped with the rebuttal.

Revised: PH had the original idea and provided supervision.

Revised: Acknowledgements

Original: CHANGE OF VARIABLE FOR PDF
Revised: The authors gratefully acknowledge financial support by the European Research Council through ERC StG Action 757275 / PANAMA; the DFG Cluster of Excellence “Machine Learning - New Perspectives for Science”, EXC 2064/1, project number 390727645; the German Federal Ministry of Education and Research (BMBF) through the Tübingen AI Center (FKZ: 01IS18039A); and funds from the Ministry of Science, Research and Arts of the State of Baden-Württemberg.

Revised: MH & AK are grateful to Alexander

Revised: Meinke for the pre-trained models and the International Max

Revised: Planck Research School for Intelligent Systems (IMPRS-IS) for support.

Revised: MH & AK would also like to thank all members of Methods of Machine Learning group for helpful feedback.

Revised: CHANGE OF VARIABLE FOR PDF

Original: Let z be an n -dimensional continuous random variable with joint density function p x .
Revised: Let z be an n -dimensional continuous random variable with joint density function p x .

Original: If y = G ( x ) , where G is a differentiable function, then y has density p y : 
Revised: If y = G ( x ) , where G is a differentiable function, then y has density p y : [Equation]

Revised: (cid:12)(cid:12)(cid:12)(cid:12)(cid:12) z

Original: [Equation]
Revised: [Equation]

Original:  where the differential is the Jacobian of the inverse of G evaluated at y .
Revised:  [Equation] (cid:12)(cid:12)(cid:12)(cid:12)(cid:12)(cid:12) (15) where the differential is the Jacobian of the inverse of G evaluated at y .

Original: This procedure, also known as ‘change of basis’, is at the core of the Laplace bridge since it is used to transform the Dirichlet into the softmax basis.
Revised: This procedure, also known as ‘change of basis’, is at the core of the Laplace bridge since it is used to transform the Dirichlet into the softmax basis.

Original: CORRECTION FOR SUM(Y)= 
Revised: CORRECTION FOR SUM(Y)=

Original: We know that the product rule of Gaussians yields
Revised: We know that the product rule of Gaussians yields

Original: [Equation]
Revised: [Equation]

Original: In our particular setup we have
Revised: In our particular setup we have

Original: [Equation] [Equation]
Revised: [Equation] [Equation]

Original: with constraint
Revised: with constraint

Original: [Equation] [Equation]
Revised: [Equation] [Equation]

Original: Therefore we get
Revised: Therefore we get

Original: [Equation]
Revised: [Equation]

Revised: VARIANCE CORRECTION

Revised: As described in the main text, the original Laplace Bridge scales worse with Σ than sampling and applying the softmax.

Revised: In Figure 8 you can see a contourplot that shows the scaling of mean and variance with and without correction.

Revised: [Figure]

Revised: As suggested, the Variance has nearly no influence on the result before the correction but our correction fixes that.

Revised: Some reviewers wanted to understand how we derived the equations for our correction, so here is a short informal explanation.

Revised: During the experimentation with the LB, we found that it doesn’t approximate the sample distribution well when Σ gets large.

Revised: We then understood why (as detailed in the limitations section) and proposed a fix for these scenarios without damaging its behavior in all other scenarios.

Revised: We experimented with multiple fixes and the result you see in the paper is the one that fulfilled most of our criteria.

Revised: Therefore, the correction doesn’t come from a principled theoretical derivation but is motivated by the theoretical findings.

Original: B APPENDIX (DERIVATION OF LB)
Revised: B APPENDIX (DERIVATION OF LB)

Original: Assume we have a Dirichlet in the standard basis with parameter vector α and probability density function:
Revised: Assume we have a Dirichlet in the standard basis with parameter vector α and probability density function:

Original: [Equation]
Revised: [Equation]

Original: We aim to transform the basis of this distribution via the softmax transform to be in the new base π : 
Revised: We aim to transform the basis of this distribution via the softmax transform to be in the new base π : π

Original: [Equation]
Revised: [Equation]

Original: Usually, to transform the basis we would need the inverse transformation H − 1 ( z ) as described in the main paper.
Revised: Usually, to transform the basis we would need the inverse transformation H − 1 ( z ) as described in the main paper.

Original: However, the softmax does not have an analytic inverse.
Revised: However, the softmax does not have an analytic inverse.

Original: Therefore David JC MacKay uses the following trick.
Revised: Therefore David JC MacKay uses the following trick.

Original: Assume we know that the distribution in the transformed basis is:
Revised: Assume we know that the distribution in the transformed basis is:

Original: [Equation]
Revised: [Equation]

Original: then we can show that the original distribution is the result of the basis transform by the softmax.
Revised: then we can show that the original distribution is the result of the basis transform by the softmax.

Original: The Dirichlet in the softmax basis: We show that the density over π shown in Equation 24 transforms into the
Revised: The Dirichlet in the softmax basis: We show that the density over π shown in Equation 24 transforms into the

Original: Dirichlet over z .
Revised: Dirichlet over z .

Original: First, we consider the special case where π is confined to an I − 1 dimensional subspace satisfying (cid:80) i
Revised: First, we consider the special case where π is confined to an I − 1 dimensional subspace satisfying (cid:80) i

Original: π
Revised: π

Original: i = c .
Revised: i = c .

Original: In this subspace we can represent φ by an I −dimensional vector φ such that
Revised: In this subspace we can represent φ by an I −dimensional vector φ such that

Original: [Equation]
Revised: [Equation]

Original: and similarly we can represent z by an I − 1 dimensional vector a :
Revised: and similarly we can represent z by an I − 1 dimensional vector a :

Original: [Equation]
Revised: [Equation]

Original: then we can find the density over a (which is proportional to the required density over z ) from the density over φ (which is proportional to the given density over π ) by finding the determinant of the ( I − 1) ×
Revised: then we can find the density over a (which is proportional to the required density over z ) from the density over φ (which is proportional to the given density over π ) by finding the determinant of the ( I − 1) ×

Original: ( I − 1) Jacobian J given by
Revised: ( I − 1) Jacobian J given by

Original: [Equation]
Revised: [Equation]

Original: We define two additional I − 1 dimensional helper vectors z + k :=
Revised: We define two additional I − 1 dimensional helper vectors z + k :=

Original: z k − z
Revised: z k − z

Original: I and n k := 1 , and use det( I − xy T )
Revised: I and n k := 1 , and use det( I − xy T )

Original: = 1 − x · y from linear algebra.
Revised: = 1 − x · y from linear algebra.

Original: It follows that
Revised: It follows that

Original: [Equation]
Revised: [Equation]

Revised: i =

Revised: [Equation]

Original: Therefore, using Equation 24 we find that
Revised: Therefore, using Equation 24 we find that

Original: [Equation]
Revised: [Equation]

Original: This result is true for any constant c since it can be put into the normalizing constant.
Revised: This result is true for any constant c since it can be put into the normalizing constant.

Original: Thereby we make sure that the integral of the distribution is 

Revised: Thereby we make sure that the integral of the distribution is

Original: and we have a valid probability distribution.
Revised: and we have a valid probability distribution.

Original: C APPENDIX (DERIVATION OF INVERSION)
Revised: C APPENDIX (DERIVATION OF INVERSION)

Original: Through the figures of the 1D Dirichlet approximation in the main paper we have already established that the mode of the Dirichlet lies at the mean of the Gaussian distribution and therefore π ( y ) =
Revised: Through the figures of the 1D Dirichlet approximation in the main paper we have already established that the mode of the Dirichlet lies at the mean of the Gaussian distribution and therefore π ( y ) =

Original: α (cid:80) i α i .
Revised: [Equation] i α i .

Original: Additionally, the elements of y must sum to zero.
Revised: Additionally, the elements of y must sum to zero.

Original: These two constraints combined yield only one possible solution for µ .
Revised: These two constraints combined yield only one possible solution for µ .

Original: [Equation]
Revised: [Equation]

Original: Calculating the covariance matrix Σ is more complicated but layed out in the following.
Revised: Calculating the covariance matrix Σ is more complicated but layed out in the following.

Original: The logarithm of the Dirichlet is, up to additive constants
Revised: The logarithm of the Dirichlet is, up to additive constants

Original: [Equation]
Revised: [Equation]

Original: Using π k as the softmax of y as shown in Equation 23 we can find the elements of the Hessian L
Revised: Using π k as the softmax of y as shown in Equation 23 we can find the elements of the Hessian L

Original: [Equation]
Revised: [Equation]

Original:  where ˆ α := (cid:80) k
Revised: ) where ˆ α := (cid:80) k

Original: α k and ˆ π = α k
Revised: α k and ˆ π = α k

Original: ˆ α for the value of π at the mode.
Revised: ˆ α for the value of π at the mode.

Original: Analytically inverting L is done via a lengthy derivation using the fact that we can write L =
Revised: Analytically inverting L is done via a lengthy derivation using the fact that we can write L =

Original: A + XBX ⊤ and inverting it with the Schur-complement.
Revised: A + XBX ⊤ and inverting it with the Schur-complement.

Revised: You can find the derivation in [Hennig, 2010].

Original: This process results in the inverse of the Hessian
Revised: This process results in the inverse of the Hessian

Original: [Equation]
Revised: [Equation]

Original: We are mostly interested in the diagonal elements, since we desire a sparse encoding for computational reasons and we otherwise needed to map a K × K covariance matrix to a K × 1 Dirichlet parameter vector which would be a very overdetermined mapping.
Revised: We are mostly interested in the diagonal elements, since we desire a sparse encoding for computational reasons and we otherwise needed to map a K × K covariance matrix to a K × 1 Dirichlet parameter vector which would be a very overdetermined mapping.

Original: Note that K is a scalar not a matrix.
Revised: Note that K is a scalar not a matrix.

Original: The diagonal elements of Σ =
Revised: The diagonal elements of Σ =

Original: L − 1 can be calculated as
Revised: L − 1 can be calculated as

Original: [Equation]
Revised: [Equation]

Original: To invert this mapping we transform Equation 32 to
Revised: To invert this mapping we transform Equation 32 to

Original: [Equation]
Revised: [Equation]

Original: by applying the logarithm and re-ordering some parts.
Revised: by applying the logarithm and re-ordering some parts.

Original: Inserting this into Equation 36 and re-arranging yields
Revised: Inserting this into Equation 36 and re-arranging yields

Revised: K (cid:89) l

Original: [Equation]
Revised: [Equation]

Original: which can be re-inserted into Equation 37 to give
Revised: which can be re-inserted into Equation 37 to give

Original: [Equation]
Revised: [Equation]

Original: which is the final mapping.
Revised: which is the final mapping.

Original: With Equations 32 and 36 we are able to map from Dirichlet to Gaussian and with Equation 39 we are able to map the inverse direction.
Revised: With Equations 32 and 36 we are able to map from Dirichlet to Gaussian and with Equation 39 we are able to map the inverse direction.

Original: D APPENDIX (EXPERIMENTAL DETAILS)
Revised: D APPENDIX (EXPERIMENTAL DETAILS)

Original: The exact experimental setups, i.e. network architectures, learning rates, random seeds, etc. can be found in the accompanying GitHub repository 4 .
Revised: The exact experimental setups, i.e. network architectures, learning rates, random seeds, etc. can be found in the accompanying GitHub repository 4 .

Original: This section is used to justify some of the decisions we made during the process in more detail, highlight some miscellaneous interesting things and showcase the additional experiments promised in the main paper.[Table]
Revised: This section is used to justify some of the decisions we made during the process in more detail, highlight some miscellaneous interesting things and showcase the additional experiments promised in the main paper.

Original: MATHEMATICAL DESCRIPTION OF THE SETUP
Revised: MATHEMATICAL DESCRIPTION OF THE SETUP

Original: In principle, the Gaussian over the weights required by the Laplace Bridge for BNNs can be constructed by any
Revised: In principle, the Gaussian over the weights required by the Laplace Bridge for BNNs can be constructed by any

Original: Gaussian approximate Bayesian method such as variational
Revised: Gaussian approximate Bayesian method such as variational

Original: Bayes [Graves, 2011, Blundell et al., 2015] and Laplace approximations for NNs [MacKay, 1992a, Ritter et al., 2018].
Revised: Bayes [Graves, 2011, Blundell et al., 2015] and Laplace approximations for NNs [MacKay, 1992a, Ritter et al., 2018].

Original: We will focus on the Laplace approximation, which uses the same principle as the Laplace Bridge.
Revised: We will focus on the Laplace approximation, which uses the same principle as the Laplace Bridge.

Original: However, in the
Revised: However, in the

Original: Laplace approximation for neural networks, the posterior distribution over the weights of a network is the one that is approximated as a Gaussian, instead of a Dirichlet distribution over the outputs as in the Laplace Bridge.
Revised: Laplace approximation for neural networks, the posterior distribution over the weights of a network is the one that is approximated as a Gaussian, instead of a Dirichlet distribution over the outputs as in the Laplace Bridge.

Original: Given a dataset D := { ( x i , t i ) } Di =1 and a prior p ( θ ) , let p( θ |D
Revised: Given a dataset D := { ( x i , t i ) } Di =1 and a prior p ( θ ) , let p

Original: [Equation]
Revised: [Equation]

Original: ∝ p( θ [Equation] p ( [Equation]
Revised: p

Original: |

Original: [Equation]
Revised: [Equation]

Original: ) = p

Original: [Equation]

Original: [Equation]

Original: be the posterior over the parameter θ of an L -layer network f θ .
Revised: be the posterior over the parameter θ of an L -layer network f θ .

Original: Then we can get an approximation of the posterior p ( θ |D ) by fitting a Gaussian N ( θ | µ θ , Σ θ ) where
Revised: Then we can get an approximation of the posterior p ( θ |D ) by fitting a Gaussian N ( θ | µ θ , Σ θ ) where

Original: [Equation]
Revised: [Equation]

Original: That is, we fit a Gaussian centered at the mode θ MAP of p ( θ |D ) with the covariance determined by the curvature at that point.
Revised: That is, we fit a Gaussian centered at the mode θ MAP of p ( θ |D ) with the covariance determined by the curvature at that point.

Original: We assume that the prior p ( θ ) is a zero-mean isotropic Gaussian N ( θ | 0 , σ 2 I ) and the likelihood function is the Categorical density
Revised: We assume that the prior p ( θ ) is a zero-mean isotropic Gaussian N ( θ | 0 , σ 2 I ) and the likelihood function is the Categorical density

Original: [Equation]
Revised: [Equation]

Original: For various applications in Deep Learning, an approximation with full Hessian is often computationally too expensive.
Revised: For various applications in Deep Learning, an approximation with full Hessian is often computationally too expensive.

Original: Indeed, for each input x ∈ R N , one has to do K backward passes to compute the Jacobian J ( x ) .
Revised: Indeed, for each input x ∈ R N , one has to do K backward passes to compute the Jacobian J ( x ) .

Original: Moreover, it requires an O ( PK ) storage which is also expensive since P is often in the order of millions.
Revised: Moreover, it requires an O ( PK ) storage which is also expensive since P is often in the order of millions.

Original: A cheaper alternative is to fix all but the last layer of f θ and only apply the Laplace approximation on W L , the last layer’s weight matrix.
Revised: A cheaper alternative is to fix all but the last layer of f θ and only apply the Laplace approximation on W L , the last layer’s weight matrix.

Original: This scheme has been used successfully by Snoek et al.
Revised: This scheme has been used successfully by Snoek et al.

Original: [2015], Wilson et al.
Revised: [2015], Wilson et al.

Original: [2016], Brosse et al.
Revised: [2016], Brosse et al.

Original: [2020], etc.
Revised: [2020], etc.

Original: and has been shown theoretically that it can mitigate overconfidence problems in
Revised: and has been shown theoretically that it can mitigate overconfidence problems in

Original: ReLU networks [Kristiadi et al., 2020].
Revised: ReLU networks [Kristiadi et al., 2020].

Original: In this case, given the approximate last-layer posterior
Revised: In this case, given the approximate last-layer posterior

Original: [Equation]
Revised: [Equation]

Original: one can efficiently compute the distribution over the logits.
Revised: one can efficiently compute the distribution over the logits.

Original: That is, let ϕ : R N → R Q be the first L − 1 layers of f θ , seen as a feature map.
Revised: That is, let ϕ : R N → R Q be the first L − 1 layers of f θ , seen as a feature map.

Original: Then, for each x ∈ R N , the induced distribution over the logit W L ϕ ( x ) =:
Revised: Then, for each x ∈ R N , the induced distribution over the logit W L ϕ ( x ) =:

Original: z is given by
Revised: z is given by

Original: [Equation]
Revised: [Equation]

Original: where ⊗ denotes the Kronecker product.
Revised: where ⊗ denotes the Kronecker product.

Original: An even more efficient last-layer approximation can be obtained using a Kronecker-factored matrix normal distribution [Louizos and Welling, 2016, Sun et al., 2017, Ritter et al., 2018].
Revised: An even more efficient last-layer approximation can be obtained using a Kronecker-factored matrix normal distribution [Louizos and Welling, 2016, Sun et al., 2017, Ritter et al., 2018].

Original: That is, we assume the posterior distribution to be
Revised: That is, we assume the posterior distribution to be

Original: [Equation]
Revised: [Equation]

Original: where U ∈ R K × K and V ∈ R Q × Q are the Kronecker factorization of the inverse Hessian matrix H − 1 W L [Martens and Grosse, 2015].
Revised: where U ∈ R K × K and V ∈ R Q × Q are the Kronecker factorization of the inverse Hessian matrix H − 1 W L [Martens and Grosse, 2015] and MN denotes the Matrix Normal [Table]

Original:  In this case, for any x ∈ R N , one can easily show that the distribution over logits is given by
Revised: distribution. In this case, for any x ∈ R N , one can easily show that the distribution over logits is given by

Original: [Equation]
Revised: [Equation]

Original: which is easy to implement and computationally cheap.
Revised: which is easy to implement and computationally cheap.

Revised: Finally

Revised: [Table]

Original: Finally, and even more efficient, is a last-layer approximation scheme with a diagonal Gaussian approximate posterior, i.e. the so-called mean-field approximation.
Revised: and even more efficient, is a last-layer approximation scheme with a diagonal Gaussian approximate posterior, i [Table] e. the so-called mean-field approximation.

Original: In this case, we assume the posterior distribution to be
Revised: In this case, we assume the posterior distribution to be

Original: [Equation]
Revised: [Equation]

Original: where σ 2 is obtained via the diagonal of the Hessian of the log-posterior w.r.t.
Revised: where σ 2 is obtained via the diagonal of the Hessian of the log-posterior w.r.t.

Original: vec ( W L ) at vec ( W L MAP ) .
Revised: vec ( W L ) at vec ( W L MAP ) .

Original: OOD DETECTION
Revised: OOD DETECTION

Original: The test scenarios are: A two-layer convolutional network trained on the MNIST dataset [LeCun, 1998].
Revised: The test scenarios are: A two-layer convolutional network trained on the MNIST dataset [LeCun, 1998].

Original: The OOD datasets for this case are FMNIST [Xiao et al., 2017], notMNIST [Bulatov, 2011], and KMNIST [Clanuwat et al., 2018].
Revised: The OOD datasets for this case are FMNIST [Xiao et al., 2017], notMNIST [Bulatov, 2011], and KMNIST [Clanuwat et al., 2018].

Original: For larger datasets, i.e. CIFAR-10 [Krizhevsky et al., 2014],
Revised: For larger datasets, i.e. CIFAR-10 [Krizhevsky et al., 2014],

Original: SVHN [Netzer et al., 2011], and CIFAR-100 [Krizhevsky et al., 2014], we use a ResNet-18 network [He et al., 2016].
Revised: SVHN [Netzer et al., 2011], and CIFAR-100 [Krizhevsky et al., 2014], we use a ResNet-18 network [He et al., 2016].

Original: In all scenarios, the networks are well-trained with 99% test accuracy on MNIST, 95 . 4% on CIFAR-10, 76 .
Revised: In all scenarios, the networks are well-trained with 99% test accuracy on MNIST, 95 . 4% on CIFAR-10, 76 .

Original: 6% on
Revised: 6% on

Original: CIFAR-100, and 100% on SVHN.
Revised: CIFAR-100, and 100% on SVHN.

Original: For the sampling baseline, we use 100 posterior samples.
Revised: For the sampling baseline, we use 100 posterior samples.

Revised: All network have been trained with conventional setups, i.e. we use ADAM with learning rate 1 e − 3 and weight decay 5 e − 4 for the MNIST experiments and SGD with a cosine annealing scheduler starting at learning rate 0 . 1 and momentum 0 .

Revised: 9 for the CIFAR and SVHN experiments.

Original: PROBIT VS LB
Revised: PROBIT VS LB

Original: The KFAC setting of the probit comparison can be found in Table 4.
Revised: The KFAC setting of the probit comparison can be found in Table 6.

Original: Especially in the MNIST scenario the probit approximation seems to break down since even in-dist detection is at chance level.
Revised: Especially in the MNIST scenario the probit approximation seems to break down since even in-dist detection is at chance level.

Original: The LB, on the other hand, yields reasonable results.
Revised: The LB, on the other hand, yields reasonable results.