CodeBrain: Bridging Decoupled Tokenizer and Multi-Scale Architecture for EEG Foundation Model

TL;DR

CodeBrain is a novel EEG foundation model that decouples signals into tokens and employs multi-scale architecture to improve interpretability and performance across diverse EEG tasks.

Contribution

It introduces the TFDual-Tokenizer and multi-scale EEGSSM architecture, enhancing EEG representation discriminability and capturing global and local dependencies efficiently.

Findings

Achieves strong generalization across eight downstream tasks and ten datasets.

Enhances interpretability by linking tokens to neural events and spectral rhythms.

Supports distribution shifts with comprehensive ablation and scaling analyses.

Abstract

Electroencephalography (EEG) provides real-time insights into brain activity and supports diverse applications in neuroscience. While EEG foundation models (EFMs) have emerged to address the scalability issues of task-specific models, current approaches still yield clinically uninterpretable and weakly discriminative representations, inefficiently capturing global dependencies and neglecting important local neural events. We present CodeBrain, a two-stage EFM designed to fill this gap. In the first stage, we introduce the TFDual-Tokenizer, which decouples heterogeneous temporal and frequency EEG signals into discrete tokens, quadratically expanding the representation space to enhance discriminative power and offering domain-specific representation-level interpretability by suggesting potential links to neural events and spectral rhythms. In the second stage, we propose the multi-scale EEGSSM architecture, which combines structured global convolution with sliding window attention to efficiently capture both sparse long-range and local dependencies, reflecting the brain's small-world topology. Pretrained on the largest public EEG corpus, CodeBrain achieves strong generalization across eight downstream tasks and ten datasets under distribution shifts, supported by comprehensive ablations, scaling-law analyzes, and interpretability evaluations. The code and the pretrained weights are available at https://github.com/jingyingma01/CodeBrain.