Source-linked AI summary

Fast Predictive Uncertainty for Classification with Bayesian Deep Networks

Marius Hobbhahn, Agustinus Kristiadi, Philipp Hennig

arXiv:2003.01227v4cs.LGstat.ML

TL;DR

Point softmax predictions are overconfident, and Bayesian predictive uncertainty is costly to estimate by Monte Carlo. The paper adapts the Laplace Bridge to map Gaussian logits to a Dirichlet approximation, with a correction for high-variance cases. The resulting method is reported to preserve predictive uncertainty with O(K) computation and is demonstrated in an ImageNet ranking use case.

  • Problem

    Point softmax outputs are overconfident, while Bayesian predictive distributions over softmax outputs are typically estimated by costly Monte Carlo sampling.

    Source: §1 · p. 1
  • Method

    The paper analytically maps Gaussian logits to a Dirichlet approximation and adds a normalization correction for high-variance cases.

    Source: §4 · p. 4, §7 · p. 9
  • Results

    O(K) construction preserves predictive uncertainty as a low-cost alternative to Monte Carlo, and the paper demonstrates uncertainty-aware output ranking on ImageNet.

    Source: §1 · p. 2, §7 · p. 9
  • Takeaways & Limitations

    With a low-cost last-layer Bayesian approximation, the method is useful for real-time applications requiring uncertainty.

    Source: §7 · p. 9

    The variance correction lacks a principled theoretical derivation, and the Laplace approximation is suboptimal in extremely sparse, high-variance cases.

    Source: §4 · p. 4, VARIANCE CORRECTION · p. 12

Abstract

from arXiv · show

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. This is costly. 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. Importantly, the vanilla Laplace Bridge comes with certain limitations. We analyze those and suggest a simple solution that compares favorably to other commonly used estimates of the softmax-Gaussian integral. 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. We further demonstrate the usefulness of this Dirichlet approximation by using it to construct a lightweight uncertainty-aware output ranking for ImageNet.

1 INTRODUCTION

Standard softmax point predictions are overconfident, while Bayesian predictive uncertainty is costly because the Gaussian-logit output is typically estimated by Monte Carlo. The paper adapts the Laplace Bridge to approximate that output with a Dirichlet and evaluates its practical use, including ImageNet ranking.

1 INTRODUCTION

  • Bayesian predictive distributions are typically approximated by Monte Carlo because the Gaussian approximate posterior over weights induces a non-analytic predictive distribution over softmax outputs.Multiple samples make prediction in Bayesian neural networks comparatively expensive.Source: §1 · p. 1
  • The Laplace Bridge maps a Gaussian distribution over logits to an approximate Dirichlet distribution over softmax outputs.The paper revisits an idea originally proposed by MacKay [1998] and calls the inverse map the Laplace Bridge [Hennig et al., 2012].Source: §1 · p. 1
  • The paper analyzes vanilla Laplace Bridge limitations, proposes a solution, evaluates approximation quality, speed, and out-of-distribution performance, and demonstrates uncertainty-aware output ranking on ImageNet.The ImageNet use case leverages Dirichlet uncertainty to improve the popular top-k metric.Source: §1 · p. 2

2 THE LAPLACE BRIDGE

The Laplace Bridge changes basis to make a Dirichlet amenable to a Gaussian Laplace approximation, then uses a pseudo-inverse to map Gaussian parameters back to Dirichlet parameters. For Bayesian deep learning, the paper uses this inverse map despite its reduced expressiveness.

2 THE LAPLACE BRIDGE

  • A Laplace approximation sets its mean to a mode and its covariance to the inverse negative Hessian, and is suited to unimodal distributions on real vector spaces.The Hessian at the mode must be positive definite.Source: §2 · p. 2
  • Because a Dirichlet is defined on the simplex and may diverge at simplex corners, the method transforms it through the softmax into logit space before applying a Laplace approximation.The transformed density includes the absolute Jacobian determinant of the change of basis.Source: §2 · p. 2
  • The transformed Dirichlet density can be approximated by a Gaussian, giving an analytic map from Dirichlet parameters α to Gaussian parameters µ and Σ.This is the forward change-of-basis map described in the section.Source: §2 · p. 2
  • A pseudo-inverse maps Gaussian parameters back to Dirichlet parameters; for Bayesian deep learning, the paper uses this direction despite reduced expressiveness, reporting sufficient accuracy.The Gaussian parameter space is larger than the Dirichlet space and is not fully identified by the transformation.Source: §2 · p. 2

3 THE LAPLACE BRIDGE FOR BNNS

For Bayesian neural-network classification, the Laplace Bridge converts a Gaussian logit prediction into a Dirichlet over softmax outputs. This provides an O(K) construction and analytic distributional operations, rather than relying on repeated Monte Carlo samples.

3 THE LAPLACE BRIDGE FOR BNNS

  • The method uses a last-layer Laplace approximation of the network, which supports fast inference and is presented as a natural fit for the speed-oriented bridge.The paper uses diagonal and Kronecker approximations to the Hessian.Source: §3 · pp. 2–3
  • Constructing the K Dirichlet parameters costs O(K), whereas Monte Carlo costs O(MJ) and has sampling error O(1/M); the bridge is faster than one Monte Carlo sample while returning a full distribution.Here M is the number of samples and J is the cost of sampling from the Gaussian logit distribution.Source: §3 · p. 3
  • The Dirichlet mean gives a point estimate of the posterior predictive distribution, and Dirichlet marginals provide closed-form distributions for subsets of outputs.These properties support uncertainty queries about subsets of classes, beyond a point estimate of the softmax output.Source: §3 · p. 3

4 LIMITATIONS OF THE LAPLACE BRIDGE

The vanilla Laplace Bridge has a zero-sum constraint and can fit poorly in sparse, high-variance regions. The proposed normalization addresses these cases, bringing its predictions closer to Monte Carlo estimates.

4 LIMITATIONS OF THE LAPLACE BRIDGE

  • The Laplace Bridge assumes its Gaussian random variable sums to zero, so arbitrary Gaussian inputs require a projection onto that constraint.The correction uses the one-vector of size K.Source: §4 · p. 3
  • In extremely sparse, high-variance cases, the softmax-Dirichlet becomes asymmetric and the Gaussian Laplace approximation is suboptimal.This occurs where the logistic transform is nearly flat.Source: §4 · pp. 3–4
  • The corrected Laplace Bridge brings its approximation close to Monte Carlo predictions where the vanilla version is overconfident far from data.Figure 3 compares the vanilla and corrected versions with softmax applied to Gaussian samples.Source: Figure 3 · p. 4
  • As variance and the mean grow, the LB’s α is dominated by the mean’s exponential influence, while Monte Carlo results are driven by variance amplified through softmax.The LB incorporates Σ linearly and μ through exp(μ_k), so its variance contribution cannot sufficiently correct the difference.Source: §4 · p. 4
  • The additional normalization rescales the Gaussian after the zero-sum correction, pulling it into a lower-variance, more symmetric region; the paper calls it LB-norm.Figure 4 depicts the lower-variance projection before applying the LB.Source: §4 · pp. 4–5, Figure 4 · p. 4
  • The correction is motivated by experiments and theoretical insights, but has no derivation from first principles.Source: §4 · pp. 4–5

5 RELATED WORK

The related work describes the open challenge of analytically approximating the multiclass softmax under a Gaussian measure. The Laplace Bridge addresses it by approximating the full distribution over network softmax outputs without additional training.

5 RELATED WORK

  • For multiclass classification, a good analytic approximation of the Gaussian-measured expected softmax remains an open problem; the LB approximates this integral and the full softmax-output distribution.Source: §5 · p. 5
  • Prior Dirichlet-output approaches learn the distribution directly for non-Bayesian networks, while some distillation methods produce a similar Dirichlet from larger models.Source: §5 · p. 5
  • Unlike prior networks and distillation, the LB requires no additional training procedure and can be applied to pre-trained weights on large datasets.Source: §5 · p. 5

6 EXPERIMENTS

The experiments use last-layer Laplace approximations and compare the Laplace Bridge (LB) with Monte Carlo and other approximations across toy, OOD, runtime, and ImageNet settings. They report comparable OOD results to Monte Carlo, faster computation, and ImageNet uncertainty-aware top-k accuracy above 0.85 with average lists under three.

6.1 2D TOY EXAMPLE

  • On the 2D half-moon task, the normalized LB yields a near-perfect fit, unlike the overconfident MAP estimate and vanilla LB.Source: §6.1 · p. 5

6.2 OOD DETECTION

  • Across the OOD benchmark suite, LB or normalized LB gives results comparable to 100-sample Monte Carlo estimates while being much faster.Source: §6.2 · pp. 5–6, Table 1 · p. 6

6.3 COMPARISON TO THE PROBIT APPROXIMATION

  • The normalized LB outperforms the multi-class probit approximation on most OOD tasks, including when using a KFAC Hessian approximation.Source: Table 2 · p. 6, §6.3 · p. 7

6.2 OOD DETECTION

6.4 TIME COMPARISON

  • At least 100 times more wall-clock time is needed for sampling to match LB quality, while LB takes 4% of prediction time versus sampling's 17–94%.For ResNet-18 on CIFAR-10, sampling accounts for 17% of total prediction time with 10 samples and 94% with 1,000; LB takes 4%.Source: §6.4 · p. 7, Figure 5 · p. 7, Table 5 · p. 7

6.5 UNCERTAINTY-AWARE OUTPUT RANKING ON IMAGENET

  • 0.744 original DenseNet top-1 accuracy contrasts with uncertainty-aware top-k accuracies over 0.85 while average prediction lists stay below 3.Figure 7 shows the distribution of prediction-list lengths, with lists over 10 grouped in the 10 bin.Source: §6.5 · p. 8, Figure 7 · p. 9

7 CONCLUSION

The Laplace Bridge maps a Gaussian prediction on neural-network logits to a Dirichlet distribution over softmax vectors, providing a low-cost alternative to Monte Carlo sampling. The paper applies this approximation to uncertainty-aware ImageNet ranking and addresses limitations with a correction.

7 CONCLUSION

  • The Laplace Bridge constructs a Dirichlet approximation in O(K) for K-class prediction, offers a low-cost alternative to Monte Carlo sampling, and supports uncertainty-aware top-k ranking on ImageNet.The conclusion says the correction outperforms alternatives including multi-class probit and that the method is useful for large-scale Bayesian inference.Source: §7 · p. 9

CORRECTION FOR SUM(Y)=0

  • The conclusion states that a simple correction addresses limitations of the vanilla Laplace Bridge and that the resulting method supports uncertainty-aware ImageNet ranking.The supplied passages do not give further details of the correction-for-sum derivation.Source: §7 · p. 9

CHANGE OF VARIABLE FOR PDF

  • A change-of-basis procedure transforms the Dirichlet distribution into the softmax basis using the Jacobian of the inverse transformation.The paper identifies this transformation as central to the Laplace Bridge.Source: CHANGE OF VARIABLE FOR PDF · p. 12

VARIANCE CORRECTION

The vanilla Laplace Bridge scales worse with covariance than sampling, with variance having little influence before correction. The proposed correction restores variance influence, though it was selected experimentally rather than derived from a principled theory.

VARIANCE CORRECTION

  • The vanilla Laplace Bridge scales worse with Σ than sampling and poorly approximates the sample distribution when Σ becomes large.In this regime, its resulting α is dominated by the mean, while Monte Carlo results are mostly determined by the large variance amplified through softmax.Source: VARIANCE CORRECTION · p. 12
  • The correction restores variance influence in the scaling behavior that was nearly absent with the vanilla Laplace Bridge.Figure 8 contrasts the uncorrected and corrected contours, with levels showing the first entry of α on a log scale.Source: VARIANCE CORRECTION · p. 12, Figure 8 · p. 12
  • The correction was chosen after comparing multiple fixes and was motivated by theoretical findings, not a principled theoretical derivation.The authors report that it did not damage behavior in other scenarios.Source: VARIANCE CORRECTION · p. 12

B APPENDIX (DERIVATION OF LB)

The appendix derives the Laplace Bridge through a change of basis from a Dirichlet distribution to the softmax basis. The derivation uses a Jacobian to obtain the transformed density and normalizes it as a valid probability distribution.

B APPENDIX (DERIVATION OF LB)

  • The derivation starts with a Dirichlet in the standard basis and transforms its density to the softmax basis.Because softmax has no analytic inverse, the derivation uses a proposed transformed-basis density to establish the corresponding original density.Source: B APPENDIX (DERIVATION OF LB) · pp. 12–13
  • For the special case where the components sum to a constant, the appendix represents the variables in a lower-dimensional subspace and obtains the transformed density using a Jacobian determinant.The determinant identity det(I−xyT)=1−x·y is used in the derivation.Source: B APPENDIX (DERIVATION OF LB) · p. 13
  • The resulting density is normalized so that its integral is one, making it a valid probability distribution.The derivation notes that a constant can be absorbed into the normalizing constant.Source: B APPENDIX (DERIVATION OF LB) · p. 13

C APPENDIX (DERIVATION OF INVERSION)

The appendix derives the Gaussian parameters associated with a Dirichlet approximation and then inverts the mapping. It obtains the mean from the mode subject to constraints and computes covariance through the inverse Hessian, focusing on its diagonal for a sparse encoding.

C APPENDIX (DERIVATION OF INVERSION)

  • The mean is determined from the Dirichlet mode together with the requirement that the elements of y sum to zero.The two constraints yield only one possible solution for μ.Source: C APPENDIX (DERIVATION OF INVERSION) · p. 13
  • The covariance is derived from the Hessian of the log-Dirichlet, with its inverse obtained using a Schur complement.The appendix points to Hennig (2010) for the inversion derivation.Source: C APPENDIX (DERIVATION OF INVERSION) · p. 13
  • The derivation focuses on diagonal covariance elements to obtain a sparse encoding, avoiding an overdetermined mapping from a K × K covariance matrix to K Dirichlet parameters.The appendix states that K is a scalar, not a matrix.Source: C APPENDIX (DERIVATION OF INVERSION) · pp. 13–14
  • The appendix rearranges the equations to obtain the inverse direction of the mapping, completing the Gaussian–Dirichlet parameter conversion.It identifies the preceding equations as the mappings in each direction.Source: C APPENDIX (DERIVATION OF INVERSION) · p. 14

D APPENDIX (EXPERIMENTAL DETAILS)

The appendix records experimental details and supplements the main paper with methodological rationale and additional experiments. It outlines Gaussian approximations used for logits and reports OOD and probit-comparison experiments.

D APPENDIX (EXPERIMENTAL DETAILS)

MATHEMATICAL DESCRIPTION OF THE SETUP

  • Unlike the Laplace Bridge’s output-space Dirichlet approximation, the Laplace approximation fits a Gaussian over weights centered at the posterior mode, with curvature setting covariance.The setup assumes a zero-mean isotropic Gaussian prior and a Categorical likelihood.Source: MATHEMATICAL DESCRIPTION OF THE SETUP · p. 14
  • A full Hessian requires K backward passes per input and O(PK) storage; fixing all but the last layer is cheaper, and KFAC offers an even more efficient matrix-normal approximation.The last-layer logits have an induced distribution computed from the feature map.Source: MATHEMATICAL DESCRIPTION OF THE SETUP · p. 14
  • A still more efficient last-layer option is mean-field with a diagonal Gaussian posterior, whose variance comes from the diagonal Hessian of the log-posterior.The passage describes this approximation as easy to implement and computationally cheap.Source: MATHEMATICAL DESCRIPTION OF THE SETUP · p. 15

OOD DETECTION

  • OOD tests use a two-layer convolutional network on MNIST and ResNet-18 on CIFAR-10, CIFAR-100, and SVHN, with FMNIST, notMNIST, and KMNIST as MNIST OOD sets.The sampling baseline uses 100 posterior samples.Source: OOD DETECTION · p. 15

PROBIT VS LB

  • In the KFAC comparison, the probit approximation appears to break down in the MNIST scenarios.Table 6 compares extended probit with the normalized Laplace Bridge norm.Source: PROBIT VS LB · p. 15, Table 6 · p. 15
Loading 2003.01227v4…