SPD Learn Pipeline and Trivialization#
This page connects the theory to what the package actually implements, aligned with the library paper. For a longer walkthrough, see User Guide.
SPDNet Building Blocks#
SPD Learn is centered on the SPDNet pipeline:
Input SPD -> BiMap -> ReEig -> BiMap -> ReEig -> LogEig -> Linear -> Output
Key layers:
Following the original paper, SPD Learn defines key architectural components as neural network layers specifically designed to handle and operate on Riemannian geometries. These layers ensure that the geometric structure of the data is preserved or appropriately transformed throughout the network:
BiMap: Stiefel-constrained bilinear mapping for SPD-preserving dimension change
ReEig: eigenvalue rectification (nonlinearity)
LogEig: maps SPD matrices to the tangent space for Euclidean classifiers
These layers are exposed as BiMap,
ReEig, and LogEig.
The model zoo builds on these blocks and includes SPDNet-based architectures
for neural decoding, such as TensorCSPNet,
TSMNet, MAtt,
Green, EEGSPDNet,
and PhaseSPDNet.
Trivialization in SPD Learn#
The paper and the code use trivialization-based parametrizations [Lezcano-Casado, 2019], to keep manifold-valued parameters valid during training.
Trivialization maps points on a manifold to vectors in a Euclidean space via a diffeomorphism (smooth, invertible map):
This allows optimizing on the manifold using standard gradient descent in the Euclidean parameterization:
In SPD Learn:
Stiefel parameters (BiMap weights) are represented in Euclidean space and mapped to the Stiefel manifold through polar decomposition: \(\phi^{-1}(\mathbf{X}) = \mathbf{X}(\mathbf{X}^T\mathbf{X})^{-1/2}\)
SPD parameters (e.g., batch norm bias) are represented in Euclidean space and mapped to SPD using matrix exponential trivialization.
This allows standard optimizers (SGD, Adam) to train models without explicit Riemannian solvers, while preserving manifold constraints by construction.
Correspondence to CNNs:
CNN Component |
SPDNet Analog |
Function |
|---|---|---|
Conv layer |
BiMap |
Feature extraction/dimension change |
ReLU |
ReEig |
Nonlinearity |
Flatten |
LogEig + Vech |
Trivialization to vector |
FC layer |
Linear |
Classification |
Batch Normalization#
SPD Learn provides SPDBatchNormMean and
SPDBatchNormMeanVar. These layers normalize SPD-valued
features while preserving geometric structure and are central to domain
adaptation models such as TSMNet.
Example: Building an SPDNet#
import torch
import torch.nn as nn
from spd_learn.modules import BiMap, ReEig, LogEig, SPDBatchNormMean
class ManualSPDNet(nn.Module):
"""SPDNet built from individual layers."""
def __init__(self, n_channels=32, n_classes=4):
super().__init__()
# SPD dimension reduction pipeline
self.bimap1 = BiMap(n_channels, n_channels // 2)
self.reeig1 = ReEig()
self.bn1 = SPDBatchNormMean(n_channels // 2)
self.bimap2 = BiMap(n_channels // 2, n_channels // 4)
self.reeig2 = ReEig()
# Trivialization to tangent space
self.logeig = LogEig(upper=True, flatten=True)
# Euclidean classifier
out_dim = (n_channels // 4) * (n_channels // 4 + 1) // 2
self.classifier = nn.Linear(out_dim, n_classes)
def forward(self, x):
# x: (batch, n_channels, n_channels) SPD matrices
x = self.reeig1(self.bimap1(x))
x = self.bn1(x)
x = self.reeig2(self.bimap2(x))
x = self.logeig(x) # Trivialization
return self.classifier(x)
# Usage
model = ManualSPDNet(n_channels=32, n_classes=4)
spd_batch = torch.randn(16, 32, 32)
spd_batch = spd_batch @ spd_batch.mT + 0.1 * torch.eye(32)
output = model(spd_batch)
print(f"Output shape: {output.shape}") # (16, 4)
Or use the pre-built SPDNet model directly.
Where to Go Next#
User Guide for a step-by-step model walkthrough
API Reference for full layer and function references