Year: 2021

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 are hard to specify. Instead of weight-space priors, we leverage function-space variational inference to apply a Dirichlet predictive prior in function space, resulting in an implicit variational Dirichlet posterior which facilitates easier specification of predictive 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 smalland 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 architecture or model size.

heading: 1 INTRODUCTION
text: Deep learning (Goodfellow et 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 et 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 inference1 (Sun et al., 2019) (fVI) for classification. 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 a 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 is Dirichlet distributed, rather than solely categorical, we gain a richer predictive uncertainty representation while retaining the same mean predictions as 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; Joo et al., 2020; Sensoy et al., 2018; 2021), through a principled objective that does not require a 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.
1While prior work uses the term functional VI, we use function-space VI for clarity.

heading: 1.1 CONTRIBUTIONS
text: 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 a function-space variational inference method for classification, using an effec-
tive KL divergence estimator from softmax samples via Newton-based Dirichlet estimation.

heading: 2 RELATED WORK
text: We build on prior work that considers classification in the Bayesian setting and function-space inference of Bayesian neural networks.

heading: 2.1 BAYESIAN CLASSIFICATION
text: Compared to regression, classification is non-trivial for Bayesian methods due to 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-Lobato & 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 a Dirichlet predictive using a log-normal approximation (Milios et al., 2018). Classification with Bayesian neural networks is possible through a wide range of approximate inference methods, including Markov chain Monte Carlo (Neal, 1995; Zhang et al., 2020), (mean-field) variational inference (MFVI) (Graves, 2011; Blundell et al., 2015; Zhang et al., 2018), Laplace approximations (MacKay, 1992; Denker & LeCun, 1991; Ritter et al., 2018; Khan et al., 2019; Immer et al., 2021; Daxberger et al., 2021), ensembles (Lakshminarayanan et al., 2017; Osband et al., 2018; Barber & Bishop, 1998; Pearce et al., 2020), expectation propagation (Hernández-Lobato & Adams, 2015) and Monte Carlo dropout (Gal & Ghahramani, 2016; Kingma et 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 1 covariance parameterizations, Rank 1 BNNs have a 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 Gaussian-predictive BNN can then be trained using this latent approximation.
Alternative 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 to distill a
trained ensemble into a single model (Malinin et 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 KL regularization against a function-space prior (Sensoy et al., 2018; 2021). Compared to these methods, we use fVI as a 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.

heading: 2.2 FUNCTION-SPACE VARIATIONAL INFERENCE
text: 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 (fBNN) by Sun et al. (2019) uses the fKL to use explicit or implicit stochastic processes as function-space priors. They use a GP trained on the data as a prior, which can be viewed as a 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 fVI, as they have an attractive closed-form Gaussian predictive (Watson et al., 2021a;b). Concurrent work by Rudner et al. (2021) has also adopted fVI for classification, by linearizing the network about a Gaussian weight distribution to estimate the fKL. This model works with a Gaussian (latent) predictive prior and posterior which loses the intuitive aspect of function-space priors. Moreover, the linearization 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-space posterior through the gradient flow of the log posterior. Function-space inference is also an attractive approach to continual learning, e.g. Pan et al. (2020).

heading: 3 DIRICHLET PREDICTIVE PRIORS FOR BAYESIAN CLASSIFICATION
text: In this section, we discuss how function-space variational inference can be applied to classification using variational implicit processes (VIP) (Ma et al., 2019).

heading: 3.1 VARIATIONAL IMPLICIT PROCESSES FOR CLASSIFICATION
text: LetD ⊂ (X ×Y) = {(xn,yn)}Nn=1 be the training data consisting of N observed pairs of input data xn and corresponding K-dimensional, one-hot class label vectors yn. Deep learning approaches to classification typically assume that a class label y follows a 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 x ∈ X to an element fx ∈ ∆K−1, where ∆K−1 denotes the K−1 simplex and elements from ∆K−1 are K-dimensional probability vectors. More precisely, we write fx = g(x,w) = σ (φ (x;w)) and y ∼ Cat (·|fx). The softmax function σ ensures that fx is a valid probability vector, i.e. that fxk ≥ 0 and ∑K k=1 fxk = 1. In conventional ML training, the
weightsw are optimized by maximizing log p(D|f) = log ∏ D Cat (y|fx) = ∑ D ∑K k=1 yk log fxk,
where ∏ D and ∑ D denote ∏ (x,y)∈D and ∑ (x,y)∈D respectively, and φ, σ, andw 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 (1)x , . . . ,f (M) x . To capture this uncertainty, we study these implicit predictions and instead posit they are drawn from a posterior predictive Dirichlet 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.
Sampling the weights and evaluating g(x, ·) is equivalent to sampling the Dirichlet predictive, while varying x with fixed weight w corresponds to evaluating a classifier g(·,w). Together, the joint variation ofw and x defines an implicit stochastic process which we assume describes an implicit K-dimensional Dirichlet distribution for each x ∈ X . When parameterized by a 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) is a set of random variables f , such that any finite subset fx1:L = {fx1 ,fx2 , ...,fxL} with L ∈ N has a joint distribution which is implicitly defined by the generative process,
w ∼ p(w), fxl = g(xl,w), ∀xl ∈ X , 1 ≤ l ≤ L.
Applying Definition 3.1 to stochastic network classifiers requires two key considerations. Firstly, the weight distribution p(w) takes a parameterized form, qθ(w), with parameters θ we wish to optimize. The function g(·, ·) represents a specific network architecture, such as a ResNet. Further details regarding the variations of qθ(w) that we consider can be found in Section D in the Appendix. Secondly, in the classification setting we are now reasoning over a joint distribution of Dirichlet distributions on a fixed K−1 simplex. Note this is distinct from the Dirichlet process (Teh, 2010), which is a trans-dimensional prior over discrete distributions, such that K is inferred. Informally, in our setting, a finite collection of L random variables would be an element from the K−1 simplex raised to the power of L. The stochastic process can be defined as a random variable from the K−1 simplex to the power of the data domain X . This is not a distribution we can reason about easily, in contrast to the Gaussian process where the joint and marginal distributions are all Gaussian (Rasmussen & Williams, 2005). However, as 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 q(f |θ), allowing us to adopt high-performance BNN representations while avoiding distributions over infinite parameter spaces, as in e.g. Bayesian nonparametrics (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 a 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-space inference to optimize qθ(w).

heading: 3.2 FUNCTION-SPACE VARIATIONAL INFERENCE AND INDEX SETS
text: To optimize our variational implicit processes, we can use the function-space evidence lower bound objective (fELBO) (Sun et al., 2019),
L(θ) = E f∼q(·|θ) [log p(D|f)]−DKL[q(f |θ) || 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 a prior over model parameters p(w). The expected log-likelihood term in the fELBO is unchanged from the regular ELBO optimization for BNNs, where we use M sample estimates using the categorical likelihood,
Ef∼q(·|θ) [log p(D|f)] ≈ 1M ∑N n=1 ∑M m=1 log p ( yn,f (m) xn ) . (2)
Since the fELBO involves distributions over functions, fVI requires computing a KL divergence between stochastic processes. Sun et al. (2019) derived this divergence as the supremum over all finite measurement setsX∈XL of index sets XL of size L,
DKL[q(f |θ) || p(f)] = sup L∈N,X∈XL DKL[q(fX |θ) || 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 et al., 2021) replaces the supremum with an expectation,
DKL[q(f |θ) || p(f)] ≈ ES∼S DKL[q(fS |θ) || p(fS)] (4)
where S ⊂ XL is a finite set of (now) fixed size L, which we refer to as the ‘index set’. Note that Sun et al. (2019) refer to S as a distribution c. 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 an adversarial approach that optimized a 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 S 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 it 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.

heading: 3.3 POSTERIOR DIRICHLET AND KL DIVERGENCE ESTIMATION
text: In contrast to prior work (Joo et al., 2020; Sensoy et al., 2018; Malinin & Gales, 2018), we do not directly predict α using a neural network, but instead assume f (m)x ∼ Dir(·|αx). This allows us to compute αx through a maximum likelihood estimate (MLE) of M samples f (m) x from the VIP. Unfortunately, the Dirichlet MLE cannot be computed in closed form and requires fixed-point iteration or other optimization techniques to be computed, increasing the computational burden (Minka, 2000). Consequently, we replace the full MLE with a computationally cheaper alternative. To this end, we consider αx in terms of two separate but dependent parameters: the Dirichlet mean ᾱx = αx/zx and the Dirichlet precision zx. By matching the first moment of the empirical distribution of fx, we obtain an estimate of ᾱx ≈ 1M ∑M m=1 f (m) x . We then fix ᾱx and estimate the precision zx using a fast, iterative, quasi-Newton algorithm (Minka, 2000) using f (1:M)x = {f (1)x , . . . ,f (M)x },( z(t+1)x )−1 = ( z(t)x )−1 + ( z(t)x )−2( d2 dzx2 LDir ( f (1:M)x ,α (t) x ))−1 d dzx LDir ( α(t)x ) , (5)
where z(t)x andα (t) x = ᾱx/z (t) x are the Dirichlet precision and concentration at iteration t of the quasiNewton algorithm and LDir ( f (1:M) x ,αx ) is the Dirichlet log-likelihood log ∏M m=1 Dir ( f (m) x |αx ) . This corresponds to breaking the dependence between ᾱx and zx that would be considered in the full MLE case in order to accelerate the resulting approximation. With ᾱx and zx estimated, αx and thus our variational posterior q(fx|θ) = Dir (fx|αx) is now available. Moreover, for tractability we choose a factorized prior and assume a factorized posterior, i.e. q(fS)= ∏L l=1 q(fsl),
2 so that the fKL term from Equation 3 can be estimated using M samples with the Dirichlet KL divergence
DKL[q(fS |θ) || p(fS)]≈ 1M ∑L l=1 ∑M m=1 ( log q ( f (m) sl ∣∣∣θ)− log p(f (m)sl )) , (6) where f (m)s is the m-th implicit function evaluated at s ∈ S and log q(f (m)s |θ) and log p(f (m)s ) are the log-likelihood of f (m)s under the variational Dirichlet posterior and Dirichlet prior respectively. Note, in this work our predictive prior p(fx) = Dir(· | β), where each βk is chosen to be 1 unless
2While this is not a rich prior, we are assuming our VIP, e.g. CNNs, will capture meaningful correlations.
otherwise stated. During optimization, we treat q(fx|θ) as constant with respect to θ due to the αx 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 θ on Dir(fx|αx) on the KL in the gradient update, learning is not hindered due to the iterative nature of SGD training in practice. From the perspective of the wake-sleep inference procedure adopted by Ma et al. (2019), optimizing the fELBO represents the wake phase and the Dirichlet estimation represents the sleep phase. Nonetheless, we optimize θ 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. MFVI) the reparameterization trick is required. In the case where only a single sample is available (M = 1), we approximate the precision zx with the size of the training data. Further practical considerations, mini-batching and prior regularization, are discussed in Section B in the Appendix. Both considerations manifest as a scaling of the KL term.

heading: 4 EXPERIMENTS
text: In this section, we present an empirical evaluation of our proposed fVI 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 et al., 2010), CIFAR10 and CIFAR100 (Krizhevsky, 2009) data sets. 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 error (ECE) (Naeini et al., 2015), which estimates the calibration of accuracy vs confidence through binning the predicted class probabilities. Additional details are provided in Appendix F.

heading: 4.1 TOY PROBLEM
text: To visualize the effects of function-space variational inference, we conducted a 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 our inference approach combined with a uniform Dirichlet prior adequately increases predictive uncertainty outside of the observed data, whereas conventional weight-space training leads to overconfident extrapolation. While the weight-
space approach learns a 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 process-based classification, e.g. Liu et al. (2020). Furthermore, fVI facilitates specifying class-biased priors, which would be difficult to achieve with weight priors.

heading: 4.2 ROTATED MNIST
text: 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 to 180◦, which simulates a challenging OOD scenario due to the absence of data augmentation. For this experiment, we used the same model archetypes and MLP architecture as for the toy problem in Section 4.1.
The log-likelihood between models is shown in Figure 4 (left), accuracy and expected calibration error are reported in Figure 9 in the Appendix. In terms of classification error rate, both approaches yield the same performance. In terms of log-likelihood, fVI consistently outperforms their weightspace counterparts as the data becomes more OOD, indicating better OOD uncertainty quantification. Subnetwork linearized Laplace (Daxberger et al., 2021) is also reported as a competitive baseline, achieving better accuracy and in-distribution log-likelihood. However, these results were obtained using a ResNet-18 architecture rather than an MLP.

heading: 4.3 ASSESSING INDEX SET DESIGN
text: To illustrate the importance of the index set, we train the fVI 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 fVI already outperforms the weight-space counterparts, a direct comparison in Figure 4 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 fELBO. 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, we use the training data as the index set, demonstrating performance without this inductive bias.

heading: 4.4 LARGER-SCALE IMAGE CLASSIFICATION UNDER CORRUPTION
text: We used the regular train splits of the CIFAR10 and CIFAR100 (Krizhevsky, 2009) image classification data sets as training data and their corrupted versions (Hendrycks & Dietterich, 2019) as OOD test data. CIFAR10 and CIFAR100 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,
contrast and saturation, or adding noise, blur or 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 et al., 2016). In addition to the previously used MAP, MC Dropout and Ensemble models, we also evaluate our fVI 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 5 reports the performance metrics for CIFAR10 and CIFAR100 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 0, but this underfitting provided superior ECE when OOD. This suggests the function-space prior should be fine-tuned (i.e. empirical Bayes) if superior ECE is desired in-distribution. For CIFAR100, 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-space prior.

heading: 4.5 LARGER-SCALE IMAGE CLASSIFICATION UNDER ADVERSARIAL ATTACKS
text: 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 fVI compared to standard weight-prior approaches on the CIFAR10 and CIFAR100 datasets using the fast gradient sign method (FGSM) (Goodfellow et al., 2014). Figure 6 compares the accuracy and the log-likelihood of the test data with increasing amounts of perturbation, ranging from = 0 (no attack) to = 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-likelihood, whereas the weight-space LLH performance drops significantly. We also observe the accuracy vs
robustness trade-off in the fVI models discussed in Section 4.4. We attribute this behavior to the quality of the uncertainty quantification at 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 fVI. This richness is illustrated in the toy example in Figure 3.

heading: 5 CONCLUSION
text: We 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 a unified formalism of prior work on Dirichlet-based classifiers with function-space regularization, and can be applied to a 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, e.g. through using data augmentation or unlabeled data.

heading: CIFAR10 CIFAR100

heading: ETHICS STATEMENT
text: 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.

heading: REPRODUCIBILITY STATEMENT
text: Code for reproducing experiment results is provided in the supplementary material. Further, Appendices F, G and H contain additional details and results of the experiments.

heading: A RELATION TO PRIOR WORK
text: 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, we first explain prior work using their own notation and then discuss the correspondence to our fVI perspective.

heading: A.1 PRIOR NETWORKS
text: Malinin & Gales (2018; 2019) use a neural network with parameters θ to directly predict the concentration parameters αc of a Dirichlet distribution p(µ|x;θ), given input x, which is distinct from our approach of estimating a posterior Dirichlet from M categorical predictions. This model is not a Bayesian network in practice, as only point estimates for the weights are learned. To ensure αc > 0, an element-wise exponential operation is applied as the final layer of the neural network. Additionally, prior networks minimize an optimization objective consisting of two separate KL divergences, representing in- and out-of-distribution data respectively,
L(θ) = Epin(x)[DKL[Dir(µ|α̂) || p(µ|x;θ)]] + Epout(x)[DKL[Dir(µ|α̃) || p(µ|x;θ)]]. (7) The first expectation Epin accounts for the actual learning, i.e. fitting the training data, whereas the second expectation Epout is supposed to regularize the model by matching a prior distribution. Accordingly, the first expectation is computed for the training data and can be compared to the expected log-likelihood term in our fVI approach. Instead of maximizing the categorical loglikelihood of M observations, prior networks construct Dirichlet targets by smoothing categorical ground truth labels to define the Dirichlet mean and setting the precision as a 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 fVI approach, where the OOD data is 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 KL divergence from the prior distribution (right) to the variational posterior (left) is considered.

heading: A.2 BELIEF MATCHING
text: Joo et al. (2020) assumes a 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 qz|x by replacing the final softmax layer with an element-wise exponential operation. To learn the target posterior, belief matching maximizes
lEB(y, α W (x)) = Eqz|x [log zy]−DKL[q W z|x || pz|x], (8)
where Eqz|x [log zy] is the expected log-likelihood of the training data and DKLqWz|xpz|x is the KL divergence between the predicted Dirichlet posterior and the Dirichlet prior. Therefore, their objective matches our fELBO objective equation 1 except for two differences: Firstly, belief matching computes both the expected log-likelihood and the KL divergence with respect to 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 KL divergence using the training data, which resembles our fKL where the index set is constrained to be the training data.

heading: A.3 EVIDENTIAL DEEP LEARNING
text: Sensoy et al. (2018) directly predicts the concentration parameter of a 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-II maximum likelihood by integrating over a Dirichlet prior and the
sum of squares between target labels yi and predicted probabilities pi. Furthermore, a regularizing KL divergence term is added, resulting in a total loss function,
L(Θ) = N∑ i=1  K∑ j=1 ( (yij − p̂ij)2 + p̂ij(1− p̂ij) Si + 1 ) + λtDKL[Dir(pi|α̃i) || Dir(pi|1)]  , (9) where yij are individual 0-1 target labels, p̂ij are components of the predicted Dirichlet mean, Si is the predicted Dirichlet precision, α̃i is the predicted Dirichlet concentration parameter, 1 is a vector of ones and λt is a annealing coefficient for optimization. The first part of their loss is responsible for fitting the training data and can thus be compared to the maximum likelihood objective in Section 3.1. The ML objective can be derived from the categorical log-likelihood via type-I maximum likelihood, whereas their objective is derived by minimizing the sum of squares via type-II 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 case 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 is 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.

heading: A.4 EXPERIMENTAL COMPARISON
text: In this subsection, we compare our MAP and MAP fVI 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, we used the official open-source implementations 3 4.
Figure 8 illustrates the test accuracy, log-likelihood, 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 in Section F. The Belief Matching model corresponds closely to the MAP fVI model, as 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)) or the hyperparameters used in Section F. It is uncertain whether this is due to the model, implementation bugs or unrecorded hyperparameters 5.
3github.com/tjoo512/belief-matching-framework 4github.com/KaosEngineer/PriorNetworks 5The authors did not respond to personal correspondence regarding this matter.

heading: B PRACTICAL CONSIDERATIONS
text: This section discusses practical issues with the approach in more detail.

heading: B.1 MINI-BATCH TRAINING
text: In weight-space variational inference, batch training requires an appropriate scaling term to the KL divergence of the ELBO during each batch iteration. This is to ensure the weight KL is correctly balanced against the batch expected log-likelihood. The fKL approximation is equally mini-batched, as we restrained into using a small portion of the index set each iteration. Therefore, like Sun et al. (2019), we divide both the batched expected 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 in variational inference is a topic of active debate (Wenzel et al., 2020; Aitchison, 2021) and is frequently scaled or annealed for numerical reasons (e.g. Dusenberry et al. (2020)).

heading: B.2 PRIOR SPECIFICATION FOR HIGH OUTPUT DIMENSION
text: When scaling to higher dimensional classification tasks, specificallyK ≥ 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, as 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, we adopt the strategy of Joo et al. (2020) and apply additional scaling to the KL term in the fELBO. This scaling can be shown to be numerically equivalent to a certain prior, i.e. β (Section 3.4, Joo et al. (2020)). Therefore, optimizing this scaling is a form of model selection. For our CIFAR100 experiments, we simply chose a scaling such that that the fKL magnitude was close to the CIFAR10 values. We found this to be about 0.1, which matches a 10x scaling suggested by the Dirichlet KL due to the summation terms.

heading: C DEEP STOCHASTIC CLASSIFIERS AS STOCHASTIC PROCESSES
text: Stochastic classifiers are stochastic processes instantiated as a feedforward or convolutional neural networks with stochastic weights. Formally, our stochastic process f is defined on the sample space Ω with an index set X defined by the data type (e.g. real vectors, RGB images), so f : X ×Ω→ ∆K−1, where ∆K−1 is the state space, which is the K−1 simplex. A random variable f(x) : Ω→ ∆K−1 can be defined for each x ∈ X and we write f(x) = fx. Kolmogorov’s extension theorem (Tao, 2013) guarantees the existence of a stochastic process f if for each L ∈ N the finite marginal joint distributions p1:L(fx1:L), where fx1:L = {fx1 , . . . ,fxL}, satisfy exchangeability and consistency. These finite marginal joint distributions are implicitly defined according to Definition 3.1. Exchangeability: For any permutation π of 1, . . . , L, pπ(x1:L)(fπ(x1:L)) = px1:L(fx1:L). This requires that the process behavior is invariant to the order of inputs. For a feedforward neural network, this is satisfied because the respective predictions do not change if the order of inputs changes. Consistency: For any 1 ≤ L′ < L, p1:L′(fx1:L′ ) = ∫ p1:L(fx1:L) dfxL′+1:L . This requires that future evaluations are independent of past evaluations. For a feedforward neural network, this is satisfied because predictions do not depend on previous predictions.
D VARIATIONAL IMPLICIT PROCESS PARAMETERIZATIONS
This work uses 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 1 summarizes the connection between popular BNN models (and the MAP baseline) and weight distribution qθ(w).

heading: E DIRICHLET MUTUAL INFORMATION
text: Smith & Gal (2018) proposed to use the mutual informationMI (θ;y) of model parameters θ and class label y to detect adversarial examples. Intuitively, this quantity measures how much the model could improve its parameters θ if y were observed. Therefore, it quantifies epistemic uncertainty and is also known as information gain. However, Smith & Gal (2018) note that, in general,MI (θ;y) is intractable for deep learning models. Therefore, they provide a tractable approximation which applies to models with stochastic parameters and requires sampling-based estimation.
Fortunately, in our function-space variational setting, where the model parameters θ are deterministic, the one-hot encoded class label y follows a categorical distribution y ∼ Cat (·|fx) and the latent probability vector fx follows a Dirichlet distribution fx ∼ Dir (·|α), the equivalent mutual information between class labels and model predictionsMI (fx;y) = H (y)−H (y|fx) can be derived in closed-form from the Dirichlet approximation using
H (y) = − K∑ k=1 αk z log (αk z ) , (10)
H (y|fx) = − K∑ k=1 αk z (ψ (αk + 1)− ψ (z + 1)) . (11)
The complete derivation is provided in Appendix E.1 on the next page.
The first termH (y) corresponds to the categorical entropy of the expected value of fx. It is equivalent to the standard measure of predictive entropy for softmax-based ML deep learning classifiers. The second termH (y|fx) is the conditional entropy of the class label y given a Dirichlet observation fx and it is influenced by the spread of fx, whereas the first term only depends on its mean. Although this function-space variational formulation admits closed-form solutions forMI (fx;y), the resulting values will still be approximations in our setting because the variational posterior Dir (fx|αx) was estimated from M observations. Nonetheless, it is not limited to models with stochastic weights, but instead provides a unified approach to measuring epistemic uncertainty for any classifier with a Dirichlet predictive distribution.

heading: E.1 MUTUAL INFORMATION DERIVATION
text: Let y ∼ Cat (·|fx), fx ∼ Dir (·|α), fxk ∼ Beta (·|ak, bk), where ak = αk and a′k = ak + 1, while bk = z − αk and z = ∑K k=1 αk,
MI (fx;y) = H (y)−H (y|fx) , (12)
H (y) = − K∑ k=1 Efx [fxk] log (Efx [fxk]) , (13)
= − K∑ k=1 Efxk [fxk] log ( Efxk [fxk] ) , (14)
= − K∑ k=1 αk z log (αk z ) , (15)
H (y|fx) = Efx
[ −
K∑ k=1 fxk log (fxk)
] , (16)
= − K∑ k=1 Efx [fxk log (fxk)] , (17)
= − K∑ k=1 Efxk [fxk log (fxk)] , (18)
= − K∑ k=1 ∫ 1 0 fxk log(fxk) fxk ak−1(1− fxk)bk−1 B(ak, bk) dfxk, (19)
= − K∑ k=1 ∫ 1 0 log(fxk) fxk ak(1− fxk)bk−1 B(ak, bk) dfxk, (20)
= − K∑ k=1 ∫ 1 0 log(fxk) fxk a′k−1(1− fxk)bk−1 B(ak, bk) dfxk, lim c→0 xc − 1 c = log x (21)
= − K∑ k=1
1
B(ak, bk) ∫ 1 0 ∂ ∂a′k fxk a′k−1(1− fxk)bk−1dfxk, (22)
= − K∑ k=1
1
B(ak, bk)
∂
∂a′k ∫ 1 0 fxk a′k−1(1− fxk)bk−1dfxk, (23)
= − K∑ k=1
1
B(ak, bk)
∂
∂a′k B(a′k, bk), (24)
= − K∑ k=1
1
B(ak, bk) B(a′k, bk)(ψ(a ′ k)− ψ(a′k + bk)), (25)
= − K∑ k=1
1
B(ak, bk) ak ak + bk B(ak, bk)(ψ(ak + 1)− ψ(ak + 1 + bk)), (26)
= − K∑ k=1 ak ak + bk (ψ(ak + 1)− ψ(ak + 1 + bk)), (27)
= − K∑ k=1 αk z (ψ (αk + 1)− ψ (z + 1)) . (28)
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, a GTX1070 GPU and a GTX2080 GPU with less than 300 hours of total runtime.
The Two Moons data was generated by the make_moons function from the scikit-learn 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.1. This toy experiment is the only exception in this regard.
For Rotated MNIST, we used all 60000 images of shape 28x28x1 reshaped to 784 from the train split with pixel values normalized to [−1, 1] and no other pre-processing or data augmentation. All 10000 images from the test set were used during evaluation, rotated by a 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 CIFAR10 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 a total of 960000 test images. The CIFAR10 and CIFAR100 (Krizhevsky, 2009) data is available under the terms of the MIT License and the corrupted CIFAR10 and corrupted CIFAR100 (Hendrycks & Dietterich, 2019) data is 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 σ = log(1 + exp(ρ)) and ρ was initialized to -5 while the means were initialized using the PyTorch default initialization scheme for CNNs. For Radial BNNs without fVI, we used a closedform Gaussian weight KL divergence with a Gaussian prior with a mean of 0 and a standard deviation of 0.1. For Radial fVI, we used our fKL instead of the weight-space KL.
For Rank1 BNNs (Dusenberry et al., 2020) and Rank1 fVI, we used 4 ensemble members and 250 training epochs instead of 200 due to slow convergence and scaled the learning rate by 0.1 at 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 σ = log(1 + exp(ρ)) and ρ was initialized to -3 while the means were initialized to 1. For Rank1 BNNs without fVI, following Dusenberry et al. (2020), the Rank1 priors were Gaussian with a mean of 1 and a standard deviation of 0.1, and weight decay of
0.0001 was used. We did not use KL annealing epochs. For Rank1 fVI, we 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: zmin = K and zmax = 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 z = zmax because M = 1. Similarly, we used M = 1 for the MC Dropout and Radial BNN models during training and also set z = 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 Rank1 models, we replicated the input M times during evaluation, which results in M distinct predictions. Furthermore, we applied a small amount of label smoothing fx (m) k ≈ (1 − γ)fx (m) k + γ 1 K throughout all steps of the KL divergence estimation, where γ was set to 10−4 for our experiments.
The equations from Minka’s quasi-Newton maximum likelihood Dirichlet precision estimator (Minka, 2000), which we used for our implementation, translated to our notation are
1 z = 1 z + 1 z2 ∆1 ∆2 , ᾱk = 1 M M∑ m=1 fx (m) k , ᾰk = 1 M M∑ m=1 log fx (m) k , (29)
∆1 = M ( ψ0(z)−
K∑ k=1 ᾱkψ0(zᾱk) + K∑ k=1 ᾱkᾰk
) , (30)
∆2 = M ( ψ1(z)−
K∑ k=1 ᾱ2kψ1(zᾱk)
) , (31)
where ψ0 is the digamma function and ψ1 is the trigamma function.
We initialized the algorithm with an approximate maximum likelihood solution using Stirling’s approximation to the gamma function Γ, which was also presented by Minka (2000),
z(0) = K − 1 −2 ∑K k=1 ᾱk(log ᾰk − log ᾱk) . (32)
We stopped the algorithm once the change per step is less than 10−5. Counting the number of iterations until convergence for a trained Ensemble fVI 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 1172. 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 O(1), i.e. a 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(BMK) time per mini-batch iteration. Assuming a constant number of quasi-Newton steps, the Dirichlet precision estimation takes O(SK +M) time for an index set of size S. Computing the fKL for an 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 log-likelihood and fKL calculation, resulting in an overall asymptotic time complexity of O(BMK) per mini-batch iteration. In case of a different index set, the asymptotic time complexity becomes O((B+S)MK).

heading: G EXPERIMENTAL RESULTS
text: All experimental results are reported here.

heading: G.1 ROTATED MNIST

heading: G.2 CIFAR10
text: G.3 CIFAR100

heading: H ABLATION STUDIES

heading: H.1 SAMPLES DURING TRAINING
text: In Section 3.3, a Dirichlet estimation procedure was proposed using M samples. In the single sample case M = 1, motivated by MAP models, a 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 is used in practice. This means that the 1 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 CIFAR10 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 is closer to the 1 sample MC Dropout performance than the Ensemble. This result indicates that the VIP model, rather than samples-during-training, has greater impact. The similarity in performance between 1 and 5 sample MC Dropout suggests the 1 sample approximation is reasonable for the training time savings.

heading: H.2 SCALING ISSUES WITH HIGH LABEL DIMENSIONALITY
text: The CIFAR100 experiments revealed an issue with the fELBO objective that caused underfitting for the larger label dimension. To examine why, recall that the categorical likelihood is log fxk when yk = 1. Therefore, the dimensionality of y does not directly influence the value. Conversely, the KL divergence between two Dirichlet densities DKL(p1||p2) does incorporate the label dimensionality K (Rauber et al., 2008),
DKL(p1||p2) = log Γ(z(1))− K∑ k=1 log Γ(α (1) k )− log Γ(z (2)) + K∑ k=1 log Γ(α (2) k )
+ K∑ k=1 (α (1) k −α (2) k )(Ψ(α (1) k )−Ψ(α (2) k )).
To counteract this linear increase due to the summation terms, we can assess a heuristic annealing scale factor on the fKL during training of 1/K = λ,
L(θ) = E f∼q(·|θ) [log p(D|f)]− λDKL[q(f |θ) || p(f)].
To investigate this relationship between the Dirichlet KL divergence and the number of classes of the classification, we conducted a toy experiment in a hypercube [−1, 1]D with fixed number of input dimensions D and increasing number of classes K. The classes were created by using each dimension as the decision boundary, i.e. xd = 0, and leveraging all permutations to create up to K = 2D classes, where D was set to 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 a 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 of 0.005 and default parameters otherwise.
Figure 11 illustrates the model’s test log-likelihood 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-likelihood of the MAP fVI 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 test log-likelihood.

heading: H.3 CHANGING PRIOR DIRICHLET PARAMETERS
text: The uniform Dirichlet distribution with concentration parameters βk = 1 is a natural choice for an uninformed prior over the simplex. However, potentially interesting cases to consider are priors where all βk are set to another value which is greater or smaller than 1. While the Dirichlet mean remains the same, βk > 1 corresponds to greater confidence that the class probabilities are uniformly distributed and βk < 1 prefers dominance of any particular class. It was also hypothesized that scaling βk could yield results comparable to scaling the fKL as discussed in the previous subsection. To test this hypothesis, we repeated the hypercube experiment from the previous subsection with the MAP fVI model without fKL scaling while using different βk as prior parameters.
Figure 12 show the test log-likelihood of the MAP fVI model without fKL scaling after training with Dirichlet priors using varying prior concentration parameters βk. However, there are no significant differences when using different βk and no particular βk achieves test log-likelihoods which would be comparable to the improvements due to fKL scaling discussed in the previous subsection.
To further investigate the Dirichlet prior with different concentration parameters, we repeated the visualizable Two Moons toy problem using the MC Dropout fVI model with different βk. Figure 13 depicts the predicted class probabilities of the toy problem with K = 2 classes. For βk < 1, the areas of confident prediction enlarge but quickly fall back to uniformity, whereas for βk > 1, the confident predictions or more locally concentrated, tapering slowly towards uniformity and ’leaking’ into the unobserved area.

Review summaries:
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. Concretely, 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. The 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). The method produces good results.

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. 

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.

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.