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.
[Figure]
[Table]
AUROC ↑ DTACC ↑ AUPR-in/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.
Checklist
For all authors... (a) Do the main claims made in the abstract and introduction accurately reflect the paper’s contributions and scope?
[Yes] (b) Did you describe the limitations of your work?
[Yes] (c) Did you discuss any potential negative societal impacts of your work?
[Yes] (d) Have you read the ethics review guidelines and ensured that your paper conforms to them?
[Yes] 2.
If you are including theoretical results... (a) Did you state the full set of assumptions of all theoretical results?
[Yes] (b) Did you include complete proofs of all theoretical results?
[Yes] Main claims are presented in the paper, further derivations are given in the supplement 3.
If you ran experiments... (a) Did you include the code, data, and instructions needed to reproduce the main experimental results (either in the supplemental material or as a URL)?
[Yes] (b) Did you specify all the training details (e
g., data splits, hyperparameters, how they were chosen)?
[Yes] (c) Did you report error bars (e.g., with respect to the random seed after running experiments multiple times)?
[Yes] (d) Did you include the total amount of compute and the type of resources used (e
g., type of GPUs, internal cluster, or cloud provider)?
[N/A] 4.
If you are using existing assets (e
g., code, data, models) or curating/releasing new assets... (a) If your work uses existing assets, did you cite the creators?
[Yes] (b) Did you mention the license of the assets?
[N/A] We are using standard open sourced datasets and models (c) Did you include any new assets either in the supplemental material or as a URL?
[No] (d) Did you discuss whether
and how consent was obtained from people whose data you’re using/curating?
[N/A] (e) Did you discuss whether the data you are using/curating contains personally identifiable information or offensive content?
[N/A] 5.
If you used crowdsourcing or conducted research with human subjects... (a) Did you include the full text of instructions given to participants and screenshots, if applicable?
[N/A] (b) Did you describe any potential participant risks, with links to Institutional Review Board (IRB) approvals, if applicable?
[N/A] (c) Did you include the estimated hourly wage paid to participants and the total amount spent on participant compensation?
[N/A]
APPENDIX A Derivation for shifted training on NTK
We continue the derivation from the main here in more detail.
Recall, 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]
where X is the matrix of all training data samples.
As before, we consider the case where the domain is shifted by c .
Using (7):
[Equation]
Where we utilize Woodbury’s Identity [21] for expanding the inverse of the difference between two matrices as:
[Equation]
Using (9)
[Equation]
we can expand (8) as:
[Equation]
Next
[Equation]
we consider expanding the first term in (11).
Since the only term dependent on c is the evaluation of the network with the initial weights θ 0 ,
i
[Equation]
e., of the general form f 0 (x − c) .
We will expand this using a Taylor series approximation by evaluating it at c = 0 , as following:
[Equation]
By substituting (12) in (11)
[Equation]
and grouping all the terms that do not depend on c
[Equation]
we can separate the deterministic and stochastic (in c ) which gives us our final result as:
[Equation]
Perturbations with different anchors The analysis above can be easily applied to the case where different anchors are used with different input samples.
Let us consider two randomly chosen anchors: c 1 , c 2 and study the dot product between two points shifted using these anchors.
[Equation]
Where in the last step we exploit the fact that c 1 , c 2 are normalized to be on the hypersphere.
We can see that by setting v = c
1 c ⊤ 2 x 1 + x 2 − 2c 2 , we can get a similar form of perturbation as x ⊤ 1 x 2 − c ⊤ 1 v by combining all c 2 related terms as before (and equivalently for c 1 ).
In this paper, we argue that the proposed stochastic data centering technique is effective at estimating epistemic uncertainties with deep networks.
For demonstration, let us consider the 1 D regression example showed in Figure 8 and train ∆ − UQ models under different train sample sizes ( 5 , 10 ,and 200 respectively).
The figure illustrates the predicted function and the associated uncertainty estimates (shaded region around the predictions).
As the training sample size increases, we notice that the uncertainties shrink to trivial values (very close to 0), thus validating that our estimates are strongly correlated with epistemic uncertainties.
B Corruptions in the anchoring process
When scaling ∆ − UQ to image data, especially using powerful base networks such as ResNets, we observe a consistency training significantly improves the uncertainty estimates.
To achieve this, we make the input transformation less trivial as x → T (c) , x − c , where T is a standard augmentation technique already used in training such as random crops or blurs.
Such that, during inference we set
T =
I , to be identity.
We apply this corruption once every 10 iterations – this is a hyper parameter, doing this more frequently makes the process much harder resulting in a worse mean estimate whereas making it less frequently results in uncertainties that are slightly worse.
We outline the exact set of corruptions used in the pytorch pseudo code listed below.
B.1 Pytorch implementation of our approach
Algorithm 1 lists the Pytorch pseduo-code for implementing ∆ − UQ for image classification.
In algorithm ??
we show pseudo code for inference using ∆ − UQ .
The example assumes a model as defined in algorithm 1 and a set of anchors drawn from the training distribution at random.
Once predictions per anchor are obtained, the mean and standard deviation are returned as the final prediction of the model, and the corresponding uncertainty on the test samples.
B.2 ImageNet-C corruptions for OOD and Calibration
Table 3 lists the set of corruptions used to construct the ImageNet-C benchmark.
[Table]
Algorithm 1 PyTorch-style example for ∆ − UQ with ResNet-50.
def create_anchored_model(model):model.conv1=nn.
Conv2d(in_channels=6, 64) return model
Tx = transforms.
Compose([transforms.
RandomResizedCrop(size=224), transforms.
RandomHorizontalFlip(),transforms.
RandomApply([color_jitter,blurr], p=0.8), ])
## load model and change the first conv layer model_basic = ResNet50(pre_trained=
False,n_class=1000) model = create_anchored_model(model_basic)
## load datasets , setup optimizer , define criterion etc.
for i, (images, targets) in enumerate(train_loader): anchors = Shuffle(images) diff = images-anchors if i % 10 ==0: tx_anchors = Tx(anchors) else:tx_anchors = anchors batch =
torch.cat([tx_anchors ,diff],axis=1) output = model(batch) loss =
criterion(output, target) optimizer.zero_grad() loss.backward() optimizer.step()
[Table]
C Additional Results: Prediction Performance on UCI Benchmarks
While our outlier rejection, calibration and sequential optimization experiments clearly established the effectiveness of the proposed uncertainty estimator, we also evaluate the quality of ∆ − UQ models, in terms of standard prediction fidelity metrics (regression in this section and classification in next).
For this study, we used a suite of regression datasets typically adopted for evaluating deep models, evaluated using the standard experiment protocol in the benchmark defined by [32].
For each of the datasets, we fit networks with a single hidden layer ( 50 neurons) and ReLU activation.
We trainedindependent models with different random 80 − 20 train-test splits and report the average performance across the trials.
For evaluation, we used the negative log-likelihood metric (lower the better).
In addition to our approach, we include the results for MCD, DEns, BNN (variational inferencing) and Probabilistic Backpropagation [31] (with a Matrix-Variate Gaussian prior).
Furthermore, for a holistic evaluation, we also report the average rank (across the 5 methods) from the suite of datasets considered.
As showed in Table 4, ∆ − UQ performs competitively over other baselines, and achieves an average rank of 2 . 0 .
Overall, we find that, in addition to producing high-quality uncertainty estimates, the proposed approach also produces high-quality predictive models.
[Table]
D Additional Results: Prediction Performance on Imagenet and CIFAR-
ImageNet-C accuracy.
We provide results for classification accuracy of our ImageNet model on the validation set and the distribution shifted variants of ImageNet-C, in table 5.
Here, at each severity level (“1”–“5”)
we compute the accuracy of the model across all 16 corruptions for that severity outlined in 3 and report the mean accuracy.
We also report the accuracy numbers for the corresponding uncertainty baselines, and see that ∆ − UQ does not compromise on accuracy on the clean data, while being highly competitive to Deep Ensembles even on the most severe corruptions. and its distribution shifted variants.
[Table]
CIFAR-10C/
ResNet-20 We perform detailed analysis of calibration and accuracy on CIFARand its corrupted variants CIFAR-10C [26] using the experimental protocol followed by [7, 25], where we use a ResNet-20 [23] and report the calibration scores across all 5 corruption levels and the validation set – the calibration metrics are reported by averaging the performance across 5 random seeds of the model.
We report the average accuracy for each corruption level in table 6, and display the calibration metrics – ECE, NLL and Brier Score in figure 9.
In both the accuracy and calibration metrics we find that ∆ − UQ outperforms all the comparable baselines, including Deep Ensembles, though using only a single model.
Ablation studies.
We mainly perform ablation on the inference part of ∆ − UQ – we try to quantify the variability in performance when anchoring is not used, and when uncertainty is not used in terms of the calibration metrics for the ImageNet-C/ResNet-50 experiment considered in the main paper.
We consider two main ablations of the main model as stated next.
Note, in all three cases the only [Table] [Figure] Note, we calibrate our predictions as before using a scaling strategy as: ¯ µ = µ (1 − σ ) .
We compare against standard baselines obtained from [25, 7].
difference is the procedure for inference.
The training procedure is kept fixed across all models, and we use the same ResNet-50 model to perform these ablation studies reported in table 7.
• Naïve : We consider the case where a model trained with anchoring as usual, is used during inference without anchoring , i.e., instead of passing { c , x − c } as before for an anchor c ,
[Table]
we pass { 0 , x } , as this behaves as a naive model that does not have the benefit of obtaining uncertainties or ensembling like behavior as ∆ − UQ .
• Ensemble mean : Next, we consider the version where anchoring is done during inference and compute the mean of predictions from different anchors as before, but we do not use the uncertainties obtained – i.e., the final prediction is simply the mean of the predictions obtained with different anchors.
• ∆ − UQ: This is our final model that takes the mean and scales it by the uncertainties during inference.
We observe from the results in table 7 first that simply training with anchoring shows benefits in model performance even if anchoring is not used during inference, as seen in improvement in the calibration performance of the naïve model (shown as { 0 , x } in the table) over the vanilla model.
Next, we see that using anchoring to compute the mean improves performance further as seen in the next column (shown as µ ), and finally factoring in the uncertainties performs the best.
Even the ablated versions perform competitively compared to some of the other uncertainty baselines, indicating the effectiveness of ∆ − UQ .
E Details on GAN-based Optimization Experiment
In this experiment, we evaluated the utility of the proposed uncertainty estimator in guiding sequential optimization in the latent space of a pre-trained GAN network.
We begin by assuming access to a generative model G(z) , which maps a latent noise vector z onto a realization on the training image manifold.
Denoting the latent space as Z , our goal is to maximize a scalar function defined for an image, i.e., f (x) by performing optimization in the latent space.
[Equation]
In our experiment, we used a GAN trained on MNIST hand-written images and defined the thickness function (total number of non-zero pixels in an image) for optimization.
The dimensionality of the noise latent space was set to 100 .
Similar to our design optimization experiments with synthetic data, we started with an initial random sample (uniform random in the latent space) of 25 , synthesized the corresponding images using the generator and computed the thickness function for each of them.
We performed optimization for 50 steps (with 1 sample in each round) and evaluated the maximum thickness achieved as the metric of choice (the global optimum is not known).
We repeated the experiment across 5 random seeds and 5 independent trials for each seed.
The results in Figure 7 (main paper) illustrate the maximum thickness obtained using different uncertainty estimators across 25 experiments.
We find that our approach consistently produced the highest function value and outperformed other approaches.
As expected, DEns performed the second best, followed by
MCD.
Interestingly, with BNNs, the variational inferencing technique is known to lead to underfitting and we find out, despite achieving reasonably higher function values, the reconstructions were of significantly poorer quality (off the image manifold).
F Benchmark Functions
Figure 10 illustrates the different benchmark functions for evaluating the proposed approach in blackbox optimization.
In addition to the functions listed, we also considered the Hartmann functions, in dimensions 3 and 6 respectively, defined as follows.
[Equation] A ij ( x [Equation]
− P ij [Equation] (cid:19) , where α = (1 . 0 , 1 . 2 , 3 . 0 , 3 . 2) ⊤
A [Equation]   3 . 0 {+
 [Equation]+}
[Equation]
10  ,
P [Equation] [Equation]   3689 11704699 43871091 8732 5743  (16)
Hartmann6: f (x [Equation] = 4 (cid:88) i =
[Equation] exp (cid:18) − 6 (cid:88)
j =
A ij ( x j − P ij )(cid:19) , where α = (1 . 0 , 1 . 2 , 3 . 0 , 3 . 2) ⊤
A =     3 .
1  ,
P [Equation] [Equation]   1312 1696 82832329 4135 8307 3736 10042348 1451 3511 2883 30474047 8828 8732 5743   (17)
G Detailed Results for Sequential Optimization
At the core of AI-powered applications in science and engineering lies the need to perform design optimization for maximizing a chosen target objective, and to enable automated exploration in high-dimensional parameter spaces.
When f is Lipschitz continuous, i.e., ∥ f (x) − f (x ′ )
∥
≤ c ∥ x − x ′ ∥ , and its first- and second-order information are accessible, this can be solved using first-order optimization methods such as stochastic gradient descent (SGD) [33] or second-order methods such as L-BFGS [34].
However, in practice, while f can be explicitly evaluated for any x , its first- and second-order information are unknown, thus making such an optimization very challenging.
Commonly referred to as black-box optimization [35], this formulation is adopted in applications ranging from drug design [36] to additive manufacturing [37] and optimizing financial investments [38] to hyper-parameter tuning in neural networks [39].
[Equation]
In each step of this optimization, we approximate the function f
using the samples observed so far to obtain the surrogate ˆ f .
Assuming that our goal is to maximize f , one can acquire more samples in regimes of D where the mean estimate from the surrogate is high ( exploitation ) or the uncertainty is large ( exploration ).
In order to balance between these two objectives and to guide the progressive search for the optima, BO utilizes an appropriate acquisition function [30].
In this study, we use the popular expected improvement (EI) score to perform candidate selection.
[Equation]
where Z = ( µ (x) − f (x [Equation] ξ )
[Equation] [Equation]
Function
Bounds
[Figure]
[Figure] [Figure] [Figure] Here, µ (x) and σ (x) are the mean and uncertainty estimates from the surrogate ˆ f for any sample, and f (x + ) is the best known function value so far during any iteration of the optimization.
Further, Φ( . ) and ϕ ( . ) denote the cumulative distribution and probability density functions corresponding to the normal distribution.
Finally, the hyper-parameter ξ controls the exploration-exploitation trade-off.
Figure 11 shows the convergence curves for each of the black-box functions.
[Figure]