Residual blocks

yetanotherspdnet.nn.rresnet_layers — the building blocks of RResNet and GBWBNRResNet. SpectralVectorField computes a tangent vector from the spectrum of each matrix; ResidualBlock moves the matrix along it with an affine-invariant (unit step) or log-Euclidean exponential map. See Riemannian residual networks.

ResidualBlock

Riemannian residual block on SPD manifold.

SpectralVectorField

Spectral vector field on SPD manifold.

affine_invariant_norm

Affine-invariant norm of tangent vectors.
class ResidualBlock(
n_features,
spectrum_type='conv1d',
spectrum_hidden_dim=3,
spectrum_n_layers=2,
spectrum_kernel_size=5,
stiefel_parametrization_mode='static',
stiefel_n_steps_ref_update=100,
use_autograd=False,
metric='affine_invariant',
device=device(type='cpu'),
dtype=torch.float64,
generator=None,
)[source]

Riemannian residual block on SPD manifold.

Applies one residual step along the vector field V = VF(X) (a SpectralVectorField), following Katsman et al., Riemannian Residual Neural Networks (NeurIPS 2023) and their reference implementation:

  • metric="affine_invariant": X_new = projx(Exp_X(V / ||V||_X)), a step of unit affine-invariant length. Without this normalization, ||V||_X = ||X^{-1/2} V X^{-1/2}||_F grows like 1 / lambda_min(X) and the matrix exponential overflows on ill-conditioned inputs.

  • metric="log_euclidean": X_new = exp(log(X) + V).

Parameters:

  • n_features (int) – Dimension of SPD matrices

  • spectrum_type (str, optional) – Type of spectrum mapping. Default is “conv1d”

  • spectrum_hidden_dim (int, optional) – Hidden dimension for spectrum mapping. Default is 3

  • spectrum_n_layers (int, optional) – Hidden layers in spectrum mapping. Default is 2

  • spectrum_kernel_size (int, optional) – Kernel size for Conv1d. Default is 5

  • stiefel_parametrization_mode (str, optional) – Parametrization mode for Q matrix. Default is “static”

  • stiefel_n_steps_ref_update (int, optional) – Steps between reference updates. Default is 100

  • use_autograd (bool, optional) – Use autograd for exp map gradient. Default is False

  • metric (str, optional) – “affine_invariant” (default) or “log_euclidean”, see above

  • device (torch.device, optional) – Device. Default is torch.device(“cpu”)

  • dtype (torch.dtype, optional) – Data type. Default is torch.float64

  • generator (torch.Generator | None, optional) – Generator. Default is None

Initialize internal Module state, shared by both nn.Module and ScriptModule.

forward(data)[source]

Forward pass: one residual step (see the class docstring).

Parameters:

data (torch.Tensor of shape (..., n, n)) – Batch of SPD matrices

Returns:

result (torch.Tensor of shape (..., n, n)) – Updated SPD matrices after one residual step

Return type:

Tensor

register_optimizer_hook(optimizer)[source]

Register optimizer hooks for dynamic parametrizations.

class SpectralVectorField(
n_features,
spectrum_type='conv1d',
spectrum_hidden_dim=3,
spectrum_n_layers=2,
spectrum_kernel_size=5,
stiefel_parametrization_mode='static',
stiefel_n_steps_ref_update=100,
use_autograd=False,
device=device(type='cpu'),
dtype=torch.float64,
generator=None,
)[source]

Spectral vector field on SPD manifold.

Computes a tangent vector at X as:

V = Q diag(f(spec(X))) Q^T

where: - spec(X): eigenvalues of the input SPD matrix - f: learnable spectrum mapping (Conv1d or MLP) - Q: learnable orthogonal matrix (Stiefel-parametrized)

The output is a symmetric matrix in the tangent space at X.

Parameters:

  • n_features (int) – Dimension of SPD matrices (n x n)

  • spectrum_type (str, optional) – Type of spectrum mapping network. Default is “conv1d”. Choices are: “conv1d” and “mlp”

  • spectrum_hidden_dim (int, optional) – Hidden dimension for spectrum mapping. Default is 3

  • spectrum_n_layers (int, optional) – Number of hidden layers in spectrum mapping. Default is 2

  • spectrum_kernel_size (int, optional) – Kernel size for Conv1d spectrum mapping. Default is 5

  • stiefel_parametrization_mode (str, optional) – Parametrization mode for orthogonal matrix Q. Default is “static”. Choices are: “static” and “dynamic”

  • stiefel_n_steps_ref_update (int, optional) – Steps between reference point updates for dynamic parametrization. Default is 100

  • use_autograd (bool, optional) – Use torch autograd for gradient computation. Default is False

  • device (torch.device, optional) – Device. Default is torch.device(“cpu”)

  • dtype (torch.dtype, optional) – Data type. Default is torch.float64

  • generator (torch.Generator | None, optional) – Generator for reproducibility. Default is None

Initialize internal Module state, shared by both nn.Module and ScriptModule.

forward(data)[source]

Compute tangent vector from SPD matrix spectrum.

Parameters:

data (torch.Tensor of shape (..., n, n)) – Batch of SPD matrices

Returns:

tangent (torch.Tensor of shape (..., n, n)) – Symmetric tangent vectors at each input point

Return type:

Tensor

register_optimizer_hook(optimizer)[source]

Register optimizer hook for dynamic Stiefel parametrization.

affine_invariant_norm(base, tangent)[source]

Affine-invariant norm of tangent vectors.

\(\lVert X^{-1/2} V X^{-1/2} \rVert_F = \lVert L^{-1} V L^{-\top} \rVert_F\) with \(X = LL^\top\): Cholesky and triangular solves only, so the gradient does not go through an eigendecomposition.

Parameters:

  • base (torch.Tensor of shape (..., n, n)) – SPD base points

  • tangent (torch.Tensor of shape (..., n, n)) – Symmetric tangent vectors at base

Returns:

norm (torch.Tensor of shape (...)) – Affine-invariant norms

Return type:

Tensor