Scalable Linearized Laplace Approximation via Surrogate Neural Kernel

TL;DR

This work tackles scalable uncertainty estimation for deep neural networks by addressing the computational bottleneck of the Linearized Laplace Approximation (LLA), which relies on Jacobians to form the NTK. It introduces a surrogate kernel learned by a compact network g_phi that imitates the Neural Tangent Kernel via Jacobian–vector products, enabling scalable, Jacobian-free LLA for large pretrained DNNs. Key contributions include a Jacobian-free LLA approximation, demonstrated scalability and calibration improvements over existing LLA variants, and a mechanism to bias the learned kernel to boost out-of-distribution detection. The approach broadens the practical applicability of Bayesian uncertainty estimation in industrial-scale models by providing efficient, calibrated post-hoc uncertainty estimates.

Abstract

We introduce a scalable method to approximate the kernel of the Linearized Laplace Approximation (LLA). For this, we use a surrogate deep neural network (DNN) that learns a compact feature representation whose inner product replicates the Neural Tangent Kernel (NTK). This avoids the need to compute large Jacobians. Training relies solely on efficient Jacobian-vector products, allowing to compute predictive uncertainty on large-scale pre-trained DNNs. Experimental results show similar or improved uncertainty estimation and calibration compared to existing LLA approximations. Notwithstanding, biasing the learned kernel significantly enhances out-of-distribution detection. This remarks the benefits of the proposed method for finding better kernels than the NTK in the context of LLA to compute prediction uncertainty given a pre-trained DNN.