Laplace Dense Fisher

posteriors.laplace.dense_fisher.build(log_posterior, per_sample=False, init_prec=0.0)

Builds a transform for dense empirical Fisher information Laplace approximation.

The empirical Fisher is defined here as: $$ F(θ) = \sum_i ∇_θ \log p(y_i, θ | x_i) ∇_θ \log p(y_i, θ | x_i)^T $$ where \(p(y_i, θ | x_i)\) is the joint model distribution (equivalent to the posterior up to proportionality) with parameters \(θ\), inputs \(x_i\) and labels \(y_i\).

More info on empirical Fisher matrices can be found in Martens, 2020 and their use within a Laplace approximation in Daxberger et al, 2021.

Parameters:

Name	Type	Description	Default
log_posterior	LogProbFn	Function that takes parameters and input batch and returns the log posterior value (which can be unnormalised) as well as auxiliary information, e.g. from the model call.	required
per_sample	bool	If True, then log_posterior is assumed to return a vector of log posteriors for each sample in the batch. If False, then log_posterior is assumed to return a scalar log posterior for the whole batch, in this case torch.func.vmap will be called, this is typically slower than directly writing log_posterior to be per sample.	False
init_prec	Tensor | float	Initial precision matrix. If it is a float, it is defined as an identity matrix scaled by that float.	0.0

Returns:

Type	Description
Transform	Empirical Fisher information Laplace approximation transform instance.

Source code in posteriors/laplace/dense_fisher.py

def build(
    log_posterior: LogProbFn,
    per_sample: bool = False,
    init_prec: torch.Tensor | float = 0.0,
) -> Transform:
    """Builds a transform for dense empirical Fisher information
    Laplace approximation.

    The empirical Fisher is defined here as:
    $$
    F(θ) = \\sum_i ∇_θ \\log p(y_i, θ | x_i) ∇_θ \\log p(y_i, θ | x_i)^T
    $$
    where $p(y_i, θ | x_i)$ is the joint model distribution (equivalent to the posterior
    up to proportionality) with parameters $θ$, inputs $x_i$ and labels $y_i$.

    More info on empirical Fisher matrices can be found in
    [Martens, 2020](https://jmlr.org/papers/volume21/17-678/17-678.pdf) and
    their use within a Laplace approximation in [Daxberger et al, 2021](https://arxiv.org/abs/2106.14806).

    Args:
        log_posterior: Function that takes parameters and input batch and
            returns the log posterior value (which can be unnormalised)
            as well as auxiliary information, e.g. from the model call.
        per_sample: If True, then log_posterior is assumed to return a vector of
            log posteriors for each sample in the batch. If False, then log_posterior
            is assumed to return a scalar log posterior for the whole batch, in this
            case torch.func.vmap will be called, this is typically slower than
            directly writing log_posterior to be per sample.
        init_prec: Initial precision matrix.
            If it is a float, it is defined as an identity matrix
            scaled by that float.

    Returns:
        Empirical Fisher information Laplace approximation transform instance.
    """
    init_fn = partial(init, init_prec=init_prec)
    update_fn = partial(update, log_posterior=log_posterior, per_sample=per_sample)
    return Transform(init_fn, update_fn)

posteriors.laplace.dense_fisher.DenseLaplaceState

Bases: NamedTuple

State encoding a Normal distribution over parameters, with a dense precision matrix

Attributes:

Name	Type	Description
params	TensorTree	Mean of the Normal distribution.
prec	Tensor	Precision matrix of the Normal distribution.
aux	Any	Auxiliary information from the log_posterior call.

Source code in posteriors/laplace/dense_fisher.py

class DenseLaplaceState(NamedTuple):
    """State encoding a Normal distribution over parameters,
    with a dense precision matrix

    Attributes:
        params: Mean of the Normal distribution.
        prec: Precision matrix of the Normal distribution.
        aux: Auxiliary information from the log_posterior call.
    """

    params: TensorTree
    prec: torch.Tensor
    aux: Any = None

posteriors.laplace.dense_fisher.init(params, init_prec=0.0)

Initialise Normal distribution over parameters with a dense precision matrix.

Parameters:

Name	Type	Description	Default
params	TensorTree	Mean of the Normal distribution.	required
init_prec	Tensor | float	Initial precision matrix. If it is a float, it is defined as an identity matrix scaled by that float.	0.0

Returns:

Type	Description
DenseLaplaceState	Initial DenseLaplaceState.

Source code in posteriors/laplace/dense_fisher.py

def init(
    params: TensorTree,
    init_prec: torch.Tensor | float = 0.0,
) -> DenseLaplaceState:
    """Initialise Normal distribution over parameters
    with a dense precision matrix.

    Args:
        params: Mean of the Normal distribution.
        init_prec: Initial precision matrix.
            If it is a float, it is defined as an identity matrix
            scaled by that float.

    Returns:
        Initial DenseLaplaceState.
    """

    if is_scalar(init_prec):
        num_params = tree_size(params)
        init_prec = init_prec * torch.eye(num_params, requires_grad=False)

    return DenseLaplaceState(params, init_prec)

posteriors.laplace.dense_fisher.update(state, batch, log_posterior, per_sample=False, inplace=False)

Adds empirical Fisher information matrix of covariance summed over given batch.

Parameters:

Name	Type	Description	Default
state	DenseLaplaceState	Current state.	required
batch	Any	Input data to log_posterior.	required
log_posterior	LogProbFn	Function that takes parameters and input batch and returns the log posterior value (which can be unnormalised)	required
per_sample	bool	If True, then log_posterior is assumed to return a vector of log posteriors for each sample in the batch. If False, then log_posterior is assumed to return a scalar log posterior for the whole batch, in this case torch.func.vmap will be called, this is typically slower than directly writing log_posterior to be per sample.	False
inplace	bool	If True, the state is updated in place. Otherwise, a new state is returned.	False

Returns:

Type	Description
DenseLaplaceState	Updated DenseLaplaceState.

Source code in posteriors/laplace/dense_fisher.py

def update(
    state: DenseLaplaceState,
    batch: Any,
    log_posterior: LogProbFn,
    per_sample: bool = False,
    inplace: bool = False,
) -> DenseLaplaceState:
    """Adds empirical Fisher information matrix of covariance summed over
    given batch.

    Args:
        state: Current state.
        batch: Input data to log_posterior.
        log_posterior: Function that takes parameters and input batch and
            returns the log posterior value (which can be unnormalised)
        per_sample: If True, then log_posterior is assumed to return a vector of
            log posteriors for each sample in the batch. If False, then log_posterior
            is assumed to return a scalar log posterior for the whole batch, in this
            case torch.func.vmap will be called, this is typically slower than
            directly writing log_posterior to be per sample.
        inplace: If True, the state is updated in place. Otherwise, a new
            state is returned.

    Returns:
        Updated DenseLaplaceState.
    """
    if not per_sample:
        log_posterior = per_samplify(log_posterior)

    with torch.no_grad(), CatchAuxError():
        fisher, aux = empirical_fisher(
            lambda p: log_posterior(p, batch), has_aux=True, normalize=False
        )(state.params)

    if inplace:
        state.prec.data += fisher
        return state._replace(aux=aux)
    else:
        return DenseLaplaceState(state.params, state.prec + fisher, aux)

posteriors.laplace.dense_fisher.sample(state, sample_shape=torch.Size([]))

Sample from Normal distribution over parameters.

Parameters:

Name	Type	Description	Default
state	DenseLaplaceState	State encoding mean and precision matrix.	required
sample_shape	Size	Shape of the desired samples.	Size([])

Returns:

Type	Description
TensorTree	Sample(s) from the Normal distribution.

Source code in posteriors/laplace/dense_fisher.py

def sample(
    state: DenseLaplaceState,
    sample_shape: torch.Size = torch.Size([]),
) -> TensorTree:
    """Sample from Normal distribution over parameters.

    Args:
        state: State encoding mean and precision matrix.
        sample_shape: Shape of the desired samples.

    Returns:
        Sample(s) from the Normal distribution.
    """
    samples = torch.distributions.MultivariateNormal(
        loc=torch.zeros(state.prec.shape[0], device=state.prec.device),
        precision_matrix=state.prec,
        validate_args=False,
    ).sample(sample_shape)
    samples = samples.flatten(end_dim=-2)  # ensure samples is 2D
    mean_flat, unravel_func = tree_ravel(state.params)
    samples += mean_flat
    samples = torch.vmap(unravel_func)(samples)
    samples = tree_map(lambda x: x.reshape(sample_shape + x.shape[-1:]), samples)
    return samples