Function-Space MCMC for Bayesian Wide Neural Networks

TL;DR

This paper explores scalable Bayesian neural network sampling methods, demonstrating that the preconditioned Crank-Nicolson algorithm becomes more effective as network width increases, with higher acceptance rates and better mixing properties.

Contribution

It introduces the use of the preconditioned Crank-Nicolson algorithm for wide neural networks and proves its acceptance probability approaches 1 as width grows, showing improved efficiency.

Findings

Acceptance probabilities approach 1 with increasing width

Preconditioned Crank-Nicolson outperforms other samplers in efficiency

Higher effective sample size in wide network configurations

Abstract

Bayesian Neural Networks represent a fascinating confluence of deep learning and probabilistic reasoning, offering a compelling framework for understanding uncertainty in complex predictive models. In this paper, we investigate the use of the preconditioned Crank-Nicolson algorithm and its Langevin version to sample from a reparametrised posterior distribution of the neural network's weights, as the widths grow larger. In addition to being robust in the infinite-dimensional setting, we prove that the acceptance probabilities of the proposed algorithms approach 1 as the width of the network increases, independently of any stepsize tuning. Moreover, we examine and compare how the mixing speeds of the underdamped Langevin Monte Carlo, the preconditioned Crank-Nicolson and the preconditioned Crank-Nicolson Langevin samplers are influenced by changes in the network width in some real-world cases. Our findings suggest that, in wide Bayesian Neural Networks configurations, the preconditioned Crank-Nicolson algorithm allows for a scalable and more efficient sampling of the reparametrised posterior distribution, as also evidenced by a higher effective sample size and improved diagnostic results compared with the other analysed algorithms.