spd_learn.models.SPDNetAE#

class spd_learn.models.SPDNetAE(n_chans: int = 22, n_outputs: int = 2, bimap_dims: Sequence[int] = (16, 8), eps: float = 0.0001, drop_prob: float = 0.0, input_type: str = 'raw', cov_method=<function covariance>, threshold: float = 0.0001, upper: bool = True, device=None, dtype=None)[source]#

Bases: Module

Riemannian Autoencoder for SPD Matrices (SPDNet-AE).

This class implements the Riemannian Autoencoder model proposed by Boucherie, de Surrel, and Yger [Boucherie et al., 2026]. SPDNet-AE learns compact latent SPD representations while preserving the intrinsic Riemannian geometry of Symmetric Positive Definite matrices throughout both the encoding and decoding pathways.

Architecture Overview

  1. Covariance Frontend: When raw signal data is provided, it is mapped to an SPD covariance matrix via CovLayer(), followed by scale normalization TraceNorm() and shrinkage regularization Shrinkage(n_chans, init_shrinkage=-2.0).

  2. SPD Encoder: Stacks BiMap and ReEig blocks across (n_chans, *bimap_dims) to produce a compact latent SPD representation \(Z \in \mathcal{S}_{++}^{d_{\text{latent}}}\).

  3. SPD Decoder: Reconstructs the original \((n_{\text{chans}} \times n_{\text{chans}})\) SPD matrix from \(Z\) using transposed bilinear congruence blocks:

    \[S_{\text{recon}} = W Z W^\top + \epsilon I\]

    with \(W \in \text{St}(d_{\text{out}}, d_{\text{in}})\) followed by ReEig, guaranteeing that the reconstructed matrix strictly remains on the SPD manifold \(\mathcal{S}_{++}^{n_{\text{chans}}}\).

  4. Reconstruction Loss: Evaluates the Log-Euclidean geodesic distance between the input covariance and reconstructed covariance:

    \[\mathcal{L}_{\text{recon}}(S_{\text{in}}, S_{\text{recon}}) = \|\log(S_{\text{in}}) - \log(S_{\text{recon}})\|_F^2\]

  5. Classification: When used in a classification pipeline (e.g. with braindecode.EEGClassifier), the latent representation \(Z\) is projected to Euclidean space via LogEig, followed by dropout and a linear classification head.

Parameters:

  • n_chans (int, default=22) – Number of EEG channels (dimension of the input covariance matrix).

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

  • bimap_dims (tuple[int, ...], default=(16, 8)) – Sequence of subspace dimensions for the SPD encoder layers. The latent SPD dimension is bimap_dims[-1].

  • eps (float, default=1e-4) – Eigenvalue rectification threshold and diagonal regularization parameter.

  • drop_prob (float, default=0.0) – Dropout probability applied to the vectorized tangent features before the final linear classification layer.

  • input_type ({"raw", "cov"}, default="raw") – Format of input data (‘raw’ time series or precomputed ‘cov’ matrices).

  • cov_method (callable, default=covariance) – Method to compute covariance when input_type is ‘raw’.

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

  • upper (bool, default=True) – If True, only upper-triangular elements are used in tangent projection.

References

[1]

Charlotte Boucherie, Thibault de Surrel, and Florian Yger. SPDNet-AE: a compact SPD representation through Riemannian autoencoding. In European Symposium on Artificial Neural Networks, Computational Intelligence and Machine Learning (ESANN). 2026. hal-05600912v1.

Examples

>>> import torch
>>> from spd_learn.models import SPDNetAE
>>> model = SPDNetAE(n_chans=22, n_outputs=2, bimap_dims=(16, 8))
>>> x = torch.randn(4, 22, 1000)  # (batch, channels, times)
>>> logits = model(x)
>>> logits.shape
torch.Size([4, 2])
>>> Z = model.encode(x)
>>> Z.shape
torch.Size([4, 8, 8])
>>> S_recon = model.decode(Z)
>>> S_recon.shape
torch.Size([4, 22, 22])
>>> S_in = model.cov_frontend(x)
>>> loss = model.reconstruction_loss(S_in, S_recon)
>>> loss.item() >= 0.0
True
autoencode(x: Tensor) → Tensor[source]#

Encode and reconstruct input to produce reconstructed SPD matrices.

Parameters:

x (torch.Tensor) – Input tensor of raw signals or SPD covariance matrices.

Returns:

Reconstructed SPD matrices with shape (…, n_chans, n_chans).

Return type:

torch.Tensor

decode(Z_spd: Tensor) → Tensor[source]#

Reconstruct full-dimensional SPD matrices from latent SPD matrices.

Parameters:

Z_spd (torch.Tensor) – Latent SPD matrices with shape (…, d_latent, d_latent).

Returns:

Reconstructed SPD matrices with shape (…, n_chans, n_chans).

Return type:

torch.Tensor

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

Encode input signals or covariance matrices into latent SPD representations.

Parameters:

x (torch.Tensor) – Input tensor. If raw signal data, expected shape is (batch, n_chans, n_times). If already an SPD covariance matrix, expected shape is (…, n_chans, n_chans).

Returns:

Latent SPD matrix with shape (…, d_latent, d_latent) where d_latent = bimap_dims[-1].

Return type:

torch.Tensor

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

Forward pass of the SPDNet-AE model as a classifier.

Parameters:

x (torch.Tensor) – Input tensor of raw signals with shape (batch, n_chans, n_times) or covariance matrices with shape (batch, n_chans, n_chans).

Returns:

Class logits with shape (batch, n_outputs).

Return type:

torch.Tensor

reconstruction_loss(S_in: Tensor, S_recon: Tensor, reduction: str = 'mean') → Tensor[source]#

Compute the Log-Euclidean geodesic reconstruction loss.

\[\mathcal{L}_{\text{recon}}(S_{\text{in}}, S_{\text{recon}}) = \|\log(S_{\text{in}}) - \log(S_{\text{recon}})\|_F^2\]
Parameters:

  • S_in (torch.Tensor) – Input SPD matrices with shape (…, n_chans, n_chans).

  • S_recon (torch.Tensor) – Reconstructed SPD matrices with shape (…, n_chans, n_chans).

  • reduction ({"mean", "sum", "none"}, default="mean") – Reduction applied to the loss across the batch.

Returns:

Reconstruction loss.

Return type:

torch.Tensor