F UNCTION - SPACE V ARIATIONAL I NFERENCE FOR D EEP B AYESIAN C LASSIFICATION
Anonymous authors
Paper under double-blind review
A BSTRACT
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 overconﬁdent and unpredictable behavior.
However, weight-space priors are model-speciﬁc, can be difﬁcult 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 speciﬁcation of predictive epistemic uncertainty.
This is achieved through the perspective of stochastic neural network classiﬁers 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 classiﬁcation tasks demonstrate that our function-space inference scales to large-scale tasks and models, improves adversarial robustness and boosts uncertainty quantiﬁcation across models, without inﬂuencing architecture or model size.
I NTRODUCTION
Deep learning (Goodfellow et al., 2016) has enabled powerful classiﬁcation 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 quantiﬁcation (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 efﬁciency (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 quantiﬁcation (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 inference 1 (Sun et al., 2019)
(fVI) for classiﬁcation.
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 classiﬁer is Dirichlet distributed, rather than solely categorical, we gain a richer predictive uncertainty representation while retaining the same mean predictions as categorical classiﬁers.
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 speciﬁc model or data transformations.
We demonstrate that this method improves uncertainty quantiﬁcation and adversarial robustness across a range of popular models and datasets for both small- and large-scale inference.
[Figure]
1. C ONTRIBUTIONS
We incorporate Dirichlet predictive priors into deep Bayesian networks using function-space inference.
Concretely, our contributions are:
We demonstrate an extension of variational implicit processes into the classiﬁcation setting.
We propose a function-space variational inference method for classiﬁcation, using an effective KL divergence estimator from softmax samples via Newton-based Dirichlet estimation.
R ELATED WORK
We build on prior work that considers classiﬁcation in the Bayesian setting and function-space inference of Bayesian neural networks.
2. B AYESIAN CLASSIFICATION
Compared to regression, classiﬁcation 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 classiﬁers (Wenzel et al., 2019).
Gaussian processes have also been used with a Dirichlet predictive using a log-normal approximation (Milios et al., 2018).
Classiﬁcation 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-ﬁeld) 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.
2. F UNCTION - 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 (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-deﬁned objectives in many cases Burt et al.
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 classiﬁcation, 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 ﬁnite representations of functions to learn a particle representation of the function-space posterior through the gradient ﬂow of the log posterior.
Function-space inference is also an attractive approach to continual learning, e.g. Pan et al.
D IRICHLET PREDICTIVE PRIORS FOR B AYESIAN CLASSIFICATION
In this section, we discuss how function-space variational inference can be applied to classiﬁcation using variational implicit processes (VIP) (Ma et al., 2019).
3. V ARIATIONAL IMPLICIT PROCESSES FOR CLASSIFICATION
Let D ⊂ ( X ×Y ) = { ( x n , y n ) }
Nn =1 be the training data consisting of N observed pairs of input data x n and corresponding K -dimensional, one-hot class label vectors y n .
Deep learning approaches to classiﬁcation 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 deﬁnes a deterministic function f which maps an input x ∈ X to an element f x ∈ ∆ K − 1 , where ∆ K − 1 denotes the K −simplex and elements from ∆ K − 1 are K -dimensional probability vectors.
More precisely, we write f x = g ( x , w ) =
σ ( φ ( x ; w )) and y ∼ Cat ( ·| f x ) .
The softmax function σ ensures that f x is a valid probability vector,
i.e. that f x k ≥ 0 and (cid:80) Kk =
1 f x k = 1 .
In conventional ML training, the weights w are optimized by maximizing log p ( D| f )
= log (cid:81)
D Cat ( y | f x ) =
(cid:80) D (cid:80)
Kk
=1 y k log f x k , where (cid:81) D and (cid:80) D denote (cid:81) ( x , y ) ∈D and (cid:80) ( x , y ) ∈D respectively, and φ , σ , 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 classiﬁer, 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 ( f x ) .
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.
[Figure]
Sampling the weights and evaluating g ( x , · ) is equivalent to sampling the Dirichlet predictive, while varying x with ﬁxed weight w corresponds to evaluating a classiﬁer g ( · , w ) .
Together, the joint variation of w and x deﬁnes 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.
Deﬁnition 3.1 (Implicit stochastic processes, from Ma et al.
An implicit stochastic process (IP) is a set of random variables f , such that any ﬁnite subset f x 1: L = { f x 1 , f x 2 , ..., f x L } with
L ∈ N has a joint distribution which is implicitly deﬁned by the generative process, [Equation] ∼ p ( [Equation]
Applying Deﬁnition 3.1 to stochastic network classiﬁers 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 speciﬁc 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 classiﬁcation setting we are now reasoning over a joint distribution of Dirichlet distributions on a ﬁxed 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 ﬁnite 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 deﬁned 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 inﬁnite parameter spaces, as in e.g. Bayesian nonparametrics (Hjort et al., 2010).
Further details regarding the formal stochastic process deﬁnition 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 ) .
3. F UNCTION - 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),
[Equation]
which resembles the conventional evidence lower bound objective (ELBO) (Jordan et al., 1999; Hoffman et al., 2013).
This enables the explicit speciﬁcation of the desired behavior of the predictive distribution directly in function space, whereas the conventional ELBO relies on indirect speciﬁcation 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,
[Equation]
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 ﬁnite measurement sets X ∈X L of index sets X L of size L ,
[Equation]
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,
[Equation]
where S ⊂ X L is a ﬁnite set of (now) ﬁxed 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 beneﬁt of this index set design in Section 4, while methods of designing effective index sets is a topic of future work.
3. P OSTERIOR DIRICHLET AND KL DIVERGENCE ESTIMATION
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 ﬁxed-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 /z x
and the Dirichlet precision z x .
By matching the ﬁrst moment of the empirical distribution of f x , we obtain an estimate of ¯ α x ≈ 1 M (cid:80) Mm =
1 f ( m ) x .
We then ﬁx ¯ α x and estimate the precision z x using a fast, iterative, quasi-Newton algorithm (Minka, 2000) using f (1: M ) x
[Equation]
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 L Dir (cid:16) f (1: M ) x , α x (cid:17) is the Dirichlet log-likelihood log (cid:81) Mm
[Equation]
1 Dir (cid:16) f ( m ) x | α x (cid:17) .
This corresponds to breaking the dependence between ¯ α x and z x that would be considered in the full MLE case in order to accelerate the resulting approximation.
With ¯ α x and z x estimated, α x and thus our variational posterior q ( f x | θ )
=
Dir ( f x | α x ) is now available.
Moreover, for tractability we choose a factorized prior and assume a factorized posterior, i.e. q ( f S )= (cid:81) Ll = 1 q ( f s l ) , 2 so that the fKL term from Equation 3 can be estimated using M samples with the Dirichlet KL divergence
[Equation]
(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 ( f x ) =
Dir ( · | β ) , where each β k is chosen to be 1 unless
Weight-space prior fVI, β = [1 , 1] fVI, β = [2 , 1] fVI, β = [1 , 2]
M A
P
M C
D r opou t
E n s e m b l e otherwise stated. During optimization, we treat q ( f x | θ ) 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 simpliﬁcation removes the inﬂuence of θ on Dir ( f x | α 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 ﬁxed 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 z x 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.
E XPERIMENTS
In this section, we present an empirical evaluation of our proposed fVI approach on classiﬁcation tasks for several models, comparing performance against their conventional training procedure.
We ﬁrst visualize the effects of a Dirichlet prior on a toy problem, and then perform image classiﬁcation 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 classiﬁcation accuracy, log-likelihood (LLH) and expected calibration error (ECE) (Naeini et al., 2015), which estimates the calibration of accuracy vs conﬁdence through binning the predicted class probabilities.
Additional details are provided in Appendix F.
4. T OY PROBLEM
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 overconﬁdent extrapolation.
While the weight-
Subnetwork
[Figure]
MC
[Figure]
Ensemble fVI
Uniform prior 0 45 90 135−−−−
Angle °
L ogL i k e li hood
Standard Index Set 0 45 90 135−−−−
Angle ° + 90° aug.
0 45 90 135−−−−
Angle ° + 90°, 180° aug.
Augmenting the index set further with 180° transforms ﬁxes the degradation, demonstrating an overall improvement in UQ over the initial training data index set. The denotes the performance of the uniform prior.
The subnetwork performance is taken from Daxberger et al.
(2021), and uses a ResNet-18 rather than an MLP.
Further metrics are reported in Figure 9.
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 classiﬁcation, e.g. Liu et al.
Furthermore, fVI facilitates specifying class-biased priors, which would be difﬁcult to achieve with weight priors.
4. R OTATED MNIST
Following Ovadia et al.
(2019), we train on the MNIST handwritten digit classiﬁcation 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 classiﬁcation 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 quantiﬁcation.
Subnetwork linearized Laplace (Daxberger et al.,
) 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.
4. A SSESSING INDEX SET DESIGN
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 classiﬁcation experiments, we use the training data as the index set, demonstrating performance without this inductive bias.
4. L ARGER - SCALE IMAGE CLASSIFICATION UNDER CORRUPTION
We used the regular train splits of the CIFAR10 and CIFAR100 (Krizhevsky, 2009) image classiﬁcation 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, MAP MC Dropout
Ensemble
Radial
Rank
Subnetwork MAP fVI MC Dropout fVI
Ensemble fVI Radial fVI Rank1 fVI Uniform prior
C I F A R
A cc u r ac y −−−− L ogL i k e li hood 0 .0 .0 .0 . E C E 0 1 2 3 
 Corruption
[Figure] [Figure]
cc
[Figure]
r ac
[Figure]
0 1 2 3 4−−− Corruption L ogL i k e li hood 0 1 2 3 4 0 .0 . Corruption
E C E
All models use a ResNet-18 architecture.
For CIFAR10, there is a clear beneﬁt of fVI priors over weight-space for log-likelihood.
Accuracy performance is constant, while fVI underﬁts 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 a result, weight-space ensembles achieve better performance over all corruptions.
Moreover, fVI models underﬁt w.r.t.
ECE in-distribution due to the prior, but improve in calibration OOD, whereas weight-space models worsen.
This result indicates prior speciﬁcation requires more care for high-dimensional classiﬁcation.
Interestingly, the fVI 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.
contrast and saturation, or adding noise, blur or other artifacts, such that it becomes ever more difﬁcult for the models to classify them.
For this experiment, we used CNN models based on the ResNetarchitecture (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 signiﬁcantly higher for corruption 0, but this underﬁtting provided superior ECE when OOD.
This suggests the function-space prior should be ﬁne-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 beneﬁt 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.
4. L ARGER - 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 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 (cid:15)
=
(no attack) to (cid:15) = 0 .
Although both weight-space and function-space models lose their classiﬁcation accuracy when the FGSM attack is introduced, the fVI models only suffer small decreases in log-likelihood, whereas the weight-space LLH performance drops signiﬁcantly.
We also observe the accuracy vs
MAP
MC Dropout
Ensemble
Radial
Rank
Uniform prior
MAP fVI
MC Dropout fVI
Ensemble fVI
Radial fVI
Rank1 fVI
C I F
A
R
A cc u r ac y −−−−−
L ogL i k e li hood 0 .0 .0 .0 .
E C E 0 .0 .0 .
Epsilon
C I F A R
A cc u r ac y 0 .0 .0 .−−
Epsilon
L ogL i k e li hood 0 .0 .0 .
Epsilon
E C E robustness trade-off in the fVI models discussed in Section 4.4. We attribute this behavior to the quality of the uncertainty quantiﬁcation 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.
C ONCLUSION
We propose a function-space variational inference approach to deep Bayesian classiﬁcation, which enables the use of Dirichlet predictive priors to improve uncertainty quantiﬁcation.
Our approach provides a uniﬁed formalism of prior work on Dirichlet-based classiﬁers 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 quantiﬁcation and adversarial robustness.
By considering the design of the index set, we show how one can trade-off scalability against OOD uncertainty quantiﬁcation by extending the fKL evaluation beyond the training distribution.
Future research should develop index sets for fVI, speciﬁcally effective methods for constructing them to reﬂect the test distribution, e.g. through using data augmentation or unlabeled data.
[Figure]
E THICS S TATEMENT
As this paper looks at the abstract problem of classiﬁcation under a Bayesian lens, we are not aware of direct ethical implications of this work.
R EPRODUCIBILITY S TATEMENT
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
A R ELATION TO P RIOR W ORK
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 ﬁrst explain prior work using their own notation and then discuss the correspondence to our fVI perspective.
A. P RIOR NETWORKS
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 ﬁnal 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,
[Equation]
The ﬁrst expectation E p in accounts for the actual learning, i.e. ﬁtting the training data, whereas the second expectation E p out is supposed to regularize the model by matching a prior distribution.
Accordingly, the ﬁrst 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 deﬁne 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.
A. B ELIEF MATCHING
Joo et al.
(2020) assumes a Dirichlet prior which, together with the categorical ground truth class labels, deﬁne a target Dirichlet posterior.
A neural network is used to directly predict concentration parameters of a Dirichlet posterior q z | x by replacing the ﬁnal softmax layer with an element-wise exponential operation.
To learn the target posterior, belief matching maximizes
[Equation]
[Equation] where E q z | x [log z y ] is the expected log-likelihood of the training data and
D KL q Wz | x p z | 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.
A. E VIDENTIAL DEEP LEARNING
Sensoy et al.
(2018) directly predicts the concentration parameter of a Dirichlet distribution by using a neural network with ReLU activations as ﬁnal 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 y i and predicted probabilities p i .
Furthermore, a regularizing
KL divergence term is added, resulting in a total loss function,
[Equation]
where y ij are individual 0-1 target labels, ˆ p ij are components of the predicted Dirichlet mean, S i is the predicted Dirichlet precision,
˜ α i is the predicted Dirichlet concentration parameter, 1 is a vector of ones and λ t is a annealing coefﬁcient for optimization.
The ﬁrst part of their loss is responsible for ﬁtting 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.
A. E XPERIMENTAL C OMPARISON
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 ofﬁcial open-source implementations 3 4 .
Figure 8 illustrates the test accuracy, log-likelihood, and expected calibration error for the corrupted
CIFAR10 image classiﬁcation 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 .
[Figure]
[Figure]
[Figure]
[Figure]
[Figure]
B P RACTICAL C ONSIDERATIONS
This section discusses practical issues with the approach in more detail.
B. M INI - BATCH TRAINING
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.
B. P RIOR S PECIFICATION FOR H IGH O UTPUT D IMENSION
When scaling to higher dimensional classiﬁcation tasks, speciﬁcally K ≥ 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.
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.
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 10 x scaling suggested by the Dirichlet KL due to the summation terms.
C D EEP STOCHASTIC CLASSIFIERS AS STOCHASTIC PROCESSES
Stochastic classiﬁers are stochastic processes instantiated as a feedforward or convolutional neural networks with stochastic weights.
Formally, our stochastic process f is deﬁned on the sample space Ω with an index set X deﬁned 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 −can be deﬁned for each x ∈ X and we write f ( x ) =
f x .
Kolmogorov’s extension theorem (Tao, 2013) guarantees the existence of a stochastic process f if for each L ∈ N the ﬁnite marginal joint distributions p 1: L ( f x 1: L ) , where f x 1: L = { f x 1 , . .
, f x L } , satisfy exchangeability and consistency.
These ﬁnite marginal joint distributions are implicitly deﬁned according to Deﬁnition 3.1.
Exchangeability :
For any permutation π of 1 , . . .
, L , p π ( x 1: L ) ( f π ( x 1: L ) )
=
p x 1: L ( f x 1: L ) .
This requires that the process behavior is invariant to the order of inputs.
For a feedforward neural network, this is satisﬁed because the respective predictions do not change if the order of inputs changes.
Consistency : For any 1 ≤ L (cid:48) <
L , p 1: L (cid:48) ( f x 1: L (cid:48) )
= (cid:82) p 1: L ( f x 1: L ) d f x L (cid:48)
+1:
L .
This requires that future evaluations are independent of past evaluations.
For a feedforward neural network, this is satisﬁed because predictions do not depend on previous predictions.
D V ARIATIONAL 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 ) .
[Table]
[Figure]
E D IRICHLET MUTUAL INFORMATION
Smith & Gal (2018) proposed to use the mutual information MI ( θ ; 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 quantiﬁes 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 ( ·| f x ) and the latent probability vector f x follows a Dirichlet distribution f x ∼ Dir ( ·| α ) , the equivalent mutual information between class labels and model predictions MI ( f x ; y )
=
H ( y ) − H
( y | f x ) can be derived in closed-form from the Dirichlet approximation using
[Equation]
The complete derivation is provided in Appendix E.1 on the next page.
The ﬁrst term H ( y ) corresponds to the categorical entropy of the expected value of f x .
It is equivalent to the standard measure of predictive entropy for softmax-based ML deep learning classiﬁers.
The second term H ( y | f x ) is the conditional entropy of the class label y given a Dirichlet observation f x
and it is inﬂuenced by the spread of f x , whereas the ﬁrst term only depends on its mean.
Although this function-space variational formulation admits closed-form solutions for MI ( f x ; y ) , the resulting values will still be approximations in our setting because the variational posterior Dir ( f x | α x ) was estimated from M observations.
Nonetheless, it is not limited to models with stochastic weights, but instead provides a uniﬁed approach to measuring epistemic uncertainty for any classiﬁer with a
Dirichlet predictive distribution.
E. M UTUAL I NFORMATION D ERIVATION
Let y ∼ Cat ( ·| f x ) , f x ∼ Dir ( ·| α ) , f x k ∼ Beta ( ·| a k , b k ) , where a k = α k and a (cid:48) k
= a k +1 , while b k =
z − α k and z = (cid:80)
Kk
=1 α k ,
MI (
[Equation]
x
[Equation]
[Equation] [Equation]
[Equation] [Equation]
[Equation] [Equation]
[Equation] [Equation]
[Equation] [Equation]
[Equation] [Equation]
−
[Equation]
a
[Equation]
a
[Equation] [Equation]
−
[Equation]
k
[Equation]
a k
[Equation]
a
[Equation]
a
[Equation]
[Equation] [Equation]
−
K
[Equation]
k = α k
[Equation]
ψ
[Equation]
α k + 1) − ψ ( z + 1)) .
F I MPLEMENTATION D ETAILS AND C OMPUTATIONAL C OMPLEXITY
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 thanhours 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 ﬁxed 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.
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.
, we normalized pixel values using the empirical mean and standard deviation, and employed data augmentation during training by ﬁrst selecting random crops of size 32x32x3 after adding 4 pixels of zero padding to each side and then randomly ﬂipping 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.
, 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 ﬁnal 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 andtraining 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 Rank
Gaussian distributions over all convolutional weights but not over the ﬁnal 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 deﬁned a minimum and maximum precision for the posterior
Dirichlet estimation: z min = K and z max = 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
=
z max because M = 1 .
Similarly, we used M = 1 for the MC Dropout and Radial BNN models during training and also set z
= z max , 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 f x ( m ) k
[Equation]
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
[Equation]
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),
[Equation]
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 95 th quantile was 3, the 99 th 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 ) .
G E XPERIMENTAL R ESULTS
All experimental results are reported here.
G. R OTATED MNIST
[Table]
[Table] [Table]
[Table] [Table]
[Table] [Table]
[Table] [Table]
G. CIFAR
[Table] [Table]
[Table] [Table]
[Table] [Table]
G. CIFAR
[Table]
[Table]
[Table]
[Table]
[Table]
[Table]
H A BLATION S TUDIES H. S AMPLES DURING TRAINING
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 onlysample 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 reﬂect 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 thesample approximation is reasonable for the training time savings. [Figure]
H. S CALING ISSUES WITH HIGH LABEL DIMENSIONALITY
The CIFAR100 experiments revealed an issue with the fELBO objective that caused underﬁtting for the larger label dimension.
To examine why, recall that the categorical likelihood is log f x k when y k = 1 .
Therefore, the dimensionality of y does not directly inﬂuence the value.
Conversely, the KL divergence between two Dirichlet densities D KL ( p 1 || p 2 ) does incorporate the label dimensionality
K (Rauber et al., 2008),
[Equation]
{+To counteract this linear increase due to the summation terms, we can assess a heuristic annealing scale factor on the fKL during training of 
+}
/K =
λ ,
[Equation]
To investigate this relationship between the Dirichlet KL divergence and the number of classes of the classiﬁcation, we conducted a toy experiment in a hypercube [ − 1 , 1] D with ﬁxed 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. x d = 0 , and leveraging all permutations to create up to
K = 2 D 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 ﬁnal fKL at convergence consistently low and signiﬁcantly improves the model’s test log-likelihood.
[Figure]
H. C HANGING PRIOR D IRICHLET PARAMETERS
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 conﬁdence 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 signiﬁcant 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.
[Figure]
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 .
Figuredepicts the predicted class probabilities of the toy problem with K = 2 classes.
For β k < 1 , the areas of conﬁdent prediction enlarge but quickly fall back to uniformity, whereas for β k > 1 , the conﬁdent predictions or more locally concentrated, tapering slowly towards uniformity and ’leaking’ into the unobserved area.
[Figure]
[Equation]
problem (Figure 3) with varying uniform prior precision.