spd_learn.models.SymSPDAttention#

class spd_learn.models.SymSPDAttention(in_features: int, out_features: int, metric: str = 'LEM', device: device | None = None, dtype: dtype | None = None)[source]#

Bases: Module

Symmetric Tangent Space Self-Attention Layer for SPD Matrices.

Implements the tangent-space attention mechanism for SPD matrices. Uses three SPDLinear projections with return_tangent=True to map patch SPD tokens directly into symmetric tangent matrices \(Q, K, V \in \mathcal{S}^{d_{\text{out}}}\), computes scaled Frobenius distance attention:

\[A_{ij} = \operatorname{softmax}_j \left( -e^\alpha \lVert Q_i - K_j \rVert_F - \beta \right)\]

aggregates features in tangent space:

\[\bar{V}_i = \sum_j A_{ij} V_j\]

and maps back to the SPD manifold via matrix exponential \(\exp(\bar{V}_i)\).

Parameters:

  • in_features (int) – Dimension of the input SPD matrices (\(d_{\text{in}}\)).

  • out_features (int) – Dimension of the output SPD matrices (\(d_{\text{out}}\)).

  • metric ({"LEM", "LCM", "AIM", "PEM"}, default="LEM") – Riemannian metric for the SPDLinear projections.

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

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

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 SymSPDAttention
>>> att = SymSPDAttention(in_features=4, out_features=3, 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, 3, 3])
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 SymSPDAttention.

Parameters:

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

Returns:

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

Return type:

torch.Tensor