Sampling from Bayesian Neural Network Posteriors with Symmetric Minibatch Splitting Langevin Dynamics

2024-10-14

arXiv.org

AI Summary

To address the high bias and low efficiency of posterior sampling in Bayesian neural networks (BNNs) under large-scale data regimes, this paper proposes Symmetric Mini-batch Splitting Langevin Dynamics (SMB-LD). SMB-LD introduces a novel coupling mechanism that integrates symmetric forward/backward mini-batch traversal with UBU-type Langevin splitting, achieving O(h²d¹ᐟ²) parameter-space sampling bias per iteration—using only a single mini-batch—thereby substantially outperforming conventional stochastic-gradient MCMC methods. The algorithm unifies kinetic Langevin dynamics, symmetric splitting integrators, and BNN posterior modeling. Extensive experiments on Fashion-MNIST, CelebA, and chest X-ray datasets demonstrate that SMB-LD significantly improves posterior predictive calibration, surpassing both standard training and stochastic weight averaging (SWA). This work advances scalable, accurate Bayesian inference for deep learning by reconciling computational efficiency with theoretical fidelity in posterior approximation.

Technology Category

• Machine Learning: Bayesian Learning
• Search and Optimization: Sampling/Simulation-based Search
• Reasoning under Uncertainty: Stochastic Optimization

Application Category

• Graph Algorithms and Modeling for the Web: Graph neural networks and deep learning approaches for Web-related graphs
• Search and Retrieval-Augmented AI: Large language models for search
• User Modeling, Personalization and Recommendation: Large Language Models (LLM) for user modeling and recommendation

Abstract

We propose a scalable kinetic Langevin dynamics algorithm for sampling parameter spaces of big data and AI applications. Our scheme combines a symmetric forward/backward sweep over minibatches with a symmetric discretization of Langevin dynamics. For a particular Langevin splitting method (UBU), we show that the resulting Symmetric Minibatch Splitting-UBU (SMS-UBU) integrator has bias $O(h^2 d^{1/2})$ in dimension $d>0$ with stepsize $h>0$, despite only using one minibatch per iteration, thus providing excellent control of the sampling bias as a function of the stepsize. We apply the algorithm to explore local modes of the posterior distribution of Bayesian neural networks (BNNs) and evaluate the calibration performance of the posterior predictive probabilities for neural networks with convolutional neural network architectures for classification problems on three different datasets (Fashion-MNIST, Celeb-A and chest X-ray). Our results indicate that BNNs sampled with SMS-UBU can offer significantly better calibration performance compared to standard methods of training and stochastic weight averaging.

Problem

• Scalable sampling for Bayesian Neural Networks
• Control sampling bias with symmetric minibatch splitting
• Improve calibration in neural network predictions

Innovation

• Symmetric Minibatch Splitting Langevin Dynamics
• UBU integrator with O(h^2 d^1/2) bias
• Improved calibration in Bayesian Neural Networks