Bayesian Neural Network with Variational Inference and Energy Loss

Theoretic Foundation

BNN with VI and the energy loss is a simplicifcation extend BNNs with LVs. The BNN+LV model is proposed in Depeweg, 2018 and this simplification follows by omitting the latent variables.

The likelihood is given by

\[p(Y |\theta ,X) = \prod_{i=1}^K p(y_i|\theta,z_i,x_i) = \prod_{i=1}^K \mathcal{N}(y_{i} | f_{\theta}(x_i), \Sigma).\]

The prior on the weights is obtained as for BNNs with VI and is given by,

\[p(\theta) = \prod_{l=1}^L \prod_{h=1}^{V_l }\prod_{j=1}^{V_{l-1}+1} \mathcal{N}(w_{hj, l} \vert 0, \lambda)\]

where \(w_{hj, l}\) is the h-th row and the j-th column of weight matrix \(\theta_L\) at layer index \(L\) and \(\lambda\) is the prior variance. Note that as we use partially stochastic networks, the above may contain less factors \(\mathcal{N}(w_{hj, l} \vert 0, \lambda)\) depending on how many layers are stochastic.

Then, with the assumed likelihood function and prior, a posterior over the weights \(\theta\) is obtained via Bayes’ rule:

\[p(\theta|\mathcal{D}) = \frac{p(Y|\theta,X) p(\theta)} {p(Y|X)}\]

The approximate the posterior is given by

\[q(\theta) = \left[ \prod_{l=1}^L\! \prod_{h=1}^{V_l}\! \prod_{j=1}^{V_{l\!-\!1}\!+\!1} \mathcal{N}(w_{hj,l}\vert m^w_{hj,l},v^w_{hj,l})\right].\]

Now the parameters \(m^w_{hj,l}\),\(v^w_{hj,l}\) can be obtained by minimizing a divergence between \(p(\theta| \mathcal{D})\). Here the following approximation of the \(\alpha\) divergence, as proposed in Lobato, 2016 and Depeweg, 2016, is used,

\[E_\alpha(q) = -\log Z_q - \frac{1}{\alpha} \sum_{n=1}^N \log \mathbf{E}_{\Theta,\sim\, q}\left[ \left( \frac{p(\mathbf{y}_n | \Theta, \mathbf{x}_n, \mathbf{\Sigma)}} {f(\Theta)}\right)^\alpha \right],\]

where \(Z_n\) is the normalising constant of the exponential form of the approximate the posterior and \(f(\Theta)\) is a function depending on the parameters of the distributions of prior on the weights, see Depeweg, 2016 for details. In order to make this optimization problem scalable, SGD is used with mini-batches, and the expectation over \(q\) is approximated with an average over \(K\) samples drawn from \(q\).

The posterior predictive distribution is given by,

\[p(y_{\star}\vert x_{\star}, \mathcal{D}) = \int \mathcal{N}(y_{\star} \vert f_{\theta}(x_{\star}), \Sigma) p(\theta \vert \mathcal{D}) d\theta.\]

Again, the above posterior predictive distribution is intractable in this form. So instead we use sampling from the posterior distribution of the weights. The mean prediction is then given by the mean prediction of samples and the predictive uncertainty is obtained as standard deviation of samples from the approximation to the posterior predictive distribution.

%%capture
%pip install git+https://github.com/lightning-uq-box/lightning-uq-box.git

Imports

import os

import tempfile
from functools import partial

import matplotlib.pyplot as plt
import torch
import torch.nn as nn
from lightning import Trainer
from lightning.pytorch import seed_everything
from lightning.pytorch.loggers import CSVLogger

from lightning_uq_box.datamodules import ToyHeteroscedasticDatamodule
from lightning_uq_box.models import MLP
from lightning_uq_box.uq_methods import BNN_VI_Regression
from lightning_uq_box.viz_utils import (
    plot_calibration_uq_toolbox,
    plot_predictions_regression,
    plot_toy_regression_data,
    plot_training_metrics,
)

plt.rcParams["figure.figsize"] = [14, 5]

seed_everything(0)  # seed everything for reproducibility

Seed set to 0

0

We define a temporary directory to look at some training metrics and results.

my_temp_dir = tempfile.mkdtemp()

Datamodule

To demonstrate the method, we will make use of a Toy Regression Example that is defined as a Lightning Datamodule. While this might seem like overkill for a small toy problem, we think it is more helpful how the individual pieces of the library fit together so you can train models on more complex tasks.

dm = ToyHeteroscedasticDatamodule(batch_size=50, n_points=500)

X_train, Y_train, train_loader, X_test, Y_test, test_loader, X_gtext, Y_gtext = (
    dm.X_train,
    dm.Y_train,
    dm.train_dataloader(),
    dm.X_test,
    dm.Y_test,
    dm.test_dataloader(),
    dm.X_gtext,
    dm.Y_gtext,
)

fig = plot_toy_regression_data(X_train, Y_train, X_test, Y_test)

Model

For our Toy Regression problem, we will use a simple Multi-layer Perceptron (MLP) that you can configure to your needs. For the documentation of the MLP see here.

network = MLP(n_inputs=1, n_hidden=[50, 50], n_outputs=1, activation_fn=nn.Tanh())
network

MLP(
  (model): Sequential(
    (0): Linear(in_features=1, out_features=50, bias=True)
    (1): Tanh()
    (2): Dropout(p=0.0, inplace=False)
    (3): Linear(in_features=50, out_features=50, bias=True)
    (4): Tanh()
    (5): Dropout(p=0.0, inplace=False)
    (6): Linear(in_features=50, out_features=1, bias=True)
  )
)

With an underlying neural network, we can now use our desired UQ-Method as a sort of wrapper. All UQ-Methods are implemented as LightningModule that allow us to concisely organize the code and remove as much boilerplate code as possible.

bnn_vi_model = BNN_VI_Regression(
    network,
    optimizer=partial(torch.optim.Adam, lr=1e-2),
    n_mc_samples_train=10,
    n_mc_samples_test=50,
    output_noise_scale=1.3,
    prior_mu=0.0,
    prior_sigma=1.0,
    posterior_mu_init=0.0,
    posterior_rho_init=-6.0,
    alpha=1e-03,
    bayesian_layer_type="reparameterization",
    stochastic_module_names=[-1],
)

Trainer

Now that we have a LightningDataModule and a UQ-Method as a LightningModule, we can conduct training with a Lightning Trainer. It has tons of options to make your life easier, so we encourage you to check the documentation.

logger = CSVLogger(my_temp_dir)
trainer = Trainer(
    accelerator="cpu",
    max_epochs=200,  # number of epochs we want to train
    logger=logger,  # log training metrics for later evaluation
    log_every_n_steps=1,
    enable_checkpointing=False,
    enable_progress_bar=False,
    limit_val_batches=0.0,  # no validation runs
    default_root_dir=my_temp_dir,
)

GPU available: False, used: False
TPU available: False, using: 0 TPU cores
💡 Tip: For seamless cloud logging and experiment tracking, try installing [litlogger](https://pypi.org/project/litlogger/) to enable LitLogger, which logs metrics and artifacts automatically to the Lightning Experiments platform.

Training our model is now easy:

trainer.fit(bnn_vi_model, dm)

  | Name          | Type             | Params | Mode  | FLOPs
-------------------------------------------------------------------
0 | model         | MLP              | 2.8 K  | train | 0
1 | train_metrics | MetricCollection | 0      | train | 0
2 | val_metrics   | MetricCollection | 0      | train | 0
3 | test_metrics  | MetricCollection | 0      | train | 0
4 | nll_loss      | GaussianNLLLoss  | 0      | train | 0
  | other params  | n/a              | 1      | n/a   | n/a
-------------------------------------------------------------------
2.8 K     Trainable params
0         Non-trainable params
2.8 K     Total params
0.011     Total estimated model params size (MB)
21        Modules in train mode
0         Modules in eval mode
0         Total Flops
/home/docs/checkouts/readthedocs.org/user_builds/lightning-uq-box/envs/latest/lib/python3.12/site-packages/lightning/pytorch/utilities/_pytree.py:21: `isinstance(treespec, LeafSpec)` is deprecated, use `isinstance(treespec, TreeSpec) and treespec.is_leaf()` instead.
`Trainer.fit` stopped: `max_epochs=200` reached.

Training Metrics

To get some insights into how the training went, we can use the utility function to plot the training loss and RMSE metric.

fig = plot_training_metrics(
    os.path.join(my_temp_dir, "lightning_logs"), ["train_loss", "trainRMSE"]
)

Evaluate Predictions

The constructed Data Module contains two possible test variable. X_test are IID samples from the same noise distribution as the training data, while X_gtext (“X ground truth extended”) are dense inputs from the underlying “ground truth” function without any noise that also extends the input range to either side, so we can visualize the method’s UQ tendencies when extrapolating beyond the training data range. Thus, we will use X_gtext for visualization purposes, but use X_test to compute uncertainty and calibration metrics because we want to analyse how well the method has learned the noisy data distribution.

preds = bnn_vi_model.predict_step(X_gtext.to(bnn_vi_model.device))

fig = plot_predictions_regression(
    X_train,
    Y_train,
    X_gtext,
    Y_gtext,
    preds["pred"],
    preds["pred_uct"],
    epistemic=preds["epistemic_uct"],
    show_bands=False,
    title="Bayesian NN with Alpha Divergence Loss",
)

For some additional metrics relevant to UQ, we can use the great uncertainty-toolbox that gives us some insight into the calibration of our prediction. For a discussion of why this is important, see …

preds = bnn_vi_model.predict_step(X_test.to(bnn_vi_model.device))
fig = plot_calibration_uq_toolbox(
    preds["pred"].cpu().numpy(),
    preds["pred_uct"].cpu().numpy(),
    Y_test.cpu().numpy(),
    X_test.cpu().numpy(),
)