Loss Decorrelation

decorrelation

CLASS	DESCRIPTION
DecorrelationLoss	Decorrelation loss for learning independent latent representations.

FUNCTION	DESCRIPTION
decorrelation_loss	Compute decorrelation loss for learning independent representations.

Classes

DecorrelationLoss(method: Literal['correlation', 'covariance', 'total_correlation'] = 'correlation', reduce: Literal['sum', 'mean'] = 'mean', sign: float = 1.0, eps: float = torch.finfo(torch.float32).eps, context_prefix: str | None = None, **kwargs)

Bases: Loss

Decorrelation loss for learning independent latent representations.

Penalizes statistical dependencies between latent dimensions to encourage learning of disentangled or decorrelated features.

Supports three decorrelation methods:

• "correlation": Penalizes off-diagonal correlation coefficients (default)
• "covariance": Penalizes off-diagonal covariance elements (faster)
• "total_correlation": Information-theoretic measure of dependence (advanced)

The correlation loss computes: \(L = \sum_{i \neq j} \rho_{ij}^2\), where \(\rho\) is the correlation matrix.

The covariance loss computes: \(L = \sum_{i \neq j} \text{Cov}(z_i, z_j)^2\)

The total correlation loss approximates: \(\text{TC}(Z) = \text{KL}(q(z) \| \prod_i q(z_i))\)

PARAMETER	DESCRIPTION
method	Decorrelation method to use. "correlation": Most interpretable, computes sum/mean of squared off-diagonal correlations. Range: [0, n_features*(n_features-1)] for reduce="sum", [0, 1] for reduce="mean". "covariance": Faster, no normalization step. Range: [0, inf) (unbounded, depends on data scale). "total_correlation": Information-theoretic, penalizes all dependencies. Range: [0, inf) (non-negative, bounded by entropy). TYPE: str, by default "correlation" DEFAULT: 'correlation'
reduce	Reduction method for off-diagonal elements. TYPE: Literal["sum", "mean"], by default "mean" DEFAULT: 'mean'
sign	Loss sign multiplier. Use 1.0 to minimize correlation (typical). TYPE: float, by default 1.0 DEFAULT: 1.0
eps	Small constant for numerical stability in variance/correlation computation. TYPE: float, by default torch.finfo(torch.float32).eps DEFAULT: eps
**kwargs	Additional keyword arguments for the loss function. DEFAULT: {}

Examples:

Basic usage with correlation penalty:

>>> from spectre.loss import DecorrelationLoss
>>> import torch
>>>
>>> loss_fn = DecorrelationLoss(method="correlation", reduce="mean")
>>> Z = torch.randn(100, 5)  # 100 samples, 5 latent dimensions
>>> loss = loss_fn({"Z": Z})

With sample weights:

>>> weights = torch.rand(100)
>>> loss = loss_fn({"Z": Z, "weights": weights})

Notes

• For highly correlated inputs, use method="correlation"
• Total correlation requires larger sample sizes for accurate estimation

METHOD	DESCRIPTION
forward	Compute decorrelation loss.

Source code in spectre/loss/decorrelation.py

def __init__(
    self,
    method: Literal[
        "correlation", "covariance", "total_correlation"
    ] = "correlation",
    reduce: Literal["sum", "mean"] = "mean",
    sign: float = 1.0,
    eps: float = torch.finfo(torch.float32).eps,
    context_prefix: str | None = None,
    **kwargs,
):
    super().__init__(
        reduce=reduce, sign=sign, context_prefix=context_prefix, **kwargs
    )

    if method not in ["correlation", "covariance", "total_correlation"]:
        raise ValueError(
            f"Method must be 'correlation', 'covariance', or 'total_correlation', "
            f"got '{method}'."
        )
    self.method = method

    if not isinstance(eps, float) or eps <= 0:
        raise ValueError(f"eps must be a positive float, got {eps}.")
    self.eps = eps

Functions

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

Compute decorrelation loss.

PARAMETER	DESCRIPTION
context	Context dictionary containing: "Z": torch.Tensor of shape (n_samples, n_features). Latent representation or features to decorrelate. "weights": torch.Tensor of shape (n_samples,), optional. Sample weights for weighted correlation/covariance computation. TYPE: dict
**kwargs	Additional keyword arguments). DEFAULT: {}

RETURNS	DESCRIPTION
Tensor	Decorrelation loss.

Source code in spectre/loss/decorrelation.py

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

    Parameters
    ----------
    context : dict
        Context dictionary containing:

        - "Z": torch.Tensor of shape `(n_samples, n_features)`. Latent
          representation or features to decorrelate.
        - "weights": torch.Tensor of shape `(n_samples,)`, optional. Sample
          weights for weighted correlation/covariance computation.

    **kwargs
        Additional keyword arguments).

    Returns
    -------
    torch.Tensor
        Decorrelation loss.
    """
    Z = self._get_context(context, "Z", None)
    weights = self._get_context(context, "weights", None)

    if Z is None:
        raise ValueError("`DecorrelationLoss` requires `Z` in context.")

    check_2d(Z)

    if weights is not None:
        check_same_len(Z, weights)
    else:
        weights = torch.ones(Z.shape[0], device=Z.device, dtype=Z.dtype)

    return _decorrelation_loss_impl(
        Z=Z,
        weights=weights,
        method=self.method,
        reduce=self.reduce,
        sign=self.sign,
        eps=self.eps,
    )

Functions

decorrelation_loss(Z: torch.Tensor, weights: torch.Tensor | None = None, method: Literal['correlation', 'covariance', 'total_correlation'] = 'correlation', reduce: Literal['sum', 'mean'] = 'mean', sign: float = 1.0, eps: float = torch.finfo(torch.float32).eps) -> torch.Tensor

Compute decorrelation loss for learning independent representations.

Functional interface for DecorrelationLoss. Penalizes statistical dependencies between latent dimensions.

PARAMETER	DESCRIPTION
Z	Latent representation or features to decorrelate. TYPE: torch.Tensor of shape (n_samples, n_features)
weights	Sample weights for weighted computation. If None, uniform weights used. TYPE: torch.Tensor of shape (n_samples,) or None, by default None DEFAULT: None
method	Decorrelation method. TYPE: str, by default "correlation" DEFAULT: 'correlation'
reduce	Reduction method for off-diagonal elements. TYPE: Literal["sum", "mean"], by default "mean" DEFAULT: 'mean'
sign	Loss sign multiplier. TYPE: float, by default 1.0 DEFAULT: 1.0
eps	Numerical stability constant. TYPE: float, by default torch.finfo(torch.float32).eps DEFAULT: eps

RETURNS	DESCRIPTION
Tensor	Decorrelation loss.

Examples:

>>> import torch
>>> from spectre.loss import decorrelation_loss
>>>
>>> Z = torch.randn(100, 5)
>>> loss = decorrelation_loss(Z, method="correlation")
>>>
>>> # With weights
>>> weights = torch.rand(100)
>>> loss = decorrelation_loss(Z, weights=weights, method="covariance")

Source code in spectre/loss/decorrelation.py

def decorrelation_loss(
    Z: torch.Tensor,
    weights: torch.Tensor | None = None,
    method: Literal["correlation", "covariance", "total_correlation"] = "correlation",
    reduce: Literal["sum", "mean"] = "mean",
    sign: float = 1.0,
    eps: float = torch.finfo(torch.float32).eps,
) -> torch.Tensor:
    """
    Compute decorrelation loss for learning independent representations.

    Functional interface for [DecorrelationLoss][spectre.loss.DecorrelationLoss].
    Penalizes statistical dependencies between latent dimensions.

    Parameters
    ----------
    Z : torch.Tensor of shape (n_samples, n_features)
        Latent representation or features to decorrelate.

    weights : torch.Tensor of shape (n_samples,) or None, by default None
        Sample weights for weighted computation. If None, uniform weights used.

    method : str, by default "correlation"
        Decorrelation method.

    reduce : Literal["sum", "mean"], by default "mean"
        Reduction method for off-diagonal elements.

    sign : float, by default 1.0
        Loss sign multiplier.

    eps : float, by default torch.finfo(torch.float32).eps
        Numerical stability constant.

    Returns
    -------
    torch.Tensor
        Decorrelation loss.

    Examples
    --------
    >>> import torch
    >>> from spectre.loss import decorrelation_loss
    >>>
    >>> Z = torch.randn(100, 5)
    >>> loss = decorrelation_loss(Z, method="correlation")
    >>>
    >>> # With weights
    >>> weights = torch.rand(100)
    >>> loss = decorrelation_loss(Z, weights=weights, method="covariance")
    """
    check_2d(Z)

    if weights is not None:
        check_same_len(Z, weights)
    else:
        weights = torch.ones(Z.shape[0], device=Z.device, dtype=Z.dtype)

    if method not in ["correlation", "covariance", "total_correlation"]:
        raise ValueError(
            f"Method must be 'correlation', 'covariance', or 'total_correlation', "
            f"got '{method}'."
        )

    if reduce not in ["sum", "mean"]:
        raise ValueError(f"reduce must be 'sum' or 'mean', got '{reduce}'.")

    return _decorrelation_loss_impl(
        Z=Z,
        weights=weights,
        method=method,
        reduce=reduce,
        sign=sign,
        eps=eps,
    )