[Paper Review] EEGFormer: Towards Transferable and Interpretable Large-Scale EEG Foundation Model

Yuqi Chen, Kan Ren+5 more|arXiv (Cornell University)|Jan 11, 2024

EEG and Brain-Computer Interfaces6 citations

TL;DR

EEGFormer pretrains a discrete-vector-quantized Transformer on 1.7TB of large-scale unlabeled EEG data to learn universal, transferable representations with interpretable codebook tokens, and demonstrates strong downstream performance and transferability to neonatal seizure detection.

ABSTRACT

Self-supervised learning has emerged as a highly effective approach in the fields of natural language processing and computer vision. It is also applicable to brain signals such as electroencephalography (EEG) data, given the abundance of available unlabeled data that exist in a wide spectrum of real-world medical applications ranging from seizure detection to wave analysis. The existing works leveraging self-supervised learning on EEG modeling mainly focus on pretraining upon each individual dataset corresponding to a single downstream task, which cannot leverage the power of abundant data, and they may derive sub-optimal solutions with a lack of generalization. Moreover, these methods rely on end-to-end model learning which is not easy for humans to understand. In this paper, we present a novel EEG foundation model, namely EEGFormer, pretrained on large-scale compound EEG data. The pretrained model cannot only learn universal representations on EEG signals with adaptable performance on various downstream tasks but also provide interpretable outcomes of the useful patterns within the data. To validate the effectiveness of our model, we extensively evaluate it on various downstream tasks and assess the performance under different transfer settings. Furthermore, we demonstrate how the learned model exhibits transferable anomaly detection performance and provides valuable interpretability of the acquired patterns via self-supervised learning.

Motivation & Objective

• Motivate learning universal EEG representations from massive unlabeled data rather than dataset-specific pretraining.
• Develop a discrete representation (vector-quantized) pretraining framework for EEG signals.
• Investigate transferability of the EEG foundation model to multiple downstream tasks and datasets.
• Provide interpretable insights by analyzing the learned discrete codebook and tokens.
• Evaluate performance on TUH-based downstream tasks and transfer to neonatal seizure detection.

Proposed method

• Represent EEG as multivariate time series and segment into patches.
• Encode patches with a Transformer encoder per channel, using position embeddings.
• Quantize encoder outputs with a vector-quantized codebook to obtain discrete tokens.
• Decode discrete tokens with a shallow Transformer decoder to reconstruct inputs (reconstruction loss plus VQ penalties).
• Train with a joint objective combining reconstruction error and vector-quantization losses (codebook usage and commitment terms).
• Fine-tune the pretrained encoder/decoder (and optionally the codebook) for downstream tasks.

Figure 1: Overview of EEGFormer . Initially, multi-variate EEG signals are segmented into patches, which are then passed through a Transformer encoder. Subsequently, a vector-quantized model is employed to generate discrete indices. These indices are then fed into a shallow Transformer decoder.

Experimental results

Research questions

• RQ1 Can a vector-quantized, Transformer-based pretraining on large-scale EEG data yield universal representations transferable across multiple EEG tasks?
• RQ2 Does discrete representation learning improve interpretability of EEG models compared to end-to-end methods?
• RQ3 How well does EEGFormer transfer to datasets outside the TUH corpus (e.g., neonatal seizure detection)?

Key findings

• EEGFormer variants outperform several baselines on TUH-downstream tasks in AUROC and AUPRC metrics.
• Pretraining epoch count positively affects downstream performance.
• Linear probing with EEGFormer already rivals some supervised baselines, with end-to-end fine-tuning yielding the best results.
• The discrete codebook enables interpretability, and n-gram analysis on learned indices can localize seizure patterns without fine-tuning.
• On Neonate and TUSZ datasets, EEGFormer achieves substantial gains in AUPRC and AUROC compared to baselines.

Figure 2: Influence of pretrain epochs on two TUH corpus.