Large Brain Model for Learning Generic Representations with Tremendous EEG Data in BCI

Wei-Bang Jiang, Li-Ming Zhao, Bao-Liang Lu

Seminal May 2024

TL;DR LaBraM is an EEG foundation model, pre-trained by masked prediction of discrete neural tokens on about 2,500 hours of EEG from around 20 datasets with differing electrode layouts. Fine-tuned, it beats all compared baselines on abnormal detection, event classification, emotion recognition and gait prediction, e.g. 0.64 vs 0.53 balanced accuracy on TUEV with the 5.8M-parameter Base model.

What they do

EEG deep learning models are usually built for one dataset and task, because recordings differ in electrode count, length and task design, and labeled sets are small. LaBraM (Large Brain Model) is a single Transformer pre-trained without labels on 2,534.78 hours of EEG from about 20 public and self-collected datasets, then fine-tuned for each downstream task.

How it works

Channel patches. Each EEG channel is cut into 1-second patches (200 samples at 200 Hz). A temporal convolutional encoder embeds each patch, and learnable temporal and spatial (electrode) embeddings are added, so one model accepts any channel set and length, up to 256 patches.
Neural tokenizer. A vector-quantized tokenizer maps patches to a codebook of 8192 codes. It is trained to regress the Fourier amplitude and phase of each patch, since directly reconstructing raw EEG did not converge.
Masked EEG modeling. Half of the patches are masked and the Transformer predicts their tokens; the inverted (symmetric) mask is also used, reusing the same tokens.
Three sizes. Base (5.8M parameters), Large (46M) and Huge (369M).

Results

TUAB abnormal detection: AUROC 0.9022 (Base) to 0.9162 (Huge), against 0.8815 for the best baseline, BIOT.
TUEV event classification: balanced accuracy 0.6409 (Base) to 0.6616 (Huge), against 0.5281 for BIOT.
SEED-V emotion recognition: accuracy 0.4095 to 0.4102, against 0.3802 for BIOT. MoBI gait prediction: the Base model is competitive with the best baseline (R2 0.2876 vs 0.2911), and Large and Huge exceed it (R2 up to 0.3145).
Training from scratch without pre-training causes a steep performance drop. On TUAB the Base model with 500 hours of pre-training data already exceeds its 2,500-hour result, while the Huge model keeps improving as data grows.

Limitations

The data and model are still far smaller than those of large vision and language models. LaBraM needs full fine-tuning for each downstream task, which is costly in computation and memory; linear probing is much worse (0.3461 balanced accuracy on TUEV). It is trained on EEG only, with no other modalities. The tokenizer reconstructs phase less well than amplitude.

Figures

From Figure 2: Top, a neural tokenizer is trained to turn EEG patches into discrete codebook tokens by reconstructing... — From Figure 2: Top, a neural tokenizer is trained to turn EEG patches into discrete codebook tokens by reconstructing the Fourier amplitude and phase. Bottom, during pre-training part of the EEG patches are masked (along with the inverted, symmetric mask) and the Transformer predicts the tokens of the masked patches.

From Figure 1: The LaBraM architecture. EEG signals are cut into fixed-length patches per channel, a temporal encoder... — From Figure 1: The LaBraM architecture. EEG signals are cut into fixed-length patches per channel, a temporal encoder embeds each patch, temporal and spatial embeddings are added, and the sequence goes through a Transformer encoder.

From Figure 5: Performance of the Base, Large and Huge models on TUAB and TUEV as the amount of pre-training data inc... — From Figure 5: Performance of the Base, Large and Huge models on TUAB and TUEV as the amount of pre-training data increases.

From Figure 8: LaBraM on TUAB and TUEV with pre-training versus trained from scratch on the downstream data.

Abstract

The current electroencephalogram (EEG) based deep learning models are typically designed for specific datasets and applications in brain-computer interaction (BCI), limiting the scale of the models and thus diminishing their perceptual capabilities and generalizability. Recently, Large Language Models (LLMs) have achieved unprecedented success in text processing, prompting us to explore the capabilities of Large EEG Models (LEMs). We hope that LEMs can break through the limitations of different task types of EEG datasets, and obtain universal perceptual capabilities of EEG signals through unsupervised pre-training. Then the models can be fine-tuned for different downstream tasks. However, compared to text data, the volume of EEG datasets is generally small and the format varies widely. For example, there can be mismatched numbers of electrodes, unequal length data samples, varied task designs, and low signal-to-noise ratio. To overcome these challenges, we propose a unified foundation model for EEG called Large Brain Model (LaBraM). LaBraM enables cross-dataset learning by segmenting the EEG signals into EEG channel patches. Vector-quantized neural spectrum prediction is used to train a semantically rich neural tokenizer that encodes continuous raw EEG channel patches into compact neural codes. We then pre-train neural Transformers by predicting the original neural codes for the masked EEG channel patches. The LaBraMs were pre-trained on about 2,500 hours of various types of EEG signals from around 20 datasets and validated on multiple different types of downstream tasks. Experiments on abnormal detection, event type classification, emotion recognition, and gait prediction show that our LaBraM outperforms all compared SOTA methods in their respective fields. Our code is available at https://github.com/935963004/LaBraM.

Summaries are written automatically from each paper's text. Figures and videos come from the paper and its project page; all credit goes to the authors.

Matt Strong