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 define 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 difficult 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 defined 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 quantification allows machine learning models to express their confidence 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 quantification—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 difficult to define 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 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 parameters.
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-overparameterized neural networks [Sun et al., 2019], we propose a simple estimator of the Kullback-Leibler divergence between distributions over functions that allows for stochastic variational inference.
The proposed estimation procedure allows defining priors that explicitly encourage high predictive uncertainty away from 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 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 figures, the predictive distributions fit 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 ∈ X ⊆ R D and targets y n ∈ Y , where Y ⊆ R Q for regression and Y ⊆ { 0 , 1 } Q for classification 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 Θ ∈ R P and letting the function mapping defined 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 finite 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 find 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 find a distribution q Θ in a variational family Q q Θ that solves the variational problem min q Θ ∈Q q Θ D KL ( q Θ ∥ p Θ |D )
[Wainwright and Jordan, 2008].
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 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 finding 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-defined for suitably chosen prior distributions over functions.
Specifically, 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-defined (i.e., finite) if the KL divergence between the distributions over parameters is finite, since, by the strong data processing inequality
[Polyanskiy and Wu, 2017],
[Equation]
As a result, if D KL ( q Θ ∥ p Θ ) < ∞ , which is the case for finite-dimensional parameter vectors Θ and q Θ absolutely continuous with respect to p Θ , then the function-space KL divergence is finite and thus well-defined as a variational objective.
Hence, for a likelihood function defined on a finite set of training targets y D and a suitably defined prior distribution over functions, we can express the variational problem above equivalently as the well-defined maximization problem max q Θ ∈Q θ F ( q Θ ) 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 (6) is in general intractable for arbitrary mappings f .
To obtain a tractable objective, Sun et al.
[2019] showed that D KL ( q f ( · ; Θ ) ∥
p f ( · ; Θ ) ) can be expressed as the supremum of the KL divergence from q f ( · ; Θ ) to p f ( · ; Θ ) over all finite sets of evaluation points, 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 cannot be obtained analytically and the KL divergence term itself is analytically intractable and difficult 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.
Deriving a Tractable Function-Space Variational Objective
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 finite 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 first approach the problem of computing the KL divergence between two BNN s evaluated at a finite set of points.
To do so, we first 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 first- order 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 ( · ; Θ ) defined 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 ′ ∈
X are given by
[Equation]
For a derivation of this result, see Appendix A. Since Gaussianity is preserved under affine 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 compute a Monte Carlo estimate.
To construct a suitable Monte Carlo estimator, 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]
12) 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 (affine 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.
In our empirical evaluation, we chose α to be the set of parameters in neural network layers 1 : L − 1 and β to be the set of parameters in the final 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 finite 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 ; Θ ) ∥ ˜
p ˜ f
( X ; Θ ) ) .
Moreover, 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 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 ∈ X N
but instead propose to estimate the supremum at every gradient step via a simple finite-sample estimator.
Specifically, letting I ( X ) ˙= D KL (˜ q ˜ f
( X ; Θ ) ∥ ˜
p ˜ f
( X ; Θ ) ) , we estimate G = sup X ∈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-specific 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 in the empirical evaluation in Section 5, see Appendix E.
3.3 Stochastic Estimation of the Approximnate 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]
with ϵ ( j ) ∼ N ( 0 , I P ) and X S C as defined above.
Crucially, to avoid having to perform expensive matrix operations as part of the KL computation, we make a diagonalization approximation to the covariance matrices of ˜ q and ˜ p , respectively.
Selection of Prior.
For all experiments that involve uncertainty quantification, 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 reflects this information can be specified.
For further details, see Appendix E, and for sensitivity 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 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 distribution is constructed from the training data and only contains randomly sampled monochrome images, and one is constructed from a related real-world dataset generated from a data distribution related to that of the training data.
For further details, see Appendix E, and for sensitivity studies on 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 q
[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 and space complexity for stochastic gradient estimation scale linearly in the mini-batch size n B and the number 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 divergence is evaluated, SK ), yielding a time and space complexity of O ( n B p ( M + SK )) .
The time complexity of a 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, 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 define 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-defined (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-defined.
A parallel line of research showed that posterior predictive distributions of shallow BNN s with mean-field 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 findings of Farquhar et al.
[2020b] that mean-field variational distributions are able to represent complex covariance structures in function space.
[Table]
LAPLACE [Immer et al., 2020]
[Table]
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 neural networks 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 approach proposed 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-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, and datasets used, see Appendix E.
For a comparison to Sun et al.
[2019] on small-scale regression problems, 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 significantly 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 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 metrics that elucidate the reliability of different methods’ predictive uncertainty under distribution shift.
[Table]
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 distributions that reflect our knowledge about the data-generating process can significantly improve uncertainty quantification 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 significantly 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-confidence 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
[Figure]
u e P
o s i
t i v e
R a t
[Figure]
decreases the least, while FSVI ’s uncertainty is significantly higher than other models’.
To assess the reliability of different uncertainty quantification 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—defined as moderate diabetic retinopathy or worse (classes { 2 , 3 , 4 } )—and non-sight-threatening diabetic retinopathy—defined 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 E.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 quantification 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, and we thank Joost van Amersfoort for helpful discussions on 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.
Table
of
Contents A Proofs & Derivations
A.1 Function-Space Variational Objective . . . . . . . . . . . . . . . . . . . . . . . .
A.2 Distribution under Linearized Function Mapping . .
B Further Empirical Results
B.1 Tabular Results for Diabetic Retinopathy Diagnosis Tasks . . . . . . . . . . . . . .
B.2 UCI Regression . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . .
C Illustrative Examples
C.1 Two Moons Classification Task . . . . . . . . . . . . . . . . . . . . . . . . . . . .
C.2 Synthetic 1D Regression Datasets . .
D
Sensitivity
Studies
D.1 Ablation Study on the Effect of Different Function-Space Priors . . . . . . . . . .
D.2 Context Distribution . .
D.3 Prior Variance . .
D.4 Number of Sets of Context Points . . . . . . . . . . . . . . . . . . . . . . . . . .
D.5 Number of Context Points . . . . . . . . . . . . . . . . . . . . . . . . . . . . . .
E Implementation, Training, and Evaluation Details
E.1 Hyperparameter Selection Protocol . .
E.2 FashionMNIST vs. MNIST/NotMNIST . . . . . . . . . . . . . . . . . . . . . . .
E.
3 CIFAR-10 vs. SVHN
E
Diabetic Retinopathy Diagnosis
E
Two Moons
E.
6 1D
Regression . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . .
E.
7 Further Implementation Details
E.8 Selection of Context Distribution
E
Compute Resources
Appendix A Proofs & Derivations A.1 Function-Space Variational Objective
This proof follows steps from Matthews et al.
Consider measures ˆ P and P both of which define distributions over some function f , indexed by an infinite index set X .
Let D be a dataset and let X D denote a set of inputs and y D a set of targets.
Consider the measure-theoretic version of
Bayes’ Theorem [Schervish, 1995]:
[Equation]
where p X ( Y | f ) is the likelihood and p ( Y )
= (cid:82) R X p
X
( Y | f ) dP ( f ) is the marginal likelihood.
We assume that the likelihood function is evaluated at a finite subset of the index set X .
Denote by π C : R X
→ R C a projection function that takes a function and returns the same function, evaluated at a finite set of points C , so we can write
[Equation]
and similarly, the marginal likelihood becomes p ( y D ) =
(cid:82) p y
| f X ( y D | f X D ) d P X D ( f X D ) .
Now, considering the measure-theoretic version of the KL divergence between an approximating stochastic process Q and a posterior stochastic process ˆ P , we can write
[Equation]
where P is some prior stochastic process.
Now, we can apply the measure-theoretic Bayes’ Theorem to obtain
[Equation]
where dQ π dP π ( f ) is marginally consistent given the projection π .
Rearranging, we can get
[Equation]
Finally
[Equation]
this lower bound can equivalently be expressed as
[Equation]
where X \D is an infinite index set excluding the finite index set X D
[Equation]
that is, X \D
∩ X D = ∅
[Equation]
or by
Theorem 1 in Sun et al.
[Equation] we can write
[Equation]
, (A.11) 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.
A.2 Distribution under Linearized Function Mapping
Proposition 1 (Distribution under Linearized Mapping) .
For a stochastic function f ( · ; Θ ) defined in terms of stochastic parameters Θ distributed according to distribution g Θ with m ˙= E g Θ [ Θ ] and
S ˙= Cov g Θ [ Θ ] , 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 .
Then the mean and co-variance of the distribution over the linearized mapping ˜ f at X , X ′ ∈
X are given by
[Equation]
Proof.
We wish to find E [ ˜ f
( X ; m , Θ )] and
[Equation]
To see that E [ ˜ f
( X ; m , θ )]
= f
( X ; m ) , note that, by linearity of expectation, we have
[Equation]
To see that Cov ( ˜ f
( X ; m , θ ) , ˜
[Equation]
( X ′
[Equation] [Equation]
; m ) ⊤ , note that in general, for a multivariate random variable Z ,
Cov
( Z , Z )
=
E [ ZZ ⊤ ] + E
[ Z ] E [ Z ] ⊤ , and hence,
Cov ( ˜ f
( X ; m , Θ ) , ˜ f
( X ′ ; m , Θ )) =
E [
[Equation] [Equation]
; m ,
[Equation] [Equation]
[Equation] [Equation]
(A.14)
We already know that E [ ˜ f
( X ; Θ )]
= f
( X ; m ) , so we only need to find E [ ˜ f
( X ; Θ ) ˜ f
( X ′ ; Θ ) ⊤ ] :
E g Θ [ ˜ f
( X ; m , Θ ) ˜
f
( X ′ ; m , Θ ) ⊤ ] =
E g Θ
[( f ( X ; m )
+ J
( X ; m )( Θ − m ))
( f ( X ′ ; m )
+ J
( X ′ ; m )( Θ − m ))
⊤ ] (A.15) =
E g Θ [ f ( X ; m ) f
( X ′ ; m ) ⊤ +
( J ( X ; m )( Θ − m )) ( J ( X ′ ; m )( Θ − m ))
⊤ + f
( X ; m )( J ( X ′ ; m )( Θ − m ))
⊤ + J
( X ; m )( Θ − m )
f ( X ′ ; m ) ⊤ ] (A.16) =
E g Θ [ f ( X ; m ) f
( X ′ ; m ) ⊤ + J
( X ; m )( Θ − m )
( Θ − m )
⊤ J ( X ′ ; m ) ⊤ + f
( X ; m )( J ( X ′ ; m )( Θ − m ))
⊤ + J
( X ; m )( Θ − m )
f ( X ′ ; m ) ⊤ ] (A.17) =
f ( X ; m ) f
( X ′ ; m ) ⊤ + J
( X ; m ) E g Θ
[( Θ − m )
( Θ − m )
⊤ ] J ( X ′ ; m ) ⊤ + f
( X
[Equation]
where the last line follows from the definition of g Θ .
By definition of the variance
[Equation]
we then obtain
[Equation]
With this result
[Equation]
we obtain the covariance function
[Equation]
Finally, Cov ( Θ ) =
S yields Cov
[Equation]
Θ ) ,
˜ f
( X ′ ; m , Θ )) =
J ( X ; m ) S J ( X ′ ; m ) ⊤ .
This concludes the proof.
Proposition 2 (Approximate Distribution under Linearized Mapping) .
For a stochastic function f ( · ; Θ ) defined in terms of stochastic parameters Θ distributed according to distribution g Θ = N ( m , S ) , 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 .
Then, for a partition of the set of parameters into sets α and β , a distribution g Θ =
N ( m , S ) with Θ α ⊥ Θ β , the distribution ˜ g ˜ f
( X ; Θ ) can be approximated via the Monte Carlo estimator
[Equation]
A.25) where g Θ α = N ( m α , S α ) , g Θ β = N ( m β , S β ) , and ˜
[Equation]
A
[Equation]
26) with J α ( · ; m ) denoting the columns of the Jacobian matrix corresponding to the sets of parameters α and ˜ f
α ( X ; m , Θ α ) ( j ) for j = 1 , ...,
R obtained by sampling parameters from the distribution g Θ α = N
( m α , S α ) .
Proof.
Consider a partition of the set of parameters into sets α and β and express the linearized mapping as
[Equation]
with
[Equation]
and
[Equation]
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 (A.27) expresses ˜ f
as a sum of (affine 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, for Z = X + Y ,
[Equation]
Letting X =
˜
f α ( X ; m , Θ α ) , Y = ˜
f β ( X ; m , Θ β ) , and X = ˜ f
( X ; Θ ) , we can write ˜
[Equation]
X
[Equation]
[Equation] [Equation]
with
[Equation] [Equation]
[Equation] [Equation]
and
[Equation] [Equation]
where we have used the fact that for a Gaussian distribution with mean m and covariance S ,
N ( z − y ; m, S ) =
N ( z ; m + y, S ) .
We can then approximate the probability density function ˜ g ˜ f
[Equation]
X ; m , Θ ) ( ˜ f
( X ; θ )) via the Monte Carlo estimator
[Equation]
with ˜ f
α ( X ; m , θ α ) ( j ) ∼ ˜ g ˜ f
α ( X ; m , Θ α ) .
Finally, we can express ˆ˜ g ˜ f
( X ; m , Θ ) ( ˜ f
( X ; θ )) equivalently as
[Equation]
where g Θ β = N ( m β , S β ) and samples ˜
f α ( X ; m , Θ α ) ( j ) are obtained by sampling parameters from the distribution g Θ α = N ( m α , S α ) .
This concludes the proof.
Proposition 3 (Covariance of Distribution over Linearized Mapping Induced by Distribution over
Parameters with Mean Zero) .
Let σ : R m → R n be an elementwise mapping with σ (0) = 0 , let f : X
× R P
→ R Q be a mapping that can be decomposed as
[Equation]
where Φ : X × R P ′
→ R P ′′ and Θ W,b = vec ( { Θ Φ W,b
[Equation]
and let the Jacobian of f evaluated at m ∈ R P be denoted by J ( · ; m ) ˙= ∂f ( · ; Θ )
∂ Θ | Θ
= m .
Then the covariance matrix of f evaluated x ∈ R D and m = 0 induced by a distribution over the network parameters with variance Σ is equal to the identity operator, that is, for output dimension j
[Equation]
where Σ linear b,j is a variance parameter associated with the final-layer bias for output dimension j .
Proof.
The mapping J ( x ; m ) gives the change in f at evaluation points x and m for infinitesimally small changes in each m i .
If m = 0 , an increase of ϵ in any m i will not change the image of f at x , since for any m i
=
ϵ there exists some m j =
with i ̸
=
j that is multiplied with h ( m i ) (for any mapping h ); the only exception being the bias term Θ linear b .
Since ∂ Θ linear b ∂ Θ linear b
[Equation]
where 1 ≤ k ≤ Q , [ J ( x ; 0 )]
k is the k th row of J ( x ; 0 ) , and
(cid:104) ∂
Θ linear b ∂ Θ linear b (cid:105) k is the k th row of ∂ Θ linear b ∂ Θ linear b =
I Q , we get
[Equation]
and hence, for output dimension j ,
[Equation]
where Σ linear b,j is a variance parameter associated with the final-layer bias for output dimension j .
For
X containing n evaluation points, we get
[Equation]
Appendix B Further Empirical Results B.1 Tabular Results for Diabetic Retinopathy Diagnosis Tasks
The results below were reproduced from Band et al.
[2021] using the RETINA benchmark.
[Table]
B.2 UCI Regression
[Table]
Appendix C Illustrative Examples C.1 Two Moons Classification Task
−−   −   0 .0 .0 .0 .0 .0 .0 .
E [ y | D ; x ] (a) FSVI :
Posterior Predictive Mean −−   −   0 .0 .0 .0 .0 .0 .0 .0 .0 .
V [ y | D ; x ] (b) FSVI :
Posterior Predictive Variance −−   −   0 .0 .0 .0 .0 .0 .0 .
E [ y | D ; x ] (c) MFVI :
Posterior Predictive Mean −−   −   0 .0 .0 .0 .0 .0 .0 .0 .0 .
V [ y | D ; x ] (d) MFVI :
Posterior Predictive Variance −−   −   0 .0 .0 .0 .0 .0 .0 .
E [ y | D ; x ] (e) MAP Ensemble: Predictive Mean −−   −   0 .0 .0 .0 .0 .0 .0 .0 .0 .
V [ y | D ; x ] (
[Figure] [Figure]
: Predictive Variance
(Figure 5c and Figure 5d), and an ensemble of MAP models (Figure 5e and Figure 5f).
The predictive means represent the expected class probabilities and the predictive variance the model’s epistemic uncertainty over the class probabilities.
With FSVI , the predictive distribution is able to faithfully capture the geometry of the data manifold and exhibits high uncertainty over the class probabilities in areas of the data space of which the data is not informative.
In contrast, neither MFVI , nor MAP ensembles are unable to accurately capture the geometry of the data manifold only exhibit high uncertainty around the decision boundary.
C.
−−−    −−
Predictive Mean
Function Draw
Training Data −−   −
Predictive Mean
Function Draw
Training Data (a) “Snelson” Dataset (Snelson and Ghahramani [2006])
−−−    −−
Predictive Mean
Function Draw
Training Data −−   −
Predictive Mean
Function Draw
Training Data (b) “OAT-1D” Dataset (van Amersfoort et al.
−−−    −−
Predictive Mean
Function Draw
Training Data −−   −
Predictive Mean
Function Draw
Training Data (c) “Subspace Inference” Dataset (Izmailov et al.
Appendix D Sensitivity Studies
We performed a comprehensive set of sensitivity studies to elucidate how different model and variational parameters affect the resulting predictive distributions.
Specifically, we used a validation holdout set (10% of the training data) to assess the effect of the prior covariance, the number of context points used per gradient steps, the context distribution, and the number of Monte Carlo samples for evaluating the expected log-likelihood on in- and out-of-distribution performance metrics.
All results below are obtained from ten random seeds.
D.1 Ablation Study on the Effect of Different Function-Space Priors
To better understand FSVI , we investigate what hyperparameter choice leads to good predictive performance and reliable out-of-distribution predictive uncertainty.
Figure 7 and Figure 8 show plots that demonstrate the link between different choices of prior variance and the resulting OOD-AUROC, test ECE
[Figure]
and test log-likelihood for FashionMNIST.
As can be seen in Figure 7b, test ECE is lowest for a prior variance of Σ 0 = 10 , while Figure 7a show that OOD-AUROC, test ECE
[Figure]
and test log-likelihood are highly negatively correlated.
This insight is useful, since it means that in real-world settings, where there are no real out-of-distribution validation sets, one can choose the prior variance that minimize test ECE and test negative log-likelihood.
We follow this approach to select the hyperparameters for Table 1 and Table
[Figure]
), the negative log-likelihood (NLL), expected calibration error (ECE), and the prior variance.
OODAUROC is very negatively correlated with both NLL and ECE.
For further results, see Appendix B.
[Figure]
D.2 Context Distribution
[Figure]
D.3 Prior Variance
[Figure]
models trained on CIFAR-10.
D.4 Number of Sets of Context Points
[Figure]
NLL 0 .0 .0 .0 .0 .0 .
ECE
[Figure]
TotalUncertainty)
[Figure] [Figure]
[Figure] [Figure]
[Figure] [Figure]
[Figure] [Figure]
[Figure] [Figure]
[Figure] [Figure]
NLL
[Figure] [Figure]
[Figure] [Figure]
[Figure] [Figure]
AUROC(
[Figure]
[Figure] [Figure]
[Figure] [Figure]
[Table]
Appendix E Implementation, Training, and Evaluation Details E.1 Hyperparameter Selection Protocol
For FSVI , we used a holdout validation set (10% of the training set) to conduct a hyperparameter search over the prior variance, the number of context points used to evaluate the KL divergence, the context distribution, and the number of Monte Carlo samples used to evaluate the expected log-likelihood.
The full results can be found in Appendix D. We selected the set of hyperparameters that yielded the lowest validation log-likelihood for all experiments.
We state the hyperparameters selected for the different datasets below.
We used 5 Monte Carlo samples and 10 context points per gradient step for all of our experiments.
For other methods, we used a holdout validation set of the same size and selected the best-performing hyperparameters.
We used implementations provided by the authors of MFVI (radial) and SWAG .
All other methods were implemented from scratch unless stated otherwise.
E.2 FashionMNIST vs. MNIST/NotMNIST
We train all model on the FashionMNIST dataset and evaluate the models’ predictive uncertainty performance on out-of-distribution data on the MNIST dataset.
Both datasets consist of images of size 28 × 28 pixels.
The FashionMNIST dataset is normalized to have zero mean and a standard deviation of one.
The MNIST dataset is normalized with the same transformation, that is, using the same mean and standard deviation used for the in-distribution data.
We chose FashionMNIST/MNIST instead of MNIST/notMNIST because the latter is notably easier than the former.
In this experiment, a network architecture with two convolutional layers of 32 and 64 3 × 3 filters and a fully-connected final layer of 128 hidden units is used.
A max pooling operation is placed after each convolutional layer and ReLU activations are used.
We do not use batch normalization.
All models are trained for 30 epochs with a mini-batch size of 200 and using the Adam optimizer with a learning rate 5 × 10 − 4 .
For FSVI with p X C = random monochrome, we used a prior variance of Σ 0 = 10 and sampled 50% of the context points for each gradient step from the mini-batch and the other 50% according to the method described in Appendix E.8.
For FSVI with p X C = KNIST, we used the KMNIST dataset.
E.3 CIFAR-10 vs. SVHN
We train all model on the CIFAR-10 dataset and evaluate the models’ predictive uncertainty performance on out-of-distribution data on the SVHN dataset.
Both datasets consist of images of size 32 × 32 × 3 , with RBG channels.
The CIFAR-10 dataset is normalized to have zero mean and a standard deviation of one.
The SVHN dataset is normalized with the same transformation, that is, using the same mean and standard deviation used for the in-distribution data.
The training data is augmented with random horizontal flips (with a probability of 0.5) and random crops (4 zero pixels on all sides).
In this experiment, a standard ResNet-18 network architecture is used.
We do not use batch normalization.
All models are trained for 50 epochs and using the Adam optimizer with a mini-batch size of 128 and a learning rate 5 × 10 − 4 .
For FSVI with p X C = random monochrome, we used a prior variance of Σ 0 = 0 .
1 and sampled 100% of the context points for each gradient step from the mini-batch and the other 50% according to the method described in Appendix E.8.
For FSVI with p X C = CIFAR-100, we used the CIFARdataset.
E.4 Diabetic Retinopathy Diagnosis
Prediction and Expert Referral.
In real-world settings where the evaluation data may be sampled from a shifted distribution, incorrect predictions may become increasingly likely.
To account for that possibility, predictive uncertainty estimates can be used to identify datapoints where the likelihood of an incorrect prediction is particularly high and refer them for further review.
We consider a corresponding selective prediction task, where the predictive performance of a given model is evaluated for varying expert referral rates.
That is, for a given referral rate of γ ∈ [0 , 1] , a model’s predictive uncertainty is used to identify the γ proportion of images in the evaluation set for which the model’s predictions are most uncertain.
Those images are referred to a medical professional for further review, and the model is assessed on its predictions on the remaining (1 − γ ) proportion of images.
By repeating this process for all possible referral rates and assessing the model’s predictive performance on the retained images, we estimate how reliable it would be in a safety-critical downstream task, where predictive uncertainty estimates are used in conjunction with human expertise to avoid harmful predictions.
Importantly, selective prediction tolerates out-of-distribution examples.
For example, even if unfamiliar features appear in certain images, a model with reliable uncertainty estimates will perform better in selective prediction by assigning these images high epistemic (and predictive) uncertainty, therefore referring them to an expert at a lower γ .
For all methods, experiments are performed using a ResNet-50 network architecture.
Training and evaluation scripts as well as model checkpoints can be found at github.com/google/uncertainty-baselines/.../diabetic_retinopathy_detection .
E.5 Two Moons
In this experiment, a network architecture with with two fully-connected layer with 30 hidden units each is used.
We train all models with a learning rate of 10 − 3 .
For FSVI , we used a prior variance of Σ 0 = 10 and sampled context points randomly from [ − 10 , 10] 2 .
E.6 1D Regression
In this experiment, we use a model consisting of two fully-connected layers with 100 hidden units each and Tanh activations.
We train all models with a learning rate of 10 − 3 .
For FSVI , we used a prior variance of Σ 0 = 10 and sampled context points randomly from [ − 10 , 10] .
E.7 Further Implementation Details
We use the Adam optimizer with default settings of β 1 = 0 .
9 , β 2 = 0 .
99 and ϵ = 10 − 8 for all experiments.
The deterministic neural networks that were used for the ensemble were trained with a weight decay of λ = 1e-1.
MFVI (tempered) was trained with a KL scaling factor of 0.1 to obtain a cold posterior.
E.8 Selection of Context Distribution
We estimate the supremum at every gradient step by sampling a set of context points X C from a distribution p X C at every gradient step.
For tasks with image inputs, we construct a distribution p X C , defined as a uniform distribution over images with monochromatic channels.
To generate a sample from this “monochrome images” distribution, we first take all images in the training data, flatten each channel, and stack the flattened image channels into a single vector each.
We then draw a random element (i.e., a pixel) from each channel vector and then use these pixels to generate a monochrome image of a given resolution by setting every channel equal to the value of the pixel that was drawn.
For regression tasks with a D -dimensional input space, p X C is defined as a uniform distribution with lower and upper bounds set to the empirical lower and upper bounds of the training data.
For further details on the effect of different sampling schemes on the posterior predictive distribution’s performance, see Appendix B.
E.9 Compute Resources
All experiments were carried out on an Nvidia V-100 GPU with 32GB of memory.