Note
Go to the end to download the full example code.
Manifold Attention with MAtt#
This tutorial demonstrates how to use MAtt (Manifold Attention Network) for EEG classification. MAtt applies attention mechanisms on the SPD manifold to weight temporal segments by their discriminative importance.
Introduction#
MAtt [Pan et al., 2022] introduces attention mechanisms on the SPD manifold:
Patch-based Processing: Segments the signal into temporal patches
Covariance per Patch: Computes SPD matrices for each segment
Manifold Attention: Weights patches using Log-Euclidean distances
Aggregation: Combines weighted SPD matrices for classification
This allows the model to focus on the most discriminative time periods within each trial, improving classification and interpretability.
Setup and Imports#
import warnings
import matplotlib.pyplot as plt
import torch
from braindecode import EEGClassifier
from moabb.datasets import BNCI2014_001
from moabb.paradigms import MotorImagery
from sklearn.metrics import accuracy_score
from sklearn.preprocessing import LabelEncoder
from skorch.callbacks import EpochScoring, GradientNormClipping
from skorch.dataset import ValidSplit
from spd_learn.models import MAtt
warnings.filterwarnings("ignore")
/opt/hostedtoolcache/Python/3.11.16/x64/lib/python3.11/site-packages/braindecode/models/eegpt.py:497: FutureWarning: Montage name 'standard_1020' is deprecated and will be removed in MNE 1.14. Use 'colin27_1020' instead.
montage = mne.channels.make_standard_montage("standard_1020")
/opt/hostedtoolcache/Python/3.11.16/x64/lib/python3.11/site-packages/braindecode/models/eegpt.py:1452: FutureWarning: Montage name 'standard_1020' is deprecated and will be removed in MNE 1.14. Use 'colin27_1020' instead.
montage = make_standard_montage("standard_1020")
Loading the Dataset#
dataset = BNCI2014_001()
paradigm = MotorImagery(n_classes=4)
print(f"Dataset: {dataset.code}")
print("Paradigm: 4-class motor imagery")
Choosing from all possible events
Dataset: BNCI2014-001
Paradigm: 4-class motor imagery
Creating the MAtt Model#
MAtt architecture:
Spatial Conv: Learns spatial filters
Temporal Conv: Extracts temporal features
Patch Embedding: Segments into n_patches temporal windows
Covariance + TraceNorm: SPD matrix per patch
AttentionManifold: Queries, keys, values on SPD manifold
ReEig + LogEig: Project to tangent space
Linear: Classification
The attention mechanism computes:
\[\text{attention}(Q, K) = \text{softmax}\left(\frac{1}{1 + \log(1 + d_{LE}(Q, K))}\right)\]
where \(d_{LE}\) is the Log-Euclidean distance.
n_chans = 22
n_outputs = 4
model = MAtt(
n_chans=n_chans,
n_outputs=n_outputs,
n_patches=6, # Number of temporal segments
temporal_out_channels=32, # Temporal feature dimension
temporal_kernel_size=25, # ~100ms at 250Hz
temporal_padding=12, # Keep time dimension
attention_in_features=32, # Input to attention (must match temporal_out_channels)
attention_out_features=24, # Output from attention
)
print("MAtt Architecture:")
print(model)
MAtt Architecture:
MAtt(
(add_extra_dim): Rearrange('batch chan time -> batch 1 chan time')
(spatial_filter): Conv2d(1, 22, kernel_size=(22, 1), stride=(1, 1))
(spatial_batch_norm): BatchNorm2d(22, eps=1e-05, momentum=0.1, affine=True, bias=True, track_running_stats=True)
(temporal_feature_extractor): Conv2d(22, 32, kernel_size=(1, 25), stride=(1, 1), padding=(0, 12))
(temporal_batch_norm): BatchNorm2d(32, eps=1e-05, momentum=0.1, affine=True, bias=True, track_running_stats=True)
(squeeze_channel): Rearrange('batch kernel 1 time -> batch kernel time')
(patch_cov_layer): Sequential(
(0): PatchEmbeddingLayer()
(1): CovLayer()
(2): TraceNorm()
)
(manifold_attention): AttentionManifold(
(q_trans): ParametrizedBiMap(
(parametrizations): ModuleDict(
(weight): ParametrizationList(
(0): _Orthogonal()
)
)
)
(k_trans): ParametrizedBiMap(
(parametrizations): ModuleDict(
(weight): ParametrizationList(
(0): _Orthogonal()
)
)
)
(v_trans): ParametrizedBiMap(
(parametrizations): ModuleDict(
(weight): ParametrizationList(
(0): _Orthogonal()
)
)
)
)
(re_eig): ReEig()
(tangent): LogEig()
(flatten): Flatten(start_dim=1, end_dim=-1)
(linear): Linear(in_features=1800, out_features=4, bias=True)
)
Training the Model#
subject_id = 1
batch_size = 32
max_epochs = 100
learning_rate = 1e-4 # Low learning rate for stable SPD learning
device = "cuda" if torch.cuda.is_available() else "cpu"
print(f"\nUsing device: {device}")
# Cache configuration
cache_config = dict(
save_raw=True,
save_epochs=True,
save_array=True,
use=True,
overwrite_raw=False,
overwrite_epochs=False,
overwrite_array=False,
)
# Load data
X, labels, meta = paradigm.get_data(
dataset=dataset, subjects=[subject_id], cache_config=cache_config
)
# Encode labels
le = LabelEncoder()
y = le.fit_transform(labels)
# Split by session
train_idx = meta.query("session == '0train'").index.to_numpy()
test_idx = meta.query("session == '1test'").index.to_numpy()
print(f"\nData shape: {X.shape}")
print(f"Training samples: {len(train_idx)}")
print(f"Test samples: {len(test_idx)}")
# Create classifier
# Note: SPD networks benefit from gradient clipping to prevent
# divergence during training on the Riemannian manifold.
clf = EEGClassifier(
model,
criterion=torch.nn.CrossEntropyLoss,
optimizer=torch.optim.Adam,
optimizer__lr=learning_rate,
train_split=ValidSplit(0.1, stratified=True, random_state=42),
batch_size=batch_size,
max_epochs=max_epochs,
callbacks=[
(
"train_acc",
EpochScoring(
"accuracy", lower_is_better=False, on_train=True, name="train_acc"
),
),
("gradient_clip", GradientNormClipping(gradient_clip_value=1.0)),
],
device=device,
verbose=1,
)
# Train
clf.fit(X[train_idx], y[train_idx])
# Evaluate
y_pred_train = clf.predict(X[train_idx])
y_pred_test = clf.predict(X[test_idx])
train_acc = accuracy_score(y[train_idx], y_pred_train)
test_acc = accuracy_score(y[test_idx], y_pred_test)
print(f"\n{'=' * 50}")
print(f"Results for Subject {subject_id}")
print(f"{'=' * 50}")
print(f"Train Accuracy: {train_acc * 100:.2f}%")
print(f"Test Accuracy: {test_acc * 100:.2f}%")
Using device: cpu
Data shape: (576, 22, 1001)
Training samples: 288
Test samples: 288
epoch train_acc train_loss valid_acc valid_loss dur
------- ----------- ------------ ----------- ------------ ------
1 0.2461 1.3908 0.2759 1.3608 2.3506
2 0.3867 1.3498 0.5517 1.3224 2.2597
3 0.4648 1.3118 0.4138 1.2988 2.3442
4 0.4023 1.2892 0.2414 1.2891 2.3222
5 0.5156 1.2551 0.4483 1.2521 2.3294
6 0.5977 1.2283 0.5517 1.2242 2.3202
7 0.5781 1.2035 0.4828 1.2104 2.3271
8 0.6250 1.1705 0.5517 1.1713 2.3137
9 0.6602 1.1476 0.6207 1.1635 2.3284
10 0.6602 1.1230 0.5517 1.1401 1.8245
11 0.7031 1.0946 0.5862 1.1027 1.6366
12 0.7109 1.0622 0.7241 1.0889 1.6485
13 0.7461 1.0365 0.6552 1.0693 1.6777
14 0.7578 1.0176 0.6897 1.0403 1.6723
15 0.7266 0.9922 0.6207 1.0236 1.6731
16 0.7734 0.9635 0.6552 1.0041 1.6730
17 0.7461 0.9437 0.7241 0.9813 1.6574
18 0.7852 0.9137 0.6552 0.9603 1.6997
19 0.8203 0.8833 0.7931 0.9403 1.6689
20 0.7617 0.8771 0.7241 0.9354 1.6474
21 0.7734 0.8555 0.6207 0.9099 1.7158
22 0.7812 0.8305 0.7241 0.8973 1.6422
23 0.7734 0.8161 0.6552 0.8866 1.6869
24 0.8320 0.7889 0.7241 0.8740 1.6656
25 0.8242 0.7832 0.7931 0.8593 1.6852
26 0.8242 0.7726 0.7931 0.8604 1.6691
27 0.8516 0.7480 0.7931 0.8293 1.6624
28 0.8320 0.7408 0.8276 0.8416 1.6550
29 0.8398 0.7273 0.8966 0.8110 1.6526
30 0.8359 0.7246 0.6552 0.8554 1.6721
31 0.8438 0.7092 0.7586 0.7943 2.0326
32 0.8281 0.6882 0.7931 0.8080 2.2946
33 0.8320 0.6907 0.6552 0.8033 2.2711
34 0.8516 0.6737 0.8621 0.7725 2.3053
35 0.8320 0.6598 0.7931 0.8007 2.3334
36 0.8320 0.6478 0.7931 0.7621 2.3245
37 0.8711 0.6333 0.8621 0.7626 2.3250
38 0.8633 0.6275 0.8276 0.7601 2.3201
39 0.8477 0.6233 0.7586 0.7531 2.3272
40 0.8516 0.6079 0.8621 0.7480 2.3306
41 0.8477 0.6062 0.7586 0.7451 2.3242
42 0.8750 0.5974 0.7931 0.7367 2.3238
43 0.8828 0.5779 0.7586 0.7295 2.3406
44 0.8828 0.5711 0.8276 0.7357 2.3249
45 0.8750 0.5632 0.7931 0.7027 2.3260
46 0.8867 0.5516 0.8621 0.7196 2.3334
47 0.8750 0.5484 0.8966 0.6972 2.3313
48 0.8945 0.5375 0.7586 0.7110 2.3428
49 0.8516 0.5445 0.8621 0.6903 2.2463
50 0.9023 0.5120 0.7586 0.7106 2.4289
51 0.9102 0.5101 0.8621 0.6743 2.4022
52 0.8789 0.5068 0.8966 0.6759 2.2351
53 0.9102 0.4923 0.8276 0.6804 2.2396
54 0.8867 0.5042 0.8276 0.6629 1.8379
55 0.9102 0.4972 0.7241 0.6963 1.8019
56 0.8867 0.4916 0.6207 0.6779 2.0346
57 0.8945 0.4832 0.7586 0.6851 1.9128
58 0.8867 0.4804 0.8966 0.6595 2.3253
59 0.9141 0.4791 0.7931 0.6588 2.3545
60 0.8945 0.4603 0.7931 0.6838 2.4402
61 0.9180 0.4505 0.8276 0.6442 2.4334
62 0.9414 0.4428 0.8966 0.6477 2.4294
63 0.9062 0.4451 0.8966 0.6318 2.4378
64 0.9375 0.4268 0.8621 0.6398 2.4450
65 0.8984 0.4350 0.7931 0.6534 2.3282
66 0.9297 0.4245 0.8966 0.6169 2.3946
67 0.9297 0.4186 0.8276 0.6464 1.8354
68 0.9375 0.4057 0.8276 0.6261 1.7481
69 0.9414 0.4109 0.8621 0.6301 2.3717
70 0.9375 0.4052 0.8621 0.6128 1.9061
71 0.9336 0.3917 0.8276 0.6316 1.6791
72 0.9336 0.3925 0.8621 0.6000 1.7487
73 0.9336 0.3892 0.8621 0.6118 2.0997
74 0.9609 0.3762 0.8966 0.5997 2.3183
75 0.9492 0.3743 0.7931 0.6045 2.3353
76 0.9492 0.3678 0.7586 0.6333 2.3579
77 0.9375 0.3637 0.8621 0.5898 2.4186
78 0.9414 0.3576 0.8276 0.6124 2.4803
79 0.9570 0.3558 0.8276 0.6038 2.3829
80 0.9531 0.3477 0.8276 0.6008 2.4074
81 0.9727 0.3436 0.8276 0.5943 2.4000
82 0.9688 0.3366 0.9310 0.5877 2.3913
83 0.9727 0.3360 0.8276 0.6075 2.4221
84 0.9570 0.3342 0.9310 0.5659 2.4253
85 0.9492 0.3364 0.8276 0.5867 2.4066
86 0.9609 0.3186 0.8621 0.5810 2.4167
87 0.9727 0.3126 0.8621 0.5683 2.4082
88 0.9766 0.3143 0.7586 0.6167 2.4258
89 0.9727 0.3179 0.8621 0.5595 2.3969
90 0.9648 0.3063 0.8276 0.6031 2.4219
91 0.9766 0.3088 0.8621 0.5725 2.4034
92 0.9492 0.3110 0.8276 0.5662 2.4106
93 0.9766 0.2975 0.8966 0.5734 2.3469
94 0.9844 0.2941 0.7931 0.5915 2.4069
95 0.9727 0.2873 0.7241 0.6126 2.3946
96 0.9727 0.2860 0.8276 0.5701 2.4101
97 0.9727 0.2773 0.7241 0.5979 2.4771
98 0.9688 0.2842 0.8276 0.5531 2.5228
99 0.9766 0.2784 0.8276 0.5603 2.4232
100 0.9844 0.2691 0.8276 0.5625 2.3492
==================================================
Results for Subject 1
==================================================
Train Accuracy: 97.57%
Test Accuracy: 80.21%
Visualizing Results#
fig, axes = plt.subplots(1, 2, figsize=(12, 4))
# Training history
history = clf.history
epochs = range(1, len(history) + 1)
ax1 = axes[0]
ax1.plot(epochs, history[:, "train_loss"], "b-", label="Train Loss", linewidth=2)
ax1.plot(epochs, history[:, "valid_loss"], "r--", label="Valid Loss", linewidth=2)
ax1.set_xlabel("Epoch", fontsize=12)
ax1.set_ylabel("Loss", fontsize=12)
ax1.set_title("Training and Validation Loss", fontsize=14)
ax1.legend(fontsize=10)
ax1.grid(True, alpha=0.3)
ax2 = axes[1]
ax2.plot(epochs, history[:, "train_acc"], "b-", label="Train Acc", linewidth=2)
ax2.plot(epochs, history[:, "valid_acc"], "r--", label="Valid Acc", linewidth=2)
ax2.set_xlabel("Epoch", fontsize=12)
ax2.set_ylabel("Accuracy", fontsize=12)
ax2.set_title("Training and Validation Accuracy", fontsize=14)
ax2.legend(fontsize=10)
ax2.grid(True, alpha=0.3)
ax2.set_ylim([0, 1])
plt.tight_layout()
plt.show()
Understanding Manifold Attention#
The attention mechanism in MAtt operates differently from standard attention:
Standard Attention (Euclidean):
\[\text{Attention}(Q, K, V) = \text{softmax}\left(\frac{QK^T}{\sqrt{d}}\right) V\]
Manifold Attention (Log-Euclidean):
\[\text{energy}_{ij} = d_{LE}(Q_i, K_j) = \|\log(Q_i) - \log(K_j)\|_F\]
\[\text{weights}_{ij} = \frac{1}{1 + \log(1 + \text{energy}_{ij})}\]
\[\text{output}_i = \sum_j \text{softmax}(\text{weights})_{ij} \odot V_j\]
This respects the Riemannian geometry of SPD matrices, computing meaningful distances on the manifold rather than in Euclidean space.
Summary#
In this tutorial, we demonstrated:
Creating a MAtt model with patch-based temporal segmentation
Training for motor imagery classification
Understanding the manifold attention mechanism
MAtt is particularly useful when:
Different time segments have varying discriminative power
You want interpretable attention weights
The data has complex temporal dynamics
Total running time of the script: (3 minutes 43.197 seconds)
