spd_learn.models.SPDMetaFormer#

class spd_learn.models.SPDMetaFormer(n_patches: int = 6, n_chans: int = 22, n_outputs: int = 2, depth: int = 2, temporal_out_channels: int = 20, temporal_kernel_size: int = 25, temporal_padding: int | None = None, attention_in_features: int = 20, attention_out_features: int = 20, ffn_dim: int | None = None, eps: float = 0.0001, drop_prob: float = 0.0, covariance_method=<function sample_covariance>, init_shrinkage: float = -2.0, pooling: str = 'mean', spatial_filters: int | None = None, temporal_filters: int | None = None, temp_kernel_length: int | None = None, device: device | None = None, dtype: dtype | None = None)[source]#

Bases: MAtt

SPD-MetaFormer for Small-Data Brain Decoding.

This class implements the SPD-MetaFormer model [Wang et al., 2026], an attention-free Riemannian architecture designed for small-data brain decoding from electroencephalography (EEG) signals.

In standard manifold attention models such as MAtt [Pan et al., 2022], the bounded similarity metric \(\frac{1}{1 + \log(1 + d)}\) causes softmax attention weights over short EEG patch sequences to remain nearly uniform, while query/key bilinear projections over-parameterize low-data regimes. SPD-MetaFormer replaces learned attention with an attention-free Fréchet token mixer and shared spectral feedforward maps equipped with Log-Euclidean geodesic skip connections.

The model consists of three stages:

  1. Frontend & Patch Covariance Extraction: Raw EEG signals are processed by spatial and temporal 2D convolutions (inherited from MAtt), segmented into \(P\) temporal patches via PatchEmbeddingLayer, and mapped to SPD covariance matrices regularized by TraceNorm and learnable Shrinkage.

  2. SPD-MetaFormer Backbone: A stack of \(\text{depth}\) SPDMetaFormerBlock layers that perform attention-free token mixing via uniform Fréchet averaging and channel mixing via a shared SPD spectral feedforward network (SPD-FFN).

  3. Learned Geometric Readout: Patch tokens are aggregated via a learnable softmax-weighted Log-Euclidean tangent mean \(\bar{L}_{\text{readout}} = \sum_{p=1}^P w_p \log(Y_p)\), vectorized via Vech, and classified using a linear projection with dropout.

Parameters:

  • n_patches (int, default=6) – Number of temporal patches extracted from the signal.

  • n_chans (int, default=22) – Number of input EEG channels (electrodes).

  • n_outputs (int, default=2) – Number of output classes for classification.

  • depth (int, default=2) – Number of stacked SPDMetaFormerBlock layers.

  • temporal_out_channels (int, default=20) – Number of temporal convolution filters.

  • temporal_kernel_size (int, default=25) – Kernel length of the temporal convolution layer.

  • temporal_padding (int, optional) – Padding along time for temporal convolution.

  • attention_in_features (int, default=20) – Input feature dimension for SPD tokens.

  • attention_out_features (int, default=20) – Output feature dimension for SPD tokens.

  • ffn_dim (int, optional) – Hidden dimension of the SPD feedforward network within each block.

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

  • drop_prob (float, default=0.0) – Dropout probability applied before the final classification head.

  • covariance_method (callable, default=sample_covariance) – Covariance estimation method.

  • init_shrinkage (float, default=-2.0) – Initial shrinkage pre-sigmoid logit for covariance regularization.

  • pooling ({"mean", "concat"}, default="mean") – Token readout pooling method in the tangent space.

  • spatial_filters (int, optional) – Alias for temporal_out_channels.

  • temporal_filters (int, optional) – Alias for intermediate feature dimension.

  • temp_kernel_length (int, optional) – Alias for temporal_kernel_size.

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

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

References

[1]

Yue-Ting Pan, Jing-Lun Chou, and Chun-Shu Wei. Matt: a manifold attention network for eeg decoding. In Advances in Neural Information Processing Systems, volume 35, 31116–31129. 2022. URL: https://proceedings.neurips.cc/paper_files/paper/2022/hash/c981fd12b1d5703f19bd8289da9fc996-Abstract-Conference.html.

[2]

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

SPDMetaFormerBlock

Attention-free SPD MetaFormer building block.

MAtt

Manifold attention network for EEG decoding.

Examples

>>> import torch
>>> from spd_learn.models import SPDMetaFormer
>>> model = SPDMetaFormer(n_chans=22, n_outputs=2, n_patches=6, depth=2)
>>> x = torch.randn(2, 22, 1000)
>>> out = model(x)
>>> out.shape
torch.Size([2, 2])
forward(x: Tensor) → Tensor[source]#

Forward pass of the SPD-MetaFormer model.

Parameters:

x (torch.Tensor) – Input EEG tensor of shape (batch, n_chans, n_times) or (batch, 1, n_chans, n_times).

Returns:

Output class logits of shape (batch, n_outputs).

Return type:

torch.Tensor