Richer Bayesian Last Layers with Subsampled NTK Features

Bayesian Last Layers (BLLs) provide a convenient and computationally efficient way to estimate uncertainty in neural networks. However, they underestimate epistemic uncertainty because they apply a Bayesian treatment only to the final layer, ignoring uncertainty induced by earlier layers. We propose a method that improves BLLs by leveraging a projection of Neural Tangent Kernel (NTK) features onto the space spanned by the last-layer features. This enables posterior inference that accounts for variability of the full network while retaining the low computational cost of inference of a standard BLL. We show that our method yields posterior variances that are provably greater or equal to those of a standard BLL, correcting its tendency to underestimate epistemic uncertainty. To further reduce computational cost, we introduce a uniform subsampling scheme for estimating the projection matrix and for posterior inference. We derive approximation bounds for both types of subsampling. Empirical evaluations on UCI regression, contextual bandits, image classification, and out-of-distribution detection tasks in image and tabular datasets, demonstrate improved calibration and uncertainty estimates compared to standard BLLs and competitive baselines, while reducing computational cost.

Neural networks often make accurate predictions, but it is hard to know when they are uncertain. Bayesian methods can help quantify uncertainty, but applying them to large neural networks is usually expensive. A cheaper alternative is to use a Bayesian model only on the last layer of the network, but this ignores uncertainty from earlier layers and can be overconfident. We propose a method that improves Bayesian last-layer uncertainty estimates by incorporating information from earlier layers through a scalable approximation to the Neural Tangent Kernel. This keeps the cost close to standard last-layer methods while capturing more of the network’s uncertainty. We provide theoretical guarantees for when this approximation is accurate and show empirically that it improves uncertainty estimates on regression, image classification, out-of-distribution detection, and decision-making tasks.