spd_learn.models.SPDMetaFormerBlock#

class spd_learn.models.SPDMetaFormerBlock(dim: int, ffn_dim: int | None = None, eps: float = 0.0001, device: device | None = None, dtype: dtype | None = None)[source]#

Bases: Module

SPD MetaFormer Block.

This module implements an attention-free MetaFormer block operating on sequences of Symmetric Positive Definite (SPD) patch tokens [Wang et al., 2026].

The block replaces computationally heavy and over-parameterized manifold attention mechanisms with two sub-blocks:

  1. Uniform Fréchet Token Mixer with Geodesic Residual: Tokens are mapped to the Log-Euclidean tangent space via matrix logarithm \(L = \log(X)\). The uniform Log-Euclidean Fréchet mean across patches is computed as \(\bar{L} = \frac{1}{P} \sum_{p=1}^P L_p\). Tokens are updated via learnable geodesic interpolation with rate \(\alpha = \sigma(\text{mix\_logit})\):

    \[U = \exp((1 - \alpha) L + \alpha \bar{L})\]

  2. Shared Spectral Feedforward Map (SPD-FFN) with Geodesic Residual: A shared 2-layer SPD spectral feedforward network processes each token:

    \[\text{FFN}(U) = W_2^\top \text{ReEig}(W_1^\top U W_1) W_2\]

    where \(W_1 \in \text{St}(d, d_{\text{ffn}})\) and \(W_2 \in \text{St}(d_{\text{ffn}}, d)\) are (semi-)orthogonal bilinear mappings. The output is combined with \(U\) via a second learnable Log-Euclidean geodesic skip connection with rate \(\beta = \sigma(\text{ffn\_logit})\):

    \[Y = \exp((1 - \beta) \log(U) + \beta \log(\text{FFN}(U)))\]
Parameters:

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

  • ffn_dim (int, optional) – Hidden dimensionality \(d_{\text{ffn}}\) of the spectral feedforward network. If None, defaults to dim.

  • eps (float, default=1e-4) – Eigenvalue rectification threshold for ReEig.

  • device (torch.device, optional) – Target device for layer parameters.

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

References

[1]

Zhida Wang, Wei Lyu, Guo Yu, and Sui Tang. SPD-MetaFormer is what you need for small-data brain decoding. arXiv preprint arXiv:2610.10952, 2026. URL: https://arxiv.org/abs/2610.10952.

See also

SPDMetaFormer

Full SPD-MetaFormer model for EEG decoding.

BiMap

Bilinear mapping layer for SPD matrices.

ReEig

Rectified eigenvalue layer.

Examples

>>> import torch
>>> from spd_learn.models import SPDMetaFormerBlock
>>> block = SPDMetaFormerBlock(dim=8, ffn_dim=12)
>>> A = torch.randn(2, 6, 8, 8)
>>> X = A @ A.mT + 0.1 * torch.eye(8)  # SPD tokens
>>> Y = block(X)
>>> Y.shape
torch.Size([2, 6, 8, 8])
forward(X: Tensor) → Tensor[source]#

Forward pass through the SPD MetaFormer block.

Parameters:

X (torch.Tensor) – Sequence of patch SPD tokens of shape (…, P, dim, dim).

Returns:

Transformed SPD tokens of shape (…, P, dim, dim).

Return type:

torch.Tensor