A Simple Approach to Improve Single-Model Deep Uncertainty via Distance-Awareness
Jeremiah Zhe Liu, Shreyas Padhy, Jie Ren, Zi Lin, Yeming Wen, Ghassen Jerfel, Zack Nado, Jasper Snoek, Dustin Tran, Balaji Lakshminarayanan
Introduction
In recent years, deep neural network (DNN) models have become ubiquitous in large-scale, real-world applications. The examples range from self-driving, image recognition, natural language understanding, to scientific applications including genomic sequence identification, genetic variant calling, drug discovery and protein design (Larson et al., 2019; Gupta et al., 2021; Jumper et al., 2021; Poplin et al., 2018; Ren et al., 2019; Han et al., 2021; Kivlichan et al., 2021; Roy et al., 2022). A key characteristic shared by these real-world tasks is their risk sensitivity: a confidently wrong decision from the deep learning model can lead to ethical violations, misleading scientific conclusions, and even fatal accidents (Amodei et al., 2016). Therefore, to ensure safe and responsible deployment of AI technologies to the real world, it is of utmost importance to develop efficient approaches that reliably improves a deep neural network’s uncertainty quality without compromising its practical utility (i.e., in terms of accuracy and scalability).
It is well-known that a naively trained modern deep network tends to perform poorly in uncertainty tasks. They can be poorly calibrated (Guo et al., 2017) or assign high confidence predictions to out-of-domain (OOD) inputs (Nguyen et al., 2015; Hendrycks and Gimpel, 2017; Lakshminarayanan et al., 2017). This has led to the development of probabilistic methods tailored for deep neural networks, with examples including Bayesian neural networks (BNNs) (Blundell et al., 2010; Osawa et al., 2019; Wenzel et al., 2020), Monte Carlo (MC) Dropout (Gal and Ghahramani, 2016), and Deep Ensembles (Lakshminarayanan et al., 2017). A shared characteristic of these approaches is that they are ensemble-based methods: they need to maintain distribution samples over millions of model parameters, or require multiple forward passes to produce a final prediction. Consequently, despite their success in academic studies, these approaches can be computationally prohibitive for real-world usage (Ovadia et al., 2019; Wilson and Izmailov, 2020). Furthermore, ensembles and BNNs may not necessarily fix all the problems with the underlying neural network in the first place. For instance, if all the ensemble members consistently make same mistakes or produce high confidence predictions far away from the data, the ensemble would inherit this behavior and also produce high confidence predictions far away from the data.
In this work, we seek to address this gap by exploring an alternative direction to improve the uncertainty quality of a single DNN, so that the model achieves high-quality uncertainty with only a single, deterministic representation. That is, instead of improving uncertainty by ensembling over multiple representations, we focus on addressing the design choices in a single DNN that hinders its uncertainty performance. Specifically, we study principled inductive biases (i.e., mathematical conditions) that a model’s hidden representation and output layer should satisfy to obtain good uncertainty performance. Then, we propose a simple, general-purpose algorithm to implement such inductive bias into a deep neural network. We expect this direction to be fruitful: an effective single-model uncertainty approach unlocks high-quality uncertainty estimation in resource-constrained, real-world settings where Monte Carlo sampling is not feasible (e.g., on-device, real-time prediction), and also can be combined with existing ensemble techniques to produce new probabilistic models that improves the state-of-the-art.
Specifically, we first propose distance awareness, i.e., a model’s output prediction is aware of the distance between a new test example and the previously trained-upon examples, as an important property for the neural network to achieve high-quality uncertainty. Formally, this means that a predictive model has the ability to output a scalar uncertainty estimate that is monotonic with respect to an appropriate distance metric between and the training examples (Definition 2). In the literature, several prior investigations have suggested the connection between a neural model’s uncertainty quality and its ability to distinguish distances in the input space. For example, Kristiadi et al. (2020) shows that deep classifiers degrades in performance and becomes increasing overconfident when making predictions on examples that further from the support of the training set. On the other hand, van Amersfoort et al. (2021) and Behrmann et al. (2021) observe that deep models often collapse the representations of distinct examples onto the same location in the hidden space (i.e., “feature collapse”), preventing the model from distinguishing between the familiar and unfamiliar examples and causing difficulty in detecting OOD inputs. In this work, we formalize the intuition behind these empirical findings by providing a precise definition of the distance awareness. We further provide rigorous arguments for the importance of the distance awareness property for a model’s uncertainty performance, by casting uncertainty estimation as a minimax decision making problem and showing that distance awareness constitutes a necessary condition for obtaining the optimal solution (Section 2).
We then move to identify simple, general-purpose algorithm to implement the distance-awareness principle into a DNN model. Figure 1 gives an high level overview. We take inspiration from the classic probablistic learning literature, where a shallow Gaussian process (GP) model (when equipped with certain kernel) is known to perfectly achieve distance awareness (Rasmussen and Williams, 2006). However, the shallow GP model tends to generalize poorly for high-dimensional data, since the common choices of kernel function lack the ability to adaptively capture the intrinsic structure underlying the data’s high surface dimension, leading to the issue of curse of dimensionality (Bach, 2017). Later works mitigate this issue by equipping the GP with a dimension-reduction feature extractor (e.g., implemented by a DNN) (Salakhutdinov and Mnih, 2008; Damianou and Lawrence, 2013; Wilson et al., 2016a, b; Calandra et al., 2016; Salimbeni and Deisenroth, 2017; Bradshaw et al., 2017). However, naively combining a DNN with a GP layer does not automatically guarantee distance awareness even with end-to-end training, as the hidden representation can still suffer the issue of feature collapse (as we will show in experiments, c.f. Figures 1 and 4).
To this end, we propose a simple and efficient algorithm for probabilistic deep learning which we term Spectral-normalized Neural Gaussian Process (SNGP). As shown in Figure 1, given an existing DNN architecture, SNGP makes two simple modifications to the neural network, namely
adding spectral normalization to the model’s hidden layers, and
replacing the typical dense output layer with a distance-aware Gaussian Process.
We show that spectral normalization improves the distance-preservation property of learned representations by bounding the hidden-space distance with respect to {\color[rgb]{0,0,0}{\color[rgb]{0,0,0}d_{X}}(\bx,\bx^{\prime})}, where and are two inputs to the feature extractor of the DNNs , and is a suitable distance metric defined for the data manifold (Proposition 3). We then build a GP output layer on these representations . To ensure computational scalability, we use a Laplace approximation to the random feature expansion of the GP. This results in a model posterior that can be learned scalably and in closed form, and lets us efficiently compute the predictive uncertainty on individual inputs without having to resort to computationally expensive methods such as multiple forward passes (Section 3).
We conduct a comprehensive study to investigate the behavior of SNGP model across data modalities. As we will show, the SNGP approach improves the calibration and OOD detection performance of a deterministic DNN, and outperforms or is competitive with other single model uncertainty approaches. We further illustrate the method’s scalability by adapting it to large-scale recognition tasks (i.e., ImageNet), and illustrate the method’s generality by extending it to additional tasks in natural language understanding and genomics sequence identification (Section 6). Finally, we show that SNGP can serve as a strong building block for probabilistic deep learning approaches and provides complementary benefits to other state-of-the-art techniques such as ensembling (e.g., MC Dropout, Deep Ensemble) and data augmentation (e.g., AugMix), providing orthogonal improvements in uncertainty quantification (Section 7).
Theoretical Motivation for Distance Awareness
Let us consider a data-generating distribution , where is the -dimensional simplex of labels, and is the input data present on a manifold with a suitable metric . In particular, such a metric is tailored toward the geometry of such that, intuitively, the distance between a pair of examples reflects a meaningful difference in the input space (e.g., for a pair of sentences , describes their semantic similarity rather than the token-level edit distance (Cer et al., 2017)). In a supervised learning setup, the goal is often to learn the conditional distribution , and the training data is often a subset of the full input space . Due to this fact, we can represent the conditional data-generating distribution as a mixture of an in-domain (IND) distribution and an OOD distribution with non-overlapping support (Meinke and Hein, 2020; Scheirer et al., 2014):
During the process of training, the model learns the in-domain predictive distribution from the training data , but does not have knowledge about .
In the example of an image classification model trained on MNIST, the out-of-domain space is the space of all images that do not contain handwritten characters with the classes 0 to 9, which can include other datasets such as CIFAR-10 or ImageNet, with the potential for some overlap whenever numbers are present in such image examples.
In the example of a weather-service chatbot, the out-of-domain space is the space of all natural utterances not related to weather queries, whose elements usually do not have a meaningful correspondence with the in-domain intent labels . The out-of-domain distribution can in general be very different from the in-domain distribution , and it is usually expected that the model will only generalize well within . However, during testing and deployment, the model is expected to construct a predictive distribution for the entire input space, , since the gamut of input data that the model can be deployed on can come from anywhere, and not just a curated training distribution.
1 Uncertainty Estimation as a Minimax Learning Problem
In order to formulate uncertainty estimation as a learning problem under (1), we need to define a loss function to measure a model ’s quality of predictive uncertainty. One popular metric, the Expected Calibration Error (ECE), is defined as C(p,p^{*})=E\big{[}|E(y^{*}=\hat{y}|\hat{p}=p)-p|\big{]}, and measures the difference in expectation between the model’s predictive confidence (e.g., the maximum probability score) and its actual accuracy (Guo et al., 2017; Nixon et al., 2019). However, ECE is not suitable as a loss function, since it does not have a unique minimum at the true solution . Using the ECE directly can result in trivial counterexamples where a predictor ignores the input example and achieve perfect calibration by predicting randomly according to the marginal distribution of the labels (Gneiting et al., 2007).
To this end, a more theoretically sound uncertainty metric needs to be obtained, which we accomplish by examining the rich literature of strictly proper scoring rules (Gneiting and Raftery, 2007) , which are loss functions that are uniquely minimized by the true distribution . These family of loss functions include a lot of the more commonly used examples such as log-loss and Brier score. Another nice property of proper scoring rules is that they’re related to ECE; both the log-loss and the Brier score are upper bounds of the calibration error, and this can be shown by the classic calibration-refinement decomposition (Bröcker, 2009). Therefore, it follows that if a proper scoring rule is minimized, it implies that the calibration error of the model is also being minimized. We can now formalize the problem of uncertainty quantification as the problem of constructing an optimal predictive distribution that minimizes the expected risk over all , i.e., an Uncertainty Risk Minimization problemIt is interesting to note that, as a special case, (LABEL:eq:brier_risk) reduces back to the familiar maximum-likelihood (MLE) objective when is the logarithm score. See Gneiting and Raftery (2007) for further detail. However, the analysis here goes beyond the scope of empirical risk parameter estimation, as it focuses on risk functional in terms of a infinite-dimensional parameter and concerns the out-of-domain situations where the training data is not available.: Unfortunately, we cannot minimize (LABEL:eq:brier_risk) over the entire input space , even with access to infinite amounts of in-domain data. This is due to the fact that during the training process, the data is collected only from , and the true OOD distribution can never be learned by the model, and therefore generalization is not guaranteed since we do not make the assumption that and are similar. In practice, using a model trained only with in-domain data to make predictions on OOD can lead to arbitrarily bad results, since nature can contain many OOD distributions that are very dissimilar with the training data. This is clearly undesirable for safety-critical applications.
We therefore reformulate the problem using a more prudent strategy; minimize instead the worst-case risk with respect to all possible , where is the space of possible data-generating distributions whose in-domain component generates the observational data, while whose out-of-domain component is unconstrained and can be arbitrary. That is, we seek to construct a to minimize the Minimax Uncertainty Risk:
where is the expected risk as defined in (LABEL:eq:brier_risk). This reformulation can be viewed from a game-theoretic lens; the uncertainty estimation task is acting as a two-player game with the model and nature, where the goal of the model is to produce a minimax strategy that minimizes the risk against all possible (even adversarial) moves of nature. Under the task of classification using the Brier score as a proper scoring rule, the solution to the minimax problem (2) adopts a simple and elegant form:
(3) has a very intuitive understanding; we should trust the model for an input point that lies in the training data domain, and otherwise make a maximum entropy (uniform) prediction.Noted that this uniform strategy is conservative and is derived under the minimax assumption that (1) the OOD distribution is adversarial, and (2) all the classes in are semantically distinct. That is, in the OOD regions, is always as far away from as possible and puts its probability mass on an class that is the most different from the model prediction. For example, for a cancer prognosis model, the puts the probability on a different cancer type that requires a completely different treatment. In practice, there exist situations where the OOD class is semantically related to an in-domain class (e.g., pickup truck v.s. truck for the CIFAR100 v.s. CIFAR10 OOD detection problem) such that the data is not completely adversarial. In this case, a non-uniform distribution is still sensible in the practical context. However that falls outside the scope of the minimax analysis that we consider here. For the practice of uncertainty estimation, (3) is conceptually important in that it verifies that there exists a unique optimal solution to the uncertainty estimation problem (2). Furthermore, this optimal solution can be constructed conveniently as a mixture of a discrete uniform distribution and the in-domain predictive distribution that the model has already learned from data, assuming one can quantify well. In fact, the expression (3) can be shown to be optimal for a broad family of scoring rules known as the Bregman scores, which includes the Brier score and the widely used log score as the special cases.
The proof for (3) relies on the following key lemma (proved in Appendix E.2):
Consider the Bregman score in (132). At a location where the model has no information about other than , the solution to the minimax problem
is the discrete uniform distribution, i.e., . Using Lemma 2.1, (3) can be proved easily by decomposing into its in-domain and out-of-domain components and apply standard minimax arguments. A proof sketch is included below:
Proof Sketch. Denote . Decompose the overall Bregman risk by domain:
where we have denoted S_{\texttt{IND}}(p,p^{*})=E_{\bx\in\Xsc_{\texttt{IND}}}\big{(}s(p,p^{*}|\bx)\big{)} and S_{\texttt{OOD}}(p,p^{*})=E_{\bx\in\Xsc_{\texttt{OOD}}}\big{(}s(p,p^{*}|\bx)\big{)}. Noting that and have disjoint support, we can decompose the minimax risk as follows:
Since the model’s predictive distribution is learned from data, the in-domain minimax risk is fixed. Therefore, we only need to show is the optimal and unique solution to the out-of-domain minimax risk \inf_{p}\sup_{p^{*}}\big{[}S_{\texttt{OOD}}(p,p^{*})\big{]}. To this end, notice that for a given :
due to the fact that we don’t impose assumption on (therefore is free to attain the global supreme by maximizing at every single location ). Furthermore, there exists that minimize at every location of , then it minimizes the integral (Berger, 1985). By Lemma 2.1, such exists and is unique, i.e.:
In conclusion, we have shown that is the unique solution to , and therefore that the unique solution to (4) is (3).
A complete version of the proof is available in Appendix B.
2 Distance Awareness as a Necessary Condition
Following from Equation (3), it stands to reason that a key recipe for a deep learning model to be able to reliably estimate predictive uncertainty is its ability to quantify (explicitly or implicitly) the domain probability . This requires that the model have a good notion of the distance (or dissimilarity) between a testing example and the training data with respect to a suitable metric for the data manifold (e.g., semantic textual similarity (Cer et al., 2017) for language data). Definition 2 makes this notion more precise:
Consider a predictive distribution trained on a domain , where is the input data manifold equipped with a suitable metric . We say is input distance aware if there exists a summary statistic of that quantifies model uncertainty (e.g., entropy, predictive variance, etc) and reflects the distance between and the training data with respect to , i.e.,
Here, is a monotonic function and is the distance between and the training data domain.
for positive and bounded constants . For a deep learning model, the bi-Lipschitz condition (6) usually leads the model’s hidden space to preserve a meaningful distance in the input data manifold that is, e.g., effective for determining if an observation is in-distribution. This is due to the fact that the upper Lipschitz bound is an important condition for the adversarial robustness of a deep network, which prevents the hidden representations from being overly sensitive to the meaningless perturbations in the pixel space (e.g., Gaussian noise) (Ruan et al., 2018; Weng et al., 2018; Tsuzuku et al., 2018; Jacobsen et al., 2019a; Sokolic et al., 2017). On the other hand, the lower Lipschitz bound prevents the hidden representation from collapsing the representations of distinct examples together, which otherwise leads to undesired invariance to the meaningful differences between examples (i.e., the concept of feature collapse) (Jacobsen et al., 2019b; Van Amersfoort et al., 2020). When we combine these two inequalities and the properties they enforce, the bi-Lipschitz condition essentially encourages to be an approximately isometric mapping, thereby ensuring that the learned representation space has a robust and meaningful correspondence with the geometry of the input data manifold . Heuristically as well, machine learning methods usually tend to attempt to learn an approximately isometric and geometry-preserving mapping (Hauser and Ray, 2017b; Perrault-Joncas and Meila, 2012; Rousseau et al., 2020). For example, deep image classifiers usually strive to learn a mapping from the image manifold to a hidden representation space where the input data is easily separable using a set of linear decision boundaries. Similarly, sentence encoders aim to project sentences into a vector space where the cosine distance can be used to measure the semantic textual similarity between natural language sentences (Cer et al., 2018). It has also been shown that preserving such approximate isometry in a neural network is possible even after significant dimensionality reduction (Blum, 2006).
Our Proposed Method: Spectral-normalized Neural Gaussian Process (SNGP)
In this section we formalize the definition of the Spectral-normalized Neural Gaussian Process (SNGP) algorithm for modern residual-based DNN (e.g., ResNet, Transformer). The SNGP algorithm provides a simple approach to encode distance awareness into the output layer, and distance preservation into the hidden layers (i.e., the two properties introduced in Section 2.2). We summarize the method in Algorithms 1-2.
SNGP achieves distance-awareness in the output layer by replacing a dense output layer with an approximate Gaussian process (GP) conditioned on the learned hidden representations of the DNN, where the posterior variance of a test input is proportional to its distance from the training datapoints in the hidden spaceIn this work, we focus on the RBF kernel for its simplicity, however it is easy to extend our framework to other kernels (e.g., Matérn, MLP, or arccos kernel, etc) by modifying the activation function and the distribution of the random-feature mapping in Equation (9) (Choromanski et al., 2018; Liu et al., 2021)..
Given a training dataset with data points, we denote the hidden representation of the DNN defined for each training point as , and denote the Gaussian Process conditioned on the hidden representation as . Then, the prior distribution of a GP model equipped with an RBF kernel is a multivariate normal:
where is the kernel amplitude (Rasmussen and Williams, 2006). The model’s posterior distribution is then calculated using Bayes Rule:
However, in the context of large-scale neural modeling, performing exact GP inference (7)-(8) is difficult for two reasons: (i) inference with the exact prior (7) is computationally intractable, due to the need of inverting a kernel matrix which has a cubic time complexity . (ii) Inference with a non-Gaussian likelihood in (8) is analytically intractable, due to the difficulty of deriving the analytical integration over a non-conjugate likelihood. The existing literature commonly handles (i)-(ii) by deploying sophisticated inference algorithms such as structured Variational Inference (VI) or Markov chain Monte Carlo (MCMC) (see Section 5 for a full review). However, this invariably imposes nontrivial engineering and computational complexity to an existing deep learning system, rendering itself infeasible in many practical scenarios where the system scalability and maintainability is of high importance.
In this work, we propose to tackle the above challenge by performing Laplace-approximation inference to the random Fourier features (RFFs) expansion of the GP model (Rasmussen and Williams, 2006), by (i) converting the GP model (7) into a featurized Bayesian linear model via the random-feature expansion (Rahimi and Recht, 2008b). Then (ii) approximate the intractable posterior (8) using Laplace approximation (Gelman et al., 2013). Both are well-established methods in the statistical machine learning literature with rigorous theoretical guarantees (Rahimi and Recht, 2008a; DeGroot and Schervish, 2012). We re-iterate that the goal here is not to develop another algorithm to approximate the exact GP posterior, but to identify a simple, practical baseline approach to implement the distance-awareness property (Definition 2) into a neural model with minimal modification to the training pipeline for a deterministic DNN. To this end, as we will show, the proposed approach leads to a closed-form posterior that can be trained end-to-end using the same pipeline for a deterministic DNN and with comparable time complexity. Empirically, the trained model maintains the generalization performance of a deterministic DNN while illustrating improved uncertainty performance in calibration and in OOD detection.
First, we approximate the prior defined in (7) by defining a low-rank approximation of the kernel matrix by using random features (Rahimi and Recht, 2008b), which gives us a random-feature Gaussian process:
Here, represents the hidden representations of the penultimate layer with a dimensionality of . is the final layer with dimension , which is represented by (1) a fixed weight matrix whose entries are sampled i.i.d from , and (2) a fixed bias term whose entries are sampled i.i.d from . We can therefore write the logits using the RFF expansion to the GP prior in (7) as a neural network layer consisting of fixed hidden weights and learnable output weights :
Notice that conditional on , the output weights are the only learnable parameters in the model. As a result, the random Fourier feature (RFF) approximation in (10) reduces an infinite-dimensional GP to a standard Bayesian linear model, for which many posterior approximation methods (e.g., expectation propagation (EP)) can be applied (Minka, 2001). In this work, we choose the Laplace method due to its simplicity and the fact that its posterior variance has a convenient closed form (Rasmussen and Williams, 2006). Briefly, the Laplace method approximates the model posterior using a Gaussian distribution that is centered around the maximum a posterior (MAP) estimate , such that is the Hessian matrix of the log posterior likelihood evaluated at the MAP estimates. For a binary classification task, the posterior precision matrix (i.e., the inverse covariance matrix) adopts a simple expression , where is the model prediction under the MAP estimates (Rasmussen and Williams, 2006). We introduce the extensions to regression and multi-class classification in Appendix A.1.
To summarize, for a classification task, the Laplace posterior for an approximate GP under the RFF expansion is:
where is the model’s MAP estimate conditioned on the RFF hidden representation in (9), and is the prior variance. To obtain the posterior distribution (11), one only need to first obtain the MAP estimate by training the entire network with respect to the MAP objective using stochastic gradient descent (SGD):
where are the hidden weights of the network and is the negative log likelihood for the task (e.g., the cross entropy loss for a classification task). Then, in the final epoch, we can compute this covariance in an incremental fashion by initializing covariance with and accumulating covariance contributions from each batch as in (11) i.e., to update the posterior precision matrix as for minibatches of size . Notice that to obtain a MAP estimate that is strictly conditioned on the hidden representations , one should freeze the hidden representations in the final epoch and only update the output weights . However in practice, we find this has no significant impact on the final performance. This is most likely due to the fact that the hidden representation has already become stabilized in the final epochs. In the online learning setting where the model is processing a infinite stream of data, this expression can be modified to where is a small scaling coefficient. This is analogous to the exponential moving average estimator for batch variance as used in batch normalization (Ioffe and Szegedy, 2015). As a result, the approximate GP posterior (11) can be learned scalably and in closed-form with minimal modification to the training pipeline of a deterministic DNN. It is worth noting that under the RFF posterior, the Laplace approximation is in fact asymptotically exact by the virtue of the Bernstein-von Mises (BvM) theorem and the fact that (10) is a finite-rank model (Freedman, 1999; LeCam, 1973; Panov and Spokoiny, 2015; Dehaene, 2019).
2 Approximately Distance-preserving Hidden Mapping via Spectral Normalization
Consider a hidden mapping with residual architecture where and and are of equal dimension. If for , all ’s are -Lipschitz, i.e., . Then:
where and , i.e., is distance preserving.
Proof is in Appendix E.1. The ability of a residual network to construct a geometry-preserving metric transform between the input space and the hidden space is well-established in learning theory and generative modeling literature, but the application of these results in the context of uncertainty estimation for DNN appears to be new (Bartlett et al., 2018; Behrmann et al., 2019; Hauser and Ray, 2017b; Rousseau et al., 2020).
However in practice, a strict preservation of distance is both impossible and likely undesirable. The reason is that per common neural network practice, the model often projects a high-dimensional example into a lower-dimensional representation (i.e., for the hidden mapping ). This necessitates information loss and precludes the possibility of exact isometry (i.e., invertibility) (Smith et al., 2021). Furthermore, an overly strict bound on the Lipschitz constant of can push the model toward identity mapping, greatly restricting the expressiveness of the network and leading to suboptimal generalization (Behrmann et al., 2019). To this end, a more prudent strategy would be to pursue approximate distance preservation, i.e., by finetuning the magnitude of Lipschitz constant to a more relaxed extent, so as to balance the tradeoff between the expressiveness of the network and its distance-preservation ability (see Section 8.1 for further discussion).
Algorithm-wise, to ensure the hidden mapping is approximately distance preserving, it is sufficient to ensure that the weight matrices for the nonlinear residual block to have spectral norm (i.e., the largest singular value) to be upper-bounded, since . In this work, we enforce the aforementioned Lipschitz constraint on ’s by applying the spectral normalization (SN) on the weight matrices as recommended in Behrmann et al. (2019). Briefly, at every training step, the SN method first estimates the spectral norm using the power iteration method (Gouk et al., 2021; Miyato et al., 2018), and then normalizes the weights as (13) where is a hyperparameter used to adjust the exact spectral norm upper bound on (so that ). Therefore, (13) allows us more flexibility in controlling the spectral norm of the neural network weights so it is the most compatible with the architecture at hand. In this work, we treat as a tunable hyperparameter which can be selected based on the validation data (Section A.2).
We summarize the method in Algorithms 1-2. As shown, during training, the model updates the hidden-layer weights and the trainable output weights via minibatch SGD (i.e., exactly the same way as a deterministic DNN). Then, in the final epoch, it also performs an update of the precision matrix using Equation (11). During inference, the model first performs the conventional forward pass to compute the final hidden features , and then compute the posterior mean (time complexity ) and the predictive variance (time complexity ) To extend Algorithm 1-2 to regression or multi-class classification task, one just need to use the appropriate precision matrix update as introduced in A.1, and replace the sigmoid output activation to identity or softmax.. To see how this predictive distribution of SNGP reflects the distance-awareness property, notice that conditional on the penultimate embedding and assuming squared loss, the predictive logit of the SNGP follows a Gaussian process:
where is the random feature embedding of the training data. Then, straightforward application of Woodbury matrix identity reveals that (14) can equivalently be expressed in the familiar dual form of the Gaussian process (Rasmussen and Williams, 2006):
where , and are kernel matrices approximating those under the RBF kernel . Consequently, under a distance-preserving mapping , for a test example that is moving away from the training-data manifold, its predictive kernel matrix systematically approaches 0, while the testing-data kernel and the training-data kernel remain in a constant range. This causes the predictive mean to approach zero, and the predictive variance to approach its maximum This conclusion also holds for non-Gaussian outcome. Where the predictive mean and variance can be expressed as and , where is the model prediction on the training data, and is the inverse covariance matrix calculated using the Laplace method (i.e., via Equation (11)). (Rasmussen and Williams (2006), Chapter 3.6). As a result, the SNGP model achieves the behavior as motivated in (3), i.e., generating a maximum-entropy predictive distribution for OOD inputs that are far outside the training domain. We include a proof of this behavior in Appendix C.3.
where is a scaling factor that is commonly set to (Lu et al., 2020). In our experiments, we have found the two approaches to perform similarly. Therefore we report the results from mean-field approximation in our experiments due to its wider applicability.
Interestingly, Equation (17) highlights the role of the kernel amplitude in improving a classifier’s calibration performance: as the predictive variance is proportional to (see Equation (16)), it effectively re-scales the magnitude of the logits and functions similarly to the temperature parameter in the temperature scaling technique (Guo et al., 2017). We also note that in our experiment, the value of generally does not have a significant impact on model performance in OOD detection. This is likely due to the fact that scaling the model uncertainty by a constant does not change the ranking of the uncertainty scores. In practice, we recommend estimating on a small amount of in-distribution validation data by minimizing with respect to a proper scoring rule. This is because an overparameterized model such as a DNN is known to overfit the training data, rendering the in-sample estimate of the second-order statistics not reliable (Wahba, 1990). In all experiments, we estimate by minimizing it against the logarithm score (i.e., the marginalized negative log likelihood (Gneiting et al., 2007)) on a small held-out validation dataset. We summarize all other model hyperparameters (e.g., the spectral norm upper bound ) and provide practical recommendations in Appendix A.2.
Where does SNGP fit in the current landscape of methods?
We have introduced SNGP as a method to improve single model uncertainty. Before describing the experimental setup, we first explain where SNGP fits in the current landscape of uncertainty methods in deep learning.
SNGP is a method for deterministic uncertainty quantification, that is, given a single, non-random representation, SNGP improves the base model by enhancing its distance-awareness property. This is a orthogonal direction to what was taken by popular ensemble-based approaches (e.g., Monte-Carlo dropout, BatchEnsemble (Wen et al., 2020) or Deep Ensemble), who improve performance primarily through integrating over multiple diverse representations (Fort et al., 2019). However, when base models consistently make overconfident predictions far away from the data, their ensemble can inherit such behavior as well. As a result, while ensembling vanilla DNNs are effective in improving accuracy and calibration under shift, they can be lacking in providing as significant a boost in OOD detection (e.g., due to the lack of distance awareness).
To see this point visually, in Figure 2 where we highlight the two orthogonal axes of improvement. The y-axis denotes better ensembling strategies: the left column shows that ensembling vanilla DNNs does not necessarily change the confidence landscape far away from the inputs. Even though efficient ensembles reduce compute and memory, they are still solving a different problem from SNGP by improving representation diversity but in an efficient manner. The x-axis denotes better single model uncertainty; the bottom row shows how DNN-GP and SNGP improve deterministic uncertainty quantification by improving distance preservation and distance awareness.
Finally, we note that these two axes provide complementary benefits, so they can be combined to further improve performance (see SNGP ensemble in top right). Since SNGP primarily improves single model uncertainty, we focus our comparisons with other methods for deterministic uncertainty quantification in the experiments (Section 6). Whenever possible, we recommend to ensemble SNGP models build toward a strong probabilistic DNN model that simultaneously quantifies representation uncertainty and ensures distance awareness. We explore ensembles of SNGP in Section 7.1. As an aside, data augmentation is another orthogonal axis of improvement (Hendrycks* et al., 2020). We also explore SNGP as a building block for data augmentation methods in Section 7.2.
Related Work
Single-model approaches to deep classifier uncertainty. Recent work examines uncertainty methods that add few additional parameters or runtime cost to the base model. The state-of-the-art on large-scale tasks are efficient ensemble methods (Wen et al., 2020; Dusenberry et al., 2020), which cast a set of models under a single one, encouraging independent member predictions using low-rank perturbations. These methods are parameter-efficient but still require multiple forward passes from the model. SNGP investigates an orthogonal approach that improves the uncertainty quantification by imposing suitable regularization on a single model, and therefore requires only a single forward pass during inference. There exists other runtime-efficient, single-model approaches to estimate predictive uncertainty, achieved by either replacing the loss function (Hein et al., 2019b; Malinin and Gales, 2018a, b; Sensoy et al., 2018b; Shu et al., 2017), the output layer (Bendale and Boult, 2016; Tagasovska and Lopez-Paz, 2019; Calandra et al., 2016; Macedo et al., 2020; Macêdo and Ludermir, 2021; Padhy et al., 2020), computing a closed-form posterior for the output layer (Riquelme et al., 2018; Snoek et al., 2015; Kristiadi et al., 2020) , or predicting a-priori uncertainty away from training data (Skafte et al., 2019). SNGP builds on these approaches by also considering the intermediate representations which are necessary for good uncertainty estimation, and proposes a simple method (spectral normalization) to achieve it. A recent method named Deterministic Uncertainty Quantification (DUQ) also regulates the neural network mapping but uses a two-sided gradient penalty (Van Amersfoort et al., 2020). However, empirically, the two-sided gradient penalty can lead to unstable training dynamics for a deep residual network, and is observed to over-constrain the model capacity in more difficulty tasks (e.g., CIFAR-100) in our experiments (Section 6.2.1). Following the initial conference publication of this work (Liu et al., 2020a) and in a similar vein, some later work also investigated building other class of probabilistic models on top of a spectral-regularized network, e.g., variational Gaussian process, Gaussian mixture model, or stochastic differential equations (Mukhoti et al., 2021; van Amersfoort et al., 2021; Cui et al., 2021). While obtaining encouraging results on small-scale benchmarks (e.g., CIFAR), these approaches are still computationally prohibitive for large-scale tasks with high number of output classes (e.g., ImageNet) which SNGP can easily scale to (Section 6.2.2).
Laplace approximation and GP inference with DNN. Laplace approximation has a long history in GP and (Bayesian) NN literature (Tierney et al., 1989; Denker and LeCun, 1991; Rasmussen and Williams, 2006; MacKay, 1992; Ritter et al., 2018; Hobbhahn et al., 2022; Kristiadi et al., 2022, 2021; Eschenhagen et al., 2021; Daxberger et al., 2021), and the theoretical connection between a Laplace-approximated DNN and GP has being explored recently (Khan et al., 2019). Differing from these works, SNGP applies the Laplace approximation to the posterior of a neural GP, rather than to a shallow GP or a dense-output-layer DNN. Earlier works that combine a GP with a DNN learns the hidden representation separately Salakhutdinov and Hinton (2007), or performs end-to-end learning via MAP estimation (Calandra et al., 2016) or structured VI (Hensman et al., 2015; Bradshaw et al., 2017; Wilson et al., 2016b). These approaches were shown to lead to poor calibration by recent work (Tran et al., 2019), which proposed a simple fix by combing Monte Carlo Dropout (MC Dropout) with random Fourier features, which we term MC Dropout Gaurssian Process (MCD-GP). SNGP differs from MCD-GP in that it considered a different regularization approach (spectral normalization) and can compute its posterior uncertainty more efficiently in a single forward pass. We compare with MCD-GP in our experiments (i.e., the DNN-GP + Dropout method in Section 7).
Distance-preserving neural networks and bi-Lipschitz condition. The theoretic connection between distance preservation and the bi-Lipschitz condition is well-established (Searcod, 2006), and learning an approximately isometric, distance-preserving transform has been an important goal in the fields of dimensionality reduction (Blum, 2006; Perrault-Joncas and Meila, 2012), generative modeling (Lawrence and Quinonero-Candela, 2006; Dinh et al., 2014, 2016; Jacobsen et al., 2018), and adversarial robustness (Jacobsen et al., 2019a; Ruan et al., 2018; Sokolic et al., 2017; Tsuzuku et al., 2018; Weng et al., 2018). This work is a novel application of the distance preservation property for uncertainty quantification. There are several methods for controlling the Lipschitz constant of a DNN (e.g., gradient penalty or norm-preserving activation (An et al., 2015; Anil et al., 2019; Chernodub and Nowicki, 2017; Gulrajani et al., 2017)), and we chose spectral normalization in this work due to its simplicity and its minimal impact on a DNN’s architecture and the optimization dynamics (Bartlett et al., 2018; Behrmann et al., 2019; Rousseau et al., 2020; Behrmann et al., 2021). Finally, Smith et al. (2021) studied the effect of spectral regularization on the residual networks in the context of image recognition, and concluded that, under mild assumptions, residual networks will be approximately distance preserving on the low-passed portion of their input.
Open Set Classification. The uncertainty risk minimization problem in Section 2 assumes a data-generation mechanism similar to the open set recognition problem (Scheirer et al., 2014), where the whole input space is partitioned into known and unknown domains. However, our analysis is unique in that it focuses on measuring a model’s behavior in uncertainty quantification and takes a rigorous, decision-theoretic approach to the problem. As a result, our analysis works with a special family of risk functions (i.e., the strictly proper scoring rule) that measure a model’s performance in uncertainty calibration. Furthermore, it handles the existence of unknown domain via a minimax formulation, and derives the solution by using a generalized version of maximum entropy theorem for the Bregman scores (Grünwald and Dawid, 2004; Landes, 2015). The form of the optimal solution we derived in (3) takes an intuitive form, and has been used by many empirical work as a training objective to leverage adversarial training and generative modeling to detect OOD examples (Hafner et al., 2020; Harang and Rudd, 2018; Lee et al., 2018a; Malinin and Gales, 2018b; Meinke and Hein, 2020; Hendrycks et al., 2018). Our analysis provides theoretical support for these practices in verifying rigorously the uniqueness and optimality of this solution, and also provides a conceptual unification of the notion of calibration and the notion of OOD generalization. Furthermore, it is used in this work to motivate a design principle (distance awareness) that enables strong OOD performance in discriminative classifiers without the need of explicit generative modeling.
Benchmarking Experiments
In this section, we benchmark the performance of the SNGP model by applying it on a variety of both toy datasets and real-world tasks across different modalities Specifically, we first benchmark the behavior of the approximate GP layer versus an exact GP, and illustrate the impact of spectral normalization on toy regression and classification tasks (Section 6.1). We then conduct a thorough benchmark study to compare the performance of SNGP against the other state-of-the-art methods on popular benchmarks such as CIFAR-10 and CIFAR-100 (Section 6.2.1. Finally, we illustrate the scalability of and the generality of the SNGP approach by applying it to a large-scale image recognition task (ImageNet, Section 6.2.2), and highlight the broad usefulness by applying SNGP to uncertainty tasks in two other data modalities, namely conversational intent understanding and genomics sequence identification (Section 6.3).
We first benchmark the performance of the Random-Feature Gaussian process layer (RFGP), which form a key component of the SNGP model, versus both an exact GP formulation and a variational Gaussian process (VGP) (Titsias, 2009) on a 1D toy regression task. We consider the bimodal toy regression example from van Amersfoort et al. (2021), and use use the RBF kernel for all three models . As seen from Figure 3(a), the GP posterior mean for the exact GP tends to zero in the absence of training data, with low variance near training points and high variance otherwise. The RFGP mimics this behavior, with the addition of some periodic artifacts, which arises naturally due to the periodic nature of the approximation. In comparison, the prediction from the VGP model at locations that are further away from data tend to depend on the location of the inducing points, and does not revert back to the original zero-mean Gaussian process prior (e.g, between the two modes and at two ends of Figure 3). Although a more modern treatment of VGP may address this issue (Dutordoir et al., 2020) Note that the goal of this toy example is not to claim RFGP outperforms VGP, but just to illustrate that the RFGP, despite its simplicity, indeed returns a valid uncertainty surface that is distance-aware. To this end, van Amersfoort et al. (2021) studied the extension of SNGP method under VGP models, and obtained promising result on toy datasets. .
1.2 Classification
After having benchmarked the predictive and uncertainty behavior of the RFF-GP algorithm compared to an exact GP formulation, in the subsequent subsections we use it as a subcomponent in the SNGP algorithm.
In this subsection, we compare the behavior of SNGP on a suite of 2D classification tasks. Specifically, we consider the two ovals benchmark (Figure 4, row 1) and the two moons benchmark (Figure 4, row 2). The two ovals benchmark consists of two near-flat Gaussian distributions, which represent the two in-domain classes (orange and blue) that are separable by a linear decision boundary. There also exists an OOD distribution (red) that the model doesn’t observe during training. Similarly, the two moons dataset consists of two moon-shaped distributions separable by a non-linear decision boundary. For both benchmarks, we sample 500 observations from each of the two in-domain classes (orange and blue), and consider a deep architecture ResFFN-12-128, which contains 12 residual feedforward layers with 128 hidden units and dropout rate 0.01. The input dimension is projected from 2 dimensions to the 128 dimensions using a dense layer.
In addition to SNGP, we also visualize the uncertainty surface of the below approaches: Gaussian process (GP) is a standard Gaussian process directly taking as input and was trained with Hamiltonian Monte Carlo (HMC). In low-dimensional datasets, GP is often considered the gold standard for uncertainty quantification. Deep Ensemble is an ensemble of 10 ResFFN-12-128 models with dense output layers. MC Dropout uses a single ResFFN-12-128 model with dense output layer and 10 dropout samples at test time. DNN-GP uses a single ResFFN-12-128 model with the GP Layer (described in Section 3.1) without spectral normalization. Finally, SNGP uses a single ResFFN-12-128 model with the GP layer and with spectral normalization. The full experimental details for these algorihms are in Appendix C.
Figure 4 shows the results of training these algorithms to achieve high (close to accuracy) on the test data, and visualizes the uncertainty surface output by each model. Each algorithm returns a predictive distribution of the form , following which the confidence of the prediction can be written as . We plot the uncertainty surface by defining , and normalize it to the range of $\Xsc_{\texttt{IND}}\bx\Xsc_{\texttt{IND}}h(.)h(\bx)$ is free to discard information that is not relevant for prediction. On the other hand, SNGP is able to maintain the distance-awareness property via its bi-Lipschitz constraint, and exhibits a uncertainty surface that is analogous to the gold-standard model (exact GP) despite the fact that SNGP is based on a deep 12-layer network.
2 Image Classification
All methods included in the vision and language understanding experiments are summarized in Table 1. We compare SNGP against a suite of algorithms; a deterministic baseline, two single-model approaches: Deterministic Uncertainty Quantification (DUQ), and Enhanced Isotropic Maximization (IsoMax+). The original DUQ consists of two novel components, one is a RBF kernel based loss function which computes the distance between the feature vector and the class centroids, and the other is adding gradient penalty to avoid feature collapse. As pointed in the paper, RBF networks prove difficult to optimize and scale to large number of output classes (e.g., CIFAR-100), so we instead replace the RBF layer with our GP layer. Thus we call this variant as Deterministic Uncertainty Quantification with GP Last Layer (DUQ-GP). IsoMax+ replaces the softmax cross-entropy loss with a IsoMax+ loss which minimizes the distance to the correct class prototype based on the normalized feature embeddings in the last layer (Macêdo and Ludermir, 2021). We also include two ablations of SNGP: DNN-SN which uses spectral normalization on its hidden weights and a dense output layer (i.e. distance preserving hidden mapping without distance-aware output layer), and DNN-GP which uses the GP as output layer but without spectral normalization on its hidden layers (i.e., distance-aware output layer without distance-preserving hidden mapping). Finally, to evaluate the SNGP’s efficacy as a base model for ensemble approaches, we also train ensemble of SNGP and compare it with two popular approaches that quantifies representation diversity: MC Dropout (with 10 dropout samples) and Deep Ensemble (with 10 models), both are trained with a dense output layer and no spectral regularization. Section 7 further explores the interplay between SNGP and various ensemble approaches.
For all models that use GP layer, we keep and compute predictive distribution by performing mean-field approximation to the softmax Gaussian posterior. Further experiment details and recommendations for practical implementation are in Appendix C. All experiments are built on the uncertainty_baselines framework https://github.com/google/uncertainty-baselines (Nado et al., 2021). We open-source our code there.
2.1 CIFAR-10 and CIFAR-100
For the CIFAR-10 and CIFAR-100 image classification benchmarks, we use a Wide ResNet 28-10 model as the base for all methods (Zagoruyko and Komodakis, 2017). Following the benchmarking setup first suggested in Ovadia et al. (2019), we evaluate the model’s predictive accuracy, negative log-likelihood (NLL), and expected calibration error (ECE) under both clean CIFAR testing data and its corrupted versions termed CIFAR-*-C (Hendrycks and Dietterich, 2018). To evaluate the model’s OOD detection performance, we consider two tasks: a standard far-OOD task using SVHN as the OOD dataset for a model trained on CIFAR-10/-100, and a difficult near-OOD task using CIFAR-100 as the OOD dataset for a model trained on CIFAR-10, and vice versa. We use the in-distribution training data statistics to preprocess and normalize the OOD datasets. In Tables 2 and 3, we use the maximum softmax probability (MSP) as the uncertainty score while performing OOD evaluation, and report the Area Under the Reciever Operator Curve (AUROC) metric.
Tables 2-3 report the results. As shown, for predictive accuracy, SNGP is competitive with other single-model approaches. For calibration error, GP-based models clearly outperform the other single-model approaches and are competitive with Deep Ensemble and MC Dropout approaches. We also observe that under MSP metric, IsoMax+ method produces underconfident predictions, leading to the highest calibration error. Furthermore, its performance degrades quickly on more complex tasks (e.g., CIFAR-100) both in terms of accuracy and in uncertainty qualityContemporary to our work, Macêdo et al. (2022) shows the performance of IsoMax++ in OOD detection can be improved by using more specialized metrics (e.g., minimum distance score (MDS)) instead of MSP.. That is consistent with the findings of Padhy et al. (2020). For OOD detection, SNGP mostly outperforms other single-model approaches that are based on a dense output layer (except for CIFAR-10 vs. CIFAR-100 where IsoMax+ does best), and is competitive with deep ensembles, MC Dropout approaches and DUQ-GP. Finally, ensemble using SNGP as base model strongly outperforms all the other approaches, illustrating the importance of the distance-awareness property for high-quality performance in uncertainty quantification and the composability of SNGP as a building block toward the state-of-the-art probabilistic deep models.
In Appendix C.2 we studied the performance of other uncertainty scores including Dempster-Shafer (Sensoy et al., 2018a), Mahalanobis distance (Lee et al., 2018c), and Relative Mahalanobis distance (Ren et al., 2021) for OOD detection. We found in general the Dempster Shafer metric attains a better OOD performance (e.g., AUROC for CIFAR-10 v.s. SVHN) than the widely-used MSP metric reported in the main text.
2.2 ImageNet
We illustrate the scalability of SNGP by experimenting on the large-scale ImageNet dataset (Russakovsky et al., 2015) using a ResNet-50 model as the base for all methods. Similar to the CIFAR benchmarking setup, we follow the procedure from (Ovadia et al., 2019) and evaluate the model’s predictive accuracy and calibration error under both clean ImageNet testing data and its corrupted versions termed ImageNet-C. For this complex, large-scale task with high-dimensional output, we find it difficult to scale some of the other single-model methods (e.g., IsoMax or DUQ), which tend to over-constrain the model expressiveness and lead to lower accuracy than a baseline DNN (a phenomenon we already observe in the CIFAR-100 experiment). Therefore we omit them from results. Table 4 reports the results of methods that achieve competitive performance on ImageNet. As shown, for predictive accuracy, SNGP is competitive with that of a deterministic network, despite using a last-layer GP and spectral normalization on the weights of the residual network. For calibration error, SNGP not only outperforms the other single-model approaches and also, interestingly, the Deep Ensemble. Finally, the ensemble using SNGP as the base model attains the strongest predictive accuracy and NLL across all methods. These results illustrate the importance of the distance-awareness property for high-quality performance in uncertainty quantification, and the proposed components (i.e., spectral normalization and the GP-layer) perform well in complex, large-scale tasks by maintaining model accuracy while improving the quality of its uncertainty quantification.
3 Generalization to other data modalities
To validate the hypothesis that distance awareness is a crucial component for good predictive uncertainty on data modalities beyond images, we also evaluate SNGP on a practical language understanding task where uncertainty quantification is of natural importance: dialog intent detection (Larson et al., 2019; Vedula et al., 2019; Yaghoub-Zadeh-Fard et al., 2020; Zheng et al., 2020). In a goal-oriented dialog system (e.g. chatbot) built for a collection of in-domain services, it is important for the model to understand if an input natural utterance from a user is in-scope (so it can activate one of the in-domain services) or out-of-scope (where the model should abstain). To this end, we consider training an intent understanding model using the CLINC out-of-scope (OOS) intent detection benchmark dataset (Larson et al., 2019). Briefly, the OOS dataset contains data for 150 in-domain services with 150 training sentences in each domain, and also 1500 natural out-of-domain utterances. We train a model only on in-domain data, and evaluate their predictive accuracy on the in-domain test data, their calibration and OOD detection performance on the combined in-domain and out-of-domain data. Due the lack of empirically validated implementations of other single-model methods for non-image modalities, we focus on comparing ablated versions of SNGP and the two general-purpose methods: MC Dropout and Deep Ensemble. The results are in Table 5. As shown, consistent with the previous vision experiments, SNGP outperforms other single model approaches. It is competitive with Deep Ensemble in predictive accuracy and calibration, and outperforms all the approaches in OOD detection.
3.2 Bacteria Genomics Sequence Identification
Here we apply SNGP to a genomic sequence prediction task as another data modality. Ren et al. (2019) proposed the genomics OOD benchmark dataset motivated by the real-world problem of bacteria identification based on genomic sequences, which can be useful for diagnosis and treatment of infectious diseases. A classification model can be trained for classifying known bacteria species with decent test accuracy. However, when deploying the model to real data, the model will be inevitably exposed to the genomic sequences from unknown bacteria species. In fact, the real data can contain approximately of sequences from unknown classes that have not been studied before. The model needs to be able to detect those out-of-distribution inputs from unknown species, and abstain from making predictions for them. The genomic OOD benchmark dataset contains 10 bacteria classes as in-distribution for training, and 60 bacteria classes as out-of-distribution for testing. The in-distribution and out-of-distribution bacteria classes were chosen and separated naturally by the year of discovery, to mimic the real scenario that new bacteria species are discovered gradually over the years. Following (Ren et al., 2019), we train a 1D CNN but with SN and GP components added on the 10 in-domain classes and evaluate the models’ accuracy, ECE and NLL on the in-domain test data, and also evaluate the model’s OOD performance for detecting the 60 OOD classes. In addition to SNGP, we also consider Monte Carlo dropout and Deep Ensemble as another two baselines for comparison. We also study the effect of each component of SN and GP on the model’s calibration and OOD detection performance. Section C.1 contains further model detail.
The results are shown in Table 6. Comparing with all the other baseline models, SNGP model achieves the best OOD performance and best in-distribution accuracy, ECE, and NLL, among the single models. Each of the two components, SN and GP helps to reduce ECE and NLL, and improve OOD detection over the baseline DNN model. SNGP Ensemble is better than DNN Ensemble and MC Dropout in terms of in-distribution accuracy and NLL, and OOD detection.
SNGP as a Building Block for Probabilistic Deep Learning
From a probabilistic machine learning perspective, the SNGP algorithm provides an efficient way to learn a single high-quality deterministic model by, 1) improving the representation learning quality of the underlying neural network model through spectral normalization, and 2) introducing distance awareness in conjunction with an approximate GP random feature layer. However, in the literature, there exist other competitive and state-of-the-art probabilistic methods that improve a deep learning model’s predictive uncertainty via different approaches.
The first class of such approaches is ensembling that quantifies a model’s uncertainty in the hidden layers by marginalizing over an (implicit) probabilistic distribution of model parameters. The most notable examples of this class is Deep Ensemble that averages parallel-trained deep learning models that are initialized from distinct random seeds (Lakshminarayanan et al., 2017; Fort et al., 2019), and also MC Dropout (Gal and Ghahramani, 2016) that samples from a functional-space model distribution that is generated by perturbing the model’s dropout masks. As these approaches are uniquely capable of quantifying hidden-representation uncertainty, we hypothesize that they would synergize very well with SNGP, which is a single-model, last-layer-oriented approach.
Data augmentation (Thulasidasan et al., 2019; Hendrycks* et al., 2020; Cubuk et al., 2020; Chen et al., 2020a, b) represents a second class of approach that improves the representation learning of neural networks by explicitly injecting expert knowledge about the types of surface-form perturbations that a model should be sensitive to or invariance against, thereby encouraging the model to learn a meaningful distance in its representation space. Compared to the spectral normalization technique which provides a global guarantee in distance preservation that applies to the entire input space (Proposition 3), the data augmentation methods provides a local guarantee in the neighborhood of the training data. However, this guarantee is also expected to be stronger (i.e., a better correspondence between the hidden-space distance and a suitable metric for the input data manifold) since we have explicitly instructed the model representation to be sensitive to or invariant against the semantically meaningful or meaningless directions as specified by human experts, respectively (see Section 8.1 for further discussion). Consequently, in the case where the coverage of the training data is sufficiently large, and a well-designed augmentation library is available for the task, data augmentation is expected to complement the spectral-normalization technique to provide a even stronger guarantee in distance-awareness.
In this section, we investigate the efficacy of the SNGP model as a building block for other uncertainty estimation techniques, and show that it leads to improvements that are complementary to those achieved by the other methods. Specifically, we consider the following methods from literature:
MC Dropout (Gal and Ghahramani, 2016) uses dropout regularization at test-time as a way to obtain samples from a predictive posterior distribution. SNGP can be easily combined with MC Dropout by enabling dropout at the inference time.
Deep Ensemble (Lakshminarayanan et al., 2017) train multiple SNGP models with different random seed initializations and then average the predictions during test-time.
Data Augmentation We consider AugMix (Hendrycks* et al., 2020), a data augmentation algorithm that is known to improve the robustness and uncertainty estimation properties of a network by improving the learned representations of data through a combination of expert-designed data augmentations.It is easy to augment the SNGP training procedure by simply plugging in the Augmix augmentations into the data processing pipeline before feeding the data into the SNGP model For ease of composability, we did not add the further consistency regularization through the Jensen-Shannon divergence as done in the original AugMix paper, as we do not notice further significant improvements on adding the loss on top of the Augmix augmentations..
We find that the SNGP approach complements extremely well with the ensemble methods in practice. From Tables 7 and 8, we can see that ensembling SNGP members either through MC Dropout or through Deep Ensemble results in improved performance on all metrics when compared to a single model, and in particular improves both the in-distribution calibration and the OOD detection for both CIFAR-10 and CIFAR-100 datasets. Though there are some benefits to adding either spectral normalization or a last-layer GP separately to ensemble members, the combined effect of both parts of the SNGP algorithm consistently results in the best calibration and OOD performance, especially in CIFAR-100 (which is a harder task). This behavior suggests that ensembling models does not inherently address the caveats in the representations of the underlying members, and further improvement can be obtained by imposing spectral normalization and the GP layer to the base models to improve their distance awareness property. To put it another way, SNGP offers an orthogonal improvement to the model’s uncertainty quality that cannot be obtained from ensembling, and combining these two approaches can result in the best overall performance. Therefore, wherever resource permits, we suggest using an ensemble of SNGP base models to compound the benefits of distance awareness and representation diversity in improving the quality of uncertainty quantification (Section 4). To conclude, even though ensembling different neural network instantiations improves epistemic uncertainty quantification by marginalizing from different points in the posterior, improving the quality of each member of the ensemble is crucial to further improve the predictive uncertainty.
2 Data Augmentation for Improved Distance Awareness
As shown in the introduction (Figure 1), applying spectral normalization to residual neural networks encourages the model representation to be distance preserving and prevents it from feature collapse. Interestingly, this overlaps with the design goal of another well-known uncertainty technique: data augmentation (e.g., AugMix (Hendrycks* et al., 2020)). Specifically, data augmentation improves the distance preservation ability of the representation by forcing the model to be invariant against the semantic-preserving perturbations (e.g., rotating an image), and also being sensitive to semantic-modifying perturbations (e.g., adding a word “not” to a natural language sentence). To empirically investigate the composability of these two techniques, we perform ablation experiments of training SNGP models on the CIFAR-10 and CIFAR-100 datasets with AugMix, and show that the improvements in representations learnt with AugMix complement SNGP. Results are shown in Table 7-8. If data augmentation improves distance awareness, we would expect to see DNN-GP + AugMix significantly improve over DNN-GP; this is indeed the case, which confirms our hypothesis that increasing distance awareness (either through smoothness or data augmentation) improve DNN-GP. It is interesting to observe that in this task, while the SNGP + AugMix on average outperforms its ablated counterparts, the gap between DNN-GP + AugMix vs. SNGP + AugMix is lower than DNN-GP vs. SNGP, illustrating the effectiveness of AugMix in improving model uncertainty via enhancing the distance awareness. Therefore, for domains where a suite of well-designed augmentation is available for the task, and when the test data is in a neighborhood of the augmented training examples (which is the case for CIFAR-C), the data augmentation method can provide a strong guarantee in preserving a meaningful distance in the input space.
In summary, these experiments illustrate the importance of the representation’s distance-awareness quality in improving the model’s uncertainty quality. Consequently, for datasets and modalities where a well-designed augmentation library is available (e.g., AugMix for image modality), it is advantageous to complement SNGP with the data augmentation approaches to obtain a stronger guarantee in preserving a semantically meaningful distance in its representation space. However, when it is difficult to obtain high-quality augmentations (e.g. the genomics example), or it is computationally intractable to sufficiently augment the training data, the SNGP alone provides a cheap, modality-invariant approach to improve trained representations that provides a global guarantee for distance preservation (see Section 8.1 for further discussion).
Conclusions and Discussion
We propose SNGP, a simple approach to improve the predictive uncertainty estimation of a single deterministic DNN. It makes minimal changes to the architecture and training/prediction pipeline of a deterministic DNN, only adding spectral normalization to the hidden mapping, and replacing the dense output layer with a random feature layer that approximates a GP. We theoretically motivate distance awareness, the key design principle behind SNGP, via a decision-theoretic analysis of the uncertainty estimation problem. We also propose a closed-form approximation method to make the GP posterior end-to-end trainable in linear time with the rest of the neural network. On a suite of vision and language understanding tasks and on modern architectures (ResNet and BERT), SNGP is competitive in prediction, calibration and out-of-domain detection, outperforms other single-model approaches, and combines well with other state-of-the-art uncertainty techniques.
A central goal of this work is to provide theoretical and empirical evidences for the importance of incorporating distance awareness (i.e., the distance of a test example from the training data) into a model’s uncertainty estimate. This provides a complementary view to the classic approaches to deep learning uncertainty, where the model uncertainty is primarily quantified by a test example’s distance from the decision boundary (e.g., Figure 1). Indeed, both the distance to training data and the distance to decision boundary are reasonable quantifiers of model uncertainty, and should be treated as equally important components of a practitioner’s uncertainty quantification toolbox. In practice, the choice between the two distances should be made based on the nature of the data generating mechanism and the optimality criteria that the practitioner wish to pursue. For example, under the classic i.i.d. assumption where the test examples always stays in-domain (i.e., identically distributed as the training data) and one wish to use model uncertainty to detect ambiguous examples, then distance to the decision boundary is a suitable choice. On the other hand, in a safe-critical setting where it is important to guard against worst-case risk (i.e., Section 2.1) in the presence of likely distributional shift, distance to the training data would be a more appropriate choice. Alternatively, one may consider selecting a uncertainty metric that incorporates both types of distances. For example, as revealed by Equation 17, the posterior predictive mean of SNGP in fact incorporates both the magnitude of the predictive logit (i.e., distance to the decision boundary) and predictive variance (i.e., distance to the training data).
Furthermore, an important observation we made in this work is that learning smooth representations is important for good uncertainty quantification. In particular, we highlighted bi-Lipschitz (Equation (6)) as an important condition for the learned representation of a DNN to attain high-quality uncertainty performance. We proposed spectral normalization as a simple approach to ensure such property in practice, and illustrated its practical effectiveness across a wide range of data modalities. Perhaps surprisingly, the improvement is observed even in the high-dimensional regime (e.g., image and text), where the true “semantic” distance seemingly diverges from that based on the surface-form representation (Section 6), and the benefit of SNGP seems to be not fully overlapping with those provided by the other state-of-the-art uncertainty techniques (Section 7). To this end, we find it theoretically relevant to initiate a discussion about the role that the bi-Lipschitz condition may play in high-dimensional and overparameterized learning, and how it compares with those of the other state-of-the-art approaches. The following discussion is in no way comprehensive, and is intended to serve as potential starting points for future work.
At a first glance, the technique proposed in Section 3.2 seems to act on a naive, surface-level distance in the input space (e.g., distance in the pixel space), hence raising the question whether the “semantic” distance (i.e., a meaningful distance metric that is appropriate for the input data manifold) is being preserved as well. Indeed, in common machine learning applications, the data points reside on an underlying low-dimensional manifold behind their high-dimensional surface representation Appendix D provides an example mathematical formalization that further elaborates this idea.. Here, the correspondence between the true manifold distance (i.e., the “semantic distance”) and the surface-level distance (e.g., distance in the pixel space) is complex and dynamic, and can be described by the concept of metric distortion from metric embedding theory (Abraham et al., 2011; Matoušek, 2013; Chennuru Vankadara and von Luxburg, 2018). Formally, consider a mapping between the metric space equipped with surface-level distance to that equipped with the true distance . Then, can be understood as a distorted version of the surface-level distance as implemented by , with the type and degree of the distortion changes depending on the location in the product feature space . More specifically, defining the distance ratio that measures the degree of local distortion at , two types of distortions may occur:
(Distance Contraction) , such that the magnitude of is high.
(Distance Expansion) , such that the magnitude of is high.
The first scenario (distance contraction) is common in both image and text modalities, where the underlying content of an image / text is invariant to the noisy movements in the pixel / token space. The second scenario (distance expansion) occurs often in language understanding, where the addition of a single token (e.g., “not”) drastically changes the meaning of a sentence. Furthermore, the type and degree of distortion is location- and direction-dependent. For example, in the language domain, a sentence’s meaning is generally invariant to its surface-level syntactic forms (i.e., distance contraction), however some sentences can be drastically changed by the modification of a few key tokens (i.e., distance expansion).
Interestingly, despite this complicated and dynamic distortion of distance, it is possible to show that, as long as the “semantic” distance is well defined (i.e., ), the global bound on the surface-level distance still provides a basic guarantee in preserving the local true distances at every (albeit at a looser degree):
Assume that the mapping preserves the surface-form distance , i.e., there exist non-negative constants such that:
then, for a well-defined surface-form distance that is a distortion of , we have:
(Distance Contraction) For a local region where the distance contraction occurs with magnitude , there exists a such that and
(Distance Expansion) For a local region where the distance expansion occurs with magnitude , there exists a such that and
As a result, by bounding the neural network model from warping the representation space distance to an extreme degree, the Lipschitz bound (124) provides a basic guarantee in the preservation of the “semantic” distance that is appropriate for the data manifold. Consequently, the SNGP approach concretely improves the unregularized training of the overparameterized network by helping it to reach a better balance between dimension reduction and information preservation, which we discuss the next.
As the information of machine learning data tends to concentrate on a low-dimensional manifold, a high-dimensional learning model is often confronted with a tradeoff between the need of dimension reduction and the need of information preservation: dimension reduction is necessary for reducing the statistical complexity of the learning problem and for ensuring superior generalization under finite data. However, an excessive reduction of information leads the model to ignore semantically meaningful features that are less correlated with the training label, creating a challenge for uncertainty quantification.
In the context of modern overparameterized networks, naive training without regularization is often observed to lead to one extreme end of this tradeoff (i.e., excessive dimension reduction). Specifically, by minimizing the training cross entropy toward zero, the model pushes its logits to a large magnitude at the location of training labels , leading to a high alignment between training logits and the corresponding labels. Viewing from the representation space , this manifests into the training data representations being pushed far away from the decision boundary , which is often achieved by stretching the hidden-space distance along the task-relevant directions (i.e., those orthogonal to the decision boundaries) to an extreme degree. This essentially leads to a warped hidden-space geometry with extremely high weighting on a handful of task-relevant principal directions (see, e.g., Figure 1 top), implying low effective degrees of freedom (Papyan, 2020; Hauser and Ray, 2017a). Although not detrimental to the in-domain generalization, this warped geometry leads the model to be overly sensitive to the hidden-space movements along the label-relevant direction, while being insensitive to semantic change in the directions that are less correlated with the training labels, thereby creating challenges for uncertainty tasks such as out-of-domain detection (Hein et al., 2019a).
To this end, the SNGP approach helps the model to strike a better balance between dimension reduction v.s. information preservation by modulating the degree of geometry distortion via the Lipschitz bound. As a result, the model is still able to disentangle input features through a cascade of layer-wise transformations, but is protected from extreme distortion in to a degree that causes feature collapse (see, e.g., Figure 4 bottom). As a result, the model is still capable of learning abstract and task-relevant features by exploiting the expressive power of overparameterization, while not blind to semantically meaningful movements that are less correlated with training labels, thereby leading to a better balance between generalization and uncertainty performance.
It is worth noting that there exist other representation learning techniques, e.g., data augmentation, constrastive learning or unsupervised pretraining, that are known to also improve a network’s uncertainty performance (Hendrycks et al., 2019; Hendrycks* et al., 2020). From the perspective of distance awareness, these methods help the model representation to preserve a semantically meaningful distance by injecting external knowledge into the learning pipeline. For example, data augmentation and constrastive learning instructs the model to learn distance contraction or expansion by creating semantically similar or dissimilar pairs using expert knowledge, so that the DNN representation is invariant against the semantic-preserving perturbations (e.g., rotating an image), and is sensitive to semantic-modifying perturbations (e.g., adding a word “not” to a natural language sentence). On the other hand, pre-training injects the model with a prior about the similarity or dissimilarity between examples that was learned (e.g., via contrastive masked language modeling) from a large external corpus. Compared to these data-intensive approaches, spectral normalization provides a global guarantee in distance preservation and does not require external resources such as pretraining datasets or an understanding of which augmentations are semantically meaningful and useful. However, the guarantee is also looser since it only places a uniform upper and lower bound on along all directions of the perturbation, rather than explicitly training to be contractive or expansive with respect to the direction of semantic perturbation. Interestingly, as shown in Sections 6 - 7, SNGP combines well with these additional techniques, indicating that the benefits they provide to the model do not conflict with each other.
Finally, we note that the Gaussian process algorithm presented this work is optimized for practicality (i.e., scalability to extremely large datasets, and seamless integration into minibatch SGD-based neural training pipeline). Such design goal invariably necessitates several approximations for the exact Gaussian process posterior: (1) random-feature expansion for the kernel function, (2) Laplace approximation for the posterior variance, and (3) mean-field approximation for the softmax posterior mean. In our preliminary experiments, we found that a higher-quality approximation to the kernel function (e.g., using orthogonal random feature (Yu et al., 2016) rather than random Fourier feature) and an near-exact Monte Carlo approximation to the softmax posterior mean did not lead to a meaningful improvement to model performance. However, this does not preclude the theoretical possibility that an exact Gaussian process posterior can further improve uncertainty quality when compared to the current algorithm choice. Therefore, identifying GP algorithms that can better implement the distance-awareness principle without sacrificing practicality is an important direction for future work.
Acknowledgments and Disclosure of Funding
We would like to thank Rodolphe Jenatton, D. Sculley, Kevin Murphy, Deepak Ramachandran for the insightful comments and fruitful discussions.
References
Appendix A Method Summary
Adding spectral normalization to the hidden weights , and
Replacing the dense output layer with a GP layer. Under the RFF approximation, the GP layer is simply a one-layer network with hidden units . Here are frozen weights that are initialized from a Gaussian and a uniform distribution, respectively (as described in Equation (10)).
Algorithm 1 summarizes the training step. As shown, for every minibatch step, the model first updates the hidden-layer weights and the trainable output weights via SGD, then performs spectral normalization using power iteration method ((Gouk et al., 2021) which has time complexity ), and finally performs precision matrix update (Equation (11), time complexity ). Since are fixed for a given architecture and usually , the computation scales linearly with respect to the sample size. We use in the experiments.
Algorithm 2 summarizes the prediction step. The model first performs the conventional forward pass to compute the final hidden feature , and then compute the posterior mean (time complexity ) and the predictive variance (time complexity ).
To estimate the predictive distribution where m_{k}\sim N\big{(}\hat{m}_{k}(\bx),\hat{\sigma}_{k}^{2}(\bx)\big{)}, we calculate its posterior mean using either Monte Carlo averaging or mean-field approaximation. Notice that this Monte Carlo averaging is computationally cheap since it only involves sampling from a closed-form distribution whose parameters are already computed by the single feed-forward pass (i.e., a single call to tf.random.normal). This is different from the full Monte Carlo sampling used by MC Dropout or deep ensembles which require multiple forward passes and are computationally expensive.
However, in applications where the inference latency is of high priority (e.g., in the on-device settings), we can reduce the computational overhead further by replacing the Monte Carlo averaging with the mean-field approximation as described in the main text (Equation 17) (Daunizeau, 2017).
A.1 Extension to Regression and Multi-class Classification
As introduced in the main text (LABEL:eq:laplace), for an arbitrary model posterior , its corresponding Laplace posterior is:
Consequently, to compute the Laplace posterior for any data likelihood , it is sufficient to first compute the MAP estimate as done in (12), and then derive the model Hessian .
In the main text, we introduced the expression for for a sigmoid cross entropy likelihood \log p(\beta|\Dsc)=-\sum_{i=1}^{n}\big{[}I(y_{i}=1)\log p_{i}+I(y_{i}=0)\log(1-p_{i})\big{]}+\frac{1}{2}||\bbeta||^{2}_{2} based on the classic result from Rasmussen and Williams (2006). In this section, we show how to apply (LABEL:eq:laplace) to other common likelihood functions such as regression (using a squared loss) and -class classification (using a softmax cross entropy loss) by deriving their Hessian.
Under squared loss, The model posterior likelihood is:
where . Then the model Hessian is:
Under a -class softmax cross entropy loss, the model posterior likelihood is:
Then, computing the Hessian with respect to yields (Kunstner et al., 2019):
This leads to a matrix with diagonal and off-diagonal blocks as:
Consequently, the Laplace posterior for each class is:
In the case where computing the class-specific covariance matrix is either infeasible (e.g., the number of output classes is simply too large) or not of interest (e.g., we are just interested in computing an scalar uncertainty statistic for an input example), one can consider computing an upper-bound of the covariance matrices in (125) as:
where is the maximum class probability. From an information theoretic perspective, quantifying model uncertainty using (126) can be understood adopting the maximum entropy distribution among the family of all class-specific distributions. To ensure computational feasibility for tasks with high-dimensional output (e.g., CIFAR100, CLINC, and ImageNet which correspond to 100, 150 and 1000 classes), we use (126) in our experiments.
A.2 Hyperparameter Configuration
SNGP is composed of two components: Spectral Normalization (SN) and Gaussian Process (GP) layer, both are available at the open-source edward2 probabilistic programming library https://github.com/google/edward2.
Spectral normalization contains two hyperparameters: the number of power iterations and the upper bound for spectral norm (i.e., in Equation (13)). In our experiments, we find it is sufficient to fix power iteration to 1. The value for the spectral norm bound controls the trade-off between the expressiveness and the distance awareness of the residual block, where a small value of may shrink the residual block back to identity mapping hence harming the expressiveness, while a large value of may lead to the loss of bi-Lipschitz property (Proposition 3). Furthermore, the proper range of depends on the layer type: for dense layers (e.g., the intermediate and the output dense layers of a Transformer), it is sufficient to set to a value between . For the convolutional layers, the norm bound needs to be set to a larger value to not impact the model’s predictive performance. This is likely caused by the fact that the current spectral normalization technique does not have a precise control of the true spectral norm of the convolutional kernel, in conjuction with the fact that the other regularization mechanisms (e.g., BatchNorm and Dropout) may rescale a layer’s spectral norm in unexpected ways (Gouk et al., 2021; Miyato et al., 2018). In general, we recommend performing a grid search for to identify the smallest possible values of that still retains the predictive performance of the original model. In the experiments, we set the norm bound to for a WideResNet model.
The Gaussian process layer (Equation 10) contains 3 hyperparameters, which are (1) the hidden dimension (, i.e., the number of random features), (2) the length-scale parameter for the RBF kernel, and (3) the kernel amplitude . In the experiments, we find the model’s performance to be not very sensitive to the hidden dimension or the length-scale parameter. Setting in a range between $l=2.0\sigma$, since it functions in a manner that is similar to the temperature parameter in temperature scaling (Guo et al., 2017).
The posterior covariance (Equation 11) does not contain hyperparameter. However, if one wishes to use a moving average estimator for the covariance matrix, i.e.,
Then there will be two additional hyperparameters: the ridge factor and the discount factor . The ridge factor serves to control the stability of matrix inverse (if the number of sample size is small), and controls how fast the moving average update converges to the population value . Similar to other moving-average update method, these two parameters can impact the quality of learned covariance matrix in non-trivial ways. In general, we recommend conducting some small scale experiments on the data to validate the learning quality of the moving average update in approximating the population covariance. In the experiments, we set and , which is sufficient for our setting where the number of minibatch steps per epoch is large. We use the exact update formula (11) in the CIFAR and Genomics experiments, and use the moving average update as described above for larger tasks such as ImageNet.
In summary, when applying SNGP method to a new dataset it is sufficient to sweep the spectral norm bound and kernel amplitude parameter on a holdout validation dataset, and fix all other parameters to their default values (Table 9). We report the values for each experiment in Section C.1.
As an aside, we also implemented two additional functionalities for GP layers: input dimension projection and input layer normalization. The input dimension project serves to project the hidden dimension of the penultimate layer to a lower value (using a random Gaussian matrix ), it can be projected down to a smaller dimension. Input layer normalization applies Layer Normalization to the input hidden features, which is akin to performing automatic relevance determination (ARD)-style variable selection to the input features. Ablation studies revealed that the model performance is not sensitive to these changes.
Appendix B Formal Statements
The expression in (3) seeks to answer the following question: assuming we know the true domain probability , and given a model that we have already learned from data, what is the best solution we can construct to minimize the minimax objective (2)? The interest of this conclusion is not to construct a practical algorithm, but to highlight the theoretical necessity of taking into account the domain probability in constructing a good solution for uncertainty quantification. If the domain probability is not necessary, then the expression of the unique and optimal solution to the minimax probability should not contain even if it is available. However the expression of (3) shows this is not the case.
To make the presentation clear, we formalize the statement about (3) into the below proposition:
the model’s predictive distribution learned from data
the true domain probability,
then there exists an unique optimal solution to the minimax problem (2), and it can be constructed using (a) and (b) as:
where is a discrete uniform distribution for classes.
As discussed in Section 2.1, the solution (131) is not only optimal for the minimax Brier risk, but is in fact optimal for a wide family of strictly proper scoring rules known as the (separable) Bregman score (Parry et al., 2012):
where is a strictly concave and differentiable function. Bregman score reduces to the log score when , and reduces to the Brier score when .
Therefore we will show (131) for the Bregman score. The proof relies on the following key lemma:
Consider the Bregman score in (132). At a location where the model has no information about other than , the solution to the minimax problem
is the discrete uniform distribution, i.e., . The proof for Lemma B is in Section E.2. It is worth noting that Lemma B only holds for a strictly proper scoring rule (Gneiting et al., 2007). For a non-strict proper scoring rule (e.g., the ECE), there can exist infinitely many optimal solutions, making the minimax problem ill-posed.
We are now ready to prove Proposition B: Proof. Denote . Decompose the overall Bregman risk by domain:
where we have denoted S_{\texttt{IND}}(p,p^{*})=E_{\bx\in\Xsc_{\texttt{IND}}}\big{(}s(p,p^{*}|\bx)\big{)} and S_{\texttt{OOD}}(p,p^{*})=E_{\bx\in\Xsc_{\texttt{OOD}}}\big{(}s(p,p^{*}|\bx)\big{)}.
Now consider decomposing the sup risk for a given . Notice that sup risk is separable by domain for any . This is because and has disjoint support, and we do not impose assumption on :
We are now ready to decompose the minimax risk . Notice that the minimax risk is also separable by domain due to the disjoint in support:
also notice that the in-domain minimax risk \inf_{p}\sup_{p^{*}}\big{[}S_{\texttt{IND}}(p,p^{*})\big{]} is fixed due to condition .
Therefore, to show that (131) is the optimal and unique solution to (133), we only need to show is the optimal and unique solution to \inf_{p}\sup_{p^{*}}\big{[}S_{\texttt{OOD}}(p,p^{*})\big{]}. To this end, notice that for a given :
due to the fact that we don’t impose assumption on (therefore is free to attain the global supreme by maximizing at every single location ). Furthermore, there exists that minimize at every location of , then it minimizes the integral (Berger, 1985). By Lemma B, such exists and is unique, i.e.:
In conclusion, we have shown that is the unique solution to . Combining with condition (a)-(b), we have shown that the unique solution to (133) is (131).
Appendix C Experiment Details
For CIFAR-10 and CIFAR-100, we followed the original Wide ResNet work to apply the standard data augmentation (horizontal flips and random crop-ping with 4x4 padding) and used the same hyperparameter and training setup (Zagoruyko and Komodakis, 2017). The only exception is the learning rate and training epochs, where we find a smaller learning rate ( for CIFAR-10 and for CIFAR100, v.s. for the original WRN model) and longer epochs ( for SNGP v.s. for the original WRN model) leads to better performance.
For CLINC OOS intent understanding data, we pre-tokenized the sentences using the standard BERT tokenizerhttps://github.com/google-research/bert with maximum sequence length 32, and created standard binary input mask for the BERT model that returns 1 for valid tokens and 0 otherwise. Following the original BERT work, we used the Adam optimizer with weight decay rate and warmup proportion . We initialize the model from the official BERT checkpointhttps://storage.googleapis.com/bert_models/2020_02_20/uncased_L-12_H-768_A-12.zip. For this fine-tuning task, we using a smaller step size ( for SNGP .v.s. for the original BERT model) but shorter epochs ( for SNGP v.s. for the original BERT model) leads to better performance. When using spectral normalization, we set the hyperparameter and apply it to the pooler dense layer of the classification token. We do not apply spectral normalization to the hidden transformer layers, as we find the pre-trained BERT representation is already competent in preserving input distance due to the masked language modeling training, and further regularization may in fact harm its predictive and calibration performance.
For Genomics sequence data, we consider a 1D CNN model following the prior work (Ren et al., 2019). Specifically, the model is composed by one convolutional layer of 2,000 filters of length 20, one max-pooling layer, one dense layer of 2,000 units, and a final dense layer with softmax activation for predicting class probabilities. To build the SNGP model, we add spectral normalization to both the convolutional layer and the dense layer, and we replace the last layer with Gaussian process layer. For MC Dropout, we add filter-wise dropout before the convolutional layer with dropout rate 0.1. For Deep Ensemble, we ensemble 5 models trained based on random initialization of network parameters and random shuffling of training inputs. The model is trained using the batch size 128, the learning rate 1e-4, and Adam optimizer. The training step is set for 1 million, but we choose the best step when validation loss is at the lowest value. Due to the large size of test OOD dataset, we randomly select 100,000 OOD samples to pair with the same number of in-distribution samples.
For each of the experiment, we fix the SNGP hyperparameters to their recommended values as in Table 9, and sweep the Spectral Norm Bound and the kernel amplitude with respect to the negative log likelihood on the validation data. The final values found for each experiment is summarized at Table 10. As shown, the optimal value for spectral norm bound usually depends on the layer type (1D Convolution v.s. 2D Convolution v.s. Dense), while the optimal value of kernel amplitude seems to be sensitive to the data modality, the model type, and also the type of covariance estimator (i.e., exact estimator v.s. moving average estimator) being used (also see earlier discussion in Section A.2).
All models are implemented in TensorFlow and are trained on 8-core Cloud TPU v2 with 8 GiB of high-bandwidth memory (HBM) for each TPU core. We use batch size 32 per core.
C.2 Evaluation
For CIFAR-10 and CIFAR-100, we evaluate the model’s predictive accuracy and calibration error under both clean and corrupted versions of the CIFAR testing data. The corrupted data, termed CIFAR10-C, includes 15 types of corruptions, e.g., noise, blurring, pixelation, etc, over 5 levels of corruption intensity (Hendrycks and Dietterich, 2018).
We assess the model’s calibration performance using the empirical estimate of ECE: which estimates the difference in model’s accuracy and confidence by partitioning model prediction into bins (Guo et al., 2017). In this work, we choose .
We also evaluate the model performance in OOD detection by using the CIFAR-10/CIFAR-100 model’s uncertainty estimate as a predictive score for OOD classification, where we consider a standard OOD task by testing CIFAR-10/CIFAR-100 model’s ability in detecting samples from the Street View House Numbers (SVHN) dataset (Netzer et al., 2011), and a more difficult OOD task by testing CIFAR-10’s ability in detecting samples from the CIFAR-100 dataset, and vice versa. For CIFAR dataset, since the input data is normalized by the sample mean and variance, we perform the same normalization for the test OOD data. We evaluate the OOD uncertainty scores for the inputs in in-distribution test set, and the inputs in OOD test set, and we use AUROC to measure the separation between the two sets based on the OOD uncertainty score.
The Dempster-Shafer metric computes its uncertainty for a test example as:
Table 11 shows the OOD performance for CIFAR dataset based on the above four different confidence scores. First, it shows that Dempster Shafer score in general has better performance than MSP. Second, we noticed that Mahalanobis distance has much worse performance on GP based models, while Relative Mahalanobis distance corrects for the degradation. For example, for CIFAR-10 vs. CIFAR-100 task, Mahalanobis distance based on SNGP model’s embeddings has only 0.742 AUROC, while other OOD methods achieve around 0.90 AUROC based on SNGP model, including Relative Mahalanobis distance. This suggests that the GP based models probably preserved background features which confounds the raw Mahalanobis distance score, and the relative Mahalanobis distance corrects for the background features and fix the performance. It is interesting to see that the Relative Mahalanobis distance has the best performance for the most challenging near-OOD task CIFAR-100 vs. CIFAR-10 for all models.
We also additionally evaluate another two simple OOD datasets, random Gaussian noise and texture dataset (i.e. DTD), used in the prior work (Hendrycks et al., 2018; Liang et al., 2018; Liu et al., 2020b). The results are included in Table 11. As shown, SNGP provides the best performance when using the default MSP method for OOD detection: 0.999 AUROC for detecting random Gaussian noise, and 0.959 AUROC for detecting DTD. When combined with more advanced OOD detection signals, such as Dempster Shafer, the performance can be further improved to 1.000 for random Gaussian and 0.988 for DTD. Furthermore, for the uncertainty metrics whose performance hinders on the qualities of both the hidden representation and the last-layer (i.e., MSP and Dempster-Shafer), the SNGP model attains the strongest performance when compared to its ablated counterparts. Finally, we observed some vanilla methods (e.g., Mahalanobis distance based on vanilla DNN) also achieves strong performance for these simple datasets. However, this advantage starts to break down on the more difficult datasets (e.g., CIFAR100 v.s. SVHN).
In conclusion, when comparing to the baseline approaches, SNGP provides the best out-of-box performance when using the default MSP method for OOD detection (Section 6.2.1), and it can be combined with more advanced OOD detection signals to further improve performance.
C.3 Theoretical Convergence to Optimal Behaviour
In this section, we discuss the asymptotic behaviour of the SNGP algorithm on OOD datapoints far from the training distribution, and show how the predictive distribution converges to the optimal distribution suggested by Equation (3). We formalize this in the following proposition:
A distance-preserving hidden mapping , so that
for two positive constants.
A dual form formulation of the SNGP model assuming a Laplace approximation over the posterior, with posterior mean and variance given by
where are kernel matrices approximating those under the RBF kernel .
The predictive distribution of the SNGP model for classification given by
then as test points tend away from the training manifold (i.e ), the limit of the predictive distribution is either exact equal to the optimal distribution as given by (3) (under the mean-field approximation), or can be closely approximated by (under the Monte Carlo approximation).
The proof for Proposition 8 follows: Proof. Given condition (a), is a bi-Lipschitz, distance-preserving function, and therefore we can write down the asymptotic convergence of the RBF kernel as follows
where and are constants, and the convergence is guaranteed by the sandwich theorem, since by the distance preservation property, .
Leveraging the above result, we can now reason about the asymptotic behaviour of the posterior moments :
As a result, we can now consider the asymptotic behaviour of the predictive distribution by leveraging the asymptotic behaviour of the moments.
First notice that by mean-field theorem, the sigmoid and softmax transformation of Gaussian is approximated as below (Bishop, 2011; Lu et al., 2020):
Then, since and as , the predictive distribution converges to the uniform distribution:
where the second equality follows from the continuous mapping theorem.
Notice that since the Monte Carlo approximation of is an unbiased estimator of the true integral, it is sufficient to investigate the asymptotic behavior of the integral:
where denotes the probability density function of a Gaussian-distributed random variable with mean and standard deviation . In the above, the second equality (i.e., switching of limit and integral) follows by the bounded convergence theorem since the integrand is a bounded function. Intuitively, Equation 170 is again a Gaussian-softmax integral whose expectation can be approximated closely via the mean-field approximation (Bishop, 2011; Lu et al., 2020). That is, by following the same line of argument as above, we now that the limiting distribution of the Monte Carlo approximation can be closely approximated by a uniform distribution, i.e.,
Alternatively, we can argue about convergence of Equation 170 by appealing to the multivariate continuous mapping theorem. Specifically, since , we have and , where denotes the convergence in distribution. Consequently, . Now, for a -dimensional , consider the coordinate of :
As shown, since , we expect to be equal . Consequently, denoting and notice sum to 1, we have:
which leads to the unique solution of , i.e., , which again shows:
Appendix D An Example Formalization of “Semantic Distance”.
In this section, we develop an example formalization of the intuition notion of semantic distance using languages from the metric embedding theory (Abraham et al., 2011; Matoušek, 2013; Chennuru Vankadara and von Luxburg, 2018). Here, our goal is to supply an example formalization of this often intuitive notion, with then goal of facilitating a rigorous understanding of Section 8.1 (Abraham et al., 2011; Matoušek, 2013; Chennuru Vankadara and von Luxburg, 2018). Indeed, the term “semantic distance” has a long history in the literature, and it is out of the scope of this work to provide an authoritative, all-encompassing mathematical construction that unifies its diverse usages across many fields such as manifold learning, representation learning, natural language processing, and cognitivle psychology (Tenenbaum et al., 2000; Mohammad and Hirst, 2006; Deselaers and Ferrari, 2011; Hashimoto et al., 2016; Higgins et al., 2018; Khemakhem et al., 2020; Chandrasekaran and Mago, 2021; Reilly et al., 2022).
We consider the setting where all examples within a problem domain can be sufficiently described by a large collection of attributes, with each attribute being either discrete (e.g., types of entities that appeared in a image and their relationships) or continuous (e.g., color intensity or camera angle). The number of attributes is finite but allowed to be very large. Furthermore, we assume different attribute impacts the semantic similarity between examples to a different degree, so that the variations along only a subset of attributes constitute a meaningful difference between the examples.
We can formalize the above intuition terms of metric space (Rudin et al., 1976). Specifically, we can assume the examples has a semantic representation that reside in a semantic space , which is a -dimensional metric space with its dimensions correspond to the discrete and continuous attributes. Specifically, can be expressed as a product of attribute subspaces :
where each attribute subspace is a metric space that corresponds to a continuous or discrete attribute, and is equipped with a well-defined metric . For example, can represent a continuous attribute such as the color, which implies and is the standard Euclidean metric for the RGB space. can also represent a discrete attribute such as entity types, which implies is a discrete space that is supported on a large amount of candidate entities (e.g., from a knowledge graph), and the metric between entities can be defined by a certain graph metric with respect to a pre-established concept hierarchy (e.g., WordNet) (Chandrasekaran and Mago, 2021). As a result, we can write the semantic representation of an example as where , and the differences between two examples in the attribute can be described as .
Consequently, a “semantic distance” can be defined in terms of a metric function for the product space , so that is a valid metric space. To this end, a proper choice of should satisfy the metric axioms (positivity, symmetry, and triangle inequality) while aligns well with the intuitive notion of “semantic similarity” between examples. For example, consider the below definition of {\color[rgb]{0,0,0}d_{X}}:
i.e., {\color[rgb]{0,0,0}d_{X}} is a weighted sum of attribute-specific metrics and their pairwise productsIt is also possible to define {\color[rgb]{0,0,0}d_{X}} with even higher-order products, e.g., , which we don’t explore here for the simplicity of exposition.. Here is the set of positive weights that sum to 1. We see the definition of semantic distance {\color[rgb]{0,0,0}d_{X}} in Equation 172 is flexible, as it not only allows the domain expert to define what constitutes a “semantically-meaningful difference” by assigning different weights among attributes, but also allows the attributes to interact to define the overall metric. We also see that {\color[rgb]{0,0,0}d_{X}} is a valid metric, as the positive guarantees positivity, and triangle inequality is closed under the summation and multiplication of positive termsThat is, if the attribute-specific metrics satisfy positivity and triangle inequality, then their sum and product also satisfy triangle inequality. For example, and .
So far, we have described a formal definition of semantic space and the associated semantic distance {\color[rgb]{0,0,0}d_{X}} in terms of a metric space (\Xsc,{\color[rgb]{0,0,0}d_{X}}). Further, (\Xsc,{\color[rgb]{0,0,0}d_{X}}) is constructed as the product of a collection of attribute metric spaces , so that the coordinates of adopt meaningful interpretations in terms of well-defined attributes. In practice, we often do not have direct access to (\Xsc,{\color[rgb]{0,0,0}d_{X}}), and can only work with its surface-form data representation , where the coordinates of and the surface-form distance is less meaningful. However, the elements in (\Xsc,{\color[rgb]{0,0,0}d_{X}}) still has a unique (i.e., one-to-one) correspondence with respect to ’s in the semantic space. For example, an image is often represented as a tensor of dimension (i.e., ). However, by visually inspecting this tensor, a human can still discern the various attributes underlying the image, implying that an inverse mapping exists there exists from the surface-form data space to the semantic space .
Formally, this means we can define an embedding function , which is a mapping from the semantic space (\Xsc,{\color[rgb]{0,0,0}d_{X}}) to the data space . As a result, every example adopts a surface-form representation , which means when measuring based on the data space, the distance between a pair of examples becomes:
Consequently, the distortion (introduced in the Section 8.1 in the discussion section) between the semantic distance {\color[rgb]{0,0,0}d_{X}} and the surface-form distance can be expressed as:
As shown, the distortion is induced by both the embedding function and the difference in the metrics (i.e., {\color[rgb]{0,0,0}d_{X}}\,v.s.\,d_{S}). From this perspective, when learning a neural network model based on the observed data , a semantic distance preserving representation should ideally satisfying the bi-Lipschitz condition (Equation 6) with respect to {\color[rgb]{0,0,0}d_{X}}:
so that has a reasonable correspondence with the semantic distance.
Appendix E Proof
The proof for Proposition 3 is an adaptation of the classic result of (Bartlett et al., 2018) to our current context:
Proof. First establish some notations. We denote the identity function such that for , we can write . For , denote ||h||=\sup\Big{\{}\frac{||f(\bx)||_{H}}{||\bx||_{X}}\mbox{ for }\bx\in\Xsc,||\bx||>0\Big{\}}. Also denote the Lipschitz seminorm for a function as:
It is worth noting that by the above definitions, for two functions and who shares the same input space, the Lipschitz inequality can be expressed using the norm, i.e., ||h(\bx)-h(\bx^{\prime})||_{H}\leq\alpha{\color[rgb]{0,0,0}d_{X}}(\bx,\bx^{\prime}) implies , and vice versa.
Now assume , . We will show Proposition 3 by first showing:
which is the bi-Lipschitz condition for a single residual block.
where the last line follows by the assumption . Rearranging, we get:
Combining (175)-(E.1), we have shown (174), which also implies:
Now show the bi-Lipschitz condition for a -layer residual network . It is easy to see that by induction:
Denoting and , we have arrived at expression in Proposition 3.
E.2 Proof of Lemma B
Proof Sketch. This proof is an application of the generalized maximum entropy theorem to the case of Bregman score. We shall first state the generalized maximum entropy theorem to make sure the proof is self-contained. Briefly, the generalized maximum entropy theorem verifies that for a general scoring function with entropy function , the maximum-entropy distribution attains the minimax optimality :
Let be a convex, weakly closed and tight set of distributions. Consider a general score function with an associated entropy function defined as . Assume below conditions on hold:
(Well-defined) For any , exists and is finite.
(Lower-semicontinous) For a weakly converging sequence where is bounded below, we have for all .
Then there exists an maximum-entropy distribution such that
Above theorem states that the maximum-entropy distribution attains the minimax optimality for a scoring function , assuming its entropy function satisfying certain regularity conditions. Authors of (Grünwald and Dawid, 2004) showed that the entropy function of a Bregman score satisfies conditions in Theorem 1. Consequently, to show that the discrete uniform distribution is minimax optimal for Bregman score at , we only need to show discrete uniform distribution is the maximum-entropy distribution.
Recall the definition of the strictly proper Bregman score (Parry et al., 2012):
where is differentiable and strictly concave. Moreover, its entropy function is:
Our interest is to show that for , the maximum-entropy distribution for the Bregman score is the discrete uniform distribution . To this end, we notice that in the absence of any information, the only constraint on the predictive distribution is that . Therefore, denoting , we can set up the optimization problem with respect to Bregman entropy (179) using the Langrangian form below:
Taking derivative with respect to and :
Notice that since is strictly concave, the function is monotonically decreasing and therefore invertible. As a result, to solve the maximum entropy problem, we can solve the above systems of equation by finding a inverse function , which lead to the simplification:
Above expression essentially states that all ’s should be equal and sum to 1. The only distribution satisfying the above is the discrete uniform distribution, i.e., .