g_vbll_loss

probly.train.vbll.torch.g_vbll_loss(layer: GVBLLLayer, features: Tensor, targets: Tensor, regularization_weight: float) Tensor[source]

Negative generative ELBO (the Jensen bound) of a GVBLLLayer.

Implements the discriminative-free generative training objective of [HWS24]: the Jensen lower bound on the expected class-conditional log-likelihood, plus the class-mean KL term and a Wishart term on the shared noise precision.

Parameters:
  • layer – The generative 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.