Original: D EEP C LASSIFIERS WITH L ABEL N OISE M ODELING AND D ISTANCE A WARENESS
Revised: D EEP C LASSIFIERS WITH L ABEL N OISE M ODELING AND D ISTANCE A WARENESS

Original: Anonymous authors
Revised: Anonymous authors

Original: Paper under double-blind review
Revised: Paper under double-blind review

Original: A BSTRACT
Revised: A BSTRACT

Original: Uncertainty estimation in deep learning has recently emerged as a crucial area of interest to advance reliability and robustness in safety-critical applications.
Revised: Uncertainty estimation in deep learning has recently emerged as a crucial area of interest to advance reliability and robustness in safety-critical applications.

Original: 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.
Revised: 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.

Original: In this work, we propose the HetSNGP method for jointly modeling the model and data uncertainty.
Revised: In this work, we propose the HetSNGP method for jointly modeling the model and data uncertainty.

Original: 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
Revised: 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

Original: 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.
Revised: 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.

Original:  I NTRODUCTION
Revised:  I NTRODUCTION

Original: While deep learning has led to impressive advances in predictive accuracy, models often still suffer from overconﬁdence and ill-calibrated uncertainties (Ovadia et al., 2019).
Revised: While deep learning has led to impressive advances in predictive accuracy, models often still suffer from overconﬁdence and ill-calibrated uncertainties (Ovadia et al., 2019).

Original: 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).
Revised: 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).

Original: Predictive uncertainties generally come in two ﬂavors: model uncertainty (also known as epistemic) and data uncertainty (also known as aleatoric) (Murphy, 2012).
Revised: Predictive uncertainties generally come in two ﬂavors: model uncertainty (also known as epistemic) and data uncertainty (also known as aleatoric) (Murphy, 2012).

Original: Model uncertainty measures how conﬁdent 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.
Revised: Model uncertainty measures how conﬁdent 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.

Original: Data uncertainty measures the uncertainty intrinsic in the data itself, for example due to fundamental noise in the labelling process.
Revised: Data uncertainty measures the uncertainty intrinsic in the data itself, for example due to fundamental noise in the labelling process.

Original: 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.
Revised: 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.

Original: 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 conﬁdent about.
Revised: 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 conﬁdent about.

Original: Many recently proposed models for uncertainty estimation excel at one or the other of these uncertainty types.
Revised: Many recently proposed models for uncertainty estimation excel at one or the other of these uncertainty types.

Original: 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
Revised: 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

Original: OOD detection.
Revised: OOD detection.

Original: Conversely, the heteroscedastic classiﬁcation 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.
Revised: Conversely, the heteroscedastic classiﬁcation 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.

Original: However, there have not been many attempts to combine the complementary beneﬁts of these two types of uncertainty modeling (e.g., Depeweg et al., 2018).
Revised: However, there have not been many attempts to combine the complementary beneﬁts of these two types of uncertainty modeling (e.g., Depeweg et al., 2018) (see related work in Appendix A.1).

Original: 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.
Revised: 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.

Original: We show that HetSNGP gives good in-distribution and OOD accuracy and calibration, yielding a model with uncertainties suitable for deployment in critical applications.
Revised: We show that HetSNGP gives good in-distribution and OOD accuracy and calibration, yielding a model with uncertainties suitable for deployment in critical applications.

Original: Our main contributions are:
Revised: Our main contributions are:

Original: • We propose a new model, the heteroscedastic spectral-normalized Gaussian process (HetSNGP), which provides both distance-aware model and data uncertainties. 
Revised: • We propose a new model, the heteroscedastic spectral-normalized Gaussian process (HetSNGP), which provides both distance-aware model and data uncertainties. (a) Deterministic 4 3 2 1 0 1 2 3 4 311 0.0 0. (b) Heteroscedastic 4 3 2 1 0 1 2 3 4 311 0.0 0. (c) SNGP 4 3 2 1 0 1 2 3 4 311 0.0 0.

Revised: (d) Posterior Net 4 3 2 1 0 1 2 3 4 311 0.0 0.

Original: • We describe an efﬁcient approximate inference scheme that allows training HetSNGP with a computational budget that is comparable to standard neural network training.
Revised: (e) HetSNGP 4 3 2 1 0 1 2 3 4 311 0.0 0. (f) Deterministic 4 3 2 1 0 1 2 3 4 311 0.0 0. (g) Heteroscedastic 4 3 2 1 0 1 2 3 4 311 0.0 0. (h) SNGP 4 3 2 1 0 1 2 3 4 311 0.0 0.

Revised: (i) Posterior Net 4 3 2 1 0 1 2 3 4 311 0.0 0.

Revised: (j) HetSNGP

Original: [Figure]
Revised: points are unseen OOD points, while the other colors represent different classes of training data. The background color shows the normalized predictive uncertainty for the different methods. The deterministic and heteroscedastic methods are overconﬁdent on the OOD data, while the other methods offer distanceaware uncertainties. The uncertainties of the Posterior Network do not grow quickly enough away from the data, such that only the SNGP and proposed HetSNGP provide effective OOD detection.

Revised: • We describe an efﬁcient approximate inference scheme that allows training HetSNGP with a computational budget that is comparable to standard neural network training.

Original: • We show empirically on different benchmark datasets that HetSNGP offers a favorable combination of model and data uncertainty.
Revised: • We show empirically on different benchmark datasets that HetSNGP offers a favorable combination of model and data uncertainty.

Original: 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.
Revised: 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.

Original: • We propose an ensembled version of our model, the HetSNGP Ensemble, which additionally accounts for parameter uncertainty and outperforms other ensemble baselines.
Revised: • We propose an ensembled version of our model, the HetSNGP Ensemble, which additionally accounts for parameter uncertainty and outperforms other ensemble baselines.

Original:  B ACKGROUND 2. M ODEL UNCERTAINTY AND SPECTRAL - NORMALIZED G AUSSIAN PROCESS
Revised:  B ACKGROUND 2. M ODEL UNCERTAINTY AND SPECTRAL - NORMALIZED G AUSSIAN PROCESS

Original: Model uncertainty (or epistemic uncertainty) captures all the uncertainty about whether a given model is correctly speciﬁed for a certain task, given the training data.
Revised: Model uncertainty (or epistemic uncertainty) captures all the uncertainty about whether a given model is correctly speciﬁed for a certain task, given the training data.

Original: 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).
Revised: 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).

Original: 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.
Revised: 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.

Original: We will generally call these points in-distribution (ID).
Revised: We will generally call these points in-distribution (ID).

Original: 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.
Revised: 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.

Original: 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).
Revised: 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).

Original: 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.
Revised: 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.

Original: 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).
Revised: 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).

Original: Note that this approach only partially captures the model uncertainty, namely in form of the uncertainty over the latents.
Revised: Note that this approach only partially captures the model uncertainty, namely in form of the uncertainty over the latents.

Original: 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).
Revised: 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).

Original: 2. D ATA UNCERTAINTY AND THE HETEROSCEDASTIC METHOD
Revised: 2. D ATA UNCERTAINTY AND THE HETEROSCEDASTIC METHOD

Original: As opposed to the model uncertainty described above, data uncertainty is intrinsic in the data and thus irreducible with increasing amounts of data.
Revised: As opposed to the model uncertainty described above, data uncertainty is intrinsic in the data and thus irreducible with increasing amounts of data.

Original: In the case of continuous data (e.g., regression problems), data uncertainty often comes in the form of random noise on the measurements.
Revised: In the case of continuous data (e.g., regression problems), data uncertainty often comes in the form of random noise on the measurements.

Original: For discrete data (e.g., classiﬁcation 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.
Revised: For discrete data (e.g., classiﬁcation 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.

Original: This label noise can be class- and input-dependent (Beyer et al., 2020).
Revised: This label noise can be class- and input-dependent (Beyer et al., 2020).

Original: 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.
Revised: 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.

Original: Modeling this type of data uncertainty can improve the calibration and robustness of predictive models (Collier et al., 2020).
Revised: Modeling this type of data uncertainty can improve the calibration and robustness of predictive models (Collier et al., 2020).

Original: A model that does explicitly handle the input- and class-dependent label noise is the heteroscedastic method (Collier et al., 2021).
Revised: A model that does explicitly handle the input- and class-dependent label noise is the heteroscedastic method (Collier et al., 2021).

Original: 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 classiﬁer.
Revised: 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 classiﬁer.

Original: The covariance matrix of this latent distribution is a function of the input (heteroscedastic) and models inter-class correlations in the logit noise.
Revised: The covariance matrix of this latent distribution is a function of the input (heteroscedastic) and models inter-class correlations in the logit noise.

Original:  M ETHOD 3. S ETUP AND N OTATION
Revised:  M ETHOD 3. S ETUP AND N OTATION

Original: Let us consider a dataset D
Revised: Let us consider a dataset D

Original: = { ( x i , y i ) }
Revised: = { ( x i , y i ) }

Original: Ni =1 of input-output pairs,where x
Revised: Ni =1 of input-output pairs,

Original:  i 2 R d and y i{ 1 , .. .
Revised: where x i ∈ R d and y i ∈ { 1 , .

Original: , K } , that is, a classiﬁcation problem with K classes.
Revised: , K } , that is, a classiﬁcation problem with K classes.

Original: The data examples are assumed to be sampled i.i.d.
Revised: The data examples are assumed to be sampled i.i.d.

Original: from some true data-generating distribution as ( x i , y i ) ⇠ p ⇤ ( x , y ) .
Revised: from some true data-generating distribution as ( x i , y i ) ∼ p ∗ ( x , y ) .

Original: 3. G ENERATIVE PROCESS
Revised: 3. G ENERATIVE PROCESS

Original: 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 .
Revised: 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 .

Original: 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.
Revised: 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.

Original: (2020), which was inspired by the econometrics literature (Train, 2009).
Revised: (2020), which was inspired by the econometrics literature (Train, 2009).

Original: Similarly to
Revised: Similarly to

Original: Liu et al.
Revised: Liu et al.

Original: (2020), we place a latent Gaussian process (GP) prior over f , as p ( f ) =
Revised: (2020), we place a latent Gaussian process (GP) prior over f , as p ( f ) =

Original: GP (0 , k ✓ ( · , · )) , where k ✓ is a parameterized kernel function with parameters ✓ .
Revised: GP (0 , k θ ( · , · )) , where k θ is a parameterized kernel function with parameters θ .

Original: Evaluating this kernel on all pairwise combinations of datapoints in x yields the kernel matrix K ✓ ( x , x )) .
Revised: Evaluating this kernel on all pairwise combinations of datapoints in x yields the kernel matrix K θ ( x , x )) .

Original: We then deﬁne u as a (heteroscedastically) noisy observation of f .
Revised: We then deﬁne u as a (heteroscedastically) noisy observation of f .

Original: Stacking the variables across the whole dataset gives us the matrices F , U 2 R N ⇥ K .
Revised: Stacking the variables across the whole dataset gives us the matrices F , U ∈ R N × K .

Original: We refer to their respective rows as f i , u i 2 R K and their columns as f c , u c 2 R N .
Revised: We refer to their respective rows as f i , u i ∈ R K and their columns as f c , u c ∈ R N .

Original: The columns are independent under the GP prior, but the rows are not.
Revised: The columns are independent under the GP prior, but the rows are not.

Original: Conversely, the columns are correlated in the heteroscedastic noise model, while the rows are independent.
Revised: Conversely, the columns are correlated in the heteroscedastic noise model, while the rows are independent.

Original: A hierarchical model using both uncertainties therefore leads to logits that are correlated across both data points and classes.
Revised: A hierarchical model using both uncertainties therefore leads to logits that are correlated across both data points and classes.

Original: The full generative process is
Revised: The full generative process is

Original: [Equation]
Revised: [Equation]

Original: The predictive can then be computed my marginalizing, that is p ( y | x ) =
Revised: The predictive can then be computed my marginalizing, that is p ( y | x ) =

Original: E u
Revised: E u

Original: [ p ( y | u )]=
Revised: [ p ( y | u )]

Original: R p ( y | u ) p ( u | x ) d u .
Revised: = (cid:82) p ( y | u ) p ( u | x ) d u .

Original: Intuitively, f captures the model uncertainty, that is, the uncertainty about the functional mapping between x and y on the level of the latents.
Revised: Intuitively, f captures the model uncertainty, that is, the uncertainty about the functional mapping between x and y on the level of the latents.

Original: 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.
Revised: 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.

Original: 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).
Revised: 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).

Original: It does not itself capture the model uncertainty, but inherits it through its hierarchical dependence on f , such that the resulting probability p ( y
Revised: It does not itself capture the model uncertainty, but inherits it through its hierarchical dependence on f , such that the resulting probability p ( y

Original: | u ) ultimately jointly captures both types of uncertainty.
Revised: | u ) ultimately jointly captures both types of uncertainty.

Original: In practice, we usually learn the kernel using deep kernel learning (Wilson et al., 2016), that is, we deﬁne it as the RBF kernel k ✓ ( x i , x j ) =
Revised: In practice, we usually learn the kernel using deep kernel learning (Wilson et al., 2016), that is, we deﬁne it as the RBF kernel k θ ( x i , x j ) =

Original: k RBF ( h i , h j ) =
Revised: k RBF ( h i , h j ) =

Original: exp( k h i  h
Revised: exp( (cid:107) h

Original:  j k 22 /  ) with length scale  and h
Revised: i − h j (cid:107) 22 /λ ) with length scale λ and h

Original: i = h ( x i ; ✓ ) 2 R m with h ( · ; ✓ ) being a neural network model (e.g., a ResNet) parameterized by ✓ .
Revised: i = h ( x i ; θ ) ∈ R m with h ( · ; θ ) being a neural network model (e.g., a ResNet) parameterized by θ .

Original: This kernel is then shared between classes.
Revised: This kernel is then shared between classes.

Original: Moreover, following Liu et al.
Revised: Moreover, following Liu et al.

Original: (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.
Revised: (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.

Original: Additionally, ⌃ ( · ; ' ) is usually also a neural network parameterized by ' .
Revised: Additionally, Σ ( · ; ϕ ) is usually also a neural network parameterized by ϕ .

Original: To ease notation, we will typically drop these parameters in the following.
Revised: To ease notation, we will typically drop these parameters in the following.

Original: Note that the prior over f c (and thus also u c ) has zero mean, thus leading to a uniform output distribution away from the data.
Revised: Note that the prior over f c (and thus also u c ) has zero mean, thus leading to a uniform output distribution away from the data.

Original: 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.
Revised: 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.

Original: 3. C OMPUTATIONAL APPROXIMATIONS
Revised: 3. C OMPUTATIONAL APPROXIMATIONS

Original: The generative process above requires computing two integrals (i.e., over u and f ) and inferring an exact GP posterior, making it computationally intractable.
Revised: The generative process above requires computing two integrals (i.e., over u and f ) and inferring an exact GP posterior, making it computationally intractable.

Original: We make several approximations to ensure tractability.
Revised: We make several approximations to ensure tractability.

Original: Low-rank approximation.
Revised: Low-rank approximation.

Original: The heteroscedastic covariance matrix ⌃ ( x i ; ' ) is a K ⇥ K matrix which is a function of the input.
Revised: The heteroscedastic covariance matrix Σ ( x i ; ϕ ) is a K × K matrix which is a function of the input.

Original: Parameterizing this full matrix is costly in terms of computation, parameter count, and memory.
Revised: Parameterizing this full matrix is costly in terms of computation, parameter count, and memory.

Original: Therefore, following Collier et al.
Revised: Therefore, following Collier et al.

Original: (2021), we make a low-rank approximation ⌃ ( x i ; ' )
Revised: (2021), we make a low-rank approximation Σ ( x i ; ϕ )

Original: =
Revised: =

Original: V ( x i ) V ( x i ) > + d 2 ( x i ) .
Revised: V ( x i ) V ( x i ) (cid:62) + d 2 ( x i ) .

Original: Here, V ( x i ) is a K ⇥ R matrix with R << K and d 2 ( x i ) is a K dimensional vector added to the diagonal of V ( x i ) V
Revised: Here, V ( x i ) is a K × R matrix with R << K and d 2 ( x i ) is a K dimensional vector added to the diagonal of V ( x i ) V

Original: ( x i ) > .
Revised: ( x i ) (cid:62) .

Revised: 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.

Revised: Note that R can be tuned as a hyperparameter and can thus be made large enough, such that this assumption holds.

Original: Also following Collier et al.
Revised: Also following Collier et al.

Original: (2021), we introduce a parameter-efﬁcient version of our method to enable scaling to problems with many classes, e.g., Imagenet21k which has 21,843 classes.
Revised: (2021), we introduce a parameter-efﬁcient version of our method to enable scaling to problems with many classes, for instance, Imagenet21k which has 21,843 classes.

Original: In the standard version of our method V ( x i ) is parameterized as an afﬁne transformation of the shared representation h 
Revised: In the standard version of our method, V ( x i ) is parameterized as an afﬁne transformation of the shared representation h i .

Original: i , in this way HetSNGP can be added as a single output layer on top of a base network.
Revised: In this way, HetSNGP can be added as a single output layer on top of a base network.

Original: For the parameter-efﬁcient HetSNGP we parameterize V ( x i ) 
Revised: For the parameter-efﬁcient HetSNGP, we parameterize V ( x i ) =

Original: = v ( x i ) [-
 >-] R  V where v ( x i ) is a vector of dimension R , 1 R is a vector of ones of dimension R , V is a K ⇥ R matrix of free parameters, and  denotes element-wise multiplication.
Revised: v ( x i ) [-
 >-]1 (cid:62) R (cid:12) V where v ( x i ) is a vector of dimension R , 1 R is a vector of ones of dimension R , V is a K × R matrix of free parameters, and (cid:12) denotes element-wise multiplication.

Revised: This additional approximation assumes that the similarity between classes is largely data-point-independent (i.e., different classes are intrinsically similar), while the speciﬁc data point can be responsible for scaling certain dimensions of similarity in the R -dimensional space.

Original: Random Fourier feature approximation.
Revised: Random Fourier feature approximation.

Original: Computing the exact GP posterior requires O ( N 3 ) operations (because ones needs to invert an N ⇥ N kernel matrix), which can be prohibitive for large datasets.
Revised: Computing the exact GP posterior requires O ( N 3 ) operations (because ones needs to invert an N × N kernel matrix), which can be prohibitive for large datasets.

Original: Following Liu et al.
Revised: Following Liu et al.

Original: (2020), we thus use a random Fourier feature (RFF) approximation (Rahimi et al., 2007), leading to a low-rank approximation of the kernel matrix as K ✓ ( x , x )=  > (  2 R N ⇥ m ), with random features  i =
Revised: (2020), we thus use a random Fourier feature (RFF) approximation (Rahimi et al., 2007), leading to a low-rank approximation of the kernel matrix as K θ ( x , x )

Original: q 2 m cos
Revised: =

Original:  (  W h i + b ) where W is a ﬁxed weight matrix with entries sampled from N (0 , 1) and b is a ﬁxed bias vector with entries sampled from U (0 , 2 ⇡ ) .
Revised: ΦΦ (cid:62) ( Φ ∈ R N × m ), with random features Φ i = (cid:113) 2 m cos( W h i + b ) where W is a ﬁxed weight matrix with entries sampled from N (0 , 1) and b is a ﬁxed bias vector with entries sampled from U (0 , 2 π ) .

Original: This approximates the inﬁnite-dimensional reproducing kernel Hilbert space (RKHS) of the RBF kernel with a subspace spanned by m randomly sampled basis functions.
Revised: This approximates the inﬁnite-dimensional reproducing kernel Hilbert space (RKHS) of the RBF kernel with a subspace spanned by m randomly sampled basis functions.

Original: This reduces the GP inference complexity to O ( Nm 2 ) , where m is the dimensionality of the latent space.
Revised: This reduces the GP inference complexity to O ( Nm 2 ) , where m is the dimensionality of the latent space.

Original: We can then write the model as a linear model in this feature space, namely
Revised: We can then write the model as a linear model in this feature space, namely

Original: [Equation]
Revised: [Equation]

Original: with f 
Revised: with f c = Φ β

Original: [Equation]
Revised: [Equation]

Original: and
Revised: and

Original: [Equation]
Revised: [Equation]

Original: Here,   c is a linear regressor in the space of the random Fourier features and the other two terms of u i are the low-rank approximation to the heteroscedastic uncertainties (as described above). 
Revised: Here, Φ β c is a linear regressor in the space of the random Fourier features and the other two terms of u i 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 ﬁdelity of the model.

Original: Laplace approximation.
Revised: Since most data sets contain a lot of redundant data points, one often ﬁnds an m (cid:28) N that models the similarities in the data faithfully. Require: dataset D = { ( x i , y i ) } Ni = Initialize θ, ϕ, W , b , (cid:98) Σ for train step = 1 to max step do

Revised: Take minibatch ( X i , y i ) from D for s = 1 to S do (cid:15) sK ∼ N ( 0 , I K ) , (cid:15) sR ∼ N ( 0 , I R )

Revised: u si,

Revised: c

Revised: =

Revised: Φ (cid:62) i ˆ β c + d ( x i ) (cid:12) (cid:15) sK +

Revised: V ( x i ) (cid:15) sR end for

Revised: L = − 1 S (cid:80) Ss =1 log p ( X i , y i | u s ) −(cid:107) ˆ β (cid:107)

Revised: Update { θ, ϕ, ˆ β } via SGD on L if ﬁnal epoch then

Revised: Compute { (cid:98) Σ − 1 c } Kc =1 as per Eq.

Revised: (5) end if end for

Revised: Algorithm

Revised: Require: test example x ∗ Φ ∗ =

Revised: (cid:113) 2 m cos( W h ( x ∗ ) + b ) for s = 1 to S do β sc ∼ p ( β c | D ) (cid:15) sK ∼ N ( 0 , I K ) , (cid:15) sR ∼ N ( 0 , I R )

Revised: u s ∗ ,c =

Revised: Φ ∗(cid:62) β sc + d ( x ∗ ) (cid:12) (cid:15) sK +

Revised: V ( x ∗ ) (cid:15) sR end for p ( y ∗ =

Revised: c | x ∗ )

Revised: = 1 S (cid:80)

Revised: Ss = exp( u s ∗ ,c /τ ) (cid:80)

Revised: Kk =1 exp( u s ∗ ,k /τ )

Original:  When using a Gaussian likelihood (i.e., in a regression setting), the GP posterior inference can be performed in closed form.
Revised: Predict y [Equation] Laplace approximation. When using a Gaussian likelihood (i.e., in a regression setting), the GP posterior inference can be performed in closed form.

Original: However, for classiﬁcation problems this is not possible, because the Categorical likelihood used is not conjugate to the Gaussian prior.
Revised: However, for classiﬁcation problems this is not possible, because the Categorical likelihood used is not conjugate to the Gaussian prior.

Original: We thus need to resort to approximate posterior inference.
Revised: We thus need to resort to approximate posterior inference.

Original: Again following Liu et al.
Revised: Again following Liu et al.

Original: (2020), we perform a
Revised: (2020), we perform a

Original: Laplace approximation (Rasmussen & Williams, 2006) to the RFF-GP posterior, which yields the closed-form approximate posterior for  c 
Revised: Laplace approximation (Rasmussen & Williams, 2006) to the RFF-GP posterior, which yields the closed-form approximate posterior for β c p

Original: [Equation]
Revised: [Equation]

Original: where p i,c is shorthand for the softmax output p ( y i =
Revised: where p i,c is shorthand for the softmax output p ( y i =

Original: c | ˆ u i ) as deﬁned in Eq.
Revised: c | ˆ u i ) as deﬁned in Eq.

Original: (3), where ˆ u i,c
Revised: (3), where ˆ u i,c

Original: =  > i ˆ  c .
Revised: = Φ (cid:62) i ˆ β c .

Original: The derivation of this is deferred to Appendix A.2.
Revised: The derivation of this is deferred to Appendix A.3.

Original: Here, b ⌃  1 c can be cheaply computed over minibatches of data by virtue of being a sum over data points.
Revised: Here, (cid:98) Σ − 1 c can be cheaply computed over minibatches of data by virtue of being a sum over data points.

Original: Moreover, ˆ  c is the MAP solution, which can be obtained using gradient descent on the unnormalized log posterior,  log p (  | D ) /  log p ( D |  )  k  k 2 , where the squared norm regularizer stems from the standard Gaussian prior on  .
Revised: Moreover, ˆ β c is the MAP solution, which can be obtained using gradient descent on the unnormalized log posterior, − log p ( β | D ) ∝ − log p ( D | β ) − (cid:107) β (cid:107) 2 , where the squared norm regularizer stems from the standard Gaussian prior on β .

Original: Using our likelihood, the training objective takes the form of a ridge-regularized cross-entropy loss.
Revised: Using our likelihood, the training objective takes the form of a ridge-regularized cross-entropy loss.

Original: We also use this objective to train the other trainable parameters ✓ and ' of our model.
Revised: We also use this objective to train the other trainable parameters θ and ϕ of our model.

Revised: Note that this approximation is necessarily ignoring multi-modality of the posterior, since it can by deﬁnition only ﬁt one mode around the MAP solution.

Revised: 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).

Original: Monte Carlo approximation.
Revised: Monte Carlo approximation.

Original: Finally, we are ultimately interested in E u
Revised: Finally, we are ultimately interested in E u

Original: [ p ( y | u )]=
Revised: [ p ( y | u )]

Original: R p ( y | u ) p ( u | x ) d u , which again would require solving a high-dimensional integral and is thus computationally intractable.
Revised: = (cid:82) p ( y | u ) p ( u | x ) d u , which again would require solving a high-dimensional integral and is thus computationally intractable.

Original: Following Collier et al.
Revised: Following Collier et al.

Original: , we approximate it using Monte Carlo samples.
Revised: , we approximate it using Monte Carlo samples.

Original: Moreover, we approximate the argmax in Eq.
Revised: Moreover, we approximate the argmax in Eq.

Original: (3) with a softmax.
Revised: (3) with a softmax.

Original: 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
Revised: 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

Original: [Equation]
Revised: [Equation]

Original: with u si,c =  > i  sc + d ( x i )  ✏ sK + V ( x i ) ✏ sR and  sc ⇠ p (  c | D ) and ✏ sK ⇠ N ( 0 , I K ) and ✏ sR = N ( 0 , I R ) 
Revised: with u si,c = Φ (cid:62) i β sc + d ( x i ) (cid:12) (cid:15) sK + V ( x i ) (cid:15) sR and β sc ∼ p ( β c | D ) and (cid:15) sK ∼ N ( 0 , I K ) and (cid:15) sR = N ( 0 , I R ) [Figure]

Original: It should be noted that we only use Eq.
Revised: It should be noted that we only use Eq.

Original: (6) at test time, while during training we replace  sc with its mean ˆ  c instead of sampling it (see Section 3.4).
Revised: (6) at test time, while during training we replace β sc with its mean ˆ β c instead of sampling it (see Section 3.4).

Original: Note also that, while a larger number S of Monte
Revised: Note also that, while a larger number S of Monte

Original: 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.3. I MPLEMENTATION Require: dataset D = { ( x i , y i ) } Ni =
Revised: 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.

Original: Initialize ✓ , ' , W , b , b ⌃ for train step = 1 to max step do

Original: Take minibatch ( X i , y i ) from D for s = 1 to S do ✏

Original: [Equation]

Original: u

Original: [Equation]

Original: V ( x i ) ✏ sR end for

Original: L =  1 S

Original: P

Original: Ss =1 log p ( X i , y i | u s )

Original: k ˆ  k

Original: Update { ✓ , ' , ˆ  } via SGD on L if ﬁnal epoch then Compute { b ⌃  1 c } Kc =1 as per Eq. (5) end if end for Require: test example x
Revised: 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).

Original: ⇤  ⇤ =

Original: q 2 m cos( W h ( x ⇤ )

Original: + b ) for s = 1 to S do  sc ⇠ p (  c | D ) ✏ sK ⇠ N ( 0 , I K ) , ✏ sR ⇠ N ( 0 , I R ) u s ⇤ ,c =  ⇤>  sc + d ( x ⇤ )  ✏ sK +
Revised: 3. I MPLEMENTATION

Original: V ( x ⇤ ) ✏ sR end for p ( y ⇤ =

Original: c | x ⇤ )

Original: = 1 S

Original: P

Original: Ss = exp( u s ⇤ ,c / ⌧ )

Original: P Kk =1 exp( u s ⇤ , [Equation] The training of our proposed model is described in Algorithm 1 and the prediction in Algorithm 2.
Revised: The training of our proposed model is described in Algorithm 1 and the prediction in Algorithm 2.

Original: Intuition.
Revised: Intuition.

Original: 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.
Revised: 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.

Original: This means that when we have an input x far that is far away from the training data,[Figure] the values of the RBF kernel exp( k h  h far k 22 /  ) will be small, and thus the uncertainty of f c will be large.
Revised: This means that when we have an input x far that is far away from the training data, the values of the RBF kernel exp( (cid:107) h − h far (cid:107) 22 /λ ) will be small, and thus the uncertainty of f c will be large.

Original: The resulting output uncertainties will therefore be large regardless of the heteroscedastic variance ⌃ ( x i ) (since they are additive) allowing for effective OOD detection.
Revised: The resulting output uncertainties will therefore be large regardless of the heteroscedastic variance Σ ( x i ) (since they are additive) allowing for effective OOD detection.

Original: Conversely, if we have an in-distribution input, the kernel values will be large and the model uncertainty captured in f c will be small.
Revised: Conversely, if we have an in-distribution input, the kernel values will be large and the model uncertainty captured in f c will be small.

Original: In this case, the additive output uncertainty will be dominated by the heteroscedastic variance ⌃ ( x i ) , allowing for effective label noise modeling in-distribution.
Revised: In this case, the additive output uncertainty will be dominated by the heteroscedastic variance Σ ( x i ) , allowing for effective label noise modeling in-distribution.

Revised: While at ﬁrst 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.

Revised: Rather, it is the only design we found that robustly achieves a disentanglement of the two uncertainty types during training.

Revised: 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.

Revised: Only in this hierarchical model setting, where the SNGP adds uncertainty to the mean ﬁrst 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.

Original: Challenges and limitations.
Revised: Challenges and limitations.

Original: As outlined above, the main challenge in this model is the inference.
Revised: As outlined above, the main challenge in this model is the inference.

Original: We chose an approximate inference scheme that is computationally efﬁcient, utilizing a series of approximations, and the fact that we model the distance-aware and heteroscedastic variances additively.
Revised: We chose an approximate inference scheme that is computationally efﬁcient, utilizing a series of approximations, and the fact that we model the distance-aware and heteroscedastic variances additively.

Original: 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 .
Revised: 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 .

Original: All of these ideas are interesting avenues for future research. 
Revised: All of these ideas are interesting avenues for future research. [Table]

Original:  E XPERIMENTS
Revised:  E XPERIMENTS

Original: We conducted experiments on synthetic data and standard image classiﬁcation benchmarks.
Revised: We conducted experiments on synthetic data and standard image classiﬁcation benchmarks.

Original: 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).
Revised: 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).

Original: 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.
Revised: 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.

Revised: Note that our main motivation for the experiments is to assess whether combining the SNGP and heteroscedastic method can successfully trade off the complementary beneﬁts of the two methods.

Revised: 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.

Revised: We would however expect that, depending on the particular task and required uncertainty type, it would generally outperform one of the two baselines.

Original: 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).
Revised: 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).

Original: Implementation details are deferred to Appendix A.1.
Revised: Implementation details are deferred to Appendix A.2.

Revised: To measure the predictive performance of the models, we use their accuracy ( Acc ) and negative log-likelihood ( NLL ).

Revised: To measure their uncertainty calibration, we use the expected calibration error ( ECE ) (Naeini et al., 2015).

Revised: 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 ).

Original: 4. S YNTHETIC EXPERIMENTS
Revised: 4. S YNTHETIC EXPERIMENTS

Original: Synthetic OOD data Following the SNGP paper (Liu et al., 2020), to assess the distanceawareness property of HetSNGP in a visually veriﬁable way, we performed experiments on twodimensional synthetic datasets; (i) two moons and (ii) a Gaussian mixture.
Revised: Synthetic OOD data Following the SNGP paper (Liu et al., 2020), to assess the distanceawareness property of HetSNGP in a visually veriﬁable way, we performed experiments on twodimensional synthetic datasets; (i) two moons and (ii) a Gaussian mixture.

Original: We see in Fig.
Revised: We see in Fig.

Original: 1 that[Table] [Table] neither the deterministic or heteroscedastic models offer distance-aware uncertainties and are therefore overconﬁdent on the OOD data (red points).
Revised: 1 that neither the deterministic or heteroscedastic models offer distance-aware uncertainties and are therefore overconﬁdent on the OOD data (red points).

Original: The Posterior network does offer distance-aware uncertainties, which on these datasets grow slower away from the training data than the SNGP methods.
Revised: The Posterior network does offer distance-aware uncertainties, which on these datasets grow more slowly away from the training data than the SNGP methods.

Original: However, we will see later that scaling  Posterior network to datasets with many classes is too expensive to be practical.
Revised: However, we will see later that scaling the Posterior network to datasets with many classes is too expensive to be practical.

Original: Only the SNGP and our proposed HetSNGP are highly uncertain on the OOD points, thus allowing for effective OOD detection.
Revised: Only the SNGP and our proposed HetSNGP are highly uncertain on the OOD points, thus allowing for effective OOD detection.

Original: 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.
Revised: 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.

Original: We use the noisy concentric circles from Berthon et al.
Revised: We use the noisy concentric circles from Berthon et al.

Original: (2021), where the three circular classes have the same mean, but different amounts of label noise.
Revised: (2021), where the three circular classes have the same mean, but different amounts of label noise.

Original: 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.
Revised: 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.

Original: Based on these two synthetic experiments, it becomes apparent that our proposed HetSNGP successfully combines the desirable
Revised: Based on these two synthetic experiments, it becomes apparent that our proposed HetSNGP successfully combines the desirable

Original: OOD uncertainties of the SNGP with the heteroscedastic uncertainties on these simple datasets.
Revised: OOD uncertainties of the SNGP with the heteroscedastic uncertainties on these simple datasets.

Original: We now proceed to evaluate these properties on more challenging higher dimensional datasets. 
Revised: We now proceed to evaluate these properties on more challenging higher-dimensional datasets. [Table]

Original: 4. CIFAR EXPERIMENT
Revised: 4. CIFAR EXPERIMENT

Original: 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.
Revised: 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.

Original: We measure the OOD detection performance in terms of area under the receiver-operator-characteristic curve (ROC) and false-positive-rate at 95% conﬁdence (FPR95).
Revised: We measure the OOD detection performance in terms of area under the receiver-operator-characteristic curve (ROC) and false-positive-rate at 95% conﬁdence (FPR95).

Original: We also evaluated the methods’ generalization performance on corrupted CIFAR-100 (Hendrycks & Dietterich, 2019).
Revised: We also evaluated the methods’ generalization performance on corrupted CIFAR-100 (Hendrycks & Dietterich, 2019).

Original: 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.
Revised: 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.

Original: Moreover, it performs on par with the best performing models on both near- and far-OOD detection.
Revised: Moreover, it performs on par with the best performing models on both near- and far-OOD detection.

Revised: 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.

Revised: 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.

Original: 4. I MAGENET EXPERIMENT
Revised: 4. I MAGENET EXPERIMENT

Original: A large-scale dataset with natural label noise and established OOD benchmarks is the Imagenet dataset (Deng et al., 2009; Beyer et al., 2020).
Revised: A large-scale dataset with natural label noise and established OOD benchmarks is the Imagenet dataset (Deng et al., 2009; Beyer et al., 2020).

Original: The heteroscedastic method has been shown to improve in-distribution performance on Imagenet (Collier et al., 2021).
Revised: The heteroscedastic method has been shown to improve in-distribution performance on Imagenet (Collier et al., 2021).

Original: 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.
Revised: 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.

Original: 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.[Table] [Table] 5 runs.
Revised: 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.

Original: Bold numbers are within one standard error of the best performing model.

Original: We use standard Imagenet, Imagenet-C, Imagenet-A, Imagenet-R, and Imagenet-V2 as OOD datasets.

Original: HetSNGP outperforms the baselines on all OOD datasets.

Original: However, the true beneﬁts of our model become apparent when looking at the Imagenet OOD datasets (Table 2, right side).
Revised: However, the true beneﬁts of our model become apparent when looking at the Imagenet OOD datasets (Table 2, right side).

Original: 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.
Revised: 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.

Original: 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.
Revised: 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.

Original: Additional OOD results on Imagenet-R and Imagenet-V2 are shown in Table 3.
Revised: Additional OOD results on Imagenet-R and Imagenet-V2 are shown in Table 5 in Appendix A.4.

Original: 4. I MAGENET -21 K
Revised: 4. I MAGENET -21 K

Original: We introduce a new large-scale OOD benchmark based on Imagenet-21k.
Revised: We introduce a new large-scale OOD benchmark based on Imagenet-21k.

Original: Imagenet-21k is a larger version of the standard Imagenet dataset used above (Deng et al., 2009).
Revised: Imagenet-21k is a larger version of the standard Imagenet dataset used above (Deng et al., 2009).

Original: It has over 12.8 million training images and 21,843 classes.
Revised: It has over 12.8 million training images and 21,843 classes.

Original: Each image can have multiple labels, whereas for standard
Revised: Each image can have multiple labels, whereas for standard

Original: Imagenet a single label is given per image.
Revised: Imagenet a single label is given per image.

Original: Note that 999 of the 1,000 Imagenet classes are contained in Imagenet-21k (class n04399382 is missing).
Revised: Note that 999 of the 1,000 Imagenet classes are contained in Imagenet-21k (class n04399382 is missing).

Original: We train a ViT-B/16 (Dosovitskiy et al., 2020) vision transformer model on the Imagenet-21k dataset.
Revised: We train a ViT-B/16 (Dosovitskiy et al., 2020) vision transformer model on the Imagenet-21k dataset.

Original: See Appendix A.1 for training details.
Revised: See Appendix A.2 for training details.

Original: We then evaluate the model on the 1,000 Imagenet classes (setting the predictive probability of class n04399382 to zero).
Revised: We then evaluate the model on the 1,000 Imagenet classes (setting the predictive probability of class n04399382 to zero).

Original: 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.
Revised: 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.

Original: This assesses the scalability of our method.
Revised: This assesses the scalability of our method.

Original: We also hope this new benchmark will prove useful for future OOD research.
Revised: We also hope this new benchmark will prove useful for future OOD research.

Original:  We see in Table 5 that the parameter-efﬁcient heteroscedastic method has the best in-distribution precision@1.
Revised: [Table] [Table] We see in Table 3 that the parameter-efﬁcient heteroscedastic method has the best in-distribution precision@1.

Original: However, its generalization performance to the OOD benchmarks is the weakest of all the methods.
Revised: However, its generalization performance to the OOD benchmarks is the weakest of all the methods.

Original: To scale to over 21k output classes, we use the parameter-efﬁcient HetSNGP method.
Revised: To scale to over 21k output classes, we use the parameter-efﬁcient HetSNGP method.

Original: Our method recovers almost the full in-distribution performance of the heteroscedastic method, signiﬁcantly outperforming the deterministic and SNGP methods.
Revised: Our method recovers almost the full in-distribution performance of the heteroscedastic method, signiﬁcantly outperforming the deterministic and SNGP methods.

Original: Notably, it also clearly outperforms all other methods on the OOD datasets.
Revised: Notably, it also clearly outperforms all other methods on the OOD datasets.

Original: We found that SNGP and  HetSNGP underﬁt to the data when using the posterior p (  c | D ) at test time, so instead we use the posterior mode
Revised: 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.

Original: ˆ 
Revised: Our method therefore achieves the optimal tradeoff.

Original:  at both training and test time.
Revised: We found that SNGP and HetSNGP underﬁt 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.

Original: 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.
Revised: 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.

Original: 4. E NSEMBLING EXPERIMENT
Revised: 4. E NSEMBLING EXPERIMENT

Original: Deep ensembles are popular methods for model uncertainty estimation due to their simplicity and good performance (Lakshminarayanan et al., 2016).
Revised: Deep ensembles are popular methods for model uncertainty estimation due to their simplicity and good performance (Lakshminarayanan et al., 2016).

Original: 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.
Revised: 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.

Original: 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.
Revised: 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.

Original: We compare the HetSNGP Ensemble to a determinstic ensemble as well as ensembles of heteroscedastic and SNGP models.
Revised: We compare the HetSNGP Ensemble to a determinstic ensemble as well as ensembles of heteroscedastic and SNGP models.

Original: We see in Table 4 that in the case of an ensemble of size four,
Revised: We see in Table 4 that in the case of an ensemble of size four,

Original: 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.
Revised: 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.

Original: Due to computational constraints we do not have error bars in this experiment.
Revised: Due to computational constraints we do not have error bars in this experiment.

Original: However, we do not expect them to be much larger than in Table 2.
Revised: However, we do not expect them to be much larger than in Table 2.

Original:   R ELATED W ORK
Revised:  E THICS AND R EPRODUCIBILITY S TATEMENTS

Revised: Ethics. We present a generic classiﬁcation algorithm. HetSNGP is not application-speciﬁc and has the computational cost of the heteroscedastic and SNGP methods combined.

Revised: We do not foresee ethical issues speciﬁc to our method beyond any already associated with the heteroscedastic and

Revised: SNGP methods or deep neural network classiﬁers in general.

Revised: Reproducibility.

Revised: Hyperparameters and further details required to reproduce our results are provided in Appendix A.2.

Revised: We have open-sourced the code to reproduce our CIFAR and Imagenet results.

Revised: We plan to open-source the code to reproduce the Imagenet-21k experiments.

Revised: Further details on our use of the Laplace approximation are provided in Appendix A.3.

Revised: A A PPENDIX A. R ELATED W ORK

Original: Uncertainty estimation.
Revised: Uncertainty estimation.

Original: 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).
Revised: 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).

Original: There have also been attempts to directly optimize specialized loss functions to improve model uncertainty (Tagasovska & Lopez-Paz, 2018; Pearce et al., 2018).
Revised: There have also been attempts to directly optimize specialized loss functions to improve model uncertainty (Tagasovska & Lopez-Paz, 2018; Pearce et al., 2018).

Original: 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).
Revised: 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).

Original: Another popular idea is to use a hierarchical output distribution, for instance a Dirichlet distribution (Milios et al., 2018; Sensoy et al., 2018;
Revised: Another popular idea is to use a hierarchical output distribution, for instance a Dirichlet distribution (Milios et al., 2018; Sensoy et al., 2018;

Original: Malinin & Gales, 2018;
Revised: Malinin & Gales, 2018;

Original: 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.
Revised: 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.

Original: This idea was also used in our Posterior Network baseline (Charpentier et al., 2020).
Revised: This idea was also used in our Posterior Network baseline (Charpentier et al., 2020).

Original: 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).
Revised: 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).

Original: Bayesian neural networks and ensembles.
Revised: Bayesian neural networks and ensembles.

Original: 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;
Revised: 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;

Original: Blundell et al., 2015), or Markov Chain Monte Carlo (Welling & Teh, 2011; Garriga-Alonso & Fortuin, 2021; Fortuin et al., 2021a).
Revised: Blundell et al., 2015), or Markov Chain Monte Carlo (Welling & Teh, 2011; Garriga-Alonso & Fortuin, 2021; Fortuin et al., 2021a).

Original: 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.
Revised: 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.

Original: 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; Antor´an et al., 2020).
Revised: 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; Antor´an et al., 2020).

Original: SNGP and heteroscedastic method.
Revised: SNGP and heteroscedastic method.

Original: The models most relevant to our approach are the SNGP (Liu et al., 2020) and the heteroscedastic method (Collier et al., 2020; 2021).
Revised: The models most relevant to our approach are the SNGP (Liu et al., 2020) and the heteroscedastic method (Collier et al., 2020; 2021).

Original: SNGP yields distanceaware model uncertainties, which grow away from the training data.
Revised: SNGP yields distanceaware model uncertainties, which grow away from the training data.

Original: The heteroscedastic method models input- and class-dependent label noise data uncertainty inside the training distribution.
Revised: The heteroscedastic method models input- and class-dependent label noise data uncertainty inside the training distribution.

Original: 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).
Revised: 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).

Original: Prior work has primarily focused on modelling data and distance-aware model uncertainty separately.
Revised: Prior work has primarily focused on modelling data and distance-aware model uncertainty separately.

Original: But safety-critical practical applications require both types of uncertainty, HetSNGP is a novel method which fulﬁlls this requirement.
Revised: But safety-critical practical applications require both types of uncertainty, HetSNGP is a novel method which fulﬁlls this requirement.

Original: C ONCLUSION
Revised: A. I MPLEMENTATION DETAILS

Original: We have proposed a new model, the HetSNGP, that jointly captures distance-aware model uncertainties and heteroscedastic data uncertainties. HetSNGP allows for a favorable combination of these two complementary types of uncertainty and allows for out-of-distribution detection and generalization as well as in-distribution performance and calibration on different benchmark datasets.  E THICS AND R EPRODUCIBILITY S TATEMENTS
Revised: 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.

Original: Ethics. We present a generic classiﬁcation algorithm. HetSNGP is not application-speciﬁc and has the computational cost of the heteroscedastic and SNGP methods combined.
Revised: 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.

Original: We do not foresee ethical issues speciﬁc to our method beyond any already associated with  the heteroscedastic  andSNGP methods or deep neural network classiﬁers in general.
Revised: 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

Original: Reproducibility.{+S = 
+}
Revised: {+S = 
+}

Original: Hyperparameters and further details required to reproduce our results are provided in Appendix A.1.
Revised: Monte Carlo samples.

Original: 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.
Revised: On Imagenet, we trained with a learning rate of 0.07 forepochs 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

Revised: Units (TPUs) in the Google Cloud.

Revised: We train all Imagnet-21k models for 90 epochs with batch size 1024 on 8 × 8 TPU slices.

Revised: 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.

Revised: We train using the sigmoid cross-entropy loss function and L2 weight decay with multiplier 0.03.

Revised: The heteroscedastic method uses a temperature of 0.4, 1,000 Monte Carlo samples and R = 50 for the low rank approximation.

Revised: HetSNGP has the same heteroscedastic hyperparameters except the optimal temperature is 1.5.

Revised: For SNGP and HetSNGP the GP covariance is approximated using the momentum scheme presented in Liu et al.

Revised: (2020) with momentum parameter 0.999.

Original: Further details on our use of the Laplace approximation are provided in Appendix A.
Revised: In this section, we will derive the Laplace posterior in Eq. 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

Revised: [Equation]

Revised: where Z is a normalization constant that does not depend on β c .

Revised: Following Rasmussen & Williams (2006), we will denote the unnormalized log posterior as

Revised: [Equation]

Revised: Recall that the ﬁrst term is the likelihood and the second term is our prior from Eq.

Revised: The Laplace approximation now approximates the posterior with a local second-order expansion around the MAP solution, that is

Revised: [Equation]

Revised: with the MAP solution ˆ β c = arg max β c Ψ( β c ) and the Hessian Λ = −∇

Revised: [Equation]

Revised: Ψ( β

Revised: [Equation]

Revised: ) | β c =ˆ β c .

Revised: The MAP solution can be found using standard (stochastic) gradient descent, while the Hessian is given by

Revised: [Equation]

Revised: 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, W ii = −∇ 2 log p ( y i | u i ) (Rasmussen & Williams, 2006).

Revised: For instance, in the case of the logistic likelihood, W ii = p i (1 − p i ) , where p i is a vector of output probabilities for logits u i .

Revised: To get the Hessian at the MAP, we then just need to compute this quantity for ˆ u

Revised: = Φ ˆ β .

Revised: The approximate posterior is therefore

Revised: [Equation]

Revised: where the precision matrix can be computed over data points (recovering Eq.

Revised: (5)) as

Revised: [Equation]

Revised: A. A DDITIONAL EXPERIMENTAL RESULTS

Revised: We add additional baselines for our Imagenet experiment from the epistemic uncertainty literature.

Revised: In particular we benchmark MC Dropout (Gal & Ghahramani, 2016) and MIMO (Havasi et al., 2020).

Revised: We use the implementations and Imagenet hyperparameters from the Uncertainty Baselines codebase (Nado et al., 2021).

Revised: [Table]

Revised: [Table]

Revised: [Table]

Revised: [Table]

Revised: [Table]