torch_blue.vi.KullbackLeiblerLoss

class torch_blue.vi.KullbackLeiblerLoss(predictive_distribution: Distribution, dataset_size: int | None = None, heat: float = 1.0, track: bool = False)

Bases: torch.nn.Module

Kullback-Leibler (KL) divergence loss.

Calculates the Evidence Lower Bound (ELBO) loss which minimizes the KL-divergence between the variational distribution and the true posterior. Requires external calculation of prior and variational log probability, i.e., modules must have return_log_probs = True.

Parameters:

  • predictive_distribution (Distribution) – Assumed distribution of the outputs. Typically, Categorical for classification and MeanFieldNormal for regression.

  • dataset_size (Optional[int]) – Size of the training dataset. Required for loss calculation. If not provided, it must be provided to the forward() method.

  • heat (float, default: 1.0) – A factor multiplied with the prior matching term. Set smaller than 1. to produce a “cold posterior” loss. Set to 0. to disable the prior matching term to imitate a non-Bayesian loss.

  • track (bool, default: False) –

    Set True to track the loss components. The log is stored as a dictionary in log. Loss history is stored for three components as lists accessible via the respective keys:

    • data_fitting: data log likelihood

    • prior_matching: the Kullback-Leibler divergence of prior anc variational distribution

    • log_probs: the raw prior and variational distribution log probabilities of the sampled weights.

track(mode: bool = True) → None

Enable or disable loss tracking.

Any existing loss history is kept and continued if tracking is re-enabled.

Parameters:

mode (bool, default: True) – If True, enable loss tracking if False disable it.

forward(model_output: torch_blue.vi.utils.vi_return.VIReturn, target: torch.Tensor, dataset_size: int | None = None) → torch.Tensor

Calculate the negative ELBO loss from sampled evaluations, a target and the weight log probs.

Accepts a Tensor of N samples, the associate log probabilities and a target to calculate the loss.

Parameters:

  • model_output (VIReturn) – The model output including the log probabilities. Shape (N, *)

  • target (Tensor) – Target prediction. Shape (*)

  • dataset_size (Optional[int], default: None) – Total number of samples in the dataset. Used in place of dataset_size if provided.

Returns:

Negative ELBO loss. Shape: (1,)

Return type:

Tensor