Loss Orthogonal

orthogonal

CLASS	DESCRIPTION
OrthogonalLoss	Orthogonality regularization loss. Encourages orthogonality by penalizing

FUNCTION	DESCRIPTION
orthogonal_loss	Compute orthogonality regularization loss with optional sample weighting.

Classes

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

Bases: Loss

Orthogonality regularization loss. Encourages orthogonality by penalizing deviation from identity matrix.

Optionally supports sample weighting when computing the orthogonality penalty.

The loss computes \(\|(Z Z^\top - I)\|^2\) where \(I\) is the identity matrix. When weights are provided, the loss becomes \(\|(W^{1/2} Z Z^\top W^{1/2} - I)\|^2\) where \(W\) is a diagonal matrix of weights.

PARAMETER	DESCRIPTION
reduce	Reduction method. 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 orthogonality loss.

Source code in spectre/loss/orthogonal.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 orthogonality loss.

PARAMETER	DESCRIPTION
context	Context dictionary. "Z": torch.Tensor. Input matrix to regularize for orthogonality of shape (n_samples, n_features). "weights": torch.Tensor | None, optional. Sample weights of shape (n_samples, ). If provided, applies per-sample weighting to the orthogonality penalty. TYPE: dict
**kwargs	Additional keyword arguments to pass to the loss function. DEFAULT: {}

RETURNS	DESCRIPTION
Tensor	Orthogonality loss.

Source code in spectre/loss/orthogonal.py

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

    Parameters
    ----------
    context : dict
        Context dictionary.

        - "Z": torch.Tensor. Input matrix to regularize for orthogonality of shape
          `(n_samples, n_features)`.
        - "weights": torch.Tensor | None, optional. Sample weights of shape
          `(n_samples, )`. If provided, applies per-sample weighting to
          the orthogonality penalty.

    **kwargs
        Additional keyword arguments to pass to the loss function.

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

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

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

    return _orthogonal_loss_impl(
        Z, weights=weights, reduce=self.reduce, sign=self.sign
    )

Functions

orthogonal_loss(Z: torch.Tensor, weights: torch.Tensor | None = None, reduce: str = 'mean', sign: float = 1.0) -> torch.Tensor

Compute orthogonality regularization loss with optional sample weighting.

The loss computes \(\|(Z Z^\top - I)\|^2\) where \(I\) is the identity matrix. When weights are provided, the loss becomes \(\|(W^{1/2} Z Z^\top W^{1/2} - I)\|^2\) where \(W\) is a diagonal matrix of weights.

PARAMETER	DESCRIPTION
Z	Input matrix to regularize for orthogonality, shape (n_samples, n_features). TYPE: Tensor
weights	Sample weights of shape (n_samples, ). If provided, applies per-sample weighting to the orthogonality penalty. TYPE: torch.Tensor | None, optional, by default None DEFAULT: None
reduce	Reduction method. TYPE: Literal["mean", "sum"], optional, by default "mean" DEFAULT: 'mean'
sign	Loss sign multiplier. TYPE: float, optional, by default 1.0 DEFAULT: 1.0

RETURNS	DESCRIPTION
Tensor	Orthogonality loss value.

Source code in spectre/loss/orthogonal.py

def orthogonal_loss(
    Z: torch.Tensor,
    weights: torch.Tensor | None = None,
    reduce: str = "mean",
    sign: float = 1.0,
) -> torch.Tensor:
    """
    Compute orthogonality regularization loss with optional sample weighting.

    The loss computes $\\|(Z Z^\\top - I)\\|^2$ where $I$ is the
    identity matrix. When weights are provided, the loss becomes
    $\\|(W^{1/2} Z Z^\\top W^{1/2} - I)\\|^2$ where $W$ is a
    diagonal matrix of weights.

    Parameters
    ----------
    Z : torch.Tensor
        Input matrix to regularize for orthogonality, shape `(n_samples, n_features)`.

    weights : torch.Tensor | None, optional, by default None
        Sample weights of shape `(n_samples, )`. If provided, applies per-sample
        weighting to the orthogonality penalty.

    reduce : Literal["mean", "sum"], optional, by default "mean"
        Reduction method.

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

    Returns
    -------
    torch.Tensor
        Orthogonality loss value.
    """
    if weights is not None:
        check_1d(weights)
        check_same_len(Z, weights)
    else:
        weights = torch.ones(Z.shape[0], device=Z.device, dtype=Z.dtype)

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

    return _orthogonal_loss_impl(Z, weights=weights, reduce=reduce, sign=sign)