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

2025-05-01

AI Summary

To address the poor scalability and low training efficiency of partially Bayesian neural networks (pBNNs) in high-dimensional settings, this paper proposes a gradient-enhanced sequential Monte Carlo (SMC) method. It is the first to jointly incorporate gradient information into both the proposal distribution and the adaptive Markov transition kernel, enabling efficient posterior inference for nonparametric pBNNs. The method significantly improves sampling efficiency and convergence speed in high-dimensional parameter spaces, supports substantially larger batch sizes, and markedly reduces training time. It outperforms state-of-the-art approaches in both predictive accuracy and optimal loss. The core contribution lies in the novel gradient-guided SMC framework, which—uniquely for pBNNs—simultaneously optimizes inference accuracy and computational scalability.

Technology Category

• Machine Learning: Bayesian Learning
• Reasoning under Uncertainty: Probabilistic Programming
• 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
• Search and Retrieval-Augmented AI: Web learning to rank, online learning, and counterfactual learning for ranking
• Semantics and Knowledge: Methods to enhance, augment, integrate or synergize semantic models such as knowledge graphs and LLMs

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.

Problem

• Improving training of partial Bayesian neural networks
• Enhancing scalability in high-dimensional problems
• Optimizing predictive performance and training efficiency

Innovation

• Gradient-based Markov kernels in SMC
• Guided proposals for pBNNs training
• Scalable SMC method for high dimensions