🤖 machine learning

bde: A Python Package for Bayesian Deep Ensembles via MILE

The paper introduces "bde," a user-friendly Python package built on JAX that implements Microcanonical Langevin Ensembles (MILE) to provide scikit-learn compatible Bayesian Deep Ensembles for efficient training, sampling, and uncertainty quantification in tabular data regression and classification tasks.

Original authors: Vyron Arvanitis, Angelos Aslanidis, Emanuel Sommer, David Rügamer

Published 2026-05-15

📖 4 min read☕ Coffee break read

CC BY 4.0

Original authors: Vyron Arvanitis, Angelos Aslanidis, Emanuel Sommer, David Rügamer

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are trying to predict the weather. You could ask a single expert meteorologist for a forecast. They might be right, but if they are wrong, you have no idea how wrong they are, or why they made that mistake.

Now, imagine instead that you hire a team of eight different meteorologists. You ask them all to study the data and give you their best guess. If they all agree, you feel very confident. If they disagree, you know there is uncertainty, and you can look at the spread of their answers to understand the risk.

bde is a new computer tool (a Python package) that does exactly this, but for artificial intelligence. It helps researchers and data scientists build "teams" of AI models to make predictions and, crucially, to tell you how confident those predictions are.

Here is a breakdown of how it works, using simple metaphors:

1. The Problem: The "Black Box" of AI

Most modern AI tools are like a "black box." You put data in, and a number comes out. But the box doesn't tell you if it's guessing wildly or if it's absolutely sure. In the real world, knowing how sure the AI is (called Uncertainty Quantification) is just as important as the prediction itself.

2. The Solution: A Team of AI Experts (Bayesian Deep Ensembles)

The bde package creates a Bayesian Deep Ensemble. Think of this as a committee of AI models.

The Team: Instead of one AI, it trains eight (or more) slightly different versions of the same AI.
The Goal: Each version learns the data in its own unique way. When they vote on a prediction, the group gives you a much smarter answer than any single member could.

3. How It Works: The Two-Step Dance (MILE)

The paper introduces a specific method called MILE (Microcanonical Langevin Ensembles). You can think of this as a two-step dance routine the AI performs to find the best answers:

Step 1: The Sprint (Optimization)
First, the AI team runs a fast sprint to find the "high ground." They quickly scan the data to find the most likely answers (the peaks of a mountain). This is done using a standard, fast training method.
Step 2: The Hike (Sampling)
Once they find the high ground, they don't just stop there. They start "hiking" around the peak. This is the sampling part. They explore the area around the best answer to see how the landscape looks. Are there other peaks nearby? Is the ground shaky?
- Why do this? This hiking phase allows the AI to understand the shape of the uncertainty. It doesn't just give you one number; it gives you a map of possibilities.

4. Why It's Fast and Powerful

Usually, doing this "hiking" (sampling) is incredibly slow and requires supercomputers. bde solves this by using a special engine called JAX.

The Analogy: Imagine you have a single chef trying to cook eight different meals one by one. It takes a long time. bde is like having a kitchen with eight chefs working in perfect sync at the same time. It uses your computer's power (CPUs, GPUs, or TPUs) to run all the AI models simultaneously, making the process fast enough for everyday use.

5. Who Is It For?

The paper states that bde is designed for tabular data—think of spreadsheets with rows and columns (like sales records, weather data, or bike-sharing logs).

It connects easily to scikit-learn, which is the most popular toolkit for data scientists. This means if you already know how to use standard data tools, you can use bde without needing to learn complex math or how to program the "engine" from scratch.
It handles both Regression (predicting a number, like "how many bikes will be rented?") and Classification (predicting a category, like "will it rain or not?").

6. What Did They Prove?

The authors tested bde on real-world datasets (like wind speed on an airfoil and bike-sharing demand).

The Result: Their tool performed as well as, or better than, other top-tier AI models in terms of accuracy.
The Big Win: More importantly, it provided better uncertainty estimates. It was better at telling you when it was unsure, which is vital for making safe, reliable decisions.

Summary

bde is a user-friendly tool that lets you build a "committee" of AI models to make predictions. It uses a clever two-step method (fast sprinting followed by careful hiking) to not only predict the future but also to tell you how confident it is in that prediction. It does this quickly by using parallel computing, making advanced "uncertainty-aware" AI accessible to anyone who works with spreadsheet-style data.

Technical Summary of "bde: A Python Package for Bayesian Deep Ensembles via MILE"

Problem Statement
Reliable uncertainty quantification (UQ) is a critical component of modern machine learning systems. While Bayesian Deep Learning (BDL) offers a principled framework for achieving this, accessible tools for sampling-based inference in Bayesian neural networks remain scarce. Existing probabilistic programming frameworks typically support Markov Chain Monte Carlo (MCMC) but require substantial manual configuration to achieve competitive performance on neural network models. Furthermore, there is a lack of user-friendly implementations that bridge high-performance MCMC research with practical data science workflows, particularly for tabular supervised learning tasks.

Methodology: The MILE Approach
The paper introduces bde, a Python package that implements Microcanonical Langevin Ensembles (MILE), a hybrid sampling technique designed for Bayesian Deep Ensembles (BDEs). The methodology follows a two-stage inference process:

Optimization Stage: The system optimizes $n_{members}$ independent instances of a flexibly configurable feed-forward neural network. This is achieved via regularized empirical risk minimization (using negative log-likelihood as the loss function) with the AdamW optimizer. The goal is to find diverse high-likelihood modes.
Sampling Stage: The process transitions to a sampling phase using Microcanonical Langevin Monte Carlo (MCMC). This phase is enhanced with a tuning phase specifically adapted for Bayesian Neural Networks.

The combination of these stages allows the ensemble to first locate diverse modes and subsequently explore local posterior structures, generating an ensemble of samples that constitute an implicit posterior approximation.

Software Design and Implementation
The bde package is engineered to leverage the speed and high-performance capabilities of JAX and blackjax, combined with the user-friendly API of scikit-learn. Key design features include:

Parallelization: Since optimization and sampling across ensemble members are independent, the package exploits JAX's parallelization and just-in-time compilation to scale efficiently across CPUs, GPUs, and TPUs.
Interface: It provides a unified interface for point predictions, credible intervals, coverage estimates, and other uncertainty metrics.
Workflow: The package automates the orchestration of optimization, sampling, parallelization, and predictive inference, requiring no specialized knowledge of MCMC dynamics, initialization strategies, or JAX internals from the user.

Key Contributions

First User-Friendly MILE Implementation: bde provides the first accessible implementation of MILE, a technique shown to deliver strong predictive accuracy and calibrated uncertainty for Bayesian neural networks.
Scikit-learn Compatibility: By adhering to the scikit-learn API, the package enables seamless integration into existing machine learning workflows.
Tabular Focus: It specifically targets tabular supervised learning tasks, including distributional regression and (multi-class) classification.
Reproducibility: The package includes automated configurations and experimental setups to ensure reproducible UQ in tabular machine learning.

Results
The authors benchmarked bde on the airfoil and bikesharing datasets, comparing it against Linear Models, Random Forests, XGBoost, CatBoost, Deep Ensembles, and TabPFN (V2).

Performance: On the airfoil dataset, bde achieved an RMSE of $0.1215 \pm 0.0042$ , outperforming most competitors including TabPFN ($0.1359$) and Deep Ensembles ($0.3650$). On the bikesharing dataset, it achieved an RMSE of $0.2261 \pm 0.0016$ .
Uncertainty Quantification: The results highlight competitive out-of-the-box performance in UQ, particularly through its native distributional regression capability. bde demonstrated strong performance in Negative Log-Likelihood (NLL) metrics, achieving $-0.9126$ (distr. regr.) on airfoil and $-0.7315$ on bikesharing.
Computational Cost: The authors acknowledge that while highly optimized BDEs incur higher computational costs than optimization-based competitors due to the iterative nature of MCMC sampling, this investment enables the construction of a flexible, non-parametric posterior approximation.

Significance and Claims
The paper positions bde as a bridge between high-performance MCMC research and practical data science. Its inclusion in the scikit-learn-contrib organization ensures adherence to rigorous software standards and API consistency, making it a trusted, community-ready tool. The authors claim that bde enables researchers and practitioners to easily apply Bayesian Deep Ensembles for tabular data without needing specialized knowledge of the underlying complex dynamics, thereby facilitating reproducible uncertainty quantification in machine learning.