Function-Space MCMC for Bayesian Wide Neural Networks

2024-08-26

arXiv.org

AI Summary

To address the degradation of weight posterior sampling in wide Bayesian neural networks (BNNs) as width increases, this paper proposes performing MCMC inference over the reparameterized weight posterior in function space. Theoretically, we establish, for the first time, that the preconditioned Crank–Nicolson (pCN) algorithm achieves unit acceptance rate in the infinite-width limit, with step size independent of network width—ensuring dimension-agnostic scalability—and develop a rigorous MCMC convergence analysis framework in the infinite-dimensional function space. Empirically, pCN significantly improves effective sample size and convergence diagnostics on real-world datasets, outperforming baselines including pCN-Langevin and underdamped Langevin dynamics. Our approach enables efficient and robust uncertainty quantification for wide BNNs.

Technology Category

• Machine Learning: Calibration & Uncertainty Quantification
• Reasoning under Uncertainty: Probabilistic Inference
• Search and Optimization: Sampling/Simulation-based Search

Application Category

• Graph Algorithms and Modeling for the Web: Algorithms and analysis for incomplete, noisy, or partially observed Web-related graphs
• Web Mining and Content Analysis: Large pretrained models with web data
• User Modeling, Personalization and Recommendation: Practical large-scale studies of user experience

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.

Problem

• Efficient sampling in wide Bayesian Neural Networks
• Analyzing mixing speeds of different MCMC algorithms
• Scalable posterior sampling with preconditioned Crank-Nicolson

Innovation

• Preconditioned Crank-Nicolson algorithm for sampling
• Langevin version enhances sampling efficiency
• Scalable sampling in wide Bayesian networks