Richer Bayesian Last Layers with Subsampled NTK Features

TL;DR

The paper tackles underestimation of epistemic uncertainty in Bayesian Last Layers by introducing Rich-BLL, a scalable method that corrects the BLL posterior by projecting non-last-layer NTK features onto the last-layer feature space. It constructs a low-dimensional transformation $\phi^{B}(x)=B\phi^{r}(x)$ (or $\phi^{L}(x)=L^{\top}\phi^{r}(x)$) to capture variations from earlier layers, yielding a posterior covariance $S^{B}$ that satisfies $S^{B}\succeq S^{\mathrm{bll}}$, with a concentration bound on the projection matrix and a subsampling scheme that preserves accuracy at a reduced cost. The method significantly reduces computational complexity to $O(Nr^2 + r^3)$ and provides finite-sample guarantees for both the projection and subsampling steps. Empirically, Rich-BLL improves calibration and uncertainty estimates across UCI regression, contextual bandits, and image classification with strong OOD performance, while matching or approaching more expensive uncertainty baselines even when using subsampling. This offers a practical, theoretically grounded enhancement for uncertainty quantification in large neural networks.

Abstract

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 sub-sampling. 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.

Figures (2)

• Figure 1: Predictive uncertainty comparison for a 1D regression problem. NNGP (BLL) underestimates epistemic uncertainty relative to the NTK-GP, while our NTK approximation (Rich-BLL) recovers richer uncertainty at BLL cost, even when using uniform subsampling.
• Figure 2: Predictive mean and uncertainty for a 1D classification toy problem. From left to right: deterministic MAP backbone, Bayesian last layer (NNGP), Rich-BLL (ours), and Rich-BLL subsampling. Shaded regions indicate predictive uncertainty.

Theorems & Definitions (23)

• Theorem 3.1
• Theorem 3.2
• Theorem 3.3
• Theorem 3.4
• Theorem 3.5
• Theorem 1.1
• proof : Proof of Theorem \ref{['theorem: Woodbury NNGP']}
• Theorem 1.1
• proof : Proof of Theorem \ref{['thm: choleskywoodbury']}
• Theorem 1.1