Loss Elbo Log Cosh

elbo_log_cosh

CLASS	DESCRIPTION
ELBOLogCoshLoss	Variational Autoencoder Log-Cosh loss function.

FUNCTION	DESCRIPTION
elbo_log_cosh_loss	Compute ELBO log-cosh loss for variational autoencoders.

Classes

ELBOLogCoshLoss(log_cosh_weight: float = 100.0, kl_div_weight: float = 10.0, reduce_mse: Literal['sum', 'mean'] = 'mean', reduce_kl_div: Literal['sum', 'mean'] = 'mean', sign: float = 1.0, context_prefix: str | None = None, **kwargs)

Bases: Loss

Variational Autoencoder Log-Cosh loss function.

Combines log-cosh reconstruction loss with KL divergence for VAE training. Log-cosh provides a smooth alternative to MSE that is less sensitive to outliers.

PARAMETER	DESCRIPTION
log_cosh_weight	Log-cosh scaling parameter. TYPE: float, optional, by default 100 DEFAULT: 100.0
kl_div_weight	KL divergence weighting factor. TYPE: float, optional, by default 10 DEFAULT: 10.0
reduce_mse	Reduction method for reconstruction loss. TYPE: Literal["sum", "mean"], optional, by default "mean" DEFAULT: 'mean'
reduce_kl_div	Reduction method for KL divergence. TYPE: Literal["sum", "mean"], optional, by default "mean" DEFAULT: 'mean'
sign	Loss sign multiplier. TYPE: float, optional, by default 1.0 DEFAULT: 1.0

METHOD	DESCRIPTION
forward	Compute ELBO log-cosh loss.

Source code in spectre/loss/elbo_log_cosh.py

def __init__(
    self,
    log_cosh_weight: float = 100.0,
    kl_div_weight: float = 10.0,
    reduce_mse: Literal["sum", "mean"] = "mean",
    reduce_kl_div: Literal["sum", "mean"] = "mean",
    sign: float = 1.0,
    context_prefix: str | None = None,
    **kwargs,
) -> None:
    # Don't pass reduce to parent since this loss has dual reduction
    super().__init__(
        reduce="mean", sign=sign, context_prefix=context_prefix, **kwargs
    )

    if not isinstance(log_cosh_weight, float):
        raise ValueError("`log_cosh_weight` must be a float.")
    check_in_interval(log_cosh_weight, "[0, inf)")
    self.log_cosh_weight = log_cosh_weight

    if not isinstance(kl_div_weight, float):
        raise ValueError("`kl_div_weight` must be a float.")
    check_in_interval(kl_div_weight, "[0, inf)")
    self.kl_div_weight = kl_div_weight

    if reduce_mse not in self.allowed_reduce:
        raise ValueError(f"reduce_mse must be one of {self.allowed_reduce}.")
    self.reduce_mse = reduce_mse

    if reduce_kl_div not in self.allowed_reduce:
        raise ValueError(f"reduce_kl_div must be one of {self.allowed_reduce}.")
    self.reduce_kl_div = reduce_kl_div

Functions

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

Compute ELBO log-cosh loss.

Arguments

kwargs: dict Additional keyword arguments.

RETURNS	DESCRIPTION
loss	ELBO log-cosh loss. TYPE: Tensor

Source code in spectre/loss/elbo_log_cosh.py

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

    Arguments
    ---------
    context: dict
        A dictionary containing X (input data), Y (reconstruction),
        mean (latent mean), logvar (latent log variance), and weights (optional).

    kwargs: dict
        Additional keyword arguments.

    Returns
    -------
    loss: torch.Tensor
        ELBO log-cosh 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(
            "`ELBOLogCoshLoss` requires `X`, `Y`, `mean`, and `logvar` in context."
        )

    return elbo_log_cosh_loss(
        X,
        Y,
        mean=mean,
        logvar=logvar,
        weights=weights,
        log_cosh_weight=self.log_cosh_weight,
        kl_div_weight=self.kl_div_weight,
        reduce_mse=self.reduce_mse,
        reduce_kl_div=self.reduce_kl_div,
        sign=self.sign,
    )

Functions

elbo_log_cosh_loss(X: torch.Tensor, Y: torch.Tensor, mean: torch.Tensor, logvar: torch.Tensor, weights: torch.Tensor | None = None, log_cosh_weight: float = 100, kl_div_weight: float = 10, reduce_mse: Literal['sum', 'mean'] = 'mean', reduce_kl_div: Literal['sum', 'mean'] = 'mean', sign: float = 1.0) -> torch.Tensor

Compute ELBO log-cosh loss for variational autoencoders.

Combines log-cosh reconstruction loss with KL divergence. Log-cosh provides a smooth alternative to MSE that is less sensitive to outliers.

PARAMETER	DESCRIPTION
X	Input data tensor. TYPE: Tensor
Y	Reconstructed data tensor (same shape as X). TYPE: Tensor
mean	Mean of latent Gaussian distribution. TYPE: Tensor
logvar	Log variance of latent Gaussian distribution (same shape as mean). TYPE: Tensor
weights	Sample weights. TYPE: Optional[torch.Tensor], optional, by default None DEFAULT: None
log_cosh_weight	Log-cosh scaling parameter. TYPE: float, optional, by default 100 DEFAULT: 100
kl_div_weight	KL divergence weighting factor. TYPE: float, optional, by default 10 DEFAULT: 10
reduce_mse	Reduction method for reconstruction loss. TYPE: Literal["sum", "mean"], optional, by default "mean" DEFAULT: 'mean'
reduce_kl_div	Reduction method for KL divergence. TYPE: Literal["sum", "mean"], optional, by default "mean" DEFAULT: 'mean'
sign	Loss sign multiplier. TYPE: float, optional, by default 1.0 DEFAULT: 1.0

RETURNS	DESCRIPTION
Tensor	ELBO log-cosh loss value.

Source code in spectre/loss/elbo_log_cosh.py

def elbo_log_cosh_loss(
    X: torch.Tensor,
    Y: torch.Tensor,
    mean: torch.Tensor,
    logvar: torch.Tensor,
    weights: torch.Tensor | None = None,
    log_cosh_weight: float = 100,
    kl_div_weight: float = 10,
    reduce_mse: Literal["sum", "mean"] = "mean",
    reduce_kl_div: Literal["sum", "mean"] = "mean",
    sign: float = 1.0,
) -> torch.Tensor:
    """
    Compute ELBO log-cosh loss for variational autoencoders.

    Combines log-cosh reconstruction loss with KL divergence. Log-cosh provides
    a smooth alternative to MSE that is less sensitive to outliers.

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

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

    mean : torch.Tensor
        Mean of latent Gaussian distribution.

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

    weights : Optional[torch.Tensor], optional, by default None
        Sample weights.

    log_cosh_weight : float, optional, by default 100
        Log-cosh scaling parameter.

    kl_div_weight : float, optional, by default 10
        KL divergence weighting factor.

    reduce_mse : Literal["sum", "mean"], optional, by default "mean"
        Reduction method for reconstruction loss.

    reduce_kl_div : Literal["sum", "mean"], optional, by default "mean"
        Reduction method for KL divergence.

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

    Returns
    -------
    torch.Tensor
        ELBO log-cosh loss value.
    """
    check_same_shape(X, Y)
    check_same_shape(mean, logvar)

    if X.ndim == 1:
        X = X.unsqueeze(1)
    if Y.ndim == 1:
        Y = Y.unsqueeze(1)

    mse = X - Y
    mse = (
        log_cosh_weight * mse
        + torch.log(1 + torch.exp(-2 * log_cosh_weight * mse))
        - torch.log(torch.tensor(2.0))
    )

    if weights is not None:
        mse = mse * weights.unsqueeze(1)
    else:
        mse = mse

    if reduce_mse == "sum":
        mse = mse.sum()
    elif reduce_mse == "mean":
        mse = mse.mean()
    else:
        raise ValueError(
            f"Unknown reduce operation '{reduce_mse}'. Options: 'sum', 'mean'."
        )

    kl_div = -0.5 * (logvar - logvar.exp() - mean**2 + 1).sum(dim=1)

    if weights is not None:
        check_same_shape(weights.squeeze(), kl_div)
        kl_div = kl_div * weights.squeeze()

    if reduce_kl_div == "sum":
        kl_div = kl_div.sum()
    elif reduce_kl_div == "mean":
        kl_div = kl_div.mean()
    else:
        raise ValueError(
            f"Unknown reduce operation '{reduce_kl_div}'. Options: 'sum', 'mean'."
        )

    return sign * (mse / log_cosh_weight + kl_div_weight * kl_div)