ELBOLoss#

class torch_uncertainty.losses.ELBOLoss(model, inner_loss, kl_weight, num_samples, dist_family=None)[source]#

The (negative) Evidence Lower Bound (ELBO) loss for Bayesian Neural Networks.

Combines an inner data-fit loss (e.g., cross-entropy or a distribution NLL) with the Kullback-Leibler regulariser \(\mathrm{KL}[q_\phi(\mathbf{w}) \| p(\mathbf{w})]\), estimated by Monte Carlo over num_samples weight samples:

\[\mathcal{L}_{\text{ELBO}} = \frac{1}{S} \sum_{s=1}^{S} \mathcal{L}_\text{inner}\!\left(f_{\mathbf{w}^{(s)}}(\mathbf{x}), y\right) + \beta_\text{KL} \cdot \mathrm{KL}\!\left[ q_\phi(\mathbf{w}) \;\|\; p(\mathbf{w}) \right],\]

with \(\mathbf{w}^{(s)} \sim q_\phi\). The KL weight \(\beta_\text{KL}\) is typically set to the inverse of the number of training points (or a manually annealed schedule).

Parameters:

  • model (Module | None) – The Bayesian Neural Network to compute the loss for.

  • inner_loss (Module) – The data-fit loss to use during training.

  • kl_weight (float) – The weight \(\beta_\text{KL}\) of the KL-divergence term.

  • num_samples (int) – The number of weight samples \(S\) used to estimate the expectation.

  • dist_family (str | None) – The distribution family to use for the output of the model. None means point-wise prediction. Defaults to None.

Note

Set model to None when using ELBOLoss inside a ClassificationRoutine — it will be filled in automatically.

forward(inputs, targets)[source]#

Gather the KL divergence from the Bayesian modules and aggregate the ELBO loss for a given network.

Parameters:

  • inputs (Tensor) – The inputs of the Bayesian Neural Network

  • targets (Tensor) – The target values

Returns:

The aggregated ELBO loss

Return type:

Tensor