SAMBA: Toward an EEG Foundation Model via Differential Mamba and Coordinate-Driven Spatial Embedding

Update: Jan. 2026

SAMBA is a scalable self-supervised framework for long-sequence EEG foundation modeling. It integrates 3D coordinate-driven spatial embedding and differential Mamba modules to enable robust, efficient, and generalizable EEG representation learning across diverse recording configurations and cognitive tasks.

---

Abstract

Modeling long EEG sequences is critical for developing generalizable neural representations, particularly due to the high temporal resolution and extended durations often required to capture brain dynamics. While transformer-based models have shown success on short EEG segments, their quadratic complexity prevents effective scaling to long contexts. Additionally, the diversity in EEG montages and subject variability presents significant generalization challenges.

We introduce SAMBA, a self-supervised learning framework featuring a U-shaped encoder-decoder architecture built on the linear-time Mamba module. SAMBA incorporates:

1. Temporal Semantic Random Masking — to reconstruct semantically masked segments in long sequences;
2. Multi-Head Differential Mamba — to reduce redundancy and enhance salient temporal features;
3. Coordinate-Driven Spatial Embedding (CDSE) — to learn robust 3D spatial representations across heterogeneous EEG devices.

Evaluations across 13 EEG datasets covering a range of tasks, montages, and sequence lengths demonstrate that SAMBA consistently outperforms state-of-the-art baselines, while maintaining low memory usage and fast inference. In addition, the learned spatial weights exhibit strong alignment with task-relevant neurophysiological regions, suggesting that SAMBA is both learnable and interpretable.

---

Learnable Spatial Embedding Visualization

The figure below illustrates the Coordinate-Driven Spatial Embedding (CDSE) module and its alignment with neurophysiological topology:

Eyes Close/Open Emphasizes the frontal lobe, consistent with alpha modulation during eye closure[1].	Driver Distraction Highlights the left temporal lobe, associated with auditory processing (conversation during driving) and cognitive control under distraction[2].
Workload Estimation Focuses on the left frontal region, linked to stress-induced workload processing[3].	Abnormal Detection (Seizure) Shows occipital dominance, consistent with findings that occipital electrodes yield better performance for distinguishing abnormal EEG[4].
Motor Imagery Localizes around the motor cortex, reflecting sensorimotor activity during imagery tasks[5].

References

_{[1] S. Hoffmann et al. The correction of eye blink artefacts in the EEG: a comparison of two prominent methods. PLOS ONE, 3(8):e3004, 2008.
[2] G. Li et al. Drivers’ EEG responses to different distraction tasks. Auto. Innov., 6(1):20–31, 2023.
[3] G. Berretz et al. Acute stress increases left hemispheric activity measured via changes in frontal alpha asymmetries. iScience, 25(2), 2022.
[4] S. Lopez et al. Automated identification of abnormal adult EEGs. In Proc. IEEE SPMB, pp. 1–5, 2015.
[5] J. Hong et al. A deep learning framework based on dynamic channel selection for early classification of left and right hand motor imagery tasks. In Proc. IEEE EMBC, pp. 3550–3553, 2022.}

Model and Checkpoints

• The full architecture is implemented under Models/.
• Two pretrained checkpoints are provided under Checkpoints/:
  • SAMBA-E: trained using PyTorch Lightning
  • SAMBA-T: trained using native PyTorch
• The accompanying paper is currently in preparation. More features and documentation will be released soon.

---

Repository Structure

SAMBA/
├── Env-requirement/     # Environment configs with dated backups
├── Checkpoints/         # Pretrained SAMBA models (SAMBA-E, SAMBA-T)
├── Models/              # Model architecture implementations
├── Experiments/         # PyTorch Lightning training modules
├── utility/             # Supporting functions: data loading, masking, loss, evaluation
├── Figures/             # Diagrams and visualizations
├── Montage/             # EEG montage metadata for multiple devices
└── README.md            # Project overview