EEG-DINO

<div align="center"> <h1>EEG-DINO: Learning EEG Foundation Models via Hierarchical Self-Distillation</h1> <div align="center"> <img src="assets/eeg-dino.png" alt="positions" width="500"/> </div>

We propose EEG-DINO, a novel foundation model for EEG encoding based on a hierarchical self-distillation framework. By multi-view semantic alignment, the model is able to extract multi-level semantic features from EEG data, which captures a wide range of semantic information, increasing the robustness against noise and variances inherent in complex EEG signals. Moreover, acknowledging the unique heterogeneous spatial-temporal dependencies in EEG signals, we design a channel-aware sampling mechanism and a decoupled positional coding scheme. They independently address spatial and temporal dimensions, enabling the model to capture the intricate structural characteristics of EEG signals. We pre-train EEG-DINO on a large-scale EEG corpus spanning over 9000 hours, which consistently achieves state-of-the-art performance on multiple downstream tasks. These results demonstrate the great effectiveness of our self-distillation framework for EEG encoding.

Pre-trained Models

Usage

CUDA_VISIBLE_DEVICES=0 python /path/to/run_finetuning.py

The default settings are for EEG-DINO-Small, if you want to use medium or large, you could change the embedding model in /path/to/models/eeg_encoder.py:

from models.embedding_small import PatchEmbedding

and change the default settings in /path/to/run_finetuning.py:

    parser.add_argument('--feature_size', default=200, type=int)
    parser.add_argument('--num_layers', default=12, type=int)
    parser.add_argument('--dim_feedforward', default=512, type=int)

512/16/1024 for medium and 1024/24/2048 for large.

Evaluation Results

We evaluate the performance of EEG-DINO on multiple downstream tasks, including TUEV, SEED-V and TUAB. The results consistently demonstrate the effectiveness of our model.

TUEV

Linear Probing:

Full-parameter fine-tuning:

SEED-V

Linear Probing:

Full-parameter fine-tuning:

TUAB

Linear Probing:

Full-parameter fine-tuning: