Year: 2021

Abstract:
Uncertainty estimation in deep learning has recently emerged as a crucial area of interest to advance reliability and robustness in safety-critical applications. While there have been many proposed methods that either focus on distance-aware model uncertainties for out-of-distribution detection or on input-dependent label uncertainties for in-distribution calibration, both of these types of uncertainty are often necessary. In this work, we propose the HetSNGP method for jointly modeling the model and data uncertainty. We show that our proposed model affords a favorable combination between these two complementary types of uncertainty and thus outperforms the baseline methods on some challenging out-of-distribution datasets, including CIFAR-100C, Imagenet-C, and Imagenet-A. Moreover, we propose HetSNGP Ensemble, an ensembled version of our method which provides an additional type of uncertainty and also outperforms other ensemble baselines.

heading: 1 INTRODUCTION
text: While deep learning has led to impressive advances in predictive accuracy, models often still suffer from overconfidence and ill-calibrated uncertainties (Ovadia et al., 2019). This is particularly problematic in safety-critical applications (e.g., healthcare, autonomous driving), where uncertainty estimation is crucial to ensure reliability and robustness (Filos et al., 2019; Dusenberry et al., 2020).
Predictive uncertainties generally come in two flavors: model uncertainty (also known as epistemic) and data uncertainty (also known as aleatoric) (Murphy, 2012). Model uncertainty measures how confident the model should be based on what it knows about the world, that is, how much it can know about certain test data given the training data it has seen. Data uncertainty measures the uncertainty intrinsic in the data itself, for example due to fundamental noise in the labelling process. Good model uncertainty estimates allow for out-of-distribution (OOD) detection, that is, for recognizing data examples that are substantially different from the training data. On the other hand, good data uncertainty estimates allow for in-distribution calibration, that is, knowing which training (or testing) data examples the model should be more or less confident about.
Many recently proposed models for uncertainty estimation excel at one or the other of these uncertainty types. For instance, the spectral-normalized Gaussian process (SNGP) (Liu et al., 2020) uses a latent Gaussian process to achieve distance-aware model uncertainties and thus affords excellent OOD detection. Conversely, the heteroscedastic classification method (Collier et al., 2020; 2021) offers excellent calibration and improved accuracy thanks to modeling input- and class-dependent label noise in the training data. However, there have not been many attempts to combine the complementary benefits of these two types of uncertainty modeling (e.g., Depeweg et al., 2018) (see related work in Appendix A.1).
In this work, we propose the heteroscedastic SNGP (HetSNGP) model, which allows for joint modeling of model and data uncertainties based on a hierarchy of two latent random variables. We show that HetSNGP gives good in-distribution and OOD accuracy and calibration, yielding a model with uncertainties suitable for deployment in critical applications.
Our main contributions are:
• We propose a new model, the heteroscedastic spectral-normalized Gaussian process (HetSNGP), which provides both distance-aware model and data uncertainties.
• We describe an efficient approximate inference scheme that allows training HetSNGP with a computational budget that is comparable to standard neural network training.
• We show empirically on different benchmark datasets that HetSNGP offers a favorable combination of model and data uncertainty. It generally preserves the SNGP’s OOD performance and the heteroscedastic in-distribution performance and even outperforms these baselines on some datasets, where OOD and heteroscedastic uncertainties are helpful.
• We propose an ensembled version of our model, the HetSNGP Ensemble, which additionally accounts for parameter uncertainty and outperforms other ensemble baselines.

heading: 2 BACKGROUND

heading: 2.1 MODEL UNCERTAINTY AND SPECTRAL-NORMALIZED GAUSSIAN PROCESS
text: Model uncertainty (or epistemic uncertainty) captures all the uncertainty about whether a given model is correctly specified for a certain task, given the training data. Before any training data has been observed, this uncertainty depends only on the prior knowledge about the task, which can for instance be encoded into distributions about the model parameters or into the architecture of the model (i.e., the model class) (Fortuin, 2021). After training data has been observed, one should expect the model uncertainty to decrease within the support of the training data distribution, that is, on points that are close to the training data in the input space. We will generally call these points in-distribution (ID). Conversely, on points that are far away from the training points and thus out-ofdistribution (OOD), we should not expect the model uncertainty to decrease, since the training data points are not informative enough to make any assertions about the correctness of the given model on these points. In this sense, the optimal model uncertainty should be distance-aware, that is, it should grow away from the training data (Liu et al., 2020).
The spectral-normalized Gaussian process (SNGP) (Liu et al., 2020) provides such distance-aware model uncertainties by specifying a Gaussian process prior (Rasmussen & Williams, 2006) over the latent data representations in the penultimate layer of the neural network. Distance-awareness is ensured by using spectral normalization on the hidden layers (Behrmann et al., 2019), which encourages bi-Lipschitzness of the mapping from data to latent space, constraining the model from collapsing any input feature dimensions in the latent representations and approximately preserving the distances between data points in the latent space (Liu et al., 2020). Note that this approach only partially captures the model uncertainty, namely in form of the uncertainty over the latents. It does not however capture the uncertainty over the model parameters, such as for instance Bayesian neural networks (MacKay, 1992; Neal, 1993) or ensemble methods (Lakshminarayanan et al., 2016).

heading: 2.2 DATA UNCERTAINTY AND THE HETEROSCEDASTIC METHOD
text: As opposed to the model uncertainty described above, data uncertainty is intrinsic in the data and thus irreducible with increasing amounts of data. In the case of continuous data (e.g., regression problems), data uncertainty often comes in the form of random noise on the measurements. For discrete data (e.g., classification problems), it usually comes as label noise, that is, a certain number of data points will be assigned the wrong label in the training data. This label noise can be class- and input-dependent (Beyer et al., 2020). For instance, the Imagenet dataset (Deng et al., 2009), contains 100 different classes of dog breeds, which are often harder to tell apart for human labelers than the other classes in the dataset. Modeling this type of data uncertainty can improve the calibration and robustness of predictive models (Collier et al., 2020).
A model that does explicitly handle the input- and class-dependent label noise is the heteroscedastic method (Collier et al., 2021). The heteroscedastic method models input- and class-dependent noise by introducing a latent multivariate Gaussian distribution on the softmax logits of a standard neural network classifier. The covariance matrix of this latent distribution is a function of the input (heteroscedastic) and models inter-class correlations in the logit noise.

heading: 3 METHOD

heading: 3.1 SETUP AND NOTATION
text: Let us consider a dataset D = {(xi, yi)}Ni=1 of input-output pairs, where xi ∈ Rd and yi ∈ {1, . . . ,K}, that is, a classification problem with K classes. The data examples are assumed to be sampled i.i.d. from some true data-generating distribution as (xi, yi) ∼ p∗(x, y).

heading: 3.2 GENERATIVE PROCESS
text: To jointly model the two different types of uncertainty, we propose a hierarchical model of two latent random variables, which we denote by f and u. f is a latent function value associated to the input x (as in the Gaussian process literature) and is designed to capture the model uncertainty, while u is a latent vector of logits (or utilities) that capture the data uncertainty, similar to the setup in Collier et al. (2020), which was inspired by the econometrics literature (Train, 2009). Similarly to Liu et al. (2020), we place a latent Gaussian process (GP) prior over f , as p(f) = GP(0, kθ(·, ·)), where kθ is a parameterized kernel function with parameters θ. Evaluating this kernel on all pairwise combinations of datapoints in x yields the kernel matrix Kθ(x,x)). We then define u as a (heteroscedastically) noisy observation of f .
Stacking the variables across the whole dataset gives us the matrices F ,U ∈ RN×K . We refer to their respective rows as fi,ui ∈ RK and their columns as fc,uc ∈ RN . The columns are independent under the GP prior, but the rows are not. Conversely, the columns are correlated in the heteroscedastic noise model, while the rows are independent. A hierarchical model using both uncertainties therefore leads to logits that are correlated across both data points and classes.
The full generative process is
fc ∼ N (0,Kθ(x,x)) (1) ui ∼ N (fi,Σ(xi;ϕ)) (2)
p(yi = c |ui) = 1 [ c = arg max
k uik
] (3)
The predictive can then be computed my marginalizing, that is p(y |x) = Eu[p(y |u)] =∫ p(y |u) p(u |x) du. Intuitively, f captures the model uncertainty, that is, the uncertainty about the functional mapping between x and y on the level of the latents. It uses the covariance between data points to achieve this distance-awareness, namely it uses the kernel function to assess the similarity between data points, yielding an uncertainty estimate that grows away from the data. On the other hand, u captures the data uncertainty, by explicitly modelling the per-class uncertainty on the level of the logits. ϕ can learn to encode correlations in the noise between different classes (e.g., the dog breeds in Imagenet). It does not itself capture the model uncertainty, but inherits it through its
hierarchical dependence on f , such that the resulting probability p(y |u) ultimately jointly captures both types of uncertainty.
In practice, we usually learn the kernel using deep kernel learning (Wilson et al., 2016), that is, we define it as the RBF kernel kθ(xi,xj) = kRBF (hi,hj) = exp(‖hi − hj‖22/λ) with length scale λ and hi = h(xi; θ) ∈ Rm with h(·; θ) being a neural network model (e.g., a ResNet) parameterized by θ. This kernel is then shared between classes. Moreover, following Liu et al. (2020), we typically encourage bi-Lipschitzness of h using spectral normalization (Behrmann et al., 2019), which then leads to an approximate preservation of distances between the input space and latent space, thus allowing for distance-aware uncertainty modeling in the latent space. Additionally, Σ(·;ϕ) is usually also a neural network parameterized by ϕ. To ease notation, we will typically drop these parameters in the following. Note that the prior over fc (and thus also uc) has zero mean, thus leading to a uniform output distribution away from the data. Note that this is also reminiscent of multi-task GPs (Williams et al., 2007), where separate kernels are used to model covariances between data points and tasks and are then combined into a Kronecker structure.

heading: 3.3 COMPUTATIONAL APPROXIMATIONS
text: The generative process above requires computing two integrals (i.e., over u and f ) and inferring an exact GP posterior, making it computationally intractable. We make several approximations to ensure tractability.
Low-rank approximation. The heteroscedastic covariance matrix Σ(xi;ϕ) is a K × K matrix which is a function of the input. Parameterizing this full matrix is costly in terms of computation, parameter count, and memory. Therefore, following Collier et al. (2021), we make a low-rank approximation Σ(xi;ϕ) = V (xi)V (xi)> + d2(xi). Here, V (xi) is a K ×R matrix with R << K and d2(xi) is a K dimensional vector added to the diagonal of V (xi)V (xi)>. This approximation assumes that the similarities between different classes in the data, which in turn lead to nonzero covariance entries, can be approximately described as being linear in a suitable R-dimensional space. Note that R can be tuned as a hyperparameter and can thus be made large enough, such that this assumption holds.
Also following Collier et al. (2021), we introduce a parameter-efficient version of our method to enable scaling to problems with many classes, for instance, Imagenet21k which has 21,843 classes. In the standard version of our method, V (xi) is parameterized as an affine transformation of the shared representation hi. In this way, HetSNGP can be added as a single output layer on top of a base network. For the parameter-efficient HetSNGP, we parameterize V (xi) = v(xi)1>R V where v(xi) is a vector of dimension R, 1R is a vector of ones of dimension R, V is a K × R matrix of free parameters, and denotes element-wise multiplication. This additional approximation assumes that the similarity between classes is largely data-point-independent (i.e., different classes are intrinsically similar), while the specific data point can be responsible for scaling certain dimensions of similarity in the R-dimensional space.
Random Fourier feature approximation. Computing the exact GP posterior requires O(N3) operations (because ones needs to invert anN×N kernel matrix), which can be prohibitive for large datasets. Following Liu et al. (2020), we thus use a random Fourier feature (RFF) approximation (Rahimi et al., 2007), leading to a low-rank approximation of the kernel matrix asKθ(x,x) = ΦΦ>
(Φ ∈ RN×m), with random features Φi = √ 2 m cos(Whi + b) where W is a fixed weight matrix with entries sampled from N (0, 1) and b is a fixed bias vector with entries sampled from U(0, 2π). This approximates the infinite-dimensional reproducing kernel Hilbert space (RKHS) of the RBF kernel with a subspace spanned by m randomly sampled basis functions. This reduces the GP inference complexity to O(Nm2), where m is the dimensionality of the latent space. We can then write the model as a linear model in this feature space, namely
ui = fi + d(xi) K + V (xi) R with fc = Φβc and (4) p(βc) = N (0, Im); p( K) = N (0, IK); p( R) = N (0, IR)
Here, Φβc is a linear regressor in the space of the random Fourier features and the other two terms of ui are the low-rank approximation to the heteroscedastic uncertainties (as described above). Again,
m can be tuned as a hyperparameter to trade off computational accuracy with fidelity of the model. Since most data sets contain a lot of redundant data points, one often finds an m N that models the similarities in the data faithfully.
Algorithm 1 HetSNGP training
Require: dataset D = {(xi, yi)}Ni=1 Initialize θ, ϕ,W , b, Σ̂ for train step = 1 to max step do
Take minibatch (Xi,yi) from D for s = 1 to S do sK ∼ N (0, IK), sR ∼ N (0, IR) usi,c = Φ > i β̂c + d(xi) sK +
V (xi) s R
end for L=− 1S ∑S s=1 log p(Xi,yi |us)−‖β̂‖2
Update {θ, ϕ, β̂} via SGD on L if final epoch then
Compute {Σ̂−1c }Kc=1 as per Eq. (5) end if
end for
Algorithm 2 HetSNGP prediction Require: test example x∗
Φ∗ = √
2 m cos(Wh(x ∗) + b)
for s = 1 to S do βsc ∼ p(βc | D) sK ∼ N (0, IK), sR ∼ N (0, IR) us∗,c = Φ ∗>βsc + d(x ∗) sK +
V (x∗) sR end for p(y∗=c |x∗)= 1S ∑S s=1 exp(us∗,c/τ)∑K k=1 exp(u s ∗,k/τ) Predict y∗ = arg maxc p(y ∗ = c |x∗)
Laplace approximation. When using a Gaussian likelihood (i.e., in a regression setting), the GP posterior inference can be performed in closed form. However, for classification problems this is not possible, because the Categorical likelihood used is not conjugate to the Gaussian prior. We thus need to resort to approximate posterior inference. Again following Liu et al. (2020), we perform a Laplace approximation (Rasmussen & Williams, 2006) to the RFF-GP posterior, which yields the closed-form approximate posterior for βc
p(βc | D) = N (β̂c, Σ̂c) with Σ̂−1c = Im + N∑ i=1 pi,c(1− pi,c)ΦiΦ>i (5)
where pi,c is shorthand for the softmax output p(yi = c | ûi) as defined in Eq. (3), where ûi,c = Φ>i β̂c. The derivation of this is deferred to Appendix A.3. Here, Σ̂ −1 c can be cheaply computed over minibatches of data by virtue of being a sum over data points. Moreover, β̂c is the MAP solution, which can be obtained using gradient descent on the unnormalized log posterior, − log p(β | D) ∝ − log p(D |β)− ‖β‖2, where the squared norm regularizer stems from the standard Gaussian prior on β. Using our likelihood, the training objective takes the form of a ridge-regularized cross-entropy loss. We also use this objective to train the other trainable parameters θ andϕ of our model. Note that this approximation is necessarily ignoring multi-modality of the posterior, since it can by definition only fit one mode around the MAP solution. Nonetheless, recent works have shown that it can achieve suprisingly good performance in many settings and sometimes even be competitive with the true posterior (Immer et al., 2021b;a).
Monte Carlo approximation. Finally, we are ultimately interested in Eu[p(y |u)] =∫ p(y |u) p(u |x) du, which again would require solving a high-dimensional integral and is thus computationally intractable. Following Collier et al. (2020), we approximate it using Monte Carlo samples. Moreover, we approximate the argmax in Eq. (3) with a softmax. To allow for a more controlled trade-off between bias and variance in this approximation, we introduce an additional temperature parameter τ into this softmax, leading to
p(yi = c |xi) = 1
S S∑ s=1 p(yi = c |usi ) = 1 S S∑ s=1 exp(usi,c/τ)∑K k=1 exp(u s i,k/τ)
(6)
with usi,c = Φ > i β s c + d(xi) sK + V (xi) sR and βsc ∼ p(βc | D) and
sK ∼ N (0, IK) and sR = N (0, IR)
Train labels
Deterministic (acc = 0.855 ± 0.004)
SNGP (acc = 0.855 ± 0.001)
Heteroscedastic (acc = 0.868 ± 0.002)
HetSNGP (acc = 0.867 ± 0.002)
Figure 2: Predicted labels on synthetic data with label noise. We see that the proposed method performs on par with the heteroscedastic method and that they both outperform the other baselines thanks to their label noise modeling capabilities.
It should be noted that we only use Eq. (6) at test time, while during training we replace βsc with its mean β̂c instead of sampling it (see Section 3.4). Note also that, while a larger number S of Monte Carlo samples increases the computational cost, these samples can generally be computed in parallel on modern hardware, such that the wallclock runtime stays roughly the same. Again, the number of samples S can be tuned as a hyperparameter, such that this approximation can be made arbitrarily exact with increased runtime (or, as mentioned, increased parallel computational power).

heading: 3.4 IMPLEMENTATION
text: The training of our proposed model is described in Algorithm 1 and the prediction in Algorithm 2.
Intuition. The distance-aware uncertainties of the SNGP are modeled through Φ, which itself depends on the latent representations h, which are approximately distance-preserving thanks to the bi-Lipschitzness. This means that when we have an input xfar that is far away from the training data, the values of the RBF kernel exp(‖h− hfar‖22/λ) will be small, and thus the uncertainty of fc will be large. The resulting output uncertainties will therefore be large regardless of the heteroscedastic variance Σ(xi) (since they are additive) allowing for effective OOD detection. Conversely, if we have an in-distribution input, the kernel values will be large and the model uncertainty captured in fc will be small. In this case, the additive output uncertainty will be dominated by the heteroscedastic variance Σ(xi), allowing for effective label noise modeling in-distribution.
While at first it might seem straightforward to combine the two types of uncertainty, it should be noted that the proposed hierarchical model is by no means an arbitrary choice. Rather, it is the only design we found that robustly achieves a disentanglement of the two uncertainty types during training. When using other architectures, for instance, using the heteroscedastic uncertainty directly in the likelihood of the GP or using two additive random variables with respective GP and heteroscedastic distributions, the two models will compete for the explained uncertainty and the heteroscedastic prediction network will typically try to explain as much uncertainty as possible. Only in this hierarchical model setting, where the SNGP adds uncertainty to the mean first and the heteroscedastic uncertainty is added in the second step, can we robustly achieve a correct assignment of the different uncertainties using gradient-based optimization.
Challenges and limitations. As outlined above, the main challenge in this model is the inference. We chose an approximate inference scheme that is computationally efficient, utilizing a series of approximations, and the fact that we model the distance-aware and heteroscedastic variances additively. If we wanted to make the model more powerful, at the cost of increased computation, we could thus consider several avenues of improvement: (i) use inducing points or even full GP inference for the GP posterior; (ii) use a more powerful approximation than Laplace, for instance, a variational approximation; or (iii) model the variances jointly, such that we do not only have covariance between data points in f and between classes in u, but have a full covariance between different classes of different data points in u. All of these ideas are interesting avenues for future research.

heading: 4 EXPERIMENTS
text: We conducted experiments on synthetic data and standard image classification benchmarks. As baselines we compare against a standard deterministic model (He et al., 2016; Dosovitskiy et al., 2020), the heteroscedastic method (Collier et al., 2021) and SNGP (Liu et al., 2020). We also compare against the Posterior Network model (Charpentier et al., 2020), which also offers distanceaware uncertainties, but it only applies to problems with few classes. Note that our main motivation for the experiments is to assess whether combining the SNGP and heteroscedastic method can successfully trade off the complementary benefits of the two methods. Since not all tasks and data sets require both types of uncertainty at the same time, we do not expect our proposed model to outperform the baselines on all tasks. We would however expect that, depending on the particular task and required uncertainty type, it would generally outperform one of the two baselines. Our non-synthetic experiments are developed within the open source codebases uncertainty baselines (Nado et al., 2021) and robustness metrics (Djolonga et al., 2020) (to assess the OOD performances). Implementation details are deferred to Appendix A.2.
To measure the predictive performance of the models, we use their accuracy (Acc) and negative log-likelihood (NLL). To measure their uncertainty calibration, we use the expected calibration error (ECE) (Naeini et al., 2015). To measure OOD detection, we use the area under the receiver-operatorcharacteristic curve when using the predictive uncertainty as a score (ROC) and the false-positive rate (rate of classifying an OOD data point as being in-distribution) at 95% recall (FPR95).

heading: 4.1 SYNTHETIC EXPERIMENTS
text: Synthetic OOD data Following the SNGP paper (Liu et al., 2020), to assess the distanceawareness property of HetSNGP in a visually verifiable way, we performed experiments on twodimensional synthetic datasets; (i) two moons and (ii) a Gaussian mixture. We see in Fig. 1 that neither the deterministic or heteroscedastic models offer distance-aware uncertainties and are therefore overconfident on the OOD data (red points). The Posterior network does offer distance-aware uncertainties, which on these datasets grow more slowly away from the training data than the SNGP methods. However, we will see later that scaling the Posterior network to datasets with many classes is too expensive to be practical. Only the SNGP and our proposed HetSNGP are highly uncertain on the OOD points, thus allowing for effective OOD detection.
Synthetic heteroscedastic data We now wish to verify in a low-dimensional setting that HetSNGP also retains the in-distribution heteroscedastic modelling property of the heteroscedastic method. We use the noisy concentric circles from Berthon et al. (2021), where the three circular classes have the same mean, but different amounts of label noise. We see that the heteroscedastic model and our proposed method are able to capture this label noise and thus achieve a better performance, while the deterministic baseline and the SNGP are not. Based on these two synthetic experiments, it becomes apparent that our proposed HetSNGP successfully combines the desirable OOD uncertainties of the SNGP with the heteroscedastic uncertainties on these simple datasets. We now proceed to evaluate these properties on more challenging higher-dimensional datasets.

heading: 4.2 CIFAR EXPERIMENT
text: We start by assessing our method on a real-world image dataset; we trained it on CIFAR-100 and used CIFAR-10 as a near-OOD dataset and Places365 (Zhou et al., 2017) as far-OOD. We measure the OOD detection performance in terms of area under the receiver-operator-characteristic curve (ROC) and false-positive-rate at 95% confidence (FPR95). We also evaluated the methods’ generalization performance on corrupted CIFAR-100 (Hendrycks & Dietterich, 2019). In Table 1 we see that HetSNGP yields a performance between the heteroscedastic and SNGP methods in terms of in-distribution accuracy, but outperforms all baselines in accuracy on the corrupted data. Moreover, it performs on par with the best performing models on both near- and far-OOD detection. This suggests that in-distribution, only the heteroscedastic uncertainty is needed, such that both the heteroscedastic method and our HetSNGP outperform the standard SNGP in terms of accuracy. However, on the corrupted data, which is outside of the training distribution, the SNGP outperforms the heteroscedastic method in terms of accuracy and our HetSNGP outperforms both baselines, since both types of uncertainty are useful in this setting.

heading: 4.3 IMAGENET EXPERIMENT
text: A large-scale dataset with natural label noise and established OOD benchmarks is the Imagenet dataset (Deng et al., 2009; Beyer et al., 2020). The heteroscedastic method has been shown to improve in-distribution performance on Imagenet (Collier et al., 2021). We see in Table 2 that HetSNGP outperforms the SNGP in terms of accuracy and likelihood on the in-distribution Imagenet validation set and performs almost on par with the heteroscedastic model. The slight disadvantage compared to the heteroscedastic model suggests a small trade-off due to the restricted parameterization of the output layer and application of spectral normalization.
However, the true benefits of our model become apparent when looking at the Imagenet OOD datasets (Table 2, right side). Here, we still have the noisy label properties from the original Imagenet dataset, such that heteroscedastic uncertainties are useful, but we are also outside of the training distribution, such that distance-aware model uncertainties become crucial. On Imagenet-C and Imagenet-A, we see that our proposed model makes good use of both of these types of uncertainties and thus manages to outperform all the baselines in terms of accuracy. Additional OOD results on Imagenet-R and Imagenet-V2 are shown in Table 5 in Appendix A.4.

heading: 4.4 IMAGENET-21K
text: We introduce a new large-scale OOD benchmark based on Imagenet-21k. Imagenet-21k is a larger version of the standard Imagenet dataset used above (Deng et al., 2009). It has over 12.8 million training images and 21,843 classes. Each image can have multiple labels, whereas for standard Imagenet a single label is given per image. Note that 999 of the 1,000 Imagenet classes are contained in Imagenet-21k (class n04399382 is missing).
We train a ViT-B/16 (Dosovitskiy et al., 2020) vision transformer model on the Imagenet-21k dataset. See Appendix A.2 for training details. We then evaluate the model on the 1,000 Imagenet classes (setting the predictive probability of class n04399382 to zero). Despite now being in a setting where the model is trained on an order of magnitude more data and greater than 21× more classes, we can use the standard Imagenet OOD datasets. This assesses the scalability of our method. We also hope this new benchmark will prove useful for future OOD research.
We see in Table 3 that the parameter-efficient heteroscedastic method has the best in-distribution precision@1. However, its generalization performance to the OOD benchmarks is the weakest of all the methods. To scale to over 21k output classes, we use the parameter-efficient HetSNGP method. Our method recovers almost the full in-distribution performance of the heteroscedastic method, significantly outperforming the deterministic and SNGP methods. Notably, it also clearly outperforms all other methods on the OOD datasets. Similarly to the CIFAR experiment, we thus see again that in-distribution, the heteroscedastic method and our HetSNGP both outperform the SNGP, since OOD uncertainties are not that important there, while on the OOD datasets, the SNGP and our HetSNGP both outperform the heteroscedastic method. Our method therefore achieves the optimal tradeoff.
We found that SNGP and HetSNGP underfit to the data when using the posterior p(βc | D) at test time, so instead we use the posterior mode β̂ at both training and test time. Additionally, we found spectral normalization was not necessary to preserve distance awareness for the ViT architecture; hence Imagenet-21k experiments are run without spectral normalization.

heading: 4.5 ENSEMBLING EXPERIMENT
text: Deep ensembles are popular methods for model uncertainty estimation due to their simplicity and good performance (Lakshminarayanan et al., 2016). We propose a variant of our HetSNGP method which captures a further source of model uncertainty in the form of parameter uncertainty by a deep ensemble of HetSNGP models. HetSNGP Ensemble is therefore a powerful yet simple method for capturing three major sources of uncertainty; 1) data uncertainty in the labels, 2) model uncertainty in the latent representations and 3) model uncertainty in the parameters.
We compare the HetSNGP Ensemble to a determinstic ensemble as well as ensembles of heteroscedastic and SNGP models. We see in Table 4 that in the case of an ensemble of size four, HetSNGP Ensemble outperforms the baselines on the Imagenet-C OOD dataset in terms of all metrics, while also outperforming them on the standard in-distribution Imagenet dataset in terms of accuracy. Due to computational constraints we do not have error bars in this experiment. However, we do not expect them to be much larger than in Table 2.

heading: 5 ETHICS AND REPRODUCIBILITY STATEMENTS
text: Ethics. We present a generic classification algorithm. HetSNGP is not application-specific and has the computational cost of the heteroscedastic and SNGP methods combined. We do not foresee
ethical issues specific to our method beyond any already associated with the heteroscedastic and SNGP methods or deep neural network classifiers in general.
Reproducibility. Hyperparameters and further details required to reproduce our results are provided in Appendix A.2. We have open-sourced the code to reproduce our CIFAR and Imagenet results. We plan to open-source the code to reproduce the Imagenet-21k experiments. Further details on our use of the Laplace approximation are provided in Appendix A.3.

heading: A APPENDIX
text: A.1 RELATED WORK
Uncertainty estimation. There has been a lot of research on uncertainty estimation in recent years (Ovadia et al., 2019), including the early observation that one needs data and model uncertainties for successful use in real-world applications (Kendall & Gal, 2017). There have also been attempts to directly optimize specialized loss functions to improve model uncertainty (Tagasovska & Lopez-Paz, 2018; Pearce et al., 2018). However, if one wants to directly implement prior knowledge regarding the OOD behavior of the model, one usually needs access to OOD samples during training (Yang et al., 2019; Hafner et al., 2020). Another popular idea is to use a hierarchical output distribution, for instance a Dirichlet distribution (Milios et al., 2018; Sensoy et al., 2018; Malinin & Gales, 2018; Malinin et al., 2019; Hobbhahn et al., 2020), such that the model uncertainty can be encoded in the Dirichlet and the data uncertainty in its Categorical distribution samples. This idea was also used in our Posterior Network baseline (Charpentier et al., 2020). Similarly to our idea of learning separate variances, it has also been proposed to treat the output variance variationally and to specify a hierarchical prior over it (Stirn & Knowles, 2020).
Bayesian neural networks and ensembles. The gold-standard for capturing model uncertainty over the parameters are Bayesian neural networks (Neal, 1993; MacKay, 1992), but they are computationally quite expensive, require a careful choice of prior distribution (Fortuin et al., 2021b; Fortuin, 2021), and require specialized approximate inference schemes, such as Laplace approximation (Immer et al., 2021b;a), variational inference (Hinton & Van Camp, 1993; Graves, 2011; Blundell et al., 2015), or Markov Chain Monte Carlo (Welling & Teh, 2011; Garriga-Alonso & Fortuin, 2021; Fortuin et al., 2021a). A more tractable model class are deep ensemble methods (Lakshminarayanan et al., 2016; Ciosek et al., 2019; D’Angelo et al., 2021; D’Angelo & Fortuin, 2021), although they are computationally still expensive. There are however some ideas to make them less expensive by distilling their uncertainties into simpler models (Malinin et al., 2019; Tran et al., 2020; Havasi et al., 2020; Antorán et al., 2020).
SNGP and heteroscedastic method. The models most relevant to our approach are the SNGP (Liu et al., 2020) and the heteroscedastic method (Collier et al., 2020; 2021). SNGP yields distanceaware model uncertainties, which grow away from the training data. The heteroscedastic method models input- and class-dependent label noise data uncertainty inside the training distribution. Heteroscedastic data uncertainty modelling can also improve segmentation models (Monteiro et al., 2020) and has been linked to the cold posterior effect in Bayesian neural networks (Wenzel et al., 2020; Adlam et al., 2020; Aitchison, 2020). Prior work has primarily focused on modelling data and distance-aware model uncertainty separately. But safety-critical practical applications require both types of uncertainty, HetSNGP is a novel method which fulfills this requirement.
A.2 IMPLEMENTATION DETAILS
To assess our proposed model’s predictive performance and uncertainty estimation capabilities, we conducted experiments on synthetic two moons data (Pedregosa et al., 2011), a mixture of Gaussians, the CIFAR-100 dataset (Krizhevsky et al., 2009), and the Imagenet dataset (Deng et al., 2009). We compare against a standard deterministic ResNet model as a baseline (He et al., 2016), against the heteroscedastic method (Collier et al., 2020; 2021) and the SNGP (Liu et al., 2020) (which form the basis for our combined model) and against the recently proposed Posterior Network model (Charpentier et al., 2020), which also offers distance-aware uncertainties, similarly to the SNGP. We used the same backbone neural network architecture for all models, which was a fully-connected ResNet for the synthetic data, a WideResNet18 on CIFAR and a ResNet50 in Imagenet.
For most baselines, we used the hyperparameters from the uncertainty baselines library (Nado et al., 2021). On CIFAR, we trained our HetSNGP with a learning rate of 0.1 for 300 epochs and used R = 6 factors for the heteroscedastic covariance, a softmax temperature of τ = 0.5 and S = 5000 Monte Carlo samples. On Imagenet, we trained with a learning rate of 0.07 for 270 epochs and used R = 15 factors, a softmax temperature of τ = 1.25 and S = 5000 Monte Carlo samples. We implemented all models in TensorFlow in Python and trained on Tensor Processing Units (TPUs) in the Google Cloud.
We train all Imagnet-21k models for 90 epochs with batch size 1024 on 8 × 8 TPU slices. We train using the Adam optimizer with initial learning rate of 0.001 using a linear learning rate decay schedule with termination point 0.00001 and a warm-up period of 10,000 steps. We train using the sigmoid cross-entropy loss function and L2 weight decay with multiplier 0.03. The heteroscedastic method uses a temperature of 0.4, 1,000 Monte Carlo samples and R = 50 for the low rank approximation. HetSNGP has the same heteroscedastic hyperparameters except the optimal temperature is 1.5. For SNGP and HetSNGP the GP covariance is approximated using the momentum scheme presented in Liu et al. (2020) with momentum parameter 0.999.
A.3 LAPLACE APPROXIMATION
In this section, we will derive the Laplace posterior in Eq. (5). The derivation follows mostly from the sections 3.4 and 3.5 in Rasmussen & Williams (2006).
First note that the log posterior of βc given the data is
log p(βc |x, y) = log p(y |βc) + log p(βc)− Z (7)
where Z is a normalization constant that does not depend on βc. Following Rasmussen & Williams (2006), we will denote the unnormalized log posterior as
Ψ(βc) = log p(y |βc) + log p(βc) (8)
Recall that the first term is the likelihood and the second term is our prior from Eq. (4).
The Laplace approximation now approximates the posterior with a local second-order expansion around the MAP solution, that is
p(βc |x, y) ≈ N (β̂c,Λ−1) (9)
with the MAP solution β̂c = arg maxβc Ψ(βc) and the Hessian Λ = −∇ 2Ψ(βc)|βc=β̂c .
The MAP solution can be found using standard (stochastic) gradient descent, while the Hessian is given by
∇2Ψ(βc) = ∇2 log p(y |βc) +∇2 log p(βc) = ∇β(∇u log p(y |u)∇βu)− Im = ∇β(∇u log p(y |u)Φ)− Im = Φ>∇2u log p(y |u)Φ− Im = −WΦ>Φ− Im
where we used the chain rule and the fact that u = Φβ and W is a diagonal matrix of point-wise second derivatives of the likelihood, that is, Wii = −∇2 log p(yi |ui) (Rasmussen & Williams, 2006). For instance, in the case of the logistic likelihood, Wii = pi (1−pi), where pi is a vector of output probabilities for logits ui. To get the Hessian at the MAP, we then just need to compute this quantity for û = Φβ̂.
The approximate posterior is therefore
p(βc |x, y) ≈ N (β̂c, (WΦTΦ + Im)−1) (10)
where the precision matrix can be computed over data points (recovering Eq. (5)) as
Λ = Im + N∑ i=1 pi(1− pi)ΦiΦ>i (11)
A.4 ADDITIONAL EXPERIMENTAL RESULTS
We add additional baselines for our Imagenet experiment from the epistemic uncertainty literature. In particular we benchmark MC Dropout (Gal & Ghahramani, 2016) and MIMO (Havasi et al., 2020). We use the implementations and Imagenet hyperparameters from the Uncertainty Baselines codebase (Nado et al., 2021).

Review summaries:
This paper studies the combination between model uncertainty and data uncertainty based on spectral-normalized Gaussian process. Theoretical results show that heteroscedastic SNGP allows for joint modeling of model and data uncertainties. The paper then proposes an approximate inference scheme for efficient model training. Experiment results support these claims on synthetic datasets and typical classification benchmarks. 

This paper presents a new method for model and data uncertainty estimation in deep neural networks combining the heteroscedastic method (Collier et al. 2020) and the Spectral-Normalised Gaussian Process (SNGP) method (Liu et al. 2020). It is shown that the two methods are complementary and their combination outperforms state-of-the-art methods on Out-Of-Distribution (OOD) detection on common image classification benchmarks.