Sampling-based inference for large linear models, with application to linearised Laplace

TL;DR

This paper introduces a scalable, sample-based Bayesian inference method for large linear models, enabling uncertainty quantification in neural networks like ResNet-18 and ResNet-50 on large datasets, overcoming previous computational limitations.

Contribution

The authors develop a scalable inference technique for conjugate Gaussian multi-output linear models and a hyperparameter selection method, improving the application of linearised Laplace in large neural networks.

Findings

Enabled linearised neural network inference on CIFAR100 with ResNet-18

Performed inference on ImageNet with ResNet-50

Applied to high-resolution tomographic reconstruction with U-Net

Abstract

Large-scale linear models are ubiquitous throughout machine learning, with contemporary application as surrogate models for neural network uncertainty quantification; that is, the linearised Laplace method. Alas, the computational cost associated with Bayesian linear models constrains this method's application to small networks, small output spaces and small datasets. We address this limitation by introducing a scalable sample-based Bayesian inference method for conjugate Gaussian multi-output linear models, together with a matching method for hyperparameter (regularisation) selection. Furthermore, we use a classic feature normalisation method (the g-prior) to resolve a previously highlighted pathology of the linearised Laplace method. Together, these contributions allow us to perform linearised neural network inference with ResNet-18 on CIFAR100 (11M parameters, 100 outputs x 50k datapoints), with ResNet-50 on Imagenet (50M parameters, 1000 outputs x 1.2M datapoints) and with a U-Net on a high-resolution tomographic reconstruction task (2M parameters, 251k output~dimensions).