disc_vbll_loss

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

Negative discriminative ELBO of a VBLLLayer using the double-Jensen bound.

Implements the discriminative classification objective of [HWS24]: the closed-form double-Jensen lower bound on the expected log-likelihood, regularized by the weight-posterior kl_divergence and a Wishart term on the learnable noise precision. Both ingredients of the bound - the logit mean and the logit variance phi^T S_k phi + sigma_k^2 - are exactly the (mean, var) returned by the layer’s forward pass.

Parameters:

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