bde: A Python Package for Bayesian Deep Ensembles via MILE

# A Python Package for Bayesian Deep Ensembles via MILE
Source: [https://arxiv.org/html/2605.14146](https://arxiv.org/html/2605.14146)
Vyron ArvanitisAngelos AslanidisFaculty of Physics, LMU Munich, Munich, GermanyThese authors contributed equally\.Emanuel SommerDepartment of Statistics, LMU Munich, Munich, GermanyMunich Center for Machine Learning, Munich, GermanyThese authors contributed equally\.David RügamerDepartment of Statistics, LMU Munich, Munich, GermanyMunich Center for Machine Learning, Munich, Germany

## Summary

bde\([https://github\.com/scikit\-learn\-contrib/bde](https://github.com/scikit-learn-contrib/bde)\) is a Python package designed to bring state\-of\-the\-art sampling\-based Bayesian Deep Learning \(BDL\) to practitioners and researchers\. The package combines the speed and high\-performance capabilities of JAX andblackjax\(Bradburyet al\.,[2018](https://arxiv.org/html/2605.14146#bib.bib5); Cabezaset al\.,[2024](https://arxiv.org/html/2605.14146#bib.bib4)\)with the user\-friendly API of scikit\-learn\(Pedregosaet al\.,[2011](https://arxiv.org/html/2605.14146#bib.bib9)\)\. It targets tabular supervised learning tasks, including distributional regression and \(multi\-class\) classification, providing a seamless interface for Bayesian Deep Ensembles \(BDEs\)\(Sommeret al\.,[2024](https://arxiv.org/html/2605.14146#bib.bib2)\)specifically implementingMicrocanonical Langevin Ensembles \(MILE\)\(Sommeret al\.,[2025](https://arxiv.org/html/2605.14146#bib.bib1)\)\.

The workflow ofbdeimplements the two\-stage BDE inference process of MILE\. First, it optimizesn\_membersindependent instances of a flexibly configurable feed\-forward neural network using regularized empirical risk minimization \(with the negative log\-likelihood as loss\) via the AdamW optimizer\(Loshchilov and Hutter,[2019](https://arxiv.org/html/2605.14146#bib.bib8)\)\. Second, it transitions to a sampling phase using Microcanonical Langevin Monte Carlo\(Robniket al\.,[2023](https://arxiv.org/html/2605.14146#bib.bib6); Robnik and Seljak,[2024](https://arxiv.org/html/2605.14146#bib.bib7)\), enhanced with a tuning phase adapted for Bayesian Neural Networks\. This combination is referred to as MILE\(Sommeret al\.,[2025](https://arxiv.org/html/2605.14146#bib.bib1)\)\. In essence, the optimization of the ensemble of neural networks first finds diverse high\-likelihood modes, from which sampling then explores local posterior structure\. This process generates an ensemble of samples \(models\) that constitute an implicit posterior approximation\.

## Software design

Because optimization and sampling across ensemble members are independent,bdeexploits JAX’s parallelization and just\-in\-time compilation to scale efficiently across CPUs, GPUs, and TPUs\. Given new test data, the package approximates the posterior predictive, enabling point predictions, credible intervals, coverage estimates, and other uncertainty metrics through a unified interface\.

## State of the field

Reliable uncertainty quantification \(UQ\) is increasingly viewed as a critical component of modern machine learning systems, and BDL provides a principled framework for achieving it\(Papamarkouet al\.,[2024](https://arxiv.org/html/2605.14146#bib.bib3)\)\. While several libraries support optimization\-based approaches such as variational inference or classical Bayesian modeling, accessible tools for sampling\-based inference in Bayesian neural networks remain scarce\. Existing probabilistic programming frameworks offer MCMC but require substantial manual configuration to achieve competitive performance on neural network models\.

## Statement of need

bdeaddresses the gap outlined above by providing the first user\-friendly implementation of MILE \- a hybrid sampling technique shown to deliver strong predictive accuracy and calibrated uncertainty for Bayesian neural networks\(Sommeret al\.,[2025](https://arxiv.org/html/2605.14146#bib.bib1)\)\. By providing full scikit\-learn compatibility, the package enables seamless integration into existing machine learning workflows, allowing users to obtain principled Bayesian uncertainty estimates without specialized knowledge of MCMC dynamics, initialization strategies, or JAX internals\.

Through automated orchestration of optimization, sampling, parallelization, and predictive inference,bdeoffers a fast, reproducible, and practical solution for applying sampling\-based BDL methods to tabular supervised learning tasks\.

## Research impact statement

bdebridges the gap between high\-performance MCMC research and practical data science\. By providing a curated implementation of MILE for tabular data, it enables researchers and practitioners alike to easily apply BDEs\. Its inclusion in thescikit\-learn\-contriborganization ensures adherence to rigorous software standards and API consistency, making it a trusted, community\-ready tool for reproducible UQ in tabular machine learning\.

## Usage example

The following example illustrates a regression task with UQ usingbdein only a few lines of code\. Training inputs are assumed to be preprocessed and normalized\. The workflow specifies the ensemble and sampling hyperparameters, fits the model, and obtains posterior predictive quantities, including predictive moments, credible intervals, and raw ensemble outputs\.

frombdeimportBdeRegressor

regressor=BdeRegressor\(

n\_members=8,

hidden\_layers=\[16,16\],

activation="relu",

epochs=1000,

validation\_split=0\.15,

lr=1e\-3,

weight\_decay=1e\-4,

patience=20,

warmup\_steps=5000,

n\_samples=200,

n\_thinning=10,

desired\_energy\_var\_start=0\.5,

desired\_energy\_var\_end=0\.1,

\)

regressor\.fit\(x=X\_train,y=y\_train\)

means,sigmas=regressor\.predict\(X\_test,mean\_and\_std=True\)

means,intervals=regressor\.predict\(X\_test,credible\_intervals=\[0\.1,0\.9\]\)

raw=regressor\.predict\(X\_test,raw=True\)

Classification follows analogously usingBdeClassifier\.

## Regression benchmark

We provide a small benchmark ofbdeon theairfoil\(Dua and Graff,[2017](https://arxiv.org/html/2605.14146#bib.bib10)\)and thebikesharing\(Fanaee\-T,Hadi,[2013](https://arxiv.org/html/2605.14146#bib.bib11)\)datasets\. We report mean predictive performance \(RMSE\), UQ metrics \(NLL in the distributional and mean regression formulation\), reported as mean±\\pmstandard deviation over 5 independent runs\. The results show competitive out\-of\-the\-box performance of BDE especially in UQ with its native distributional regression capability\.

All models used 10 CPU cores without additional tuning\. For BDE, we generated 1000 posterior samples from a feed\-forward neural network with four hidden layers of width 16\. While even highly optimized BDEs generally incur a higher computational cost than optimization\-based competitors due to the iterative nature of MCMC sampling, this investment enables the construction of a flexible, non\-parametric posterior approximation\. This trade\-off, as demonstrated above, yields strong performance and rigorous epistemic UQ, e\.g\. through calibrated credible intervals\. All experimental configurations are shipped in the released package to ensure reproducibility\.

## AI usage disclosure

Generative AI was used via GitHub Copilot for local, line\- or block\-level code autocompletion during software development, using the Claude Sonnet 3\.7 and 4 models\. Codex was additionally used to assist with the initial draft of the package documentation and minor refactors to ensure compatibility with thescikit\-learnAPI\. No AI tools were used for ideation, architectural or design decisions, code review, or testing strategy\. All AI\-generated suggestions were critically reviewed, modified where necessary, and fully validated by the authors, who retain complete responsibility\.

## Acknowledgments

DR’s research is funded by the Deutsche Forschungsgemeinschaft \(DFG, German Research Foundation\) – 578966082\.

## References

- J\. Bradbury, R\. Frostig, P\. Hawkins, M\. J\. Johnson, C\. Leary, D\. Maclaurin, G\. Necula, A\. Paszke, J\. VanderPlas, S\. Wanderman\-Milne, and Q\. Zhang \(2018\)JAX: composable transformations of Python\+NumPy programs\.External Links:[Link](http://github.com/google/jax)Cited by:[Summary](https://arxiv.org/html/2605.14146#Sx1.p1.1)\.
- A\. Cabezas, A\. Corenflos, J\. Lao, and R\. Louf \(2024\)BlackJAX: Composable Bayesian inference in JAX\.External Links:2402\.10797Cited by:[Summary](https://arxiv.org/html/2605.14146#Sx1.p1.1)\.
- D\. Dua and C\. Graff \(2017\)UCI machine learning repository\.University of California, Irvine, School of Information and Computer Sciences\.External Links:[Link](http://archive.ics.uci.edu/ml)Cited by:[Regression benchmark](https://arxiv.org/html/2605.14146#Sx7.p1.1)\.
- Fanaee\-T,Hadi \(2013\)Bike Sharing Dataset\.Note:UCI Machine Learning RepositoryCited by:[Regression benchmark](https://arxiv.org/html/2605.14146#Sx7.p1.1)\.
- I\. Loshchilov and F\. Hutter \(2019\)Decoupled weight decay regularization\.InInternational Conference on Learning Representations,Cited by:[Summary](https://arxiv.org/html/2605.14146#Sx1.p2.1)\.
- T\. Papamarkou, M\. Skoularidou, K\. Palla, L\. Aitchison, J\. Arbel, D\. Dunson, M\. Filippone, V\. Fortuin, P\. Hennig, J\. M\. Hernández\-Lobato, A\. Hubin, A\. Immer, T\. Karaletsos, M\. E\. Khan, A\. Kristiadi, Y\. Li, S\. Mandt, C\. Nemeth, M\. A\. Osborne, T\. G\. J\. Rudner, D\. Rügamer, Y\. W\. Teh, M\. Welling, A\. G\. Wilson, and R\. Zhang \(2024\)Position: Bayesian Deep Learning is Needed in the Age of Large\-Scale AI\.InProceedings of the 41st International Conference on Machine Learning,Cited by:[State of the field](https://arxiv.org/html/2605.14146#Sx3.p1.1)\.
- F\. Pedregosa, G\. Varoquaux, A\. Gramfort, V\. Michel, B\. Thirion, O\. Grisel, M\. Blondel, P\. Prettenhofer, R\. Weiss, V\. Dubourg, J\. Vanderplas, A\. Passos, D\. Cournapeau, M\. Brucher, M\. Perrot, and E\. Duchesnay \(2011\)Scikit\-learn: machine learning in Python\.Journal of Machine Learning Research12,pp\. 2825–2830\.Cited by:[Summary](https://arxiv.org/html/2605.14146#Sx1.p1.1)\.
- J\. Robnik, G\. B\. De Luca, E\. Silverstein, and U\. Seljak \(2023\)Microcanonical hamiltonian monte carlo\.The Journal of Machine Learning Research24\(1\),pp\. 14696–14729\.Cited by:[Summary](https://arxiv.org/html/2605.14146#Sx1.p2.1)\.
- J\. Robnik and U\. Seljak \(2024\)Fluctuation without dissipation: microcanonical langevin monte carlo\.InSymposium on Advances in Approximate Bayesian Inference,pp\. 111–126\.Cited by:[Summary](https://arxiv.org/html/2605.14146#Sx1.p2.1)\.
- E\. Sommer, J\. Robnik, G\. Nozadze, U\. Seljak, and D\. Rügamer \(2025\)Microcanonical langevin ensembles: advancing the sampling of bayesian neural networks\.InThe Thirteenth International Conference on Learning Representations,Cited by:[Summary](https://arxiv.org/html/2605.14146#Sx1.p1.1),[Summary](https://arxiv.org/html/2605.14146#Sx1.p2.1),[Statement of need](https://arxiv.org/html/2605.14146#Sx4.p1.1)\.
- E\. Sommer, L\. Wimmer, T\. Papamarkou, L\. Bothmann, B\. Bischl, and D\. Rügamer \(2024\)Connecting the dots: is mode\-connectedness the key to feasible sample\-based inference in bayesian neural networks?\.InProceedings of the 41st International Conference on Machine Learning,Cited by:[Summary](https://arxiv.org/html/2605.14146#Sx1.p1.1)\.