Microcanonical Langevin Ensembles: Advancing the Sampling of Bayesian Neural Networks

TL;DR

This paper introduces Microcanonical Langevin Ensembles, a novel sampling method for Bayesian Neural Networks that achieves faster, more reliable inference with better uncertainty quantification, addressing key challenges in probabilistic deep learning.

Contribution

The paper proposes an ensembling approach combined with Microcanonical Langevin Monte Carlo to improve sampling efficiency, robustness, and predictability in Bayesian Neural Networks.

Findings

Achieves up to tenfold speedup over No-U-Turn Sampler

Maintains or improves predictive performance and uncertainty quantification

Enhances method predictability and parallelization capabilities

Abstract

Despite recent advances, sampling-based inference for Bayesian Neural Networks (BNNs) remains a significant challenge in probabilistic deep learning. While sampling-based approaches do not require a variational distribution assumption, current state-of-the-art samplers still struggle to navigate the complex and highly multimodal posteriors of BNNs. As a consequence, sampling still requires considerably longer inference times than non-Bayesian methods even for small neural networks, despite recent advances in making software implementations more efficient. Besides the difficulty of finding high-probability regions, the time until samplers provide sufficient exploration of these areas remains unpredictable. To tackle these challenges, we introduce an ensembling approach that leverages strategies from optimization and a recently proposed sampler called Microcanonical Langevin Monte Carlo (MCLMC) for efficient, robust and predictable sampling performance. Compared to approaches based on the state-of-the-art No-U-Turn Sampler, our approach delivers substantial speedups up to an order of magnitude, while maintaining or improving predictive performance and uncertainty quantification across diverse tasks and data modalities. The suggested Microcanonical Langevin Ensembles and modifications to MCLMC additionally enhance the method's predictability in resource requirements, facilitating easier parallelization. All in all, the proposed method offers a promising direction for practical, scalable inference for BNNs.