spd_learn.models.GyroSPDAttention#

class spd_learn.models.GyroSPDAttention(dim: int, metric: str = 'LEM', power: float = 1.0, device: device | None = None, dtype: dtype | None = None)[source]#

Bases: Module

Gyrovector Self-Attention Layer for SPD Matrices.

Implements the general gyrovector attention mechanism on the SPD manifold introduced in GyroAtt [Wang et al., 2025].

Query, Key, and Value representations are parameterized via orthogonal gyro-automorphisms \(O \in \mathrm{O}(d)\) generated through the Cayley transform of skew-symmetric matrices:

\[S = W - W^\top, \quad O = (I - S)(I + S)^{-1}\]

Each patch SPD token \(X_i \in \mathcal{S}_{++}^d\) is transformed as:

\[Q_i = O_Q X_i O_Q^\top, \quad K_i = O_K X_i O_K^\top, \quad V_i = O_V X_i O_V^\top\]

The pairwise manifold distance \(D_{ij}\) between patch tokens is computed under either the Log-Euclidean Metric (LEM) or the Log-Cholesky Metric (LCM):

\[D_{ij} = \lVert \log(Q_i) - \log(K_j) \rVert_F\]

Attention weights are computed via logarithmic decay:

\[A_{ij} = \operatorname{softmax}_j \left( \frac{1}{1 + \log(1 + D_{ij})} \right)\]

The aggregated output is the weighted Fréchet mean in tangent/Cholesky space, followed by a learnable Riemannian left gyro-translation bias \(B \in \mathcal{S}_{++}^d\):

\[\bar{V}_i = \exp \left( \sum_j A_{ij} \log(V_j) \right), \quad Y_i = B^{1/2} \bar{V}_i B^{1/2}\]
Parameters:

  • dim (int) – Dimension of the input SPD matrices (\(d \times d\)).

  • metric ({"LEM", "LCM"}, default="LEM") – Riemannian metric for pairwise distance and Fréchet aggregation.

  • power (float, default=1.0) – Power parameter for metric deformation.

  • device (torch.device, optional) – Device on which to allocate parameters.

  • dtype (torch.dtype, optional) – Data type of parameters.

bias#

Learnable SPD bias matrix \(B \in \mathcal{S}_{++}^d\), parameterized via SymmetricPositiveDefinite. Initialized to the identity matrix.

Type:

nn.Parameter

References

[1]

Rui Wang, Chen Hu, Xiaoning Song, Xiaojun Wu, Nicu Sebe, and Ziheng Chen. Towards a general attention framework on gyrovector spaces for matrix manifolds. In Advances in Neural Information Processing Systems, volume 38. 2025. URL: https://openreview.net/forum?id=ChenHu-ML/GyroAtt.

Examples

>>> import torch
>>> from spd_learn.models import GyroSPDAttention
>>> att = GyroSPDAttention(dim=4, metric="LEM")
>>> X = torch.randn(2, 6, 4, 4)
>>> X = X @ X.mT + 0.1 * torch.eye(4)
>>> out = att(X)
>>> out.shape
torch.Size([2, 6, 4, 4])
extra_repr() → str[source]#

Return the extra representation of the module.

To print customized extra information, you should re-implement this method in your own modules. Both single-line and multi-line strings are acceptable.

forward(X: Tensor) → Tensor[source]#

Forward pass of GyroSPDAttention.

Parameters:

X (torch.Tensor) – Input SPD tokens of shape (…, n_patches, dim, dim).

Returns:

Attended SPD tokens of shape (…, n_patches, dim, dim).

Return type:

torch.Tensor

reset_parameters() → None[source]#

Reset layer parameters to initial state.