torch_blue.vi.distributions.Distribution

class torch_blue.vi.distributions.Distribution

Base class for distributions.

Distributions for variational implemented as subclasses of this. However, not all distributions can function or make sense in all available roles. Furthermore, each role has distinct requirements, which can be implemented independently. Available roles are prior, variational distribution, and predictive distribution.

To simplify detection of correct usage this class has the three attributes is_prior, is_variational_distribution, and is_predictive_distribution, which are False by default. When implementing custom distributions the flags for the intended usages must be set to True. This also implements initialization time checking that provides feedback if any required components for the flagged uses is missing.

A prior specifies knowledge about the parameter distribution before training. In Bayesian training, weights are generally drawn towards the prior, unless they take an important role. Mathematically, this prior pull can take the same role as weight decay.

Each prior must name the distribution_parameters that define it as well as the way to calculate the log likelihood of a weight configuration in the prior_log_prob() method.

Generally, each name in distribution_parameters should also be an attribute of the class storing that parameter. This is necessary since the prior typically needs to be rescaled based on the layer width. By default, each parameter is assumed to require scaling, but in certain cases like the shape parameter of a Gamma distribution, this might not be the case. It that case the subset of scaling parameters must be specified in _scaling_parameters. Non-scaling parameters are technically not required to be a class attribute.

Furthermore, parameters might only assume positive values. These should be stored as logarithm of their true value, mapping them to the whole real line. Their parameter name should begin with the prefix “log_”, which is automatically detected and handled during rescaling.

To enable the prior_initialization functionality, the class must implement the reset_parameters() method. Which initializes the parameters for one random variable of a model, whose variational parameters are supported by the prior, to the prior values.

Parameters:

distribution_parameters (Tuple[str, ...]) – Parameters characterizing the prior and can be set during prior based initialization.
_required_parameters (Tuple[str, ...], default: ()) – Parameters besides a sample needed to calculate log_prob().
_scaling_parameters (Tuple[str, …], default: distribution_parameters) – Parameters that need to be rescaled for prior rescaling.

prior_log_prob: Callable[..., Tensor]: Function to calculate the log likelihood of a weight configuration under this prior.

A variational distribution specifies the parametrization used to fit the true weight distribution, i.e., the weight posterior.

Each prior must name the distribution_parameters that will be optimized during training it as well as a default for each parameter in _default_variational_parameters. It is important to note that the first parameter is assumed to be a form of mean or mode of the distribution that might be used as the weight in a non-Bayesian version of the network. While a default value for it must be given, it will usually be ignored in initialization in favor of initializing it similar to non-Bayesian weights.

Additionally, a sample() method must be defined that accepts one tensor for each variational parameter and returns a sample from the specified distributions. Finally, the way to calculate the log likelihood of a weight configuration in the variational_log_prob() method is required.

Variables:

distribution_parameters (Tuple[str, ...]) – The names of the variational parameters that characterize the distribution. These are fit during training.
_default_variational_parameters (Tuple[float, ...]) – Default initialization values for the variational parameters. If the parameter is “mean”, “mode” or “loc”, it is initialized analogously to non-Bayesian weights and the default is ignored.

sample : Callable[[Tensor, ...], Tensor]: This method is used to sample the weigh matrices in the forward pass. The sample method needs to be implemented by each subclass. It accepts one Tensor for each variational parameter in the order specified in variational_parameters and returns a sample from the distribution of the same shape. All input Tensors must have the same shape.

variational_log_prob : Callable[[Tensor, ...], Tensor]: This method is used to calculate the log likelihood of a weight configuration under this distribution. The log_prob method needs to be implemented by each subclass. It accepts one Tensor containing the weight configuration and one Tensor for each variational parameter in the order specified in variational_parameters and returns the log likelihood of the weight configuration. All input Tensors must have the same shape.

A predictive distribution is the assumed distribution of the model outputs. Its parameters should be derivable from sufficient samples for the same prediction. Each distribution must define which parameters are used to represent a prediction. For example, regression might use a predictive mean and standard deviation, while classification might use a probability for each class.

Furthermore, the distribution must be able to assign a probability to each possible prediction given the expected prediction. This is required for loss calculation. Typically, it is enough for subclasses to define log_prob_from_parameters() and predictive_parameters_from_samples(), which the class automatically uses to first calculate the predictive parameters from the provided samples and then the log likelihood of those samples from the parameters. In case this detour does not work log_prob_from_samples() can be overwritten. However, predictive_parameters_from_samples() should still be defined to allow extracting predictions.

Parameters:: distribution_parameters (Tuple[str, ...]) – String names of the predictive parameters. Mainly for documentation purposes.

predictive_parameters_from_samples: Callable[[Tensor], Union[Tensor, Tuple[Tensor, ...]]]: Abstract method that accepts the output of a model as Tensor of shape (S, *), where S is the number of samples. Calculates the predictive parameters implied by the samples.

Parameters:: samples (Tensor) – The model output as Tensor of shape (S, B, *), where S is the number of samples and B is the batch size.
Returns:: One Tensor for each predictive parameter in the order specified in distribution_parameters.
Return type:: Union[Tensor, Tuple[Tensor, …]]

log_prob_from_parameters: Callable[[Tensor, Union[Tensor, Tuple[Tensor, …]]], Tensor]: Abstract method that accepts a reference and the predictive parameters as calculated by predictive_parameters_from_samples() and calculates the log likelihood of the reference under the predicted distribution.

Parameters:

reference (Tensor) – The ground truth or label as Tensor of shape (B, *), where B is the batch size.
parameters (Union[Tensor, Tuple[Tensor, ...]]) – The parameters specifying the predicted distribution in the order specified in distribution_parameters.

Returns:

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

Return type:

Tensor

reset_variational_parameters(module: VIModule, variable: str, fan_in: int, kaiming_scaling: bool = True) → None

Reset the variational parameters of module.

Parameters equivalent to non-Bayesian weights (currently “mean”, “mode”, or “loc”) are reset accordingly using Kaiming uniform initialization based on fan_in (cf. torch.init._calculate_fan_in_and_fan_out()). Other parameters are initialized to the fixed values specified by class defaults, i.e., _default_variational_parameters. If kaiming_scaling is True , the defaults are scaled with scale * default. Any parameter beginning with “log” is assumed to be in log space and scaled with default + log(scale). The scale is 1 / sqrt(fan_in) for vectors and 1 / sqrt(3 * fan_in) for matrices.

Parameters:

module (VIModule) – Module to reset parameters.
variable (str) – Name of the variable to reset.
fan_in (int) – Size of the input parameter map.
kaiming_scaling (bool, default: True) – Whether th scale all parameters according to input map size.

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

Initialize the parameters of a VIModule according to the prior distribution.

To enable the prior_initialization functionality, the class must implement this method. It initializes the parameters for one random variable of a VIModule, whose variational parameters are supported by the prior, to the prior values. To that end the name of the random variable to initialize is passed to the method.

This method is called separately for each submodule and therefore does not have to consider any further submodules. It is also called separately for each random variable but should manage all variational parameters the prior can provide. By convention, a prior whose distribution parameters are a true subset of the variational parameters initializes the parameters it can handle, e.g. a MeanFieldNormalPrior can can be used to initialize parameters for any variational distribution that uses a mean and log_std.

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

match_parameters(distribution_parameters: Tuple[str, ...]) → Tuple[Dict[str, int], Dict[str, int]]

Compare distribution parameters to another set of parameters.

Typically, this is used to compare to the distribution parameters of another Distributions.

Parameters:: distribution_parameters (Tuple[str]) – Tuple of parameter names to compare.
Returns:: The first dictionary maps the names of the shared parameters to their index in variational_parameters. The second dictionary maps the names of parameters exclusive to variational_parameters to their index.
Return type:: Tuple[Dict[str, int], Dict[str, int]]

kaiming_rescale(fan_in: int, eps: float = 1e-05) → None

Rescale the prior based on layer width, for normalization.

Parameters from _scaling_parameters are scaled linearly based on the square root of the layer width, unless their name begins with “log_”, in which case they are scaled such that their exponential scales in the same way.

Parameters:

fan_in (int) – The relevant layer width.
eps (float, default: 1e-5) – Epsilon for numerical stability.

Return type:

None

property primary_parameter: str: The distribution parameter that is closest to a non-Bayesian weight.

log_prob_from_samples(reference: torch.Tensor, samples: torch.Tensor) → torch.Tensor

Calculate the log probability for reference given a set of samples.

Usually combines predictive_parameters_from_samples() and log_prob_from_parameters(), but can be redefined, if needed.

Parameters:

reference (Tensor) – Expected prediction as Tensor of shape (*)
samples (Tensor) – Model prediction as Tensor of shape (S, *), where S is the number of samples.

Returns:

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

Return type:

Tensor