FastBNNs

Introduction

FastBNNs implements fast and flexible Bayesian inference of neural networks based on propagation of statistical moments [1] and the unscented transform [2] in PyTorch. FastBNNs enables one-line conversion of many PyTorch-based models to Bayesian counterparts, whereby learnable model parameters are treated as Normal random variables.

Installation

FastBNNs is primarily built around PyTorch. To install FastBNNs, first install PyTorch

pip install torch==2.8.0 torchvision==0.23.0 torchaudio==2.8.0 --index-url https://download.pytorch.org/whl/cu129

then install FastBNNs as

pip install fastbnns

Alternatively, a requirements.txt file is provided to support installation of FastBNNs. To install in a virtual environment, run

python -m venv .venv
source .venv/bin/activate  # bash
# .venv/Scripts/Activate.ps1  # Windows PowerShell
pip install -r requirements.txt

Basic usage

A neural network nn: torch.nn.Module can be converted to a Bayesian neural network using the bnn.base.BNN wrapper class:

import torch

from fastbnns.bnn.base import BNN
from fastbnns.models.mlp import MLP

hidden_features = 32
n_hidden_layers = 1
in_features = 1
out_features = 1
nn = MLP(
    in_features=in_features,
    out_features=out_features,
    n_hidden_layers=n_hidden_layers,
    hidden_features=hidden_features,
    activation=torch.nn.LeakyReLU,
)
bnn = BNN(nn=nn, convert_in_place=False)

Forward calls through bnn can be made identically to nn:

data = torch.randn((1, in_features), dtype=torch.float32)
out_nn = nn(data)
out_bnn = bnn(data)

In this usage, a single sample of each network parameter is made in bnn before computing the forward computation identically to nn. As such, multiple network samples can be made to characterize the output distribution as

n_samples = 100
out_bnn_mc = torch.stack([bnn(data) for _ in range(n_samples)])
out_bnn_mc_mean = out_bnn_mc.mean(dim=0)
out_bnn_mc_var = out_bnn_mc.var(dim=0)

Alternatively, to leverage the fast inference methods (i.e., non-sampling-based), the network input can be wrapped in the custom type bnn.types.MuVar:

out_bnn_fast = out_bnn(bnn.types.MuVar(data))
out_bnn_fast_mean = out_bnn_fast.mu
out_bnn_fast_var = out_bnn_fast.var

Bayesian treatment

The wrapped model bnn: torch.nn.Module is still an instance of torch.nn.Module and can be trained using standard PyTorch or PyTorch Lightning strategies. However, a Bayesian treatment of bnn training requires use of a custom loss function, such as the evidence lower bound (ELBO) used in Bayes-by-backprop [3]. Examples of training the Bayesian MLP from Basic Usage using the ELBO loss are provided in PyTorch and PyTorch Lightning.

References

[1] David J. Schodt, Ryan Brown, Michael Merritt, Samuel Park, Delsin Menolascino, and Mark A. Peot. A framework for variational inference of lightweight bayesian neural networks with heteroscedastic uncertainties. 2024. arXiv:2402.14532 [cs].

[2] David J. Schodt. Few-sample Variational Inference of Bayesian Neural Networks with Arbitrary Nonlinearities. 2024. arXiv:2405.02063 [cs].

[3] Charles Blundell, Julien Cornebise, Koray Kavukcuoglu, and Daan Wierstra. Weight Uncertainty in Neural Networks, May 2015. arXiv:1505.05424 [cs, stat]

Copyright

LANL O4956

© 2025. Triad National Security, LLC. All rights reserved.

This program was produced under U.S. Government contract 89233218CNA000001 for Los Alamos National Laboratory (LANL), which is operated by Triad National Security, LLC for the U.S. Department of Energy/National Nuclear Security Administration. All rights in the program are reserved by Triad National Security, LLC, and the U.S. Department of Energy/National Nuclear Security Administration. The Government is granted for itself and others acting on its behalf a nonexclusive, paid-up, irrevocable worldwide license in this material to reproduce, prepare. derivative works, distribute copies to the public, perform publicly and display publicly, and to permit others to do so.