Technical Summary: Is the Last Layer Sufficient for Uncertainty Quantification?

Problem Statement

The adoption of Deep Neural Networks (DNNs) in mission-critical settings is hindered by a lack of reliable Epistemic Uncertainty Quantification (UQ). While Bayesian methods offer a natural framework for modeling uncertainty via the predictive posterior distribution, leading approaches such as Deep Ensembles (DE) or full Bayesian Neural Networks (BNNs) often incur prohibitive computational costs. A common approximation to reduce this burden is to linearize the DNN around its parameters to form a Bayesian Generalized Linear Model (GLM). However, a prevailing assumption in the field is that linearizing only the final connected layer (the "last-layer approximation") significantly degrades UQ performance compared to linearizing the entire network (full linearization). This paper investigates whether this trade-off between computational efficiency and UQ fidelity is necessary.

Methodology

The authors employ a dual approach combining theoretical analysis via Random Matrix Theory (RMT) and large-scale empirical evaluation.

Theoretical Analysis

The study compares two Bayesian GLMs derived from a randomly initialized DNN:

1. LL-GLM (Last-Layer GLM): Formed by linearizing only the final connected layer parameters. This induces the Conjugate Kernel (CK).
2. DNN-GLM (Full-Linearization GLM): Formed by linearizing all network parameters. This induces the Neural Tangent Kernel (NTK).

The authors analyze the Bayes Free Energy (BFE), defined as the negative log of the marginal likelihood, which serves as a key measure of model quality and UQ capability. They derive the limiting BFE under a double-asymptotic regime where both the network width and the number of training data points tend to infinity proportionally. This analysis utilizes tools from RMT to characterize the deterministic limiting spectra of the CK and NTK Gram matrices.

Empirical Evaluation

To validate theoretical findings in practical settings, the authors introduce LinearSampling, a lightweight PyTorch package. This framework provides an algorithm to sample from the predictive posterior of both LL-GLMs and DNN-GLMs. The algorithm employs a "sample-then-optimize" strategy, initializing parameters near the trained weights and using (Stochastic) Gradient Descent to solve the linearized optimization problem. This approach scales to large models, including those with millions of parameters.

The empirical evaluation spans regression, classification, and language modeling tasks, comparing LL-GLMs against DNN-GLMs and other state-of-the-art UQ methods (e.g., Deep Ensembles, SWAG, MC-Dropout, LLA). Metrics include Expected Calibration Error (ECE), Gaussian Negative-Log Likelihood (NLL), and specialized uncertainty metrics like VARROC (AUC-ROC using variance of softmax predictions) for both in-distribution and out-of-distribution (OOD) detection.

Key Contributions

1. Theoretical Characterization of Limiting BFE: The authors derive the limiting BFE for GPs employing CK and NTK kernels under mild assumptions on network components and data distributions. They characterize this limit as the unique solution to a system of implicit fixed-point equations.
2. Robustness and Descent Analysis: Theoretically, they demonstrate that both kernels exhibit robustness to input perturbations and descent phenomena (including double- and triple-descent behavior). They identify that the NTK only shows a theoretical advantage in the highly over-sampled regime (where data points far exceed network width) at initialization, but this advantage vanishes after training.
3. Large-Scale Empirical Validation: Through extensive experiments across diverse tasks, the paper demonstrates that LL-GLMs consistently match the UQ performance of full DNN-GLMs. In many cases, the LL-GLM performs comparably to the gold-standard Deep Ensembles.
4. Scalable Sampling Framework: The release of LinearSampling, a scalable algorithm that enables efficient posterior sampling from Bayesian GLMs for large language models and deep vision architectures, overcoming previous computational bottlenecks.

Results

• Theoretical Findings: The analysis reveals no meaningful improvement in UQ capabilities (measured by BFE) for full linearization over last-layer linearization in most regimes. The only regime where the full NTK theoretically outperforms the CK is when the number of data points is significantly larger than the network width, a scenario the authors suggest may be an artifact of random initialization rather than a fundamental limitation, as the gap disappears after training.
• Empirical Findings:
  • Regression: On UCI regression datasets, LL-GLMs achieved comparable or sometimes superior performance to DNN-GLMs in terms of RMSE, NLL, and ECE.
  • Classification: On image classification tasks (MNIST, FashionMNIST, CIFAR-10/100, ImageNet), LL-GLMs matched the UQ performance of DNN-GLMs. Both methods successfully captured "distance-aware" variance, increasing uncertainty for points far from training data.
  • Language Modeling: On a GPT-2 model fine-tuned for sentiment analysis, the LL-GLM achieved statistically indistinguishable VARROC-ID scores compared to the DNN-GLM.
  • Efficiency: Crucially, the LL-GLM achieved these results with substantially lower computational and memory costs. In the GPT-2 experiment, the LL-GLM required 10x less wall-time and 3.5x less memory than the DNN-GLM.

Significance and Conclusion

The paper concludes that the last-layer approximation is sufficient for high-quality Uncertainty Quantification. The authors argue that this finding suggests epistemic uncertainty in DNNs arises primarily from uncertainty in how the network features are used (the final layer weights) rather than uncertainty in the features themselves (the representations learned by earlier layers).

This result challenges the assumption that full-network linearization is necessary for accurate UQ. By validating that the computationally cheaper last-layer approach yields comparable performance, the work advocates for the broader adoption of last-layer linearization in Bayesian deep learning, particularly for large-scale models where full-network Bayesian inference is infeasible. The authors note limitations regarding the theoretical connection to trained networks under realistic assumptions and the current focus on epistemic rather than aleatoric uncertainty, identifying these as avenues for future work.