Tractable Function-Space Variational Inference in Bayesian Neural Networks
Abstract
Reliable predictive uncertainty estimation plays an important role in enabling the deployment of neural networks to safety-critical settings.
A popular approach for estimating the predictive uncertainty of neural networks is to deﬁne a prior distribution over the network parameters, infer an approximate posterior distribution, and use it to make stochastic predictions.
However, explicit inference over neural network parameters makes it difﬁcult to incorporate meaningful prior information about the data-generating process into the model.
In this paper, we pursue an alternative approach.
Recognizing that the primary object of interest in most settings is the distribution over functions induced by the posterior distribution over neural network parameters, we frame Bayesian inference in neural networks explicitly as inferring a posterior distribution over functions and propose a scalable function-space variational inference method that allows incorporating prior information and results in reliable predictive uncertainty estimates.
We show that the proposed method leads to state-of-the-art uncertainty estimation and predictive performance on a range of prediction tasks and demonstrate that it performs well on a challenging safety-critical medical diagnosis task in which reliable uncertainty estimation is essential.
Introduction
Machine learning models succeed at an increasingly wide range of narrowly deﬁned tasks [Krizhevsky et al., 2012, Mnih et al., 2013, Silver et al., 2016, Jumper et al., 2021] but may fail without warning when used on inputs that are meaningfully different from the data they were trained on [Amodei et al., 2016, Hendrycks et al., 2021, Rudner and Toner, 2021a,b].
To deploy machine learning models in safety-critical environments where failures are costly or may endanger human lives, machine learning methods must be reliable and have the ability to ‘fail gracefully.’
A promising tool for incorporating fail-safe mechanisms into machine learning systems, predictive uncertainty quantiﬁcation allows machine learning models to express their conﬁdence in the correctness of their predictions.
In this paper we develop a method for obtaining reliable uncertainty estimates in Bayesian neural networks ( BNN s, Neal [1996]).
While BNN s have promised to combine the advantages of deep learning and Bayesian inference, existing approaches for approximate inference in BNN s fall short of this promise and have been demonstrated to result in approximate posterior predictive distributions that underperform ‘non-Bayesian’ methods both in terms of predictive accuracy and uncertainty quantiﬁcation—making them of limited use in practice [Ovadia et al., 2019, Foong et al., 2019,
Farquhar et al., 2020a, Band et al., 2021].
A potential reason for this shortcoming is that commonly used parameter-space inference methods make it difﬁcult to deﬁne meaningful priors that effectively incorporate information about the data-generating process into inference.
PredictiveMean FunctionDraw TrainingData (a) Predictive Distribution (b) Predictive Mean       0 .
V [ y | D ; x ] (c) Predictive Variance
To avoid this limitation, we follow Sun et al.
[2019] and consider a variational objective deﬁned explicitly in terms of distributions over functions induced by distributions over parameters.
In contrast to prior works that rely on approximation techniques that prevent such function-space variational objectives to be used with high-dimensional inputs and highly-overparameterized neural networks, we propose a simple estimator of the Kullback-Leibler divergence between distributions over functions that enables us to perform stochastic variational inference.
The proposed estimation procedure allows deﬁning priors that explicitly encourage high predictive uncertainty away from the training data as well as priors that reﬂect relevant information about the task at hand.
We demonstrate that this approach leads to posterior approximations that exhibit signiﬁcantly improved predictive uncertainty estimates compared to a wide array of state-of-the-art Bayesian and non-Bayesian methods.
Figure 1 shows examples of predictive distributions obtained via functionspace variational inference on low-dimensional, easy-to-visualize datasets.
As can be seen in the ﬁgures, the predictive distributions ﬁt the training data well while also exhibiting a high degree of predictive uncertainty in parts of the input space far away from the training data, as desired.
Contributions.
We propose a simple estimation procedure for performing function-space variational inference in BNN s.
The variational method allows for the incorporation of meaningful prior information about the data-generating process into the inference and produces reliable predictive uncertainty estimates.
We perform a thorough empirical evaluation in which we compare the proposed approach to a wide array of competing methods and show that it consistently results in high predictive performance and reliable predictive uncertainty estimates, outperforming other methods in terms of predictive accuracy, robustness to distribution shifts, and uncertainty-based detection of distributionally-shifted data samples.
We evaluate the proposed method on standard benchmarking datasets as well as on a safety-critical medical diagnosis task in which reliable uncertainty estimation is essential.
Preliminaries
We consider supervised learning tasks on data D ˙= { ( x n , y n ) }
Nn =1 =
( X D , y D ) with inputs x n 2 X ✓ R D and targets y n 2 Y , where Y ✓ R Q for regression and Y ✓ { 0 , 1 } Q for classiﬁcation tasks.
Bayesian neural networks ( BNN s) are stochastic neural networks trained using (approximate)
Bayesian inference.
Denoting the parameters of such a stochastic neural network by the multivariate random variable ⇥ 2 R P and letting the function mapping deﬁned by a neural network architecture be given by f : X ⇥ R P ! R Q , we obtain a random function f ( · ; ⇥ ) .
For a parameter realization ✓ , we obtain a corresponding function realization, f ( · ; ✓ ) .
When evaluated at a ﬁnite collection of points X =
{ x i } mi =1 , f
( X ; ⇥ ) is a multivariate random variable and f ( X ; ✓ ) is a vector.
Letting p y | f ( X ; ⇥ ) be a likelihood function and p y | f
( X ; ⇥ ) ( y D | f
( X D ; ✓ )) be the likelihood of observing the targets y D under the stochastic function f ( · ; ⇥ ) evaluated at inputs X D and letting p ⇥ be a prior distribution over the stochastic network parameters ⇥ , we can use Bayes’ Theorem to ﬁnd the posterior distribution, p ⇥ |D
[MacKay, 1992, Neal, 1996].
However, since the mapping f is a nonlinear function of the stochastic parameters ⇥ , exact inference is analytically intractable.
Variational inference is an approach that seeks to sidestep this intractability by framing posterior inference as a variational optimization problem, where the goal is to ﬁnd a distribution q ⇥ in a variational family Q q ⇥ that solves the variational problem min q ⇥ 2 Q q [Equation] D KL ( q ⇥ k p ⇥ |D )
[Wainwright and Jordan, 2008].
If Q q ⇥ is the family of mean-ﬁeld Gaussian distributions and the prior distribution over parameters p ⇥ given by a diagonal Gaussian distribution, the resulting variational objective is amenable to stochastic variational inference and can be optimized using gradient-based methods [Hinton and van Camp, 1993, Graves, 2011, Hoffman et al., 2013, Blundell et al., 2015].
2.1 A Function-Space Perspective on Variational Inference in Bayesian Neural Networks
Instead of seeking to infer an approximate posterior distribution over parameters, we frame variational inference in stochastic neural networks as inferring an approximation to the posterior distribution over functions p f ( · ; ⇥ ) |D induced by the posterior distribution over parameters p ⇥ |D , that is,
[Equation]
where  ( · ) is the Dirac delta function [Wolpert, 1993].
Considering the prior distribution over functions p f ( · ; ⇥ ) induced by a prior distribution over parameters p ⇥ ,
[Equation]
and the variational distribution over functions q f ( · ; ⇥ ) induced by a variational distribution over parameters q ⇥ ,
[Equation]
we can express the problem of ﬁnding a posterior distribution over functions variationally as
[Equation]
which allows us to effectively incorporate meaningful prior information about the underlying datagenerating process into training.
As discussed by Burt et al.
[2021], this variational objective is guaranteed to be well-deﬁned for suitably chosen prior distributions over functions.
Speciﬁcally, the
KL divergence between two distributions over functions generated from different distributions over parameters applied to the same mapping (e.g., the same neural network architecture) is well-deﬁned (i.e., ﬁnite) if the KL divergence between the distributions over parameters is ﬁnite, since, by the strong data processing inequality [Polyanskiy and Wu, 2017],
[Equation]
As a result, if D KL ( q ⇥ k p ⇥ ) < 1 , which is the case for ﬁnite-dimensional parameter vectors ⇥ and q ⇥ absolutely continuous with respect to p ⇥ , then the function-space KL divergence is ﬁnite and thus well-deﬁned as a variational objective.
Hence, for a likelihood function deﬁned on a ﬁnite set of training targets y D and a suitably deﬁned prior distribution over functions, we can express the variational problem above equivalently as the well-deﬁned maximization problem max q ⇥ 2 Q ✓ F ( q ⇥ ) with
[Equation]
where D KL ( q f ( · ; ⇥ )
k p f ( · ; ⇥ ) ) is also a KL divergence between distributions over functions.
Unfortunately, evaluating the KL divergence in Equation (6) is in general intractable for arbitrary mappings f .
To obtain a tractable objective, Sun et al.
[2019] showed that D KL ( q f ( · ; ⇥ )
k p f ( · ; ⇥ ) ) can be expressed as the supremum of the KL divergence from q f ( · ; ⇥ ) to p f ( · ; ⇥ ) over all ﬁnite sets of evaluation points, resulting in the objective function [Equation] where X N ˙= [Equation]
[Equation]
✓ R
[Equation]
⇥ D } is the collection of all ﬁnite sets of evaluation points.
However, this objective function is still challenging to optimize in practice: The supremum cannot be obtained analytically and the KL divergence term itself is analytically intractable and difﬁcult to estimate in high dimensions—even for a single evaluation point.
In the next section, we will describe an approximation and estimation procedure that allows scaling function-space variational inference to large neural networks and high-dimensional input data.
The primary obstacle to computing the objective in Equation (6) is the KL divergence from q f ( · ; ⇥ ) to p f ( · ; ⇥ ) .
There are two reasons why the KL divergence in Equation (7) is intractable:
First, for
BNN s or other non-linear models, we do not have access to the probability density functions of the multivariate distributions q f
( X ; ⇥ ) and p f ( X ; ⇥ ) ; second, for all but extremely simple input spaces, we are unable to compute the supremum over all possible ﬁnite sets of evaluation points.
In the remainder of this section, we outline an approach for obtaining an estimator of a locally accurate approximation to the KL divergence that allows for scalable gradient-based optimization of Equation (7).
We ﬁrst approach the problem of computing the KL divergence between two BNN s evaluated at a ﬁnite set of points.
To do so, we ﬁrst derive tractable approximations to the distributions over functions q f ( X ; ⇥ ) and p f
( X ; ⇥ )
Next, we show that under these approximations, we are able to obtain a closed-form approximation to the KL divergence and describe a simple Monte Carlo estimator of the supremum in the function-space KL divergence.
3.1 Approximating Distributions over Functions via Local Linearization
To obtain an approximation to the probability distributions of q f ( X ; ⇥ ) and p f ( X ; ⇥ ) , we use a ﬁrstorder Taylor expansion of the mapping f about the mean parameters of q ⇥ and p ⇥ , respectively, and derive the induced distributions under the linearized mapping.
For a stochastic function f ( · ; ⇥ ) deﬁned in terms of stochastic parameters ⇥ distributed according to distribution g ⇥ with m ˙= E g ⇥ [ ⇥ ] and S ˙= Cov g ⇥ [ ⇥ ] , we denote the linearization of the stochastic function f ( · ; ⇥ ) about m by
[Equation]
where J ( · ; m ) ˙= ( @ f ( · ; ⇥ ) / @ ⇥ )
| ⇥ = m is the Jacobian of f ( · ; ⇥ ) evaluated at ⇥ = m , and the mean and covariance of the distribution over the linearized mapping ˜ f
at X , X 0 2 X are given by
[Equation]
For a derivation of this result, see Appendix A. Since Gaussianity is preserved under afﬁne transformations, if g ⇥ is a multivariate Gaussian distribution with mean m and diagonal co-variance S , then the distribution ˜ g over ˜ f
( X ; ⇥ ) is given by
[Equation]
For stochastic functions parameterized by many millions of parameters, obtaining the covariance of ˜ g ˜ f
( X ; ⇥ ) —which requires computing an inner product of two Jacobian matrices—can be computationally expensive.
Instead of computing the distribution over the linearized mapping exactly, we can construct a suitable Monte Carlo estimator.
To do so, we consider a partition of the set of parameters into sets ↵ and  (with |  | ⌧ | ↵ | ) and note that the linearized mapping can then be expressed as ˜ [Equation]
[Equation]
with
[Equation]
(13) where J ↵ ( · ; m ) and J  ( · ; m ) are the columns of the Jacobian matrix corresponding to the sets of parameters ↵ and  , respectively, and ⇥ ↵ and ⇥  are the corresponding random parameter vectors.
Noting that Equation (12) expresses ˜ f
as a sum of (afﬁne transformations of) random variables, we can use the fact that for independent Gaussian random variables X and Y , the distribution h Z of
Z = X + Y is equal to the convolution of the distributions h X
and h Y to obtain an approximation to ˜ f .
In particular, we can show that if g ⇥ is a multivariate Gaussian distribution with ⇥ ↵ ? ⇥  , the distribution ˜ g ˜ f
( X ; ⇥ ) can be approximated by the Monte Carlo estimator
[Equation]
where g ⇥  = N ( m  , S  ) and samples ˜ f ↵ ( X ; m , ⇥ ↵ ) ( j ) are obtained by sampling parameters from the distribution g ⇥ ↵ = N ( m ↵ , S ↵ ) .
For a derivation of this result, see Appendix A.
This approximation is exact if m = 0 , and if m 6 = 0 , it converges to ˜ g ˜ f
( X ; m , ⇥ ) as S ↵ decreases.
That is, for S ↵ with very small [ S ↵ ] ii , the approximation will be accurate.
In our empirical evaluation, we used a single Monte Carlo sample, R = 1 , and chose ↵ to be the set of parameters in neural network layers 1 : L  1 and  to be the set of parameters in the ﬁnal neural network layer.
3.2 Approximating the Function-Space Kullback-Leibler Divergence
From Section 3.1, we know that if q ⇥ and p ⇥ are both Gaussian distributions, then the induced distributions under the linearized mapping ˜ f evaluated at a ﬁnite set of evaluation points will be Gaussian as well.
This means that for Gaussian variational and prior distributions over ⇥ , we can obtain locally accurate approximations to the induced distributions q f ( · ; ⇥ ) to p f ( · ; ⇥ ) and use them to approximate the KL divergence in the variational objective by D KL (˜ q ˜ f
( X ; ⇥ )
k ˜ p
˜ f
( X ; ⇥ ) ) .
Moreover, for an isotropic
Gaussian prior and a mean-ﬁeld Gaussian variational distribution, D KL (˜ q ˜ f
( X ; ⇥ )
k ˜ p
˜ f
( X ; ⇥ ) ) is a KL divergence between two multivariate Gaussians and can be obtained analytically.
Using this approximation, we obtain an estimator of the variational objective given by
[Equation]
15) where the arguments of the KL divergence have been replaced by the (locally accurate) approximations to the variational and prior distributions over functions evaluated at X , respectively.
Since the stochastic functions ˜ f ( · ; ⇥ ) induced by q ⇥ and p ⇥ under the linearized mapping will be closer to the stochastic function under f the smaller the variance of q ⇥ and p ⇥ , respectively, the approximation to the KL divergence will be more accurate the smaller the variance of q ⇥ and p ⇥ .
Next, we turn to computing the supremum.
Unlike Sun et al.
[2019], who consider the supremum as a separate optimization problem, we do not seek to compute the supremum by searching over points
X 2 X N
but instead propose to estimate the supremum at every gradient step via a simple ﬁnite-sample estimator.
Speciﬁcally, letting I ( X ) ˙= D KL (˜ q ˜ f ( X ; ⇥ ) k ˜ p
˜ f
( X ; ⇥ ) ) , we estimate G = sup X 2 X N I ( X ) using the Monte Carlo estimator
[Equation]
16) where X S C ˙= { X ( i ) C } Si =1 is a collection of S sets of context points
X ( i ) C ˙= { x ( j ) } Kj =1 jointly sampled from a context distribution p X C .
Each context set X ( i ) C can be viewed as a single Monte Carlo sample from the input space so that the estimator ˆ G ( X S C ) provides an S -sample Monte Carlo estimate of the supremum.
While this estimator is crude and only provides a rough approximation to the true supremum, it encourages the variational distribution over functions to match the prior distribution over functions on the sets of context points.
The choice of the context distribution p X C can be informed by knowledge about the prediction task and should be viewed as a problem-speciﬁc modeling choice.
Similarly, the numbers of samples S and K are hyperparameters to be optimized with a validation set.
For details on how p X C is chosen for the empirical evaluation in Section 5, see Appendix D.
3.3 Stochastic Estimation of the Approximnate Function-Space Variational Objective
Let q ⇥ be a Gaussian mean-ﬁeld variational distribution, let p ⇥ be an isotropic Gaussian prior, let ( X B , y B ) be a mini-batch of the training data, and reparameterize ⇥ as ˆ ⇥ ( µ , ⌃ , ✏ ( j ) )
˙= µ + ⌃  ✏ ( j ) .
Using the estimator ˆ G ( X S C ) deﬁned above and estimating the expected log-likelihood via Monte
Carlo sampling, we obtain a Monte Carlo estimator for the function-space variational objective:
[Equation]
[Equation]
M
[Equation]
with ✏ ( j ) ⇠ N ( 0 , I P ) and X S C as deﬁned above.
This Monte Carlo estimator is biased due to the linearization and context-set approximations but allows for scalable gradient-based stochastic optimization.
Selection of Prior.
For all experiments that involve uncertainty quantiﬁcation, we chose a prior distribution over parameters that induces a prior distribution over functions p f ( · ; ⇥ ) and a prior predictive distribution that exhibit a high degree of predictive uncertainty at evaluation points from regions in input space where p X C has non-zero support and, under smoothness constraints, on evaluation points in nearby regions.
For settings where prior information is encoded in data—for example, in the form of expert demonstrations of robotic manipulation tasks [Rudner et al., 2021] or in the form of pre-trained networks in continual or transfer learning [Rudner et al., 2022]—an empirical prior that reﬂects this information can be speciﬁed.
For further details, see Appendix D.
Selection of Context Distribution.
The distribution p X C allows us to incorporate information about the data-generating process into training and encourage the variational distribution to match the prior over functions in relevant parts of the input space.
By taking advantage of the abundance of data available in real-world settings, context distributions can be constructed from large datasets like ImageNet [Krizhevsky et al., 2012], from small but diverse datasets like CIFAR-100, or by using any set of task-related unlabeled data.
In our experiments, we choose two types of context distributions.
One of the context distributions is constructed from the training data and only contains randomly sampled monochrome images, and one is constructed from a real-world dataset generated from a data distribution related to that of the training data.
For example, when training on FashionMNIST, we use KMNIST as the context distribution, and when training on CIFAR-10, we use CIFAR-100 as the context distribution. For further details, see Appendix D.
Posterior Predictive Distribution.
After optimizing the variational objective with respect to the parameters of the variational distribution q ⇥ , we use the fact that we can obtain function draws by sampling from the distribution over parameters to obtain an approximate posterior predictive distribution
[Equation]
where M ⇤ is the number of Monte Carlo samples used to estimate the predictive distribution.
Related Work
There is a growing body of work on function-space approaches to inference in BNN s, deep learning, and applications such as continual learning [Benjamin et al., 2019, Sun et al., 2019, Titsias et al., 2020, Burt et al., 2021, Pan et al., 2020, Ma and Hernández-Lobato, 2021, Rudner et al., 2022].
Function-Space Inference in Bayesian Neural Networks.
Previously proposed methods for FSVI in BNN s are based on approximate gradient estimators and either replace the supremum in Equation (7) with an expectation [Sun et al., 2019] or do not deﬁne an explicit variational objective [Wang et al., 2019].
Sun et al.
[2019] and Carvalho et al.
[2020] use Gaussian process priors over functions for which the function-space variational inference problem is not well-deﬁned (see Section 2.
and Burt et al.
More recent work has attempted to circumvent the intractability of the variational objective in Equation (6) by proposing alternative objectives for function-space inference in BNN s [Ma et al., 2019, Ober and Aitchison, 2020, Ma and Hernández-Lobato, 2021].
Rudner et al.
[2022] extend the approach presented in Section 3 to sequential inference problems and apply it to continual learning.
Linear Models.
Immer et al.
[2020] and Khan et al.
[2019] show that approximate BNN posterior distribution via the Laplace and Generalized-Gauss-Newton approximation corresponds to exact posteriors under linearizations of different models.
Unlike in our approach, they use a Laplace approximation and do not perform variational inference and do not optimize the variance parameters.
Furthermore, Immer et al.
[2020] and Khan et al.
[2019] use a neural network model to obtain a parameter maximum a posteriori estimate, but then use a linearization of the neural network model to compute a posterior predictive distribution.
In contrast, our work only uses the linearization to obtain an estimator of the variational objective but uses the unlinearized model to construct a posterior predictive distribution.
Pathologies of Variational Inference in Bayesian Neural Networks.
Burt et al.
[2021] consider the function-space variational objective in Equation (6) and show that the KL divergence between
BNN s with different networks architectures are not well-deﬁned.
A parallel line of research showed that posterior predictive distributions of shallow BNN s with mean-ﬁeld variational distributions have a limited ability to represent complex covariance structures in function space [Foong et al., 2019, 2020] but that deep BNN s do not suffer from this limitation [Farquhar et al., 2020b].
Our results are consistent with the ﬁndings of Farquhar et al.
[2020b] that mean-ﬁeld variational distributions are able to represent complex covariance structures in function space.
[Table]
Empirical Evaluation
In this section [Table]
we evaluate FSVI on high-dimensional classiﬁcation tasks that were out of reach for function-space variational inference methods proposed in prior works and compare FSVI to several well-established and state-of-the-art Bayesian deep learning and deterministic uncertainty quantiﬁcation methods.
We show that FSVI (sometimes signiﬁcantly ) outperforms existing Bayesian and non-Bayesian methods in terms of their in-distribution uncertainty calibration and out-of-distribution predictive uncertainty estimation.
For a details on models, training and validation procedures
[Table] and datasets used, see Appendix D. For a comparison to Sun et al.
[2019] on small-scale regression tasks, see Appendix B.2.
5.1 Predictive Performance, Uncertainty Estimation, and Distribution Shift Detection
In this set of experiments, we assess the reliability of the uncertainty estimates generated by FSVI .
If a
BNN trained via FSVI is able to perform reliable uncertainty estimation, its predictive uncertainty will be signiﬁcantly higher on input points that were generated according to a different data-generating distribution than the training data.
For models trained on the FashionMNIST dataset, we use the
MNIST and NotMNIST datasets as out-of-distribution evaluation points, while for models trained on the CIFAR-10 dataset, we use the SVHN dataset as out-of-distribution evaluation points.
For models trained on either FashionMNIST or CIFAR-10, we evaluate their in-distribution performance in terms of test accuracy, test log-likelihood, and test calibration.
To evaluate the quality of different models’ uncertainty estimates, we compute uncertainty estimates for the pairs FashionMNIST/MNIST, FashionMNIST/NotMNIST, and CIFAR-10/SVHN to and measure for a range of thresholds how well the datasets in each pair can be separated solely based on the uncertainty estimates.
This experiment setup follows prior work by van Amersfoort et al.
[2020] and Immer et al.
We report the area under the receiver operating characteristic (ROC) curve in Tables 1 and 2.
Predictive Performance and Calibration.
To assess in-distribution predictive performance and calibration, we report the test accuracy, negative log-likelihood (NLL), and expected calibration error (ECE) for models trained on FashionMNIST and CIFAR-10 in Tables 1 and 2.
On both
FashionMNIST and CIFAR-10, FSVI achieves the lowest NLL and either the best or second-best predictive accuracy and ECE, respectively, across all methods.
Notably, FSVI signiﬁcantly outperforms
SPG [Ma and Hernández-Lobato, 2021], an alternative function-space variational inference method.
Predictive Uncertainty under Distribution Shift.
In Tables 1 and 2, we report evaluation metrics that elucidate the reliability of different methods’ predictive uncertainty under distribution shift.
FSVI exhibits reliable predictive uncertainty estimates that allow distinguishing between in- and out-of-distribution inputs with high accuracy.
As would be expected, we observe that using context
[Table]
distributions that reﬂect our knowledge about the data-generating process can signiﬁcantly improve uncertainty quantiﬁcation under FSVI .
For the FashionMNIST experiment, we used the KMNIST dataset, which contains grayscale images of Kuzushiji letters, and for the CIFAR-10 experiment, we used the CIFAR-100 dataset, which contains RGB images of 100 classes.
Both KMNIST and
CIFAR-100 differ from the OOD datasets (MNIST and NotMNIST and SVHN, respectively) used to compute OOD-AUROC metrics in Tables 1 and 2, but using them as context distributions signiﬁcantly increased the ability of BNN s trained via FSVI to identify distributionally shifted samples.
Since the variational objective encourages matching the prior (which we chose to have high variance) on samples from the context distribution can improve uncertainty estimation in regions of the input space far from the training data.
5.2 Generalization and Reliability of Predictive Uncertainty under Distribution Shift
To assess the reliability of predictive models in deep learning, Ovadia et al.
[2019] propose the following desiderata:
In order for a model to be considered reliable, it ought to (i) exhibit low predictive uncertainty on training data and high predictive uncertainty on out-of-distribution inputs, (ii) generate predictive uncertainty estimates that allow distinguishing in- from out-of-distribution inputs, and (iii) if possible, maintain high predictive accuracy even under distribution shift.
Models that satisfy these desiderata are less likely to make poor, high-conﬁdence predictions and more amenable for use in safety-critical downstream tasks.
[Figure]
To illustrate these desiderata, we follow Ovadia et al.
[2019] and consider the rotated MNIST task, where a model is trained on MNIST and evaluated on rotated MNIST digits.
The goal is to maintain a high level of predictive accuracy (measured in terms of
Brier scores) while exhibiting an increasing level of predictive uncertainty on distribution shifts of increasing magnitude.
Figure 2 shows Brier scores (lower is better) and predictive entropy estimates (higher means more uncertain) of four different models.
As rotating the MNIST digits gradually shifts the data distributions, we would expect Brier scores to increase (corresponding to worse predictive accuracy) as the rotation angle increases.
A model with reliable predictive entropy estimates would only experience a small decrease under distribution shift while exhibiting a large increase in predictive uncertainty.
As can be seen in the plot, the Brier scores of FSVI decreases the least, while FSVI ’s uncertainty is signiﬁcantly higher than other models’.
To assess the
[Figure]
reliability of different uncertainty quantiﬁcation methods on a more challenging distribution-shift task, we consider corrupted CIFAR-10 inputs under the second-mildest corruption level used in [Ovadia et al., 2019] and report our results in Table 2.
Consistent with the rotated MNIST results, FSVI achieves the highest accuracy on the corrupted data.
5.3 Safety-Critical Uncertainty-Aware Selective Prediction: Diabetic Retinopathy Diagnosis
To evaluate the reliability of the predictive uncertainty of FSVI in a realworld safety-critical setting, we consider the task of diagnosing diabetic retinopathy (DR), a medical condition that can lead to impaired vision, from retina scans [Leibig et al., 2017, Filos et al., 2019, Band et al., 2021].
We use two publicly available datasets, EyePACS [2015] and APTOS [2019], each containing RGB images of a human retina graded by a medical expert on the following scale: 0 (no DR), 1 (mild DR), 2 (moderate
DR), 3 (severe DR), and 4 (proliferative DR).
The Kaggle dataset was collected from patients in the United States, while the APTOS dataset was collected from patients in India using cheaper but more modern scanning devices.
We follow Leibig et al.
[2017], Filos et al.
[2019], and Band et al.
[2021] and binarize all examples from both the EyePACS and APTOS datasets by dividing the classes up into sight-threatening diabetic retinopathy—deﬁned as moderate diabetic retinopathy or worse (classes { 2 , 3 , 4 } )—and non-sight-threatening diabetic retinopathy—deﬁned as no or mild diabetic retinopathy (classes { 0 , 1 } ).
This results in a binary prediction task.
To assess the reliability of predictive models when medical training and test data are obtained from different patient populations or collected with the same medical equipment, we follow Band et al.
[2021] and use the Kaggle dataset for training and the distributionally shifted APTOS dataset for evaluation.
The results are shown in Figure 4, which plot the ROC curves for the binary prediction problems as well as the area under the ROC curve for an uncertainty aware selective prediction task.
For further details about the uncertainty-aware selective prediction evaluation protocol, see Appendix D.4.
Figure 4 shows that FSVI performs well on all four tasks and is only outperformed by MC DROPOUT .
For full tabular results, see Appendix B.1.
Conclusion
The paper proposed a scalable and effective approach to function-space variational inference in BNN s.
We demonstrated that the proposed estimator of the function-space variational objective can be scaled up to high-dimensional data and large neural network architectures and that FSVI exhibits consistently reliable in- and out-of-distribution predictive performance on a wide range of datasets when compared to well-established and state-of-the-art uncertainty quantiﬁcation methods.
We hope that this work will lead to further research into function-space variational inference and the development of more sophisticated data-driven prior distributions over functions.
Acknowledgements
We thank Bryn Elesedy, Bobby He, and Andrew Jesson for feedback on an early draft of this paper.
We thank Joost van Amersfoort for helpful discussions about experiment design and implementations.
Tim G. J. Rudner is funded by the Rhodes Trust and the Engineering and Physical Sciences Research
Council (EPSRC).
We gratefully acknowledge donations of computing resources by the Alan Turing
Institute.
[-ϵ there exists some m j = 
-]
[-ECE   
-]