het_vbll_loss

probly.losses.torch.het_vbll_loss(layer: HetVBLLLayer, features: Tensor, targets: Tensor, regularization_weight: float) → Tensor[source]

Negative heteroscedastic VBLL ELBO from [HWS24].

Implements the heteroscedastic discriminative objective of a HetVBLLLayer, combining the reduced Knowles-Minka softmax bound with the input-dependent noise KL and the weight-posterior KL.

Parameters:

  • layer – The heteroscedastic variational Bayesian last layer to fit.

  • features – Backbone features feeding the layer, shape (batch, in_features).

  • targets – Integer class labels, shape (batch,).

  • regularization_weight – Weight on the regularization terms (typically 1 / dataset_size).

Returns:

A scalar tensor with the negative ELBO to minimize.