het_vbll_loss

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

Negative ELBO of a HetVBLLLayer using the reduced Knowles-Minka bound.

Implements the heteroscedastic discriminative objective of [HWS24], 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.