CBraMod: A Criss-Cross Brain Foundation Model for EEG Decoding

Jiquan Wang; Sha Zhao; Zhiling Luo; Yangxuan Zhou; Haiteng Jiang; Shijian Li; Tao Li; Gang Pan

Source-linked AI summary

CBraMod: A Criss-Cross Brain Foundation Model for EEG Decoding

Jiquan Wang, Sha Zhao, Zhiling Luo, Yangxuan Zhou, Haiteng Jiang, Shijian Li, Tao Li, Gang Pan

arXiv:2412.07236v6Dec 10, 2024eess.SPcs.AIcs.LGq-bio.NC

TL;DR

Existing EEG decoding and foundation models have limited generalizability and do not fully account for heterogeneous spatial and temporal dependencies. CBraMod combines criss-cross attention, adaptive positional encoding, and masked EEG reconstruction, achieving state-of-the-art performance across up to 10 downstream BCI tasks and 12 public datasets.

Problem
Existing EEG models have limited generalizability across downstream BCI tasks and often model spatial and temporal dependencies together despite their heterogeneous structure.
Method
CBraMod uses a criss-cross transformer with parallel spatial and temporal attention, asymmetric conditional positional encoding, and patch-based masked EEG reconstruction pre-training.
Results
CBraMod achieves state-of-the-art performance across up to 10 downstream BCI tasks using 12 public datasets.
Takeaways & Limitations
The results support CBraMod’s capability and generalizability across a broad range of EEG BCI tasks.
Takeaways & Limitations
The study uses the open-source SEED-V version with 16 subjects rather than the private 20-subject version used in LaBraM.

Abstract

from arXiv · show

Electroencephalography (EEG) is a non-invasive technique to measure and record brain electrical activity, widely used in various BCI and healthcare applications. Early EEG decoding methods rely on supervised learning, limited by specific tasks and datasets, hindering model performance and generalizability. With the success of large language models, there is a growing body of studies focusing on EEG foundation models. However, these studies still leave challenges: Firstly, most of existing EEG foundation models employ full EEG modeling strategy. It models the spatial and temporal dependencies between all EEG patches together, but ignores that the spatial and temporal dependencies are heterogeneous due to the unique structural characteristics of EEG signals. Secondly, existing EEG foundation models have limited generalizability on a wide range of downstream BCI tasks due to varying formats of EEG data, making it challenging to adapt to. To address these challenges, we propose a novel foundation model called CBraMod. Specifically, we devise a criss-cross transformer as the backbone to thoroughly leverage the structural characteristics of EEG signals, which can model spatial and temporal dependencies separately through two parallel attention mechanisms. And we utilize an asymmetric conditional positional encoding scheme which can encode positional information of EEG patches and be easily adapted to the EEG with diverse formats. CBraMod is pre-trained on a very large corpus of EEG through patch-based masked EEG reconstruction. We evaluate CBraMod on up to 10 downstream BCI tasks (12 public datasets). CBraMod achieves the state-of-the-art performance across the wide range of tasks, proving its strong capability and generalizability. The source code is publicly available at https://github.com/wjq-learning/CBraMod.

1 INTRODUCTION

EEG supports non-invasive, real-time brain–computer interfacing across applications, but existing decoding and foundation-model approaches face limits in generalizability and in modeling heterogeneous spatial-temporal dependencies. CBraMod addresses these challenges with criss-cross modeling, adaptive positional encoding, and broad downstream evaluation.

EEG provides a non-invasive, real-time measure of brain activity for decoding intentions or commands in BCI systems.
Traditional and deep EEG decoders have commonly been developed for specific tasks and datasets, limiting general-purpose applicability.
EEG foundation models pre-train on large EEG corpora and fine-tune on clinical or BCI applications, commonly using channel-based patches.
CBraMod uses a criss-cross transformer to model spatial and temporal dependencies separately through parallel attention mechanisms.
Asymmetric conditional positional encoding dynamically represents patch positions and adapts CBraMod to EEG datasets with different channel formats.
CBraMod is pre-trained with patch-based masked EEG reconstruction and evaluated across up to 10 downstream BCI tasks using 12 public datasets.

2 METHOD

CBraMod processes EEG through patch encoding, asymmetric positional encoding, criss-cross attention, and masked reconstruction. Its design separately models spatial and temporal dependencies while accommodating varied EEG channel configurations and sequence lengths.

Masked EEG Reconstruction: CBraMod is pre-trained with patch-based masked EEG reconstruction, replacing randomly selected patches with a mask token and minimizing MSE only on masked patches.The framework patches and masks EEG, learns representations through criss-cross transformer blocks, and reconstructs the masked originals.
Patching & Masking: EEG samples are divided into fixed-length temporal patches across channels, producing C×n patches for signals with diverse channel counts and time lengths.A sample S ∈ R^(C×T) is segmented into patches X ∈ R^(C×n×t), where n is the number of patches per channel.
Time-Frequency Patch Encoding: Each masked patch is encoded using parallel time-domain convolution and frequency-domain FFT branches, whose embeddings are added.The time branch extracts within-patch temporal features, while the frequency branch derives an energy vector with FFT before projection.
Asymmetric Conditional Positional Encoding: ACPE dynamically encodes spatial and temporal patch positions with an asymmetric design: longer-range spatial information and shorter-range temporal information.A convolutional positional encoder generates positional embeddings from spatial-temporal neighborhoods and adds them to patch embeddings.
Criss-Cross Transformer: Criss-cross attention uses parallel spatial and temporal attention to capture dependencies within the same time interval and within the same channel.Attention heads are divided between S-Attention and T-Attention, then their outputs are concatenated in each criss-cross transformer block.

3 EXPERIMENTS

Experiments evaluate CBraMod across diverse downstream BCI settings, compare it with established baselines, and ablate its pre-training, attention, and positional-encoding choices.

Experiment Setup: CBraMod is pre-trained on TUEG, a diverse corpus containing 69,652 recordings from 14,987 subjects across 26,846 sessions.TUEG contains 27,062 hours of EEG and over 40 channel configurations.
Experiment Setup: Downstream evaluation standardizes EEG to 200 Hz with one-second patches and covers up to 10 BCI tasks across 12 public datasets.The experiments use consistent training, validation, and test splits across methods.
Experiment Setup: The evaluation compares CBraMod with non-foundation-model and foundation-model baselines using classification and regression metrics reported over five random seeds.Metrics include Balanced Accuracy, AUROC, Cohen’s Kappa, Weighted F1, Pearson’s Correlation, R2 Score, and RMSE.
Results: CBraMod achieves state-of-the-art performance on both FACED and SEED-V emotion-recognition datasets.FACED and SEED-V differ in channel counts, sampling rates, subjects, and emotion categories.
Ablation Studies: Positional encoding performance ranks ACPE above CPE and APE, while no positional encoding performs worst.The comparison supports asymmetric positional designs for EEG’s differing spatial and temporal dependencies.
Ablation Studies: The pre-training ablation finds clean pre-training best, with significant gains and smaller variance than dirty pre-training or training without pre-training.Dirty pre-training performs only slightly better than no pre-training, suggesting bad samples weaken pre-training effectiveness.

4 CONCLUSION

CBraMod is an EEG foundation model designed to learn generic EEG representations through masked reconstruction, criss-cross spatial-temporal modeling, and flexible positional encoding. It achieves state-of-the-art performance across up to 10 downstream BCI tasks and 12 public datasets.

CBraMod learns generic EEG representations through patch-based masked EEG reconstruction on a large EEG corpus.
Its criss-cross transformer models spatial and temporal dependencies in parallel, while asymmetric convolutional positional encoding handles diverse EEG formats.
CBraMod achieves state-of-the-art performance across up to 10 downstream BCI tasks and 12 public datasets.
Existing EEG foundation models have limited generalizability across downstream tasks, while earlier decoding methods often depend on handcrafted features or task-specific learning.

B MORE DETAILS FOR EXPERIMENTAL SETTINGS ON PRE-TRAINING

The pre-training setup uses 30-second EEG segments and documents the associated hyperparameters. The longer segment duration is selected to support learning long-term dependencies and downstream performance.

CBraMod pre-training uses EEG recordings segmented into 30-second samples.
The 30-second duration is longer than the 10-second BIOT and 4- or 8-second LaBraM pre-training samples.
The authors choose longer segments because they may help the model learn long-term dependencies and potentially improve downstream performance.

C PRE-TRAINING VISUALIZATION

CBraMod’s pre-training loss generally decreases over 40 epochs, while downstream evaluations cover diverse BCI and clinical tasks. Reported results show state-of-the-art performance on several datasets, with additional controls addressing pre-training overlap.

C PRE-TRAINING VISUALIZATION: The pre-training loss generally decreases across 40 epochs, with minor fluctuations between epochs 10 and 14.
Downstream evaluations: The downstream evaluation spans emotion recognition, motor imagery, sleep staging, seizure detection, imagined speech, diagnosis, vigilance, stress, events, and abnormal detection.
Reported results: CBraMod achieves state-of-the-art performance on seizure detection, mental disorder diagnosis, mental stress detection, event classification, and abnormal detection.
Reported results: For imagined speech classification, CBraMod reaches 0.4216 in Cohen’s Kappa versus 0.3800 for LaBraM.
Reported results: For mental disorder diagnosis, CBraMod reaches 0.9560 versus 0.9409 in balanced accuracy against LaBraM.
Evaluation controls: Additional CBraMod variants excluding TUEV or TUAB from pre-training were evaluated to mitigate potential data leakage effects.

E.9 MOTOR IMAGERY CLASSIFICATION

CBraMod is evaluated on motor imagery classification using the BCIC-IV-2a dataset, which contains four imagined movement classes from nine subjects. Across the paper’s downstream tasks, the model shows strong capability and generalizability.

E.9 MOTOR IMAGERY CLASSIFICATION: The BCIC-IV-2a dataset contains four motor imagery tasks recorded from nine subjects using 22 electrodes.
E.9 MOTOR IMAGERY CLASSIFICATION: The tasks involve imagining left-hand, right-hand, both-feet, and tongue movements.
E.9 MOTOR IMAGERY CLASSIFICATION: Results across the wide range of downstream BCI tasks indicate that CBraMod learns generic EEG representations with strong capability and generalizability.

F SCALING DATA SIZE AND MODEL SIZE

The study examines how CBraMod performance changes with pretraining data and model size. Performance improves with both scales, though data-scaling gains slow beyond 1000 hours.

Model scale: Larger CBraMod models lead to better performance on downstream tasks across variants ranging from 0.1M to 4M parameters.The model-size experiments evaluate multiple variants within this parameter range.
Evaluation scope: Including the additional BCIC-IV-2a experiment, the study evaluated 13 downstream datasets rather than the 12 counted in the main text.The authors excluded BCIC-IV-2a from the main-text count for narrative clarity.

G ARCHITECTURE COMPARISON ON OUR PRE-TRAINING DATASETS

The architecture comparison evaluates BIOT and LaBraM after pretraining on CBraMod’s dataset under matched settings. CBraMod outperforms LaBraM under this comparison, while architecture differences contribute more to improvement than pretraining-dataset size.

Architecture comparison: CBraMod performs significantly better than LaBraM retrained on the same pretraining dataset across all evaluated datasets.LaBraM retrained on the authors’ dataset generally slightly exceeds original LaBraM, especially on FACED.
Architecture comparison: BIOT retrained on the authors’ dataset and original BIOT exhibit very similar performance.The comparison uses the BIOT and LaBraM architectures with CBraMod’s pretraining dataset and settings.
Interpretation: The authors conclude that model architecture contributes more to performance improvement than the larger pretraining dataset.

I ABLATION STUDY ON TIME-DOMAIN AND FREQUENCY-DOMAIN SIGNALS

The ablation study compares time-domain signals, frequency-domain signals, and their combination. Combining both signal domains performs better than using either domain alone.

Signal-domain ablation: Combining time-domain and frequency-domain signals achieves better performance than using only time-domain or only frequency-domain features.The ablation study evaluates the three signal-input settings in Table 17.

J ABLATION STUDY ON FINE-TUNING

The fine-tuning ablation tests fully fine-tuned and fixed-parameter settings for CBraMod and competing models. Fixing pretrained parameters causes a large performance decline, although fixed CBraMod remains stronger than fixed BIOT and LaBraM.

Fine-tuning impact: Fixing CBraMod’s pretrained parameters and training only the classifier causes a very large performance decline.The fully fine-tuned setting adjusts all CBraMod parameters on downstream datasets.
Limitation: CBraMod currently requires fine-tuning and cannot serve as a fixed-parameter feature extractor like CLIP or SAM.
Fixed-parameter comparison: In the fixed-parameter setting, CBraMod outperforms fixed BIOT and fixed LaBraM on unseen datasets.

K LOW-RESOURCE COMPARISON WITH EXISTING METHODS

CBraMod is evaluated in low-resource fine-tuning with 30% of labeled data and compared with other foundation models. The section also examines how splitting criss-cross attention heads affects performance.

Low-resource comparison: 30% of labeled data is used to compare CBraMod, BIOT, and LaBraM in low-resource fine-tuning.The comparison is reported in Table 19.
Low-resource comparison: CBraMod consistently outperforms existing foundation models across the evaluated low-resource datasets.Its low-resource performance is only marginally lower than full-data performance on SEED-V and SHU-MI.
Attention split analysis: The study also evaluates how the criss-cross attention head split influences performance.The criss-cross transformer uses equal spatial and temporal head allocation in its intermediate layer.

M MASK RATIO ANALYSIS

The analyses identify a generally effective mask ratio, compare implementation and computational choices, and examine learned representations, convergence, and model limitations. Together, these experiments document CBraMod’s behavior across downstream evaluation and deployment-related settings.

Mask ratio analysis: Mask ratios from 0.3 to 0.7 usually yield better CBraMod performance across downstream datasets.On FACED, the best mask ratio is 0.5, while SHU-MI performs best at 0.4 and nearly as well at 0.5.
Mask ratio analysis: A 0.5 mask ratio is selected as an appropriate choice for good performance across multiple downstream datasets.The mask-ratio comparison is presented in Figure 9.
Mask-token analysis: The mask-token comparison reports no significant performance difference between full-zero and learnable mask tokens.The two masking options use, respectively, a zero vector or a learnable vector matching the patch-embedding dimension.
Computational comparison: CBraMod has higher parameter count and computational complexity than non-foundation-model baselines, despite lower complexity than other EEG foundation models.The comparison uses the CHB-MIT dataset with 16 channels and 10 seconds of EEG.
Representation analysis: Criss-cross visualizations show spatial-temporal correlation patterns, with stronger associations involving the central patch in deeper transformer layers.The analysis uses 19-channel, 30-second TUEG samples and indicates learning of both criss-cross and non-criss-cross dependencies.
Downstream evaluation: CBraMod outperforms EEG-SimpleConv under both evaluated LOSO settings and converges faster than EEGConformer on downstream tasks.The comparison reports higher balanced accuracy, Cohen’s Kappa, and weighted F1 for CBraMod; it reaches convergence within 10 epochs.
Limitations: The authors identify reduced clean pre-training data, deployment difficulty on low-power devices, and unexplored billion-scale EEG pre-training scaling laws as limitations.These constraints arise from crude filtering of dirty TUEG data, model size and computational demands, and limited computational resources.