Utilising Gradient-Based Proposals Within Sequential Monte Carlo Samplers for Training of Partial Bayesian Neural Networks

TL;DR

The paper addresses scalable Bayesian inference for partial Bayesian neural networks (pBNNs) by integrating gradient-based proposals into sequential Monte Carlo (SMC) samplers, allowing non-parametric posterior estimation of the stochastic parameters $p(m{ heta} | extbf{y}_{1:N}, m{ ho})$. It introduces guided open-horizon SMC (GOHSMC) with gradient-based Langevin kernels, deriving a tractable weight update by leveraging the reversibility of Langevin dynamics to leave invariant the current posterior $p(m{ heta} | m{ ho})$. Empirical results on UCI regression and image classification show improved predictive performance and faster training with larger batch sizes, with larger first-layer stochasticity amplifying the benefits of Langevin dynamics. Overall, pBNNs under gradient-based GOHSMC provide competitive uncertainty quantification relative to VI and SGHMC while enabling scalable, data-efficient Bayesian training for neural networks.

Abstract

Partial Bayesian neural networks (pBNNs) have been shown to perform competitively with fully Bayesian neural networks while only having a subset of the parameters be stochastic. Using sequential Monte Carlo (SMC) samplers as the inference method for pBNNs gives a non-parametric probabilistic estimation of the stochastic parameters, and has shown improved performance over parametric methods. In this paper we introduce a new SMC-based training method for pBNNs by utilising a guided proposal and incorporating gradient-based Markov kernels, which gives us better scalability on high dimensional problems. We show that our new method outperforms the state-of-the-art in terms of predictive performance and optimal loss. We also show that pBNNs scale well with larger batch sizes, resulting in significantly reduced training times and often better performance.

Figures (7)

• Figure 1: Validation loss comparison between different methods on different batch sizes for the Red Wine Quality and California Housing Dataset respectively, averaged over 5 runs.
• Figure 2: Validation loss and accuracy over the lifetime of each pBNN method and batch size for the MNIST Dataset, averaged over 10 runs.
• Figure 3: Validation loss and accuracy over the lifetime of each pBNN method and batch size for the FashionMNIST Dataset, averaged over 10 runs.
• Figure 4: Validation loss and accuracy over the lifetime of each pBNN method and batch size for the CIFAR10 Dataset, averaged over 5 runs.
• Figure 5: AUROC curves for each method averaged over 5 runs.