torch_blue.vi.distributions.MeanFieldNormal

class torch_blue.vi.distributions.MeanFieldNormal(mean: float = 0.0, std: float = 1.0, eps: float = 1e-10)

Bases: torch_blue.vi.distributions.base.Distribution

Distribution assuming uncorrelated, normal distributed values.

This distribution is implemented as prior, variational distribution, and predictive distribution. Its distribution parameters are “mean” and “log_std”.

As prior it becomes equivalent to an L2-weight decay term int the Kullback-Leibler loss.

As variational distribution it is often the default assumption.

As predictive distribution makes the Kullback-Leibler loss similar to MSE loss.

Parameters:

  • mean (float, default: 0.0) – The mean of the normal distribution before potential rescaling. Ignored if used as predictive distribution.

  • std (float, default: 1.0) – The standard deviation of the normal distribution before potential rescaling. This is converted to a log std internally. Ignored if used as predictive distribution.

  • eps (float, default: 1e-10) – Epsilon for numerical stability.

property std: float

Standard deviation of the distribution.

sample(mean: torch.Tensor, log_std: torch.Tensor) → torch.Tensor

Sample from a Gaussian distribution.

Parameters:

  • mean (Tensor) – The mean for each sample as Tensor.

  • log_std (Tensor) – The log standard deviation for each sample as Tensor. Must have the same shape as mean.

Returns:

The sampled Tensor of teh same shape as mean.

Return type:

Tensor

variational_log_prob(sample: torch.Tensor, mean: torch.Tensor, log_std: torch.Tensor) → torch.Tensor

Compute the log probability of sample based on a normal distribution.

Calculates the log probability of sample based on the provided mean and log standard deviation. All Tensors must have the same shape as sample.

This calculation is affected by _globals._USE_NORM_CONSTANTS, which can be set with use_norm_constants().

Parameters:

  • sample (Tensor) – The weight configuration to calculate the log probability for.

  • mean (Tensor) – The means of the reference distribution.

  • log_std (Tensor) – The log standard deviations of the reference distribution.

Returns:

The log probability of sample based on the provided mean and log_std.

Return type:

Tensor

prior_log_prob(sample: torch.Tensor) → torch.Tensor

Compute the Gaussian log probability of a sample using the prior parameters.

All Tensors have the same shape.

This calculation is affected by _globals._USE_NORM_CONSTANTS, which can be set with use_norm_constants().

Parameters:

sample (Tensor) – A Tensor of values to calculate the log probability for.

Returns:

The log probability of the sample under the prior.

Return type:

Tensor

reset_parameters_to_prior(module: VIModule, variable: str) → None

Reset the parameters of a module to prior mean and standard deviation.

Parameters:

  • module (VIModule) – The module containing the parameters to reset.

  • variable (str) – The name of the random variable to reset as given by variational_parameters of the associated Distribution.

Return type:

None

static predictive_parameters_from_samples(samples: torch.Tensor) → Tuple[torch.Tensor, torch.Tensor]

Calculate predictive mean and standard deviation of samples.

Parameters:

samples (Tensor) – The model output as Tensor of shape (S, *), where S is the number of samples.

Returns:

  • Tensor – The predictive mean as Tensor of shape (*), i.e., the average along the sample dimension.

  • Tensor – The predictive standard deviation as Tensor of shape (*), i.e., the standard deviation along the sample dimension.

log_prob_from_parameters(reference: torch.Tensor, parameters: Tuple[torch.Tensor, torch.Tensor]) → torch.Tensor

Calculate the log probability of reference given the predictive mean and standard deviation.

This calculation is affected by _globals._USE_NORM_CONSTANTS, which can be set with use_norm_constants().

Parameters:

  • reference (Tensor) – The ground truth label as Tensor of the same shape as each Tensor in parameters.

  • parameters (Tuple[Tensor, Tensor]) – A tuple containing the predictive means and standard deviation as two Tensors as returned by predictive_parameters_from_samples().

Returns:

The log probability of the reference under the predicted normal distribution. Shape: (1,).

Return type:

Tensor