laplace.baselaplace

BaseFunctionalLaplace

BaseFunctionalLaplace(
    model: Module,
    likelihood: Likelihood | str,
    sigma_noise: float | Tensor = 1.0,
    prior_precision: float | Tensor = 1.0,
    prior_mean: float | Tensor = 0.0,
    temperature: float = 1.0,
    enable_backprop: bool = False,
    dict_key_x: str = "input_ids",
    dict_key_y: str = "labels",
    backend: type[CurvatureInterface] | None = None,
    backend_kwargs: dict[str, Any] | None = None,
    asdl_fisher_kwargs: dict[str, Any] | None = None,
)

Bases: BaseLaplace

Provide prediction and sampling shared by function-space estimators.

Subclasses supply latent predictive moments and refresh their own posterior state. This class handles output links, sampling, and checkpoint settings for FunctionalLaplace, ELLA, and VaLLA.

Source code in laplace/baselaplace.py

def __init__(
    self,
    model: nn.Module,
    likelihood: Likelihood | str,
    sigma_noise: float | torch.Tensor = 1.0,
    prior_precision: float | torch.Tensor = 1.0,
    prior_mean: float | torch.Tensor = 0.0,
    temperature: float = 1.0,
    enable_backprop: bool = False,
    dict_key_x: str = "input_ids",
    dict_key_y: str = "labels",
    backend: type[CurvatureInterface] | None = None,
    backend_kwargs: dict[str, Any] | None = None,
    asdl_fisher_kwargs: dict[str, Any] | None = None,
) -> None:
    if likelihood not in [lik.value for lik in Likelihood]:
        raise ValueError(f"Invalid likelihood type {likelihood}")

    self.model: nn.Module = model
    self.likelihood: Likelihood | str = likelihood

    # Only do Laplace on params that require grad
    self.params: list[torch.Tensor] = []
    self.is_subset_params: bool = False
    for p in model.parameters():
        if p.requires_grad:
            self.params.append(p)
        else:
            self.is_subset_params = True

    self.n_params: int = sum(p.numel() for p in self.params)
    self.n_layers: int = len(self.params)
    self.prior_precision: float | torch.Tensor = prior_precision
    self.prior_mean: float | torch.Tensor = prior_mean
    if sigma_noise != 1 and likelihood != Likelihood.REGRESSION:
        raise ValueError("Sigma noise != 1 only available for regression.")

    self.sigma_noise: float | torch.Tensor = sigma_noise
    self.temperature: float = temperature
    self.enable_backprop: bool = enable_backprop

    # For models with dict-like inputs (e.g. Huggingface LLMs)
    self.dict_key_x = dict_key_x
    self.dict_key_y = dict_key_y

    if backend is None:
        backend = CurvlinopsGGN
    else:
        if self.is_subset_params and issubclass(
            backend, (BackPackInterface, AsdfghjklInterface)
        ):
            raise ValueError(
                "If some grad are switched off, the BackPACK and Asdfghjkl backends"
                " are not supported."
            )

    self._backend: CurvatureInterface | None = None
    self._backend_cls: type[CurvatureInterface] = backend
    self._backend_kwargs: dict[str, Any] = (
        dict() if backend_kwargs is None else backend_kwargs
    )
    self._asdl_fisher_kwargs: dict[str, Any] = (
        dict() if asdl_fisher_kwargs is None else asdl_fisher_kwargs
    )

    # log likelihood = g(loss)
    self.loss: float = 0.0
    self.n_outputs: int = 0
    self.n_data: int = 0

    # Declare attributes
    self._prior_mean: torch.Tensor
    self._prior_precision: torch.Tensor
    self._sigma_noise: torch.Tensor
    self._posterior_scale: torch.Tensor | None

log_likelihood

log_likelihood: Tensor

Compute log likelihood on the training data after .fit() has been called. The log likelihood is computed on-demand based on the loss and, for example, the observation noise which makes it differentiable in the latter for iterative updates.

Returns:

• log_likelihood ( Tensor ) –

prior_precision_diag

prior_precision_diag: Tensor

Obtain the diagonal prior precision \(p_0\) constructed from either a scalar, layer-wise, or diagonal prior precision.

Returns:

• prior_precision_diag ( Tensor ) –

_uses_backpack_backend

_uses_backpack_backend: bool

Whether the curvature backend extends modules using BackPACK.

_uses_asdl_backend

_uses_asdl_backend: bool

Whether the curvature backend computes Jacobians using ASDL.

optimize_prior_precision

optimize_prior_precision(
    pred_type: PredType | str,
    method: TuningMethod | str = TuningMethod.MARGLIK,
    n_steps: int = 100,
    lr: float = 0.1,
    init_prior_prec: float | Tensor = 1.0,
    prior_structure: (
        PriorStructure | str
    ) = PriorStructure.DIAG,
    val_loader: DataLoader | None = None,
    loss: (
        Metric | Callable[[Tensor], Tensor | float] | None
    ) = None,
    log_prior_prec_min: float = -4,
    log_prior_prec_max: float = 4,
    grid_size: int = 100,
    link_approx: LinkApprox | str = LinkApprox.PROBIT,
    n_samples: int = 100,
    verbose: bool = False,
    progress_bar: bool = False,
    log_prior_prec_base: float = 10.0,
) -> None

Optimize the prior precision post-hoc using the method specified by the user.

Parameters:

• pred_type (PredType or str in {'glm', 'nn'}) –

  type of posterior predictive, linearized GLM predictive or neural network sampling predictiv. The GLM predictive is consistent with the curvature approximations used here.

• method (TuningMethod or str in {'marglik', 'gridsearch'}, default: PredType.MARGLIK ) –

  specifies how the prior precision should be optimized.

• n_steps (int, default: 100 ) –

  the number of gradient descent steps to take.

• lr (float, default: 1e-1 ) –

  the learning rate to use for gradient descent.

• init_prior_prec (float or tensor, default: 1.0 ) –

  initial prior precision before the first optimization step.

• prior_structure (PriorStructure or str in {'scalar', 'layerwise', 'diag'}, default: PriorStructure.SCALAR ) –

  if init_prior_prec is scalar, the prior precision is optimized with this structure. otherwise, the structure of init_prior_prec is maintained.

• val_loader (DataLoader, default: None ) –

  DataLoader for the validation set; each iterate is a training batch (X, y).

• loss (callable or Metric, default: None ) –

  loss function to use for CV. If callable, the loss is computed offline (memory intensive). If torchmetrics.Metric, running loss is computed (efficient). The default depends on the likelihood: RunningNLLMetric() for classification and reward modeling, running MeanSquaredError() for regression.

• log_prior_prec_min (float, default: -4 ) –

  exponent of the lower bound of the gridsearch interval.

• log_prior_prec_max (float, default: 4 ) –

  exponent of the upper bound of the gridsearch interval.

• grid_size (int, default: 100 ) –

  number of values to consider inside the gridsearch interval.

• link_approx (LinkApprox or str in {'mc', 'probit', 'bridge'}, default: LinkApprox.PROBIT ) –

  how to approximate the classification link function for the 'glm'. For pred_type='nn', only 'mc' is possible.

• n_samples (int, default: 100 ) –

  number of samples for link_approx='mc'.

• verbose (bool, default: False ) –

  if true, the optimized prior precision will be printed (can be a large tensor if the prior has a diagonal covariance).

• progress_bar (bool, default: False ) –

  whether to show a progress bar; updated at every batch-Hessian computation. Useful for very large model and large amount of data, esp. when subset_of_weights='all'.

• log_prior_prec_base (float, default: 10.0 ) –

  base of the logarithmic gridsearch interval. The bounds are log_prior_prec_base raised to log_prior_prec_min and log_prior_prec_max.

Source code in laplace/baselaplace.py

def optimize_prior_precision(
    self,
    pred_type: PredType | str,
    method: TuningMethod | str = TuningMethod.MARGLIK,
    n_steps: int = 100,
    lr: float = 1e-1,
    init_prior_prec: float | torch.Tensor = 1.0,
    prior_structure: PriorStructure | str = PriorStructure.DIAG,
    val_loader: DataLoader | None = None,
    loss: torchmetrics.Metric
    | Callable[[torch.Tensor], torch.Tensor | float]
    | None = None,
    log_prior_prec_min: float = -4,
    log_prior_prec_max: float = 4,
    grid_size: int = 100,
    link_approx: LinkApprox | str = LinkApprox.PROBIT,
    n_samples: int = 100,
    verbose: bool = False,
    progress_bar: bool = False,
    log_prior_prec_base: float = 10.0,
) -> None:
    """Optimize the prior precision post-hoc using the `method`
    specified by the user.

    Parameters
    ----------
    pred_type : PredType or str in {'glm', 'nn'}
        type of posterior predictive, linearized GLM predictive or neural
        network sampling predictiv. The GLM predictive is consistent with the
        curvature approximations used here.
    method : TuningMethod or str in {'marglik', 'gridsearch'}, default=PredType.MARGLIK
        specifies how the prior precision should be optimized.
    n_steps : int, default=100
        the number of gradient descent steps to take.
    lr : float, default=1e-1
        the learning rate to use for gradient descent.
    init_prior_prec : float or tensor, default=1.0
        initial prior precision before the first optimization step.
    prior_structure : PriorStructure or str in {'scalar', 'layerwise', 'diag'}, default=PriorStructure.SCALAR
        if init_prior_prec is scalar, the prior precision is optimized with this structure.
        otherwise, the structure of init_prior_prec is maintained.
    val_loader : torch.data.utils.DataLoader, default=None
        DataLoader for the validation set; each iterate is a training batch (X, y).
    loss : callable or torchmetrics.Metric, default=None
        loss function to use for CV. If callable, the loss is computed offline (memory intensive).
        If torchmetrics.Metric, running loss is computed (efficient). The default
        depends on the likelihood: `RunningNLLMetric()` for classification and
        reward modeling, running `MeanSquaredError()` for regression.
    log_prior_prec_min : float, default=-4
        exponent of the lower bound of the gridsearch interval.
    log_prior_prec_max : float, default=4
        exponent of the upper bound of the gridsearch interval.
    grid_size : int, default=100
        number of values to consider inside the gridsearch interval.
    link_approx : LinkApprox or str in {'mc', 'probit', 'bridge'}, default=LinkApprox.PROBIT
        how to approximate the classification link function for the `'glm'`.
        For `pred_type='nn'`, only `'mc'` is possible.
    n_samples : int, default=100
        number of samples for `link_approx='mc'`.
    verbose : bool, default=False
        if true, the optimized prior precision will be printed
        (can be a large tensor if the prior has a diagonal covariance).
    progress_bar : bool, default=False
        whether to show a progress bar; updated at every batch-Hessian computation.
        Useful for very large model and large amount of data, esp. when `subset_of_weights='all'`.
    log_prior_prec_base : float, default=10.0
        base of the logarithmic gridsearch interval. The bounds are
        `log_prior_prec_base` raised to `log_prior_prec_min` and
        `log_prior_prec_max`.
    """
    likelihood = (
        Likelihood.CLASSIFICATION
        if self.likelihood == Likelihood.REWARD_MODELING
        else self.likelihood
    )

    if likelihood == Likelihood.CLASSIFICATION:
        warnings.warn(
            "By default `link_approx` is `probit`. Make sure to set it equals to "
            "the way you want to call `la(test_data, pred_type=..., link_approx=...)`."
        )

    if method == TuningMethod.MARGLIK:
        if val_loader is not None:
            warnings.warn(
                "`val_loader` will be ignored when `method` == 'marglik'. "
                "Do you mean to set `method = 'gridsearch'`?"
            )

        self.prior_precision = (
            init_prior_prec
            if isinstance(init_prior_prec, torch.Tensor)
            else torch.as_tensor(init_prior_prec)
        )

        if (
            len(self.prior_precision) == 1
            and prior_structure != PriorStructure.SCALAR
        ):
            self.prior_precision = fix_prior_prec_structure(
                self.prior_precision.item(),
                prior_structure,
                self.n_layers,
                self.n_params,
                self._device,
                self._dtype,
            )

        log_prior_prec = self.prior_precision.log()
        log_prior_prec.requires_grad = True
        optimizer = torch.optim.Adam([log_prior_prec], lr=lr)

        if progress_bar:
            pbar = tqdm.trange(n_steps)
            pbar.set_description("[Optimizing marginal likelihood]")
        else:
            pbar = range(n_steps)

        for _ in pbar:
            optimizer.zero_grad()
            prior_prec = log_prior_prec.exp()
            neg_log_marglik = -self.log_marginal_likelihood(
                prior_precision=prior_prec
            )
            neg_log_marglik.backward()
            optimizer.step()

        self.prior_precision = log_prior_prec.detach().exp()
    elif method == TuningMethod.GRIDSEARCH:
        if val_loader is None:
            raise ValueError("gridsearch requires a validation set DataLoader")

        interval = torch.logspace(
            log_prior_prec_min,
            log_prior_prec_max,
            grid_size,
            base=log_prior_prec_base,
        )

        if loss is None:
            loss = (
                torchmetrics.MeanSquaredError(num_outputs=self.n_outputs).to(
                    self._device
                )
                if likelihood == Likelihood.REGRESSION
                else RunningNLLMetric().to(self._device)
            )

        self.prior_precision = self._gridsearch(
            loss,
            interval,
            val_loader,
            pred_type=pred_type,
            link_approx=link_approx,
            n_samples=n_samples,
            progress_bar=progress_bar,
        )
    else:
        raise ValueError("For now only marglik and gridsearch is implemented.")

    if verbose:
        print(f"Optimized prior precision is {self.prior_precision}.")

_glm_forward_call

_glm_forward_call(
    x: Tensor | MutableMapping,
    likelihood: Likelihood | str,
    joint: bool = False,
    link_approx: LinkApprox | str = LinkApprox.PROBIT,
    n_samples: int = 100,
    diagonal_output: bool = False,
) -> Tensor | tuple[Tensor, Tensor]

Compute the posterior predictive on input data x for "glm" pred type.

Parameters:

• x (Tensor or MutableMapping) –

  (batch_size, input_shape) if tensor. If MutableMapping, must contain the said tensor.

• likelihood (Likelihood or str in {'classification', 'regression', 'reward_modeling'}) –

  determines the log likelihood Hessian approximation.

• link_approx (('mc', 'probit', 'bridge', 'bridge_norm'), default: 'mc' ) –

  how to approximate the classification link function for the 'glm'. For pred_type='nn', only 'mc' is possible.

• joint (bool, default: False ) –

  Whether to output a joint predictive distribution in regression with pred_type='glm'. If set to True, the predictive distribution has the same form as GP posterior, i.e. N([f(x1), ...,f(xm)], Cov[f(x1), ..., f(xm)]). If False, then only outputs the marginal predictive distribution. Only available for regression and GLM predictive.

• n_samples (int, default: 100 ) –

  number of samples for link_approx='mc'.

• diagonal_output (bool, default: False ) –

  whether to use a diagonalized posterior predictive on the outputs. Only works for pred_type='glm' and link_approx='mc'.

Returns:

• predictive ( Tensor or tuple[Tensor] ) –

  For likelihood='classification', a torch.Tensor is returned with a distribution over classes (similar to a Softmax). For likelihood='regression', a tuple of torch.Tensor is returned with the mean and the predictive variance. For likelihood='regression' and joint=True, a tuple of torch.Tensor is returned with the mean and the predictive covariance.

Source code in laplace/baselaplace.py

def _glm_forward_call(
    self,
    x: torch.Tensor | MutableMapping,
    likelihood: Likelihood | str,
    joint: bool = False,
    link_approx: LinkApprox | str = LinkApprox.PROBIT,
    n_samples: int = 100,
    diagonal_output: bool = False,
) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
    """Compute the posterior predictive on input data `x` for "glm" pred type.

    Parameters
    ----------
    x : torch.Tensor or MutableMapping
        `(batch_size, input_shape)` if tensor. If MutableMapping, must contain
        the said tensor.

    likelihood : Likelihood or str in {'classification', 'regression', 'reward_modeling'}
        determines the log likelihood Hessian approximation.

    link_approx : {'mc', 'probit', 'bridge', 'bridge_norm'}
        how to approximate the classification link function for the `'glm'`.
        For `pred_type='nn'`, only 'mc' is possible.

    joint : bool
        Whether to output a joint predictive distribution in regression with
        `pred_type='glm'`. If set to `True`, the predictive distribution
        has the same form as GP posterior, i.e. N([f(x1), ...,f(xm)], Cov[f(x1), ..., f(xm)]).
        If `False`, then only outputs the marginal predictive distribution.
        Only available for regression and GLM predictive.

    n_samples : int
        number of samples for `link_approx='mc'`.

    diagonal_output : bool
        whether to use a diagonalized posterior predictive on the outputs.
        Only works for `pred_type='glm'` and `link_approx='mc'`.

    Returns
    -------
    predictive: torch.Tensor or tuple[torch.Tensor]
        For `likelihood='classification'`, a torch.Tensor is returned with
        a distribution over classes (similar to a Softmax).
        For `likelihood='regression'`, a tuple of torch.Tensor is returned
        with the mean and the predictive variance.
        For `likelihood='regression'` and `joint=True`, a tuple of torch.Tensor
        is returned with the mean and the predictive covariance.
    """
    f_mu, f_var = self._glm_predictive_distribution(
        x, joint=joint and likelihood == Likelihood.REGRESSION
    )

    if likelihood == Likelihood.REGRESSION:
        if diagonal_output and not joint:
            f_var = torch.diagonal(f_var, dim1=-2, dim2=-1)
        return f_mu, f_var

    if link_approx == LinkApprox.MC:
        return self._glm_predictive_samples(
            f_mu,
            f_var,
            n_samples=n_samples,
            diagonal_output=diagonal_output,
        ).mean(dim=0)
    elif link_approx == LinkApprox.PROBIT:
        kappa = 1 / torch.sqrt(1.0 + np.pi / 8 * f_var.diagonal(dim1=1, dim2=2))
        return torch.softmax(kappa * f_mu, dim=-1)
    elif "bridge" in link_approx:
        # zero mean correction
        f_mu -= (
            f_var.sum(-1)
            * f_mu.sum(-1).reshape(-1, 1)
            / f_var.sum(dim=(1, 2)).reshape(-1, 1)
        )
        f_var -= torch.einsum(
            "bi,bj->bij", f_var.sum(-1), f_var.sum(-2)
        ) / f_var.sum(dim=(1, 2)).reshape(-1, 1, 1)

        # Laplace Bridge
        _, K = f_mu.size(0), f_mu.size(-1)
        f_var_diag = torch.diagonal(f_var, dim1=1, dim2=2)

        # optional: variance correction
        if link_approx == LinkApprox.BRIDGE_NORM:
            f_var_diag_mean = f_var_diag.mean(dim=1)
            f_var_diag_mean /= torch.as_tensor(
                [K / 2], device=self._device, dtype=self._dtype
            ).sqrt()
            f_mu /= f_var_diag_mean.sqrt().unsqueeze(-1)
            f_var_diag /= f_var_diag_mean.unsqueeze(-1)

        sum_exp = torch.exp(-f_mu).sum(dim=1).unsqueeze(-1)
        alpha = (1 - 2 / K + f_mu.exp() / K**2 * sum_exp) / f_var_diag
        return torch.nan_to_num(alpha / alpha.sum(dim=1).unsqueeze(-1), nan=1.0)
    else:
        raise ValueError(
            "Prediction path invalid. Check the likelihood, pred_type, link_approx combination!"
        )

sample

sample(
    n_samples: int = 1, generator: Generator | None = None
) -> Tensor

Sample from the Laplace posterior approximation, i.e., \( \theta \sim \mathcal{N}(\theta_{MAP}, P^{-1})\).

Parameters:

• n_samples (int, default: 100 ) –

  number of samples

• generator (Generator, default: None ) –

  random number generator to control the samples

Returns:

• samples ( Tensor ) –

Source code in laplace/baselaplace.py

def sample(
    self, n_samples: int = 1, generator: torch.Generator | None = None
) -> torch.Tensor:
    """Sample from the Laplace posterior approximation, i.e.,
    \\( \\theta \\sim \\mathcal{N}(\\theta_{MAP}, P^{-1})\\).

    Parameters
    ----------
    n_samples : int, default=100
        number of samples

    generator : torch.Generator, optional
        random number generator to control the samples

    Returns
    -------
    samples: torch.Tensor
    """
    raise NotImplementedError

_glm_predictive_samples

_glm_predictive_samples(
    f_mu: Tensor,
    f_var: Tensor,
    n_samples: int,
    diagonal_output: bool = False,
    generator: Generator | None = None,
) -> Tensor

Sample from the posterior predictive on input data x using "glm" prediction type. I.e., the inverse-link function correponding to the likelihood is applied on top of the functional sample.

Parameters:

• f_mu (Tensor or MutableMapping) –

  glm predictive mean (batch_size, output_shape)

• f_var (Tensor or MutableMapping) –

  glm predictive covariances (batch_size, output_shape, output_shape)

• n_samples (int) –

  number of samples

• diagonal_output (bool, default: False ) –

  whether to use a diagonalized glm posterior predictive on the outputs.

• generator (Generator, default: None ) –

  random number generator to control the samples (if sampling used)

Returns:

• samples ( Tensor ) –

  samples (n_samples, batch_size, output_shape)

Source code in laplace/baselaplace.py

def _glm_predictive_samples(
    self,
    f_mu: torch.Tensor,
    f_var: torch.Tensor,
    n_samples: int,
    diagonal_output: bool = False,
    generator: torch.Generator | None = None,
) -> torch.Tensor:
    """Sample from the posterior predictive on input data `x` using "glm" prediction
    type. I.e., the inverse-link function correponding to the likelihood is applied
    on top of the functional sample.

    Parameters
    ----------
    f_mu : torch.Tensor or MutableMapping
        glm predictive mean `(batch_size, output_shape)`

    f_var : torch.Tensor or MutableMapping
        glm predictive covariances `(batch_size, output_shape, output_shape)`

    n_samples : int
        number of samples

    diagonal_output : bool
        whether to use a diagonalized glm posterior predictive on the outputs.

    generator : torch.Generator, optional
        random number generator to control the samples (if sampling used)

    Returns
    -------
    samples : torch.Tensor
        samples `(n_samples, batch_size, output_shape)`
    """
    f_samples = self._glm_functional_samples(
        f_mu, f_var, n_samples, diagonal_output, generator
    )

    if self.likelihood == Likelihood.REGRESSION:
        return f_samples
    else:
        return torch.softmax(f_samples, dim=-1)

_checkpoint_configuration

_checkpoint_configuration() -> dict[str, Any]

Record input keys and backend options required by a checkpoint.

Returns:

• dict[str, Any] –

  Settings that must match when the checkpoint is restored.

Source code in laplace/baselaplace.py

def _checkpoint_configuration(self) -> dict[str, Any]:
    """Record input keys and backend options required by a checkpoint.

    Returns
    -------
    dict[str, Any]
        Settings that must match when the checkpoint is restored.
    """
    return {
        "dict_key_x": self.dict_key_x,
        "dict_key_y": self.dict_key_y,
        "backend_class": (
            f"{self._backend_cls.__module__}.{self._backend_cls.__qualname__}"
        ),
        "backend_kwargs": deepcopy(self._backend_kwargs),
        "enable_backprop": self.enable_backprop,
    }

_validate_checkpoint_configuration

_validate_checkpoint_configuration(
    state_dict: dict[str, Any]
) -> None

Reject a checkpoint created with different input or backend settings.

Parameters:

• state_dict (dict[str, Any]) –

  Saved functional posterior state.

Raises:

• ValueError –

  If a required setting differs from the current estimator.

Source code in laplace/baselaplace.py

def _validate_checkpoint_configuration(self, state_dict: dict[str, Any]) -> None:
    """Reject a checkpoint created with different input or backend settings.

    Parameters
    ----------
    state_dict : dict[str, Any]
        Saved functional posterior state.

    Raises
    ------
    ValueError
        If a required setting differs from the current estimator.
    """
    for key, current in self._checkpoint_configuration().items():
        if not same_configuration(state_dict.get(key), current):
            raise ValueError(f"Checkpoint requires matching {key}.")

_glm_predictive_distribution

_glm_predictive_distribution(
    x: Tensor | MutableMapping, joint: bool = False
) -> tuple[Tensor, Tensor]

Provide latent moments from a concrete function-space estimator.

Parameters:

• x (Tensor or MutableMapping) –

  Batched model inputs.

• joint (bool, default: False ) –

  Whether to retain covariance between input rows.

Returns:

• tuple[Tensor, Tensor] –

  Mean [B, O] and covariance [B, O, O] in model-output units. With joint=True, the mean is [B*O] and covariance is [B*O, B*O]. Observation noise is excluded.

Raises:

• NotImplementedError –

  If the subclass does not implement latent predictive moments.

Source code in laplace/baselaplace.py

def _glm_predictive_distribution(
    self, x: torch.Tensor | MutableMapping, joint: bool = False
) -> tuple[torch.Tensor, torch.Tensor]:
    """Provide latent moments from a concrete function-space estimator.

    Parameters
    ----------
    x : torch.Tensor or MutableMapping
        Batched model inputs.
    joint : bool, default=False
        Whether to retain covariance between input rows.

    Returns
    -------
    tuple[torch.Tensor, torch.Tensor]
        Mean ``[B, O]`` and covariance ``[B, O, O]`` in model-output
        units. With ``joint=True``, the mean is ``[B*O]`` and covariance
        is ``[B*O, B*O]``. Observation noise is excluded.

    Raises
    ------
    NotImplementedError
        If the subclass does not implement latent predictive moments.
    """
    raise NotImplementedError

_refresh_posterior

_refresh_posterior() -> None

Update cached posterior quantities after a hyperparameter change.

Source code in laplace/baselaplace.py

def _refresh_posterior(self) -> None:
    """Update cached posterior quantities after a hyperparameter change."""

_check_fitted

_check_fitted() -> None

Refresh a fitted posterior or reject predictions before fitting.

Raises:

• RuntimeError –

  If fit has not completed.

Source code in laplace/baselaplace.py

def _check_fitted(self) -> None:
    """Refresh a fitted posterior or reject predictions before fitting.

    Raises
    ------
    RuntimeError
        If ``fit`` has not completed.
    """
    if not getattr(self, "_fitted", False):
        raise RuntimeError("Call fit before making functional Laplace predictions.")
    self._refresh_posterior()

_check_nonempty_inputs

_check_nonempty_inputs(x: Tensor | MutableMapping) -> None

Require at least one input row for a function-space prediction.

Parameters:

• x (Tensor or MutableMapping) –

  Batched model inputs.

Raises:

• ValueError –

  If the input batch is empty.

Source code in laplace/baselaplace.py

def _check_nonempty_inputs(self, x: torch.Tensor | MutableMapping) -> None:
    """Require at least one input row for a function-space prediction.

    Parameters
    ----------
    x : torch.Tensor or MutableMapping
        Batched model inputs.

    Raises
    ------
    ValueError
        If the input batch is empty.
    """
    inputs = x[self.dict_key_x] if isinstance(x, MutableMapping) else x
    if inputs.shape[0] == 0:
        raise ValueError("Function-space predictions require at least one input.")

predictive_moments

predictive_moments(
    x: Tensor | MutableMapping, *, joint: bool = False
) -> tuple[Tensor, Tensor]

Return latent predictive moments without observation noise.

Parameters:

• x (Tensor or MutableMapping) –

  Batched model inputs.

• joint (bool, default: False ) –

  Include covariance between input rows when True.

Returns:

• tuple[Tensor, Tensor] –

  Mean [B, O] and covariance [B, O, O] by default. With joint=True, the mean is flattened and covariance has shape [B*O, B*O].

Raises:

• RuntimeError –

  If the posterior has not been fitted.

• ValueError –

  If the input batch is empty.

Source code in laplace/baselaplace.py

def predictive_moments(
    self, x: torch.Tensor | MutableMapping, *, joint: bool = False
) -> tuple[torch.Tensor, torch.Tensor]:
    """Return latent predictive moments without observation noise.

    Parameters
    ----------
    x : torch.Tensor or MutableMapping
        Batched model inputs.
    joint : bool, default=False
        Include covariance between input rows when ``True``.

    Returns
    -------
    tuple[torch.Tensor, torch.Tensor]
        Mean ``[B, O]`` and covariance ``[B, O, O]`` by default. With
        ``joint=True``, the mean is flattened and covariance has shape
        ``[B*O, B*O]``.

    Raises
    ------
    RuntimeError
        If the posterior has not been fitted.
    ValueError
        If the input batch is empty.
    """
    self._check_fitted()
    self._check_nonempty_inputs(x)
    return self._glm_predictive_distribution(x, joint=joint)

_classification_validation_nll

_classification_validation_nll(
    x: Tensor | MutableMapping, targets: Tensor
) -> Tensor

Compute summed classification NLL using the probit logit scale.

Parameters:

• x (Tensor or MutableMapping) –

  Validation inputs.

• targets (Tensor) –

  Class targets for those inputs.

Returns:

• Tensor –

  Summed negative log likelihood for the batch.

Source code in laplace/baselaplace.py

def _classification_validation_nll(
    self, x: torch.Tensor | MutableMapping, targets: torch.Tensor
) -> torch.Tensor:
    """Compute summed classification NLL using the probit logit scale.

    Parameters
    ----------
    x : torch.Tensor or MutableMapping
        Validation inputs.
    targets : torch.Tensor
        Class targets for those inputs.

    Returns
    -------
    torch.Tensor
        Summed negative log likelihood for the batch.
    """
    mean, covariance = self.predictive_moments(x)
    scale = torch.sqrt(1 + torch.pi / 8 * covariance.diagonal(dim1=-2, dim2=-1))
    return torch.nn.functional.cross_entropy(
        mean / scale, classification_targets(targets), reduction="sum"
    )

__call__

__call__(
    x: Tensor | MutableMapping,
    pred_type: PredType | str = PredType.GP,
    joint: bool = False,
    link_approx: LinkApprox | str = LinkApprox.PROBIT,
    n_samples: int = 100,
    diagonal_output: bool = False,
    generator: Generator | None = None,
    fitting: bool = False,
    **model_kwargs: dict[str, Any]
) -> Tensor | tuple[Tensor, Tensor]

Predict from the fitted function-space posterior.

Parameters:

• x (Tensor or MutableMapping) –

  Batched model inputs.

• pred_type (PredType or str, default: PredType.GP ) –

  Prediction type. Function-space estimators support only "gp".

• joint (bool, default: False ) –

  Retain cross-input covariance for regression predictions.

• link_approx (LinkApprox or str, default: LinkApprox.PROBIT ) –

  Approximation for the classification probability link.

• n_samples (int, default: 100 ) –

  Number of draws for a Monte Carlo classification link.

• diagonal_output (bool, default: False ) –

  Return only output variances for regression when supported.

• generator (Generator or None, default: None ) –

  Random generator for Monte Carlo classification predictions.

• fitting (bool, default: False ) –

  Evaluate reward-model preference pairs as classification inputs.

• **model_kwargs (Any, default: {} ) –

  Accepted for compatibility with the base prediction interface; these arguments are not forwarded to the model.

Returns:

• Tensor or tuple[Tensor, Tensor] –

  Classification probabilities [B, C], or regression/reward latent mean and covariance. With joint=True, the regression mean is flattened and covariance has shape [B*O, B*O].

Raises:

• RuntimeError –

  If the posterior has not been fitted.

• ValueError –

  If the input batch is empty, the prediction type or link is unsupported, or the generator has an incompatible type or device.

Notes

Regression and individual reward predictions exclude observation noise. Their marginal mean is [B, O] and covariance is [B, O, O], or variance [B, O] with diagonal_output=True.

Source code in laplace/baselaplace.py

def __call__(
    self,
    x: torch.Tensor | MutableMapping,
    pred_type: PredType | str = PredType.GP,
    joint: bool = False,
    link_approx: LinkApprox | str = LinkApprox.PROBIT,
    n_samples: int = 100,
    diagonal_output: bool = False,
    generator: torch.Generator | None = None,
    fitting: bool = False,
    **model_kwargs: dict[str, Any],
) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
    """Predict from the fitted function-space posterior.

    Parameters
    ----------
    x : torch.Tensor or MutableMapping
        Batched model inputs.
    pred_type : PredType or str, default=PredType.GP
        Prediction type. Function-space estimators support only ``"gp"``.
    joint : bool, default=False
        Retain cross-input covariance for regression predictions.
    link_approx : LinkApprox or str, default=LinkApprox.PROBIT
        Approximation for the classification probability link.
    n_samples : int, default=100
        Number of draws for a Monte Carlo classification link.
    diagonal_output : bool, default=False
        Return only output variances for regression when supported.
    generator : torch.Generator or None, default=None
        Random generator for Monte Carlo classification predictions.
    fitting : bool, default=False
        Evaluate reward-model preference pairs as classification inputs.
    **model_kwargs : Any
        Accepted for compatibility with the base prediction interface;
        these arguments are not forwarded to the model.

    Returns
    -------
    torch.Tensor or tuple[torch.Tensor, torch.Tensor]
        Classification probabilities ``[B, C]``, or regression/reward
        latent mean and covariance. With ``joint=True``, the regression
        mean is flattened and covariance has shape ``[B*O, B*O]``.

    Raises
    ------
    RuntimeError
        If the posterior has not been fitted.
    ValueError
        If the input batch is empty, the prediction type or link is
        unsupported, or the generator has an incompatible type or device.

    Notes
    -----
    Regression and individual reward predictions exclude observation
    noise. Their marginal mean is ``[B, O]`` and covariance is
    ``[B, O, O]``, or variance ``[B, O]`` with ``diagonal_output=True``.
    """
    self._check_fitted()
    self._check_nonempty_inputs(x)
    if pred_type != PredType.GP:
        raise ValueError("Only gp is supported as a functional prediction type.")
    if link_approx not in list(LinkApprox):
        raise ValueError(f"Unsupported link approximation {link_approx}.")
    if generator is not None and (
        not isinstance(generator, torch.Generator)
        or generator.device != self._device
    ):
        raise ValueError("Invalid random generator (check type and device).")
    likelihood = self.likelihood
    if likelihood == Likelihood.REWARD_MODELING:
        likelihood = Likelihood.CLASSIFICATION if fitting else Likelihood.REGRESSION
    if likelihood == Likelihood.CLASSIFICATION and link_approx == LinkApprox.MC:
        return self.predictive_samples(
            x,
            n_samples=n_samples,
            diagonal_output=diagonal_output,
            generator=generator,
            fitting=fitting,
        ).mean(dim=0)
    return self._glm_forward_call(
        x, likelihood, joint, link_approx, n_samples, diagonal_output
    )

_glm_functional_samples

_glm_functional_samples(
    f_mu: Tensor,
    f_var: Tensor,
    n_samples: int,
    diagonal_output: bool = False,
    generator: Generator | None = None,
) -> Tensor

Draw per-input latent samples from supplied Gaussian moments.

Parameters:

• f_mu (Tensor) –

  Latent means shaped [B, O].

• f_var (Tensor) –

  Per-input covariance shaped [B, O, O].

• n_samples (int) –

  Number of draws.

• diagonal_output (bool, default: False ) –

  Drop covariance between output coordinates before sampling.

• generator (Generator or None, default: None ) –

  Random generator for the draws.

Returns:

• Tensor –

  Latent draws shaped [n_samples, B, O].

Source code in laplace/baselaplace.py

def _glm_functional_samples(
    self,
    f_mu: torch.Tensor,
    f_var: torch.Tensor,
    n_samples: int,
    diagonal_output: bool = False,
    generator: torch.Generator | None = None,
) -> torch.Tensor:
    """Draw per-input latent samples from supplied Gaussian moments.

    Parameters
    ----------
    f_mu : torch.Tensor
        Latent means shaped ``[B, O]``.
    f_var : torch.Tensor
        Per-input covariance shaped ``[B, O, O]``.
    n_samples : int
        Number of draws.
    diagonal_output : bool, default=False
        Drop covariance between output coordinates before sampling.
    generator : torch.Generator or None, default=None
        Random generator for the draws.

    Returns
    -------
    torch.Tensor
        Latent draws shaped ``[n_samples, B, O]``.
    """
    covariance = (
        torch.diag_embed(f_var.diagonal(dim1=-2, dim2=-1))
        if diagonal_output
        else f_var
    )
    return draw_gaussian(f_mu, covariance, n_samples, generator)

functional_samples

functional_samples(
    x: Tensor | MutableMapping,
    pred_type: PredType | str = PredType.GLM,
    n_samples: int = 100,
    diagonal_output: bool = False,
    generator: Generator | None = None,
    joint: bool = False,
) -> Tensor

Draw latent outputs; use joint=True to retain cross-input covariance.

Parameters:

• x (Tensor or MutableMapping) –

  Batched model inputs.

• pred_type (PredType or str, default: PredType.GLM ) –

  Accepted for API compatibility. Every value draws from the function-space posterior; "nn" does not draw model weights.

• n_samples (int, default: 100 ) –

  Number of draws.

• diagonal_output (bool, default: False ) –

  Drop cross-output covariance before sampling.

• generator (Generator or None, default: None ) –

  Random generator on the model device.

• joint (bool, default: False ) –

  Retain covariance between input rows.

Returns:

• Tensor –

  Latent draws shaped [n_samples, B, O].

Raises:

• RuntimeError –

  If the posterior has not been fitted.

• ValueError –

  If the input batch is empty, the sampling type is unsupported, or n_samples is not positive.

Source code in laplace/baselaplace.py

def functional_samples(
    self,
    x: torch.Tensor | MutableMapping,
    pred_type: PredType | str = PredType.GLM,
    n_samples: int = 100,
    diagonal_output: bool = False,
    generator: torch.Generator | None = None,
    joint: bool = False,
) -> torch.Tensor:
    """Draw latent outputs; use `joint=True` to retain cross-input covariance.

    Parameters
    ----------
    x : torch.Tensor or MutableMapping
        Batched model inputs.
    pred_type : PredType or str, default=PredType.GLM
        Accepted for API compatibility. Every value draws from the
        function-space posterior; ``"nn"`` does not draw model weights.
    n_samples : int, default=100
        Number of draws.
    diagonal_output : bool, default=False
        Drop cross-output covariance before sampling.
    generator : torch.Generator or None, default=None
        Random generator on the model device.
    joint : bool, default=False
        Retain covariance between input rows.

    Returns
    -------
    torch.Tensor
        Latent draws shaped ``[n_samples, B, O]``.

    Raises
    ------
    RuntimeError
        If the posterior has not been fitted.
    ValueError
        If the input batch is empty, the sampling type is unsupported,
        or ``n_samples`` is not positive.
    """
    self._check_fitted()
    self._check_nonempty_inputs(x)
    if pred_type not in list(PredType):
        raise ValueError(f"Unsupported functional sampling type {pred_type}.")
    if n_samples <= 0:
        raise ValueError("n_samples must be positive.")
    mean, covariance = self._glm_predictive_distribution(x, joint=joint)
    if joint:
        batch = (
            x[self.dict_key_x].shape[0]
            if isinstance(x, MutableMapping)
            else x.shape[0]
        )
        n_outputs = mean.numel() // batch
        if diagonal_output:
            output_ids = torch.arange(mean.numel(), device=mean.device) % n_outputs
            covariance = covariance * (output_ids[:, None] == output_ids[None, :])
        draws = draw_gaussian(mean, covariance, n_samples, generator)
        return draws.reshape(n_samples, batch, n_outputs)
    return self._glm_functional_samples(
        mean, covariance, n_samples, diagonal_output, generator
    )

predictive_samples

predictive_samples(
    x: Tensor | MutableMapping,
    pred_type: PredType | str = PredType.GLM,
    n_samples: int = 100,
    diagonal_output: bool = False,
    generator: Generator | None = None,
    fitting: bool = False,
    joint: bool = False,
) -> Tensor

Draw outputs and apply the classification link when appropriate.

Parameters:

• x (Tensor or MutableMapping) –

  Batched model inputs.

• pred_type (PredType or str, default: PredType.GLM ) –

  Accepted for function-space sampling API compatibility.

• n_samples (int, default: 100 ) –

  Number of draws.

• diagonal_output (bool, default: False ) –

  Drop cross-output covariance before sampling.

• generator (Generator or None, default: None ) –

  Random generator on the model device.

• fitting (bool, default: False ) –

  Interpret reward-model inputs as preference pairs.

• joint (bool, default: False ) –

  Retain covariance between input rows.

Returns:

• Tensor –

  Predictive draws shaped [n_samples, B, O].

Raises:

• RuntimeError –

  If the posterior has not been fitted.

• ValueError –

  If the input batch is empty, the sampling type is unsupported, or n_samples is not positive.

Notes

Classification draws are class probabilities obtained by applying softmax to latent samples. Regression and individual reward draws are latent function values without observation noise.

Source code in laplace/baselaplace.py

def predictive_samples(
    self,
    x: torch.Tensor | MutableMapping,
    pred_type: PredType | str = PredType.GLM,
    n_samples: int = 100,
    diagonal_output: bool = False,
    generator: torch.Generator | None = None,
    fitting: bool = False,
    joint: bool = False,
) -> torch.Tensor:
    """Draw outputs and apply the classification link when appropriate.

    Parameters
    ----------
    x : torch.Tensor or MutableMapping
        Batched model inputs.
    pred_type : PredType or str, default=PredType.GLM
        Accepted for function-space sampling API compatibility.
    n_samples : int, default=100
        Number of draws.
    diagonal_output : bool, default=False
        Drop cross-output covariance before sampling.
    generator : torch.Generator or None, default=None
        Random generator on the model device.
    fitting : bool, default=False
        Interpret reward-model inputs as preference pairs.
    joint : bool, default=False
        Retain covariance between input rows.

    Returns
    -------
    torch.Tensor
        Predictive draws shaped ``[n_samples, B, O]``.

    Raises
    ------
    RuntimeError
        If the posterior has not been fitted.
    ValueError
        If the input batch is empty, the sampling type is unsupported,
        or ``n_samples`` is not positive.

    Notes
    -----
    Classification draws are class probabilities obtained by applying
    softmax to latent samples. Regression and individual reward draws are
    latent function values without observation noise.
    """
    samples = self.functional_samples(
        x, pred_type, n_samples, diagonal_output, generator, joint
    )
    if self.likelihood == Likelihood.CLASSIFICATION or (
        self.likelihood == Likelihood.REWARD_MODELING and fitting
    ):
        return samples.softmax(dim=-1)
    return samples

FunctionalLaplace

FunctionalLaplace(
    model: Module,
    likelihood: Likelihood | str,
    n_subset: int,
    sigma_noise: float | Tensor = 1.0,
    prior_precision: float | Tensor = 1.0,
    prior_mean: float | Tensor = 0.0,
    temperature: float = 1.0,
    enable_backprop: bool = False,
    dict_key_x="input_ids",
    dict_key_y="labels",
    backend: type[CurvatureInterface] | None = BackPackGGN,
    backend_kwargs: dict[str, Any] | None = None,
    independent_outputs: bool = False,
    seed: int = 0,
)

Bases: BaseFunctionalLaplace

Applying the GGN (Generalized Gauss-Newton) approximation for the Hessian in the Laplace approximation of the posterior turns the underlying probabilistic model from a BNN into a GLM (generalized linear model). This GLM (in the weight space) is equivalent to a GP (in the function space), see Approximate Inference Turns Deep Networks into Gaussian Processes (Khan et al., 2019)

This class implements the (approximate) GP inference through which we obtain the desired quantities (posterior predictive, marginal log-likelihood). See Improving predictions of Bayesian neural nets via local linearization (Immer et al., 2021) for more details.

Note that for likelihood='classification', we approximate \( L_{NN} \) with a diagonal matrix ( \( L_{NN} \) is a block-diagonal matrix, where blocks represent Hessians of per-data-point log-likelihood w.r.t. neural network output \( f \), See Appendix A.2.1 for exact definition). We resort to such an approximation because of the (possible) errors found in Laplace approximation for multiclass GP classification in Chapter 3.5 of R&W 2006 GP book, see the question here for more details. Alternatively, one could also resort to one-vs-one or one-vs-rest implementations for multiclass classification, however, that is not (yet) supported here.

Parameters:

• model (Module) –
• likelihood (Likelihood or str in {'classification', 'regression', 'reward_modeling'}) –

  determines the log likelihood Hessian approximation. In the case of 'reward_modeling', it fits Laplace using the classification likelihood, then does prediction as in regression likelihood. The model needs to be defined accordingly: The forward pass during training takes x.shape == (batch_size, 2, dim) with y.shape = (batch_size,). Meanwhile, during evaluation x.shape == (batch_size, dim). Reward modeling is supported by KronLaplace, DiagLaplace, ELLA, and VaLLA.

• sigma_noise (Tensor or float, default: 1 ) –

  observation noise for the regression setting; must be 1 for classification

• prior_precision (Tensor or float, default: 1 ) –

  prior precision of a Gaussian prior (= weight decay); can be scalar, per-layer, or diagonal in the most general case

• prior_mean (Tensor or float, default: 0 ) –

  prior mean of a Gaussian prior, useful for continual learning

• temperature (float, default: 1 ) –

  temperature of the likelihood; lower temperature leads to more concentrated posterior and vice versa.

• enable_backprop (bool, default: False ) –

  whether to enable backprop to the input x through the Laplace predictive. Useful for e.g. Bayesian optimization.

• dict_key_x –

  The dictionary key under which the input tensor x is stored. Only has effect when the model takes a MutableMapping as the input. Useful for Huggingface LLM models.

• dict_key_y –

  The dictionary key under which the target tensor y is stored. Only has effect when the model takes a MutableMapping as the input. Useful for Huggingface LLM models.

• backend (subclasses of `laplace.curvature.CurvatureInterface`, default: BackPackGGN ) –

  backend for access to curvature/Hessian approximations. Defaults to CurvlinopsGGN if None.

• backend_kwargs (dict, default: None ) –

  arguments passed to the backend on initialization, for example to set the number of MC samples for stochastic approximations.

• n_subset (int) –

  number of data points for Subset-of-Data (SOD) approximate GP inference.

• independent_outputs (bool, default: False ) –

  GP kernel here is product of Jacobians, which results in a \( C \times C\) matrix where \(C\) is the output dimension. If True, only a diagonal of a GP kernel is used. This is (somewhat) equivalent to assuming independent GPs across output channels.

• seed (int, default: 0 ) –

  Random seed for subset of data sampler.

Source code in laplace/baselaplace.py

def __init__(
    self,
    model: nn.Module,
    likelihood: Likelihood | str,
    n_subset: int,
    sigma_noise: float | torch.Tensor = 1.0,
    prior_precision: float | torch.Tensor = 1.0,
    prior_mean: float | torch.Tensor = 0.0,
    temperature: float = 1.0,
    enable_backprop: bool = False,
    dict_key_x="input_ids",
    dict_key_y="labels",
    backend: type[CurvatureInterface] | None = BackPackGGN,
    backend_kwargs: dict[str, Any] | None = None,
    independent_outputs: bool = False,
    seed: int = 0,
):
    assert backend in [BackPackGGN, AsdlGGN, CurvlinopsGGN]
    self._check_prior_precision(prior_precision)
    super().__init__(
        model,
        likelihood,
        sigma_noise,
        prior_precision,
        prior_mean,
        temperature,
        enable_backprop,
        dict_key_x,
        dict_key_y,
        backend,
        backend_kwargs,
    )
    self.enable_backprop = enable_backprop

    self.n_subset = n_subset
    self.independent_outputs = independent_outputs
    self.seed = seed

    self.K_MM = None
    self.Sigma_inv = None  # (K_{MM} + L_MM_inv)^{-1}
    self.train_loader = (
        None  # needed in functional variance and marginal log likelihood
    )
    self.batch_size = None
    self._prior_factor_sod = None
    self.mu = None  # mean in the scatter term of the log marginal likelihood
    self.L = None

    # Posterior mean (used in regression marginal likelihood)
    self.mean = parameters_to_vector(self.model.parameters()).detach()

    self._fitted = False
    self._recompute_Sigma = True

log_likelihood

log_likelihood: Tensor

Compute log likelihood on the training data after .fit() has been called. The log likelihood is computed on-demand based on the loss and, for example, the observation noise which makes it differentiable in the latter for iterative updates.

Returns:

• log_likelihood ( Tensor ) –

prior_precision_diag

prior_precision_diag: Tensor

Obtain the diagonal prior precision \(p_0\) constructed from either a scalar, layer-wise, or diagonal prior precision.

Returns:

• prior_precision_diag ( Tensor ) –

_uses_backpack_backend

_uses_backpack_backend: bool

Whether the curvature backend extends modules using BackPACK.

_uses_asdl_backend

_uses_asdl_backend: bool

Whether the curvature backend computes Jacobians using ASDL.

log_det_ratio

log_det_ratio: Tensor

Computes log determinant term in GP marginal likelihood

For classification we use eq. (3.44) from Chapter 3.5 from GP book R&W 2006 with (note that we always use diagonal approximation \(D\) of the Hessian of log likelihood w.r.t. \(f\)):

log determinant term := \( \log | I + D^{1/2}K D^{1/2} | \)

For regression, we use "standard" GP marginal likelihood:

log determinant term := \( \log | K + \sigma_2 I | \)

scatter

scatter: Tensor

Compute scatter term in GP log marginal likelihood.

For classification we use eq. (3.44) from Chapter 3.5 from GP book R&W 2006 with \(\hat{f} = f \):

scatter term := \( f K^{-1} f^{T} \)

For regression, we use "standard" GP marginal likelihood:

scatter term := \( (y - m)K^{-1}(y -m )^T \), where \( m \) is the mean of the GP prior, which in our case corresponds to \( m := f + J (\theta - \theta_{MAP}) \)

__call__

__call__(
    x: Tensor | MutableMapping,
    pred_type: PredType | str = PredType.GP,
    joint: bool = False,
    link_approx: LinkApprox | str = LinkApprox.PROBIT,
    n_samples: int = 100,
    diagonal_output: bool = False,
    generator: Generator | None = None,
    fitting: bool = False,
    **model_kwargs: dict[str, Any]
) -> Tensor | tuple[Tensor, Tensor]

Predict from the fitted function-space posterior.

Parameters:

• x (Tensor or MutableMapping) –

  Batched model inputs.

• pred_type (PredType or str, default: PredType.GP ) –

  Prediction type. Function-space estimators support only "gp".

• joint (bool, default: False ) –

  Retain cross-input covariance for regression predictions.

• link_approx (LinkApprox or str, default: LinkApprox.PROBIT ) –

  Approximation for the classification probability link.

• n_samples (int, default: 100 ) –

  Number of draws for a Monte Carlo classification link.

• diagonal_output (bool, default: False ) –

  Return only output variances for regression when supported.

• generator (Generator or None, default: None ) –

  Random generator for Monte Carlo classification predictions.

• fitting (bool, default: False ) –

  Evaluate reward-model preference pairs as classification inputs.

• **model_kwargs (Any, default: {} ) –

  Accepted for compatibility with the base prediction interface; these arguments are not forwarded to the model.

Returns:

• Tensor or tuple[Tensor, Tensor] –

  Classification probabilities [B, C], or regression/reward latent mean and covariance. With joint=True, the regression mean is flattened and covariance has shape [B*O, B*O].

Raises:

• RuntimeError –

  If the posterior has not been fitted.

• ValueError –

  If the input batch is empty, the prediction type or link is unsupported, or the generator has an incompatible type or device.

Notes

Regression and individual reward predictions exclude observation noise. Their marginal mean is [B, O] and covariance is [B, O, O], or variance [B, O] with diagonal_output=True.

Source code in laplace/baselaplace.py

def __call__(
    self,
    x: torch.Tensor | MutableMapping,
    pred_type: PredType | str = PredType.GP,
    joint: bool = False,
    link_approx: LinkApprox | str = LinkApprox.PROBIT,
    n_samples: int = 100,
    diagonal_output: bool = False,
    generator: torch.Generator | None = None,
    fitting: bool = False,
    **model_kwargs: dict[str, Any],
) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
    """Predict from the fitted function-space posterior.

    Parameters
    ----------
    x : torch.Tensor or MutableMapping
        Batched model inputs.
    pred_type : PredType or str, default=PredType.GP
        Prediction type. Function-space estimators support only ``"gp"``.
    joint : bool, default=False
        Retain cross-input covariance for regression predictions.
    link_approx : LinkApprox or str, default=LinkApprox.PROBIT
        Approximation for the classification probability link.
    n_samples : int, default=100
        Number of draws for a Monte Carlo classification link.
    diagonal_output : bool, default=False
        Return only output variances for regression when supported.
    generator : torch.Generator or None, default=None
        Random generator for Monte Carlo classification predictions.
    fitting : bool, default=False
        Evaluate reward-model preference pairs as classification inputs.
    **model_kwargs : Any
        Accepted for compatibility with the base prediction interface;
        these arguments are not forwarded to the model.

    Returns
    -------
    torch.Tensor or tuple[torch.Tensor, torch.Tensor]
        Classification probabilities ``[B, C]``, or regression/reward
        latent mean and covariance. With ``joint=True``, the regression
        mean is flattened and covariance has shape ``[B*O, B*O]``.

    Raises
    ------
    RuntimeError
        If the posterior has not been fitted.
    ValueError
        If the input batch is empty, the prediction type or link is
        unsupported, or the generator has an incompatible type or device.

    Notes
    -----
    Regression and individual reward predictions exclude observation
    noise. Their marginal mean is ``[B, O]`` and covariance is
    ``[B, O, O]``, or variance ``[B, O]`` with ``diagonal_output=True``.
    """
    self._check_fitted()
    self._check_nonempty_inputs(x)
    if pred_type != PredType.GP:
        raise ValueError("Only gp is supported as a functional prediction type.")
    if link_approx not in list(LinkApprox):
        raise ValueError(f"Unsupported link approximation {link_approx}.")
    if generator is not None and (
        not isinstance(generator, torch.Generator)
        or generator.device != self._device
    ):
        raise ValueError("Invalid random generator (check type and device).")
    likelihood = self.likelihood
    if likelihood == Likelihood.REWARD_MODELING:
        likelihood = Likelihood.CLASSIFICATION if fitting else Likelihood.REGRESSION
    if likelihood == Likelihood.CLASSIFICATION and link_approx == LinkApprox.MC:
        return self.predictive_samples(
            x,
            n_samples=n_samples,
            diagonal_output=diagonal_output,
            generator=generator,
            fitting=fitting,
        ).mean(dim=0)
    return self._glm_forward_call(
        x, likelihood, joint, link_approx, n_samples, diagonal_output
    )

_glm_forward_call

_glm_forward_call(
    x: Tensor | MutableMapping,
    likelihood: Likelihood | str,
    joint: bool = False,
    link_approx: LinkApprox | str = LinkApprox.PROBIT,
    n_samples: int = 100,
    diagonal_output: bool = False,
) -> Tensor | tuple[Tensor, Tensor]

Compute the posterior predictive on input data x for "glm" pred type.

Parameters:

• x (Tensor or MutableMapping) –

  (batch_size, input_shape) if tensor. If MutableMapping, must contain the said tensor.

• likelihood (Likelihood or str in {'classification', 'regression', 'reward_modeling'}) –

  determines the log likelihood Hessian approximation.

• link_approx (('mc', 'probit', 'bridge', 'bridge_norm'), default: 'mc' ) –

  how to approximate the classification link function for the 'glm'. For pred_type='nn', only 'mc' is possible.

• joint (bool, default: False ) –

  Whether to output a joint predictive distribution in regression with pred_type='glm'. If set to True, the predictive distribution has the same form as GP posterior, i.e. N([f(x1), ...,f(xm)], Cov[f(x1), ..., f(xm)]). If False, then only outputs the marginal predictive distribution. Only available for regression and GLM predictive.

• n_samples (int, default: 100 ) –

  number of samples for link_approx='mc'.

• diagonal_output (bool, default: False ) –

  whether to use a diagonalized posterior predictive on the outputs. Only works for pred_type='glm' and link_approx='mc'.

Returns:

• predictive ( Tensor or tuple[Tensor] ) –

  For likelihood='classification', a torch.Tensor is returned with a distribution over classes (similar to a Softmax). For likelihood='regression', a tuple of torch.Tensor is returned with the mean and the predictive variance. For likelihood='regression' and joint=True, a tuple of torch.Tensor is returned with the mean and the predictive covariance.

Source code in laplace/baselaplace.py

def _glm_forward_call(
    self,
    x: torch.Tensor | MutableMapping,
    likelihood: Likelihood | str,
    joint: bool = False,
    link_approx: LinkApprox | str = LinkApprox.PROBIT,
    n_samples: int = 100,
    diagonal_output: bool = False,
) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
    """Compute the posterior predictive on input data `x` for "glm" pred type.

    Parameters
    ----------
    x : torch.Tensor or MutableMapping
        `(batch_size, input_shape)` if tensor. If MutableMapping, must contain
        the said tensor.

    likelihood : Likelihood or str in {'classification', 'regression', 'reward_modeling'}
        determines the log likelihood Hessian approximation.

    link_approx : {'mc', 'probit', 'bridge', 'bridge_norm'}
        how to approximate the classification link function for the `'glm'`.
        For `pred_type='nn'`, only 'mc' is possible.

    joint : bool
        Whether to output a joint predictive distribution in regression with
        `pred_type='glm'`. If set to `True`, the predictive distribution
        has the same form as GP posterior, i.e. N([f(x1), ...,f(xm)], Cov[f(x1), ..., f(xm)]).
        If `False`, then only outputs the marginal predictive distribution.
        Only available for regression and GLM predictive.

    n_samples : int
        number of samples for `link_approx='mc'`.

    diagonal_output : bool
        whether to use a diagonalized posterior predictive on the outputs.
        Only works for `pred_type='glm'` and `link_approx='mc'`.

    Returns
    -------
    predictive: torch.Tensor or tuple[torch.Tensor]
        For `likelihood='classification'`, a torch.Tensor is returned with
        a distribution over classes (similar to a Softmax).
        For `likelihood='regression'`, a tuple of torch.Tensor is returned
        with the mean and the predictive variance.
        For `likelihood='regression'` and `joint=True`, a tuple of torch.Tensor
        is returned with the mean and the predictive covariance.
    """
    f_mu, f_var = self._glm_predictive_distribution(
        x, joint=joint and likelihood == Likelihood.REGRESSION
    )

    if likelihood == Likelihood.REGRESSION:
        if diagonal_output and not joint:
            f_var = torch.diagonal(f_var, dim1=-2, dim2=-1)
        return f_mu, f_var

    if link_approx == LinkApprox.MC:
        return self._glm_predictive_samples(
            f_mu,
            f_var,
            n_samples=n_samples,
            diagonal_output=diagonal_output,
        ).mean(dim=0)
    elif link_approx == LinkApprox.PROBIT:
        kappa = 1 / torch.sqrt(1.0 + np.pi / 8 * f_var.diagonal(dim1=1, dim2=2))
        return torch.softmax(kappa * f_mu, dim=-1)
    elif "bridge" in link_approx:
        # zero mean correction
        f_mu -= (
            f_var.sum(-1)
            * f_mu.sum(-1).reshape(-1, 1)
            / f_var.sum(dim=(1, 2)).reshape(-1, 1)
        )
        f_var -= torch.einsum(
            "bi,bj->bij", f_var.sum(-1), f_var.sum(-2)
        ) / f_var.sum(dim=(1, 2)).reshape(-1, 1, 1)

        # Laplace Bridge
        _, K = f_mu.size(0), f_mu.size(-1)
        f_var_diag = torch.diagonal(f_var, dim1=1, dim2=2)

        # optional: variance correction
        if link_approx == LinkApprox.BRIDGE_NORM:
            f_var_diag_mean = f_var_diag.mean(dim=1)
            f_var_diag_mean /= torch.as_tensor(
                [K / 2], device=self._device, dtype=self._dtype
            ).sqrt()
            f_mu /= f_var_diag_mean.sqrt().unsqueeze(-1)
            f_var_diag /= f_var_diag_mean.unsqueeze(-1)

        sum_exp = torch.exp(-f_mu).sum(dim=1).unsqueeze(-1)
        alpha = (1 - 2 / K + f_mu.exp() / K**2 * sum_exp) / f_var_diag
        return torch.nan_to_num(alpha / alpha.sum(dim=1).unsqueeze(-1), nan=1.0)
    else:
        raise ValueError(
            "Prediction path invalid. Check the likelihood, pred_type, link_approx combination!"
        )

sample

sample(
    n_samples: int = 1, generator: Generator | None = None
) -> Tensor

Sample from the Laplace posterior approximation, i.e., \( \theta \sim \mathcal{N}(\theta_{MAP}, P^{-1})\).

Parameters:

• n_samples (int, default: 100 ) –

  number of samples

• generator (Generator, default: None ) –

  random number generator to control the samples

Returns:

• samples ( Tensor ) –

Source code in laplace/baselaplace.py

def sample(
    self, n_samples: int = 1, generator: torch.Generator | None = None
) -> torch.Tensor:
    """Sample from the Laplace posterior approximation, i.e.,
    \\( \\theta \\sim \\mathcal{N}(\\theta_{MAP}, P^{-1})\\).

    Parameters
    ----------
    n_samples : int, default=100
        number of samples

    generator : torch.Generator, optional
        random number generator to control the samples

    Returns
    -------
    samples: torch.Tensor
    """
    raise NotImplementedError

functional_samples

functional_samples(
    x: Tensor | MutableMapping,
    pred_type: PredType | str = PredType.GLM,
    n_samples: int = 100,
    diagonal_output: bool = False,
    generator: Generator | None = None,
    joint: bool = False,
) -> Tensor

Draw latent outputs; use joint=True to retain cross-input covariance.

Parameters:

• x (Tensor or MutableMapping) –

  Batched model inputs.

• pred_type (PredType or str, default: PredType.GLM ) –

  Accepted for API compatibility. Every value draws from the function-space posterior; "nn" does not draw model weights.

• n_samples (int, default: 100 ) –

  Number of draws.

• diagonal_output (bool, default: False ) –

  Drop cross-output covariance before sampling.

• generator (Generator or None, default: None ) –

  Random generator on the model device.

• joint (bool, default: False ) –

  Retain covariance between input rows.

Returns:

• Tensor –

  Latent draws shaped [n_samples, B, O].

Raises:

• RuntimeError –

  If the posterior has not been fitted.

• ValueError –

  If the input batch is empty, the sampling type is unsupported, or n_samples is not positive.

Source code in laplace/baselaplace.py

def functional_samples(
    self,
    x: torch.Tensor | MutableMapping,
    pred_type: PredType | str = PredType.GLM,
    n_samples: int = 100,
    diagonal_output: bool = False,
    generator: torch.Generator | None = None,
    joint: bool = False,
) -> torch.Tensor:
    """Draw latent outputs; use `joint=True` to retain cross-input covariance.

    Parameters
    ----------
    x : torch.Tensor or MutableMapping
        Batched model inputs.
    pred_type : PredType or str, default=PredType.GLM
        Accepted for API compatibility. Every value draws from the
        function-space posterior; ``"nn"`` does not draw model weights.
    n_samples : int, default=100
        Number of draws.
    diagonal_output : bool, default=False
        Drop cross-output covariance before sampling.
    generator : torch.Generator or None, default=None
        Random generator on the model device.
    joint : bool, default=False
        Retain covariance between input rows.

    Returns
    -------
    torch.Tensor
        Latent draws shaped ``[n_samples, B, O]``.

    Raises
    ------
    RuntimeError
        If the posterior has not been fitted.
    ValueError
        If the input batch is empty, the sampling type is unsupported,
        or ``n_samples`` is not positive.
    """
    self._check_fitted()
    self._check_nonempty_inputs(x)
    if pred_type not in list(PredType):
        raise ValueError(f"Unsupported functional sampling type {pred_type}.")
    if n_samples <= 0:
        raise ValueError("n_samples must be positive.")
    mean, covariance = self._glm_predictive_distribution(x, joint=joint)
    if joint:
        batch = (
            x[self.dict_key_x].shape[0]
            if isinstance(x, MutableMapping)
            else x.shape[0]
        )
        n_outputs = mean.numel() // batch
        if diagonal_output:
            output_ids = torch.arange(mean.numel(), device=mean.device) % n_outputs
            covariance = covariance * (output_ids[:, None] == output_ids[None, :])
        draws = draw_gaussian(mean, covariance, n_samples, generator)
        return draws.reshape(n_samples, batch, n_outputs)
    return self._glm_functional_samples(
        mean, covariance, n_samples, diagonal_output, generator
    )

_glm_functional_samples

_glm_functional_samples(
    f_mu: Tensor,
    f_var: Tensor,
    n_samples: int,
    diagonal_output: bool = False,
    generator: Generator | None = None,
) -> Tensor

Draw per-input latent samples from supplied Gaussian moments.

Parameters:

• f_mu (Tensor) –

  Latent means shaped [B, O].

• f_var (Tensor) –

  Per-input covariance shaped [B, O, O].

• n_samples (int) –

  Number of draws.

• diagonal_output (bool, default: False ) –

  Drop covariance between output coordinates before sampling.

• generator (Generator or None, default: None ) –

  Random generator for the draws.

Returns:

• Tensor –

  Latent draws shaped [n_samples, B, O].

Source code in laplace/baselaplace.py

def _glm_functional_samples(
    self,
    f_mu: torch.Tensor,
    f_var: torch.Tensor,
    n_samples: int,
    diagonal_output: bool = False,
    generator: torch.Generator | None = None,
) -> torch.Tensor:
    """Draw per-input latent samples from supplied Gaussian moments.

    Parameters
    ----------
    f_mu : torch.Tensor
        Latent means shaped ``[B, O]``.
    f_var : torch.Tensor
        Per-input covariance shaped ``[B, O, O]``.
    n_samples : int
        Number of draws.
    diagonal_output : bool, default=False
        Drop covariance between output coordinates before sampling.
    generator : torch.Generator or None, default=None
        Random generator for the draws.

    Returns
    -------
    torch.Tensor
        Latent draws shaped ``[n_samples, B, O]``.
    """
    covariance = (
        torch.diag_embed(f_var.diagonal(dim1=-2, dim2=-1))
        if diagonal_output
        else f_var
    )
    return draw_gaussian(f_mu, covariance, n_samples, generator)

_glm_predictive_samples

_glm_predictive_samples(
    f_mu: Tensor,
    f_var: Tensor,
    n_samples: int,
    diagonal_output: bool = False,
    generator: Generator | None = None,
) -> Tensor

Sample from the posterior predictive on input data x using "glm" prediction type. I.e., the inverse-link function correponding to the likelihood is applied on top of the functional sample.

Parameters:

• f_mu (Tensor or MutableMapping) –

  glm predictive mean (batch_size, output_shape)

• f_var (Tensor or MutableMapping) –

  glm predictive covariances (batch_size, output_shape, output_shape)

• n_samples (int) –

  number of samples

• diagonal_output (bool, default: False ) –

  whether to use a diagonalized glm posterior predictive on the outputs.

• generator (Generator, default: None ) –

  random number generator to control the samples (if sampling used)

Returns:

• samples ( Tensor ) –

  samples (n_samples, batch_size, output_shape)

Source code in laplace/baselaplace.py

def _glm_predictive_samples(
    self,
    f_mu: torch.Tensor,
    f_var: torch.Tensor,
    n_samples: int,
    diagonal_output: bool = False,
    generator: torch.Generator | None = None,
) -> torch.Tensor:
    """Sample from the posterior predictive on input data `x` using "glm" prediction
    type. I.e., the inverse-link function correponding to the likelihood is applied
    on top of the functional sample.

    Parameters
    ----------
    f_mu : torch.Tensor or MutableMapping
        glm predictive mean `(batch_size, output_shape)`

    f_var : torch.Tensor or MutableMapping
        glm predictive covariances `(batch_size, output_shape, output_shape)`

    n_samples : int
        number of samples

    diagonal_output : bool
        whether to use a diagonalized glm posterior predictive on the outputs.

    generator : torch.Generator, optional
        random number generator to control the samples (if sampling used)

    Returns
    -------
    samples : torch.Tensor
        samples `(n_samples, batch_size, output_shape)`
    """
    f_samples = self._glm_functional_samples(
        f_mu, f_var, n_samples, diagonal_output, generator
    )

    if self.likelihood == Likelihood.REGRESSION:
        return f_samples
    else:
        return torch.softmax(f_samples, dim=-1)

_checkpoint_configuration

_checkpoint_configuration() -> dict[str, Any]

Record input keys and backend options required by a checkpoint.

Returns:

• dict[str, Any] –

  Settings that must match when the checkpoint is restored.

Source code in laplace/baselaplace.py

def _checkpoint_configuration(self) -> dict[str, Any]:
    """Record input keys and backend options required by a checkpoint.

    Returns
    -------
    dict[str, Any]
        Settings that must match when the checkpoint is restored.
    """
    return {
        "dict_key_x": self.dict_key_x,
        "dict_key_y": self.dict_key_y,
        "backend_class": (
            f"{self._backend_cls.__module__}.{self._backend_cls.__qualname__}"
        ),
        "backend_kwargs": deepcopy(self._backend_kwargs),
        "enable_backprop": self.enable_backprop,
    }

_validate_checkpoint_configuration

_validate_checkpoint_configuration(
    state_dict: dict[str, Any]
) -> None

Reject a checkpoint created with different input or backend settings.

Parameters:

• state_dict (dict[str, Any]) –

  Saved functional posterior state.

Raises:

• ValueError –

  If a required setting differs from the current estimator.

Source code in laplace/baselaplace.py

def _validate_checkpoint_configuration(self, state_dict: dict[str, Any]) -> None:
    """Reject a checkpoint created with different input or backend settings.

    Parameters
    ----------
    state_dict : dict[str, Any]
        Saved functional posterior state.

    Raises
    ------
    ValueError
        If a required setting differs from the current estimator.
    """
    for key, current in self._checkpoint_configuration().items():
        if not same_configuration(state_dict.get(key), current):
            raise ValueError(f"Checkpoint requires matching {key}.")

_check_fitted

_check_fitted() -> None

Refresh a fitted posterior or reject predictions before fitting.

Raises:

• RuntimeError –

  If fit has not completed.

Source code in laplace/baselaplace.py

def _check_fitted(self) -> None:
    """Refresh a fitted posterior or reject predictions before fitting.

    Raises
    ------
    RuntimeError
        If ``fit`` has not completed.
    """
    if not getattr(self, "_fitted", False):
        raise RuntimeError("Call fit before making functional Laplace predictions.")
    self._refresh_posterior()

_check_nonempty_inputs

_check_nonempty_inputs(x: Tensor | MutableMapping) -> None

Require at least one input row for a function-space prediction.

Parameters:

• x (Tensor or MutableMapping) –

  Batched model inputs.

Raises:

• ValueError –

  If the input batch is empty.

Source code in laplace/baselaplace.py

def _check_nonempty_inputs(self, x: torch.Tensor | MutableMapping) -> None:
    """Require at least one input row for a function-space prediction.

    Parameters
    ----------
    x : torch.Tensor or MutableMapping
        Batched model inputs.

    Raises
    ------
    ValueError
        If the input batch is empty.
    """
    inputs = x[self.dict_key_x] if isinstance(x, MutableMapping) else x
    if inputs.shape[0] == 0:
        raise ValueError("Function-space predictions require at least one input.")

predictive_moments

predictive_moments(
    x: Tensor | MutableMapping, *, joint: bool = False
) -> tuple[Tensor, Tensor]

Return latent predictive moments without observation noise.

Parameters:

• x (Tensor or MutableMapping) –

  Batched model inputs.

• joint (bool, default: False ) –

  Include covariance between input rows when True.

Returns:

• tuple[Tensor, Tensor] –

  Mean [B, O] and covariance [B, O, O] by default. With joint=True, the mean is flattened and covariance has shape [B*O, B*O].

Raises:

• RuntimeError –

  If the posterior has not been fitted.

• ValueError –

  If the input batch is empty.

Source code in laplace/baselaplace.py

def predictive_moments(
    self, x: torch.Tensor | MutableMapping, *, joint: bool = False
) -> tuple[torch.Tensor, torch.Tensor]:
    """Return latent predictive moments without observation noise.

    Parameters
    ----------
    x : torch.Tensor or MutableMapping
        Batched model inputs.
    joint : bool, default=False
        Include covariance between input rows when ``True``.

    Returns
    -------
    tuple[torch.Tensor, torch.Tensor]
        Mean ``[B, O]`` and covariance ``[B, O, O]`` by default. With
        ``joint=True``, the mean is flattened and covariance has shape
        ``[B*O, B*O]``.

    Raises
    ------
    RuntimeError
        If the posterior has not been fitted.
    ValueError
        If the input batch is empty.
    """
    self._check_fitted()
    self._check_nonempty_inputs(x)
    return self._glm_predictive_distribution(x, joint=joint)

_classification_validation_nll

_classification_validation_nll(
    x: Tensor | MutableMapping, targets: Tensor
) -> Tensor

Compute summed classification NLL using the probit logit scale.

Parameters:

• x (Tensor or MutableMapping) –

  Validation inputs.

• targets (Tensor) –

  Class targets for those inputs.

Returns:

• Tensor –

  Summed negative log likelihood for the batch.

Source code in laplace/baselaplace.py

def _classification_validation_nll(
    self, x: torch.Tensor | MutableMapping, targets: torch.Tensor
) -> torch.Tensor:
    """Compute summed classification NLL using the probit logit scale.

    Parameters
    ----------
    x : torch.Tensor or MutableMapping
        Validation inputs.
    targets : torch.Tensor
        Class targets for those inputs.

    Returns
    -------
    torch.Tensor
        Summed negative log likelihood for the batch.
    """
    mean, covariance = self.predictive_moments(x)
    scale = torch.sqrt(1 + torch.pi / 8 * covariance.diagonal(dim1=-2, dim2=-1))
    return torch.nn.functional.cross_entropy(
        mean / scale, classification_targets(targets), reduction="sum"
    )

predictive_samples

predictive_samples(
    x: Tensor | MutableMapping,
    pred_type: PredType | str = PredType.GLM,
    n_samples: int = 100,
    diagonal_output: bool = False,
    generator: Generator | None = None,
    fitting: bool = False,
    joint: bool = False,
) -> Tensor

Draw outputs and apply the classification link when appropriate.

Parameters:

• x (Tensor or MutableMapping) –

  Batched model inputs.

• pred_type (PredType or str, default: PredType.GLM ) –

  Accepted for function-space sampling API compatibility.

• n_samples (int, default: 100 ) –

  Number of draws.

• diagonal_output (bool, default: False ) –

  Drop cross-output covariance before sampling.

• generator (Generator or None, default: None ) –

  Random generator on the model device.

• fitting (bool, default: False ) –

  Interpret reward-model inputs as preference pairs.

• joint (bool, default: False ) –

  Retain covariance between input rows.

Returns:

• Tensor –

  Predictive draws shaped [n_samples, B, O].

Raises:

• RuntimeError –

  If the posterior has not been fitted.

• ValueError –

  If the input batch is empty, the sampling type is unsupported, or n_samples is not positive.

Notes

Classification draws are class probabilities obtained by applying softmax to latent samples. Regression and individual reward draws are latent function values without observation noise.

Source code in laplace/baselaplace.py

def predictive_samples(
    self,
    x: torch.Tensor | MutableMapping,
    pred_type: PredType | str = PredType.GLM,
    n_samples: int = 100,
    diagonal_output: bool = False,
    generator: torch.Generator | None = None,
    fitting: bool = False,
    joint: bool = False,
) -> torch.Tensor:
    """Draw outputs and apply the classification link when appropriate.

    Parameters
    ----------
    x : torch.Tensor or MutableMapping
        Batched model inputs.
    pred_type : PredType or str, default=PredType.GLM
        Accepted for function-space sampling API compatibility.
    n_samples : int, default=100
        Number of draws.
    diagonal_output : bool, default=False
        Drop cross-output covariance before sampling.
    generator : torch.Generator or None, default=None
        Random generator on the model device.
    fitting : bool, default=False
        Interpret reward-model inputs as preference pairs.
    joint : bool, default=False
        Retain covariance between input rows.

    Returns
    -------
    torch.Tensor
        Predictive draws shaped ``[n_samples, B, O]``.

    Raises
    ------
    RuntimeError
        If the posterior has not been fitted.
    ValueError
        If the input batch is empty, the sampling type is unsupported,
        or ``n_samples`` is not positive.

    Notes
    -----
    Classification draws are class probabilities obtained by applying
    softmax to latent samples. Regression and individual reward draws are
    latent function values without observation noise.
    """
    samples = self.functional_samples(
        x, pred_type, n_samples, diagonal_output, generator, joint
    )
    if self.likelihood == Likelihood.CLASSIFICATION or (
        self.likelihood == Likelihood.REWARD_MODELING and fitting
    ):
        return samples.softmax(dim=-1)
    return samples

_check_prior_precision

_check_prior_precision(prior_precision: float | Tensor)

Checks if the given prior precision is suitable for the GP interpretation of LLA. As such, only single value priors, i.e., isotropic priors are suitable.

Source code in laplace/baselaplace.py

@staticmethod
def _check_prior_precision(prior_precision: float | torch.Tensor):
    """Checks if the given prior precision is suitable for the GP interpretation of LLA.
    As such, only single value priors, i.e., isotropic priors are suitable.
    """
    if torch.is_tensor(prior_precision):
        if not (
            prior_precision.ndim == 0
            or (prior_precision.ndim == 1 and len(prior_precision) == 1)
        ):
            raise ValueError("Only isotropic priors supported in FunctionalLaplace")

_init_K_MM

_init_K_MM()

Allocates memory for the kernel matrix evaluated at the subset of the training data points. If the subset is of size \(M\) and the problem has \(C\) outputs, this is a list of C \((M,M\)) tensors for diagonal kernel and \((M x C, M x C)\) otherwise.

Source code in laplace/baselaplace.py

def _init_K_MM(self):
    """Allocates memory for the kernel matrix evaluated at the subset of the training
    data points. If the subset is of size \\(M\\) and the problem has \\(C\\) outputs,
    this is a list of C \\((M,M\\)) tensors for diagonal kernel and \\((M x C, M x C)\\)
    otherwise.
    """
    if self.independent_outputs:
        self.K_MM = [
            torch.empty(
                size=(self.n_subset, self.n_subset),
                device=self._device,
                dtype=self._dtype,
            )
            for _ in range(self.n_outputs)
        ]
    else:
        self.K_MM = torch.empty(
            size=(self.n_subset * self.n_outputs, self.n_subset * self.n_outputs),
            device=self._device,
            dtype=self._dtype,
        )

_init_Sigma_inv

_init_Sigma_inv()

Allocates memory for the cholesky decomposition of [ K_{MM} + \Lambda_{MM}^{-1}. ] See See Improving predictions of Bayesian neural nets via local linearization (Immer et al., 2021) Equation 15 for more information.

Source code in laplace/baselaplace.py

def _init_Sigma_inv(self):
    """Allocates memory for the cholesky decomposition of
    \\[
        K_{MM} + \\Lambda_{MM}^{-1}.
    \\]
    See See [Improving predictions of Bayesian neural nets via local linearization (Immer et al., 2021)](https://arxiv.org/abs/2008.08400)
    Equation 15 for more information.
    """
    if self.independent_outputs:
        self.Sigma_inv = [
            torch.empty(
                size=(self.n_subset, self.n_subset),
                device=self._device,
                dtype=self._dtype,
            )
            for _ in range(self.n_outputs)
        ]
    else:
        self.Sigma_inv = torch.empty(
            size=(self.n_subset * self.n_outputs, self.n_subset * self.n_outputs),
            device=self._device,
            dtype=self._dtype,
        )

_store_K_batch

_store_K_batch(K_batch: Tensor, i: int, j: int)

Given the kernel matrix between the i-th and the j-th batch, stores it in the corresponding position in self.K_MM.

Source code in laplace/baselaplace.py

def _store_K_batch(self, K_batch: torch.Tensor, i: int, j: int):
    """Given the kernel matrix between the i-th and the j-th batch, stores it in the
    corresponding position in self.K_MM.
    """
    if self.independent_outputs:
        for c in range(self.n_outputs):
            self.K_MM[c][
                i * self.batch_size : min((i + 1) * self.batch_size, self.n_subset),
                j * self.batch_size : min((j + 1) * self.batch_size, self.n_subset),
            ] = K_batch[:, :, c]
            if i != j:
                self.K_MM[c][
                    j * self.batch_size : min(
                        (j + 1) * self.batch_size, self.n_subset
                    ),
                    i * self.batch_size : min(
                        (i + 1) * self.batch_size, self.n_subset
                    ),
                ] = torch.transpose(K_batch[:, :, c], 0, 1)
    else:
        bC = self.batch_size * self.n_outputs
        MC = self.n_subset * self.n_outputs
        self.K_MM[
            i * bC : min((i + 1) * bC, MC), j * bC : min((j + 1) * bC, MC)
        ] = K_batch
        if i != j:
            self.K_MM[
                j * bC : min((j + 1) * bC, MC), i * bC : min((i + 1) * bC, MC)
            ] = torch.transpose(K_batch, 0, 1)

_build_L

_build_L(lambdas: list[Tensor])

Given a list of the Hessians of per-batch log-likelihood w.r.t. neural network output \( f \), returns the contatenation of these hessians in a suitable format for the used kernel (diagonal or not).

In this function the diagonal approximation is performed. Please refer to the introduction of the class for more details.

Parameters:

• lambdas (list of torch.Tensor of shape (C, C)) –

    Contains per-batch log-likelihood w.r.t. neural network output \( f \).
  

Returns:

• L ( list with length C of tensors with shape M or tensor (MxC) ) –

  Contains the given Hessians in a suitable format.

Source code in laplace/baselaplace.py

def _build_L(self, lambdas: list[torch.Tensor]):
    """Given a list of the Hessians of per-batch log-likelihood w.r.t. neural network output \\( f \\),
    returns the contatenation of these hessians in a suitable format for the used kernel
    (diagonal or not).

    In this function the diagonal approximation is performed. Please refer to the introduction of the
    class for more details.

    Parameters
    ----------
    lambdas : list of torch.Tensor of shape (C, C)
              Contains per-batch log-likelihood w.r.t. neural network output \\( f \\).

    Returns
    -------
    L : list with length C of tensors with shape M or tensor (MxC)
        Contains the given Hessians in a suitable format.
    """
    # Concatenate batch dimension and discard non-diagonal entries.
    L_diag = torch.diagonal(torch.cat(lambdas, dim=0), dim1=-2, dim2=-1).reshape(-1)

    if self.independent_outputs:
        return [L_diag[i :: self.n_outputs] for i in range(self.n_outputs)]
    else:
        return L_diag

_build_Sigma_inv

_build_Sigma_inv()

Computes the cholesky decomposition of [ K_{MM} + \Lambda_{MM}^{-1}. ] See See Improving predictions of Bayesian neural nets via local linearization (Immer et al., 2021) Equation 15 for more information.

<<<<<<< HEAD As the diagonal approximation is performed with \Lambda_{MM} (which is stored in self.L), ======= As the diagonal approximation is performed with \(\Lambda_{MM}\) (which is stored in self.L),

main the code is greatly simplified.

Source code in laplace/baselaplace.py

def _build_Sigma_inv(self):
    """Computes the cholesky decomposition of
            \\[
                K_{MM} + \\Lambda_{MM}^{-1}.
            \\]
            See See [Improving predictions of Bayesian neural nets via local linearization (Immer et al., 2021)](https://arxiv.org/abs/2008.08400)
            Equation 15 for more information.

    <<<<<<< HEAD
            As the diagonal approximation is performed with \\Lambda_{MM} (which is stored in self.L),
    =======
            As the diagonal approximation is performed with \\(\\Lambda_{MM}\\) (which is stored in self.L),
    >>>>>>> main
            the code is greatly simplified.
    """
    if self.independent_outputs:
        self.Sigma_inv = [
            torch.linalg.cholesky(
                self.gp_kernel_prior_variance * self.K_MM[c]
                + torch.diag(
                    torch.nan_to_num(1.0 / (self._H_factor * lambda_c), posinf=10.0)
                )
            )
            for c, lambda_c in enumerate(self.L)
        ]
    else:
        self.Sigma_inv = torch.linalg.cholesky(
            self.gp_kernel_prior_variance * self.K_MM
            + torch.diag(
                torch.nan_to_num(1 / (self._H_factor * self.L), posinf=10.0)
            )
        )
    self._recompute_Sigma = False

_get_SoD_data_loader

_get_SoD_data_loader(
    train_loader: DataLoader,
) -> DataLoader

Subset-of-Datapoints data loader

Source code in laplace/baselaplace.py

def _get_SoD_data_loader(self, train_loader: DataLoader) -> DataLoader:
    """Subset-of-Datapoints data loader"""
    return DataLoader(
        dataset=train_loader.dataset,
        batch_size=train_loader.batch_size,
        sampler=SoDSampler(
            N=len(train_loader.dataset), M=self.n_subset, seed=self.seed
        ),
        shuffle=False,
    )

fit

fit(
    train_loader: DataLoader | MutableMapping,
    progress_bar: bool = False,
)

Fit the Laplace approximation of a GP posterior.

Parameters:

• train_loader (DataLoader) –

  train_loader.dataset needs to be set to access \(N\), size of the data set train_loader.batch_size needs to be set to access \(b\) batch_size

• progress_bar (bool, default: False ) –

  whether to show a progress bar during the fitting process.

Source code in laplace/baselaplace.py

def fit(
    self, train_loader: DataLoader | MutableMapping, progress_bar: bool = False
):
    """Fit the Laplace approximation of a GP posterior.

    Parameters
    ----------
    train_loader : torch.data.utils.DataLoader
        `train_loader.dataset` needs to be set to access \\(N\\), size of the data set
        `train_loader.batch_size` needs to be set to access \\(b\\) batch_size
    progress_bar : bool
        whether to show a progress bar during the fitting process.
    """
    # Set model to evaluation mode
    self.model.eval()

    data = next(iter(train_loader))
    with torch.no_grad():
        if isinstance(data, MutableMapping):  # To support Huggingface dataset
            if "backpack" in self._backend_cls.__name__.lower():
                raise ValueError(
                    "Currently BackPACK backend is not supported "
                    + "for custom models with non-tensor inputs "
                    + "(https://github.com/pytorch/functorch/issues/159). Consider "
                    + "using AsdlGGN backend instead."
                )

            out = self.model(data)
        else:
            X = data[0]
            try:
                out = self.model(X[:1].to(self._device))
            except (TypeError, AttributeError):
                out = self.model(X.to(self._device))
    self.n_outputs = out.shape[-1]
    setattr(self.model, "output_size", self.n_outputs)
    self.batch_size = train_loader.batch_size

    if (
        self.likelihood == "regression"
        and self.n_outputs > 1
        and self.independent_outputs
    ):
        warnings.warn(
            "Using FunctionalLaplace with the diagonal approximation of a GP kernel is not recommended "
            "in the case of multivariate regression. Predictive variance will likely be overestimated."
        )

    N = len(train_loader.dataset)
    self.n_data = N

    assert self.n_subset <= N, (
        "`num_data` must be less than or equal to the original number of data points."
    )

    train_loader = self._get_SoD_data_loader(train_loader)
    self.train_loader = train_loader
    self._prior_factor_sod = self.n_subset / self.n_data

    self._init_K_MM()
    self._init_Sigma_inv()

    f, lambdas, mu = [], [], []

    if progress_bar:
        loader = enumerate(tqdm.tqdm(train_loader, desc="Fitting"))
    else:
        loader = enumerate(train_loader)

    for i, data in loader:
        if isinstance(data, MutableMapping):  # To support Huggingface dataset
            X, y = data, data[self.dict_key_y].to(self._device)
        else:
            X, y = data
            X, y = X.to(self._device), y.to(self._device)

        Js_batch, f_batch = self._jacobians(X, enable_backprop=False)

        if self.likelihood == Likelihood.REGRESSION and y.ndim != out.ndim:
            raise ValueError(
                f"The model's output has {out.ndim} dims but "
                f"the target has {y.ndim} dims."
            )

        with torch.no_grad():
            loss_batch = self.backend.factor * self.backend.lossfunc(f_batch, y)

        if self.likelihood == Likelihood.REGRESSION:
            b, C = f_batch.shape
            lambdas_batch = torch.unsqueeze(
                torch.eye(C, device=self._device, dtype=self._dtype), 0
            ).repeat(b, 1, 1)
        else:
            # second derivative of log lik is diag(p) - pp^T
            ps = torch.softmax(f_batch, dim=-1)
            lambdas_batch = torch.diag_embed(ps) - torch.einsum(
                "mk,mc->mck", ps, ps
            )

        self.loss += loss_batch
        lambdas.append(lambdas_batch)
        f.append(f_batch)
        mu.append(
            self._mean_scatter_term_batch(Js_batch, f_batch, y)
        )  # needed for marginal likelihood
        for j, (X2, _) in enumerate(train_loader):
            if j >= i:
                X2 = X2.to(self._device)
                K_batch = self._kernel_batch(Js_batch, X2)
                self._store_K_batch(K_batch, i, j)

    self.L = self._build_L(lambdas)
    self.mu = torch.cat(mu, dim=0)
    self._build_Sigma_inv()
    self._fitted = True

_refresh_posterior

_refresh_posterior() -> None

Rebuild the subset-of-data GP posterior after a prior change.

Warns:

• UserWarning –

  If prior precision changed since fitting and the cached posterior must be recomputed.

Source code in laplace/baselaplace.py

def _refresh_posterior(self) -> None:
    """Rebuild the subset-of-data GP posterior after a prior change.

    Warns
    -----
    UserWarning
        If prior precision changed since fitting and the cached posterior
        must be recomputed.
    """
    if self._recompute_Sigma:
        warnings.warn(
            "The prior precision has changed since fit; recomputing the GP posterior."
        )
        self._build_Sigma_inv()

functional_variance

functional_variance(Js_star: Tensor) -> Tensor

GP posterior variance:

\[ k_{**} - K_{*M} (K_{MM}+ L_{MM}^{-1})^{-1} K_{M*} \]

Parameters:

• Js_star (torch.Tensor of shape (N*, C, P)) –

    Jacobians of test data points
  

Returns:

• f_var ( torch.Tensor of shape (N*,C, C) ) –

  Contains the posterior variances of N* testing points.

Source code in laplace/baselaplace.py

def functional_variance(self, Js_star: torch.Tensor) -> torch.Tensor:
    """GP posterior variance:

    $$
        k_{**} - K_{*M} (K_{MM}+ L_{MM}^{-1})^{-1} K_{M*}
    $$

    Parameters
    ----------
    Js_star : torch.Tensor of shape (N*, C, P)
              Jacobians of test data points

    Returns
    -------
    f_var : torch.Tensor of shape (N*,C, C)
            Contains the posterior variances of N* testing points.
    """
    # Compute K_{**}
    K_star = self.gp_kernel_prior_variance * self._kernel_star(Js_star)

    # Compute K_{*M}
    K_M_star = []
    for X_batch, _ in self.train_loader:
        K_M_star_batch = self.gp_kernel_prior_variance * self._kernel_batch_star(
            Js_star, X_batch.to(self._device)
        )
        K_M_star.append(K_M_star_batch)
        del X_batch

    # Build_K_star_M computes K_{*M} (K_{MM}+ L_{MM}^{-1})^{-1} K_{M*}
    f_var = K_star - self._build_K_star_M(K_M_star)

    # If the considered kernel is diagonal, embed the covariances.
    # from (N*, C) -> (N*, C, C)
    if self.independent_outputs:
        f_var = torch.diag_embed(f_var)

    return f_var

functional_covariance

functional_covariance(Js_star: Tensor) -> Tensor

GP posterior covariance:

\[ k_{**} - K_{*M} (K_{MM}+ L_{MM}^{-1})^{-1} K_{M*} \]

Parameters:

• Js_star (torch.Tensor of shape (N*, C, P)) –

    Jacobians of test data points
  

Returns:

• f_var ( torch.Tensor of shape (N*xC, N*xC) ) –

  Contains the posterior covariances of N* testing points.

Source code in laplace/baselaplace.py

def functional_covariance(self, Js_star: torch.Tensor) -> torch.Tensor:
    """GP posterior covariance:

    $$
        k_{**} - K_{*M} (K_{MM}+ L_{MM}^{-1})^{-1} K_{M*}
    $$

    Parameters
    ----------
    Js_star : torch.Tensor of shape (N*, C, P)
              Jacobians of test data points

    Returns
    -------
    f_var : torch.Tensor of shape (N*xC, N*xC)
            Contains the posterior covariances of N* testing points.
    """
    # Compute K_{**}
    K_star = self.gp_kernel_prior_variance * self._kernel_star(Js_star, joint=True)

    # Compute K_{*M}
    K_M_star = []
    for X_batch, _ in self.train_loader:
        K_M_star_batch = self.gp_kernel_prior_variance * self._kernel_batch_star(
            Js_star, X_batch.to(self._device)
        )
        K_M_star.append(K_M_star_batch)
        del X_batch

    # Build_K_star_M computes K_{*M} (K_{MM}+ L_{MM}^{-1})^{-1} K_{M*}
    f_var = K_star - self._build_K_star_M(K_M_star, joint=True)

    # If the considered kernel is diagonal, embed the covariances.
    # from (N*, N*, C) -> (N*, N*, C, C)
    if self.independent_outputs:
        f_var = torch.diag_embed(f_var)

    # Reshape from (N*, N*, C, C) to (N*xC, N*xC)
    f_var = f_var.permute(0, 2, 1, 3).flatten(0, 1).flatten(1, 2)

    return f_var

_build_K_star_M

_build_K_star_M(
    K_M_star: Tensor, joint: bool = False
) -> Tensor

Computes K_{M} (K_{MM}+ L_{MM}^{-1})^{-1} K_{M} given K_{M*}.

Parameters:

• K_M_star (list of torch.Tensor) –

     Contains K_{M*}. Tensors have shape (N_test, C, C)
     or (N_test, C) for diagonal kernel.
  

• joint (boolean, default: False ) –

  Wether to compute cross covariances or not.
  

Returns:

• torch.tensor of shape (N_test, N_test, C) for joint diagonal, –
• (N_test, C) for non-joint diagonal, (N_test, N_test, C, C) for –
• joint non-diagonal and (N_test, C, C) for non-joint non-diagonal. –

Source code in laplace/baselaplace.py

def _build_K_star_M(
    self, K_M_star: torch.Tensor, joint: bool = False
) -> torch.Tensor:
    """Computes K_{*M} (K_{MM}+ L_{MM}^{-1})^{-1} K_{M*} given K_{M*}.

    Parameters
    ----------
    K_M_star : list of torch.Tensor
               Contains K_{M*}. Tensors have shape (N_test, C, C)
               or (N_test, C) for diagonal kernel.

    joint : boolean
            Wether to compute cross covariances or not.

    Returns
    -------
    torch.tensor of shape (N_test, N_test, C) for joint diagonal,
    (N_test, C) for non-joint diagonal, (N_test, N_test, C, C) for
    joint non-diagonal and (N_test, C, C) for non-joint non-diagonal.
    """
    # Shape (N_test, N, C, C) or (N_test, N, C) for diagonal
    K_M_star = torch.cat(K_M_star, dim=1)

    if self.independent_outputs:
        prods = []
        for c in range(self.n_outputs):
            # Compute K_{*M}L^{-1}
            v = torch.squeeze(
                torch.linalg.solve(
                    self.Sigma_inv[c], K_M_star[:, :, c].unsqueeze(2)
                ),
                2,
            )
            if joint:
                prod = torch.einsum("bm,am->ba", v, v)
            else:
                prod = torch.einsum("bm,bm->b", v, v)
            prods.append(prod.unsqueeze(-1))
        prods = torch.cat(prods, dim=-1)
        return prods
    else:
        # Reshape to (N_test, NxC, C) or (N_test, N, C)
        K_M_star = K_M_star.reshape(K_M_star.shape[0], -1, K_M_star.shape[-1])
        # Compute K_{*M}L^{-1}
        v = torch.linalg.solve(self.Sigma_inv, K_M_star)
        if joint:
            return torch.einsum("acm,bcn->abmn", v, v)
        else:
            return torch.einsum("bcm,bcn->bmn", v, v)

optimize_prior_precision

optimize_prior_precision(
    pred_type: PredType | str = PredType.GP,
    method: TuningMethod | str = TuningMethod.MARGLIK,
    n_steps: int = 100,
    lr: float = 0.1,
    init_prior_prec: float | Tensor = 1.0,
    prior_structure: (
        PriorStructure | str
    ) = PriorStructure.SCALAR,
    val_loader: DataLoader | None = None,
    loss: (
        Metric | Callable[[Tensor], Tensor | float] | None
    ) = None,
    log_prior_prec_min: float = -4,
    log_prior_prec_max: float = 4,
    grid_size: int = 100,
    link_approx: LinkApprox | str = LinkApprox.PROBIT,
    n_samples: int = 100,
    verbose: bool = False,
    progress_bar: bool = False,
    log_prior_prec_base: float = 10.0,
) -> None

optimize_prior_precision_base from BaseLaplace with pred_type='gp'

Source code in laplace/baselaplace.py

def optimize_prior_precision(
    self,
    pred_type: PredType | str = PredType.GP,
    method: TuningMethod | str = TuningMethod.MARGLIK,
    n_steps: int = 100,
    lr: float = 1e-1,
    init_prior_prec: float | torch.Tensor = 1.0,
    prior_structure: PriorStructure | str = PriorStructure.SCALAR,
    val_loader: DataLoader | None = None,
    loss: torchmetrics.Metric
    | Callable[[torch.Tensor], torch.Tensor | float]
    | None = None,
    log_prior_prec_min: float = -4,
    log_prior_prec_max: float = 4,
    grid_size: int = 100,
    link_approx: LinkApprox | str = LinkApprox.PROBIT,
    n_samples: int = 100,
    verbose: bool = False,
    progress_bar: bool = False,
    log_prior_prec_base: float = 10.0,
) -> None:
    """`optimize_prior_precision_base` from `BaseLaplace` with `pred_type='gp'`"""
    assert pred_type == PredType.GP  # only gp supported
    assert prior_structure == "scalar"  # only isotropic gaussian prior supported
    if method == "marglik":
        warnings.warn(
            "Use of method='marglik' in case of FunctionalLaplace is discouraged, rather use method='CV'."
        )
    super().optimize_prior_precision(
        pred_type,
        method,
        n_steps,
        lr,
        init_prior_prec,
        prior_structure,
        val_loader,
        loss,
        log_prior_prec_min,
        log_prior_prec_max,
        grid_size,
        link_approx,
        n_samples,
        verbose,
        progress_bar,
        log_prior_prec_base=log_prior_prec_base,
    )
    self._build_Sigma_inv()

_kernel_batch

_kernel_batch(jacobians: Tensor, batch: Tensor) -> Tensor

Compute K_bb, which is part of K_MM kernel matrix.

Parameters:

• jacobians (Tensor(b, C, P)) –
• batch (Tensor(b, C)) –

Returns:

• kernel ( tensor ) –

  K_bb with shape (b * C, b * C)

Source code in laplace/baselaplace.py

def _kernel_batch(
    self, jacobians: torch.Tensor, batch: torch.Tensor
) -> torch.Tensor:
    """Compute K_bb, which is part of K_MM kernel matrix.

    Parameters
    ----------
    jacobians : torch.Tensor (b, C, P)
    batch : torch.Tensor (b, C)

    Returns
    -------
    kernel : torch.tensor
        K_bb with shape (b * C, b * C)
    """
    jacobians_2, _ = self._jacobians(batch)
    P = jacobians.shape[-1]  # nr model params
    if self.independent_outputs:
        kernel = torch.empty(
            (jacobians.shape[0], jacobians_2.shape[0], self.n_outputs),
            device=jacobians.device,
            dtype=self._dtype,
        )
        for c in range(self.n_outputs):
            kernel[:, :, c] = torch.einsum(
                "bp,ep->be", jacobians[:, c, :], jacobians_2[:, c, :]
            )
    else:
        kernel = torch.einsum(
            "ap,bp->ab", jacobians.reshape(-1, P), jacobians_2.reshape(-1, P)
        )
    del jacobians_2
    return kernel

_kernel_star

_kernel_star(
    jacobians: Tensor, joint: bool = False
) -> Tensor

Compute K_star_star kernel matrix.

Parameters:

• jacobians (Tensor(b, C, P)) –

Returns:

• kernel ( tensor ) –

  K_star with shape (b, C, C)

Source code in laplace/baselaplace.py

def _kernel_star(
    self, jacobians: torch.Tensor, joint: bool = False
) -> torch.Tensor:
    """Compute K_star_star kernel matrix.

    Parameters
    ----------
    jacobians : torch.Tensor (b, C, P)

    Returns
    -------
    kernel : torch.tensor
        K_star with shape (b, C, C)

    """
    if joint:
        if self.independent_outputs:
            kernel = torch.einsum("acp,bcp->abc", jacobians, jacobians)
        else:
            kernel = torch.einsum("acp,bep->abce", jacobians, jacobians)

    else:
        if self.independent_outputs:
            kernel = torch.empty(
                (jacobians.shape[0], self.n_outputs),
                device=jacobians.device,
                dtype=self._dtype,
            )
            for c in range(self.n_outputs):
                kernel[:, c] = torch.norm(jacobians[:, c, :], dim=1) ** 2
        else:
            kernel = torch.einsum("bcp,bep->bce", jacobians, jacobians)
    return kernel

_kernel_batch_star

_kernel_batch_star(
    jacobians: Tensor, batch: Tensor
) -> Tensor

Compute K_b_star, which is a part of K_M_star kernel matrix.

Parameters:

• jacobians (Tensor(b1, C, P)) –
• batch (Tensor(b2, C)) –

Returns:

• kernel ( tensor ) –

  K_batch_star with shape (b1, b2, C, C)

Source code in laplace/baselaplace.py

def _kernel_batch_star(
    self, jacobians: torch.Tensor, batch: torch.Tensor
) -> torch.Tensor:
    """Compute K_b_star, which is a part of K_M_star kernel matrix.

    Parameters
    ----------
    jacobians : torch.Tensor (b1, C, P)
    batch : torch.Tensor (b2, C)

    Returns
    -------
    kernel : torch.tensor
        K_batch_star with shape (b1, b2, C, C)
    """
    jacobians_2, _ = self._jacobians(batch)
    if self.independent_outputs:
        kernel = torch.empty(
            (jacobians.shape[0], jacobians_2.shape[0], self.n_outputs),
            device=jacobians.device,
            dtype=self._dtype,
        )
        for c in range(self.n_outputs):
            kernel[:, :, c] = torch.einsum(
                "bp,ep->be", jacobians[:, c, :], jacobians_2[:, c, :]
            )
    else:
        kernel = torch.einsum("bcp,dep->bdce", jacobians, jacobians_2)
    return kernel

_jacobians

_jacobians(
    X: Tensor, enable_backprop: bool = None
) -> tuple

A wrapper function to compute jacobians - this enables reusing same kernel methods (kernel_batch etc.) in FunctionalLaplace and FunctionalLLLaplace by simply overwriting this method instead of all kernel methods.

Source code in laplace/baselaplace.py

def _jacobians(self, X: torch.Tensor, enable_backprop: bool = None) -> tuple:
    """A wrapper function to compute jacobians - this enables reusing same
    kernel methods (kernel_batch etc.) in FunctionalLaplace and FunctionalLLLaplace
    by simply overwriting this method instead of all kernel methods.
    """
    if enable_backprop is None:
        enable_backprop = self.enable_backprop
    return self.backend.jacobians(X, enable_backprop=enable_backprop)

_mean_scatter_term_batch

_mean_scatter_term_batch(Js: Tensor, f: Tensor, y: Tensor)

Compute mean vector in the scatter term in the log marginal likelihood

See scatter_lml property above for the exact equations of mean vectors in scatter terms for both types of likelihood (regression, classification).

Parameters:

• Js (tensor) –

  Jacobians (batch, output_shape, parameters)

• f (tensor) –

  NN output (batch, output_shape)

• y (Tensor) –

  data labels (batch, output_shape)

Returns:

• mu ( tensor ) –

  K_batch_star with shape (batch, output_shape)

Source code in laplace/baselaplace.py

def _mean_scatter_term_batch(
    self, Js: torch.Tensor, f: torch.Tensor, y: torch.Tensor
):
    """Compute mean vector in the scatter term in the log marginal likelihood

    See `scatter_lml` property above for the exact equations of mean vectors in scatter terms for
    both types of likelihood (regression, classification).

    Parameters
    ----------
    Js : torch.tensor
          Jacobians (batch, output_shape, parameters)
    f : torch.tensor
          NN output (batch, output_shape)
    y: torch.tensor
          data labels (batch, output_shape)

    Returns
    -------
    mu : torch.tensor
        K_batch_star with shape (batch, output_shape)
    """
    if self.likelihood == Likelihood.REGRESSION:
        return y - (f + torch.einsum("bcp,p->bc", Js, self.prior_mean - self.mean))
    elif self.likelihood == Likelihood.CLASSIFICATION:
        return -torch.einsum("bcp,p->bc", Js, self.prior_mean - self.mean)

log_marginal_likelihood

log_marginal_likelihood(
    prior_precision: Tensor | None = None,
    sigma_noise: Tensor | None = None,
) -> Tensor

Compute the Laplace approximation to the log marginal likelihood. Requires that the Laplace approximation has been fit before. The resulting torch.Tensor is differentiable in prior_precision and sigma_noise if these have gradients enabled. By passing prior_precision or sigma_noise, the current value is overwritten. This is useful for iterating on the log marginal likelihood.

Parameters:

• prior_precision (Tensor, default: None ) –

  prior precision if should be changed from current prior_precision value

• sigma_noise (Tensor, default: None ) –

  observation noise standard deviation if should be changed

Returns:

• log_marglik ( Tensor ) –

Source code in laplace/baselaplace.py

def log_marginal_likelihood(
    self,
    prior_precision: torch.Tensor | None = None,
    sigma_noise: torch.Tensor | None = None,
) -> torch.Tensor:
    """Compute the Laplace approximation to the log marginal likelihood.
    Requires that the Laplace approximation has been fit before.
    The resulting torch.Tensor is differentiable in `prior_precision` and
    `sigma_noise` if these have gradients enabled.
    By passing `prior_precision` or `sigma_noise`, the current value is
    overwritten. This is useful for iterating on the log marginal likelihood.

    Parameters
    ----------
    prior_precision : torch.Tensor, optional
        prior precision if should be changed from current `prior_precision` value
    sigma_noise : torch.Tensor, optional
        observation noise standard deviation if should be changed

    Returns
    -------
    log_marglik : torch.Tensor
    """
    # update prior precision (useful when iterating on marglik)
    if prior_precision is not None:
        self.prior_precision = prior_precision

    # update sigma_noise (useful when iterating on marglik)
    if sigma_noise is not None:
        if self.likelihood != Likelihood.REGRESSION:
            raise ValueError("Can only change sigma_noise for regression.")
        self.sigma_noise = sigma_noise

    return self.log_likelihood - 0.5 * (self.log_det_ratio + self.scatter)