Loss Elbo

elbo

CLASS	DESCRIPTION
ELBOGaussianLoss	Loss for Gaussian ELBO (Evidence Lower BOund).

FUNCTION	DESCRIPTION
elbo_gaussian_loss	Compute Gaussian ELBO loss for variational autoencoders.

Classes

ELBOGaussianLoss(reduce: Literal['sum', 'mean'] = 'mean', sign: float = 1.0, context_prefix: str | None = None, **kwargs)

Bases: Loss

Loss for Gaussian ELBO (Evidence Lower BOund).

Computes the negative ELBO for variational autoencoders, combining Kullback-Leibler (KL) divergence between posterior and prior with reconstruction loss.

The loss function is defined as \(L = \| X - Z \|^2 + \frac{1}{2} \sum_k (\exp(\log \sigma^2_k) + \mu_k^2 - 1 - \log \sigma^2_k)\), where:

• \(\| X - Z \|^2\) is the MSE reconstruction loss
• The sum term is the KL divergence \(\mathrm{KL}(N(\mu, \sigma^2) \| N(0, I))\)
• \(\mu\) is mean, \(\sigma^2\) is variance

PARAMETER	DESCRIPTION
reduce	Reduction operation for KL divergence. "sum": Sum over batch dimension "mean": Mean over batch dimension TYPE: Literal["sum", "mean"], optional, by default "mean" DEFAULT: 'mean'
sign	Loss sign multiplier, by default 1.0. TYPE: float DEFAULT: 1.0
**kwargs	Additional keyword arguments. DEFAULT: {}

METHOD	DESCRIPTION
forward	Compute Gaussian ELBO loss.

Source code in spectre/loss/elbo.py

def __init__(
    self,
    reduce: Literal["sum", "mean"] = "mean",
    sign: float = 1.0,
    context_prefix: str | None = None,
    **kwargs,
):
    super().__init__(
        reduce=reduce, sign=sign, context_prefix=context_prefix, **kwargs
    )

Functions

forward(context: dict, **kwargs) -> torch.Tensor

Compute Gaussian ELBO loss.

PARAMETER	DESCRIPTION
context	Dictionary containing the input data and model outputs. TYPE: dict
kwargs	Additional keyword arguments. TYPE: dict DEFAULT: {}

RETURNS	DESCRIPTION
Tensor	The computed ELBO loss.

Source code in spectre/loss/elbo.py

def forward(self, context: dict, **kwargs) -> torch.Tensor:
    """
    Compute Gaussian ELBO loss.

    Parameters
    ----------
    context : dict
        Dictionary containing the input data and model outputs.

    kwargs : dict
        Additional keyword arguments.

    Returns
    -------
    torch.Tensor
        The computed ELBO loss.
    """
    X = self._get_context(context, "X", None)
    Y = self._get_context(context, "Y", None)
    mean = self._get_context(context, "mean", None)
    logvar = self._get_context(context, "logvar", None)
    weights = self._get_context(context, "weights", None)

    if X is None or Y is None or mean is None or logvar is None:
        raise ValueError(
            "`ELBOGaussianLoss` requires `X`, `Y`, `mean`, and `logvar` in context."
        )

    if weights is None:
        weights = torch.ones(X.shape[0], device=X.device, dtype=X.dtype)
    weights = weights.squeeze()

    return _elbo_gaussian_loss_impl(
        X, Y, mean, logvar, weights, reduce=self.reduce, sign=self.sign
    )

Functions

elbo_gaussian_loss(X: torch.Tensor, Y: torch.Tensor, mean: torch.Tensor, logvar: torch.Tensor, weights: Optional[torch.Tensor] = None, reduce: str = 'mean', sign: float = 1.0) -> torch.Tensor

Compute Gaussian ELBO loss for variational autoencoders.

The ELBO loss combines KL divergence between posterior q(z|x) and prior p(z) with reconstruction loss between input and output.

The loss function is defined as \(L = \| X - Y \|^2 + \frac{1}{2} \sum_k (\exp(\log \sigma^2_k) + \mu_k^2 - 1 - \log \sigma^2_k)\), where:

• \(\| X - Y \|^2\) is the MSE reconstruction loss
• The sum term is the KL divergence \(\mathrm{KL}(N(\mu, \sigma^2) \| N(0, I))\)
• \(\mu\) is mean, \(\sigma^2\) is variance

PARAMETER	DESCRIPTION
X	Input data tensor. TYPE: Tensor
Y	Reconstructed data tensor (same shape as X). TYPE: Tensor
mean	Mean of posterior Gaussian distribution q(z|x). TYPE: Tensor
logvar	Log variance of posterior Gaussian distribution (same shape as mean). TYPE: Tensor
weights	Sample weights of shape (n_batches,) or (n_batches, 1). TYPE: Optional[torch.Tensor], optional, by default None DEFAULT: None
reduce	Reduction operation for KL divergence. Options: "sum", "mean". TYPE: str, optional, by default "mean" DEFAULT: 'mean'
sign	Loss sign multiplier. TYPE: float, optional, by default 1.0 DEFAULT: 1.0

RETURNS	DESCRIPTION
Tensor	ELBO loss (KL divergence + MSE reconstruction loss).

Source code in spectre/loss/elbo.py

def elbo_gaussian_loss(
    X: torch.Tensor,
    Y: torch.Tensor,
    mean: torch.Tensor,
    logvar: torch.Tensor,
    weights: Optional[torch.Tensor] = None,
    reduce: str = "mean",
    sign: float = 1.0,
) -> torch.Tensor:
    """
    Compute Gaussian ELBO loss for variational autoencoders.

    The ELBO loss combines KL divergence between posterior q(z|x) and prior p(z)
    with reconstruction loss between input and output.

    The loss function is defined as
    $L = \\| X - Y \\|^2 + \\frac{1}{2} \\sum_k (\\exp(\\log \\sigma^2_k)
        + \\mu_k^2 - 1 - \\log \\sigma^2_k)$,
    where:

    - $\\| X - Y \\|^2$ is the MSE reconstruction loss
    - The sum term is the KL divergence $\\mathrm{KL}(N(\\mu, \\sigma^2) \\| N(0, I))$
    - $\\mu$ is mean, $\\sigma^2$ is variance

    Parameters
    ----------
    X : torch.Tensor
        Input data tensor.

    Y : torch.Tensor
        Reconstructed data tensor (same shape as X).

    mean : torch.Tensor
        Mean of posterior Gaussian distribution q(z|x).

    logvar : torch.Tensor
        Log variance of posterior Gaussian distribution (same shape as mean).

    weights : Optional[torch.Tensor], optional, by default None
        Sample weights of shape (n_batches,) or (n_batches, 1).

    reduce : str, optional, by default "mean"
        Reduction operation for KL divergence. Options: "sum", "mean".

    sign : float, optional, by default 1.0
        Loss sign multiplier.

    Returns
    -------
    torch.Tensor
        ELBO loss (KL divergence + MSE reconstruction loss).
    """
    check_same_shape(X, Y)
    check_same_shape(mean, logvar)

    if weights is not None:
        if weights.ndim > 1:
            weights = weights.squeeze()
        check_1d(weights)
        check_same_len(X, weights)
    else:
        weights = torch.ones(X.shape[0], device=X.device, dtype=X.dtype)

    if reduce not in ["sum", "mean"]:
        raise ValueError("Unknown reduce operation. Options: 'sum', 'mean'.")

    return _elbo_gaussian_loss_impl(X, Y, mean, logvar, weights, reduce, sign)