Utilising Gradient-Based Proposals Within Sequential Monte Carlo Samplers for Training of Partial Bayesian Neural Networks
This work addresses efficient training of Bayesian neural networks for machine learning practitioners, but it is incremental as it builds on existing SMC and pBNN methods.
The paper tackles training partial Bayesian neural networks (pBNNs) by introducing a new sequential Monte Carlo (SMC) method with gradient-based proposals, resulting in improved predictive performance and reduced training times, including scalability with larger batch sizes.
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.