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

TL;DR

This paper introduces a new gradient-based SMC training method for partial Bayesian neural networks, improving scalability, predictive accuracy, and training efficiency over existing approaches.

Contribution

A novel SMC-based training approach for pBNNs that uses guided proposals and gradient-based Markov kernels, enhancing scalability and performance.

Findings

Outperforms state-of-the-art in predictive accuracy and loss

Scales well with larger batch sizes, reducing training time

Achieves better performance with high-dimensional problems

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.