Simple and Principled Uncertainty Estimation with Deterministic Deep Learning via Distance Awareness

In plain words

Neural networks are usually confident even on inputs they have never seen before. A common fix is to train many models and let them vote, which works but multiplies the cost of every prediction. Another is a Gaussian process, which admits when an input is unlike its training data, but it does not scale well to the high-dimensional inputs a deep network handles.

SNGP keeps a single network and changes two things. First, it constrains the network's inner layers so that inputs which are far apart stay far apart after they pass through the network; the network can no longer squeeze a strange input onto something familiar. Second, it replaces the last layer with a Gaussian-process-style layer, approximated so that it can be trained like any other layer. That layer reports how far an input is from the data it has seen, so the model can say "I do not know" for unfamiliar inputs. The result is a single model that gives uncertainty estimates close to those of an ensemble, at about the cost of one ordinary prediction.

TL;DR

• SNGP = spectral normalization of hidden weights + a random-feature Gaussian process output layer with a closed-form posterior covariance §3.1 §3.2.
• The design goal is "input distance awareness": the output layer must know how far an input is from the training data, and the hidden layers must not collapse distinct inputs §2 §3.
• Author-reported: competitive with a 10-model deep ensemble on accuracy, calibration and OOD detection on CIFAR-10/100 and CLINC, with latency close to a deterministic network §5 tab. 2 tab. 3 tab. 4.
• Main risks for an implementer: the paper and its released code differ in the predictive approximation, the Laplace likelihood and several hyperparameters (see Paper vs code). The paper's own equations also disagree with each other on the precision update.

Problem

Deep classifiers can be arbitrarily wrong on inputs far from the training set, so they need a notion of distance to the training data to abstain or flag out-of-domain inputs §1. Gaussian processes with suitable kernels have this property, but applying them to high-dimensional data needs a learned feature extractor, and that extractor can destroy the distance property §1 fig. 1. Ensembles and Bayesian networks give good uncertainty but cost several forward passes or several trainings; the authors name their memory and inference cost as a limit for real-time, industrial use §1. The paper's minimax analysis treats uncertainty estimation as a minimax problem under the Brier risk and identifies distance awareness as a necessary condition for minimax-optimal estimates §2.

Method

SNGP changes the architecture of a residual network logit(x) = g(h(x)), where h is the stack of hidden layers and g the output layer §3. It makes two changes, each motivated by the distance analysis in §2:

1. Distance-aware output layer §3.1. The dense output layer is replaced by a Gaussian process with an RBF kernel. Its prior over the training logits is

   \mathbf{g} \sim \mathrm{MVN}(\mathbf{0}, \mathbf{K}), \quad \mathbf{K}_{ij} = \exp(-\lVert h_i - h_j \rVert_2^2 / 2) §3.1

   The exact GP is too costly for large classification problems, so the kernel is approximated with random Fourier features. A frozen random projection W_L (entries i.i.d. N(0, 1)) and bias b_L (entries i.i.d. U(0, 2π)) produce D_L features, and only the output weights β are learned §3.1:

   \Phi_i = \sqrt{2/D_L}\,\cos(-W_L h_i + b_L) §3.1

   Training is then an ordinary Bayesian linear model in Φ. Its Laplace posterior has per-class precision

   \hat\Sigma_k^{-1} = I + \sum_i \hat p_{ik}(1 - \hat p_{ik})\,\Phi_i\Phi_i^\top §3.1

   where p̂_ik is the softmax prediction under the MAP estimate. The precision is accumulated over minibatches during one final pass over the training data, not during the normal SGD steps §3.1 alg. 1.

2. Distance-preserving hidden layers §3.2. For residual blocks h_l(x) = x + g_l(x), if every residual branch g_l is α-Lipschitz with α ≤ 1, the hidden map is bi-Lipschitz (its distances stay within constant factors of the input distances) §3.2. The paper enforces a spectral-norm bound on each hidden weight matrix W_l after every SGD step with one step of power iteration, rescaling only when the estimate exceeds the bound c §3.2 alg. 1:

   W_l \leftarrow c\,W_l / \hat\lambda \quad \text{if } c < \hat\lambda §3.2

Training alg. 1: for every minibatch, take an SGD step on the hidden weights, biases and output weights β; apply spectral normalization; in the final epoch, update the precision matrices; after the last step, invert each precision to get the posterior covariance. The extra cost over a deterministic network is one power iteration per weight matrix per step and the precision pass, which the paper says scales linearly in sample size app. A.

Prediction alg. 2: one forward pass gives the hidden features and Φ(x). The posterior mean of each logit is Φᵀβ_k, and the posterior variance is Φᵀ Σ_k Φ. The predictive distribution is the softmax averaged over Gaussian logit samples; the paper uses 10 Monte Carlo samples alg. 2 §5 app. A.

# SNGP training step, one minibatch  [alg. 1]
# Inputs: features (B, ...), labels (B,); c = spectral bound [app. C]
def train_step(features, labels, final_epoch):
    # 1. Random features of the hidden state  [§3.1]
    phi_batch = random_features(hidden(features))    # (B, D_L)
    logits = phi_batch @ betas                       # (B, K)
    # 2. SGD on hidden weights, biases and betas  [alg. 1]
    loss = cross_entropy(logits, labels)
    gradient_step(hidden_weights, biases, betas, loss)
    # 3. Spectral normalization, one power iteration  [§3.2]
    for W in hidden_weights:
        norm = estimate_spectral_norm(W)
        if norm > c:
            W = c * W / norm
    # 4. In the final epoch, add this minibatch to the precision  [§3.1]
    if final_epoch:
        batch_sum = zeros(K, D_L, D_L)
        probs = softmax(logits)                      # (B, K)
        for each example i in the minibatch:
            for each class k:
                # Laplace weight p (1 - p)  [§3.1]
                weight = probs[i, k] * (1 - probs[i, k])
                batch_sum[k] += weight * outer(phi_batch[i], phi_batch[i])
        # Discounted update, form of app. C; see Pitfalls  [app. C]
        precision = m * precision + (1 - m) * batch_sum
    # 5. After the last step, invert the precision  [alg. 1]
    if final_epoch:
        covariance = [inverse(precision[k]) for k in range(K)]

# Random Fourier features of the hidden state  [§3.1]
# W_L (D_L, D_h) ~ N(0, 1) and b_L (D_L,) ~ U(0, 2 pi), both frozen [§3.1]
def random_features(hidden):
    angles = -hidden @ W_L.T + b_L                   # (B, D_L)
    return sqrt(2 / D_L) * cos(angles)               # (B, D_L)

# SNGP prediction for one input  [alg. 2]
# Monte Carlo average over 10 samples of the logit distribution [§5]
def predict(x):
    # 1. Random features of the hidden state  [alg. 2]
    phi = random_features(hidden(x))                 # (D_L,)
    # 2. Posterior mean of each logit  [alg. 2]
    mean = phi @ betas                               # (K,)
    # 3. Posterior variance of each logit  [alg. 2]
    variance = [phi @ covariance[k] @ phi for k in range(K)]    # (K,)
    # 4. Average the softmax over Gaussian logit samples  [alg. 2] [app. A]
    samples = normal(mean, variance, size=(10, K))   # (10, K)
    return mean_over_samples(softmax(samples))       # (K,)

Hyperparameters of the GP layer and spectral normalization are given in a single table in the appendix tab. 5. The paper also states how the covariance ridge and discount factors enter the precision update app. C.

Hyperparameters

Paper values are from tab. 5 and app. C unless noted. "Code default" is the default flag value in the released scripts, which may differ from the settings used for the paper's runs (see Paper vs code).

Results

Author-reported, averaged over 10 seeds. Means are shown; the tables also give ± standard deviations over seeds, omitted here. Setup: Wide ResNet-28-10 on CIFAR-10 and CIFAR-100, with corrupted test sets (CIFAR-10-C, 15 corruption types at 5 intensities) and OOD evaluation against SVHN and against the other CIFAR set §5 app. C. BERT-Base on CLINC OOS intent detection, with in-domain intents and out-of-scope utterances §5. Hardware: 8-core Cloud TPU v2 for training app. C; the latency column is ms per example as reported, with no per-row hardware in the table. Baselines: Monte Carlo dropout and deep ensembles (both with 10 samples or models), MC-dropout GP and DUQ, and two ablations, DNN-SN and DNN-GP §5 tab. 1.

CIFAR-10, Wide ResNet-28-10 tab. 2

CIFAR-100, Wide ResNet-28-10 tab. 3

CLINC OOS, BERT-Base tab. 4

The paper's text reads these as: accuracy competitive with the deterministic network, calibration and OOD detection better than the single-model baselines, and calibration competitive with the ensemble §5. On CIFAR-100 the ensemble has the best clean accuracy and corrupted accuracy, and SNGP is not best on clean ECE or clean NLL. The 2D synthetic benchmark is qualitative only fig. 1 app. C.

Paper vs code

Code: the official repository google/uncertainty-baselines, linked from the abstract footnote, at commit 9920c89e (Apache-2.0). The SNGP layer and the spectral normalization wrappers come from google/edward2 (linked in §5 and app. C); that repository was not cloned or read for this card.

• Predictive approximation. Paper: Monte Carlo averaging with 10 samples §5 app. A. Code: the default is gp_mc_sampling = -1, which uses mean_field_logits with a tunable factor, 20 for WRN and 0.1 for BERT sngp.py sngp.py sngp.py sngp.py. The paper gives no mean-field factor, and app. A calls mean-field "future work". The WRN docstring says CIFAR-100 needs 7.5 sngp.py.
• Likelihood of the Laplace step. Paper: per-class precision from the cross-entropy likelihood, with the p(1-p) weights of §3.1 §3.1. Code: a Gaussian (squared-loss) likelihood, so all classes share one covariance sngp.py sngp.py. The code's own note calls this a difference from the paper.
• CIFAR learning rates. Paper: 0.04 for CIFAR-10 and 0.08 for CIFAR-100 app. C. Code: the default is 0.08 sngp.py, and its docstring tells CIFAR-100 users to set 0.04 sngp.py. The two sources conflict, so verify against the run configuration before reusing either.
• Length-scale. Paper: 2.0 for both tasks tab. 5. Code: 1.0 by default for WRN sngp.py, 2.0 for BERT sngp.py.
• Covariance ridge and discount. Paper: s = 0.001, m = 0.999 app. C. Code for BERT matches sngp.py. Code for WRN defaults to ridge 1.0 and discount -1, which means an exact covariance from the final epoch, with the covariance reset at each epoch start sngp.py sngp.py.
• Random features. Paper: i.i.d. N(0, 1) projections §3.1. Code for WRN defaults to orthogonal random features sngp.py. The BERT script does not set this flag in the file read here; its feature type is therefore not verified.
• Input projection and normalization for WRN. Paper: projection to 128 dimensions and input layer normalization on tab. 5 app. C. Code defaults: projection off, normalization off sngp.py.
• Batch size for BERT. Paper: 32 per core app. C. Code: 64 per core sngp.py. The CIFAR script's batch size was not checked.
• OOD score. Paper: Dempster-Shafer uncertainty for CIFAR OOD app. C. Code: the flag defaults to off and falls back to one minus the maximum probability sngp.py sngp.py. The code has a TODO noting that the Dempster-Shafer score computed from Monte Carlo logits is degenerate sngp.py.
• Spectral normalization scope. Matches: BERT normalizes only the pooler app. C sngp.py, with bound 0.95; WRN bound 6.0 sngp.py.
• Hyperparameters that match. Epochs (250 WRN, 40 BERT), BERT learning rate (5e-5), spectral bounds, and hidden width (1024 WRN, 2048 BERT) agree with the paper.

The released evaluation scripts were only partly read; the comparison above covers the flags, the predictive path and the covariance update.

Pitfalls

• Precision update, two formulas. §3.1 writes the minibatch update as (1 - m) * prev + m * batch, while app. C writes m * prev + (1 - m) * batch. With m = 0.999 the two differ by almost everything: §3.1 would keep only the latest minibatch. Use the app. C form for the paper's settings, and check it by comparing with a single exact final-epoch pass.
• Initialization of the precision. Alg. 1 initializes the precision to I; app. C initializes it to s·I with s = 0.001. The two give different posterior variances on small data, and tab. 5 lists the app. C value. Pick one and state it.
• Sign of the random projection. §3.1 writes cos(-W_L h + b_L), while alg. 2 writes cos(W_L h + b_L). Since W_L is frozen, using either sign consistently gives a valid model, but mixing them between training and prediction does not.
• L2 term. §3.1 includes ½‖β‖² in the log-posterior, but tab. 5 sets the L2 coefficient to 0. The paper does not say which governs the reported runs.
• Predictive averaging. Paper results use Monte Carlo averaging of 10 samples §5. The released code's default is mean-field, with a factor that the paper does not state. Expect different calibration if you swap one for the other, and state which one you use.
• Length-scale. §3.1 gives no formula that includes the RBF length-scale, yet tab. 5 lists ℓ = 2.0. Where the length-scale enters the random features is not written in the paper; the code delegates it to edward2, which was not read.
• Spectral bound is not a guarantee. The bound c with c ≤ 1 is only a sufficient condition for distance preservation §3.2. The authors use c = 6.0 for conv layers and say the right bound must be found by search; app. C recommends a grid over c starting at 0.9 app. C. The power iteration is a single step per training step, so the bound is approximate.
• Cost that the abstract hides. The precision update needs a full pass over the training data at the final epoch, with O(D_L²) work per example app. A. The paper's per-class form stores K matrices of size D_L × D_L; the code uses one shared matrix. The paper's latency numbers are one forward pass plus a D_L² variance term app. A tab. 2 tab. 4.
• Batch-scaled learning rate. The code multiplies the base learning rate by batch size over 128 sngp.py. The paper states a learning rate and a per-core batch size, but not how they combine, so the effective rate depends on the core count.
• Convolutional layers. The paper's bi-Lipschitz argument covers residual blocks; for convolutions the authors need a looser bound than the theory supports app. C §6. Treat the conv results as empirical.
• Dempster-Shafer score. The app. C metric is the one used in the tables, but the released code's own TODO says the Monte Carlo version is degenerate sngp.py. Check which OOD score any reimplementation uses before comparing with the tables.
• Figures. fig. 1 is qualitative (two synthetic datasets, uncertainty surfaces); it does not provide numbers.

Lineage

• Builds on: spectral normalization for residual networks, as used in Miyato et al. 2018 and in the Lipschitz analysis of Bartlett et al. 2018 and Behrmann et al. 2019 §3.2; random Fourier features (Rahimi and Recht 2008) and the Laplace approximation to GP posteriors (Rasmussen and Williams 2006) §3.1. Neither is a card here.
• Compared against: deterministic networks, Monte Carlo dropout, deep ensembles, MC-dropout GP and DUQ, the single-model methods in the related work section §5 tab. 1. Ablations DNN-SN and DNN-GP isolate the two components §5.
• Superseded by: not known from content already in this wiki.