Function-Space Variational Inference for Deep Bayesian Classification
Authors: Jihao Andreas Lin
Joe Watson
Pascal Klink
Jan Peters
Venue: ICLR 2022 Submitted (2022)
Keywords: Bayesian deep learning
image classification
functional variational inference
Dirichlet distribution

Abstract: Bayesian deep learning approaches assume model parameters to be latent random variables and infer posterior predictive distributions to quantify uncertainty, increase safety and trust, and prevent overconfident and unpredictable behavior. However, weight-space priors are model-specific, can be difficult to interpret and hard to choose. Instead of weight-space priors, we leverage function-space variational inference to apply a Dirichlet predictive prior in function space, resulting in a variational Dirichlet posterior which facilitates easier specification of epistemic uncertainty.  This is achieved through the perspective of stochastic neural network classifiers as variational implicit processes, which can be trained using function-space variational inference by devising a novel Dirichlet KL estimator. Experiments on small- and large-scale image classification tasks demonstrate that our function-space inference scales to large-scale tasks and models, improves adversarial robustness and boosts uncertainty quantification across models, without influencing the in-distribution performances, architecture or model size.

Under review as a conference paper at ICLR 2022
FUNCTION-SPACE VARIATIONAL INFERENCE FOR DEEP BAYESIAN CLASSIFICATION
Anonymous authors Paper under double-blind review
ABSTRACT
Bayesian deep learning approaches assume model parameters to be latent ran- dom variables and infer posterior predictive distributions to quantify uncertainty, increase safety and trust, and prevent overconfident and unpredictable behavior. However, weight-space priors are model-specific, can be difficult to interpret and are hard to specify. Instead of weight-space priors, we leverage function-space vari- ational inference to apply a Dirichlet predictive prior in function space, resulting in an implicit variational Dirichlet posterior which facilitates easier specification predictive epistemic uncertainty. This is achieved through the perspective stochastic neural network classifiers variational implicit processes, which can be trained using function-space variational inference by devising a novel Dirichlet KL estimator. Experiments small- and large-scale image classification tasks demonstrate that our function-space inference scales to large-scale tasks and mod- els, improves adversarial robustness and boosts uncertainty quantification across models, without influencing architecture or model size.
INTRODUCTION
Deep learning (Goodfellow al., 2016) has enabled powerful classification models capable of working with complex data modalities and scaling to large data sets (Krizhevsky et al., 2012). outperforming non-parametric models such as Gaussian processes (Rasmussen Williams, 2005) in these domains. The aim of Bayesian neural networks (BNN) is to provide these complex models with function priors for regularization, generalization and uncertainty quantification (UQ) useful in prediction tasks (Gal, 2016; Wilson & Izmailov, 2020; Fortuin, 2021; Abdar et al., 2021). Predictive uncertainty is crucial for the deployment of machine learning systems in real-world settings, as it provides a degree of safety (McAllister et al., 2017), trust (Lim et al., 2019), sample efficiency (Deisenroth & Rasmussen, 2011; Gal et al., 2017) and human-in-the-loop cooperation (Filos et al., 2019). Unfortunately, the highly redundant and nonlinear parameterization of modern deep learning systems has been shown to be challenging for Bayesian methods, sometimes resulting in inadequate uncertainty quantification (Foong et al., 2020; 2019; Osband et al., 2018; Ovadia al., 2019; Wenzel et al., 2020; Yao et al., 2019) even on toy tasks (e.g. Figure 3).
In this work, we consider function-space variational inference' (Sun al., 2019) (IVI) for classifica- tion. This inference method reasons about priors in function space, specifying the prior with respect to their predictive distribution rather than their weight distribution. By approximately representing the posterior in function space using variational implicit process (Ma et al., 2019), our approach can be used with several popular BNN representations such as deep ensembles (Lakshminarayanan et al., 2017) and Monte Carlo dropout (Gal & Ghahramani, 2016). By positing that the implicit predictive of a stochastic network classifier Dirichlet distributed, rather than solely categorical, we gain a richer predictive uncertainty representation while retaining the same mean predictions categorical classifiers. Furthermore, the approach provides a unifying perspective of prior work which use the Dirichlet distribution and function-space regularization (Malinin & Gales, 2018; Malinin et al., 2020; 1oo et al., 2020; Sensoy et al., 2018: 8: 2021). 1), through a principled objective that does not require specific model or data transformations. We demonstrate that this method improves uncertainty quantification and adversarial robustness across a range of popular models and datasets for both small- and large-scale inference.
While prior work uses the term functional VI, we use funetion-space VI for clarity.

---

Under review as a conference paper at ICLR 2022
5
5
5
h
5
5 5
1
Ply= 5)
Ply
90 Angle (*)
135
180
45
90 Angle (*)
135
180
(a) Ensemble
(b) Ensemble fVI
Figure 1: Bayesian classification (blue) on the rotated MNIST task. Although the ensemble has reasonable uncertainty quantification, it can sporadically fluctuates rotations and exhibits regions of overconfidence. Using function-space VI, the uniform prior (red) regularizes the prediction outside of the data distribution be more consistent and reasonable in its uncertainty. Here, the index set is the training data and uncertainty is expressed through the 1st, 3rd and 25th percentile.
1.1 CONTRIBUTIONS
We incorporate Dirichlet predictive priors into deep Bayesian networks using function- space inference. Concretely, our contributions are:
1. We demonstrate an extension of variational implicit processes into the classification setting. 2. We propose function -ancce variational inference method for classification, using an effec- tive KL. divergence estimator from softmax samples via Newton-basee Dirichlet estimation.
RELATED WORK
We build on prior work that considers classification in the Bayesian setting and function-space inference of Bayesian neural networks.
BAYESIAN CLASSIFICATION
Compared to regression, classification non-trivial for Bayesian methods due the nonlinear link function required to predict the class labels. As a result, closed-form Bayesian models, such as Gaussian processes (GP), require approximate inference methods such as the Laplace approximation (Rasmussen & Williams, 2005), variational inference (Gibbs & Mackay, 2000; Hensman et al., 2015; Salimbeni et al., 2018: Izmailov et al., 2018), and expectation propagation (Hernandez-Looato Hernandez-Lobato, 2016). The Pólya-Gamma data augmentation trick (Polson et al., 2013) has enabled scalable closed-form variational training of sparse Gaussian process classifiers (Wenzel et al., 2019). Gaussian processes have also been used with Dirichlet predictive using log-normal approximation (Milios et al., 2018). Classification with Bayesian neural networks possible through wide range of approximate inference methods, including Markov chain Monte Carlo (Neal, 1995; Zhang et al., 2020), (mean-fich) field) variational inference (MFVI) (Graves, 2011; Blundell et al., 2015; Zhang al., 2018). Laplace approximations (MacKay, 1992; Denker & LeCun, 1991; Ritter et al., 2018; Khan et al., 2019; Immer et al., 2021; Daxherger et al., 2021), ensembles (Lakshminarayanan al., 2017: Osband al., 2018; Barber & Bishop, 1998; Pearce et al., 2020), expectation propagation (Hernández-Lobato & Adams, 2015) and Monte Carlo dropout (Gal & Ghahramani, 2016; Kingma al., 2015). Radial BNNs (Farquhar et al., 2020) are motivated as a practical alternative to MFVI BNNs that uses Gaussian weight priors and posteriors. By sampling weights in a radial fashion, they avoid the pathologies encountered when sampling high-dimensional Gaussian distributions. Rank 1 BNNs (Dusenberry et al., 2020) combine ensembles and weight priors. Using the shared BatchEnsemble structure (Wen et al., 2020) and rank covariance parameterizations, Rank BNNs have scalable memory requirement. Alternatively, the Laplace bridge (Hobbhahn et al., 2021) approximately maps a Dirichlet predictive density backwards through the softmax into a latent Gaussian predictive. A daussian-predichte BNN can then be trained using this latent approximation.
Alternatve methods avoid propagating uncertainties by predicting Dirichlet concentrations directly with neural networks. Prior networks (Malinin & Gales, 2018; 2019) require categorical labels to be converted to Dirichlets, and mimic fVI as the objective consists of two KL divergences, for in- and outside the data distribution respectively. Another training method used prior networks distill
2

---

Under review as a conference paper at ICLR 2022
trained ensemble into single model (Malinin al., 2020). Similarly, belief matching (Joo et al., 2020) converts the labels to Dirichlets using Bayes rule. This method also echoes fVI, as the training objective is an fELBO where the index set is the training data. Another method converts the training labels to categorical probabilities and uses a Bayes risk objective with KI. regularization against function-space prior (Sensoy et al., 2018; 2021). Compared to these methods, principled approach to inference that allows us to use standard BNN parameterizations with the categorical likelihood, avoiding the need to design networks and data representations that facilitate this training approach. For a longer discussion on this prior work, see Appendix A.
fVI
2.2 FUNCTION-SPACE VARIATIONAL INFERENCE
Function-space inference is best understood through Gaussian processes, which are closed-form distributions over functions. Using this view, sparse GPs may be viewed as approximate inference over functions (de G. Matthews et al., 2016), minimizing the fKL. from its exact posterior via inducing points. The function-space variational BNN (IBNN) by Sun al. (2019) uses the fKL to use explicit or implicit stochastic processes as function- priors. They use a GP trained on the data as a prior, which can be viewed form of empirical Bayes. Similar approaches take a mirror descent view for batch training (Shi et al., 2019). However, function-space variational inference can yield ill-defined objectives in many cases Burt et al. (2021). Neural linear models have also been used with [VI, as they have an attractive closed-form Gaussian predictive (Watson et al., 2021a:b). Concurrent work by Rudner al. (2021) ) has also adopted f'I for classification, by linearizing the network about a Gaussian weight distribution to estimate the IKL. This model works with Gaussian (latent) predictive prior and posterior which loses the intuitive aspect of function-space priors. Moreover, the lincarization requires computation of the weight Jacobian during training, for which the memory requirement scales with the number of model parameters and outputs. Variational implicit processes (Ma et al., 2019) take a stochastic process perspective of stochastic parametric regressors and, for the regression setting with Gaussian likelihoods, proposes a wake-sleep procedure for inference. Wang et al. (2019) propose particle optimization methods, using finite representations of functions to learn a particle representation of the function-spact posterior through the gradient flow the log posterior. Function-space inference is also an attractive approach to continual learning, e.g. Pan et al. (2020).
3 DIRICHLET PREDICTIVE PRIORS FOR BAYESIAN CLASSIFICATION
In this section, we discuss how function-space variational inference can be applied to classification using variational implicit processes (VIP) (Ma et al., 2019).
3.1 VARIATIONAL. IMPLICIT PROCESSES FOR CLASSIFICATION
Let D c (Xx (Z b be the training data consisting of N observed pairs of input data R, and corresponding K K-ddimenional,, one-hot class label vectors yn- Deep learning approaches to classification typically assume that class label y follows categorical distribution whose parameters are predicted by a neural network. To this end, a network Φ with weights w defines a deterministic function f which maps an input æ E X' an element fa E A K-1, where A K-1 denotes the K-1 simplex and elements from A K-1 are K K-dimensional probability vectors. More precisely, we write fa g(x, w) a (Ø (z:; w)) and y Cat (- f The softmax function ensures that fa is a valid probability vector, i.e. that f= 2 0 and Ek= 1 1. In conventional ML training, the weights are optimized by maximizing log (D) f) log IIp Cat (ylfa) Ep Ek=1 y/ log fax- where II, and En denote [( and Σ (ED respectively, and Φ. o,, and w are implicit in f.
In the Bayesian setting, this categorical predictive is expected to express epistemic uncertainty in its prediction. For many BNN models, this amounts to sampling from a weight distribution p (w) and averaging the predictions to produce a posterior categorical predictive. We argue that this process throws away the epistemic uncertainty of the classifier, which is captured in M posterior predictive samples f!!"
(MM) To capture this uncertainty, we study these implicit predictions and instead posit they are drawn from posterior predictive Dirichler distribution p(fx). The Dirichlet density allows us to capture the variance in the categorical samples, as seen in Figure 2. In general, the model becomes a generative mapping from X' to the space of K-dimensional Dirichlet distributions.
3

---

Under review as a conference paper at ICLR 2022
Observations
Arithmetic Mean
Estimated Dirichlet
Reducing observations from simplex (left) to their mean (middle) discards information which is present in the variance of the predictions that reflects model uncertainty. A Dirichlet distribution (right) fitted to the same samples can capture the uncertainty with its density function.
Figure 2:
Sampling the weights and evaluating g -) equivalent to sampling the Dirichlet predictive, while varying with fixed weight w corresponds to evaluating a classifier g(,, w). Together, the joint variation of and az defines an implicit stochastic process which we assume describes an implicit K-dimennsonnal Dirichlet distribution for each a: E X. When parameterized by neural network, we can treat this stochastic process as a variational implicit process (Example 3, Ma et al. (2019)).
Definition 3.1 (Implicit stochastic processes, from Ma et al. (2019)). An implicit stochastic process (IP) a set of random variables f. such that any finite subset SML. = (fa): Smay.... say) with L E N has joint distribution which is implicitly defined by the generative process.
w ~ p(w)
Sa
g 2n, W),
aa E X
151SL.
Applying Definition 3.1 to stochastic network classifiers requires two key considerations. Firstly, the weight distribution p( w) takes a parameterized form, (ww), with parameters O we wish to optimize. The function g(::) represents specific network architecture, such as a ResNet. Further details regarding the variations of go (wr) that we consider can be found in Section D in the Appendix. Secondly, the classification setting we are now reasoning over a joint distribution of Dirichlet distributions fixed K-1 simplex. Note this is distinct from the Dirichlet process (Teh, 2010), which is a rrans-dimensional prior over discrete distributions, such that is inferred. Informally, in setting. finite collection of L random variables would be element from the K-1 simplex raised to the power of L. The stochastic process can be defined as a random variable from the reason about
K-1 simplex to the power of the data domain X. This is not distribution easily, in contrast to the Gaussian process where the joint and marginal distributions are all Gaussian (Rasmussen & Williams, 2005). However, we only interact with the joint distribution implicitly, we avoid the complexity of representing it directly. Using this VIP perspective, we can use the stochastic networks as an implicit variational posterior over functions (f))), allowing us to adopt gegh-performacc BNN representations while avoiding distributions over infinite parameter spaces, as in e.g. Bayesian nonparametries (Hjort et al., 2010). Further details regarding the formal stochastic process definition can be found in Section C in the Appendix.
While modeling the predictive distribution as Dirichlet captures the variation in the predictive samples, the inference procedure is not affected if we use weight-space priors. To use this rich predictive density to inform our inference strategy, we can design a Dirichlet predictive prior in function space and apply variational function- sprrionn-spceeeeeee inference to optimize (ww)..
3.2 FUNCTION-SPACE VARIATIONAL. INFERENCE AND INDEX SETS
To optimize our variational implicit processes, we can use the function-space evidence lower bound objective (fELBO) (Sun et al.. 2019),
(( = E f- [logr(D)f)) D p(f)),
(1)
which resembles the conventional evidence lower bound objective (ELBO) (Jordan et al., 1999; Hoffman et al., 2013). This enables the explicit specification of the desired behavior of the predictive distribution directly in function space, whereas the conventional ELBO relies on indirect specification via prior model parameters p(w). The expected logelikeli term in the fELBO is unchanged from the regular BIBBO optimization for BNNs, where we use M sample estimates using the categorical likelihood,
E [log (Dgg(DD))] ~ 1 EN EM log Yuu fa sl(m))
(2)
4

---

Under review as a conference paper at ICLR 2022
Since the fELBO involves distributions over functions. fVI requires computing KL divergence between stochastic processes, Sun et al. (2019) derived this divergence as the supremum over all finite measurement sets X X th of index sets X of size L,
DkL( p[f]]=
sup LEN,XEXL
DkLlO(fx10) || p(fx)]
(3)
However this realization is impractical as this supremum is not obtainable in closed form. A tractable approximation (Sun et al., 2019; Bruinsma al., 2021) replaces the supremum with an expectation,
( p(f)] Es-s DKLlg(fs|0) p(fs)
(4)
where c X a finite set of (now) fixed size T., which we refer to as the 'index set'. Note that Sun et al. (2019) refer to S as distribution Sun et al. (2019) motivate this approximation with the KL between two conditional stochastic processes (Section A.3, Sun et al. (2019)) which is derived for generic stochastic processes. This approximation assumes that S contains the points that the stochastic processes are conditioned on. For our posterior and prior process, this is the training data. Sun et al. (2019) also evaluated adversarial approach that optimized worst-case index set. They report that the optimized index set approximated the training data, but this index set provides inadequate regularization for out of distribution (OOD) prediction, as this predictive regularization is only applied in-distribution. The theoretical and practical requirements can be balanced by using the index set & sampled around the training data. Moreover, the observation that this regularization should be applied OOD suggests that the sampled index set should be drawn from the test-time data distribution. While the training data should capture this distribution, we can improve further using unlabeled data or data augmentation. We demonstrate the benefit of this index set design in Section 4, while methods of designing effective index sets is a topic of future work.
3.3 POSTERIOR DIRICHLET AND KL DIVERGENCE ESTIMATION
In contrast prior work (Joo et al., 2020; Sensoy al., 2018; Malinin & Gales, 2018), we do not directly predict a using neural network, but instead assume f(m) ~ Dir(-lax.. This allows us compute Oz through a maximum likelihood estimate (MLE) of M samples f(m) from the VIP. Unfortunately, the Dirichlet MLE cannot be computed in closed form and requires fixed-point iteration or other optimization techniques be computed, increasing the computational burden (Minka, 2000). Consequently, we replace the full MLE with a computationally cheaper alternative. To this end, we consider Ogg in terms of two separate but dependent parameters: the Dirichlet mean 0N = azz/ 2m and the Dirichlet precision 2n- By matching the first moment of the empirical distribution of fa, we obtain an estimate of a₂ ~ fast, iterative, quasi-Newton algorithm (Minka, 2000) using fú elI:M)
1 EM Em=1 ((m)) We then fix ase and estimate the precision z using a (.....I(MM))
M
((++1) Fi
(1)
+
(1)
-2
d2 dz CDir CDie (fL a! d LDie (a!))
(5)
dzz
where (04) and 0g az/
(t)
(t) are the Dirichlet precision and concentration iteration t of the quasi-
| (I.M) axx)) is the Dirichlet log-likelihood log I Dir (fl" (m) (O
Newton algorithm and LDi
This corresponds to breaking the dependence between &₂ and Zn that would be considered in the full MLE case in order to accelerate the resulting approximation. With 2œ and 2 estimated, az and thus our variational posterior (() = Dir (f= is now available. Moreover, for tractability we choose a factorized prior and assume a factorized posterior, i.e. q(fs)=I[[= ( IKL term from Equation can be estimated using AM samples with the Dirichlet KL divergence
ff,,, ).2 so that the
Dkll((ss(0) [[s)] + El EM (log q (S! 0) - log (Fl)))
(6)
(m)) are where f! is the m-th implicit function evaluated at s E S and log {f (m) and log (f()) the log-likelihood of f(m) under the variational Dirichlet posterior and Dirichlet prior respectively. Note, in this work our predictive prior p(fx) )= =Dir(- (), where each Bg is chosen to be unless
While this is not rich prior, we assuming our VIP, e.g. CNNs, will capture meaningful correlations.
5

---

Under review as a conference paper at ICLR 2022
Weight-space prior
IV 3 [1.1]
MV.. 3 [2
FVI. B = [1,
MAP
MC Dropout
Ensemble
Figure 3: Toy classification problem using the Two Moons dataset. Each row corresponds to different network. The first column shows the undesirable predictive certainty of standard weight- space priors outside of the data distribution. The second column illustrates how our function-space inference approach combined with a uniform Dirichlet prior prevents unreasonable extrapolation behavior and instead adequately increases model uncertainty outside of the observed data. The third and fourth column show that our approach can also combined with class-biased predictive priors that inform the OOD-prediction without sacrificing accuracy. This explicit design is not straightforward with weight-space priors.
otherwise stated. During optimization, we treat qf ()) as constant with respect to 0 due to the ag estimate, which serves the practical purpose of pruning the Dirichlet MLE from the computation graph and evokes expectation maximization-style inference. While this simplification removes the influence of 0 on Dir fa Oiz)) on the KL in the gradient update, learning is hindered due to the iterative nature of SGD training in practice. From the perspective of the wake-sleep inference procedure adopted by Ma al. (2019), optimizing the fELBO represents the wake phase and the Dirichlet estimation represents the sleep phase. Nonetheless, we optimize O using backpropagation on the fELBO objective by optimizing through the function samples w.r.t. the expected log- likelihood and fixed posterior and prior densities in the fKL approximation. For some models (e.g. ensembles) this is standard gradient descent, while for others (e.g. MFVT) the reparameterization trick is required. In the where only single sample is available (M 1), approximate the precision 2n with the size of the training data. Further practical considerations. mini- bani-batching and prior regularization, diseussed in Section in the Appendix. Both considerations manifest scaling of the KI. term.
EXPERIMENTS
In this section, we present an empirical evaluation of our proposed f'II approach on classification tasks for several models, comparing performance against their conventional training procedure. We first visualize the effects of a Dirichlet prior on a toy problem, and then perform image classification on the MNIST (LeCun al., 2010), CIFAR 10 and CIFAR 100 (Krizhevsky, 2009) data Depending on the experiment, we either employ feedforward multilayer perceptrons (MLPs) or convolutional neural networks (CNNs). Metrics include classification accuracy, log- likelihood (LLH) and expected calibration еиror (ECE) (Nacini et al., 2015), which estimates the calibration of accuracy vs confidence through binning the predicted class probabilities. Additional details are provided in Appendix F.
4.1 TOY PROBLEM
To visualize the effects of function-space variational inference, we conducted toy experiment with the Two Moons data set. We used three different model archetypes in this experiment: MAP (i.e. a deterministic baseline), MC Dropout (Gal Ghahramani, 2016), and (deep) Ensembles (Lakshminarayanan et al., 2017). Figure 3 shows that inference approach combined with uniform Dirichlet prior adequately increases predictive uncertainty outside of the observed data, whereas conventional weight-space training leads overconfident extrapolation. While the weight-

---

Under review as a conference paper at ICLR 2022
MAP
MC Dropout
Ensemble Ensemble IVI
Subnetwork
MAP IVI
MC Dropout IVI
Uniform prior
Standard Index Set
+ 90" aug.
90" 180 aug.
Log- Likelihood
-2
-4
-4
-6
-6
6
0
45
90 Angle 9
135
180
0
45 90 135 180
0
45 90 Angle "
135 180
Angle "
Figure 4: Comparison of rotated MNIST lig-lioeldoood for models trained with fVI using different index sets: training data without rotations (left), 90" augmentations (middle), and 90" and 180 augmentations. Colored lines denote the mean and shaded areas denote two standard deviations over 10 seeds. The 909 augmentation improved fVI performance for that scale of perturbation, but worsens performance at 180%. We attribute the degradation to the nature of parametric function approximation. Augmenting the index set further with 180 transforms fixes the degradation, demonstrating an overall improvement in UQ over the initial training data index set. The horizontal black line denotes the performance of the uniform prior. The subnetwork performance taken from Daxberger et al. (2021), and uses a ResNet-18 18 rather than an MLP. Further metrics are reported in Figure 9.
space approach learns decision boundary that bisects the data, the function-space approach learns richer decision boundaries that more accurately capture the data distribution. This behavior has also been observed with Gaussian mrocess-based classification, e.g. Liu al. (2020). Furthermore, fVI facilitates specifying class-biased priors, which would be difficult to achieve with weight priors.
4.2 ROTATED MNIST
Following Ovadia et al. (2019), we train on the MNIST handwritten digit classification dataset (LeCun et al., 2010) and evaluate on constructed test data with rotations of up 180", which simulates challenging OOD scenario due to the absence of data augmentation. For this experiment, we used the same model archetypes and MLP architecture for the toy problem in Section 4.1.
The log-likelihood between models is shown in Figure 4 (left), accuracy and expected calibration error reported in Figure 9 in the Appendix. In terms of classification error rate, both approaches yield the same performance. In terms of log-likelihood space counterparts as the data becomes more OOD, indicating better OOD uncertainty quantification. Subnetwork linearized Laplace (Daxberger et al., 2021) is also reported as competitive baseline, achieving better accuracy and in-distribution log- ligeehodd.ood. However, these results were obtained using a ResNet- 18 architecture rather than an MLP.
fVI consistently outperforms their weight-
4.3 ASSESSING INDEX SET DESIGN
To illustrate the importance of the index set, we train the I'II models for rotated MNIST using three different index sets: the training data, additional 90º augmentation, and additional 90 and 180" augmentation. While simply using the training data without rotations as index set for I'I already outperforms the weight-space counterparts, a direct comparison in Figure illustrates that performance can be further increased if an appropriate index set, i.e. example OOD data, is available. With the enriched index sets, the OOD performance move closer to that of the prior, indicating more accurate inclusion in the IELBO. Sets for greater OOD performance could be designed through manual data augmentation, unlabeled data or synthetic data generation. Note, for all other image classification experiments, without this inductive bias.
the training data as the index set, demonstrating performance
4.4 BARGER-SCALES IMAGE CLASSIFICATION UNDER CORRUPTION
We used the regular train splits of the CIFARIO and CIFAR100 (Krizhevsky, 2009) image classifica- tion data sets training data and their corrupted versions (Hendrycks & Dietterich, 2019) OOD test data. CIFAR1O and CIFARIO0 consist of natural color images of animals and vehicles, and the corrupted versions perturb the images at 5 increasing levels of severity by changing the brightness,
7

---

Under review as a conference paper at ICLR 2022
MAP
MC Dropout
Ensemble
Radial
Rank1
Subnetwork
MAP IVI
MC Dropout IVI
Ensemble [VT
Radial I
Rank1 IVI
Uniforn prior
100
0.4
Log Likelihhhodd
CIFARIO
Accuracy
80
0.3
ECE
-2
0.2
60
0.1
40
80
Log- Likelihood
0.4
CIFAR100
Accuracy
-2
60
ECE
-4
0.2
40
-6
0
0
2 3 5
0
2 4
2
4
Corruption
Corruption
Corruption
Figure 5: Metrics for corrupted image classification on CIFARIO (top) and CIFAR 100 (bottom). All models use a ResNet-18 architecture. For CIFARIO, there is a clear benefit of I'I priors over weight-space for -olliiklihoodd Accuracy performance is constant, while fVI underfits for ECE when in-distribution due to regularization from the prior. For CIFAR100, the higher label dimensionality results in stronger regularization from the uniform prior. As result, weight-space ensembles achieve better performance over all corruptions. Moreover, I'I models underfit w.r.t. ECE in-distribution due to the prior, but improve in calibration OOD, whereas weight-spase models worsen. This result indicates prior specification requires more care for high-dimensional classification. Interestingly, the IVI performance is consistent across models for accuracy and log-likelihood on both datasets, but are diverse for ECE. The linearized Laplace subnetwork result is taken from Daxberger et al. (2021).
contrast and saturation, or adding noise, blur other artifacts, such that it becomes ever more difficult for the models to classify them. For this experiment, we used CNN models based on the ResNet- 18 architecture (He al., 2016). In addition to the previously used MAP, MC Dropout and Ensemble models, we also evaluate our IVI objective on Radial BNNs (Farquhar et al., 2020), as an effective variant of MFVI, and Rank1 BNNs (Dusenberry et al., 2020), which combine ensembles and VI.
Figure reports the performance metrics for CIFAR 10 and CIFAR 100 under corruption. Tabulated results can be found in Section G in the Appendix. The general trend of these results is that the function-space prior frequently provides gains in OOD UQ with only a small decrease in (uncorrupted) test performance. This trade-off between accuracy and robustness has been observed and discussed in the adversarial robustness setting (Tsipras et al., 2019; Yang et al., 2020), and remains an open problem if and how both qualities can be achieved in practice. Moreover, the shared function-space prior resulted in remarkable consistency across models, compared to the variety seen in weight-space priors. Due to the predictive prior, the ECE was significantly higher for corruption O, but this underfitting provided superior ECE when OOD. This suggests the function-space prior should be fine-tuned (i.c. empirical Bayes) if superior ECE is desired in-distribution. For CIFAR 100, the higher prior regularization due to higher dimensionality (see Section B.2), resulted in reduced benefit over weight-space models compared to CIFAR10, with improved performance only evident at stronger corruptions. Figure? 7 summarizes the LLH difference due to the function- sppcti prior.
4.5 LARGER-SCALE IMAGE CLASSIFICATION UNDER ADVERSARIAL ATTACKS
Despite the success of CNNs in computer vision, adversarial attacks are one of the biggest risks when it comes to practical applications (Akhtar & Mian, 2018). We evaluate the robustness of IVI compared to standard weight-prite prrior approaches on the CIFARIO and CIFARION datasets using the fast gradient sign method (FGSM) (Goodfellow et al., 2014). Figure 6 compares the accuracy and the log- of the test data with increasing amounts of perturbation, ranging from c= 0 (no attack) to c= 0.3. Although both weight-space and function-space models lose their classification accuracy when the FGSM attack is introduced, the fVI models only suffer small decreases in log- ligeliiol whereas the weight-space L.H performance drops significantly. We also observe the accuracy vs
8

---

Under review as a conference paper at ICLR 2022
MAP
MC Dropout
Ensemble
Radial Radial fV1
Rankl Rankl IVI
Uniform prior
MAP IV
MC Dropout IVI
Ensemble TVI
100
0
Log- Log-Likelite Liikeihhodd
0.8
80
-2
CIFARIO
Accuracy
0.6
60
40
6
0.4
20
0.2
10
0
80
0
0.8
Log Likelihood
CIFARIOD
Accuracy
60
0.6
ECE
40
0.4
20
0.2
-10
0
0.1
0.2
0.3
0.1
0.2
0.3
0.1
0.2
0.3
Epsilon
Epsilon
Epsilon
Figure 6: Metrics for adversarial examples CIFAR1O (top) and CIFARIO0 (bottom). All models use a ResNet-18 18 architecture. For CIFAR10, there are significant benefits fVI over weight-space approaches across all metrics. For CIFARI00, the fVI benefits are still evident, but the higher label dimensionality results in stronger regularization from the uniform prior. As a result, the weight-space ensembles achieve slightly better performance over all epsilons.
robustness trade-off in the I'VI models discussed in Section 4.4. We attribute this behavior to the quality the uncertainty quantification the decision boundary. While both approaches have brittle boundaries due to the nature of CNNs, the predictive uncertainty at these decision boundaries is richer for IVI. This richness is illustrated in the toy example Figure 3.
CONCLUSION
propose a function-space variational inference approach to deep Bayesian classification, which enables the use of Dirichlet predictive priors to improve uncertainty quantification. Our approach provides unified formalism of prior work on Dirichlet-based classifiers with function-space reg- ularization, and can be applied to general class of variational implicit processes without altering their underlying architectures. Experiments demonstrate that our approach generally outperforms the corresponding weight-space priors in terms of uncertainty quantification and adversarial robustness. By considering the design of the index set, we show how one can trade-off scalability against OOD uncertainty quantification by extending the fKL evaluation beyond the training distribution. Future research should develop index sets for fVI, specifically effective methods for constructing them to reflect the test distribution, c.g. through using data augmentation or unlabeled data.
MAP
MC Dropout
Ensemble
Radial
Rankl
CIFARI0
CIFARI00
8
1.5
6
2
4
1
0.5
2
0
0
1
3
0.1
0.2
0.3
0
2
3
5
0
0.1
0.2
0.3
Corruption
Epsilon
Corruption
Epsilon
Figure 7: A summary of the improvement in log-likelihood for function-space inference over weight space, for CIFAR10 10 and CIFAR 100 experiments under corruption and adversarial perturbations. The general trend shows that models that use reparameterization gradients (Radial and Rankl) benefit the most from IVI. The improvement also generally increases the test data becomes more OOD. Ensembles are improved the least, and even weight-space inference outperforms function-space for CIFAR 100. This result is attributed to the regularization of the uniform prior in high dimensions.

---

Under review as a conference paper at ICLR 2022
ETHICS STATEMENT
As this paper looks at the abstract problem of classification under a Bayesian lens, we are not aware of direct ethical implications of this work.
REPRODUCIBILITY STATEMENT
Code for reproducing experiment results is provided in the supplementary material. Further, Appen- dices F, and H contain additional details and results of the experiments.
REFERENCES
Moloud Abdar. Farhad Pourpanah, Sadiq Hussain, Dana Rezazadegan, Li Liu, Mohammad Ghavamzadeh, Paul Fieguth, Xiaochun Cao, Abbas Khosravi, U. Rajendra Acharya, Vladimir Makarenkov, and Saeid Nahavandi. A review of uncertainty quantification in deep learning: Techniques, applications and challenges. Information Fusion, 2021.
Laurence Aitchison. statistical theory of cold posteriors in deep neural networks. International Conference Learning Representations, 2021.
Naveed Akhtar and A. Mian. Threat of adversarial attacks on deep learning in computer vision: A survey. IEEE Access, 6:14410-14430, 2018.
D. Barber and Christopher Bishop. Ensemble learning in bayesian neural networks. In Generalization in Neural Networks and Machine Learning, 1998.
Charles Blundell, Julien Cornebise, Koray Kavukcuoglu, and Daan Wierstra. Weight uncertainty in neural network. In International Conference Machine Leurning, 2015.
Wessel Bruinsma, James Requeima, Andrew Y. K. Foong, Jonathan Gordon, and Richard E Turner. The gaussian neural process. In Third Symposium Advances in Approximate Bayesian Inference, 2021.
David R. Burt, Sebastian W. Ober, Adrià Garriga-Alonso, and Mark van der Wilk. Understanding variational inference in function-space. In Third Symposium Advances in Approximate Bayesian Inference, 2021.
Erik Daxberger, Eric Nalisnick, James Allingham, Javier Antoran, and Jose Miguel Hernandez- Lobato. Bayesian deep learning via subnetwork inference. In International Conference on Machine Learning, 2021.
Alexander G. de G. Matthews, James Hensman, Richard Turner, and Zoubin Ghahramani. On sparse variational methods and the kullback-leibler divergence between stochastic processes. In International Conference on Artificial Intelligence and Statistics, 2016.
Mare Deisenroth and Carl E Rasmussen. Pilco: A model-hased and data-efficient approach policy search. In International Conference on Machine Learning, 2011.
John Denker and Yann LeCun. Transforming neural-net output levels to probability distributions. In Advances in Neural Information Processing Systems, 1991.
Michael W. Dusenberry, Ghassen Jerfel, Yeming Wen, Yi-An Ma, Jasper Snoek, Katherine Heller, Balaji Lakshminarayanan, and Dustin Tran. Efficient and scalable bayesian neural nets with rank-1 factors. In International Conference on Machine Learning. 2020.
Sebastian Farquhar, Michael Osborne, and Yarin Gal. Radial bayesian neural networks: Beyond discrete support in large-scale bayesian deep learning. In International Conference Artificial Intelligence and Statistics, 2020.
Angelos Filos, Sebastian Farquhar, Aidan N. Gomez, Tim G. J. Rudner. Zachary Kenton, Lewis Smith, Milad Alizadeh, Amoud de Kroon, and Yarin Gal. Benchmarking bayesian deep learning with diabetic retinopathy diagnosis. https: / github. cow/OATML///S /bdl -benchmarks, 2019.
10

---

Under review as a conference paper at ICLR 2022
Andrew Foong, Yingzhen Li, José Hemándeez-obatto and Richard Turner. 'in-between' uncertainty in bayesian neural networks. In JCML Workshop on Uncertainty and Robustness in Deep Learning. 2019.
Andrew Foong, David Burt, Yingzhen Li, and Richard Turner. On the expressiveness of approximate inference in bayesian neural networks. In Advances in Neural Information Processing Systems, 2020.
Vincent Fortuin, Priors in bayesian deep learning: A review, 2021.
Yarin Gal. Uncertainty in Deep Learning. Phl) thesis, University of Cambridge, 2016.
Yarin Gal and Zoubin Ghahramani. Dropout as a bayesian approximation: Representing model uncertainty in deep learning. In International Conference on Machine Learning, 2016.
Yarin Gal, Riashat Islam, and Zoubin Ghahramani. Deep Bayesian active learning with image data. In Proceedings of the 34th International Conference on Machine Learning. 2017.
M.N. Gibbs and DJJ.C. Mackay. Variational gaussian process classifiers. IEEE Transactions on Neural Networks, 2000.
Ian Goodfellow, Jonathon Shlens, and Christian Szegedy. Explaining and harnessing adversarial examples. In International Conference on Learning Representations, 2014.
Ian J. Goodfellow, Yoshua Bengio, and Aaron Courville. Deep Learning. MIT Press, 2016.
Alex Graves. Practical variational inference for neural networks. In Advances in Neural Information Processing Systems, 2011.
Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recognition. In Computer Vision and Pattern Recognition, 2016.
Dan Hendrycks and Thomas Dietterich. Benchmarking neural network robustness common corruptions and perturbations. In International Conference on Learning Representations, 2019.
James Hensman, Alexander Matthews, and Zoubin Ghahramani. Scalable Variational Gaussian Process Classification. In Proceedings of the Eighteenth International Conference on Artificial Intelligence and Statistics, 2015.
Daniel Hernandez-Lobald and Jose Miguel Hernandez-Lobato, Scalable gaussian process classifica- tion via expectation propagation. Proceedings of the 19th International Conference on Artificial Intelligence and Statistics, 2016.
José Miguel Hernández-Londe and Ryan P. Adams. Probabilistic backpropagation for scalable learning of bayesian neural networks. International Conference on Machine Learning, 2015.
Nils Lid Hjort, Chris Holmes, Peter Müller, and Stephen Walker. Bayesian nonparametrics. volume 28. Cambridge University Press, 2010.
Marius Hobbhahn, Agustinus Kristiadi, and Philipp Hennig. Fast predictive uncertainty for classifica- tion with bayesian deep networks, 2021.
Matthew D. Hoffman, David M. Blei, Chong Wang, and John Paisley. Stochastic variational inference. Journal of Machine Learning Research, 2013.
Alexander Immer, Maciej Korzepa, and Matthias Bauer. Improving predictions of bayesian neural nets via local linearization. In Proceedings of The 24th International Conference on Artificial Intelligence and Statistics, 2021.
Pavel Izmailov, Alexander Novikov, and Dmitry Kropotov. Scalable gaussian processes with billions of inducing inputs via tensor train decomposition. In Proceedings of the 21st International Conference on Artificial Intelligence and Statistics, 2018
Taejong Joo, Uijung Chung, and Min-Gwan Seo. Being bayesian about categorical probability. In International Conference Machine Learning, 2020.

---

Under review as a conference paper at ICLR 2022
Michael I. Jordan, Zoubin Ghahramani, Tommi S. Jaakkola, and Lawrence K. Saul. An introduction to variational methods for graphical models. Machine Learning, 1999.
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, 2019.
D. P. Kingma and L.. J. Ba. Adam: A method for stochastic optimization. International Conference on Learning Representations, 2015.
Durk Kingma, Tim Salimans, and Max Welling. Variational dropout and the local reparameterizatioo trick. In Advances in Neural Information Processing Systems, 2015.
Alex Krizhevsky. Learning multiple layers of features from tiny images. Technical report, University of Toronto, 2009.
Alex Krizhevsky, Ilya Sutskever, and Geoffrey E Hinton. Imagenet classification with deep convolu- tional neural networks. In Advances in Neural Information Processing Systems, 2012.
Balaji Lakshminarayanan, Alexander Pritzel, and Charles Blundell. Simple and scalable predictive uncertainty estimation using deep ensembles. Advances in Neural Information Processing Systems, 2017.
Yann LeCun, Corinna Cortes, and CJ Burges. Mnist handwritten digit database. ATT Labs, 2010.
Zhan Wei Lim, Mong Li Lee, Wynne Hsu, and Tien Yin Wong, Building trust in deep learning system towards automated disease detection. In Proceedings of the AAAI Conference on Artificial Intelligence, 2019.
Jeremiah Z. Liu, Zi Lin, Shreyas Padhy, Dustin Tran, Tania Bedrax-Wert and Balaji Lakshmi- narayanan. Simple and principled uncertainty estimation with deterministic deep learning via distance awareness. In Advances in Neural Information Processing Systems, 2020.
Chao Ma, Yingzhen Li, and Jose Miguel Hernandez-Lobato. Variational implicit processes. In International Conference on Machine Learning, 2019.
David J. C. MacKay. A practical bayesian framework for backpropagation networks. Neural Computation, 4(3), 1992.
Andrey Malinin and Mark Gales. Predictive uncertainty estimation via prior networks, In Advances in Neural Information Processing Systems, 2018.
Andrey Malinin and Mark Gales. Reverse kl-divergence training of prior networks: Improved uncertainty and adversarial robustness. In H. Wallach, H. Larochelle, A. Beygelzimer, F. d' d'Alché- Buc, E. Fox, and R. Garnet (eds.), Advances in Neural Information Processing Systems, 2019.
Andrey Malinin, Bruno Mlodozeniec, and Mark Gales. Ensemble distribution distillation. In International Conference on Learning Representations, 2020.
Rowan McAllister, Yarin Gal, Alex Kendall, Mark van der Wilk, Amar Shah, Roberto Cipolla, and Adrian Weller. Concrete problems for autonomous vehicle safety: Advantages of bayesian deep learning. In Proceedings of the Twenty-Sixth International Joint Conference on Artificial Intelligence, 2017.
Dimitrios Milios, Raffaello Camoriano, Pietro Michiardi, Lorenzo Rosasco, and Maurizio Filippone. Dirichlet-based gaussian processes for large-scale calibrated classification. In Advances in Neural Information Processing Systems, 2018.
Thomas Minka. Estimating dirichlet distribution, 2000.
Mahdi Pakdaman Nacini, Gregory F. Cooper, and Milos Hauskrecht. Obtaining well calibrated probabilities using bayesian binning. In Conference on Artificial Intelligence (AAAI), 2015.
Radford M. Neal. Bayesian Learning for Neural Networks. PhD thesis, University of Toronto, CAN, 1995.
12

---

Under review as a conference paper at ICLR 2022
Ian Osband, John Aslanides, and Albin Cassirer. Randomized prior functions for deep reinforcement learning. Advances in Neural Information Processing Systems, 2018.
Yaniv Ovadia, Emily Fertig, Jie Ren, Zachary Nado, D. Sculley, Sebastian Nowozin, Joshua Dillon, Balaji Lakshminarayanan, and Jasper Snoek. Can you trust your model's uncertainty? ? evaluating predictive uncertainty under dataset shift. In Advances in Neural Information Processing Systems, 2019.
Pingho Pan, Siddharth Swaroop, Alexander Immer, Runa Eschenhagen, Richard E Turner, and Mohammad Emtiyaz Khan. Continual deep learning by functional regularisation of memorable past. Advances in neural information processing systems, 2020.
Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer. James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, Alban Desmaison, Andreas Kopf. Edward Yang. Zachary DeVito, Martin Raison, Alykhan Tejani, Sasank Chilamkurthy, Benoit Steiner, Lu Fang, Junjie Bai, and Soumith Chintala. Pytorch: An imperative style, gegh-pefformncce deep learning library. In Advances in Neural Information Processing Systems, 2019.
Tim Pearce, Felix Leibfried, Alexandra Brintrup, Mohamed Zaki, and Andy Neely. Uncertainty in neural networks: Approximately bayesian ensembling. In International Conference Artificial Intelligence and Statistics, 2020.
Nicholas G. Polson, James G. Scott, and Jesse Windle. Bayesian inference for logistic models using polya-gamman gamma latent variables. Journal of the American Statistical Association, 2013.
Carl Edward Rasmussen and Christopher K. I. Williams. Gaussian Processes for Machine Learning. The MIT Press, 2005.
Thomas W Rauber, Tim Braun, and Karsten Berns. Probabilistic distance measures of the Dirichlet and Beta distributions. Pattern Recognition, 2008
Hippolyt Ritter, Aleksandar Botev, and David Barber. A scalable laplace approximation for neural networks. In International Conference on Learning Representations, 2018.
Tim G. J. Rudner, Zonghao Chen, and Yarin Gal. Rethinking function-space variational inference in bayesian neural networks. In Third Symposium on Advances in Approximate Bayesian Inference, 2021.
Hugh Salimbeni, Stefanos Eleftheriadis, and James Hensman. Natural gradients in practice: Non- conjugate variational inference in gaussian process models. In Proceedings of the 21st International Conference on Artificial Intelligence and Statistics, 2018.
Murat Sensoy, Lance Kaplan, and Melih Kandemir. Evidential deep learning to quantify classification uncertainty. In Advances In Neural Information Processing Systems, 2015 8.
Murat Sensoy, Maryam Saleki, Simon Julier, Reyhan Aydogan, and John Reid. Misclassification risk and uncertainty quantification in deep classifiers. In Proceedings of the IEFF/CVF Winter Conference on Applications of Computer Vision, 2021.
Jiaxin Shi, Mohammad Emtiyaz Khan, and Jun Zhu. Scalable training of inference networks for Gaussian- -ruusssanprocesss models. In International Conference on Machine Learning, 2019.
L. Smith and Y. Gal. Understanding Measures of Uncertainty for Adversarial Example Detection. In Conference on Uncertainty in Artificial Intelligence, 2018.
Shengyang Sun, Guodong Zhang, Jiaxin Shi, and Roger Grosse. Functional variational bayesian neural networks. In International Conference on Learning Representations, 2019.
T. Tao. An Introduction Measure Theory. Graduate studies in mathematics. American Mathematical Society, 2013. ISBN 9781470409227.
Yee Whye Teh. Dirichlet Process. Springer US, 2010.
13

---

Under review as a conference paper at ICLR 2022
Dimitris Tsipras, Shibani Santurkar, Logan Engstrom, Alexander Tumne, and Aleksander Madry. Ro- bustness may be at odds with accuracy. In International Conference on Learning Representations, 2019.
Ziyu Wang, Tongzheng Ren, Jun Zhu, and Bo Zhang. Function space particle optimization for Bayesian neural networks. In International Conference on Learning Representations, 2019.
Joe Watson, Jihao Andreas Lin, Pascal Klink, Joni Pajarinen, and Jan Peters. Latent derivative bayesian last layer networks. In Proceedings of the 24th International Conference on Artificial Intelligence and Statisticy, 2021a.
Joe Watson, Jihao Andreas Lin, Pascal Klink, and Jan Peters. Neural linear models with functional gaussian process priors. In Third Symposium on Advances in Approximate Bayesian Inference, 2021b.
Yeming Wen, Dustin Tran, and Jimmy Ba. Batchensemble: an alternative approach efficient ensemble and lifelong learning. In International Conference on Learning Representations, 2020.
Florian Wenzel, Théo Galy-Fajou, Christan Donner, Marius Kloft, and Manfred Opper. Efficient gaussian process classification using pólya-gamma data augmentation. In Proceedings of the AAAI Conference Artificial Intelligence, 2019.
Florian Wenzel, Kevin Roth, Bastiaan S Veeling, Jakub Swiętkowski, Linh Tran, Stephan Mandt, Jasper Snock, Tim Salimans, Rodolphe Jenatton, and Sebastian Nowozin. How good is the bayes posterior in deep neural networks really? In International Conference on Machine Learning. 2020.
Andrew Gordon Wilson and Pavel Izmailov. Bayesian deep learning and probabilistic perspective of generalization. In Advances in Neural Information Processing Systems, 2020.
Yao- Yuan Yang. Cyrus Rashtchian, Hongyang Zhang, Russ R. Salakhutdinov, and Kamalika Chaud- huri. A closer look at accuracy vs. robustness. In Advances in Neural Information Processing Systems, 2020.
Yao, W. Pan, S. Ghosh, and F, Doshi-Velez, Quality of uncertainty quantification for bayesian neural network inference. In ICML Workshop Uncertainty & Robustness in Deep Learning. 2019.
Guodong Zhang, Shengyang Sun, David Duvenaud, and Roger Grosse. Noisy natural gradient as variational inference. In International Conference on Machine Learning, 2018.
Ruqi Zhang, Chunyuan Li, Jianyi Zhang, Changyou Chen, and Andrew Gordon Wilson Cyclical stochastic gradient meme for bayesian deep learning. In International Conference on Learning Representations, 2020.
14

---

Under review as a conference paper at ICLR 2022
A RELATION TO PRIOR WORK
In this section, we discuss similarities and differences to prior work with Dirichlet posteriors in more detail by comparing assumptions and optimization objectives. To this end, first explain prior work using their own notation and then discuss the correspondence to our I'VI perspective.
A.I PRIOR NETWORKS
Malinin Gales (2018; 2019) use a neural network with parameters 0 to directly predict the concentration parameters ae of a Dirichlet distribution p(ulæ; 0), given input which is distinct from our approach of estimating posterior Dirichlet from M categorical predictions. This model is not Bayesian network in practice, as only point estimates for the weights are learned. To ensure 0e > 0, an element- exponential operation applied as the final layer of the neural network. Additionally, prior networks minimize optimization objective consisting of two separate KL divergences, representing in- and out-of- distrriibbttionn data respectively,
(Diriii(( [()) (m) (Dir(p à)) p(u x; 0)]] + pp.. [D (Dir(ulā) || p(u|2; 0)]]-
(7)
The first expectation Ep... accounts for the actual learning, i.e. fitting the training data, whereas the second expectation Epout is supposed to regularize the model by matching prior distribution. Accordingly, the first expectation is computed for the training data and can be compared to the expected log- ligellllioood term in fVI approach. Instead of maximizing the categorical log- likelihood of M observations, prior networks construct Dirichlet targets by smoothing categorical ground truth labels define the Dirichlet mean and setting the precision as hyperparameter during training. Although we also apply 'label smoothing' to the predictions, it is for numerical reasons and not for the construction of target distributions from labels. Additionally, prior networks treat the precision of their constructed target distribution as a hyperparameter, whereas we estimate the Dirichlet precision of our predicted variational posterior distribution via maximum likelihood. The second expectation is computed for OOD data and resembles the fKL term in our I'I approach, where the OOD data used as the index set. In contrast to prior networks, the more general fKL formulation also allows the training data or mixtures of training data and OOD data as index sets, whereas prior networks explicitly compute their second expectation for OOD data only. Furthermore, both prior network expectations consider the KL divergence from the neural network predictive distribution (right) to the target or prior distribution (left), whereas, in our fVI approach, and variational inference in general, the KI divergence from the prior distribution (right) to the variational posterior (left) is considered.
A.2 BELIEF MATCHING
Joo et al. (2020) assumes Dirichlet prior which, together with the categorical ground truth class labels, define a target Dirichlet posterior. A neural network is used to directly predict concentration parameters of a Dirichlet posterior 4z1æ by replacing the final softmax layer with an element wise exponential operation. To learn the target posterior, belief matching maximizes
LE (y. a W (x)) = E [log øn] - DRL. ly || P=]æe]
W
(8)
where Egate (log zy] is the expected log- lieelhooddo of the training data and D is the KL divergence between the predicted Dirichlet posterior and the Dirichlet prior. Therefore, their objective matches our IELBO objective equation except for two differences: Firstly, belief matching computes both the expected log-likelihood and the KL divergence with respect their single, directly predicted Dirichlet distribution, whereas we evaluate them as arithmetic averages of M stochastic categorical model outputs. Secondly, belief matching does not recognize the function- space aspect and instead only considers evaluation of the KI divergence using the training data, which resembles our fKL. where the index set is constrained to the training data.
A.3 EVIDENTIAL DEEP LEARNING
Sensoy al. (2018) directly predicts the concentration parameter of Dirichlet distribution by using a neural network with ReLU activations as final layer to assert the positive constraint. Additionally, a loss function is derived via type-If maximum likelihood by integrating over Dirichlet prior and the
15

---

Under review as a conference paper at ICLR 2022
sum of squares between target labels y, and predicted probabilities p,. Furthermore, regularizing KL divergence term is added, resulting in a total loss function,
(Vij - p.s)² Pcy(1 P (o(1-Po)) P(s)) f=1
(()) = Σ Σ
+ Ar DeDr,, Dir(p:[1)]
(9)
S. ++1
where Vij are individual 0-1 target labels, Pij are components of the predicted Dirichlet mean, S, is the predicted Dirichlet precision, a; is the predicted Dirichlet concentration parameter, is a vector of ones and X a annealing coefficient for optimization. The first part their loss is responsible for fitting the training data and can thus be compared to the maximum likelihood objective in Section 3.1. The MI. objective can be derived from the categorical log-likel log-likelihood via type-I maximum likelihood, whereas their objective is derived by minimizing the sum of squares via type- maximum likelihood. The second part of their loss resembles the fKL in our fVI approach. However, they only evaluate the KL divergence for the training data and explicitly consider the uniform Dirichlet distribution with concentration 1. Therefore, their KL divergence regularization term is a special of our fKL with the index set being the training data and the prior being the uniform Dirichlet distribution. For both parts, a major difference between evidential deep learning and our approach the realization of the predictive Dirichlet distribution. Evidential deep learning directly predicts Dirichlet concentration parameters, whereas we use M categorical predictions to estimate a Dirichlet distribution via maximum likelihood.
A.4 EXPERIMENTAL COMPARISON
In this subsection, we compare our MAP and MAP IVI models to this related work from Appendix A, namely Belief Matching and Prior Networks, which both demonstrated scalability to ResNet models. To reproduce their results, used the official open-source implementations'
Figure illustrates the test accuracy, log- ligeliklioood, and expected calibration error for the corrupted CIFAR10 image classification task (see Section 4.4). We trained the models using the same procedure and hyperparameters described Section F. The Belief Matching model corresponds closely to the MAP IVI model, both the objectives and models are similar. Unfortunately, we were not able to reproduce the Prior Networks performance described in the paper (Malinin & Gales, 2019), neither with their listed hyperparameters (Table 4, (Malinin & Gales, 2019)) the hyperparameters used
in Section F. It is uncertain whether this is due to the model, implementation bugs or unrecorded hyperparameters
MAP
MAP IVT
Belief Matching
Prior Networks
100
0.4
Log Log-Likee Likelihood
Accuracy
80
0.3
-1
ECE
0.2
60
0.1
-2
40
0
1
2
3
4
5
0
2
3
4
5
0
2
3
4
5
Corruption
Corruption
Corruption
Figure 8: Comparing Prior Networks and Belief Matching against MAP and MAP FVI models on the corrupted CIFARIO task.
com/kaosins "github. com KaosEngi iigginnes neer/ p orNetworks
'The authors did respond to personal correspondence regarding this matter.
16

---

Under review as a conference paper at ICLR 2022
B PRACTICAL CONSIDERATIONS
This section discusses practical issues with the approach in more detail.
B.I MINI-BATCH TRAINING
In weight-space variational inference, batch training requires an appropriate scaling term the KI. divergence of the ELBO during each batch iteration. This is to ensure the weight KI. is correctly balanced against the batch expected log-likelihood. The fKL approximation is equally mini-batched, as we restrained into using small portion of the index set each iteration. Therefore, like Sun et al. (2019), we divide both the batched expected log-like log-likelihood and the fKL by the mini-batch size B for numerical stability. The result of this is that the fKL weight in the total ELBO then depends on the mini-batch size, which is theoretically undesirable. However, in practice the KL divergence scaling variational inference is topic of active debate (Wenzel al., 2020; Aitchison, 2021) and frequently scaled or annealed for numerical reasons (e.g. Dusenberry et al. (2020)).
B.2 PRIOR SPECIFICATION FOR HIGH OUTPUT DIMENSION
When scaling to higher dimensional classification tasks, specifically 2 100, we observed numerical issues with the fELBO objective when using the uniform Dirichlet predictive prior. In higher dimensions, this prior would provide greater regularization. This is because the magnitude of the categorical likelihood does not change with dimensionality, it is the log probability of the label class. Conversely, the KL between two Dirichlet densities requires summing over the parameters, so the magnitude naturally increases with K. To alleviate this over-regularization, of Joo et al. (2020) and apply additional scaling to the KL term in the fELBO. This scaling can be shown be numerically equivalent to certain prior, i.e. 3 (Section 3.4, Joo al. (2020)). Therefore, optimizing this scaling is form of model selection. For our CIFAR 100 experiments, we simply chose scaling such that that the fKL magnitude was close the CIFARIO values. We found this to be about 0.1, which matches 10x scaling suggested by the Dirichlet KL due to the summation terms.
adopt the strategy
17

---

Under review as a conference paper at ICLR 2022
C DEEP STOCHASTIC CLASSIFIERS AS STOCHASTIC PROCESSES
Stochastic classifiers are stochastic processes instantiated as a feedforward or convolutional neural networks with stochastic weights. Formally, our stochastic process f defined on the sample space $2 with an index set X' defined by the data type (c.g. real vectors, RGB images), so f X x Ω + A K-1 where A K- is the state space, which the K-1 simplex. random variable f(x) Ω A k can be defined for each æ E X' and we write f(x) fa.
Kolmogorov's extension theorem (Tao, 2013) guarantees the existence of stochastic process for each I. E N the finite marginal joint distributions (f.z(fmr.z) where fax.z = {f=...... f satisfy exchangeability and consistency. These finite marginal joint distributions are implicitly defined according to Definition 3.1.
Exchangeability: For any permutation π I,..., L, Pr(an.)) (fr(an..) Pri.z (farz) ). This requires that the process behavior is invariant to the order of inputs. For feedforward neural network, this is satisfied because the respective predictions do not change the order of inputs changes. Consistency: For any 1 L' L, (f future evaluations are independent of past evaluations. For a feedforward neural network, this is satisfied because predictions do not depend previous predictions.
P1-L (f..() dfæ This requires that
D VARIATIONAL IMPLICIT PROCESS PARAMETERIZATIONS
This work popular classes of deep stochastic models used for weight-space BNNs, and interprets them as variational implicit processes in order to train them with function-space inference. To clarify this connection, Table I summarizes the connection between popular BNN models (and the MAP baseline) and weight distribution go (w)
Table 1: Summary of stochastic (and deterministic) representations for the network weights w, which correspond to popular BNN approaches. These distributions apply either to the whole network or per layer. Note Eg is typically factorized in practice. While use Radial BNNs rather than MFVI, include for completeness. Rankl references specifically the Gaussian realization.
Model
Parameterization for (-())
Scope
MAP
6(w wa) N Ee)
Network
MFVI
Layer
Radial
w 6& VE Ge. c/ccc e~ N(O,I), r (0,
Layer
MC Dropout
w ~ wg @ b, (ernouuli(pdáppoaa)
Layer
Ensemble
N7 Σ 8(w w
Network
Rankl
2m 5(w w m)) wh wa O dow 6w N(x vy v. up
Layer
18

---

Under review as a conference paper at ICLR 2022
MAP
MC Dropout
Ensemble
Subnetwork
MAP IVI
MC Dropout IVI
Ensemble IVI
100
0
0.8
Log- Lgg-Liiklihood
0.6
Accuracy
60
BCB
-4
0.4
40
0.2
20
8 0
0
0
45 90 135
180
45 90 135 Angle
180
45 90 135 180 Angle
Angle
Figure Metrics for rotated MNIST. models use a MLP architecture with two hidden layers with 50 units each. In terms of bikeliiel and expected calibration emror, the I'VI models outperform their respective baselines, which indicates improved uncertainty quantification. The lincarized laplace subnetwork metrics are taken from Daxberger et al. (2021), and uses ResNet-18 rather than MLP.
E DIRICHLET MUTUAL INFORMATION
Smith & Gal (2018) proposed to use the mutual information MI (0; y) of model parameters 0 and class label detect adversarial examples. Intuitively, this quantity measures how much the model could improve its parameters 0 y were observed. Therefore, it quantifies epistemic uncertainty and is also known information gain. However, Smith & Gal (2018) note that, in general, MI (0; y) intractable for deep learning models. Therefore, they provide a tractable approximation which applies models with stochastic parameters and requires sampling-based estimation.
Fortunately, in our function-space variational setting, where the model parameters i are deterministic, the one-hot encoded class label y follows categorical distribution y ~ Cat (-If=) and the latent probability vector fæ follows a Dirichlet distribution fa ~ Dir (-la). the equivalent mutual informa- tion between class labels and model predictions MI (fa; y) H (g) H (ylfx) can be derived in closed-fom form from the Dirichlet approximation using
K
Ok
Hly)
log
(10)
2
k=
ak (xb (ax + 1) & (++1)) 1)).
(ylfe)
(11)
z
The complete derivation provided in Appendix E.1 on the next page.
The first term H (y) corresponds to the categorical entropy of the expected value of fa- it equivalent the standard measure of predictive entropy for softmax-based ML deep learning classifiers. The second term H fflfx) is the conditional entropy of the class label given a Dirichlet observation fa and is influenced by the spread of fa, whereas the first term only depends on its mean. Although this function-space variational formulation admits closed-form solutions for MI (fa: y). the resulting values will still be approximations in our setting because the variational posterior Dir (fa (ax.) was estimated from M observations. Nonetheless, it is not limited models with stochastic weights, but instead provides a unified approach to measuring epistemic uncertainty for any classifier with Dirichlet predictive distribution.
19

---

Under review as a conference paper at ICLR 2022
E. MUTUAL INFORRMTION DERIVATION
Let y ~ Cat (-15a), fa ~ Dir (+la), Jan ~ Beta (-lxs,b)) bn). where on - or and aj - ax + 1, while by = an and 2 = Σk=1 k= ag.
MI (fwit) H (y) H (ylfa),
(12)
K
H (y) Σ
S, [f log (Esn [fmx))
(13)
K
fu
[ log (Bren [fak])
(14)
k=1
K
Okk
(Kk
log
(15)
A=1
H (y)1z)) fm
Σ Jan log (Sax)
(16)
k=1
K
Σ
Ey Ifa log (fax))
(17)
K
En Wa log (Jak)l
(18)
K
I Sax log(Swk) Jak a 1 Sank )on dlaks B(ax, bx)
Σ
(19)
K
Sax 2k (1 (mx) dsaks
log(
(00)
B(ak, bb)
K
faek a 11-1((1 (1 - feer)
x² -
log(f
dfakk
lim c+0
log (21)
B(ar, bk)
K
Jo
Dak fak ar d'a' (1 fax) 60 1
(22)
B(ax, bk)
K
θ
1'' Sax Shach-1 2 - faxt (n
(23)
B(ax, bk)
dak
K
1
Σ
B(ag. bb),
(24)
B(ag, bx) dal
K
Blak, (v)(()))))))))))))))))))())))( - (aah + bx)),
(25)
B(ax, bn)
K
4x B(ax, (b)((((x + 1) slax + + 6x)).
(26)
Baax, bu) as + by
ak (15(x + 1) saakk + +1+ bb)),,
(27)
aK + bb
K
Ok (1b (at + +1 - + (2 + 1))
(28)
z
20

---

Under review as a conference paper at ICLR 2022
F IMPLEMENTATION DETAILS AND COMPUTATIONAL COMPLEXITY
We implemented all models using the PyTorch library (Paszke et al., 2019) and all experiments were conducted using a i5-6600K CPU, GTX1070 GPU and GTX2080 GPU with less than 300 hours of total runtime.
The Two Moons data was generated by the library using 100 samples, 0.2 noise and random state 456, the PyTorch manual seed was set to 123. For this toy problem, all models were MLPs with two hidden layers consisting of 25 hidden units each, bias terms enabled and ReLU activation. For the Dropout models, the dropout rate was set to 0.2 and for the Ensemble models, we used 10 members per ensemble. All models were trained for 1000 epochs at a learning rate of 0.005 using the Adam optimizer (Kingma & Ba, 2015) with default parameters. The index set for the KL divergence was the visible 2D input plane, discretized at steps of 0.05. Since there was no mini-batch training, the KL term was not scaled according to Section B. This toy experiment is the only exception in this regard.
function from the scikit learn
For Rotated MNIST, used all 60000 images of shape 28x28x1 reshaped to 784 from the train split with pixel values normalized to [-1, 1] and no other preccesssssng or data augmentation. All 10000 images from the test set were used during evaluation, rotated by fixed degree, ranging from 0 to 180 in 10 steps, resulting in a total of 190000 test images. The MNIST (LeCun et al., 2010) data is available under the terms of the Creative Commons Attribution-Share Alike 3.0 license. All models were MLPs with two hidden layers consisting of 50 units each, bias terms enabled and ReLU activation. For the Dropout models, the dropout rate was set to 0.2 and for the Ensemble models, we used 10 members per ensemble. All models were trained for 30 epochs at a learning rate of 0.001 using a mini-batch size of 256. The index set for the KL divergence was the training data itself, except for the index set comparison section, where the different index sets are stated explicitly. Results were obtained using 10 random seeds.
For corrupted CIFARIO and CIFAR100, we used all 50000 images of shape 32x32x3 from the regular train splits. Following He et al. (2016), we normalized pixel values using the empirical mean and standard deviation, and employed data augmentation during training by first selecting random crops of size 32x32x3 after adding 4 pixels of zero padding to each side and then randomly flipping 50% of the images horizontally. All 10000 images from the regular test set were used during evaluation plus their corrupted versions (Hendrycks & Dietterich, 2019) with 19 different corruptions and 5 levels of severity, resulting in total of 960000 test images. The CIFAR1O and CIFAR 100 (Krizhevsky, 2009) data is available under the terms of the MIT License and the corrupted CIFAR10 and corrupted CIFAR 100 (Hendrycks Dietterich, 2019) data available under the terms of the Apache License 2.0. All models were CNNs following the ResNet-18 architecture (He et al., 2016), designed for CIFAR images, rather than ImageNet.
Adopting He et al. (2016), we trained with a batch size of 128 and used the SGD optimizer with momentum (0.9) for 200 epochs and scaled the learning rate by 0.1 at epochs 100 and 150. For the MAP, MC Dropout and Ensemble models without fVI, we used 0.0005 weight decay. For MC Dropout models (Gal & Ghahramani, 2016), the dropout rate was set to 0.2. For Ensemble (Lakshminarayanan et al., 2017) and Ensemble fVI, we used 5 members per ensemble.
For Radial BNNs (Farquhar et al., 2020) and Radial fVI, we implemented weight priors for all convolutional weights but not for the final linear layer. The standard deviation was parameterized using a = log(1 + exp(p)) and p was initialized to -5 while the means were initialized using the PyTorch default initialization scheme for CNNs. For Radial BNNs without IVI, we used a closed- form Gaussian weight KL divergence with a Gaussian prior with a mean of O and a standard deviation of O.1. For Radial fVI, we used our fKL instead of the weight-space KL.
For Rank1 BNNs (Dusenberry et al., 2020) and Rankl IVI, we used 4 ensemble members and 250 training epochs instead of 200 due to slow convergence and scaled the learning rate by 0.1 epochs 150 and 200. During training, we used implicit batch ensembling (Wen et al., 2020), whereas during prediction, we created explicit ensemble predictions by replicating the input. We placed Rank1 Gaussian distributions over all convolutional weights but not over the final linear layer. The standard deviation was parameterized using o = log(1 + exp(p)) and p was initialized to -3 while the means were initialized to 1. For R BNNs without IVVI, following Dusenberry et al. (2020), the Rank1 priors were Gaussian with a mean of I and a standard deviation of 0. 1, and weight decay of
21

---

Under review as a conference paper at ICLR 2022
0.0001 was used. We did not use KL annealing epochs. For Rankl IVI, used our fKL instead of the weight-space KL. For all fVI models, the index set for the KL divergence during fVI training was always the training data itself. Results were obtained using 10 random seeds.
To ensure numerical stability, we defined a minimum and maximum precision for the posterior Dirichlet estimation: Smin K and umax N, where K is the number of classes and N is the number of training examples. For the MAP models, we skipped the Dirichlet MLE and set 2 = 2ax because M 1. Similarly, we used M for the MC Dropout and Radial BNN models during training and also set = zmax, although we set M = 10 during evaluation. For the Ensemble models, M was always the number of members in the ensemble and for the Rank models, we replicated the input M times during evaluation, which results in M distinct predictions. Furthermore, we applied (m) , throughout all steps of the KL divergence estimation, where 7 was set to 10 for our experiments.
small amount of label smoothing fæk (m) R3 (1 f>
The equations from Minka's 's quasi-Newton maximum likelihood Dirichlet precision estimator (Minka, 2000), which we used for our implementation, translated our notation are
2 2
1
+
2
1
2
A
O
M
M
Jmk
[m)
M
Σ
M
log fak
(m)
(29)
K
k
A = M
vo(=) -
Σ
+
Σ
(30)
k=1
k=1
= M [1(z)) - Σ az (zāx)
(31)
where to is the digamma function and t is the trigamma function.
We initialized the algorithm with an approximate maximum likelihood solution using Stirling's approximation to the gamma function I", which was also presented by Minka (2000),
K
(0)
(32)
-2 āi (log dg log åg)
We stopped the algorithm once the change per step is less than 10-5 Counting the number of iterations until convergence for trained Ensemble f'I model with M 10 ensemble members and 10000 MNIST test examples, the mean was 3.0796, the 95th quantile was 3, the 99th quantile was 15 and the maximum was 117. Note that the number of iterations until convergence in vectorized mini-batch computation is equal to the maximum number of iterations until convergence of the items in the mini-batch.
Although the computational complexity of the underlying deep learning model depends on the model architecture, data input size, number of parameters, etc., for the following comparison, we assume that a single forward pass through the model takes ((1), i.e. constant amount of time, because the weight-space and function-space objectives share the same model. With a mini-batch size of B, computation of the standard ML objective takes O BM K) time per mini-batch iteration. Assuming a constant number of quasi-Newton steps, the Dirichlet precision estimation takes O(SK + M) time for index set of size S. Computing the fKL for index set of size S takes O(SMK) time. If the training data is used as index set the forward pass through the model can be shared between the ligeliillhood and fKL calculation, resulting in an overall asymptotic time complexity of ((MM)) per mini-batch iteration. In case of a different index set, the asymptotic time complexity becomes (KK)
22

---

Under review as a conference paper at ICLR 2022
G
EXPERIMENTAL RESULTS
All experimental results are reported here.
G.1
ROTATED MNIST
Table Accuracies for the rotated MNIST experiment. Means and standard errors over ten seeds. Best results within archetype in boldface, best results overall in blue.
MNIST Accuracy
Angle
MAP
96.82 = 0.06
91.19 0.15
87.17 = 0.43
71.60 0.76
53.11 0.76
36.03 0.52
25.18 0.46
MAP IVI
97.09 +: 0.05
95.01 0.13
88.51 = 0.32
0.51 73.53
0.57 54.32
36.03
0.40
24.36 0.35
MC Dropout
96.23
0.05
0.16 93:70
93.63
86.77 0.16
71.06 0.28
52.46
0.48
35.57
0.56
25.25 4 0.54
MC Dropout IVI
96.10
0.08
0.17
0.42 86.61
71.04 1 0.74
52.20
0.71
34.94 0.35
25.38 0.33
Ensemble
97.97 0.02
96.36 0.04
01.02-50
76.98 + 0.25
67.96
0.29
39.52 0.27
27.94 0.32
Ensemble IVI
98.09 + 0.02
96.71 7 0.02
91.95
0.15:=0.11 0.13
78 87 0.23
60.08 0.31
10.81 0.22
28. 20 0.29
Table 3: Accuracies for the rotated MNIST experiment. Mcans and standard errors ten seeds. Best results within archetype in boldface, best results overall blue.
MNIST Accuracy 1
Angle
70
80
90
100
110
120
MAP
18.60
0.39
15.01
0.34
12.81 1. 0.46
11.77 1 0.52
11.66 0.49
14.09 1. 0.33
MAP IVI
17.91 + 0.42
14.72
13.33 + 0.48
13.02 + 0.40
13.33 + 0.28
15.38 + 0.21
MC Dropout
20.26
0.48
17.22 1.22..0..4 0.47
15.11 +: 0.43
13.41 :: 0.38
12.67 1: 0.35
14.32 1: 0.23
MC Dropout IVI
20.99 + 0.39
18.22 + 0.34
16.26 + 0.39
14.82 + 0.46
13.69 = 0.38
14.44 + 0.29
Ensemble
19.90 + 0.31
16.11 + 0.32
13.62 + 0.31
12.09 = 0.26
12.31 + 0.22
14.92 + 0.21
Ensemble IVI
20.08 0.38
16.43 0.37
14.28
0.30
13.07 0.21
13.19 1 0.16
15.43 1: 0.12
Table Accuracies for the rotated MNIST experiment. Means and standard errors over ten seeds. Best results within archetype in boldface, best results overall in blue.
MNIST Accuracy f
Angle
130
140
150
160
170
180
MAP
16.92 0.43
19.05
0.56
22.36 36
0.55
25.06
0.64
26.95 +: 0.58
28.58 1: 0.57
MAP IVI
17.94
0.29
19.82 = 0.43
22.45 - 0.54
25.26 - 0.46
27.42 0.43
28.93 0.42
MC Dropout
17.35 11.35:1.000 0.34
20.11 0.49
22.82 0.43
24.37
0.39
26.03 1: 0.43
28.14 0.43
MC Dropout IVI
0.35
20.25 - 0.40
23.25 - 0.34 0.34
25.24
26.86 = 0.18
28.95 = 0.29
Ensemble
18.02 0.332
20.83 0.229
24.49
0.226
27.61 0.20
30.03 0.24
31.27 1 0.32
Ensemble IVI
18.47=0.J1
21.49 + 0.20
25.21 0.21
28.54
31.31 = 0.29
32.51 = 0.31
Table 5: Log-like log-iikels for the rotated MNIST experiment. Means and standard errors over ten seeds. Best results within archetype in boldface, best results overall in blue.
MNIST
Angle
Log-Likelibood
MAP
0.11 1 0.00
-0.20 0.00
-0.48 0.02
11.20 = 0.03
-2.36 1: 0.04
-3,80 1. 0.04
-5,05 1. 0.05
0.00
0.01
0.02
0.08
0.08
MC Dropout
-0.15 0.15::000
-0.23 0.000
-0.43 0.00
-0.88 0.01
-1.50 0.02
-2.28 0.03
-2.00
Dropout IVI
1.0001
Ensemble
-0.07 0.00
0.012:0000
-0.29 0.00
-0.73
-1.50
-..06
+..00
Ensemble
-0.17.0.00
-0.67
-1.20
-1.85
23

---

Under review as a conference paper at ICLR 2022
Table 6: Log-likelihooke for the rotated MNIST experiment. Means and standard errors over ten seeds. Best results within archetype boldface, best results overall blue.
MNIST Log-Likelhhod 1
Angle
70
MAP
-6.05
0.007
-7.22 0.09
0.00
-7.27
-6.95
0.000
MAP IVI
-3.54
0.04
-4.14
-..8
-..16
-4.05 0.00
MC Dropout
0.04
0.01
0.05
0..2
-.02
0.05
MC Dropout IVI
0.08
-.04
-.00
-.025
-3.28
0.04
-3.17
Ensemble
0.01
5.00
0.06
0.00
-5.86
-5.68 0.06
Ensemble IVI
-2.89 0.02
-3.27
0.02
-3.49 0.02
-3.57
33.57
3.57 0.02
-3.47 +
Table 7: ligllikelihoods for the rotated MNIST experiment. Means and standard errors over seeds. Best results within archetype in boldface, best results overall in blue.
MNIST LoggLiielihood f
Angle
130
140
150
160
170
180
MAP
-6.57 +: 0.09
-6.36 +: 0.08
-6.20 +: 0.07
-6.16 0.06
-6.13
0.05
-6.00 +: 0.07
MAP
-3.87 +: 0.05
-3.78 0.04
-3.68 0.04
-3.63 + 0.03
-3.61
0.04
-3.50 + 0.04
MC Dropout
-3.93 +: 0.05
-3.85 +: 0.06
-3.87 +: 0.05
-4.04 -4.04
0.04
-4.15
-4.19 0.04
MC Dropout IVI
-3.03 0.02
-2.94
-291 + 0.02
-2.97
-3.00
Ensemble
-5.37 0.05
-5.24
-5.18
-5.31 0.01
-5.37
0.02
-5.25 0.02
Ensemble [VI
3.29 0.02
3.22
3.17
3.21
3.23
0.02
3.16
Table Expected calibration errors for the rotated MNIST experiment. Means and standard errors ten seeds. Best results within archetype in boldface, best results overall in blue.
MNIST ECE
Angle 30
10
20
40
50
60
MAP
0.01 0.00
0.02 = 0.00
0.06
0.00
0.17 0.01
0.32 = 0.01
0.46 : 0.00
0.56 0.00
MAP IVI
0.05 0.00
0.05 + 0.00
0.04
0.00
0.03 + 0.00
0.15
0.01
0.28 + 0.01
0.38
0.01
MC Dropout
0.04 0.00
0.05 + 0.00
0.05
0.00
0.02 + 0.00
0.01 0.13
0.25 +: 0.01
0.33 0.01
MC Dropout
0.09 0.00
0.10
0.10 0.00
0.04 0.01
0.07 10071:00
0.19 0.01
0.27 0.01
Ensemble
0.01 0.00
0.02 = 0.00
0.01 0.00
0.04 + 0.00
0.15 0.00
0.27 + 0.00
0.36 0.00
Ensemble IVI
0.06 = 0.00
0.08 = 0.00
0.10 0.00
0.08 = 0.00
0.02 0.00
0.14 = 0.00
0.23 0.00
Table Expected calibration errors for the rotated MNIST experiment. Means and standard errors over ten seeds. Best results within archetype in boldface, best results overall in blue.
MNIST ECE
Angle a
70
80
110
120
MAP
0.62 = 0.00
0.65 + 0.01
0.68 + 0.01
0.69 + 0.01
0.69 + 0.01
0.67
0.01
MAP IVI
0.44 + 0.01
0.48 + 0.01
0.50 + 0.01
0.50 + 0.01
0.51 + 0.01
0.50
0.00
MC Dropout
0.39 0.01
0.43 + 0.01
0.46 0.46
0.48 0.01
0.50 0.01
0.48 0.01
MC Dropout IVI
0.32
0.36 0.01
0.38 0.01
0.40
0.42 0.01
0.40 0.01
Ensemble
0.43 = 0.00
0.46 + 0.00
0.49 + 0.00
0.51 + 0.01
0.32 + 0.00
0.50 0.00
Ensemble IVI
0.31 + 0.00
0.35 + 0.00
0.38 0.00
0.40 + 0.00
0.41 + 0.00
0.40 0.00
Table 10: Expected calibration errors for the rotated MNIST experiment. Means and standard errors over ten seeds. Best results within archetype in boldface, best results overall in blue.
MNIST ECE
Angle
130
140
150
160
170
180
MAP
0.65 + 0.01
0.63 + 0.01
0.60 + 0.01
0.58 + 0.00
0.57 + 0.01
0.56 + 0.01
MAP IVI
0.48 : 0.00
0.47 + 0.01
0.45 - 0.01
0.43 0.00
0.41 + 0.01
0.40 + 0.01
MC Dropout
0.45 - 0.01
0.42 1: 0.01
0.41 0.011
0.41 + 0.01
0.41 0.41 0.01
0.40 1:0::00 0.00
MC Dropout IVI
0.37 = 0.00
0.35 1 0.00
0.33 0.000
0.440.00
0.33
0.32 1 0.00
Ensemble
0.48 + 0.00
0.46
0.00
0.43 0.00
0.42 0.00
0.42 0.00
Ensemble IVI
0.38 + 0.00
0.35 + 0.00
0.33 333:00000
0.32 = 0.00
0.31 0.00
0.30 0.00
24

---

Under review as a conference paper at ICLR 2022
G.2 CIFARI0
Table 11: Accuracies for the corrupted CIFARTO experiment. Means and standard errors over ten seeds. Best results within archetype in boldface, best results overall in blue.
CIFARIO Accuracy f
Corruption Severity
0
2
MAP
94.32 +
87.65 0.08
81.76
75.97 0.13
68.86 0.20
57.28 0.21
MAP IVI
94.40 0.08
87.66 + 0.06
81.66 + 0.10
75.81 0.16
68.77 + 0.17
57.06 + 0.18
MC Dropout
94.32 0.04
88.21 0.07
82.13 0.13
75.81
0.19
67.59 0.21
55.72 0.26
MC Dropout IVI
93.38 + 0.03
87.01 0.09
80.64 0.14
74.36 0.18
66.33 + 0.18
54.87 = 0.19
Ensemble
95.30
89.37 0.03
83.77 0.06
78.21 0.07
71.05
59.33 + 0.16
Ensemble IVI
95.26 0.03
89.44 0.03
43.94
78.49 0.08
71.42 0.12
59.73 = 0.16
Radial
95.05 0.04
87.98 0.07
81.82 +: 0.10
75.93 0.12
68.93 0.16
57.24 1 0.21
Radial IVI
93.73 = 0.03
87,43 + 0.07
76.09 + 0.21
68.84 = 0.29
57.42 = 0.35
Rank1
93.68 0.05
87,60 1.:00+00.00 0.05
82.32 + 0.08
76.90 + 0.08
69.92 + 0.13
58.53 + 0.17
Rankl IVI
93.91 = 0.04
87.75 + 0.05
82.38 0.09
76.77 = 0.14
69.69 + 0.17
58.23 = 0.19
Subnetwork
91.00 = 0.00
83,00 1.00
77.00
08.00 = 1.00
6
59.00 + 0.00
Belief Matching
94.52 + 0.03
86.98 + 0.12
80.39 0.21
73.62 0.29
65.74 + 0.35
53.67 1 0.36
Prior Networks
66.65 = 0.61
61.28 + 0.43
57.87 = 0.37
54.71 = 0.35
50.24 = 0.33
42.29 = 0.31
Table 12: Log-likelihoods for the corrupted CIFARIO experiment. Means and standard errors over ten seeds. Best results within archetype in boldface, best results overall in blue.
CIFARIO Log-Likelihood f
Corruption Severity
MAP MAP IVI
0.22
-055 0.00
-0.80 0.01
1.01
-1.52
-2,20 0.02
0.00
-0.48
0.00
0.00
-1.16=0.01
-
MC Dropout
0.17
0.00
0.01
0..03
-1.36
-2.09 0.002
MC Dropout
0.00
0.00
-0.64
0..5
-1.12
-1.56
Ensemble
-0.1510000
-0.35 0.00
-0.54 0.00
0.70
=1.03
1,501
Ensemble IVI
-0.01
-0338 0.00
-0.55 0.00
0.72
-.09
-133 + 0.01
Radial
-0.21 0.00
-0.58 0.00
-0.33 + 0.01
-1 +
-1.79
-2.61 + 0.02
Radial IVI
-0.28 +:
0.49 0.00
0.699 0.01
0.90
-117=
-1.61 0.02
Rank1 Rankl IVI
-0.33 + 0.00
-0.71 + 0.01
-1.06 0.01
-1.46 0.01
-2.02 0.02
-3.00
0.03
0.27 0.00
0.477 0.00
-0.65 0.00
0.85
=.01
-153
0.01
Subnetwork
0.00
-0.51
-0.73
-1.25 -1.25
47 0.03
Belief Matching
-0.26 0.00
-0.51 0.00
-0.73
0..97
-1.26
-170 0.02
Prior Networks
-1.45 1:: 0.02
-1.59 - 0.01
-1.66 0.01
-1.73 0.01
-1.85 0.01
-2.06 + 0.01
Table 13: Expected calibration errors for the corrupted CIFAR10 experiment. Means and standard errors over ten seeds. Best results within archetype in boldface, best results overall in blue.
CIFARIO ECB
Corruption Severity
0
3
4
MAP MAP IVI
0.03 0.00
0.08 + 0.00
0.12 0.00
0.16 + 0.00
0.22 + 0.00
0.30 + 0.00
0.05 + 0.00
0.05 + 0.00
0.05 0.00
0.09
0.00
0.14 + 0.00
0.22 0.00
MC Dropout
0.01 = 0.00
0.03 + 0.00
0.06 1:06:00 0.00
0.09 = 0.00
0.15 + 0.00
0.23 + 0.00
MC Dropout IVI
0.06 = 0.00
0.03 + 0.00
0.03 + 0.00
0.05 + 0.00
0.09 + 0.00
0.17 + 0.00
Ensemble
0.01 = 0.00
+ 0.00
0.08 0.00
0.12 0.00
Ensemble
0.07 = 0.00
0.05 + 0.00
0.03 0.00
0.04 0.00
0.06 + 0.00
0.11 0.00
Radial
0.03 = 0.00
0.08 + 0.00
0.13 0.00
0.17 + 0.00
0.23 + 0.00
0.32 + 0.00
Radial IVI
0.05 + 0.00
0.05 + 0.00
0.05 + 0.00
0.09 + 0.00
0.14 0.00
0.23 0.00
Rank1
0.04 0.00
0.09 0.00
0.13
0.17 + 0.00
0.23 0.00
0.32
Rank1 IVI
0.05 = 0.00
+ 0.00
0.00
0.07 = 0.00
0.12 0.00
0.20 0.00
Subnetwork
0.01 = 0.00
0.03 0.00
0.06 1N6100 0.00
0.11
0.13
0.16
Belief Matching
0.07 = 0.00
0.07 0.00
0.07 0.00
0.00
0.14
0.44
Prior Networks
0.20 - 0.00
0.21 0.00
0.21
0.00
0.23
0.24 0.00
25

---

Under review as a conference paper at ICLR 2022
Table 14: Accuracies for the CIFARIO adversarial attack experiment. Means and standard errors over ten seeds. Best results within archetype in boldface, best results overall in blue.
CIFARIO
Adversurial Attack Epsilon
Accuracy
0,00
0.05
0.10
0.15
0.20
0.25
0.30
MAP
94.32 + 0.05
42.35
0.23
35.61
0.31
31.93 + 0.38
28.04
0.41
23.60
0.40
19.65 4 0.34
MAP IVI
94.40 0.08
70.04 1 0.17
61.10
0.20
51.32
0.47
0.68 40.86
31.60
0.77
24.80 1: 0.63
Dropout
32.76
0.44
MC Dropout IVI
93.13
0.04
56.99 0.16
0.29 49.91
0.39 13.75
0.57 36.97
29.86 0.68
23.65 0.68
Ensemble
0.133
Pasemble FVT
0.03 95.26
60.12 0.10
49.34 0.21
39.38 0.36
30.59 0.42
23.73 0.43
18.92 + 0.39
Radial
0184
27.92 0.17
18. = 0.18
14.98 + 0.22
12.69 1: 0.25
11:28
0.22
10:47 0,20
Radial FVI
93.72 0.03
07.33 0.13
59.92 0.25
02 07
0.40
43
0.60
35. 35.32
0.64
28,44 0.66
Rank1
93..5
0.05
19.63
D.21
9.66 0.18
8.34 0.16
8.47 0.17
8,77 0.20
9.11 0.19
Rank IVI
93.86 0.04
67.89
0.14
58.99 0.12
49.88 0.24
40.48 0.36
31.82 0.5.1
24.89 0.57
Table 15: Log-likelihoods for the CIFARIO adversarial attack experiment. Means and standard errors over ten seeds. Best results within archetype in boldface, best results overall in blue.
CIFARIO
Adversarial Attack Epsilon
Log-LLkkelhhock
MAP
-0.22 0.00
-4.00
+.00
-458 0.02
-4.71 + 0:003
-4.87 + 0.04
-5.16 0.00
-5.53 1.3-+-00 0.06
MC Dropout
0.000
3.26 1.0.0
1.402
4.61 1 0.02
-4.00 0.02
4.699 1.0.0
4.96 0.07
Dropout IVI
2..22
Ensemble
0.000
-..00
-3.61
-3.91 0.011
-1.07 0.01
-4.21 ..202
-4.32 0.08
Ensemble
-0.21
0.000
-.00
-1.98:0.01
-2.59
-3.00
Radial
-0.21 0.00
-0.02
0.02
-6.78 0.04
-7.01 0.066
-7.30 0.007
-7.56 0.09
Radial INI
-0.28
-1.00
0.01
-2.77 0.003
0.03
Rankl
0..03:20000
-7.81
0.03
-9.39 0.08
-9.23
-9.08 1.00
-9.00 0.07
Rank1
00.27 0.00
-146
-1.75
2.04 0.01
2.34 1
-2666
0.004
Table 16: Expected calibration errors for the CIFARIO adversarial attack experiment. Means and standard errors over ten seeds. Best results within archetype in boldface, best results overall in blue.
CIFARIO ECE
Adversarial Attack Epsilon
0.00
0.05
0.10
0.15
0.20
0.25
0.30
MAP
0.03 0.00
0.48 = 0.00
0.53 0.00
0.55 = 0.00
0.58 + 0.00
0.61 + 0.00
0.65 0.01
0.00
0.19 0.00
0.24 0.00
0.30 - 0.00
0.38 : 0.011
0.45 0.01
0.52 0.01
Dropout
0.00
4:
0.53 0.00
- 0.00
0.52 0.00
0.52 + 0.00
0.54
MC Dropout IVI
0.06 0.00
0.28 = 0.00
0.32 0.00
0.34 : 0.00
0.35 0.00
0.38 + 0.01
0.41 0.01
Ensemble
0.01 0.00
0.39 4: 0.00
0.51 0.00
0.54 = 0.00
0.56 + 0.00
0.57 + 0.00
0.58 + 0.00
Ensemble IVI
0.07 0.00
0.20 0.00
0.27 0.00
0.32 + 0.00
0.37 0.00
0.41 + 0.00
0.45 0.00
Radial
0.03 0.00
0.62 - 0.00
0.70 0.00
0.71 + 0.00
0.73 0.00
0.74 0.00
0.75 0.01
Radial
0.05 0.00
0.21
0.26 0.00
0.30 0.00
0.42 0.00
0.47 0.01
Rankl
0.04 00
0.76
0.86 0.00
0.86
0.83 0.00
0.81 + 0.00
0.80 0.00
Rankl IVI
0.05 = 0.00
0.20 0.00
0.25 0.00
0.30 0.00
0.361:00 0.00
0.43 0.01
0.49 0.49 0.01
26

---

Under review as a conference paper at ICLR 2022
G.3
CIFARI00
Table 17: Accuracies for the corrupted CIFARI00 experiment. Means and standard errors over ten seeds. Best results within archetype in boldface, best results overall in blue.
CIFARI00 Accuracy 1
Corruption Severity
MAP
75.68 = 0.07
64.17 0.06
56.41 = 0.07
49.78 = 0.07
43.11 + 0.08
33.03 = 0.08
MAP IVI
74.77 0.09
64.26 0.07
55.82
50.31 0.11
43.56 0.12
33.81 0.11
MC Dropout
74.15
MC Dropout
0.12
61.03 0.10
0.11
46.46 0.10
39.88 + 0.09
30.87 = 0.10
Ensemble
79.13 0.05
68.00
59.19
53.42
0.07
46.45
0.06
35.69 0.06
Ensemble fV1
0.06
66.38 0.07
57.97 0.09
62.23 0.09
45.14 0.09
35.19 0.11
Radial
76.40
0.07
54.68
0.06
0.05
4..29
31.89 + 0.07
Radial IVI
75.29 2.19
64.84 0.11
56.49
50196
0.10
3 0.09
Rankl
73.68 0.10
63.48 0.09
0.12
49.92
43.45
33.87 0.11
R IVT
75.56 0.10
65.49 0.08
57.55 + 0.11
52.24+0.10
45.63 + 0.11
35.73 + 0.11
Table 18: Log-liken Log-likelihoods for the corrupted CIFARI00 experiment. Mcans and standard errors over seeds. Best results within archetype in boldface, best results overall in blue.
CIFARI00
Coruuptio Severity
Log-Likell f
MAP
-1.00 +: 0.00
-1.99
0.00
-2.09
0.01
0.01
-2.99
2:99=000
-3.73 0.01
MAP IVI
-1.20 + 0.00
-1.69
0.00
-2.08
0.01
-2.35
+2.55+00.0
-2.69 1 0.01
-3.19
0.01
MC Dropout
-0.97 0.00
-151
-2.02
0.01
-2.42+00.01
-2.96
-3.76
MC Dropout IVT
-1.17 + 0.00
-1.65 0.01
-2.09 + 0.01
-.039+0.01
-2.79 0.00
-3.35 + 0.01
Ensemble
0.81
1.30
1.70
1.98
2.37=0011
2.93 0.0.1
Ensemble IVI
-1.18 + 0.00
-1.61 + 0.00
-1.98 + 0.00
+.00
-258
0.00
-3.06 0.01
Radial
0.98 1: 0.00
-1.66 1.66::0001
0.01
2.602
3321 0.02
4.00 100-10 0.03
Radial IVI
-1.21
+.00
Rank1
-1.48
-2.29 0.00
33.01
0.01
-3.50
-4.8
0.02
-.57
Rank1 IVI
-1.17 + 0.00
-1.64+0.00 0.00
-2.00 0.01
-2.25 + 0.01
-2.58
0.01
-3.06 + 0.01
Table 19: Expected calibration errors for the corrupted CIFAR 100 experiment. Means and standard errors over ten seeds. Best results within archetype in boldface, best results overall in blue.
CIFARIO0 ECE 4
Corruption Severity
0
3
4
MAP
0.08 0.00
0.12 + 0.00
0.16 + 0.00
0.20
0.00
0.23 + 0.00
0.30 0.00
MAP
0.12 + 0.00
0.11 0.00
0.09
0.00
0.06
0.00
0.04
0.00
0.04 0.00
MC Dropout
0.02 = 0.00
0.06 + 0.00
0.10 0.00
0.13 = 0.00
0.17 + 0.00
0.21 + 0.00
MC Dropout IVI
0.08 = 0.00
0.06 + 0.00
0.02 + 0.00
0.01 + 0.00
0.05 + 0.00
0.10 0.00
Ensemble
0.05 0.00
0.04 0.00
0.02 0.00
0.02 0.00
0.04 0.00
0.10 0.00
Ensemble IVI
0.18 = 0.00
0.18 0.00
0.10
0.14 0.00
0.10 = 0.00
0.04 0.00
Radial
0.09 = 0.00
0
0.00
0.00
0.27 + 0.00
0.34 + 0.00
Radial IVI
0.13 = 0.00
0.12 0.00
0.00
0.08 108-000 0.00
0.05 + 0.00
0.04 + 0.00
Rank1
0.16 = 0.00
0.22 1
0.28 0.00
0.32
0.37 0.00
0.44 0.00
Rank1 IVI
0.13 = 0.00
0.13 0.00
0.00
0.00
0.06 0.00
0.03 0.00
27

---

Under review as a conference paper at ICLR 2022
Table 20: Accuracies for the CIFARI00 adversarial attack experiment. Means and standard errors over ten seeds. Best results within archetype in boldface, best results overall in blue.
CIFARI00
Adversarial Attack Epsilon
Accuracy
0.00
0.05
0.10
0.15
0.20
0.25
0.30
MAP
118
0.13
6.70
1.88 0.08
1.18 0.07
3.62
MAP IVI
74.77 1: 0.09
15.83
0.11
9.79 0.12
7.08 1: 0.09
5.38 + 0.05
4.30 +: 0.08
3.46 1: 0.07
Dropout
74.05 0.08
19.33
0.10
10.56 0.08
7.45 7.45 0.09
5.90 0.09
4.91 = 0.06
4.01 0.07
MC Drupout IVI
71.61 0.10
21.60
1 0.07
13.22 1 0.09
9.52 0.09
7.19 0.12
5.62 1. 0.11
4.49 0.11
Ensemble
79.13 0.05
26.04 0,05
13.51 = 0.09
8.64 0.07
6,32 + 0.06
4.92 = 0.07
3.97 1: 0.U7
Ensemble
0.07
Radial
#..02
12.05 0.08
7.36 0.06
0.505
4.32 0.07
3.77 +
3.19 0.08
Radial IVI
75.29 0..99:2:0.00 0.10
16 85
+: 0.12
10.97 = 0.09
7.91 0.00
6.06 0.10
4.83 = 0.12
3.99 0.10
Rank1
73.76 0.07
14.01 = 0.09
9.51 0.08
7.63 0.12
6.18 0.12
5.09 0.14
4.20 + 0.13
Rank1 IVI
75.58 0.09
18.34 0.08
11.68 0.07
3.151
6.67 + 0.09
5.29 + 0.08
4.38 :: 0.08
Table 21: Log-likelihoods for the CIFAR100 adversarial attack experiment. Means and standard errors over ten seeds. Best results within archetype in boldface, best results overall in blue.
CIFARIOC Log-Likelitodd
Adversarial Attack Epsilon
0.00
0.05
0.10
0.15
0.20
0.25
0.30
MAP
1.00 0.00
6.64
0.03
6.72 4: 0.03
6.53 +: 0.03
6.52 4: 0.05
6.84
0.06
6.76 0.08
MAP IVI
0.00
-5.36
0.01
-5.62 0.02
-5.53 1:5+t000 0.02
-5.45 0.02
-5.40
0.02
-5.36
+
MC Dropout
0.97 0.00
1.002 0.01
-5.93 0.02
-6.08 0.02
-6233 0.02
-6.16 0.03
-6.71 0.05
MC Dropout IVI
1.17 0.00
0.01 4.36
5.17 0.02
5.35 +: 0.01
5.45 1: 0.01
0.01 5.55
5.66 0.02
Ensemble
Ensemble IVI
1.18 0.00
-3.97 0.000
1.00
-5.10 5.10
-5.15 0.00
-5,17 -5.17 0.01
-5.18.10.008
-6.73
-6.50
-6.79
-..04
-7.28
Radial IVI
-1.21
-5.11:0.01
5.366
-5.29 0.009
-5.22 0.02
-5.18
-5.16 0.02
Rankl
0.00
0.02
-10.09
0..076
-10.1100001
Rankl
-1.17
-4.55 0.001
5.28
5.2211.001
-3.17
-5.16
-5.16
Table 22: Expected calibration errors for the CIFAR100 adversarial attack experiment. Means and standard errors over seeds. Best results within archetype in boldface, best results overall in blue.
CIFAR100 ECE.
Adversurial Attack Epsilon
0.00
0.05
0.10
0.15
0.20
0.25
0.30
MAP
0.08=00
0.55 =: 0.00
0.52::0000
0.49 1: 0.00
0.49 0.01
0.49 :: 0.01
0.51 = 0.01
MAP IVI
0.12=0.00
0.23 = 0.00
0.23 :: 0.00
0.22 = 0.00
0.23 0.00
0.23 = 0.01
0.23 0.01
Dropout
0.03 0.00
0.45 - 0.00
0.47 0.00
0.44 : 0.00
0.43 0.00
0.43 = 0.00
0.44
MC Dropout IVI
0.09
0.26 0.00
0.29
0.28 0.00
0.28 0.00
0.29 - 0.00
0.30 0.00
Ensemble
0.05 0.00
0.26 + 0.00
0.31 0.00
0.30 + 0.00
0.30 0.00
0.30 + 0.01
0.29 0:01
Ensemble IVI
0.18 0.00
0.08 = 0.00
0.14 :: 0.00
0.15 = 0.00
0.16 4: 0.00
0.17 + 0.00
0.18 0.00
Radial
0.09 0.00
0.57 = 0.00
0.54 0.00
0.51 0.00
0.51 0.00
0.53 + 0.01
0.55 0.01
Radial
0.00
0.00
Rankl
0.16
0.76
0.73
0.71
0.71 0.00
0.70 0.01
Rankl IVI
0.13 113=000 0.00
0.20 0.00
0.21 0.00
0.20:0.000
0.20 0.00
0.20 4 0.00
0.20 0.01
28

---

Under review as a conference paper at ICLR 2022
H ABLATION STUDIES
SAMPLES DURING TRAINING
In Section 3.3, Dirichlet estimation procedure was proposed using M samples. In the single sample case M 1, motivated by MAP models, crude approximation was proposed to approximate the precision with the number of training data samples.
For many VIPs used in BNNs, the computational cost of sampling during training means that only 1 sample used in practice. This means that the sample approximation is used for more than just MAP, but also for MC Dropout, Radial and Rank1 BNNs as well in practice.
To assess the consequence of this approximation, we repeated the CIFARIO corruption experiment for MC dropout, but with 5 samples to reflect that of ensembles. Figure 10 shows the result. It can be seen that the 5 sample MC Dropout performance closer to the sample MC Dropout performance than the Ensemble. This result indicates that the VIP model, rather than paaples-during-raining has greater impact. The similarity in performance between and 5 sample MC Dropout suggests the sample approximation is reasonable for the training time savings.
MC Dropout (1 sample)
MC Dropout (5 samples)
Ensemble (5 samples)
MC Dropout IVI (I sample)
MC Dropout IVI (S samples)
Ensemble IVI (5 samples)
100
0.4
Log- Likelihood
Accuracy
-1
0.3
80
BCB
0.2
-2
60
-3
0.1
0
0
2
4
5
0
3
5
3
Figure 10: Reproduction of CIFARIO corruption results in Figure including MC Dropout results with 5 predictive samples during training. For 5 sample MC Dropout, 3 random seeds were used rather than 10 due to the additional training time.
29

---

Under review as a conference paper at ICLR 2022
H.2 SCALING ISSUES WITH HIGH LABEL DIMENSIONALITY
The CIFARION 100 experiments revealed an issue with the fELBO objective that caused underfitiing for the larger label dimension. To examine why, recall that the categorical likelihood is fag when Uk =1 Therefore, the dimensionality of y does not directly influence the value. Conversely, the KL divergence between two Dirichlet densities DKL(pI) p2) does incorporate the label dimensionality K (Rauber et al., 2008),
K - log [(akk - log (2()() +
DKI(PillPe) = log (()
Σ log I(akk) Arm
(1) a ag
(2) (V ((
To counteract this linear increase due to the summation terms, assess a heuristic annealing scale factor on the fKI. during training of 1/K = X,
( = E f (log((D)
ADKI |1 [[f]].
To investigate this relationship between the Dirichlet KI. divergence and the number of classes of the classification, we conducted a toy experiment input dimensions D and increasing number of classes K. The classes were created by using each dimension as the decision boundary, i.c. zz= and leveraging all permutations to create up to K 22 classes, where D was set 8. The training data, index set for fKL, and test data, consisting of 1000 data points each, were all sampled uniformly at random from the hypercube. We used MAP model with 2-layer MLP architecture with 25 hidden units each, bias terms enabled, and ReLU activation functions. In terms of optimization, we trained for 3000 epochs using Adam optimizer with a learning rate 0.005 and default parameters otherwise.
hypercube [-1, 1] with fixed number of
Figure 11 illustrates the model's test log- ligelike and fKL after training. The regular MAP model represents a decent baseline with linear decrease in performance as the label dimensionality increases. In contrast, the test log-likeliin ligeliiel of the MAP IVI model without KL scaling decreases exponentially while the fKL increases approximately linearly. Applying the above proposed scaling to the fKL during training aids the optimization, keeps the final fKL at convergence consistently low and significantly improves the model's 's test logllikelihood.
200
IKL at convergoan
Test LLH
150
-2
100
-4-
MAP
50
MAP IVI w/ IKL scaling
-6
MAP IVI w/ fKL scaling
2
2
28
Label Dimensionality
Figure 11: A toy classification example on a hypercube to illustrate the scaling issues associated with the Dirichlet density. The vanilla MAP performance acts as a baseline, and demonstrates the typical range of log likelihood values for this task across increasing label dimensionality. For fVI, the Dirichlet fKI. reports a significantly larger range that is x x100 the log likelihood range. This value imbalance affects the iELBO objective, resulting in significant underfitting. Applying heuristic to scale the iKL term, keeping the fKL invariant across label dimensionality avoids the underfitting phemonema. Plot reports mean and 2 standard deviations over 10 seeds.
30

---

Under review as a conference paper at ICLR 2022
H.3 CHANGING PRIOR DIRICHLET PARAMETERS
The uniform Dirichlet distribution with concentration parameters 85. - is a natural choice for an uninformed prior over the simplex. However, potentially interesting cases consider are priors where all 3g are set to another value which is greater or smaller than 1. While the Dirichlet mean remains the same, 3k
corresponds greater confidence that the class probabilities are uniformly distributed and BR < prefers dominance of any particular class. It was also hypothesized that scaling 3g could yield results comparable scaling the fKI. as discussed in the previous subsection. To test this hypothesis, we repeated the hypercube experiment from the previous subsection with the MAP model without fKL scaling while using different By as prior parameters.
Figure 12 show the test log-likelihood of the MAP IVI model without fKL scaling after training with Dirichlet priors using varying prior concentration parameters 3g. However, there are no significant differences when using different 33 and particular 33 achieves test log- likelihhoods which would be comparable to the improvements due to fKL scaling discussed the previous subsection.
Log- Likeliiooood
MAP 3g
-2
3e 5
-4
3g 1
-6
3 5
3u 10
27 Label Dimensionality
2
Figure 12: Insert caption.
To further investigate the Dirichlet prior with different concentration parameters, we repeated the visualizable Two Moons toy problem using the MC Dropout IVI model with different 8g. Figure 13 depicts the predicted class probabilities of the toy problem with 2 classes. For Bg < 1, the areas of confident prediction enlarge but quickly fall back uniformity, whereas for B& 1, the confident predictions or more locally concentrated, tapering slowly towards uniformity and 'leaking' into the unobserved area.
3g = .1
3 =
Bg
3x = 5
Bx = 10
Figure 13: Reproduction of the Two Moons toy problem (Figure 3) with varying uniform prior precision.
31


Reviews:
[{"review_id": "hPS_51jADpJ", "reviewer": "Reviewer_7nPR", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "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": "Pros:\n\n1. The paper is rigorously written, scholarship and attribution of related work is generally good, and the mathematical definitions for all used concepts and background material are abundantly available in the main paper and the appendix. This is quality material, I congratulate the authors for doing the work.\n\n2. The key idea ,while very close to the space of papers described in appendix Section A, is a useful and interesting extension to try to utilize Bayesian estimates of the forward passes to obtain the Dirichlet parameter. While this is a narrow extension and does not introduce a better model, it is technically well justified and distinct as an inference approach.\n\n3. In most cases the approach also seems to be empirically beneficial.\n\n4. The authors have an unusual approach to implicitly estimate the quantity of interest using a nested Newton step, I found this quite interesting and appreciate the link to Tom Minka's old note on the Dirichlet.\n\nCons:\n1. Primarily I want to note that the paper did not need and does not really benefit from the implicit processes connection. The key trick might actually be easier to convey and explain without that distraction in order to present the paper better. It would have been sufficient to mention in the related work section that this can be seen as an instance of this framework. Occasionally clarity suffers when trying to blend too many ideas into one paper. This paper here basically needs fVI and the rest can be derived and explained fine, Definition 3.1 does not seem to carry its weight in terms of prominence to help the reader understand the material.\n\n2. In the otherwise well-designed experiments, I would like to see a clear demonstration of the performance of the models described in appendix Section A. It is valuable and interesting that these models have been explained as special cases of this framework here, but I would still like to see which performance can be achieved with them vis-a-vis with the fukll Bayesian estimator propsoed here to get a full idea of rthe progress this framework offers. As the evaluation currently stands, I understand that this works compared to weight space VI, but not how it works compared to the methods in Sec. A with various Bayesian models. As such, it is hard to evaluate how much benefit the novelty of this paper offers compared to the existing evidential learning-style models if the key trick with the inner Newton step to estimate the Dirichlet parameter were not applied, but for instance the models still were used in a Bayesian way.\n\n3.\nIn the experiments, the authors repeatedly point out that in CIFAR 100 weight based models do better due to the strong regularisation effect of the Dirichlet prior used here when having this many classes. I appreciate that insight and hypothesis, but could we please test it? How would the performance in this case differ if the prior were to be set to be weaker (but still uniform), i.e. 1/lambda in order to be less informative. According to the authors' arguments, this should improve performance. Could we test that?\n\n4. The authors mention that weight priors cannot obtain estimates of OOD such as the ones shown in Figure 3 middle and right. I would argue that in Hierarchical Gaussian Process Priors For Neural Network Weights by Karaletsos & Bui in Neurips 2020 explicit weight priors were used with similar results, with the key insight that these priors were data-dependent and the sampled weights as such local variables to each datapoint. There is an interesting relationship to this property here, as the Dirichlet estimate also is local to each datapoint. So maybe a softer statement could be that global weight priors may struggle with such generalizations as opposed to local priors? \n\n5. The authors have a mathematical hack in their method which they acknowledge: they disconnect the gradient when separately performing their inner optimization for the Dirichlet parameter and as such do not directly optimize the joint objective. I would be curious to see how the method performs if that gradient were to be carried forward using bi-level optimization (here's a review https://arxiv.org/abs/2101.11517#), as in the current inference we probably do not see the method evaluated fully. I appreciate that the authors claim it still works, but it would be instructive to get a better sense of this.\n\n", "summary_of_the_paper": "The authors propose Function space variational inference for deep Bayesian classification, an inference framework which extends ideas around evidential learning for classification by applying function space Variational Inference and marginalizing over latent functions.\nConcretely, this leads to a method where the estimated Dirichlet is not directly regressed from a forward pass of a function, but is estimated as a parameter that would match the empirical distribution over sampled logits given multiple samples from the function posterior (using a nested optimization step), and hence can be plugged into the fVI framework.\nThe authors show this is a flexible way to model uncertainty for classification and can be combined with diverse models (ensembles, MCDropout and various Bayesian models).\nThe method produces good results.", "summary_of_the_review": "I like the premise of this paper and it is clear that this is a work of depth and quality. However, I feel there are some details missing that I would hope can be fixed to make the work clearer and to address some of its weaknesses. \nIf the authors can address the points I made about the paper (probably the first is hard to do now given the narrative of the manuscript, but the rest may be feasible) and the method is clarified to have clearly identifiable empirical benefits over the baselines in Sec. A, I would be happy to revise my opinion positively and strongly support publication.\n\nUpdate post-review:\nI thank the authors for their measured responses. My sense is this work could still benefit from some more iterations, such as the disconnected gradient and the somewhat dissatisfactory empirical comparisons. As such, while I still like the paper, I cannot clearly recommend acceptance because I would have these questions open as it stands. I suggest working on fixing these aspects to have a more complete version of this manuscript and hope to see it again as a reviewer when that is the case so I can strongly support it.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "WwlQ3pgdv8y", "reviewer": "Reviewer_why2", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "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": "## Strengths\n\n- **Significance.** This is a timely contribution that seeks to address significant limitations of Bayesian deep learning, namely those concerning the calibration of uncertainty and robustness to adversarial attacks.\n- **Clarity.** The paper is generally well-written and mostly self-contained. I found it clear and easy to follow. I particularly appreciated the treatment of the prior works related to this paper, which I found to be extensive while also being succinct and clear.\n- **Reproducibility.** It's excellent to see the depth of low-level implementation details that are provided in the appendix. This should greatly aid subsequent works in reproducing the results reported in this paper. It was also encouraging to see that the code is included in the supplementary material and that it seems to be clean and well-structured.\n\n## Weaknesses\n\n- **Novelty.** Being a combination of implicit variational processes, functional variational inference, and basic concepts from probabilistic multi-class classification,  the novelty of the contribution is somewhat limited. I find the claim that this work provides \"a unifying view of prior work which use the Dirichlet distribution and function-space regularization\" to be a slight over-statement. A method that combines concepts from two areas is not necessary a unifying theory thereof. One significant technical challenge that is tackled is the MC estimation and minimization of the KL divergence, which is addressed using a quasi-Newton method.\n\n## Questions\n\n- Section 4.4 considers image classification at five increasing levels of corruption. It is understood that each of these numbers corresponds to a particular setting of brightness, contrast, saturation, etc. However, what these settings are remains unclear. Is this detailed anywhere?\n- Do you actually make use of the ability to incorporate class-biased predictive prior that the proposed approach unlocks in any experiments other than the qualitative results shown in Figure 3? More generally, are there concrete problems in which having strong informative predictive prior for multiclass classification is actually helpful, and where conventional Bayesian approaches, through posterior inference with its data-fit component are unable to adapt to the class-biases. It would be compelling to see quantitative results that support the claims concerning class-biased predictive priors.\n- I am not fully convinced about the results shown in Section 4.4, in particular in Figure 5 for CIFAR100 where the accuracy and log-likelihoods are not distinctly better across corruption levels, particularly with Ensembles. Furthermore, the ECE is only lower for the highest levels of corruption. The explanation offered is that the method underfits, which is supported by the relatively low log-likelihood and accuracy. Obviously one can easily attain good ECE with OOD data by deliberately underfitting, which is exactly why the log-likelihood and accuracy are important to consider as well. Therefore, I don't understand why the following comment from the authors isn't actually done? \"This suggests the function-space prior should be fine-tuned (i.e. empirical Bayes) if superior ECE is desired in-distribution.\" When would this not be desired?\n\n## Miscellaneous Issues\n\n- Pg. 2: \"which use Gaussian weight priors and posteriors\" → \"that uses [...]\"\n- Pg. 4: \"evidence lower bound (ELBO)\" - the ELBO and, more generally, variational inference, long predates Hoffman et al. 2013.\n- Pg. 4: \"Further details [...], which we consider, [...]\" - which → that; extraneous commas\n- Pg. 8: \"LLH\" this abbreviation was never defined.\n- The abbreviation \"OOD\" has even been explicitly defined. Likewise, the abbreviation \"ECE\" has never been defined, nor is any explanation given as to what it is. Yet it is one of the pivotal metrics considered in this line of work.", "summary_of_the_paper": "This paper considers a functional variational inference approach to Bayesian deep learning for multi-class classification.  This paradigm enables the use of a Dirichlet predictive prior in function-space, which would otherwise be cumbersome under the classical weight-space treatment. The authors approximate the posterior in function space through the broad family of variational implicit processes, which are able to subsume popular BNN paradigms such as deep ensembles and Monte Carlo dropout. Experiments are provided to support the claim that this combination of function-space regularization with a Dirichlet predictive prior improves adversarial robustness and uncertainty quantification over its weight-space counterparts.", "summary_of_the_review": "Given the technical quality of the paper, I am leaning toward recommending acceptance of the paper. However, given some concerns around novelty and empirical evaluations detailed in my review, I am disinclined to recommend a clear acceptance.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "SP5Rubv72jp", "reviewer": "Reviewer_rF5E", "summary": "", "questions": "", "limitations": "", "rating": 3, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "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:\n\nSee \"Summary Of The Paper\".\n\n---\n\nReasons for score:\n\nThe idea is interesting and the paper writing is clear. However, my major concern is about the theoretical soundness of the proposed method. See my detailed comments below.\n\n---\n\nPros:\n\n- The paper is well-organized and easy to follow.\n- The discussion on related work is comprehensive and systematic.\n- Substantial experimental results are provided on several different important tasks (especially in the Appendix).\n\n---\n\nConcerns:\n\nMy main concern is about the theoretical soundness of the proposed method. In Sec. 3.2 Eq. (4), the author mentioned that _a more tractable approximation replace the supermum with an expectation_, $D_{KL}[q(f|\\theta)\\|p(f)]\\approx E_{s\\sim S} D_{KL}[q(f_s|\\theta)\\|p(f_s)]$. However, in the original fVI method (Sun et al., 2019), the sampling-based method samples **measurement sets**, i.e., $D_{KL}[q(f|\\theta)\\|p(f)]\\approx E_{S\\sim c} D_{KL}[q(f_S|\\theta)\\|p(f_S)]$, rather than **data points in the measurement set**. Note that there is a crucial difference: we will be estimating the KL divergence between the *joint* distributions induced by the variational and the prior if we sample measurement sets. In Eq. (4), however, the dependence among predictions (in both prior and variational distribution) at different input $x$ is completely ignored. This perhaps leads to meaningless results since dependence is the core of functions. Example:\n\nSuppose that in a 10-class classification problem, we have 1000 data points located nearly at $x_A$, and the class labels are evenly distributed (that is, 100 data with label 1, 100 data with label 2,..., 100 data with label 10). Assume we also have 10 data points located at $x_B$ (1 data with label 1, 1 data with label 2,..., 1 data with label 10). Intuitively the ideal posterior predictive Dirichlet mean at $x_A$ and $x_B$ should be $(1/10, \\cdots, 1/10)$, but the predictive Dirichlet precision should be different since we see more data at $x_A$. However, according to Eq. (1) and Eq. (6), the number of likelihood terms and the KL term are always equal at both $x_A$ and $x_B$, leading to equal posterior predictive Dirichlet distributions.\n\n\nMy second concern is about Sec. 3.3. I'm not sure whether it makes sense to estimate the distributions of $f^{(m)}_x$ with Dirichlet distribution. In the last line of Page 3, the author mentioned that *sampling the weights and evaluating $g(x,\\cdot)$ is equivalent to sampling the Dirichlet predictive*. Is there any intuition that predictive distribution induced by the variational implicit process is approximately Dirichlet?\n\n\n\n", "summary_of_the_paper": "This paper proposes a function-space variational inference method for classification tasks. The method is based on the recently proposed functional variational inference (fVI), which directly approximates the posterior over functions instead of the model parameters. The authors use a variational implicit process as a parameterized posterior and adopt Dirichlet predictive priors for the categorical predictive distributions. Since the predictive posterior is defined implicitly, the authors further propose an iterative approach to estimate the Dirichlet predictive posterior, and end up with a wake-sleep style inference procedure. Finally, the proposed method is verified on several standard classification tasks.", "summary_of_the_review": "In summary, I raise several concerns about the theoretical soundness of the proposed method, which currently prevents me from rating this paper higher. I will consider raising my evaluation if the authors address my concerns or point out my misunderstandings.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "2: Several of the paper’s claims are incorrect or not well-supported.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "3JqcKw_4-0A", "reviewer": "Reviewer_FSJj", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "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": "Firstly, this work is relevant for ICLR: designing function-space priors and inference for Bayesian neural networks is a hot topic right now. I also thought that this paper was written well, with good motivation and explanations throughout. \n\nI like the overall idea of taking multiple outputs from a Bayesian neural network, and instead of integrating them out to get a categorical distribution, using the samples to estimate a Dirichlet distribution (and using Minka's algorithm to estimate the parameters of the Dirichlet). \n\nThis work is closely related to prior works such as Joo et al., 2020 and others. Appendix A summarises these works and I thought nicely summarises the differences between this method and these closest prior works. However, although theoretically there are differences between these past works and the overall framework of fVI that this paper introduces, it does not seem to me like there are significant differences in practice on large datasets. More concretely:\n\n1. The authors say that their framework allows for index sets that are not just the training data. However, in all but one of their experiments, only the training data is used. It is nice to see the rotated MNIST example with a larger index set, however, this is quite a limited experiment (as it is only MNIST, and only tested in a very specific OOD regime). I would have liked to see more experiments where the authors used data augmentation or unlabelled data as part of the index set. Or, used some sampling similar to Sun et al., 2019, where they sample from some interval to get data not in the training set. By not having such experiments, the paper becomes similar to past work. \n\n2. The other difference is that fVI allows for Bayesian neural networks with M>1 samples in order to estimate the Dirichlet distribution. I like this idea. However, in practice, I was disappointed that M>1 only for ensembles and Rank1 models, and not also for MC dropout and Radial BNNs. \n\nOther comments / concerns: \n\n3. Why were z_min and z_max chosen as K and N respectively (Appendix F)? And why was z estimated to be N when M=1 (why not estimate z using Minka's algorithm)? From what I understand, as N becomes very large, the distribution would approach the same categorical that people use. N is quite large for many datasets already, so effectively this procedure is adding a very small amount of noise around what would otherwise be the categorical. I understand this reduces training cost. But it would be nice to see results for eg Radial BNNs with larger M, and also see how the estimated z compares with z=N (or z=K). (Appendix H is very short and I would like to see this discussed in more detail, including more in the main text.)\n\n4. fVI versions of algorithms are almost always doing worse in terms of loglik and ECE on the training data. I had hoped for better 'uncertainty' performance on the in-distribution set as well. This indicates to me that the fVI version still over-estimates in-distribution data (eg MAP vs MAP fVI). The real benefit is only coming when the test distribution is not the same as the training distribution. Do the authors have some insight as to why it is not helping in-distribution? In Section 4.4 the authors say, \"due to the predictive prior, the ECE was significantly higher for corruption 0\". I did not understand this argument, can the authors please elaborate?\n\n5. Did the authors try different prior parameters (rather than beta=1)? It seems like this might affect results considerably, and I would be interested to see this more (I liked the third and fourth columns of Figure 3). For example, on the toy dataset (Figure 3), what happens when using beta = [5,5] or even something like beta = [0.8,0.8]. In fact, in general, I think the ideal solution for Figure 3 for regions outside the input-data-space, would be to *slightly* predict the closest class (eg the top left corner should unconfidently predict red, and the bottom right should unconfidently predict blue). Perhaps different beta would allow for this.\n\n6. At the end of Section 3, it would be good to add a sentence moving some of Appendix B into the main text, by mentioning the authors down-weight the KL term (and that this might be justified by referencing Wenzel et al., 2020 and Aitchison, 2021). \n\nVery minor points:\n\n7. Why is label smoothing applied (Appendix F)? Is this necessary for some reason, eg for numerical stability? Why?\n\n8. Some more work the authors could refer to: (i) Burt et al., 2021, \"Understanding VI in function-space\", (ii) recent SG-MCMC methods for (weight-space) Bayesian neural networks, (iii) recent natural-gradient variational inference methods for (weight-space) Bayesian neural networks.", "summary_of_the_paper": "The paper introduces a function-space variational inference algorithm (\"fVI\") that uses a Dirichlet predictive prior, and approximates the output of a (weight-space-Bayesian) neural network as a Dirichlet distribution. Inference is achieved by viewing the network as a variational impolicit process, and optimising the functional ELBO from Sun et al., 2019. This view allows them to impose priors in function-space to many existing algorithms, such as MAP, MC Dropout, ensembles, Radial BNNs, and so on. The authors show results on a toy problem, and focus on out-of-distribution performance on larger problems such as MNIST, CIFAR10 and CIFAR100, as well as performance under a fast gradient sign method adversarial attack. The fVI variants almost always perform better when the test distribution is not the same as the training distribution. ", "summary_of_the_review": "Overall, I like the approach of the paper, and think there is potential. Therefore, although I am borderline, I am leaning towards accept. My main review outlines my biggest issues (especially points 1-4), these tend to be about the practical implementation not aligning with what I view as the biggest benefits of the theory (and the biggest differences to past works), as well as questions about experimental results.\n\n------------\nPost-rebuttal: Please see discussion. I am sticking with my score for now, despite some of my initial concerns remaining (and additional valide points raised by Reviewer rF5E). ", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}]

Author rebuttal:

Meta-review: A variational function-space prior is proposed, resulting in a variational Dirichlet posterior. After rebuttal, reviewers still had many remaining questions or concerns about the paper. For instance, rF5E outlines several concerns, many relating to factorization assumptions. Reviewer 7nPR also provides several suggestions. I will not repeat them here, but do encourage the authors to look closely at these questions and suggestions. At this particular time the paper is not strongly resonating with reviewers, but could be updated so that the value of the contributions is more obvious.
Decision: Reject