Module Orthogonal

orthogonal

CLASS	DESCRIPTION
OrthogonalLayer	Base class for orthogonal neural network layers.
QRLayer	QR decomposition-based orthogonal layer.
SVDLayer	SVD-based orthogonal layer.
PCALayer	PCA-based orthogonal layer.
RunningQRLayer	Running QR projection orthogonal layer with weighted decomposition.

Classes

OrthogonalLayer(in_features: int, out_features: int | None = None, **kwargs)

Bases: Module

Base class for orthogonal neural network layers.

Provides a common interface for layers that update_weights and store orthogonal transformations during training.

PARAMETER	DESCRIPTION
in_features	Number of input features. TYPE: int
out_features	Number of output features, if None uses in_features. TYPE: int, optional, by default None DEFAULT: None

METHOD	DESCRIPTION
update_weights	Compute and store orthogonal transformation.
forward	Apply orthogonal transformation.

Source code in spectre/module/orthogonal.py

def __init__(
    self,
    in_features: int,
    out_features: int | None = None,
    **kwargs,
) -> None:
    super().__init__(**kwargs)

    if not isinstance(in_features, int):
        raise TypeError(
            f"Number of input features `in_features` must be an integer, got"
            f"{type(in_features)}."
        )
    check_in_interval(in_features, "[1, inf)")
    self.in_features = in_features

    if out_features is not None:
        if not isinstance(out_features, int):
            raise TypeError(
                f"Number of output features `out_features` must be an integer, got"
                f"{type(out_features)}."
            )
        check_in_interval(out_features, "[1, inf)")
        self.out_features = out_features
    else:
        self.out_features = in_features

    self.register_buffer("weights", torch.eye(self.in_features, self.out_features))

Functions

update_weights(X: torch.Tensor, weights: torch.Tensor | None = None) -> None abstractmethod

Compute and store orthogonal transformation.

Source code in spectre/module/orthogonal.py

@abstractmethod
def update_weights(
    self, X: torch.Tensor, weights: torch.Tensor | None = None
) -> None:
    """Compute and store orthogonal transformation."""
    raise NotImplementedError()

forward(X: torch.Tensor, weights: torch.Tensor | None = None) -> torch.Tensor

Apply orthogonal transformation.

During training, attempts to update_weights and store orthogonal weights using the specific method implemented in update_weights(). If the input has sufficient dimensions and no linear algebra errors occur, weights are updated. Returns input unchanged in all cases.

PARAMETER	DESCRIPTION
X	Input tensor of shape (n_samples, in_features). TYPE: Tensor
weights	Weights tensor of shape (n_samples,). TYPE: torch.Tensor | None, optional, by default None DEFAULT: None

RETURNS	DESCRIPTION
Tensor	Input tensor orthogonalized if not training. If training, returns X unchanged.

Source code in spectre/module/orthogonal.py

def forward(
    self, X: torch.Tensor, weights: torch.Tensor | None = None
) -> torch.Tensor:
    """
    Apply orthogonal transformation.

    During training, attempts to update_weights and store orthogonal weights using
    the specific method implemented in `update_weights()`. If the input has
    sufficient dimensions and no linear algebra errors occur, weights are
    updated. Returns input unchanged in all cases.

    Parameters
    ----------
    X : torch.Tensor
        Input tensor of shape (n_samples, in_features).

    weights : torch.Tensor | None, optional, by default None
        Weights tensor of shape (n_samples,).

    Returns
    -------
    torch.Tensor
        Input tensor orthogonalized if not training. If training, returns
        X unchanged.
    """
    if X.size(-1) >= self.in_features:
        if self.training:
            try:
                self.update_weights(X, weights=weights)
            except torch.linalg.LinAlgError as e:
                warnings.warn(
                    f"{self.__class__.__name__} orthogonalization failed: {e}. Weights not updated.",
                    RuntimeWarning,
                    stacklevel=2,
                )
            return X

        elif not self.training:
            return X @ self.weights
    else:
        return X

QRLayer(in_features: int, out_features: int | None = None, **kwargs)

Bases: OrthogonalLayer

QR decomposition-based orthogonal layer.

Computes QR decomposition (X=QR) during training and stores the inverse of R matrix as orthogonal weights. Q matrix provides an orthogonal basis while R contains the transformation coefficients.

PARAMETER	DESCRIPTION
in_features	Number of input features. TYPE: int
out_features	Number of output features, if None uses in_features. TYPE: int, optional, by default None DEFAULT: None

METHOD	DESCRIPTION
update_weights	Compute and store QR-based orthogonal transformation.

Source code in spectre/module/orthogonal.py

def __init__(
    self,
    in_features: int,
    out_features: int | None = None,
    **kwargs,
) -> None:
    super().__init__(in_features=in_features, out_features=out_features, **kwargs)

Functions

update_weights(X: torch.Tensor, weights: torch.Tensor | None = None) -> None

Compute and store QR-based orthogonal transformation.

PARAMETER	DESCRIPTION
X	Input tensor of shape (n_samples, in_features). TYPE: Tensor
weights	Weights tensor of shape (n_samples,). TYPE: torch.Tensor | None, optional, by default None DEFAULT: None

Source code in spectre/module/orthogonal.py

def update_weights(
    self, X: torch.Tensor, weights: torch.Tensor | None = None
) -> None:
    """
    Compute and store QR-based orthogonal transformation.

    Parameters
    ----------
    X : torch.Tensor
        Input tensor of shape (n_samples, in_features).

    weights : torch.Tensor | None, optional, by default None
        Weights tensor of shape (n_samples,).
    """
    _, R = torch.linalg.qr(X.view(-1, X.size(-1)))

    # Extract square submatrix for inversion
    dim = min(R.size(0), R.size(1))
    if dim == 0:
        return

    R_square = R[:dim, :dim]
    R_inv = torch.linalg.inv(R_square)

    # Copy available dimensions to weights
    rows = min(R_inv.size(0), self.in_features)
    cols = min(R_inv.size(1), self.out_features)

    self.weights[:rows, :cols].copy_(R_inv[:rows, :cols].detach())

SVDLayer(in_features: int, out_features: int | None = None, **kwargs)

Bases: OrthogonalLayer

SVD-based orthogonal layer.

Computes singular value decomposition during training and stores the right singular vectors as orthogonal weights. Provides optimal low-rank approximation through principal components.

PARAMETER	DESCRIPTION
in_features	Number of input features. TYPE: int
out_features	Number of output features, if None uses in_features. TYPE: int, optional, by default None DEFAULT: None

METHOD	DESCRIPTION
update_weights	Compute and store SVD-based orthogonal transformation.

Source code in spectre/module/orthogonal.py

def __init__(
    self,
    in_features: int,
    out_features: int | None = None,
    **kwargs,
) -> None:
    super().__init__(in_features=in_features, out_features=out_features, **kwargs)

Functions

update_weights(X: torch.Tensor, weights: torch.Tensor | None = None) -> None

Compute and store SVD-based orthogonal transformation.

PARAMETER	DESCRIPTION
X	Input tensor of shape (n_samples, in_features). TYPE: Tensor
weights	Sample weights. TYPE: torch.Tensor | None, optional, by default None DEFAULT: None

Source code in spectre/module/orthogonal.py

def update_weights(
    self, X: torch.Tensor, weights: torch.Tensor | None = None
) -> None:
    """
    Compute and store SVD-based orthogonal transformation.

    Parameters
    ----------
    X : torch.Tensor
        Input tensor of shape (n_samples, in_features).

    weights : torch.Tensor | None, optional, by default None
        Sample weights.
    """
    _, _, Vt = torch.linalg.svd(X.view(-1, X.size(-1)), full_matrices=False)
    # Vt has shape (min(n_samples, in_features), in_features)
    # Take first out_features right singular vectors (rows of Vt)
    n_components = min(Vt.size(0), self.out_features)
    if n_components > 0 and Vt.size(1) == self.in_features:
        # Transpose to get columns: (in_features, n_components)
        weights = Vt[:n_components, :].T
        # Copy only the available components
        self.weights[:, :n_components].copy_(weights.detach())

PCALayer(in_features: int, out_features: int | None = None, standardize: bool = True, **kwargs)

Bases: OrthogonalLayer

PCA-based orthogonal layer.

Computes principal component analysis during training and stores the principal components as orthogonal weights. Provides dimensionality reduction through variance maximization with optional data standardization.

PARAMETER	DESCRIPTION
in_features	Number of input features. TYPE: int
out_features	Number of output features, if None uses in_features. TYPE: int, optional, by default None DEFAULT: None
standardize	Whether to standardize input data before PCA. TYPE: bool, optional, by default True DEFAULT: True

METHOD	DESCRIPTION
update_weights	Compute and store PCA-based orthogonal transformation.

Source code in spectre/module/orthogonal.py

def __init__(
    self,
    in_features: int,
    out_features: int | None = None,
    standardize: bool = True,
    **kwargs,
) -> None:
    super().__init__(in_features=in_features, out_features=out_features, **kwargs)
    if not isinstance(standardize, bool):
        raise TypeError(
            f"`standardize` must be a boolean, got {type(standardize)}."
        )
    self.standardize = standardize

Functions

update_weights(X: torch.Tensor, weights: torch.Tensor | None = None) -> None

Compute and store PCA-based orthogonal transformation.

PARAMETER	DESCRIPTION
X	Input tensor of shape (n_samples, in_features). TYPE: Tensor
weights	Weights tensor of shape (n_samples,). TYPE: torch.Tensor | None, optional, by default None DEFAULT: None

Source code in spectre/module/orthogonal.py

def update_weights(
    self, X: torch.Tensor, weights: torch.Tensor | None = None
) -> None:
    """
    Compute and store PCA-based orthogonal transformation.

    Parameters
    ----------
    X : torch.Tensor
        Input tensor of shape (n_samples, in_features).

    weights : torch.Tensor | None, optional, by default None
        Weights tensor of shape (n_samples,).
    """
    result = principal_component_analysis(
        X=X.view(-1, X.size(-1)),
        weights=weights,
        out_features=self.out_features,
        standardize=self.standardize,
    )
    components = result.eigenvectors
    # components has shape (in_features, n_components)
    # where n_components <= min(n_samples, in_features, out_features)
    if (
        components.size(0) == self.in_features
        and components.size(1) > 0
        and components.size(1) <= self.out_features
    ):
        n_components = components.size(1)
        # Copy only the available components
        self.weights[:, :n_components].copy_(components.detach())

RunningQRLayer(in_features: int, eps: float = torch.finfo(torch.float32).eps, **kwargs)

Bases: OrthogonalLayer

Running QR projection orthogonal layer with weighted decomposition.

Maintains a running average of \(R^{-1}\) across training batches and supports both QR decomposition mode and projection mode. During training with QR mode, update_weightss weighted QR decomposition and returns orthogonal Q matrix. In projection mode, applies the running average transformation.

PARAMETER	DESCRIPTION
in_features	Number of input features. TYPE: int
eps	Small value to avoid division by zero. TYPE: float, optional, by default torch.finfo(torch.float32).eps DEFAULT: eps

Notes

The weighted QR decomposition produces a Q matrix that is orthogonal with respect to the weighted inner product, not the standard Euclidean inner product. For data with sample weights w_i, Q satisfies: \(Q^\top \mathrm{diag}(w) Q = I\) but NOT necessarily \(Q^\top Q = I\) in the unweighted case.

The running average of \(R^{-1}\) across batches is a heuristic approximation that does not have strict mathematical guarantees for orthogonalization.

METHOD	DESCRIPTION
update_weights	Compute weighted QR decomposition and update running average.

Source code in spectre/module/orthogonal.py

def __init__(
    self,
    in_features: int,
    eps: float = torch.finfo(torch.float32).eps,
    **kwargs,
) -> None:
    super().__init__(in_features=in_features, out_features=in_features, **kwargs)

    if not isinstance(eps, (float, int)):
        raise ValueError(f"`eps` must be a number, got {type(eps)}.")
    check_in_interval(eps, "[0, 1]")
    self.eps = eps

    self.register_buffer("num_batches_tracked", torch.tensor(0, dtype=torch.long))

Functions

update_weights(X: torch.Tensor, weights: torch.Tensor | None = None) -> None

Compute weighted QR decomposition and update running average.

PARAMETER	DESCRIPTION
X	Input tensor of shape (n_samples, in_features). TYPE: Tensor
weights	Sample weights of shape (n_samples, ). TYPE: torch.Tensor, optional, by default None DEFAULT: None

Source code in spectre/module/orthogonal.py

def update_weights(
    self, X: torch.Tensor, weights: torch.Tensor | None = None
) -> None:
    """
    Compute weighted QR decomposition and update running average.

    Parameters
    ----------
    X : torch.Tensor
        Input tensor of shape (n_samples, in_features).

    weights : torch.Tensor, optional, by default None
        Sample weights of shape (n_samples, ).
    """
    if weights is not None:
        check_sample_weights(weights)
    else:
        weights = torch.ones(X.shape[0], device=X.device, dtype=X.dtype)

    weight_sum = weights.sum()
    weights = weights / weight_sum

    # Clamp to avoid negative values from numerical errors
    sqrt_w = torch.sqrt(torch.clamp(weights, min=0.0)).unsqueeze(1)
    X_weighted = X * sqrt_w

    Q, R = torch.linalg.qr(X_weighted)
    R_inv = torch.linalg.pinv(R)

    # Update running average
    self.num_batches_tracked += 1
    decay = 1.0 / max(self.num_batches_tracked.item(), 1)

    self.weights = (1 - decay) * self.weights + decay * R_inv.detach()

Functions