bde: A Python Package for Bayesian Deep Ensembles via MILE

2026-05-13

AI Summary

This work addresses the lack of efficient, user-friendly, and uncertainty-aware open-source tools for Bayesian deep ensembles on tabular data. Building upon the Microcanonical Langevin Ensembles (MILE) sampling inference framework, the authors implement highly efficient Markov chain Monte Carlo (MCMC) sampling using JAX and provide a scikit-learn–compatible Python interface. The resulting tool substantially improves training efficiency, uncertainty quantification, and scalability of Bayesian deep ensembles for both regression and classification tasks, while maintaining high predictive performance and ease of use.

Abstract

bde is a user-friendly Python package for Bayesian Deep Ensembles with a particular focus on tabular data. Built on an efficient JAX implementation of the sampling-based inference method Microcanonical Langevin Ensembles (MILE), it provides scikit-learn compatible estimators for fast training, efficient Markov Chain Monte Carlo sampling, and uncertainty quantification in both regression and classification tasks.

Problem

• Bayesian Deep Ensembles
• tabular data
• uncertainty quantification
• MILE

Innovation

• Bayesian Deep Ensembles
• MILE
• JAX
• uncertainty quantification
• tabular data