spd_learn.models.CorAtt#

class spd_learn.models.CorAtt(n_patches: int = 6, n_chans: int = 22, n_outputs: 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, metric: str = 'mix', drop_prob: float = 0.0, covariance_method=<function sample_covariance>, spatial_filters: int | None = None, temporal_filters: int | None = None, device: device | None = None, dtype: dtype | None = None)[source]#

Bases: MAtt

Correlation Attention Network for EEG Decoding (CorAtt).

This class implements the CorAtt model [Hu et al., 2025], a Riemannian self-attention network for EEG decoding operating on the correlation manifold.

The model consists of:

  1. Spatial and temporal 2D convolution feature frontend (inherited from MAtt).

  2. Segmenting features into temporal patches and computing patch covariance.

  3. Normalizing to full-rank correlation matrices and applying CorrAttention.

  4. Tangent projection (log_map_lsm() or log_map_olm()).

  5. tril_embed() isometrically vectorizing hollow tangent matrices.

  6. Dropout and Linear classification head.

Parameters:

  • n_patches (int, default=6) – Number of temporal patches.

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

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

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

  • temporal_kernel_size (int, default=25) – Kernel size for the temporal convolution.

  • temporal_padding (int, optional) – Padding along time for the temporal convolution. Defaults to temporal_kernel_size // 2.

  • attention_in_features (int, default=20) – Input feature dimension for the correlation attention module.

  • attention_out_features (int, default=20) – Output feature dimension for the correlation attention module.

  • metric ({"mix", "OLM", "LSM"}, default="mix") – Riemannian metric for the correlation attention module and tangent projection.

  • drop_prob (float, default=0.0) – Dropout probability before the linear classifier.

  • covariance_method (callable, default=sample_covariance) – Method used for computing sample covariances.

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

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

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

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

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

Forward pass of CorAtt.

Parameters:

x (torch.Tensor) – Input tensor of raw EEG signals (batch_size, n_chans, n_times) or (batch_size, 1, n_chans, n_times).

Returns:

Output class logits of shape (batch_size, n_outputs).

Return type:

torch.Tensor