Single Model Uncertainty Estimation via Stochastic Data Centering
Abstract
We are interested in estimating the uncertainties of deep neural networks, which play 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 same weight initialization, trained on datasets that are shifted by a constant bias gives rise to slightly inconsistent trained models, where the differences in predictions are a strong indicator of epistemic uncertainties.
Using the neural tangent kernel (NTK), we demonstrate that this phenomena occurs in part because the NTK is not shift-invariant.
Since this is achieved via a trivial input transformation, we show that this behavior can therefore be approximated by training a single neural network – using a technique that we call ∆ − UQ – that estimates uncertainty around prediction by marginalizing out the effect of the biases during inference.
We show that ∆ − UQ ’s uncertainty estimates are superior to many of the current methods on a variety of benchmarks– outlier rejection, calibration under distribution shift, and sequential design optimization of black box functions.
Code for ∆ − UQ can be accessed at github.com/LLNL/DeltaUQ
Introduction
Accurately estimating uncertainties in a deep neural network (DNN) is an active area of research due to its implications in a wide range of scientific and engineering problems.
Broadly, there are two kinds of uncertainties that are often considered – (a) aleatoric : uncertainties in the data generating process that are typically irreducible, and (b) epistemic : uncertainties of the model that can be reduced by observing more data, which is the focus of our work.
Some of the popular techniques for the latter include Bayesian methods [1–4] that use a prior on the network weights, Monte Carlo approximations such as MC Dropout [5] that approximate sampling from the posterior, and empirical methods such as Deep Ensembles (DEns) [6].
In particular, DEns trains an ensemble of neural networks with different initializations, such that the uncertainty estimate on a test sample is given by the inconsistency between predictions from the member models.
In practice, DEns has been found to often outperform other related methods [7, 6, 8], but it comes at the cost of training several DNNs to obtain reliable uncertainties (typically 10 − 20 ), which is a severe computational bottleneck when it comes to modern DNNs especially on large scale datasets.
In light of this, there is increased interest in developing single model estimators that can still produce high quality uncertainties.
There has been some promising work in this direction in the recent past, specific to deep classifiers [8] or regression models [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 domain using random biases, such that every model in the ensemble is trained with data that has been shifted by 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], thus implying each DNN learns a slightly different model due to the bias.
This leads to one of our key observations:
Anchor Ensemble: When an ensemble of DNNs, with the same fixed initialization, is trained on a dataset shifted by random constant biases, the variation across the 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 is made a random variable, the corresponding kernel is stochastic such that for each c , the model converges to a slightly different NTK.
Consequently, this anchoring-based ensembling provides a different approach to DEns for sampling solutions from the hypothesis space.
However, interestingly, since the anchor ensembling emerges primarily from the act of a trivial data transformation, it lends itself to be approximated easily using a single DNN:
∆ − UQ : For a randomly chosen anchor c , 1 can be approximated using a single DNN trained 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 the anchor, such that the standard deviation of the predictions is our estimate for uncertainty.
Without affecting 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 distribution shifts on ImageNet, and sequential optimization of a large suite of black-box functions.
We observe that ∆ − UQ consistently outperforms existing uncertainty estimates, while also being efficient to train as a single model estimator.
With just a few simple changes, one can easily modify the training of any existing DNN to support ∆ − UQ estimation.
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 network f ( θ ) ∈ 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 the observed data.
Given a prior distribution on the weights p ( θ ) , we can define the posterior over θ as p ( θ |D ) and subsequently, infer the posterior predictive distribution for a test sample (x t , y t ) .
This can be 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 approximations to the posterior including Monte-Carlo Dropout [5], variational inference [12, 4], and sampling methods such as Markov chain Monte Carlo [3, 13].
In parallel, it has been empirically shown that Deep Ensembles (DEns) [6] often tend to outperform Bayesian methods in terms of model calibration performance, even under challenging distribution shifts [7].
The success of DEns has been attributed to its ability to sample different functional modes [14] from the hypothesis space, and thus approximate the posterior predictive distribution [1].
However, a critical limitation of DEns is the 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 empirical evaluation based on model calibration or out-of-distribution detection, but the recent advances in the neural tangent kernel (NTK) theory [10, 15–17] provide a convenient framework for more rigorous analysis.
The basic idea of NTK is that, when the width of a neural network tends to infinity and the learning rate of stochastic gradient descent (SGD) tends to zero, the function f (x; θ ) converges to a solution obtained by kernel regression using the NTK defined as K x i x
[Equation]
When the samples x 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 have been 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, where the 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 different random initializations of ( θ 0 ), perturbing the weight space after training, or using Bayesian methods by placing a prior on p ( θ ) .
Here, we propose to use a new strategy that involves injecting multiple trivial 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 a constant bias, c , to obtain a new dataset D c using which we train the model f c .
Since we always choose c from the training distribution at random, this has the effect of zero-centering the dataset around different training points.
Let { f c 1 , f c 2 , . . .
, f c k } denote the set of models trained using different c ’s, then our goal is to characterize the relationship between them as a function of c .
If the NTK induced by f is shift-invariant (for e.g., when Fourier features [11] are used), the shifts will 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] 2 , we find that the models lead to an effective deep ensemble, wherein the variation across the predictions is a strong indicator of epistemic uncertainties.
Effect of shifted training on NTK:
We are interested in understanding how (1) changes when the entire 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.
To simplify the expansion, we use a Taylor series expansion for the cos − 1 function: cos − 1 (u − c)
≈ cos − 1 (u)
+ c √ 1 − (u − c) 2 .
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 obtain the 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 function itself.
In other words, (2) outlines the effective NTK as a function of c .
We also note that Γ does not affect 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 to infinity.
It has been shown that (c.f.
[Equation]
[Figure]
where X is the matrix of all training data samples.
As before, we consider the case where the domain is shifted by c .
Using (3):
[Equation] [Equation]
In (4)
[Equation]
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 this derivation can be found in the appendix).
For a given initialization θ 0 , the first term is deterministic, and exactly the same as (3), since all other the terms are fixed.
However, we see that the second term can vary based on the choice of c .
Existing ensembling approaches [6] rely on the randomness of the initialization θ 0 to pick diverse solutions from the hypothesis space [2].
In contrast, equations (2) and (5) suggest that even for a fixed θ 0 , it is possible to make the NTK stochastic (in c ) by shifting the entire input domain.
To better understand how 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 on several shifted domains in Figure 1(B).
The original spectrum for the MLP without any shift in the training domain is shown for comparison in 1(A).
As indicated by (2), we see that each individual shift leads to a different NTK (as indicated by the spectra), either by flattening it or making it narrower in 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 of our 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, we present 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 higher dimension as E : x → { c , x − c } , we refer to the residual by ∆ =
x − c .
This transformation
[Figure]
allows 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 of the transformed inputs takes the same form as before (except for a scaling factor).
Therefore, the expressions for the equivalent NTK for different anchor shifts, seen in (2), and the corresponding prediction on a test sample seen in (5) remain the same for ∆ − UQ , by setting v = x i + x j − 2c .
A similar form is also obtained in the more general case when different anchors are used for different samples (as shown in the appendix), which leads to our main claim stated in 2 , that the ∆ − UQ model achieves similar perturbations of the NTK as an anchor ensemble, based on the choice of c .
Training.
During training, for every input x i we choose an anchor as a random sample from the training dataset.
Subsequently, we obtain the coordinate transformation { [c , x i − c ] , y i } , using which we train the model.
With vector-valued data, this is implemented as a simple concatenation.
In the case of image 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 number of parameters in the first layer of the network, ∆ − UQ does not incur additional computational overheads.
y loss =
criterion(y
,targets)
Over the course of training, in expectation, every training pair gets combined with a large number of anchors.
Since the prediction on this training pair – regardless of anchor choice – must always be the same, this places an implicit consistency in the predictions that they must be similar no matter which anchor is chosen.
This consistency trades-off with diversity of the kinds of functions that can be learned when compared with an anchor ensemble, where the models are trained independently.
This can be seen in the comparisons of the NTK spectra for ∆ − UQ with anchor ensembles in Figures 1(C) and (F).
In practice, however, we find that the diversity from this single model is still sufficiently large, to estimate good quality uncertainties.
Inference.
For a test sample x t , we obtain the prediction from ∆ − UQ as the mean prediction across several randomly chosen anchors; and the standard deviation around these predictions is our estimate for the epistemic uncertainty.
In other words, we marginalize out the effect of anchors to obtain the final prediction mean and uncertainty.
Formally, the predictive distribution is given by p ( 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 θ ∗ , we compute the sample mean and uncertainty around it as:
[Equation]
Discussion.
In Figure 3, we show a 1 D regression example using 20 training examples along with the predicted mean and estimated uncertainties.
As it can be seen, both the anchoring-based ensemble (left) and ∆ − UQ training show higher epistemic uncertainties around regions with no training samples.
In the ensembles version, we train 20 different networks – while in ∆ − UQ we use just a single network, where the uncertainty is obtained using all 20 anchors during inference.
We always use only a single anchor for an input during every training iteration, but multiple anchors during inference time.
In theory, multiple anchors could be used during training as well, where the loss is imposed on the mean (obtained with multiple anchors).
However, we find that simply using a single random anchor in each iteration achieves a similar effect, as it enforces a consistency that the same 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 ] 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 of anchors ( 10 − 20 ) during inference is typically sufficient to obtain high quality estimates as we show in 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, and this proves to be very useful in applications such as sequential optimization.
Experiments
We validate our approach in this section using a variety of applications and benchmarks – (a) first, we consider the utility of epistemic uncertainties in object recognition problems where they have been successfully 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 consider the challenging problem of sequential design optimization of black-box functions, where the goal is 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 many competitive methods across an extensive suite of black-box functions.
a. Outlier rejection.
A popular application for epistemic uncertainties is in rejecting outliers since, by definition, they are in regions outside of the training distribution.
As such, we expect the model to produce highly uncertain predictions for these images, which should help us design
an effective
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 using standard 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 anchors with 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 simple transform like Gaussian blur.
The uncertainty for a test sample is given by the standard deviation of the logits obtained by varying the anchors.
To obtain a scalar statistic, we compute the mean across all classes as the uncertainty score for that sample.
We provide more details on the transformations used for corrupting the anchor, and the accuracies of our model across shifted variants in more detail in the appendix.
We show the results for outlier rejection in 4b, where we follow the protocol established in [25], that uses a Gaussian blur of intensity 5 from ImageNet-C [26] as the outlier set, and the clean
ImageNet validation data as inliers.
We use the estimated uncertainty obtained with 10 anchors as in (6) as our score for outlier rejection 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 outperforms many 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 the mean prediction into account, similar to existing approaches for semantic novelty detection (with scores 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 become farther away (more severe intensity) from the training distribution.
[Table] [Figure] [Table] AUROC
[Table]
DTACC ↑ AUPR-in
[Table]
out ↑ [Table] [Figure]
[Table]
b.
Calibration
under
distribution shifts.
Following our observation that our uncertainty estimates are effective in rejecting outliers in 4b, here we study if they can be leveraged to calibrate ImageNet models under distribution shifts.
To calibrate a classifier, we simply scale the logits of the mean by the uncertainties as follows: µ calib.
= µ (1 − ¯ σ ) , where ¯ σ is the standard deviation estimated from (6), where once again for classification we simply compute the average of standard deviation across all the 1000 classes, followed by min-max normalization to [0 , 1) .
This simple scaling of the mean reflects our prior belief – a highly certain prediction must remain unchanged, whereas an uncertain one gets tempered down.
Note, this scaling is applied to logits from all classes and hence the accuracy of the mean remains unchanged before or after calibration.
We evaluate how calibrated the predictions are using three commonly used metrics – Calibration error (ECE), negative log likelihood (NLL), and Brier Score.
We use the same ResNet-50 classifier trained on ImageNet as before, and measure calibration for predictions on 16 different ImageNet-C corruptions at severity 5 .
We list the corruptions and show examples of them in the appendix.
We report the 25 th (lower quartile), 50 th (median) and 75 th (upper quartile) along with the mean of the three metrics acrosscorruptions 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 adjust the prediction probabilities, further validating the quality of its uncertainty estimates.
c. Sequential Optimization.
Denoting a high-dimensional function as f : D → R , our goal is to solve the following optimization problem:
x ∗ = arg max x ∈D f(x) .
Here, D refers to a bounded 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 simulator or running an experiment) motivates the additional objective of minimizing the number of evaluations.
[Table]
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 class of 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 the so-called exploration-exploitation trade-off using an appropriate acquisition function [30].
In this study, we use the popular expected improvement (EI) score to perform candidate selection.
[Figure] [Figure]
Setup :
In this experiment, we consider a large suite of blackbox optimization functions with varying dimensionality ( 1 to 8 ) and complexity.
We provide definitions and details of these functions used in our study in the appendix.
We also perform an experiment with a pre-trained generative model (GAN) trained on MNIST handwritten digits, wherein we perform optimization in the 100 − 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 neural networks (BNN) trained via variational inferencing; and (iv) deep ensembles (DEns).
For all 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 containing 128 neurons and ReLU activation.
All methods were trained with the same set of hyperparameters: Adam optimizer learning rate 1 e −and 500 epochs, except for BNN, which required 1000 epochs for convergence.
With MCD, we used
50 forward passes at test time for each sample to 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]
Dim
[Figure]
Dim
[Figure]
Dim
[Figure]
Dim
[Figure]
Dim
[Figure]
DIM
[Figure] [Figure]
The DEns model was constructed using 5 constituent members (increasing this did not provide any benefits), each trained with a different initialization.
The number of initial samples and the number of steps in the sequential optimization were set to be the same across all methods.
In each round of the
Bayesian optimization, we used 10 , 000 samples for initialization and 15 restarts (i.e., starting points for multistart acquisition function optimization), and finally one candidate ( q = 1 ) was evaluated with the black-box function and added to the observed dataset.
We performed experiments with 5 random seeds (different initializations), each for 5 independent trials.
Since the goal is to reach the global optima with the fewest number of samples, we use the widely adopted area under the iteration vs best achieved function value curve to obtain a holistic evaluation of different approaches (see Figure 5).
Results : From table 2, we find that ∆ − UQ produces significantly higher AUC scores in comparison to existing baselines, across all benchmark functions.
While MCD and DEns behave reasonably well in low dimensions, their performance suffers when we go to higher dimensions (see figure 6).
Furthermore, we find that the performance of BNN is generally lower due to the inherent small samples sizes that we operate in.
Finally, as shown in figure 7, our approach consistently achieves higher function values in the MNIST GAN-based optimization, thus validating the quality of the uncertainties produced via anchoring.
Broader Impact
We presented a simple, scalable, and accurate single model uncertainty estimator that outperforms many existing techniques.
Uncertainty quantification (UQ) plays a significant role in a variety of applications in science and engineering, from developing safeguards in critical applications such as healthcare, 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 potential capabilities.
Our goal is to present a fundamental approach for uncertainty estimation that maybe applicable in different kinds of problems and settings.
Due to its fundamental nature, we do not foresee misuse directly.
However, we note that our method does not mitigate biases that may come with the training data, which could potentially affect the quality of uncertainties themselves.
Acknowledgement
This work was performed under the auspices of the U.S. Department of Energy by Lawrence
Livermore National Laboratory under Contract DE-AC52-07NA27344.
Supported by the LDRD
Program under projects 21-ERD-028, 22-ERD-006; with IM release number LLNL-JRNL-836221.
Appendix
The appendix contains the following additional information: (a) expanded derivations, (b) PyTorch code snippets for ∆ − UQ training & inference; (c) results on calibration in standard regression settings on the UCI benchmarks dataset; (d) results on accuracy of predictions on ImageNet and CIFARand their corrupted variants; (e) calibration results on CIFAR10-C; (f) ablations on anchor-based training; (f) details on the design optimization experiment, including definitions for the benchmark functions and convergence plots.
[-  3 . 0 
-]