Single Model Uncertainty Estimation via Stochastic Data Centering
Anonymous Author(s)
AffiliationAddress email
Abstract
We are interested in estimating the uncertainties of deep neural networks, whichplay an important role in many scientific and engineering problems.
In this paper,we present a striking new finding that an ensemble of neural networks with the sameweight initialization, trained on datasets that are shifted by a constant bias gives riseto slightly inconsistent trained models, where the differences in predictions are astrong indicator of epistemic uncertainties.
Using the neural tangent kernel (NTK),we demonstrate that this phenomena occurs in part because the NTK is not shiftinvariant.
Since this is achieved via a trivial input transformation, we show that itcan therefore be approximated using just a single neural network – using a techniquethat we call ∆ − UQ – that estimates uncertainty around prediction by marginalizingout the effect of the biases.
We show that ∆ − UQ ’s uncertainty estimates aresuperior to many of the current methods on a variety of benchmarks– outlierrejection, calibration under distribution shift, and sequential design optimization ofblack box functions.
Introduction
Accurately estimating uncertainties in a deep neural network (DNN) is an active area of research dueto its implications in a wide range of scientific and engineering problems.
Broadly, there are twokinds of uncertainties that are often considered – (a) aleatoric : uncertainties in the data generatingprocess that are typically irreducible, and (b) epistemic : uncertainties of the model that can bereduced by observing more data, which is the focus of our work.
Some of the popular techniquesfor the latter include Bayesian methods [1–4] that use a prior on the network weights, Monte Carloapproximations such as MC Dropout [5] that approximate sampling from the posterior, and empiricalmethods such as Deep Ensembles (DEns) [6].
In particular, DEns trains an ensemble of neuralnetworks with different initializations, such that the uncertainty estimate on a test sample is given bythe inconsistency between predictions from the member models.
In practice, DEns has been found tooften outperform other related methods [7, 6, 8], but it comes at the cost of training several DNNs toobtain reliable uncertainties (typically 10 − 20 ), which is a severe computational bottleneck when itcomes to modern DNNs especially on large scale datasets.
In light of this, there is increased interestin developing single model estimators that can still produce high quality uncertainties.
There has beensome promising work in this direction in the recent past, specific to deep classifiers [8] or regressionmodels [9], and has been shown to perform comparably to DEns in some use-cases.
In this work, we first begin by exploring an alternate approach for constructing deep ensembles –instead of using multiple randomized weight initializations, we propose to shift the training domainusing random biases, such that every model in the ensemble is trained with data that has been shiftedby a different constant bias c .
Formally, for a labeled dataset { (x , y ) } , the k th model is trained to fit{ (x − c k , y ) } , where c k (referred as an anchor) is of the same size as x .
Though this manipulation appears trivial, the kernel induced by deep networks is not inherently shift invariant[10, 11], thusimplying each DNN learns a slightly different model due to the bias.
This leads to one of our keyobservations:
Anchor Ensemble: When an ensemble of DNNs, with the same fixed initialization,are trained on a dataset shifted by random constant biases, the variation acrossthe ensemble’s predictions is a strong indicator of model uncertainty.
Based on analysis with the neural tangent kernel (NTK) [10], we show that when the anchor c ismade a random variable, the effective kernel is stochastic such that for each c , the model convergesto a slightly different kernel.
Consequently, this anchoring-based ensembling provides a differentapproach to DEns for sampling solutions from the hypothesis space.
However, interestingly, theproposed ensembling lends itself easily to be approximated using a single DNN:
∆ − UQ
: For a random anchor c , 1 can be approximated using a single DNNtrained on the dataset transformed as { x , y } → { [c , x − c] , y } .
During inference, we obtain multiple predictions for a given sample by varying the choice of theanchor, such that the standard deviation of the predictions is our estimate for uncertainty.
Withoutaffecting the performance of the model, we find that ∆ − UQ produces meaningful epistemic uncertainties that we validate in a variety of applications: outlier rejection, calibration under distributionshift on ImageNet, and sequential optimization of a large suite of black-box functions.
We observethat ∆ − UQ consistently outperforms existing uncertainty estimates, while also being efficient totrain as a single model estimator.
With just a few simple changes, one can easily modify the trainingof any existing DNN model to support ∆ − UQ estimation.
We include a Torch implementation in thesupplement, and will make our code public for easy benchmarking.
Background and Related Work
Denote training data as D = { (x i , y i ) } ni
=1 , where x i ∈ X and y i ∈ Y , to train a neural networkf ( θ ) ∈ H with randomly initialized weights θ 0 , such that a loss function L is minimized, i.e. ,arg min θ L ( f (x; θ ) , y ) .
Here, H denotes the hypothesis space of potential solutions for fitting theobserved data.
Given a prior distribution on the weights p ( θ ) , we can define the posterior over θ asp ( θ |D ) and subsequently, infer the posterior predictive distribution for a test sample (x t , y t ) .
This canbe used to quantify the uncertainty around the prediction as p ( y t | x t , D ) =
(cid:82) θ p ( y t | x t , θ ) p ( θ |D )
d θ .
A challenge for neural networks, however, is that the posterior p ( θ |D ) is computationally intractable.
This has motivated the use of Bayesian Neural Networks (BNNs) [3] using different approximationsto the posterior including Monte-Carlo Dropout [5], variational inference [12, 4], and samplingmethods such as Markov chain Monte Carlo [3, 13].
In parallel, it has been empirically shownthat Deep Ensembles (DEns) [6] often tend to outperform Bayesian methods in terms of modelcalibration performance, even under challenging distribution shifts [7].
The success of DEns hasbeen attributed to its ability to sample different functional modes [14] from the hypothesis space, andthus approximate the posterior predictive distribution [1].
However, a critical limitation of DEns isthe need to train multiple models (typically 10 − 20 ) in order to obtain well calibrated uncertainties,which can be impractical for complex model architectures that have become commonplace today.
Characterizing the behavior of deep uncertainty estimators has mostly been done using empiricalevaluation based on model calibration or out-of-distribution detection, but the recent advances in theneural tangent kernel (NTK) theory [10, 15–17] provide a convenient framework for more rigorousanalysis.
The basic idea of NTK is that, when the width of a neural network tends to infinity andthe learning rate of SGD tends to zero, the function f (x; θ ) converges to a solution obtained bykernel regression using the NTK defined as K x i x j = E θ (cid:68)
∂f
(x i , θ ) ∂ θ , ∂f
[Equation]
x j
[Equation]
) ∂
[Equation]
(cid:69) .
When the samplesx i , x j ∈ S d − 1 , i.e., points on the hypersphere and have unit norm, the NTK for a simple 2 layer
ReLU MLP can be simplified as a dot product kernel [15–17]:
[Equation]
The NTK framework and its extensions that enable a posterior interpretation in the infinite limit havebeen used to study deep ensembles [2].
Previous work [18–20] has also shown the existence of a distinct yet related kernel, referred to as the neural network Gaussian process (NNGP) kernel, wherethe initialization tends to a GP in infinite width limit.
Uncertainty estimation with ∆ − UQ
As discussed above, most existing frameworks sample the hypothesis space either through differentrandom initializations of ( θ 0 ), perturbing the weight space after training, or using Bayesian methodsby placing a prior on p ( θ ) .
Here, we propose to use a new strategy that involves injecting multipletrivial biases into the training data and analyze the resulting models using the NTK framework.
3.1 Anchor Ensembles: ensembling by injecting trivial biases
Let us examine the scenario where we shift an entire dataset (both train and validation) using aconstant bias, c , to obtain a new dataset D c using which we train the model f c .
Since we alwayschoose c from the training distribution at random, this has the effect of zero-centering the datasetaround different training points.
Finally, we are interested in the relationship between the models{ f c 1 , f c 2 , . . .
f c k } .
If the NTK induced by f is shift-invariant (for e.g., when Fourier features [11]are used), the shifts make no difference, resulting in identical models f c 1 = · · · = f c k .
However,since NTKs for models like MLPs and CNNs are not inherently shift-invariant [17], we find that themodels lead to an effective deep ensemble, wherein the variation across the predictions is a strongindicator of epistemic uncertainties.
Effect of shifted training on NTK:
We are interested in understanding how (1) changes when theentire training domain is shifted by c – i.e., h NTK ((x i − c) ⊤ (x j − c)) .
Without loss of generality,for the sake of notational convenience, we assume x i − c
and x j − c are also made unit norm.
Tosimplify the expansion, we use a Taylor series expansion for the cos − 1 function: cos − 1 (u − c) ≈cos − 1 (u) + c √
[Equation]
Expanding (x i − c) ⊤ (x j − c) as x ⊤ i x j − c ⊤ (x i + x j − c) and letting v = (x i + x j − c) , we obtainthe expression for h NTK under a shifted domain as follows:
[Equation]
where we combine all terms dependent on c into Γ x i , x j , c , which also behaves as a dot product kernel.
From (2), we note that a trivial shift in the domain results in a non-trivial shift in the NTK functionitself.
In other words, (2) outlines the effective NTK as a function of c .
We also note that Γ does notaffect the spectral properties of the original NTK, as we observe in Figure 1.
Let us now consider the prediction on a test sample x t in the limit as the inner layer widths grow toinfinity.
It has been shown that (c.f.
[Equation]
where X is the matrix of all training data samples.
As before, we consider the case where the domainis shifted by c .
Using (3):
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Figure]
In (4), we utilize Woodbury’s Inverse Identity [21] to expand the inverse of a sum of matrices.
Next,in (5), we combine all the terms dependent on c into a function g .
Note, we also expand f 0 (x − c)using the Taylor series approximation to represent it as a sum of f 0 (x) and other terms (details of thisderivation can be found in the supplement).
For a given initialization θ 0 , the first term is deterministic, and exactly the same as (3), since all otherthe terms are fixed.
However, we see that the second term can vary based on the choice of c .
Existingensembling approaches [6] rely on the randomness of the initialization θ 0 to pick diverse solutionsfrom the hypothesis space [2].
In contrast, equations (2) and (5) suggest that even for a fixed θ 0 , it ispossible to make the NTK stochastic (in c ) by shifting the entire input domain.
To better understandhow the NTK actually changes from these expressions, we study the spectral properties of a simple
MLP network empirically, following the analysis in [11].
Spectral properties of shifted NTKs.
We compute the Fourier spectra using the same MLP onseveral shifted domains in Figure 1(B).
The original spectrum for the MLP without any shift in thetraining domain is shown for comparison in 1(A).
As indicated by (2), we see that each individualshift leads to a different NTK (as indicated by the spectra), either by flattening it or making it narrowerin the frequency domain.
The same behavior persists even when we construct positional embeddings(PE), based on sinusoidal functions, prior to building the MLP model (1(D-E)).
This leads to one ofour main findings stated in 1 , and illustrated in Figure 3.
3.2 ∆ − UQ : rolling anchor ensembles into a single model
Since different models in an anchoring-based ensemble are trained with the same initialization, wepresent a new technique to approximate the epistemic uncertainties using a single neural network.
More specifically, we perform a simple coordinate transformation by lifting the domain to a higherdimension as E : x → { c , x − c } , we refer to the residual by ∆ =
x − c .
This transformationallows the use of multiple representations (w.r.t.
different anchors) for the same input sample x ,i.e. , f ∆ ( { c 1 , x − c 1 } ) =
f ∆ ( { c 2 , x − c 2 } ) = · · · = f ∆ ( { c k , x − c k } ) , where f ∆ refers to the∆ − UQ model that takes the tuple ( { c k , x − c k } ) and predicts the target y .
It is easy to see that [c , x i − c]
⊤ [c , x j − c] =
x ⊤ i x j − c ⊤ (x i + x j − 2c) .
That is, the dot product ofthe transformed inputs takes the same form as before (except for a scaling factor).
Therefore, theexpressions for the equivalent NTK for different anchor shifts, seen in (2), and the correspondingprediction on a test sample seen in (5) remain the same for ∆ − UQ , by setting v = x i + x j − 2c .
This leads to our main claim stated in 2 , that the ∆ − UQ model achieves similar perturbations ofthe NTK as an anchor ensemble, based on the choice of c .
Training.
During training, for every input x i we choose an anchor as random sample from the trainingdataset.
Subsequently, we obtain the coordinate transformation { [c , x i − c ] , y i } , using which we train
[Figure]
the model.
With vector-valued data, this is implemented as a simple concatenation.
In the case ofimage data, we append the channels to create a 6 − dimensional tensor (for a 3-channel RGB image).
We show simple a Pytorch snippet for training ∆ − UQ in Figure 2.
Other than increasing the numberof parameters in the first layer of the network, ∆ − UQ does not incur additional computationaloverheads.
for
A
[Figure] [Figure]
[Figure]
D =
[Figure]
− A %% Delta
X_d = torch.cat([A, D],axis=1)
y_d = model(X_d)
%% prediction loss =
criterion(y_d,targets)
Over the course of training, every training pair getscombined with a large number of anchors.
Since theprediction on this training pair – regardless of anchorchoice – must always be the same, this places a consistency in the predictions that they must be similarno matter which anchor is chosen.
This consistencytrades-off with diversity of the kinds of functions thatcan be learned when compared with an anchor ensemble, where the models are trained independently.
Thiscan be seen in the comparisons of the NTK spectra for ∆ − UQ with anchor ensembles in Figures1(C) and (F).
In practice, however, we find that the diversity from this single model is still sufficientlylarge, to estimate good quality uncertainties.
Inference.
For a test sample x t , we obtain the prediction from ∆ − UQ as the mean predictionacross several randomly chosen anchors; and the standard deviation around these predictions is ourestimate for the epistemic uncertainty.
In other words, we marginalize out the effect of anchors toobtain the final prediction mean and uncertainty.
Formally, the predictive distribution is given byp ( y t | x t ) =
(cid:82) c ∈ X p ( y t | x t , c , θ ) p (c)
d c .
In practice, for a trained ∆ − UQ model specified as θ ∗ , wecompute the sample mean and uncertainty around it as:
[Equation]
Discussion.
In Figure 3, we show an 1 D regression example using 20 training examples alongwith the predicted mean and estimated uncertainties.
As it can be seen, both the anchoring-basedensemble (left) and ∆ − UQ training show higher epistemic uncertainties around regions with notraining samples.
In the ensembles version, we train 20 different networks – while in ∆ − UQ we usejust a single network, where the uncertainty is obtained using all 20 anchors during inference.
We always only use a single anchor for an input during every training iteration, but multiple anchorsduring inference time.
In theory, multiple anchors could be used during training as well, where theloss is imposed on the mean (obtained with multiple anchors).
However, we find that simply using asingle random anchor in each iteration achieves a similar effect, as it enforces a consistency that thesame training pair (x , y ) when combined with many different anchors over the course of training as[c 1 , x i − c 1 ] , [c 2 , x i − c 2 ] , . . .
, [c k , x i − c k ] which must all produce the same prediction, y .∆ − UQ relies on randomly drawn anchors for uncertainty estimation, which is similar to DEns[6], that relies on the diversity of the base learners in an ensemble.
Arguably, sampling a random set of anchors from the training distribution is simpler than sampling from the posterior p ( θ |D ) .
Furthermore, since every anchor realizes a slightly different function, using a small number ofanchors ( 5 − 10 ) during inference is typically sufficient to obtain high quality estimates as we showin our experiments.
Finally, due to the nature of training with random anchors we also see that∆ − UQ produces particularly effective uncertainty estimates when the training set size is small, andthis proves to be very useful in applications like sequential optimization.
Experiments
We validate our approach in this section using a variety of applications and benchmarks – (a) first, weconsider the utility of epistemic uncertainties in object recognition problems where they have beensuccessfully used for outlier rejection and calibrating models under distribution shifts.
We show that∆ − UQ can be very effective even with large-scale datasets like ImageNet [22]; (b) next, we considerthe challenging problem of sequential design optimization of black-box functions, where the goalis to maximize a scalar function of interest with the fewest number of sample evaluations.
Using a
Bayesian optimization setup, we show that the uncertainties obtained using ∆ − UQ outperform manycompetitive methods across an extensive suite of black-box functions.a.
Outlier rejection.
A popular application for epistemic uncertainties is in rejecting outlierssince, by definition, they are in regions outside of the training distribution.
As such, we expect themodel to produce highly uncertain predictions for these images, which should help us design aneffective OOD detector.
To evaluate this hypothesis, we train a modified ResNet-50 [23] model on
ImageNet that accepts 6 input channels (anchor, ∆ ) as outlined earlier.
We train the model usingstandard hyperparameter settings except, training it longer for 120 epochs – our top-1 accuracy ( 76 . 1 )matches that of a standard ResNet-50.
Specifically for images, we found that corrupting the anchorswith common transforms like random crops, Gaussian blurs, and color jitter improves performance.
That is, instead of [ c, x − c ] , we use [ T ( c ) , x − c ] where T is a transform like Gaussian blur.
Wedescribe this process in more detail in the supplement.
The uncertainty for a test sample is given bythe standard deviation of the logits obtained by varying the anchors.
To obtain a scalar statistic, wecompute the mean across all classes as the uncertainty score for that sample.
We show the results for (a) Uncertainties change meaningfully as out- liers get more severe
[Table]
[Figure]
214 outlier rejection in Table 4b, where we follow the protocol established in [25], that uses a Gaussianblur of intensity 5 from ImageNet-C [26] as the outlier set, and the clean ImageNet validation data asinliers.
We use the estimated uncertainty obtained with 10 anchors as in (6) as our score for outlierrejection and report commonly used metrics such as AUROC, Detection Accuracy (DTACC), and
AUPR-in/out.
We note that, just the inconsistency of predictions obtained using ∆ − UQ outperformsmany baselines including mean-field stochastic variational inference (SVI) [12, 4], SVI-AvUC [25],
Monte Carlo dropout [5], and temp.
scaling [24].
While this can be further improved by taking themean prediction into account, similar to existing approaches for semantic novelty detection (withscores such as entropy, energy [27] etc.), our focus here is to evaluate the quality of uncertainty alone.
In Figure 4a, we observe that ∆ − UQ ’s uncertainty estimate changes smoothly as the outliers becomefarther away (more severe intensity) from the training distribution.
[Table]
b. Calibration under distribution shifts.
Following our observation that our uncertainty estimatesare effective in rejecting outliers in table 4b, here we study if they can be leveraged to calibrate
ImageNet models under distribution shift.
To calibrate a classifier, we simply scale the logits of themean by the uncertainties as follows: µ calib.
= µ (1 − ¯ σ ) , where ¯ σ is the standard deviation estimatedfrom (6), where once again for classification we simply compute the average of standard deviationacross all the 1000 classes, followed by min-max normalization to [0 , 1) .
This simple scaling ofthe mean reflects our prior belief – a highly certain prediction must remain unchanged, whereasan uncertain one gets tempered down.
Note, this scaling is applied to logits from all classes andhence the accuracy of the mean remains unchanged before or after calibration.
We evaluate howcalibrated the predictions are using three commonly used metrics – Calibration error (ECE), negativelog likelihood (NLL), and Brier Score.
We use the same ResNet-50 classifier trained on ImageNet asbefore, and measure calibration for predictions on 16 different ImageNet-C corruptions at severity5 .
We list the corruptions and show examples of them in the supplement.
We report the 25 th (lowerquartile), 50 th (median) and 75 th (upper) quantiles along with the mean of the three metrics across16 corruptions in table 1.
Across all measures we see that ∆ − UQ is able to calibrate models better,even in comparison to state-of-the-art approaches that use an explicit calibration objective to adjustthe prediction probabilities, further validating the quality of its uncertainty estimates.
c. Sequential Optimization.
Denoting a high-dimensional function as f : D →
R , our goalis to solve the following optimization problem: x ∗ = arg max x ∈D f(x) .
Here, D refers to abounded design space comprising D different parameters with their corresponding value ranges[ ℓ d , h d ] , ∀ d = 1 · · · D .
The high computational or financial cost of evaluating f (invoking a simulatoror running an experiment) motivates the additional objective of minimizing the number of evaluations.
Given the high-dimensional nature of the design spaces, a simple brute-force search or even spacefilling random sample designs [28] often require significantly large sample sizes to identify the optima,thus motivating the use of sequential optimization strategies.
In particular, Bayesian Optimization(BO) techniques based on statistical surrogates (e.g., Gaussian processes) form an important classof solutions [29].
In a nutshell, given an initial experiment design and their function evaluations,{ x i , f (x i ) } n 0 i =0 , sequential optimization techniques incrementally select candidates to achieve theso-called exploration-exploitation trade-off using an appropriate acquisition function [30].
In thisstudy, we use the popular expected improvement (EI) score to perform candidate selection.
Setup :
In this experiment, we consider a large suite of black-box optimization functions with varyingdimensionality ( 1 to 8 ) and complexity to comprehensively evaluate different epistemic uncertaintyestimation techniques with deep neural network surrogates (see Appendix 1 for description ofthe functions used in our study).
Finally, we perform an experiment with a pre-trained generative model (GAN) trained on MNIST handwritten digits, wherein we perform optimization in the100 − D latent space Z such that thickness of the resulting digit is maximized: (cid:80)
i I ( x i > 0) , ∀ i ,where I denotes the identity function.
We use the following baseline uncertainty estimation approaches in our study: (i) Gaussian processes (GP); (ii) Monte-Carlo dropout (MCD); (iii) Bayesian
[Table]
12 ± 0 .0 .
16 ± 0 .0 .
11 ± 0 .0 .
13 ± 0 .0 .
47 ± 0 .
Ackley Function
Dim. = Dim. =
Dim. =
Griewank Function
Dim. =
Dim. = Dim. =
Levy Function
Dim. = Dim. =
Dim.
= 8 DIM =
DIM = DIM =
GP
MCD
BNN
DEns
Ours neural networks (BNN) trained via variational inferencing; and (iv) deep ensembles (DEns).
Forall neural network surrogates, we computed positional embeddings (sinusoidal) of the raw parameter inputs prior to building a fully-connected network with 4 hidden layers each containing128 neurons and ReLU activation.
All methods were trained with the same set of hyperparameters: Adam optimizer learning rate 1 e − 4 and 500 epochs, except for BNN, which required 1000epochs for convergence.
With MCD, we used 50 forward passes at test time for each sampleto obtain the uncertainties.
Finally, with ∆ − UQ, we set the number of anchors for inferencing as min(20 , n ) , where n is the number of samples in the observed dataset in any iteration.
[Figure] [Figure] The DEns model was constructed using 5 constituent members(increasing this did not provide any benefits), each trained witha different initialization. The number of initial samples andthe number of steps in the sequential optimization were set tobe the same across all methods.
In each round of the Bayesianoptimization, we used 10 , 000 samples for initialization and 15restarts (i.e., starting points for multistart acquisition functionoptimization), and finally one candidate ( q = 1 ) was evaluatedwith the black-box function and added to the observed dataset(similar trends were observed even with batch sampling q =2 ).
For every function, we performed experiments with 5random seeds (different initializations), each for 5 independenttrials.
Since the goal is to reach the global optima (knownfor benchmark functions) with the fewest number of samples,we use the widely adopted area under the iteration vs bestachieved function value curve to obtain a holistic evaluation of different approaches (see Figure 6).
Results : From table 2, we find that ∆ − UQ produces significantly higher AUC scores in comparisonto existing baselines, across all benchmark functions.
While MCD and DEns behave reasonablywell in low dimensions, their performance suffers when we go to higher dimensions (see figure 5).
Furthermore, we find that the performance of BNN is generally lower due to the inherent smallsamples sizes that we operate in.
Finally, as showed in figure 7, our approach consistently achieveshigher function values in the MNIST GAN-based optimization, thus validating the quality of theuncertainties produced via anchoring.
Broader Impact
We presented a simple, scalable, and accurate single model uncertainty estimator that outperformsmany existing techniques.
Uncertainty quantification (UQ) plays a significant role in a variety ofapplications in science and engineering, from developing safeguards in critical applications such ashealthcare, security, or finance to exploration of novel design spaces using sequential optimization.
Our paper addresses both of these applications using demonstrative problems to showcase the potentialcapabilities.
Our goal is to present a fundamental approach for uncertainty estimation that maybeapplicable in different kinds of problems and settings.
Due to its fundamental nature, we do notforesee misuse directly.
However, we note that our method does not mitigate biases that may comewith the training data, which could potentially affect the quality of uncertainties themselves.