Tractable Function-Space Variational Inference in Bayesian Neural Networks
Anonymous Author(s)
AffiliationAddress
Abstract
Reliable predictive uncertainty estimation plays an important role in allowingneural networks to be deployed in safety-critical settings.
A popular approach forestimating the predictive uncertainty of neural networks is to treat the networkparameters as random variables and infer an approximate posterior that can beused to obtain a distribution over network predictions.
However, explicit inferenceover neural network parameters makes it difficult to incorporate meaningful priorinformation about the data generating process into training.
In this paper, we pursuean alternative approach.
Noting that stochastic neural networks define distributionsover functions induced by distributions over parameters, we follow prior work inframing Bayesian inference in neural networks as inferring a posterior distributionover functions and propose a scalable function-space variational inference methodthat allows incorporating prior information and encourages reliable predictiveuncertainty estimation.
We show that the proposed method leads to state-of-theart uncertainty estimation and predictive performance on a range of predictionproblems, and demonstrate that it performs well on a challenging safety-criticalmedical diagnosis task in which reliable uncertainty estimation is essential.
Introduction
Machine learning models succeed at an increasingly wide range of narrowly defined tasks [Krizhevskyet al., 2012, Mnih et al., 2013, Silver et al., 2016, Jumper et al., 2021], but fail without warning whenused on inputs that are meaningfully different from the data they were trained on [Amodei et al.
,2016, Hendrycks et al., 2021].
To deploy machine learning models in safety-critical environmentswhere failures are costly or may endanger human lives, machine learning methods must be reliableand have the ability to ‘fail gracefully.’
As a promising tool for incorporating fail-safe mechanismsinto machine learning systems, predictive uncertainty quantification allows machine learning modelsto express their confidence in the correctness of their predictions Gal and Ghahramani [2016].
In this paper we develop a method for obtaining reliable uncertainty estimates in Bayesian neuralnetworks ( BNN s, Neal [1996]).
While BNN s have promised to combine the advantages of deeplearning and Bayesian inference, existing approaches for approximate inference in BNN s fall short ofthis promise and have been demonstrated to result in approximate posterior predictive distributionsthat underperform non-Bayesian methods both in terms of predictive accuracy and uncertaintyquantification—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 commonlyused parameter-space inference methods make it difficult to define meaningful priors that effectivelyincorporate information about the data-generating process into inference.
To avoid this limitation, we follow the foundation work by Sun et al.
[2019] and consider a variational objective defined explicitly in terms of distributions over functions induced by distributions over pa-
−−−    −−
PredictiveMean FunctionDraw TrainingData (a) Predictive Distribution (b) Predictive Mean −−   −   0 .
V [ y | D ; x ] (c) Predictive Variance rameters.
In contrast to prior works that rely on approximation techniques that prevent such functionspace variational objectives to be used with high-dimensional inputs and highly-overparameterizedneural networks [Sun et al., 2019], we propose a simple estimator of the Kullback-Leibler divergencebetween distributions over functions that allows for stochastic variational inference.
The proposedestimation procedure allows defining priors that explicitly encourage high predictive uncertainty awayfrom the training data as well as priors that reflect relevant information about the task at hand.
We demonstrate that this approach leads to posterior approximations that exhibit significantly improved predictive uncertainty estimates compared to a wide array of state-of-the-art Bayesian andnon-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 thefigures, the predictive distributions fit the training data well while also exhibiting a high degreeof predictive uncertainty in parts of the input space far away from the training data, as desired.
Anonymized code to reproduce our results is available at https://github.com/fsvibnns/fsvi .
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 priorinformation about the data-generating process into the inference and produces reliable predictiveuncertainty 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 highpredictive performance and reliable predictive uncertainty estimates, outperforming other methodsin terms of predictive accuracy, robustness to distribution shifts, and uncertainty-based detection ofdistributionally-shifted data samples.
We evaluate the proposed method on standard benchmarkingdatasets as well as on a safety-critical medical diagnosis task in which reliable uncertainty estimationis essential.
Preliminaries
We consider supervised learning tasks on data D ˙= { ( x n , y n ) }
Nn =1 =
( X D , y D ) with inputs 61 x n ∈ X ⊆ R D and targets y n ∈ Y , where Y ⊆ R Q for regression and Y ⊆ { 0 , 1 } Q for classification 62 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 multivariaterandom variable Θ ∈ R P and letting the function mapping defined by a neural network architecturebe given by f : X × R P → Y , we obtain a random function f ( · ; Θ ) .
For a parameter realization 66 θ , we obtain a corresponding function realization, f ( · ; θ ) .
When evaluated at a finite collection ofpoints
X = { x i } mi =1 , f
( X ; Θ ) is a multivariate random variable and f ( X ; θ ) is a vector.
Letting p ( y D | f
( X D ; θ )) be the likelihood of observing the targets y under the stochastic functionf ( · ; Θ ) evaluated at inputs X D and letting p Θ be a prior distribution over the stochastic networkparameters
Θ , we can use Bayesian inference for finding the conditional distribution over parametersgiven the observed data, p Θ |D
[MacKay, 1992, Neal, 1996].
However, since the mapping f is a nonlinear function of the stochastic parameters this approach we can obtain a Bayesian neural network defined in terms of a variational distribution over parametersq Θ by solving min q Θ ∈Q q Θ D KL ( q Θ ∥ p Θ |D ) , where Q q Θ is a variational family of distributions.
If
Q q Θ is the family of mean-field Gaussian distributions and the prior distribution over parameters p Θgiven by a diagonal Gaussian distribution, the resulting variational objective is amenable to stochasticvariational inference and can be optimized using gradient-based methods.
[Blundell et al., 2015,
Graves, 2011, Hinton and van Camp, 1993, Hoffman et al., 2013].
2.1 A Function-Space Perspective on Variational Inference in Bayesian Neural Networks
By framing variational inference in BNN s as computing an approximation to a posterior distributionover functions , p f ( · ; Θ ) |D , we can incorporate meaningful prior information about the underlyingdata-generating process into training.
Specifically, we can express the problem of finding a posteriordistribution over functions variationally as
[Equation]
where p f ( · ; Θ ) the a prior distribution over functions induced by a prior distribution over parameters 87 p Θ , and q f ( · ; Θ ) is a variational distribution over functions induced by a variational distribution overparameters q Θ .
As discussed by Burt et al.
[2021], this variational objective is guraanteed to bewell-defined for suitably chosen prior distributions over functions.
Specifically, the KL divergencebetween two distributions over functions generated from different distributions over parametersapplied to the same mapping (e.g., the same neural network architecture) is well-defined (i.e., finite)if the KL divergence between the distributions over parameters is finite, since, by the strong dataprocessing inequality [Polyanskiy and Wu, 2017],
[Equation]
As a result, if D KL ( q Θ ∥ p Θ ) < ∞ , which is the case for finite-dimensional parameter vectors Θ and 95 q Θ absolutely continuous with respect to p Θ , then the function-space KL divergence is finite andthus well-defined as a variational objective.
Hence, for a likelihood function defined on a finite set of training targets y D and a suitably definedprior distribution over functions, we can express the variational problem above equivalently as thewell-defined maximization problem max q Θ ∈Q θ F ( q Θ , X D , y D ) with
[Equation]
where D KL ( q f ( · ; Θ ) ∥
p f ( · ; Θ ) ) is also a KL divergence between distributions over functions.
Unfortunately, evaluating the KL divergence in Equation (2) is intractable in general.
To obtaina tractable objective, Sun et al.
[2019] show that D KL ( q f ( · ; Θ ) ∥
p f ( · ; Θ ) ) can be expressed as the 103 supremum of the KL divergence from q f ( · ; Θ ) to p f ( · ; Θ ) over all finite sets of evaluation points, 104 resulting in the objective function
[Equation]
where X N ˙= (cid:83) n ∈ N { X ∈ X n | X n ⊆ R n × D } is the collection of all finite sets of evaluation points.
However, this objective function is still challenging to optimize in practice: The supremum cannotbe obtained analytically and—more problematically—the KL divergence term itself is analyticallyintractable and difficult to estimate—even for a single evaluation point.
While the supremum canbe estimated via Monte Carlo sampling, the fact that the KL divergence between BNN s cannotbe easily estimated poses has prevented the application of function-space variational inference tochallenging prediction tasks.
For example, Sun et al.
[2019] attempted to circumvent estimating the KL divergence between BNN s directly by instead approximating the gradients of the KL divergencewith respect to the variational parameters, but this approach does not scale to large neural networksor high-dimensional input data.
In the next section, we propose an estimation procedure that allows scaling function-space variationalinference to large networks and high-dimensional input data.
A Tractable Function-Space Variational Objective
The primary obstacle to making the objective in Equation (2) tractable is the KL divergence fromq f ( · ; Θ ) to p f ( · ; Θ ) .
In this section, we describe how to obtain an estimator of the KL divergence 120 amenable to gradient-based optimization.
3.1 Approximating the Function-Space Kullback-Leibler Divergence via Local Linearization
First, we note that the KL divergence in Equation (3) is intractable for two reasons: ( i ) For BNN s orother non-linear models, we do not have access to the probability density functions of the multivariatedistributions q f ( X ; Θ ) and p f
( X ; Θ ) ; ( ii ) For all but extremely simple input spaces, we are unable to 125 compute the supremum over all possible finite sets of evaluation points.
In the remainder of thissection, we outline a simple approach for obtaining an estimator of the KL divergence that allows forscalable gradient-based optimization of Equation (3).
We first approach the problem of computing the KL divergence between two BNN s evaluated at a finiteset of points.
To obtain an estimator of KL divergence D KL ( q f ( X ; Θ ) ∥
p f ( X ; Θ ) ) , we approximate 130 q f
( X ; Θ ) and p f
( X ; Θ ) by computing the first-order Taylor expansion of the mapping f about the 131 mean parameters of q Θ and p Θ , respectively.
To obtain locally accurate estimates of the probabilitydensities of q f ( X ; Θ ) and p f ( X ; Θ ) , we use the following result:
Proposition 1 (Distribution under Linearized Mapping) .
For a stochastic function f ( · ; Θ ) defined in terms of stochastic parameters Θ ∼ g Θ , and g Θ with m ˙= E g Θ [ Θ ] , and Jacobian
J ( · , m ) ˙= ∂f ( · ; Θ ) ∂ Θ | Θ = m , denote the linearization of the stochastic function f ( · ; Θ ) about m byf [Equation]
[Equation] [Equation] [Equation] m ) .
Then the mean and covariance of the distribution over the linearized mapping ˜ f at X , X
′ ∈ X aregiven by
[Equation]
Proof. See Appendix A. Corollary 1. 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 ˜ g ˜ f
( X ; Θ ) =
N ( f ( X ; m ) , J ( X , m ) S J ( X ′ , m ) ⊤ ) .
Proof. The result follows from the fact that Gaussianity is preserved under affine transformations. Proposition 1 tells us that if q Θ and p Θ are both Gaussian distributions, then the induced distributionsunder the linearized mapping ˜ f evaluated at a finite set of evaluation points will be Gaussian aswell.
This means that for Gaussian variational or prior distributions over Θ , we can obtain a locallyaccurate estimate of the probability density function of the induced multivariate random variablef ( X ; Θ ) .
1 We can use this insight to estimate the KL divergence in the variational objective by D KL (˜ q ˜ f
( X ; Θ ) ∥ ˜
p ˜ f
( X ; Θ ) ) for several different variational and prior distributions over parameters.
For example, for an isotropic Gaussian prior and a mean-field Gaussian variational distribution, D KL (˜ q ˜ f
( X ; Θ ) ∥ ˜
p ˜ f
( X ; Θ ) ) is a KL divergence between two multivariate Gaussians and can be obtained 150 analytically.
Using this approximation, we obtain an estimator of the variational objective given by
[Equation]
where the arguments of the KL divergence have been replaced by the (locally accurate) approximationsto the variational and prior distributions over functions evaluated at X , respectively.
Since thestochastic functions ˜ f
( · ; Θ ) induced by q Θ and p Θ under the linearized mapping will be closer to thestochastic function under f the smaller the variance of q Θ and p Θ , respectively, the approximation tothe 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 asa separate optimization problem, we do not seek to compute the supremum by searching over points
X ∈ X N
but instead propose to estimate the supremum at every gradient step via a simple finite-sampleestimator.
Specifically, letting I ( X ) ˙= D KL (˜ q ˜ f
( X ; Θ ) ∥ ˜
p ˜ f
( X ; Θ ) ) , we estimate G = sup X ∈X N I ( X ) 161 using the Monte Carlo estimator
[Equation]
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 sampledfrom a context distribution p X C .
Each set of context points 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 Carloestimate of the supremum.
While this estimator is crude and only provides a rough approximationto the true supremum, it encourages the variational distribution over functions to match the priordistribution over functions on the sets of context points.
The choice of the context distribution p X Ccan be informed by knowledge about the prediction task and should be viewed as a problem-specificmodeling choice.
Similarly, the numbers of samples S and K are hyperparameters to be optimizedwith a validation set.
For details on how p X C is chosen in the empirical evaluation in Section 5,see Appendix F.
For an ablation study on different choices of p X C , see Appendix D.
3.2 Stochastic Estimation of the Function-Space Variational Objective
Let q Θ be a Gaussian mean-field 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 ) defined 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] [Equation]
S C , D )
= 1 M
[Equation]
with ϵ ( j ) ∼ N ( 0 , I P ) and X S C as defined above.
Crucially, to avoid having to perform expensivematrix operations as part of the KL computation, we make a diagonalization approximation to thecovariance matrices of ˜ q and ˜ p , respectively.
Selection of Prior.
For all experiments that involve uncertainty quantification, we chose a priordistribution over parameters that induces a prior distribution over functions p f ( · ; Θ ) that, under thelikelihood, exhibits a high degree of predictive uncertainty.
Such a prior will encourage the BNNto have high uncertainty at evaluation points from regions in input space where p X C has non-zerosupport and, under smoothness constraints, on evaluation points in nearby regions.
For settings whereprior information is encoded in data—for example, in the form of expert demonstrations of roboticmanipulation tasks, or in the form of pre-trained networks in continual or transfer learning—anempirical prior that reflect this information can be specified.
For further details, see Appendix F, andfor ablation studies on the effect of different prior choices on performance, see Appendix D.
Selection of Context Distribution.
The distribution p X C allows us to incorporate information aboutthe data-generating process into training.
This way, p X C allows us obtain gradient estimates thatencourage the variational distribution to match the prior over functions in more relevant parts ofthe input space.
By taking advantage of the abundance of data available in real-world settings, wecan construct context distributions from large datasets like ImageNet [Krizhevsky et al., 2012] orfrom small but diverse datasets like CIFAR-100.
In our experiments, we choose two types of contextdistributions.
One distribution constructed from the training data and only containing randomlysampled monochrome images, and one constructed from a related real-world dataset generated froma similar but different data distribution.
For further details, see Appendix F, and for ablation studieson the effect of different context distributions on performance, see Appendix D.
Posterior Predictive Distribution.
After optimizing the variational objective with respect to the parameters of the variational distribution q Θ , we obtain an approximate posterior predictive distribution
[Equation]
where M ∗ is the number of Monte Carlo samples used to estimate the predictive distribution.
Computational Complexity.
Under the diagonal approximation used to compute ¯ F , the time andspace complexity for stochastic gradient estimation scale linearly in the mini-batch size n B and thenumber forward passes through the network per input point (i.e., the number of Monte Carlo samples
M used to estimate the expected log-likelihood plus the number of points on which the KL divergenceis evaluated, SM ), yielding a time and space complexity of O ( n B p ( M + SK )) .
In comparison,parameter-space mean-field variational inference has complexity O ( n B pM ) .
The time complexity ofa prediction on n ∗ points with M ∗ Monte Carlo samples has a complexity of O ( n ∗ M ∗ ) .
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, Burt et al., 2021, Pan et al., 2020,
Sun et al., 2019, Titsias et al., 2020, Ma and Hernández-Lobato, 2021].
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 (3)with an expectation [Sun et al., 2019] or do not define an explicit variational objective [Wang et al.
More recent work has attempted to circumvent the intractability of the variational objectivein Equation (2) 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].
Linear Models.
Immer et al.
[2020] take a similar approach as Khan et al.
[2019] and showthat approximating BNN posterior distribution via the Laplace and Generalized-Gauss-Newtonapproximation corresponds to linearizations of different models.
Crucially, unlike in our approach,they use a Laplace approximation and do not perform variational inference.
As such, they do notoptimize the variance parameters.
While they also use a linearization similar to the first-order Taylorexpansion in Proposition 1, they use it to make predictions, whereas we only use the linearizationto obtain an estimator for the KL divergence between distributions over functions.
In our approach,predictions and expected log-likelihood estimates are obtained from the unlinearized BNN .
Pathologies of Variational Inference in Bayesian Neural Networks.
Burt et al.
[2021] considerthe function-space variational objective in Equation (2) and show that the KL divergence between
BNN s with different networks architectures are not well-defined.
A parallel line of research showedthat posterior predictive distributions of shallow BNN s with mean-field variational distributions havea 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].
Empirical Evaluation
The seminal work of Sun et al.
[2019] introduced function-space variational inference as a tool for approximate inference in BNN s but the approach proposed in Sun et al.
[2019] is limited to small neuralnetworks and low-dimensional input data [Ma and Hernández-Lobato, 2021].
In this section, we evaluate FSVI on high-dimensional classification tasks that were previously out of reach under the approachproposed by Sun et al.
[2019] and compare FSVI to several well-established and state-of-the-art methods.
We show that FSVI (sometimes significantly ) outperforms existing Bayesian and non-Bayesianmethods in terms of their in-distribution uncertainty calibration and out-of-distribution predictiveuncertainty estimation.
For a details on models, training and validation procedures, and datasets used,see Appendix F. Our code can be accessed at https://github.com/fsvibnns/fsvi .
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 predictions will besignificantly higher on input points that were generated according to a different data-generatingdistribution 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 onthe 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 qualityof different models’ uncertainty estimates, we compute uncertainty estimates for the pairs FashionMNIST/MNIST, FashionMNIST/NotMNIST, and CIFAR-10/SVHN to and measure for a rangeof thresholds how well the datasets in each pair can be separated solely based on the uncertaintyestimates.
This experiment setup follows prior work by van Amersfoort et al.
[2020] and Immer et al.[2020].
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 andcalibration, we report the test accuracy, negative log-likelihood (NLL), and expected calibrationerror (ECE) for models trained on FashionMNIST and CIFAR-10 in Tables 1 and 2.
On both
[Table]
LAPLACE [Immer et al., 2020]
[Table] [Table]
FashionMNIST [Table] CIFAR
[Table]
10, FSVI achieves the lowest NLL and either the best or second-best predictive accuracy and ECE, respectively, across all methods.
Notably, FSVI significantly 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 metricsthat elucidate the reliability of different methods’ predictive uncertainty under distribution shift.
FSVI exhibits reliable predictive uncertainty estimates that allow distinguishing between in- andout-of-distribution inputs with high accuracy.
As would be expected, we observe that using contextdistributions that reflect our knowledge about the data-generating process can significantly improveuncertainty quantification under FSVI .
For the FashionMNIST experiment, we used the KMNISTdataset, 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 tocompute OOD-AUROC metrics in Tables 1 and 2, but using them as context distributions significantlyincreased the ability of BNN s trained via FSVI to identify distributionally shifted samples.
Sincethe variational objective encourages matching the prior (which we chose to have high variance) onsamples from the context distribution can improve uncertainty estimation in regions of the input spacefar 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 thefollowing desiderata:
In order for a model to be considered reliable, it ought to (i) exhibit lowpredictive 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-distributiondata points, and (iii) if possible, maintain high predictive accuracy even under distribution shift.
Models that satisfy these desiderata are less likely to make poor, high-confidence predictions and aremore 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, wherea model is trained on MNIST and evaluated on rotated MNIST digits.
The goal is to maintain a highlevel of predictive accuracy (measured in terms of
Brier scores) while exhibiting an increasing level ofpredictive uncertainty on distribution shifts of increasing magnitude.
Figure 2 shows Brier scores (loweris better) and predictive entropy estimates (highermeans more uncertain) of four different models.
Asrotating the MNIST digits gradually shifts the datadistributions, we would expect Brier scores to increase (corresponding to worse predictive accuracy)as
the rotation angle increases.
A model with reliablepredictive entropy estimates would only experiencea small decrease under distribution shift while exhibiting a large increase in predictive uncertainty.
Ascan be seen in the plot, the Brier scores of FSVI decreases the least, while FSVI ’s uncertainty issignificantly higher than other models’.
To assess the reliability of different uncertainty quantificationmethods on a more challenging distribution-shift task, we consider corrupted CIFAR-10 inputs underthe 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 diabeticretinopathy (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 wascollected from patients in the United States, while the APTOS dataset wascollected from patients in India using cheaper but more modern scanningdevices.
We follow Leibig et al.
[2017], Filos et al.
[2019], and Bandet al.
[2021] and binarize all examples from both the EyePACS and APTOS datasets by dividing the classes up into sight-threatening diabeticretinopathy—defined as moderate diabetic retinopathy or worse (classes{ 2 , 3 , 4 } )—and non-sight-threatening diabetic retinopathy—defined as 323 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 obtainedfrom different patient populations or collected with the same medical equipment, we follow Bandet al.
[2021] and use the Kaggle dataset for training and the distributionally shifted APTOS datasetfor evaluation.
The results are shown in Figure 4, which plot the ROC curves for the binaryprediction problems as well as the area under the ROC curve for an uncertainty aware selectiveprediction task.
For further details about the uncertainty-aware selective prediction evaluationprotocol, see Appendix F.4.
Figure 4 shows that FSVI performs well on all four tasks and is onlyoutperformed by MC DROPOUT .
For full tabular results, see Appendix B.1.
[Figure]
s i t i v
[Figure]
5.4 Ablation Study and Discussion
To investigate the effect of different prior distributions over functions on predictive predictiveperformance and uncertainty metrics, we performed an ablation study on the effect of changing thevariance of the induced prior distribution over functions using the FashionMNIST and CIFAR-10datasets.
In Figure 5, we relate the effect of changing the prior variance Σ 0 on uncertainty-baseddistribution shift detection (“OOD-AUROC”, higher is better) and the test negative log-likelihoodand calibration error (lower is better for both) and find that the prior variance range that results in thebest in- vs. shifted distribution uncertainty (i.e., high OOD-AUROC) also results the most accurateprediction and calibration (i.e., low NLL and ECE).
This correlation suggests that optimizing for low ECE and NLL, which can be achieved with an in-distribution held-out validation set, will result inreliable uncertainty estimation both on in-distribution data as well as under distribution shift.
For anextensive set of ablation studies on all relevant hyperparameters, see Appendix D.
[Figure]
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 scaledup to high-dimensional data and large neural network architectures and that FSVI exhibits consistentlyreliable in- and out-of-distribution predictive performance on a wide range of datasets when comparedto well-established and state-of-the-art uncertainty quantification methods.
We hope that this workwill lead to further research into function-space variational inference and the development of moresophisticated data-driven prior distributions over functions.
Checklist
For all authors...
(a) Do the main claims made in the abstract and introduction accurately reflect the paper’scontributions and scope?
[Yes] See Section 5.
(b) Did you describe the limitations of your work?
[Yes] See Section 5.
(c) Did you discuss any potential negative societal impacts of your work?
[N/A]
(d) Have you read the ethics review guidelines and ensured that your paper conforms tothem?
[Yes]
If you are including theoretical results...
(a) Did you state the full set of assumptions of all theoretical results?
[Yes] See Ap-
pendix A.
(b) Did you include complete proofs of all theoretical results?
[Yes] See Appendix A.
If you ran experiments...
(a) Did you include the code, data, and instructions needed to reproduce the main experi-
mental results (either in the supplemental material or as a URL)?
[TODO]
(b) Did you specify all the training details (e.g., data splits, hyperparameters, how theywere chosen)?
[Yes] See Appendix F.
(c) Did you report error bars (e.g., with respect to the random seed after running experi-
ments multiple times)?
[Yes] See Table 2, Table 2, and Table 3.
(d) Did you include the total amount of compute and the type of resources used (e.g., typeof GPUs, internal cluster, or cloud provider)?
[Yes] See Appendix F.9.
If you are using existing assets (e.g., code, data, models) or curating/releasing new assets...
(a) If your work uses existing assets, did you cite the creators?
[Yes]
(b) Did you mention the license of the assets?
[N/A]
(c) Did you include any new assets either in the supplemental material or as a URL?
[Yes]
See https://github.com/fsvibnns/fsvi .
(d) Did you discuss whether and how consent was obtained from people whose data you’re
using/curating?
[N/A]
(e) Did you discuss whether the data you are using/curating contains personally identifiable
information or offensive content?
[N/A]
If you used crowdsourcing or conducted research with human subjects...
(a) Did you include the full text of instructions given to participants and screenshots, if
applicable?
[N/A]
(b) Did you describe any potential participant risks, with links to Institutional Review
Board (IRB) approvals, if applicable?
[N/A]
(c) Did you include the estimated hourly wage paid to participants and the total amount
spent on participant compensation?
[N/A]