Is the Last Layer Sufficient for Uncertainty Quantification?

Epistemic uncertainty quantification (UQ) for deep neural networks (DNNs) is a requirement for safe adoption of AI in mission-critical settings. Several leading methods for UQ linearize DNNs to form Bayesian Generalized Linear Models (GLMs), where epistemic uncertainty is modeled via the predictive posterior distribution. Linearizing around the parameters of the final connected layer of a DNN is a commonly used approximation for reducing the computational burden of such GLMs, though it is often believed to come at the cost of degraded performance. In this work, we compare GLMs arising from full-network and last-layer linearization using both theoretical and empirical approaches. We first employ tools from random matrix theory to conduct a theoretical comparison; this analysis reveals no meaningful improvement in the UQ capabilities of full linearization. Coupled with a large-scale empirical evaluation across a range of modern machine learning tasks, we arrive at the following conclusion: a last-layer approximation yields comparable UQ performance while offering substantially improved computational efficiency.

How can you tell if your favourite neural network actually knows what it's talking about? To answer this question, we measure how uncertain a model is when giving a prediction. This can be measured many ways, but in this paper we focus on how different the answer would be if you queried the model many times using the same question. This is called Bayesian uncertainty quantification, and using this framework for very large models has been a recent challenge in AI research. One way to accomplish this is to replace a large neural network with a simple, linear approximation. We can construct this approximation using information from all of the layers of the neural network, or using only information from the final layer of the neural network. This second method makes computation much cheaper, but it was often assumed to give worse performance. In our work, we theoretically show that, surprisingly, there should be no difference in performance. We have validated this by comparing these methods on a large range of problems, where again we found no significant difference in performance between the two methods.