Generative Posterior Networks for Approximately Bayesian Epistemic Uncertainty Estimation
Abstract
In many real-world problems, there is a limited set of training data, but an abundance of unlabeled data.
We propose a new method, Generative Posterior Networks (GPNs), that uses unlabeled data to estimate epistemic uncertainty in highdimensional problems.
A GPN is a generative model that, given a prior distribution over functions, approximates the posterior distribution directly by regularizing the network towards samples from the prior.
We prove theoretically that our method indeed approximates the Bayesian posterior and show empirically that it improves epistemic uncertainty estimation and scalability over competing methods.
Introduction
In supervised learning tasks the distribution of labeled training data often does not match exactly with the distribution of data the model will see at deployment.
This distributional shift can cause significant problems in safety-critical environments where mistakes can be catastrophic.
Ideally, we want our learned models to be able to estimate their epistemic uncertainty – uncertainty deriving from lack of data samples.
Unfortunately, deep learning models struggle to estimate this type of uncertainty.
While there are many proposed methods for addressing this problem, epistemic uncertainty estimation in deep learning remains an open problem due to out of distribution (OOD) performance and scalability. [Figure]
[Figure]
Concurrently, many recent works have shown incredible performance using unlabeled data (Radford et al., 2019; Chen et al., 2020).
While labeled training data is expensive to collect, in many real-world problems, such as image classification, there is an abundance of unlabeled data.
Moreover, this unlabeled data is often much more representative of the deployment distribution.
In this work, we propose a new method for estimating epistemic uncertainty that leverages this unlabeled data.
Our work builds off of the extensive ensembling literature.
Neural network ensembling is a method that predicts epistemic uncertainty by looking at the disagreement between the ensemble members.
Additionally, under the right regularization, ensembles approximate samples from the Bayesian posterior Pearce et al.
; He et al.
However, these methods have one main drawback, namely that each new sample from the posterior requires training a new network from scratch.
To address this challenge, we introduce Generative Posterior Networks (GPNs), a generative neural network model that directly approximates the posterior distribution by regularizing the output of the network towards samples from the prior distribution.
By learning a low-dimensional latent representation of the posterior, our method can quickly sample from the posterior and construct confidence intervals (CIs).
We prove that our method approximates the Bayesian posterior over functions and show empirically that our method can improve uncertainty estimation over competing methods.
While not motivated by the Bayesian Inference problem, Osband et al.
(a concurrent work) propose Epistemic Neural Networks, which is very similar to our method.
The key difference between our method and Epistemic NNs is in the regularization: our method takes advantage of unlabeled data to regularize towards outputs as opposed to the parameters of the prior networks.
In our experiments we show empirically our method improves posterior sampling and OOD detection.
Generative Posterior Networks
Randomized MAP
Sampling We consider a slight variation on the usual Bayesian Inference problem: instead of finding the posterior of the parameters of a function, we will instead find the posterior of the outputs of that function.
Specifically, consider the transformed random variable ˆY = f
( x sample ; θ ) for some function f parameterized by θ and some set of unlabeled sample points x sample .
We assume a known prior over parameters P ( θ ) ∼ N ( µ θ , Σ θ ) as well as access to a noisy dataset ( x obs , y obs ) =
{ ( x 1 obs , y 1 obs ) . .
, ( x N obs , y N obs ) } , where the observations y i obs = f
( x i obs ; θ ) + ϵ depend on an unknown parameter θ and are corrupted by Gaussian noise ϵ ∼ N (0 , σ ϵ ) .
The goal of this Bayesian Inference problem is to approximate samples from the posterior,
[Equation]
Pearce et al.
(2020) showed how we can use a method called Randomized MAP Sampling (RMS) to approximately sample from the posterior distribution.
The idea behind RMS is to use the fact that the parameters that minimize the regularized MSE loss are equal to the maximum a posteriori (MAP) solution for the posterior.
Thus, using optimization techniques, like gradient descent, we can easily find MAP solutions.
However, the challenge of sampling from the posterior remains.
Instead, RMS randomly samples shifted prior distributions such that the resulting MAP solutions for each of these shifted distributions form samples from the desired posterior.
Unlike Pearce et al.
, we consider using the RMS technique on our modified Bayesian Inference problem, to estimate the posterior of the outputs ˆY as opposed to the parameters θ .
Specifically,
RMS in this setting samples “anchor points” ˆY anc from some distribution ˆY anc ∼ N ( µ anc , Σ anc ) . [Table]
[Table]
Each anchor point corresponds to the mean of a shifted prior distribution P anc ( ˆY ) =
N ( ˆY anc , Σ ˆY ) .
Now we define the MAP function which finds the MAP solution using this shifted prior:
[Equation]
By sampling different anchor points
[Table]
we can then construct a probability distribution over MAP solutions P ( Y MAP ( ˆY anc )) .
If we assume that the prior over ˆY is Gaussian, this distribution is equivalent to the posterior distribution P ( ˆY | y obs ) for a specific anchor distribution.
While this
Gaussian prior assumption does not hold exactly in general, it will hold exactly when f is linear or has an infinitely wide final layer.
Theorem 2.1.
Assume that the prior distribution of ˆY is Gaussian, P ( ˆY ) =
N ( µ ˆY , Σ ˆY ) .
If we choose the distribution over ˆY anc to be P ( ˆY anc ) =
N ( µ anc , Σ anc ) , where µ anc = µ ˆY and Σ anc = Σ ˆY + Σ ˆY Σ − 1 like Σ ˆY , then P ( ˆ Y MAP ( ˆY anc ))
=
P ( ˆY | y obs ) .
Our proof follows a similar technique to Pearce et al.
and is provided in Appendix E.1.
Just as in Pearce et al.
(2020), we use the approximation Σ anc ≈ Σ ˆY to make it tractable to sample from the anchor distribution.
1 With this assumption, the anchor point sampling distribution ˆY anc ∼ N ( µ anc , Σ anc ) becomes approximately equal to the output prior distribution P ( ˆY ) .
Generative Model
While RMS can be used to sample from the posterior, every new sample requires solving a new optimization problem.
Instead, we propose to learn a generative model to approximate the MAP function itself.
Let g be a neural network parameterized by some vector, ϕ , that takes as input a sample from the anchor distribution, ˆY anc , and outputs the MAP solution ˆY MAP ( ˆY anc ) .
Since we are using the prior distribution as our anchor distribution, we can sample anchor points using ˆY anc = f ( x sample ; θ anc ) , where θ anc ∼ N ( µ θ , Σ θ ) .
If we assume that g has enough expressive power to represent the MAP function, then g ( ˆY anc ; ϕ )
=
ˆ Y MAP ( ˆY anc ) , if ϕ is equal to:
[Equation]
[Figure] [Figure] where δ j = g ( x j sample , θ anc ; ϕ ) − f ( x j sample ; θ anc ) .
See Appendix E.2 for a step-by-step derivation.
In practice, of course, it is expensive to compute the regularization term, σ 2 ϵ δ T Σ − 1 ˆ Y δ , for a large set of sample points x sample .
Instead, we approximate this term with β ∥ δ ∥ 22 where β is a hyper-parameter.
This assumes independence and uniform variance of the prior, but we found it to work well in practice.
Finally, optimizing the generative function g over the the space of all anchor parameters θ anc is impractical, as this space is too high dimensional.
Instead, we construct a low-dimensional embedding for θ anc .
Details for how we construct this embedding are provided in Appendix B.
Experiments
The goal of our experiments is to illustrate the ability of our method to accurately model epistemic uncertainty for OOD data while retaining high performance on in distribution data.
To do so, similar to related work (Liu et al., 2020; Van Amersfoort et al., 2020), we consider two separate datasets, “In
Distribution” and “Out of Distribution” (OOD).
At training time, we provide labeled data from the In
Distribution dataset and unlabeled data from both datasets.
We evaluate each model on a hold-out set of In Distribution and OOD data.
Our first experiment is a small 1-dimensional regression problem where we can easily compute the ground truth (Figure 2).
For our high-dimensional regression task, we use a Superconductivity prediction task (Hamidieh, 2018), splitting the data into In Distribution and OOD based on target values (Table 1).
For our classification tasks, we use the MNIST and CIFAR data for In Distribution data and Fashion MNIST and SVHN for OOD data (Tables 2 and 3).
One desired property of posterior approximators is that they have low variance on in distribution data and high variance on OOD data.
Thus, along with test loss/accuracy, we report confidence interval width/correctness (for regression) and posterior entropy (for classification).
We also use posterior sample variance to construct ROC curves for OOD-detection and measure the area under the curve (AUC).
Finally, to evaluate scalability of our method, we trained increasingly larger ensembles and measured their OOD-detection performance vs. training time.
See Appendix D for full details.
While all methods are able to achieve high In Distribution performance (loss or accuracy), the GPN is able to consistently achieve the highest OOD-detection AUC.
The SNGP and OR-ensemble are also able to achieve high OOD-detection AUC, but have other drawbacks.
Specifically, SNGP has a very high CI-width on In Distribution data for the regression task and a very small contrast in sample entropy between the In Distribution and OOD data for the classification tasks.
And, while the
OR-ensemble performs well, Figure 3 shows GPN scales much better.
Conclusion
In this paper we introduce Generative Posterior Networks (GPNs), a method that uses unlabeled data to learn a generative model of the Bayesian posterior distribution.
We prove that under mild assumptions, GPNs approximate samples from the true posterior.
We then show empirically that our method significantly outperforms competing epistemic uncertainty predictions techniques on high-dimensional classification tasks and scales much better than ensembling methods, the closest performing baseline.
A Related Work
There are two primary methods of formalizing epistemic uncertainty: (1) as a Gaussian Process (GP) or (2) as a Bayesian Inference problem.
The GP framework is often used in practice, specifically in low-dimensional, low-data settings where the problem can be solved exactly (Hensman et al., 2013; Titsias, 2009).
Many GP approximation methods have been proposed to improve scalability (Wilson & Nickisch, 2015; Wilson et al., 2016), but each have their drawbacks.
We highlight Spectralnormalized Neural Gaussian Processes (SNGP)
(Liu et al., 2020) as one of the best GP approximation methods for high dimensional problems.
While SNGP shows impressive out of distribution prediction performance, we show in our experiments that the learned posterior does not adequately capture the true posterior.
The Bayesian Inference formulation, considers a different problem: given a prior distribution over functions and some labeled data, the goal is to construct a posterior distribution in the Bayesian sense.
Because this posterior distribution is usually intractable to compute exactly, approximate sampling methods are used instead (Tanner & Wong, 1987;
Blei et al., 2017).
For high-dimensional problems,
Bayesian Neural Networks (BNNs) (Blundell et al., 2015) and Bayesian Dropout (Blundell et al., 2015) are alternative methods for approximating Bayesian Inference on neural networks.
Neural network ensembling is a method that predicts out-of-distribution data well in practice, requires very little fine-tuning and, under the right regularization, approximates samples from the Bayesian posterior Pearce et al.
; He et al.
However, these methods have one main drawback, namely that each new sample from the posterior requires training a new network from scratch.
Our method, on the other hand, seeks to construct a generative posterior model using similar regularization techniques to Pearce et al.
(2020), allowing for quick sampling from the posterior.
Epistemic Neural
Networks Osband et al.
(2021) are a concurrent work to ours that, while not motivated by the
Bayesian Inference problem, arrived at a very similar technique to ours.
The key difference between our method and Epistemic NNs is in the regularization: our method requires unlabeled data from the test-distribution for our regularization, but results in improved posterior sampling.
B Algorithm details B.1 Low-Dimensional Embedding
As mentioned in the main paper, optimizing the generative function g over the the space of all anchor parameters θ anc is impractical, as this space is too high dimensional.
Instead, we seek to learn a low-dimensional representation of the anchor parameters.
There are two key reasons why we would expect to be able to decrease the representational power of the anchor distribution and maintain the same of fidelity in the posterior estimation.
(1) Because of the nature of neural networks, there are many settings of parameters θ that would result in the same output vector ˆY .
Because we are focusing on the posterior of ˆY , we need less representational power to model ˆY .
And, maybe more importantly, because the posterior parameters are highly correlated, we would expect to need less expressive power to represent the posterior distribution than the prior distribution.
Thus, we would like to construct a simpler estimate of the prior space using a low-dimensional embedding vector, z .
That is, we want our generative model to take as input the embedding vector z ∼ N (0 , I ) , instead of θ anc , in order to estimate the MAP function.
To do this, we need to learn a mapping from anchor parameters θ anc to embedding vectors z .
For our experiments, we used a 1-1 embedding scheme where we sample k parameters from the true prior θ 1 , . . .
, θ k ∼ P prior ( θ ) and k independent samples from our embedding z 1 , . . .
, z k .
We then jointly optimize over ϕ and z 1 , . . .
, z k as follows:
[Equation]
where δ j = g ( x [Equation] ) , x j sample ∼ P sample ( x ) , ϵ ∼ N (0 , I ) is a noise injection vector that improves the smoothness of interpolations between embedding vectors, and
L reg ( z 1 , . . . , z k ) is a regularizer to keep z 1 , . .
z k roughly normally distributed, which allows us to easily sample from the embedding space.
For our experiments we use the KL divergence between z
[Figure] [Figure]
and the normal distribution, L reg ( z 1 , . . .
, z k )
=
D KL ( N (¯ z , s z ) , N (0 , 1)) .
Figure 1 illustrates how we can sample from this embedding space to construct posterior functions.
B.2 Classification
Our method uses the assumption that the prior over ˆY is approximately normally distributed.
Of course, this approximation is only exact when either the network is linear or the final layer is infinitely wide and does not have an activation function (assuming the prior parameters are normally distributed).
Thus, in classification tasks, where it is beneficial to a soft-max to the final output, this assumption is violated.
To get around this, we add the anchor loss to the pre-softmax outputs for classification tasks.
C Additional Experiment Results
Figure 4 illustrates our methods ability to form credible posterior distributions.
Specifically, the figure shows two informative test images, one from the In Distribution dataset and one from the
OOD dataset, for both MNIST/Fashion MNIST and CIFAR-10/SVHN, along with box plots of sampled PMFs from each of the learned models.
While all methods predict the correct label with high precision on the in distribution example, ours has a uniquely wide sample distribution on the
OOD example, illustrating high predicted epistemic uncertainty.
Figures 5 show the full ROC curves for the OOD prediction tasks.
D Experiment Details
Small Scale Regression Our first experiment is a small 1-dimensional regression problem where we can easily compute the ground truth.
We provide each method with the same 6 observations.
We compute 100 samples from the approximate posterior for each method.
For ground truth, we use the
Metropolis-Hastings MCMC algorithm.
[Table]
[Table]
High-Dimensional Regression For a high-dimensional regression task
, we used a Superconductivity prediction dataset (Hamidieh, 2018), the goal of which is to predict the critical temperature of different superconductive materials based on 81 features.
For the In Distribution and OOD datasets, we split the full dataset based on the target values.
Specifically, the In Distribution and OOD datasets contained data with a target values in the intervals [13 . 9 , + ∞ ) and [0 , 13 . 9) , which is roughly 58% of the data in distribution and 42% out of distribution.
The unlabeled dataset we provide our method and the Output-Regularized ensemble method is the full training set.
Regularization hyper-parameters were chosen by running each method with 5 different settings, then choosing the model with both a low validation loss and high CI-correct on a validation set.
Classification We run two different classification experiments: for one, we use the MNIST as our In Distribution dataset and Fashion MNIST as the OOD dataset, and for the other, we use the
CIFAR-10 dataset as our In Distribution dataset and SVHN as the OOD dataset.
The unlabeled dataset we provide our method and the Output-Regularized ensemble method consists of unlabeled data 50% from the In Distribution and 50% from the OOD datasets.
As in the regression experiments, we ran each method with 5 different regularization parameters, then chose the model with both a high validation accuracy and OOD prediction performance on a validation set.
D.1 Baselines
We compare our method against 4 competing Bayesian methods: Bayesian Dropout (Gal & Ghahramani, 2016), GPs, and anchor-regularized neural-network ensembles (Pearce et al., 2020) with 10 ensemble members.
For the high-dimensional problems, we need to use approximate GPs.
Specifically, we use a grid-interpolated GP with Deep Kernel Learning (DKL) (Wilson et al., 2016) implemented with the GPyTorch library (Gardner et al., 2018) and the Spectral-normalized Neural
Gaussian Process (SNGP) method (Liu et al., 2020).
For the ensembles, since every member requires a unique set of parameters (and anchor points) the number of ensemble members was limited by the memory capacity of our GPUs.
As described in Equation 3, our method a method to sample unlabeled data points.
In low-dimensional problems, this can simply be a uniform sample from the interval of interest, for example [ − 2 , 2] in
Figure 2.
However, in high-dimensional problems, this requires access to an additional dataset of unlabeled training data.
From a practical perspective, this is a reasonable setup as many real-world datasets have far more unlabeled data than labeled data.
This additional data obviously favors our method over the baseline methods as our method is the only method that is able to take advantage of this unlabeled data.
To address this, we added an additional ensemble regularized using the same regularization as Equation 3; in other words, we regularize the outputs of the ensemble members towards outputs of sampled prior networks.
We denote this new method as an “Output-regularized (OR)”
Ensemble to distinguish it from the “Parameter-regularized” (PR) Ensemble.
We show that despite having the access to the same additional unlabeled data, our method outperforms and scales better than this OR Ensemble method.
Each model uses the same 4-layer (Superconductor), 5-layer (MNIST), or 7-layer (CIFAR) architecture, except the ensembles, which use slightly smaller networks for each ensemble member (to fit on a single GPU).
Full experimental details are in Appendix D.
D.2 Metrics
Because we are interested in epistemic uncertainty estimation, we cannot use common aleatoric uncertainty evaluation metrics, such as Expected Calibration Error (ECE) (Guo et al., 2017).
Instead, we evaluate each methods ability to construct useful posterior estimates using the following metrics.
OOD Detection One desired property of an epistemic uncertainty estimators is in predicting when a data point is outside the training distribution and, thus, any prediction made on such a data point will likely be incorrect.
For this reason, OOD detection is a common metric used in epistemic uncertainty literature (Liu et al., 2020; Van Amersfoort et al., 2020).
For an ideal posterior distribution, we expect high sample variance on OOD data and low sample variance from In Distribution data.
Thus, in our experiments, we use posterior sample variance to detect OOD data.
Specifically, for each model, we use posterior variance over 100 samples from each model as a score to construct ROC curves and measure the measure the area under the curve (AUC).
Confidence Intervals / Entropy Another desired property of an epistemic uncertainty estimator for a regression task is in constructing confidence intervals.
We would expect such confidence intervals (CIs) to contain the true class with high probability.
We also expect that, on in distribution examples, confidence intervals are narrow.
To construct these intervals, we take 100 samples from each model on every data point in the test data and compute the range of the middle 95% of the samples.
CI-correct refers to the percentage of true labels that fall in this interval.
Note that, unlike the other methods tested, ensembles are not a generative model.
Thus, when trying to sample from the ensemble method, we can only sample as many functions as there are members of the ensemble; for our experiments this number is 10.
For classification, CIs are not as informative.
Instead, we look at the posterior sample entropy.
For an ideal posterior distribution, we expect entropy to be low for data inside the training set (data the model should have confidence on) and high for data outside the training set (which we expect the model to be uncertain on).
To measure this sample entropy, we take 100 samples from each model on every data point and take the average of the post-softmax outputs.
We then average this sample entropy for both the In Distribution and OOD datasets.
Scalability One key benefit of a generative model over an ensemble is efficiency: every sample from an ensemble method requires learning an entire new network from scratch.
Generative models, on the other hand, are able to produce new samples with a single pass through the network.
To illustrate the benefit of this efficiency in practice, we ran an experiment to measure OOD detection AUC vs. computation time for both the Superconductor and CIFAR-10 tasks.
We trained 2 (Superconductor) and 5 (CIFAR) ensemble members at a time and computed the OOD-detection performance of the cumulative combined ensembles (both parameter-regularized and output-regularized).
Each method is trained using an Nvidia 1080TI.
Figure 3 shows our method is able to achieve high out of distribution prediction performance in significantly less computation time.
[Table]
D.3 Experiment datasets
Table 4 details the labeled and unlabeled data provided to the agent at training time and
the In
Distribution and OOD data used for evaluation at test time.
D.4 Network Architectures
For the superconductor regression problem, we used a 4-layer network with 4 128-unit wide linear layers for the ensemble methods and a 4-layer network with 4 512-unit wide linear layers for the other methods.
For the MNIST classification problem, we used a 5-layer network with 2 convolution layers and 3 128-unit wide linear layers for our the ensemble methods and a 5-layer network with 2 convolution layers and 3 512-unit wide linear layers for our the other methods.
For the CIFAR classification problem, we used a 7-layer network with 3 convolution layers and 4 256-unit wide linear layers for our the ensemble methods and a 7-layer network with 3 convolution layers and512-unit wide linear layers for our the other methods.
D.5 Hyper parameters
The non-network architecture hyper-parameters we used for GPN were roughly consistent across all experiments, except for the regularization coefficient, β .
As stated in our experiments section, we performed a small search over 5 values of β for each experiment.
The full set of hyper-parameters are shown in Table[-
-]
E Proofs E.1 Proof of Theorem 2.
Proof.
We will start by showing that P ( y obs | ˆY ) is normally distributed.
Let ˆ Y ′
i
=
f ( x i obs ; θ ) be a transformation of the random variable θ corresponding to the outputs of the function f parameterized by θ evaluated at observation points x obs .
Note that by definition of y obs , for each i ∈ 1 , . . .
, N :
[Equation]
where ϵ i ∼ N (0 , σ ϵ ) .
By assumption: [ ˆY , ˆY ′ ] T is normally distributed.
Since ϵ is normally distributed and independent of [ ˆY , ˆY ′ ] T , then the joint distribution [ ˆY , ˆY ′ , ϵ ] T is also normally distributed.
Now, by applying a linear transformation, we obtain the joint variable   ˆY ˆ Y ′ y obs 
 = (cid:34)
I
I I I (cid:35)   ˆY ˆ Y ′ ϵ   , which is also normally distributed (as it is a linear transformation of a Gaussian random variable).
Now, since the conditional and marginal distributions of a Gaussian distribution are also Gaussian,
P ( y obs , ˆY ′ | ˆY ) and P ( y obs | ˆY ) are also Gaussian.
Now, using Bayes’ rule, we can show P ( ˆY | y obs ) is also normally distributed.
[Equation]
for some mean and covariance µ like , Σ like .
Since the product of two multivariate Gaussians is a
Gaussian, we know that:
[Equation]
where
[Equation]
Now we consider at the distribution P ( ˆ Y MAP ( ˆY anc )) where ˆY anc ∼ N ( µ anc , Σ anc ) .
We will show that, when we set
[Equation]
the distribution P ( ˆ Y MAP ( ˆY anc )) becomes equal to the posterior distribution N ( µ post , Σ post ) .
There are three steps needed to show equality between these distributions:
Show that P ( ˆ Y MAP ( ˆY anc )) is normally distributed for some mean and variance, µ RMSpost , Σ RMSpost .
Show that µ RMSpost = µ post . 3.
Show that Σ RMSpost = Σ post .
Using the same reasoning as above, in the limit as M → ∞ ,
[Equation]
Since the max of a Gaussian is the mean, then
[Equation]
where we define:
[Equation]
We will now show that E [ ˆ Y MAP ( ˆY anc )] = µ post .
Because we set µ anc = µ ˆY :
[Equation]
Finally, we will show that V ar [ ˆ Y MAP ( ˆY anc )]
= Σ post .
[Equation]
So, P ( ˆ Y MAP ( ˆY anc )) ∼ N ( µ RMSpost , Σ RMSpost )
where µ RMSpost = µ post and Σ RMSpost = Σ post .
Thus, in the limit as M → ∞ , P ( ˆ Y MAP ( ˆY anc )) =
P ( ˆY | y obs ) .
E.2 Derivation of loss function
We start with
[Equation]
If we choose µ anc = µ ˆY and Σ anc =
Σ ˆY , we can simplify the logarithm of the likelihood and anchor distributions as follows:
[Equation]
where δ j = g ( x j sample , θ anc ; ϕ ) − f ( x j sample ; θ anc ) .
Putting these together, we get:
[Equation]