torch_blue.vi.AnalyticalKullbackLeiblerLoss

class torch_blue.vi.AnalyticalKullbackLeiblerLoss(model: VIModule, predictive_distribution: Distribution, dataset_size: int | None = None, divergence_type: KullbackLeiblerModule | None = None, heat: float = 1.0, track: bool = False)

Bases: torch.nn.Module

Analytical Kullback-Leibler loss function.

A version of the Kullback-Leibler loss function that calculates the prior matching term analytically from the prior and variational parameters. To that end it stores a reference to the model for access to the parameters. Furthermore, only specific combinations of Prior and VariationalDistribution are supported (see table below). Additionally, it can emulate a non-Bayesian loss, when provided a model with NonBayesian variational distribution and a NonBayesianPredictiveDistribution.

Supported class combinations

Prior

Variational distribution

Kullback-Leibler module

MeanFieldNormalPrior

MeanFieldNormalVarDist

NormalNormalDivergence

UniformPrior

MeanFieldNormalVarDist

UniformNormalDivergence
Parameters:

  • model (VIModule) – The model to be trained.

  • predictive_distribution (Distribution) – The kind of distribution to assume for the forecasts. This is closely related to the non-Bayesian losses, e.g. MeanFieldNormal corresponds to MSE loss, while Categorical corresponds to cross-entropy loss.

  • dataset_size (Optional[int], default: None) – Number of samples in the training dataset. If not provided here it must be provided during the forward call (which also takes precedence if both are provided).

  • divergence_type (Optional[KullbackLeiblerModule], default: None) – If None the correct KullbackLeiblerModule is selected from the integrated options, if available. Otherwise, this is used to manually specify the used KullbackLeiblerModule, typically to integrate custom modules.

  • 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) – If True the loss components are tracked for every forward pass in log. This can be enabled, disabled and re-enable with the track() method. Any stored data will remain even if disabled and re-enabled.

Variables:

log (Optional[Dict[str, List[Tensor]]]) – If tracking was never enabled this is None. Otherwise, it is a dictionary with the keys data_fitting and prior_matching. These two components are appended every forward pass while tracking is enabled. Their sum is the total loss.

Raises:

  • AssertionError – If divergence_type is None and prior and variational distribution are not consistent in all submodules, since this is not supported yet.

  • NotImplementedError – If divergence_type is None and the combination of prior and variational distribution is not supported.

  • ValueError – If divergence_type is None and the model does not contain any VIBaseModule, i.e. is non-Bayesian.

  • UnsupportedDistributionError – If predictive_distribution does not support being use as predictive distribution.

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.

Return type:

None

prior_matching() → torch.Tensor

Calculate the prior matching KL-Divergence of model.

Returns:

The prior matching KL-Divergence of model.

Return type:

Tensor

forward(model_output: torch.Tensor, target: torch.Tensor, dataset_size: int | None = None) → torch.Tensor

Calculate the negative ELBO loss from sampled evaluations and a target.

Accepts a Tensor of N samples and a target to calculate the loss.

Parameters:

  • model_output (Tensor) – The model output with model.return_log_probs = False, i.e. the sampled model prediction. 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. Must be specified if dataset_size is None.

Returns:

Negative ELBO loss. Shape: (1,)

Return type:

Tensor