Bensemble: Modular Bayesian Deep Learning & Ensembling

Bensemble is a production-ready, lightweight library for Bayesian Deep Learning and Neural Network Ensembling.

---

Key Resources

Resource	Description
📘 Documentation	Full API reference and user guides.
📝 Tech Report	In-depth technical details and theoretical background.
✍️ Blog Post	Summary of the project and motivation.
📊 Benchmarks	Comparison of methods on standard datasets.

---

Features

• PyTorch-Native: No hidden training loops. Use standard PyTorch to train your models, and use Bensemble for inference, ensembling, and analytics.
• Unified Ensembling API: Seamlessly combine explicit models (Deep Ensembles, NAS) and implicit methods (MC Dropout) via a single Ensemble interface.
• Neural Ensemble Search (NES): Algorithms to automatically search for diverse architectures using Stein Variational Gradient Descent (SVGD).
• Uncertainty Analytics: Principled decomposition of predictive uncertainty into aleatoric (data noise) and epistemic (model ignorance) components.
• Model Calibration & Metrics: Evaluate models using Expected Calibration Error (ECE), Brier Score, and NLL. Fix overconfident networks post-hoc with Temperature and Vector Scaling.

---

Installation

You can install bensemble using pip:

pip install bensemble

Or, using uv for lightning-fast installation:

uv pip install bensemble

---

Quick Start

Example 1: Ensembling, Calibration & Uncertainty

Easily ensemble standard PyTorch models, calibrate them, and decompose their uncertainty to detect Out-Of-Distribution data.

importtorchimporttorch.nnasnnfrombensemble.core.ensembleimportEnsemblefrombensemble.calibration.scalingimportTemperatureScalingfrombensemble.uncertaintyimportdecompose_classification_uncertaintyfrombensemble.metricsimportexpected_calibration_error# 1. Create a Deep Ensemble from standard trained PyTorch modelsmodels= [nn.Sequential(nn.Linear(10, 20), nn.ReLU(), nn.Linear(20, 3)) for_inrange(5)]
ensemble=Ensemble.from_models(models)
# 2. Calibrate the ensemble using a hold-out validation setval_logits, val_labels=torch.randn(100, 3), torch.randint(0, 3, (100,))
scaler=TemperatureScaling(init_temp=1.5).fit(val_logits, val_labels)
# 3. Predict on test datatest_x=torch.randn(10, 10)
# Returns shape: [5 models, 10 batch_size, 3 classes]logits=scaler(ensemble.predict_members(test_x)) probs=torch.softmax(logits, dim=-1)
# 4. Decompose Uncertainty & Evaluatetotal, aleatoric, epistemic=decompose_classification_uncertainty(probs)
ece=expected_calibration_error(probs.mean(dim=0), val_labels[:10])
print(f"Calibration Error (ECE): {ece:.4f}")
print(f"Epistemic Uncertainty (OOD awareness): {epistemic.mean().item():.4f}")

Example 2: Variational Inference

Build a Bayesian Neural Network from scratch using our custom layers with the Local Reparameterization Trick.

importtorchimporttorch.nnasnnfromtorch.utils.dataimportDataLoader, TensorDatasetfrombensemble.layersimportBayesianLinearfrombensemble.lossesimportVariationalLoss, GaussianLikelihoodfrombensemble.utilsimportget_total_kl, predict_with_uncertainty# 1. Define Model using Bayesian Layersmodel=nn.Sequential(
BayesianLinear(10, 50, prior_sigma=1.0),
nn.ReLU(),
BayesianLinear(50, 1, prior_sigma=1.0),
)
# 2. Define Objectives (Likelihood + Divergence)likelihood=GaussianLikelihood()
criterion=VariationalLoss(likelihood, alpha=1.0)
optimizer=torch.optim.Adam(list(model.parameters()) +list(likelihood.parameters()), lr=0.01)
# 3. Standard PyTorch Training Loopmodel.train()
forepochinrange(50): # Dummy loopx, y=torch.randn(10, 10), torch.randn(10, 1)
optimizer.zero_grad()
loss=criterion(model(x), y, get_total_kl(model))
loss.backward()
optimizer.step()
# 4. Predict with Uncertaintymean, std=predict_with_uncertainty(model, torch.randn(5, 10), num_samples=100)
print(f"Prediction: {mean[0].item():.2f} ± {std[0].item():.2f}")

---

Algorithms & Demos

We implement a wide range of Bayesian and Ensembling approaches. Check out the interactive demos in the notebooks/ directory:

Method	Description
Deep Ensembles	Naive yet powerful ensembling of independent networks with explicit uncertainty decomposition.
Monte Carlo Dropout	Implicit ensembling by keeping dropout active at test time.
Neural Ensemble Search (NES)	Automatically searches for diverse architectures (NES-RS/NES-RE).
NES via Bayesian Sampling	Extracts diverse subnetworks from a Supernet using Stein Variational Gradient Descent (SVGD).
Variational Inference	Approximates posterior using Gaussian distributions with the Local Reparameterization Trick.
Variational Rényi	Generalization of VI minimizing $\alpha$-divergence (VR-VI) for better robustness.
Laplace Approximation	Fits a Gaussian around the MAP estimate using Kronecker-Factored Curvature (K-FAC).
Probabilistic Backprop	Propagates moments through the network using Assumed Density Filtering (ADF).

---

Structure

bensemble/
├── core/ # Base protocols, ensemble abstractions, and adapters
│ ├── ensemble.py # Central `Ensemble` class
│ ├── member.py # Adapters for explicit and stochastic models
│ └── types.py # Protocols: Predictor, KLProvider, PosteriorSource
│
├── layers/ # Bayesian Layers for Variational Inference
│ ├── base.py # BaseBayesianLayer: KL divergence, SNR-based pruning
│ ├── linear.py # Bayesian Linear layer
│ └── conv.py # Bayesian Convolution layer
│
├── methods/ # Standalone posterior-approximation engines
│ ├── laplace_approximation.py # Laplace with K-FAC curvature
│ └── probabilistic_backpropagation.py # PBP (ADF-based)
│
├── search/ # Neural Ensemble Search algorithms
│ ├── nes.py # NES-RS & NES-RE (RandomSearcher, EvolutionarySearcher)
│ ├── bayesian.py # NESBS (discrete, SVGD-inspired sampler)
│ ├── selection.py # Greedy forward selection, scoring criteria
│ └── space.py # SearchSpace protocol
│
├── diversity/ # Methods to induce ensemble variation
│ └── dropout.py # Monte Carlo Dropout wrapper
│
├── uncertainty/ # Uncertainty analysis
│ └── decomposition.py # Separation of Aleatoric and Epistemic uncertainty
│
├── calibration/ # Post-hoc model calibration tools
│ └── scaling.py # Temperature Scaling and Vector Scaling
│
├── losses.py # VariationalLoss, GaussianLikelihood
├── utils.py # get_total_kl, predict_with_uncertainty, etc.
└── metrics.py # Scoring rules: ECE, NLL, Brier Score

---

Development Setup

If you want to contribute to bensemble or run tests, we recommend using uv.

# 1. Clone the repository
git clone https://github.com/intsystems/bensemble.git
cd bensemble
# 2. Create and activate virtual environment via uv
uv venv
source .venv/bin/activate # on Windows: .venv\Scripts\activate# 3. Install in editable mode with dev dependencies
uv pip install -e ".[dev]"

Run Tests

pytest tests/

Linting

We use ruff to keep code clean:

ruff check .
ruff format .

---

Authors

Developed by:

• Dmitrii Vasilenko
• Muhammadsharif Nabiev
• Fedor Sobolevsky
• Vadim Kasyuk

---

License

This project is licensed under the MIT License - see the LICENSE file for details.