Probabilistic Deep Learning with Generalised Variational Inference
Abstract
We study probabilistic Deep Learning methods through the lens of Approximate Bayesian
Inference.
In particular, we examine Bayesian Neural Networks (BNNs), which usually suﬀer from multiple ill-posed assumptions such as prior and likelihood misspeciﬁcation.
In this direction, we investigate a recently proposed approximate inference framework called
Generalised Variational Inference (GVI) in comparison to state-of-the-art methods including standard Variational Inference, Monte-Carlo Dropout, Stochastic gradient Langevin dynamics and Deep Ensembles.
Also, we expand the original research around GVI by exploring a broader set of model architectures and mathematical settings on both real and synthetic data.
Our experiments demonstrate that approximate posterior distributions derived from such a method oﬀer attractive properties with respect to uncertainty quantiﬁcation, prior speciﬁcation robustness and predictive performance, especially in the case of BNNs.
The code for all the experiments can be found in the following public Github repository: https://github.com/gfelekis/GVI-posteriors-in-Probabilistic-Deep-Learning
1. Introduction
Bayesian methods provide the gold standard method to capture uncertainty in deep learning models through their natural probabilistic representation.
The success of these methods in practice, including the recent advances of Bayesian Deep Learning ﬁeld, relies on Approximate Inference methods.
These approximation schemes perform Bayesian reasoning eﬃciently, as they are approximating the posterior distribution either in a stochastic way (MCMC) or in a deterministic one (Variational Inference), thus allowing Bayesian modelling to be applied to many practical tasks.
Recently, there’s a lot of progress on these ideas mainly through variational techniques (3),(4),(14) and (17).
In recent years, due to the availability of massive datasets, the main focus of Variational Inference is on scalable approaches (15), (5), black box algorithms (21) and Bayesian deep learning architectures such as the Variational Autoencoders (VAE) (18).
In our work, we motivate the recently proposed framework of Generalised Variational Inference (19) which is a generalisation of standard Bayesian and Variational methods and seems to be able to overcome a lot of drawbacks and pathogenies of them regarding the prior, the likelihood and the computational needs.
In this work we advance the research on GVI by exploring a broader set of divergences on more complex Bayesian Neural Network architectures and comparing these with state-of-the-art approximate inference methods.
Contributions: We carry out an extensive comparative analysis of GVI among diﬀerent discrepancy settings and model complexity.
We also compare these against multiple Approximate Inference methods ranging from standard Variational Inference to Monte Carlo
Dropout (13), Stochastic Gradient Langevin Dynamics (22) and Deep Ensembles.
We found certain settings of diﬀerent divergences, hyperparameters and neural network depths that markedly improve the performance of GVI and provide empirical evidence of the superiority of it over traditional approximate inference methods in the framework of Bayesian Neural
Networks.
Finally, we also conduct a detailed empirical analysis of uncertainty quantiﬁcation in a controlled synthetic setting for GVI with a variety of diﬀerent divergences.
We ﬁnally evaluate their epistemic and aleatoric uncertainty and asses model calibration.
2. Theoretical Background - The many views of Variational Inference 2.1. ELBO view (Model selection)
It has been shown by Csiszar (6) and Donsker (9) that the Bayesian inference objective can be seen as the solution to an inﬁnite dimensional optimisation problem and therefore every posterior distribution is the result of a well deﬁned problem of this nature (Zellner in 1998 (23)).
A few years later, Bissiri et.
al (2) extensively discussed a generalised solution inspired by the Fenchel’s conjugate of KL divergence and restated the optimisation problem over P (Θ), about a parameter θ , speciﬁed via an abstract loss function l (compared to the previous log-likelihood one) and regularised by the KL divergence, as follows:
[Equation]
where q ∗ ( θ ) =
q ( θ | κ ∗ ) for some optimal parameter κ ∗ ∈ K , the variational parameter space.
The later is the standard way of deriving Variational Inference, by maximizing the evidence in the data and picking the element from the approximation family Q that maximizes the well-known Evidence Lower Bound (ELBO).
2.2. Discrepancy-minimisation view (DVI)
The Bayesian posterior is not a unique solution to the optimisation problem and hence, we could see its solution as the minimisation of a distance metric (usually in the form of divergences).
Speciﬁcally, if we want to approximate the standard Bayesian posterior q ∗ B ( θ ) with a variational distribution q ( θ ) we could write the solution of this problem as:
[Equation]
| q
[Equation]
The fact that the same objective function that maximizes the
ELBO is the one that minimizes the distance of Q and q ∗ B ( θ ) in the KL divergence sense by just rearranging the terms of the ELBO equation has motivated a large body of research (20), (1), (8) that tries to approximate the posterior by minimizing divergences, diﬀerent from the KL, between the family Q and q ∗ B ( θ ).
In particular, for a divergence measure
D : P (Θ) × P (Θ) −→ R
+ we can deﬁne a new class of methods called Discrepancy Variational
Inference (DVI) methods .
DVI methods’ objective is of the following form:
[Equation]
2.3. Constrained optimisation view (GVI)
Under the traditional Bayesian inference framework, modern machine learning models suﬀer from certain ill-posed assumptions mainly regarding prior misspeciﬁcation (11).
In order to tackle this Knoblauch et al.
(19) proposed the Rule of Three (RoT) framework which split the inference problem into three elements: A loss, a divergence and a space of feasible solutions.
The RoT foundation provides enough ﬂexibility to tackle most of the inappropriate assumptions especially by its modularity nature.
The authors deﬁne a generalised representation of Bayesian inference as follows:
Deﬁnition:
For given observations x 1: n , a prior π ( θ ), a space Π ⊆ P (Θ), a loss function l : Θ × X −→ R and a divergence D ( ·|| π ) :
Π −→ R
+ we say that posteriors have been constructed via the Rule of Three if they can be written as 1 :
[Equation]
The Rule of Three P ( l, D, Π) has a modular interpretation and is decomposed into three separate parts, each of them serving a speciﬁc and separate from each other purpose.
Many of the existing approximation methods can be interpreted by RoT e.g. standard VI solves a problem speciﬁed by the Rule of Three P ( l, D, Π) where: q ∗ VI ( θ ) :
= P (log p ( x i | θ ) , KL ( q || π ) , Q ).
All the above lead us to the third view of VI, the constrained optimisation view (19) which treats the VI solution as the best Q -constrained solution, where Q is a variational family.
Deﬁnition: Any Bayesian inference method solving a RoT form P ( l, D, Q ) for Q = { q ( θ | κ ) : κ ∈ K } ⊆ P (Θ) is a procedure called Generalazided Variational Inference (GVI) .
Hence, GVI like VI has the form P ( l, D, Q ) which satisﬁes the RoT modularity property and like DVI targets non-standard posteriors without conﬂating l and D .
Changing the divergence in the GVI-sense aﬀects only the uncertainty quantiﬁcation and cannot interfere with the way that the best parameter is found (e.g. interfere with the loss function).
3. Experiments and Results
The black box nature of Bayesian Neural Networks (BNN) lurks a high risk of having a misspeciﬁed priors.
Even in small BNN architectures we deﬁne priors over weights and biases without having proper intuition about what they actually express.
As a result, such inappropriate assumptions could lead many times to strange posterior distributions.
In that sense approximate inference methods which can capture robustness to prior misspeciﬁcation might actually be “better”, at least for the BNN case compared to asymptotically exact inference algorithms.
For all these reasons, we motivate the use of GVI to achieve such robustness, which we then test experimentally across a range of diﬀerent discrepancies and real/synthetic datasets.
Our experiments evaluate to what extent switching to a Generalised framework through GVI can actually improve prior misspeciﬁcation, uncertainty quantiﬁcation (UQ) and predictive performance on BNNs.
3.1. Regression on UCI data sets
Initially, we conduct multiple regression experiments on four diﬀerent data sets from the
UCI Machine Learning repository (10) in order to compare diﬀerent divergence measures and diﬀerent neural network depths on the GVI setting.
Thanks to the modularity of RoT, and as a result of GVI, in order to tackle the issue of prior misspeciﬁcation in Bayesian
Neural Networks we just need to focus on trying diﬀerent discrepancy measures while keeping the loss function ﬁxed.
Indeed, in the experiments that we have conducted we expand the research of (19) to almost all the members of the f -divergence family and also Fisher distance, while we kept the loss function ﬁxed to the usual log-likelihood one.
Consequently, we made a straight comparison between the best performing divergences that we found on the previous step, the Standard VI (KL Divergence) and three approximate inference methods: Monte-Carlo Dropout, Deep Ensemble and Stochastic Gradient Langevin Dynamics.
In Figure 1, we can see this comparison of the methods across diﬀerent number of hidden layers for a ﬁxed dataset.
Overall, we empirically demonstrate the superiority of GVI compared to the the other methods in most of the scenarios that we created.
From the complete results that can be found in the Appendix C we can observe that GVI outperforms in every divergence the other methods in 9 out of 12 experiments that we conduct on diﬀerent datasets and model architectures.
The model settings can be found in Appendix A.
3.2. Regression on Gaussian Process ground truth
We evaluate aleatoric (noise) and epistemic (model) uncertainty on a GVI setting across diﬀerent discrepancies, network depths and compare this with the prior art.
From (16) we regress on a heteroscedastic toy data set generated from a GP as the ground truth.
Based on (7) we decompose uncertainty into aleatoric ( U A ) and epistemic ( U E ) as follows:
[Equation]
In Figure 2 we focus on the GVI case of Fisher distance and we can observe that in the interpolation region all the models seem to ﬁt the data really well and performance improves when the number of hidden layers increases.
In the extrapolation regions we observe two diﬀerent trends.
First, for x < − 2 the models seem to be heavily overconﬁdent in the one hidden layer case, because the heteroscedacity of the data lowers noise variance, and UQ improves as the number of hidden layers increases.
For x > 2 we can observe underconﬁdence especially when the number of hidden layers is increasing.
In this region aleatoric uncertainty dominates the epistemic one.
In Figure 3 we provide an extended comparison among the diﬀerent methods for the one hidden layer case and see how GVI behaves compared to other approximate inference methods regarding the UQ.
As mentioned before UQ is performed poorly for one hidden layer case in general due to the heteroscedacity of the data which makes all the models overconﬁdent in the left extrapolation area.
As expected epistemic uncertainty, which expresses model uncertainty, dominates in this region and thus a more complex model will be more suitable for this evaluation (see Fig. 2).
We observe that compared to KL, JS divergence seems to oﬀer both better UQ and balance of aleatoric-epistemic uncertainty in all the regions.
It is worth mentioning that KL and JS divergences both belong to the f -family and the later is a normalized and symmetrical version of the ﬁrst which makes it smoother.
For the three approximate inference methods, both Deep Ensembles and SGLD are badly ﬁt the data.
They oﬀer a poor UQ and predictive performance (black line) even in the interpolation area, contrary to the MC Dropout model which performs signiﬁcantly better.
MC Dropout eﬃciently captures aleatoric uncertainty also in the left extrapolation region and does not collapse it compared to other methods apart from JS (GVI).
This is probably due to the fact that is know that it’s eﬀectively introduces noise (12).
In Table 1 we report the predicted probability diﬀerence to ground truth and the eﬀective coverage at the 1 σ interval for GVI and MC Dropout methods.
Except KL divergence JS,
Fisher and MC Dropout provide nice calibration measures.
Overall, GVI with JS divergence and MC Dropout seem to perform better compared to all the other methods in UQ, predictive performance and calibration.
[Table]
4. Conclusions
This work oﬀers a comparative analysis of a generalised Bayesian inference framework among diﬀerent discrepancy settings, datasets, network architectures and other approximate inference methods.
We have empirically shown the usefulness of GVI on BNNs in certain settings regarding the model complexity and provide extensive insights of the performance of diﬀerent divergence measures and its beneﬁts over traditional approximate inference methods.
Overall, GVI oﬀers an alternative approach to dealing with the challenging model and prior speciﬁcation tasks in BNNs.
It handles model misspeciﬁcation via robust scoring functions and prior misspeciﬁcation through alternative divergences and we argue that these might be preferred in Bayesian Deep Learning.
Appendix A. Regression on UCI data sets A.1. Model settings & data
Across all the experiments we kept the structure of the neural network ﬁxed in order to make the comparisons as fair as possible.
In particular, we use a Multi-Layer Perceptron with [100] hidden units for the one hidden layer case, [100, 100] hidden units for the two hidden layer case and [100, 100, 100] hidden units for the three hidden layer case.
The activation function was a ReLU function and inference was performed via Bayes by backprop and the Adam optimiser.
For the training of each model we run 100 epochs and perform 30 random splits of each data set with a split of 90%-10% (train-test).
All the models were evaluated on the test sets using the average negative log likelihood (NLL) as well as the average root mean square error (RMSE).
Also, for each of the 30 splits, the predictions are computed based on 100 samples from the variational posterior.
Note that the priors and variational posteriors are both fully factorised normal distributions and thus our model was predicting the regression mean µ ( x ) and the log-standard deviation log σ ( x ).
Also, that helped us of having all of our divergences in a closed Gaussian form.
Below we present the closed form of the divergences that were used:
For a prior π ∼ N ( µ 1 , σ 21 ) and approximate posterior q ∼ N ( µ 2 , σ 22 ) distributions we have the following closed forms:
• Kullback-Leibler Divergence:
[Equation]
• Reverse Kullback-Leibler
Divergence:
[Equation]
• α -Divergence:
[Equation]
• Jensen-Shannon Divergence:
[Equation]
• Total Variation Distance:
For the TV distance there is no closed form and hence we approximate it by its bounds.
In particular, the following inequality holds:
[Equation]
Hence we deﬁned an approximation for each bound (upper U, lower L):
[Equation]
• α -R´enyi Divergence:
[Equation]
where:
[Equation]
• Fisher Distance:
[Equation]
where:
[Equation]
Appendix B. Regression on Gaussian Process ground truth B.1. Model settings & data
Here, the heteroscedastic data was generated by a Gaussian Process with an RBF kernel ( l = 1 , σ n = 0 .
3 | x + 2 | ).
Also a Multi-Layer Perceptron was used as the regressor with [100]
ReLU hidden units for the one hidden layer case, [100, 200] ReLU hidden units for the two hidden layer case and [100, 200, 100] ReLU hidden units for the three hidden layer case and in all cases it was trained for 500 epochs.
We compute epistemic and aleatoric uncertainty and we also investigate the behaviour of all the models across diﬀerent divergences and diﬀerent depths by not only visual inspection but also some model selection and information criteria that we are going to see below.
B.2. Model Calibration
Here is the complete table of the diﬀerences between the predicted probability and the true probability for each of the divergences.
Divergence
No. of hidden layers
Predicted probability
Diﬀerence
KL 0.
KL 0. KL 0.
RKL 0. RKL 0.
RKL 0.
A 0. A 0.
A 0.
AR
[Table] [Table]
We highlight the divergence measures which correspond to the two best calibrated models for each neural network depth.
Note here that we deﬁne a well-calibrated model to be one whose predicted probability diﬀers no more than 0.1 from the true probability:
• For the one hidden layer case: TVL and α AR • For the two hidden layers case: KL, AR and A • For the three hidden layers case: TVL, A and JS.
In Figure 4 we can see the calibration curves for each divergence for diﬀerent number of hidden layers.
Each curve represents how conﬁdence varies across diﬀerent sigma intervals, indicating either underconﬁdence if the curve is above the dashed line or overconﬁdence if it falls below it.
Two key observations here are that the conﬁdence increases together with the number of layers in almost all the models and that there is a general overconﬁdence trend on lower sigma intervals i.e. (0 , 0 . 7).
B.3. Model Selection
Finally, one standard way to perform model selection in Probabilistic Deep Learning is to compute the values of certain criteria from the Information Theory ﬁeld.
Speciﬁcally here we tested three of them for each model:
Bayesian Information Criterion - BIC
[Equation]
Akaike Information Criterion - AIC
[Equation]
ˆ
[Equation]
3. Hannan–Quinn Information Criterion - HQC
[Equation]
where ˆ L is the maximized value of the likelihood function of the model M, p is the number of parameters estimated by the model and n the sample size.
Information Criteria (ICs) are really useful as they can measure the eﬃciency of a model in terms of predicting the data and also the complexity of the model.
We can observe that all three of them are constituted of two terms: a complexity term, which is the one dependent to the number of model’s parameters, and a data ﬁt term.
Generally, the computation of ICs is a trade-oﬀ between these two terms as they try to answer the question: Do I need more data or a more complex model?
Therefore, when as a next step we computed all the values of the ICs for all the models we ended up observing that the complexity penalty dominated the IC value and hence the one hidden layer case was always the preferable choice (see table in the Appendix).
Here it is important to mention a few things about the information criteria.
It would be naive to conclude that the information criteria, which were used here, are useless.
Instead, they serve their purpose as they penalise the model for trying to ﬁt more parameters than the needed ones for this speciﬁc problem/data set.
In the example here, where the data set was generated from an RBF kernel of a Gaussian process which is a stationary kernel, the penalty term dominates.
It is not necessarily true that the IC value has to increase as long as the number of parameters grows.
As we said, all the ICs perform a trade-oﬀ between model complexity and data ﬁt.
Here, the data ﬁt component is roughly equivalent for all the diﬀerent neural network depths, thus the only thing that changes is the number of parameters.
Consequently, it makes perfect sense to see this monotonic deterioration.
For neural networks, an advantage of having more than one hidden layers is that they can capture non stationary kernels (Neural Network kernel, Polynomial kernel, etc).
So we encourage the reader to examine the same results for a non-stationary case where the ICs would be much more informative than here.
Overall, it is crucial to understand that the results have to do with the speciﬁc problem, which is simple, and the fact that number of parameters is not a really good complexity value for an MLP.
Having said that, it would probably be more fruitful to compare just the diﬀerences between the log-likelihoods of the models.
In the table below we present the log-likelihood value for each model
and we highlight in bold the best for each divergence.
[Table]
From this table we can notice that, although some divergences achieve the best log-likelihood value at the maximum number of hidden layers, it is not always the case and it does not justify the increase of the number of model parameters.
It is interesting to see that the best model based on the log-likelihood value agrees in most of the cases with our visual inspection conclusion about each layer.
Appendix C. Analytical results C.1. Regression on UCI data sets
In Figure 4 we can see the comparative analysis of the diﬀerent GVI settings for the multiple divergences and in Figure 5 the comparison of GVI’s top performers against common approximate inference methods:
We can make the following remarks from Figure 6:
• For the one hidden layer case: The GVI methods seem to outperform MC-Dropout and Deep Ensembles and be equally good with the SGLD method especially for the RMSE values.
MC-Dropout seems to be the worst performer here.
• For the two hidden layers case: Here, all the approximate inference methods seem to improve and come closer to the GVI RMSE values in most of the cases.
However, we see that Ensemble and SGLD outperform all the other methods when it comes to NLL values.
• For the three hidden layers case: We can notice that the uncertainty value (variance) increases in all the methods.
GVI remains the best performer and in some cases alongside SGLD and also improves its performance in the NLL scores.
C.2. Regression on Gaussian Process ground truth
In Figures 6 and 7 are the complete results from the regression workaround on the Gaussian
Process ground truth data set.
Also, here we provide the best models based on visual inspection:
• For the one hidden layer case: JS, KL and A and α AR • For the two hidden layers case: KL, F and TVU • For the three hidden layers case: TVU, A and JS.
Figure