Tractable Function-Space Variational Inference in Bayesian Neural Networks
Authors: Tim G. J. Rudner
Zonghao Chen
Yee Whye Teh
Yarin Gal
Venue: NeurIPS 2022 Accept (2022)
Keywords: Variational Inference
Bayesian Neural Networks
Uncertainty Quantification

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.

Tractable Function-Space Variational Inference in Bayesian Neural Networks
Tim G. Rudner"
Zonghao Chen University College London
Yee Whye Teh University of Oxford
Yarin Gal University of Oxford
University of Oxford
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 to define a prior distribution over the network parameters, infer approximate posterior distri- bution, and use it to make stochastic predictions. However, explicit inference over neural network parameters makes it difficult incorporate meaningful prior information about the data-generating process into the model. In this paper, we pursue 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, explicitly as inferring a posterior distribution over functions and propose scal- able 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 sfate-of-the-ar art uncertainty estimation and predictive performance on a range of prediction tasks and demonstrate that it performs well challenging safety-critical medical diagnosis task in which reliable uncertainty estimation is essential.
frame Bayesian inference in neural networks
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, a02]a.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 (BNNs, Neal [1996]). While BNNs have promised to combine the advantages of deep learning and Bayesian inference, existing approaches for approximate inference BNNs fall short of this promise and have been demonstrated to result in approximate posterior predictive distributions that underperform noon-Baesiann methods both in terms of predictive accuracy and uncertainty quantification makina 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-genn process into inference.
*Corresponding author. Email: <tim. wrim..udnerrcc..cmm ac. uk>
36th Conference on Neural Information Processing Systems (NcurIPS 2022).

---

021
0.
0.18
9.15
Prstettie Mon
0:30
Fraetion Drn
Tharing Tate
(a) Predictive Distribution
(b) Predictive Mean
(c) Predictive Variance
Figure 1: ID regression on the Snelxon dataset and binary classification on the Two Moons dataset. The plots show the predictive distributions of a BNNs, obtained via function-space variational inference (FSVI). For further illustrative exampled and comparisons to deep ensembles and BNNs learned via parameter-space variational inference, Appendix
To avoid this limitation, we follow 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 function-space variational objectives to used with high-dimensional inputs and beghly-overparameltriied neural networks, propose a simple estimator of the Kullback -allback-Leibler divergence between distributions over functions that enables to perform 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 hand.
We demonstrate that this approach leads to posterior approximations that exhibit significantly im- proved predictive uncertainty estimates compared to wide array of state-of-the-art Bayesian and non-Bayesian methods. Figure shows examples of predictive distributions obtained via function- space variational inference on low-dimensional, casy-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-spact varia- tional inference in BNNs. 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 thorough empirical evaluation in which compare the pro- posed approach to a wide array of competing methods and show that it consistently results high predictive performance and reliable predictive uncertainty estimates, outperforming other methods in terms of predictive accuracy, robustness to distribution shifts, and uncertainty- detection of curriutiionallyssiifee
data samples. We evaluate the proposed method standard benchmarking datasets as well as on a safety-critical medical diagnosis task in which reliable uncertainty estimation essential."
2 Preliminaries
We consider supervised learning tasks on data D (n)). (XD, with inputs X, E X C DR and targets yn where y R for regression and {0, 1) for classification lasks. Bayesian neural networks (BNNS) are stochastic neural networks trained using (approximate) Bayesian inference. Denoting the parameters of such stochastic neural network by the multivariate random variable θ E R and letting the function mapping defined by a neural network architecture be given by f X' R² + R we obtain random function f(- : -). For parameter realization 0, we obtain corresponding function realization, f(: 0). When evaluated at a finite collection of points | f(X; +) is multivariate random variable and f fX: 0) vector.
Letting Py)s(x:e) be a likelihood function and (y) (x; 0)) be the likelihood of ob- serving the targets yp under the stochastic function f(-; θ) evaluated at inputs X and letting pe use Bayes' Theorem to find the posterior distribution, PeyD [MacKay, 1992, Neal, 1996]. However, since the mapping f
be prior distribution the stochastic network parameters θ,
Our code can be accessed at http:://ggttub..mmmmmnnnneeeeer http./////mmmuunn////////

---

is 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 variational optimization prohlem, where the goal is to find distribution ge in vari- ational family 2g0 that solves the variational problem min TxxL. (ge) || Pojp) [Wainwright and Jordan, 2008]. If 2ne is the family of mean-field Gaussian distributions and the prior dis- tribution over parameters pe 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 (())) induced by the posterior distribution over parameters pe pp. that is,
Ps( P)( DD) [
Pelp (D) 8(f(-;)) f();
(1)
where 8(-) is the Dirac delta function [Wolpert, 1993]. Considering the prior distribution over functions :)(-)) induced by prior distribution over parameters pe,
{ po(0')((((: [(f(-; f(:;0(())))))) do',
Ps( (f()) p)(-,((((:;0)
(2)
and the variational distribution over functions (f(-()) induced by variational distribution over parameters ge.
f ()(:;())))) JRP
ge (0') 8(f(-;)) - -f(+))))
(3)
can express the problem of finding a posterior distribution over functions variationally as min ge g
(KI.( ())()))))),
(4)
which allows us effectively incorporate meaningful prior information about the underlying data- generating process into training. As discussed by Burt et al. [2021], this variational objective is guaranteed be well-defined for suitably chosen prior distributions over functions. Specifically, the KI. 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].
(5) As result, (ge) || pe) < oo, which the case for diiee-dimensoonll parameter vectors and
DK (X)((((::)) P)(=))) S DKL pP).
ge absolutely continuous with respect to pe. then the function-spar 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 yo and a suitably defined prior distribution functions, well-defined maximization problem max EQ. F(go) ) with
express the variational problem above equivalently the
F(1e)= = Byrxx, [logry Py s(x,)) (yD [(x:; 0)) (x)( || P)()))),
(6)
where [TK (9)(.r)) |l p)()))) also KI. divergence hetween distributions over functions.
Unfortunately, evaluating the KL divergence in Equation (6) is in general intractable for arbitrary mappings f. To obtain tractable objective, Sun et al. [2019] showed that (x || ()))))) can be expressed as the supremum of the KL divergence from ()(:)) to P)(:0) over all finite sets of evaluation points, resulting in the objective function
F(ge) Ey [log Py).x.e)
f(x:: 0))] sup MDx (1)(, |)(x)))
(7)
Xe
where X' = U f E Xr X, S RaxD} is the collection of all finite sets evaluation points. However, this objective function is still challenging 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 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.
3

---

m) 3 Deriving a Tractable Function-Space Variational Objective
The primary obstacle to computing the objective in Equation (6) is the KL divergence from (=(=e) to ()(:;)). There are two reasons why the KL divergence in Equation (7) is intractable: First, for BNNS or other non-linear models, we do not have access to the probability density functions of the multivariate distributions ()) and () second, for all but extremely simple input spaces, we unable 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 KI. divergence between two BNNs evaluated at finite set of points. To do so, we first derive tractable approximations to the distributions over functions () and e) Next, we show that under these approximations, we are able to obtain flosed-form approximation to the KL divergence and describe 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 ()(:e)) and p)(x,))- we use a first- order Taylor expansion of the mapping / about the mean parameters of qe and pe, respectively, and derive the induced distributions under the linearized mapping.
For stochastic function f(-; ()) defined in terms of stochastic parameters O distributed according to
distribution ge with m Eg [+) and $ = Cov [+], we denote the lincarization of the stochastic
function f(:; ① about by
(8) where m-; = (if(: /0)) @-m is the Jacobian of f(:: evaluated at O = m, and the mean and covariance of the distribution over the linearized mapping f at X, X' X are given by
f(:: = f(-: m. θ) f(::m(())))))))))))))()) [+::
EIf(X; = f(X;
(ovl 0), f(X', θ)] = m)S(XX m)
(10)
For derivation of this result, see Appendix A. Since Gaussianity is preserved under affine transfor- mations, if ge is a multivariate Gaussian distribution with mean m and diagonal co- co-variance S, then the distribution over f(X; is given by
Of(Xm..) .) =N) (X))))
(11)
For stochastic functions parameterized by many millions of parameters, obtaining the covariance of ff(x-e)- wwhich requires computing an inner product of two Jacobian matrices-cal be computa- tionally expensive. Instead of computing the distribution over the linearized mapping exactly, can construct a suitable Monte Carlo estimator. To do so, we consider partition of the set of parameters into sets a and / (with |3]| < |a)) and note that the lincarized mapping can then be expressed as f(: m. θ) = (m) + fa @a) + Ja(-: ()())))@ - -m3), m,
(12)
with
fa m, [-) - ma),
(13)
where (-; m) and Js(-; m) are the columns of the Jacobian matrix corresponding to the sets of parameters and 8, respectively, and O., and O, are the corresponding random parameter vectors. Noting that Equation (12) expresses ff as sum of (affine transformations of) random variables, we can use the fact that for independent Gaussian random variables and Y, the distribution hzz of Z = X Y is equal to the convolution of the distributions fix and fy to obtain an approximation to f. In particular, we can show that if ge a multivariate Gaussian distribution with θ., Og, the distribution 0f(x;0) can be approximated by the Monte Carlo estimator
(14) are obtained by sampling parameters from the distribution ge., = A (m.. S.). For a derivation of this result, see Appendix A. This estimator is biased for finite K but converges (m(X,m,))
ff(x...9) = Σne R
R ",=
fa (X; m, ).)()), "TTT,,Ts (X; ())ss
where ge, = =N(m, S3) and samples (x; m, ².,)
R - co. Similarly, for finite R, the smaller [S..]i. the more accurate and less biased the estimator will be. In our empirical evaluation, we use single Monte Carlo sample, R = 1, to preserve Gaussianity and choose a 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 qe and pe are both Gaussian distributions, then the induced distri- butions 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 O we can obtain locally accurate approximations to the induced distributions (-:))) to P)()))) and use them to approximate the KL divergence in the variational objective by (K (ff(x:)) ll ()(x;))) Moreover, for an isotropic Gaussian prior and a mean-field Gaussian variational distribution, DKL ((f(xe) || Pf(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 F(ge)= = B | loggyyl/yx.)))//0000 (X;: 0))]
sup Tx (xx((((x.)) ((jx)))))
(15)
xeX:
where the arguments of the KL divergence have been been replaced by the (locally accurate) approximations to the variational and prior distributions over functions evaluated X. respectively. Since the stochastic functions/ f(- ( induced by ge and pe under the linearized mapping will be closer to the stochastic function under f the smaller the variance of (e and pe, respectively, the approximation to the KL divergence will be more accurate the smaller the variance of qe and pe.
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 E Ax but instead propose to estimate the supremum every gradient step via simple finite-sample estimator. Specifically, letting / (XX = Dxi (ūf(xe)) (j(x,))), we estimate G supxex: ((X)
using the Monte Carlo estimator
G(X8)= = max I(X), XEX
(16)
where N3 - (x )= collection of S sets of context points X (x6) jointly sampled
from a context distribution px:- Each context set (6) can be viewed as single Monte Carlo sample from the input space so that the estimator ((x))) 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 the sets of context points. The choice of the context distribution px, can be informed by knowledge about the prediction task and should be viewed as a problem- wroblem-specific modeling choice. Similarly, the numbers of samples S and K are hyperparameters to be optimized with validation set. For details on how Px, is chosen for the empirical evaluation in Section 5, see Appendix D.
3.3 Stochastic Estimation of the Approximnate Function-Space Variational Objective
Let de be a Gaussian mean-field field variational distribution, let pe be an isotropic Gaussian prior, let (Xg, YB) be mini-batch of the training data, and reparameterize O as θ Σ, e())) = + E
Using the estimator ((Z)) defined above and estimating the expected log- likelhoolh via Monte Carlo sampling, we obtain a Monte Carlo estimator for the function-space variational objective:
F(u, = M Σ logPyl/(xx)) f(X E, E,()))
max DRI. ((j((((())) (() X
(17)
with e0) ~ (0, Ip) and X defined above. This Monte Carlo estimator is biased due to the lincarization and context-set approximations but allows for scalable gradient-based stochastic optimization.
Selection of Prior. For all experiments that involve uncertainty quantification, we chose prior distribution over parameters that induces prior distribution over functions Ps(-))) and a prior predictive distribution that exhibit a high degree of predictive uncertainty at evaluation points from regions in input space where Pxe has non-zero support and, under smoothness constraints, on evaluation points in nearby regions. For settings where prior information encoded in data data-for example, in the form of expert demonstrations of robotic manipulation tasks [Rudner et al., 2021] or 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 D.
5

---

Selection of Context Distribution. The distribution Px, allows us to incorporate information about the data- Mata-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 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 FashionMNIST, 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 ge. we use the fact that we can obtain function draws by sampling from the distribution over parameters to obtain an approximate posterior predictive distribution
aly. x.) / ((x.1 ()) df(x.;
(18)
1 " ply. [(x.; 0()) with θ
(())
qe:
M.
where M, is the number of Monte Carlo samples used estimate the predictive distribution.
Related Work
There is growing body of work on function-space approaches inference in BNNS, 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 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 BNNS are based on approximate gradient estimators and either replace the supremum in Equation (7) with expectation [Sun et al., 2019] or do not define 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.1 and Burt et al. [2021]]. More recent work has attempted to circumvent the intractability of the variational objective in Equation (6) by proposing alternative objectives for function-spact spacoo inference in BNNS [Ma et al., 2019, Ober and Aitchison, 2020, Ma and Hernández-L 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 Seeeralieed-Gaus--eeeeee
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 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 BNNs with different networks architectures are not well-defined. parallel line of research showed that posterior predictive distributions of shallow BNNs 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 BNNs 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 1: Comparison of in- and out-of-distribution performance metrics on FashionMNIST (mean + standard error over random seeds). The last two columns show the AUROC for binary in- vs. out-of-distribution detection on MNIST (M) and NotMNIST (NM). MNIST and NoMMIIST are used as out-of-distributoon datasets. Best overall results for single and ensemble models are printed in boldface with gray shading. Results within a 95% confidence interval of the best overall result are printed in boldface only. All methods use the same four-layer CNN architecture. For further details about model architectures and training and evaluation protocols, see A Appendix D.
Method
Accuracy
ECE +
AUROC M
AUROC NM
MAP
91.73
0.037:100001
87.00
14
MFVI [Blundell et al., 2015]
91.03+00.4
0.038
93. 10.4.
88. +8.78
MFVI (tempered)
+0.08
0.058++0.001
86,30 06.30
80.78+b..6
MFVI (radial) [Farquhar et al., 2020a]
90.31
0.035+0.001
84.40#0..8
82. 1:1.15
MC DROPOUT [Gal and Ghahramani, 2016]
90.55
0.012:0.00 0010
88.46
80.02
SWAG [Maddox et al., 2019]
92.56:10.05
0.043
18.10:0
80.31
DUQ [van Amersfoort et al., 2020]
92.40 02.40=00220
00.50
94.60
BNN-LAPLLCC [Immer et al., 2020]
92.25
0.01220.003
95.55=0.00
SPG [Ma and Hernández-l.obute 2021]
91.60
95.60
FSVI (px_- random monochrome)
93.13:0.13
0.012
96. 20:46
95.024:000
FSVI (px: KMNIST)
93. #0.12
0.010 0.0001
99. +0.20
97. 26.26...22
Deep Ensemble
92.49
0.019 .000
89.22#0.uu
83.17
ESVI Ensemble (px; random monochrome)
94.44=0.07
0.020++.0001
97.85+0.13
96.95+0.22
Empirical Evaluation
In this section, we evaluate FSVI on high-dimensional classification tasks that out of reach for function-space variational inference methods proposed in prior works and compare FSVI to several well-established and state-of-the-an Bayesian deep learning and deterministic uncertainty quantifi- cation methods. We show that FSVI (sometimes significantly) outperforms existing Bayesian and non-Bayesian methods in terms of their in-distribution uncertainty calibration and out-o--distiibution predictive uncertainty estimation. For a details on models, training and validation procedures, and datasets used, see Appendix D. For comparison to Sun et al. [2019] small-scale regression tasks, see Appendix B.2.
5.1 Predictive Performance, Uncertainty Estimation, and Distribution Shift Detection
In this set of experiments, assess the reliability of the uncertainty estimates generated by FSVI. If 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 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 10 dataset, we use the SVHN dataset as -uittf-diitributiin evaluation points.
For models trained either FashionMNIST CIFAR- 10, we evaluate their in-distribution perfor- mance in terms of test accuracy, test log-like likeelol,hood,, and test calibration. To evaluate the quality of different models' uncertainty estimates, we compute uncertainty estimates for the pairs Fash- COMMNIST/MNIST FashionMNNST/N
and CIFAR- 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. [2020]. We report the under the receiver operating characteristic (ROC) curve in Tables and 2.
to and measure for a range
Predictive Performance and Calibration. To assess in-distribution predictive performance and calibration, report the test accuracy, negative likellood (NLL), and expected calibration error (ECE) for models trained on FashionMNIST and CIFAR- 10 in Tables I and 2. On both FaShionMNIST and CIFAR- 10, FSVI achieves the lowest NLL and either the best or second-bes pre- dictive accuracy and ECE, respectively, across all methods. Notably, FSVI significantly outperforms SPG [Ma and -eemándezzLobato, 2021], alternative function-space variational inference method.
Predictive Uncertainty under Distribution Shift. Tables 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 dusttf-diittibutton inputs with high accuracy. As would be expected, we observe that using context
7

---

Table 2: Comparison of in- and out-of-distributioo performance metrics on CIFAR-10 (mean + standard enorr over ten random seeds). SVHN and corrupted CTFAR-10 (C-CIFAR) used as an out-of-distribution datasets. The penultimate column shows the AUROC for binary in- vs. out-of-distribution detection on SVHN. Best overall results for single and ensemble models are printed in boldface with gray shading. Results within a 95% confidence interval of the best overall result are printed in boldface only. All methods use ResNet- 18 architecture. For further details about model architectures and training and evaluation protocols, see Appendix D.
Method
Accuracyt
ECEJ
OOD-AUROCT C-CIFAR Acct
MAP
93.19.4..11
0.043
94.65-20.27
78.87
MFVI [Blundell et al., 2015]
89.98 09
0.040=0.001
92.14 24.34
79.36 19.36=1..3
MFVT (tempered)
90.87+0.11
0.048=0.001
91.82:20.90
79. 19.86=1.32
MC DROPOUT [Gal and Ghahramani, 2016]
93.55 00
0.040
92.44=0.57
80.13
SWAG [Maddox al., 2019]
93.13++0114
0.067 +v.002
89.79=0.60
12.
VOGN [Osawa et al., 2019]
84.27++0.20
0.040
87.60=0.20
DUQ [van Amersfoort et al., 2020]
94.10=0.20
92.70=1.30
SPG [Ma and Hernández-Lobato, 2021]
77.6910.64
88 88.30=4.00
FSVI (px; random monochrome)
93.35=0.04
0.03410.001
94.7610.22
80.81:1..4
FSVI (px; = CIFAR-I00)
93.57=0.04
0.026
98. 02.07.1 1:0.10
81.20:10.42
Deep Ensemble
95.13=0.05
0.01920.001
98.044 04.007
81.22::0.37
FSVI Ensemble (px; random monochrome)
95.19 20.00
0.013
99.1911..4
81.33 81.35=20.4
distributions that reflect our knowledge about the data-generating process can significantly improve uncertainty quantification under FSVI. For the MaSHionMNIST experiment, we used the KMNIST dataset, which contains grayscale images of Kuzushiji letters, and for the CIFAR-10 experiment, we used the CIFAR- 1-0 dataset, which contains RGB images of 100 classes. Both KMNNST and CIFAR-100 differ from the OOD datasets (MNIST and NotMNIST and SVHN, respectively) used to compute
metrics in Tables and 2. but using them context distributions significantly increased the ability of BNNS 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 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 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 -ut-of-distribution inputs, and (iii) if possible, maintain high predictive accuracy even under distribution shift. Models that satisfy these desiderata are less likely make poor, high-confidence predictions and more amenable for use in safety-critical downstream tasks.
1.4
To illustrate these desiderata, follow Ovadia et al. [2019] and consider the rotated MNIST task, where model is trained on MNIST and evaluated ro- tated 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 increas- ing magnitude. Figure 2 shows Brier scores (lower 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 in- crease (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 exhibit- ing a large increase in predictive uncertainty. As can decreases the least, while FSVI's uncertainty is significantly higher than other models'. To assess the
Briet Score
0.7
D5V1
MAP
Deep Ensemble
2.0
Entroper
1.0
0.0
30
60
90
120
150
180
Rotation Angle (*)
Figure 2: Predictive uncertainty and accuracy on rotated MNIST. Models with reliable uncertainty estimates would exhibit higher predictive uncer- tainty the more the digits are rotated. Ideally, such models would maintain high predictive accuracy (low Brier score).
seen in the plot, the Brier scores of FSVT
8

---

Deep Emsembl:
MC Dropout Esseuble
ESVI Busomble
Radal Ersemble
Runk-1 Ensemble
MPVI Eusemble
Diggout
Rank-1 BNN
True Positive Rate
True Positive Rate
58
12
24
02 Propection of Canse Refered Expert
False Positive Rate
False Positive Rate
Propertion of Canes Refered In Expen
(a) ROC: In-Domain
(b) ROC: Country Shift
(c) Selective Prediction AUROC: In-Domain
(d) Selective Prediction AUROC: Country Shift
Figure 4: We jointly assess model predictive performance and uncertainty quantification on both in do-maii and distributionally shifted data. Left: The receiver operating characteristic curve (ROC) for in- population diagnosis the (a) EyePACS [2015] test set and for (b) changing medical equipment and patient populations on the APTOS [2019] test set. The dot in black denotes the NHS-renodeddded 85% sensitivity and 80% specificity ratios [Widdowson, 2016]. Right: Selective prediction on AUROC in (c) EyePACS [2015] and (d) APTOS [2019] settings. Shading denotes standard error over six random seeds. See Appendix B.1 for tabular results.
task, reliability of different uncertainty quantification methods on a more challenging distribution-shiff
we consider corrupted CIFAR-10 inputs under the second-mildest corruption level used in [Ovadia al., 2019] and report our results Table 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 real- world safety-critical setting, we consider the task of diagnosing diabetic retinopathy (DR), medical condition that can lead to impaired vision, from retina scans [Leibig et al., 2017, Filos et al., 2019, 9, Band et al., 2021]. We use two publicly available datasets, EyePACS [2015] and APTOS [2019], each containing RGB images of human retina graded by med- ical expert on the following scale: 0 (no DR), I (mild DR), (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 AP- TOS datasets by dividing the classes up into sight-threatening diabetic retinopathy defined as moderate diabetic retinopathy worse (classes defined as (2,3, 4)) 4))))amd and non-sight-threatening diabetic reinopaaty-deffieed
Figure 3: Retina scan exam-
ples. Top: healthy, Bottom: unhealthy.
no or mild diabetic retinopathy (classes (0, 1}). This results in 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 al. [2021] and use the Kaggle dataset for training and the distributionally shifted APTOS dataset for evaluation. The results
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.
6 Conclusion
The paper proposed a scalable and effective approach to function-space variational inference in BNNs. 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 outto--distrbuution predictive performance on wide range of datasets when compared to well-established and -fate-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-drives 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.
References
Dario Amodei, Chris Olah, Jacob Steinhardt, Paul Christiano, John Schulman, and Dan Mané. Concrete problems in ai safety, 2016.
APTOS. APTOS 2019 Blindness Detection Dataset, 2019.
Neil Band, Tim G. J. Rudner, Qixuan Feng, Angelos Filos, Zachary Nado, Michael W. Dusenberry, Ghassen Jerfel, Dustin Tran, and Yarin Gal. Benchmarking Bayesian Deep Learning on Diabetic Retinopathy Detection Tasks. 2021.
Ari Benjamin, David Rolnick, and Konrad Kording. Measuring and regularizing networks in function space. In International Conference on Learning Representations, 2019.
Charles Blundell, Julien Cornebise, Koray Kavukcuoglu, and Daan Wierstra. Weight uncertainty in neural networks. volume 37 of Proceedings of Machine Learning Research, pages 1613-1622, Lille, France, 07-09 Jul 2015. PMLR.
David R. Burt, Sebastian W. Ober, Adrià Garriga-Alonso, and Mark van der Wilk. Understanding variational inference in function-space. In Third Symposium on Advances in Approximate Bayesian Inference, 2021.
Eduardo D. C. Carvalho, Ronald Clark, Andrea Nicastro, and Paul H. J. Kelly. Scalable uncertainty for computer vision with functional variational inference. In IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR), June 2020.
EyePACS. Diabetic Retinopathy Detection Dataset, 2015.
Sebastian Farquhar, Michael Osborne, and Yarin Gal. Radial Bayesian neural networks: Beyond discrete support in large-scale Bayesian deep learning. In Silvia Chiappa and Roberto Calandra, editors, Proceedings of the Twenty Third International Conference on Artificial Intelligence and Statistics, volume 108 of Proceedings of Machine Learning Research, pages 1352-1362 PMLR, 26-28 Aug 2020a.
Sebastian Farquhar, Lewis Smith, and Yarin Gal. Liberty or depth: Deep Bayesian neural nets do not need complex weight posterior approximations. In Advances in Neural Information Processing Systems 33: Annual Conference on Neural Information Processing Systems 2020, NeurlPS 2020,
December 6-12, 2020, virtual, 2020b.
Angelos Filos, Sebastian Farquhar, Aidan N. Gomez, Tim G. J. Rudner, Zachary Kenton, Lewis Smith, Milad Alizadeh, Arnoud de Kroon, and Yarin Gal. A systematic comparison of Bayesian deep learning robustness in diabetic retinopathy tasks, 2019.
Andrew Y. K. Foong, Yingzhen Li, José Miguel Hernández-Lobato, and Richard Turner. 'in- between' uncertainty in Bayesian neural networks, 2019.
Andrew Y. K. Foong, David R. Burt, Yingzhen Li, and Richard E. Turner. On the expressiveness of approximate inference in Bayesian neural networks. In Advances in Neural Information Processing Systems 33: Annual Conference on Neural Information Processing Systems 2020, NeurIPS 2020, December 6-12, 2020, virtual, 2020.
Yarin Gal and Zoubin Ghahramani. Dropout as a Bayesian approximation: Representing model uncertainty in deep learning. In Proceedings of the 33rd International Conference on International Conference on Machine Learning Volume 48, ICML 2016, pages 1050-1059, 2016.
Alex Graves. Practical variational inference for neural networks. In Proceedings of the 24th International Conference Neural Information Processing Systems, NIPS'11, 11, page 2348 2356. Red Hook, NY, USA, 2011, Curran Associates Inc. ISBN 9781618395993.
Dan Hendrycks, Nicholas Carlini, John Schulman, and Jacob Steinhardt. Unsolved problems in ml safcty, 2021.
10

---

Geoffrey E. Hinton and Drew van Camp. Keeping the neural networks simple by minimizing the description length of the weights. In Proceedings of the Sixth Anmual Conference on Computational Learning Theory, COLT '93, page 5-13, New York, NY, USA, 1993. Association for Computing Machinery. ISBN 0897916115.
Matthew D. Hoffman, David M. Blei, Chong Wang, and John Paisley, Stochastic variational inference. Journal of Machine Learning Research, 14(1):1033-144,, May 2013, ISSN 1532-4435.
Alexander Immer, Maciej Korzepa, and Matthias Bauer. Improving predictions Bayesian neural networks via local linearization, 2020.
Pavel Izmailov, Wesley J. Maddox, Polina Kirichenko, Timur Garipov, Dmitry Vetrov, and An- drew Gordon Wilson. Subspace inference for Bayesian deep learning. In Ryan P. Adams and Vibhav Gogate, editors, Proceedings The 35th Uncertainty Artificial Intelligence Conference, volume 115 of Proceedings of Machine Learning Research, pages 1169-1177 PMLR, 22-25 Jul 2020.
John Jumper, Richard Evans, Alexander Pritzel, Tim Green, Michael Figurnov, Olaf Ronneberger, Kathryn Tunyasuvunakool, Russ Bates, Augustin Zidek, Anna Potapenko, Alex Bridgland, Clemens Meyer, Simon A A Kohl, Andrew Ballard, Andrew Cowie, Bernardino Romera-Paredes, Stanislav Nikolov, Rishub Jain, Jonas Adler, Trevor Back, Stig Petersen, David Reiman, Ellen Clancy, Michal Zielinski, Martin Steinegger, Michalina Pacholska, Tamas Berghammer, Sebastian Bodenstein, David Silver. Oriol Vinyals, Andrew W Senior, Koray Kavukcuoglu, Pushmeet Kohli, and Demis Hassabis. Highly accurate protein structure prediction with AlphaFold. Nature, 596 (7873):583-58) 2021. doi: 10. 1888/441566001-00333
Mohammad Emtiyaz. E Khan, Alexander Immer, Ehsan Abedi, and Maciej Korzepa. Approximate inference turns deep networks into Gaussian processes. In Advances in Neural Information Processing Systems 32, pages 3094-3104. Curran Associates, Inc., 2019.
Alex Krizhevsky, Ilya Sutskever, and Geoffrey E. Hinton. ImageNet classification with deep con- volutional neural networks. Advances Neural Information Processing Systems 25:, pages 1106-1114, 2012.
Christian Leibig, Vanceda Allken, Murat Seçkin Ayhan, Philipp Berens, and Siegfried Wahl. Leverag- ing Uncertainty Information From Deep Neural Networks for Disease Detection. Nature Scientific Reports, 7(1):17816, 2017.
Chao Ma and José Miguel Hernández-Lobato Functional variational inference based on stochastic process generators. In A. Beygelzimer, Y. Dauphin, P. Liang, and J. Wortman Vaughan, editors, Advances in Neural Information Processing Systems, 2021.
Chao Ma. Yingzhen Li, and Jose Miguel Hernandez-Lobato Variational implicit processes. In Kamalika Chaudhuri and Ruslan Salakhutdinov, editors, Proceedings of the 36th International Conference Machine Learning, volume 97 of Proceedings of Machine Learning Research, pages 4222-4233. PMLR, 09-15 Jun 2019.
David J. C. MacKay, A practical Bayesian framework for backpropagation networks. Neural Comput., 4(3):448 May 1992. ISSN 0899-7667. doi: 10. 1662necoo. 1992.4.3.448.
Wesley Maddox, Pavel Izmailov, Timur Garipov, Dmitry P Vetrov, and Andrew Gordon Wilson. A simple baseline for Bayesian uncertainty in deep learning. In Advances in Neural Information Processing Systems, pages 13153-13164 2019.
Alexander G. de G. Matthews, James Hensman, Richard Turner, and Zoubin Ghahramani. On sparse variational methods and the Kullback-Leible divergence between stochastic processes. volume 51 of Proceedings of Machine Learning Research, pages 231-239, Cadiz, Spain, 09-11 1 May 2016. PMLR.
Volodymyr Mnih, Koray Kavukcuoglu, David Silver, Alex Graves, Ioannis Antonoglou, Daan Wierstra, and Martin Riedmiller. Playing atari with deep reinforcement learning. In NIPS Deep Learning Workshop. 2013.
Radford M Neal. Bayesian Learning for Neural Networks. 1996.
Sebastian W. Ober and Laurence Aitchison. Global inducing point variational posteriors for Bayesian neural networks and deep Gaussian processes, 2020.

---

Kazuki Osawa, Siddharth Swaroop, Mohammad Emtiyaz E Khan, Anirudh Jain, Runa Eschenhagen, Richard Turner, and Rio Yokota. Practical deep learning with Bayesian principles. In Advances in Neural Information Processing Systems, volume 32, pages 4287-4299. Curran Associates, Inc., 2019.
Yaniv Ovadia, Emily Fertig, Jie Ren, Zachary Nado, D. Sculley, Sebastian Nowozin, Joshua Dillon, Balaji Lakshminarayanan, and Jasper Snock. Can you trust your model's uncertainty? Evaluating predictive uncertainty under dataset shift. In Advances in Neural Information Processing Systems 32. 2019.
Pingbo Pan, Siddharth Swaroop, Alexander Immer, Runa Eschenhagen, Richard E. Turner, and Mohammad Emtiyaz Khan. Continual deep learning by functional regularisation of memorable past. In Advances in Neural Information Processing Systems 33: Annual Conference on Neural Information Processing Systems 2020, NeurIPS 2020, December 6-12, 2020, virtual, 2020.
Yury Polyanskiy and Yihong Wu. Strong data-processing inequalities for channels and Bayesian networks. In Eric Carlen, Mokshay Madiman, and Elisabeth M. Werner, editors, Convexity and Concentration, pages 211 1249, New York, NY, 2017. Springer New York. ISBN 978-1-4939-7005-
6.
Tim G. J. Rudner and Helen Toner. Key Concepts in AI Safety: An Overview. In CSET Issue Briefs, 2021a.
Tim G. J. Rudner and Helen Toner. Key Concepts in Al Safety: Robustness and Adversarial Examples. In CSET Issue Briefs, 2021b.
Tim G. J. Rudner, Cong Lu, Michael A. Osborne, Yarin Gal. and Yee whye Teh. On Pathologies in KL-Regularized Reinforcement Learning from Expert Demonstrations. 2021.
Tim G. J. Rudner, Freddie Bickford Smith, Qixuan Feng, Yee Whye Teh, and Yarin Gal. Continual Learning via Sequential Function-Space Variational Inference. In Proceedings of the 38th Inter- national Conference Machine Learning, Proceedings of Machine Learning Research. PMLR, 2022.
M. J. Schervish. Theory of Statistics. Springer-Verlag, New York, NY, 1995.
David Silver, Aja Huang, Chris J. Maddison, Arthur Guez, Laurent Sifre, George van den Driessche, Julian Schrittiesser, Joannis Antonogiou, Veda Panneershelvam, Marc Lanctot, Sander Dieleman, Dominik Grewe, John Nham, Nal Kalchbrenner, Ilya Sutskever, Timothy Lillicrap, Madeleine Leach, Koray Kavukcuoglu, Thore Graepel, and Demis Hassabis. Mastering the game of Go with deep neural networks and tree search. 529, 2016.
Edward Snelson and Zoubin Ghahramani. Sparse Gaussian processes using pseudo-inputs. In Y. Weiss, B. Schölkopf, and J. C. Platt, editors, Advances in Neural Information Processing Systems 18, pages 1257-1264. MIT Press, 2006.
Shengyang Sun, Guodong Zhang, Jiaxin Shi, and Roger B. Grosse. Functional variational Bayesian neural networks. In 7th International Conference on Learning Representations, ICLR 2019, New Orleans, LA, USA, May 6-9. 2019. OpenReview.net, 2019.
Michalis K. Titsias, Jonathan Schwarz, Alexander G. de G. Matthews, Razvan Pascanu, and Yee Whye Teh. Functional regularisation for continual learning with Gaussian processes. In International Conference on Learning Representations, 2020.
Joost van Amersfoort, Lewis Smith, Yee Whye Teh, and Yarin Gal. Uncertainty estimation using a single deep deterministic neural network. In International Conference on Machine Learning, 2020.
Joost van Amersfoort, Lewis Smith, Andrew Jesson, Oscar Key, and Yarin Gal. Variational determin- istic uncertainty quantification, 2021.
Martin J Wainwright and Michael Jordan. Graphical Models, Exponential Families, and Variational Inference. Now Publishers Inc., Hanover, MA, USA, 2008. ISBN 1601981843
Ziyu Wang, Tongzheng Ren, Jun Zhu, and Bo Zhang. Function space particle optimization for Bayesian neural networks. In International Conference on Learning Representations, 2019.
D. T. S. Widdowson. The Management of Grading Quality: Good Practice in the Quality Assurance of Grading. Tech. Rep., 2016.
David H. Wolpert. Bayesian backpropagation over i-o functions rather than weights. In J. Cowan, G. Tesauro, and J. Alspector, editors, Advances in Neural Information Processing Systems, vol- ume 6. Kaugrann,, 1993.
12

---

Checklist
1. For all authors...
(a) Do the main claims made in the abstract and introduction accurately reflect the paper's
contributions 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 to them? [Yes]
2. If you are including theoretical results...
(a) Did you state the full set of assumptions of all theoretical results? [Yes] pendix A.
See Ap-
(b) Did you include complete proofs of all theoretical results? [Yes] See Appendix A.
3. If you ran experiments..
(a) Did you include the code, data, and instructions needed to reproduce the main experimental results (either in the supplemental material a URL)? [Yes] coom/timuudnerFSV..
See https://githan
(b) Did you specify all the training details (e.g., data splits, hyperparameters, how they were chosen)? [Yes] See Appendix D.
(c) Did you report error bars (e-g., with respect the random seed after running experi- ments multiple times)? [Yes] See Table 1, Table 2, and Table 3.
(d) Did you include the total amount of compute and the type of resources used (e.g., type of GPUs, internal cluster, or cloud provider)? [Yes] See Appendix D.9.
4. If you are using existing assets (e.g., code, data, models) or curating/releasing new assets...
(a) your work uses existing assets, did you eite 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: /httttttttu.t com/t com/timrudnerr /FSVI.
(d) Did you discuss whether and how consent obtained from people whose data you're
using/curating? [N/A]
(c) Did you discuss whether the data you are using/curating contains personally identifiable information or offensive content? [N/A]
5. If you used crowdsourcing or conducted research with human subjects...
(a) Did you include the full text of instructions given participants and screenshots, if applicable? [N/A]
(b) Did you describe any potential participant risks, with links to Institutional Review Board (IRB) approvals, applicable? [N/A]
(c) Did you include the estimated hourly wage paid to participants and the total amount spent participant compensation? [N/A]
13


Reviews:
[{"review_id": "lfejoqyWmc8", "reviewer": "Reviewer_8CnJ", "summary": "This works presents the function-space variational inference in tractable manner. Specifically, the existing function-space variational inference [1] has computational issue in computing KL terms between variational and prior distributions over function-space. This work first presents the tractable form of KL terms by approximating the distributions over function-space linearly by the Taylor approximation, and then employs the estimate of supreme of KL terms by using the context sets, which results in the function-space variational inference in tractable manner. Through experiments, authors validate that the proposed inference could perform the reliable uncertainly estimation, and outperform the existing Bayesian deep learning model for predictive performance and uncertainty estimation.\n\n[1] Functional variational bayesian neural networks, ICLR 19 ", "questions": "$\\bullet $ Q1. could you explain some intuition on why authors choose such context distributions in Table 1 and 2? what if the KMNIST or CIFAR100 is not available, which datasets do the authors use? Please explain why authors chose such datasets.  ", "limitations": "See section Weakness above.", "rating": 6, "confidence": 4, "soundness": 4, "presentation": 4, "contribution": 3, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "$\\textbf{Strengths}$\n\n$\\bullet$ This work proposes the tractable function-space variational inference, which can easily be computed.\n\n$\\bullet$ This work has been sufficiently validated that the proposed inferences leads to the reliable uncertainly estimation and superior predictive performance on a variety of the datasets.\n\n$\\textbf{Weaknesses}$\n\n$\\bullet$ The technical novelty looks incremental.\n\n$\\bullet$ As described in methodology section and shown in Table 1 and 2, the predictive performances seem to be largely affected depending on how to chose the context points and distribution. Thus, the performance could be unstable depending on users who might have deep understanding of given task or not. Thus, I think that explaining the intuition on how to choose context distribution should be described in detail.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}, {"review_id": "R5oGk16wbKE", "reviewer": "Reviewer_UrsU", "summary": "This paper proposes a tractable method for function space Variational inference that allows the method to be applied to realistic real-world settings. The authors achieve this by means of two approximations. Firstly, rather than calculating a KL divergence between two posteriors for which we do not know the closed-form density, they calculate the KL divergence between locally accurate linear approximations to the posteriors, whose density *is* known. Secondly, they use a rough approximation to the supremum over the aforementioned KL divergence, which they argue is enough since it encourages the variational distribution to match the prior on a set of context points.", "questions": "## Major\n\n0. [Significance] To me it seems very strange to do function space inference when the function space prior does not have any correlation between inputs. Under this modelling assumption, we are saying that we do not expect one input to tell us anything about the other input. This issue is sidestepped in FSVI due to the fact that NNs have some smoothness properties and as such the approximate posterior does not treat all inputs as independent. Nevertheless, one should strive to have a correct model and then apply approximations which are hopefully not harmful, rather than have an incorrect model that is fixed by applying approximations. \n\n1. [Experimental] It would be great to have a qualitative comparison between FSVI and HMC for some of the to datasets, for example those in Figure 1. Providing only the fits for FSVI without a comparison isn't particularly illustrative, and since these problems are small, HMC can be run to get a (somewhat) good idea of the 'true' posterior. \n\n2. [Experimental] On a similar note, comparisons to Sun, et al. and Ma, et al. on some smaller settings (perhaps the UCI datasets used by Sun et al.). This would allow the reader to understand what is being gained/lost (other than improved tractability) when using FSVI rather than these previous approaches to function space BNNs.\n\n3. [Signifigance] On lines 168-170, the authors suggest that the choice of context distribution can be viewed as a problem-specific modelling choice. I think that this is a very interesting idea, however, the idea is somewhat vague. Would it be possible to elaborate on this and give a few concrete examples? \n\n4. [Experimental/Significance] On lines 170-171, the authors mention that $S$ and $K$ are hyperparameters to be optimised on a validation set. How impactful are these hyperparameters? Do the authors have any concrete pieces of advice for choosing values/ranges? A sensitivity analysis would be very helpful. \n\n5. [Experimental] The ablation study for the choice of context distribution, while a great inclusion, was somewhat unhelpful because there was no accompanying discussion and the plot labels were somewhat mysterious. It would also be interesting to see the impact of swapping the context set and OOD dataset in the OOD results from Tables 1 and 2.\n\n6. [Experimental/Significance] How restrictive is the mean-field variational assumption (c.f., line 174)? Would it be possible to include a small ablation study? \n\n7. [Experimental/Significance] Ditto for the diagonalization discussed on lines 179-180.\n\n8. [Experimental]. For all of the baselines used in Section 5, a description of what hyperparameters were optimised and how (and additionally how hyperparameters which were not optimised were chosen) needs to be supplied. Without this information, it is not clear what to make of the results. \n\n9. [Experimental] Re Figure 2. Why is the brier score being used to measure predictive accuracy, rather than accuracy itself? Also, it seems like FSVI is underfitting, since the accuracy is low for small rotations. However, this seems to contradict the results in Table 1. Is there a difference in the experimental setup that I do not appreciate? Finally, would it be possible to see a plot like this for corrupted CIFAR10? This would be much more informative than a single column in Table 2.\n\n10. [Experimental] Figure 4 is somewhat incomprehensible because there are too many lines. I would suggest removing the ensemble results (other than deep ensemble) since they are not important for the key message. \n\n11. [Experimental/Significance] It seems to me that the conclusion that the correlation between OOD-AUROC and ECE is much less strong for CIFAR10 than FashionMNIST, which calls into question the generality of the finding. Is it possible to run this experiment for additional datasets?\n\n## Minor\n\nA. On line 27, I would suggest citing MacKay in addition to, or instead of, Neal for BNNs.\n\nB. On lines 29-31, the authors suggest that BNNs underperform non-Bayesian methods. However, this is a very strong statement. For example, the Linearised Laplace approximation is known to outperform deep ensembles in some settings; see https://arxiv.org/pdf/2106.14806.pdf and http://proceedings.mlr.press/v139/daxberger21a/daxberger21a.pdf. Similarly, HMC BNNs can also outperform deep ensembles in certain settings http://proceedings.mlr.press/v139/izmailov21a/izmailov21a.pdf.\n\nC. On line 87, the prior is introduced as though it was present in the equation above, which is confusing. Perhaps the authors meant to describe the posterior here? \n\nD. In eqn 1, the authors switch to a KL between the approximate posterior and *prior* without any explanation, which could confuse readers. \n\nE. Proposition 1 does not need to be a proposition. As far as I understand, it is a well-known result (e.g., in the linearised Laplace approximation community). Similarly, Corollary 1 does not need to be a corollary. This is also a well-known result, as evidenced by the trivial proof. In my opinion, presenting these results in such a way only serves to make the paper less approachable. \n\nF. I found the use of $g$ and $\\tilde{g}$ in Corollary 1 confusing – the symbols are too similar!\n\nG. Lines 188-189, 198-199 are repeats of 171-172 and can be removed. \n\nH. In what sense is CIFAR100 diverse?  (line 195)\n\nI. Would it be possible to provide a more detailed description of the procedure for generating the monochrome images? I am afraid that the description in the text was not clear to me at all. \n\nJ. On line 2017, I believe that there is a typo. Should 'S*M*' be 'S*K*'?\n\nK. It would be great to have a slightly expanded discussion of the related work. In particular, the relation of FSVI to the approaches of, (Benjamin et al., 2019; Pan et al., 2020; Titsias et al., 2020) could be expanded upon. \n\nL. Similarly, a discussion of how the results from (Foong et al., 2019) for weight space BNNs related to function space BNNs, would be helpful. \n\nM. As far as I understand the debate about the properties of mean-field variational distributions in deep BNNs has not been settled, so the statement on line 233 implying otherwise is somewhat strong. I would rather say that \"Farquhar et al., (2020) suggest that deep BNNs do not suffer from this issue.\"\n\nN. On lines 246-247, do the authors mean that the *uncertainty* of the predictions will be significantly higher? \n\nO. Lines 258-260, seem to be redundant with the previous paragraph. \n\nP. I assume that the Kaggle dataset on line 316 is the EyePACS dataset on line 313?\n", "limitations": "The authors have not made an effort to explicitly describe the limitations of their work. For example, they have not discussed the sensitivity of their proposed method to its hyperparameters, nor have they described the trade-offs between their method and existing functional BNN methods. However, I am confident that this issue can easily be addressed. ", "rating": 4, "confidence": 2, "soundness": 2, "presentation": 3, "contribution": 2, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "## Strengths\n\n**Originality**: the paper proposes a novel solution to a challenging and important problem in the BNN literature. I particularly enjoyed the idea of using local linearization to approximate the KL divergence between the approximate posterior and prior! \n\n**Clarity**: for the most part I found the paper well written and explained. There were a few issues, which I detail later.\n\n## Weaknesses\n\n**Quality**: While the experimental evaluation is quite extensive (many benchmarks and many baselines – very nice!), I do have some concerns and suggestions, which are also listed below. My main concern is that the experiments do not add to the understanding of the performance of the method–the relative strengths and weaknesses are not apparent, and the impact of the various approximations is unclear. There is also a lack of some important details that make it difficult to interpret the results of the experiments. \n\n**Significance**: I have a number of concerns about the approximations and assumptions in this work, which have not been addressed in the text. Without having these concerns addressed it is hard for me to make any judgement about the potential significance of the work. My primary concern is with the choice of prior. This concern and others are listed in the questions below. \n", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}, {"review_id": "FzZB5Oy6d9", "reviewer": "Reviewer_UY8r", "summary": "The authors propose a change to an established approach (Sun et al., 2019) for performing inference in Bayesian neural networks, which works directly in function-space rather than weight-space. The previous work by Sun et al. (2019) derive an ELBO with an intractable KL divergence between processes, so to make the problem tractable, they estimate the gradients of this KL term using the spectral Stein gradient estimator. This estimator, however, has been shown to be less efficient for high-dimensional distributions. \nAs an alternative approach, the current paper's authors propose a simple estimator of the KL term, which is based on a Taylor expansion of the random functions induced by BNNs and an MC estimate over context points from the input space.\nThe authors empirically show how their method leads to state-of-the-art uncertainty estimation and predictive performance on several datasets.\n", "questions": "1. You mention that you choose a prior variance of 10, and based on the experiment in section 5.4 it is clear how this choice is seemingly the best. I am a little confused about whether this is the prior over weights of the BNNs or over functions they represent, however. I suppose it is over the weights, but the text keeps referring to the induced prior over functions. Which prior variance does $\\Sigma_0$ refer to?\n\n2. Further to the question above, if $\\Sigma_0$ is the variance of the prior over the weights, did you try to estimate the induced variance over functions as well? Given that the datasets are normalised, an induced function-space variance of more than 1 would be peculiar.\n\n3. In lines 154-157, you mention that, theoretically, the approximation to the KL divergence will be better for smaller variances. While a variance of 10 empirically works the best, did you investigate the quality of the approximation for this value? I'm wondering whether the variance of 10 might actually result in a rather poor approximation that happens to work well and whether this is somehow related to the cold posterior effect (this is more of a speculative/discussion point rather than a question and I might well be misunderstanding something here!).\n\n4. Did you empirically compare the performance of your approach to that of Sun et al. (2019) on problems that are within reach of both approaches?\n", "limitations": "The authors have adequately addressed the limitations of their work. A potential negative societal impact has not been discussed, but it is not necessary given the paper's theoretical nature.", "rating": 7, "confidence": 3, "soundness": 4, "presentation": 4, "contribution": 4, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "**Strengths**\n\nThis is an impressive paper; incredibly well-written and thorough, with a good and concise related work section. The proposed approach is a novel and significant contribution to the field, which will undoubtedly be of great interest to the NeurIPS community. \n\nThe proposed approach is carefully derived in a way that makes it possible for the reader to follow along all the way. The paper contains a great empirical assessment with an impressive amount of details for making the work reproducible. When a setting or a choice of, say, a hyperparameter or distribution is not immediately clear, the authors carefully discuss the options. This might be the most polished and detailed paper I have ever reviewed.\n\nThe proposed approach also works well in practice, and several studies of the effect of hyperparameters and priors (in the supplementary) provide valuable insights.\n\n**Weaknesses**\n\nIt's really quite hard to find any substantial weaknesses in the paper.\nIn my view, the greatest weakness is that the authors do not compare their approach against that of Sun et al. (2019) despite framing the entire paper as an improvement over this. It would have been interesting to see the two approaches compared side-by-side on problems that both methods can handle.\n\nAn additional (minor!) weakness of the proposed approach is that it introduces a few extra choices to be made and hyperparameters to be optimised compared to the method by Sun et al. (2019). The authors do, however, provide suggestions and thorough discussions of how to make these choices.\n\n**Comments, typos, and various minor things**\n\n* Line 100: $\\mathcal Q_\\theta$ -> $\\mathcal Q_{q_\\Theta}$.\n* Line 111: The word \"poses\" doesn't seem to fit in this sentence.\n* In the equation following line 138, the matrix $\\mathbf S$ is only defined in the appendix.\n* In line 175, if $\\Sigma$ is the variance, I think the reparametrisation should be written (informally) using $\\Sigma^{½}$ or using the Cholesky decomposition $L$.\n* Line 207: $SM$ should be $SK$, I suppose.\n* In lines 246-247, I find the phrasing of \"predictions will be significantly higher\" a bit confusing. I suppose it refers to the predictive uncertainty, which should be higher, but I'm not entirely sure.\n* I suppose the ablation studies are actually more like hyperparameter tuning studies rather than actual ablation studies.\n", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}, {"review_id": "0FhCeTnPxQV", "reviewer": "Reviewer_uYCu", "summary": "The manuscript proposes a scalable functional-space variational inference method by incorporating prior information. The main contribution of the paper is to propose a KL divergence estimator between a posterior and variational distribution over functions by approximating the first-order Tayler series expansion of the mean of the parameters.\n", "questions": "What are the potential ways to apply this approach to non-Gaussian priors?\nHow does the proposed method work for other problems such as Bandits, and Bayesian Optimization?", "limitations": "The results presented are convincing for the image classification task, but you need to compare against other variational inference approaches proposed in the literature. Some references are provided above.\n\n", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "The proposed approach seems intuitive and the proof seems reasonable. The results demonstrate the benefits of the approach.\n\nWeakness:\n-  The diagonalization approximation to the covariance matrices in Section 3.2 is a simplification that deviates from the original approach from Sun et. al. \n- The proposed method are presented only image classification problems.\n-  Authors need to compare other approaches that scale functional variation inference to even larger models: please refer to the following manuscript.\n\"Carvalho, Eduardo DC, et al. \"Scalable uncertainty for computer vision with functional variational inference.\" CVPR'2020\n\n- Other approaches which are not functional variational inference but provide well-calibrated uncertainty estimation need to be compared.\n\nKarandikar, Archit, et al. \"Soft calibration objectives for neural networks.\" Advances in Neural Information Processing Systems 34 (2021): 29768-29779.\nKrishnan, Ranganath, and Omesh Tickoo. \"Improving model calibration with accuracy versus uncertainty optimization.\" Advances in Neural Information Processing Systems 33 (2020): 18237-18248.\n\n\n\n ", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ["I don’t know"]}, "extra_text": {}}]

Author rebuttal:

Decision: Accept