[论文解读] 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 Interfaces被引用 6

一句话总结

EEGFormer 通过对1.7TB的大规模无标签EEG数据进行离散向量量化 Transformer 的预训练，以学习通用、可迁移的表示，具有可解释的码本令牌，并展示了强大的下游表现和对新生儿癫痫检测的可迁移性。

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.

研究动机与目标

• 从大规模无标签数据中激发学习通用 EEG 表示，而不是数据集特定的预训练。
• 为 EEG 信号开发一个离散表示（向量量化）预训练框架。
• 研究 EEG 基础模型在多种下游任务和数据集上的可迁移性。
• 通过分析学习到的离散码本和令牌，提供可解释的洞见。
• 在基于 TUH 的下游任务上评估性能并迁移到新生儿癫痫检测。

提出的方法

• 将 EEG 表示为多变量时间序列并分割成补丁。
• 使用每个通道的 Transformer 编码器对补丁进行编码，并使用位置嵌入。
• 用向量量化码本对编码器输出进行量化，以获得离散令牌。
• 用一个浅层 Transformer 解码器对离散令牌进行解码以重建输入（重建损失加上 VQ 惩罚）。
• 用一个联合目标进行训练，结合重建误差和向量量化损失（码本使用与承诺项）。
• 对预训练的编码器/解码器（以及可选的码本）进行下游任务的微调。

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.

实验结果

研究问题

• RQ1 在大规模 EEG 数据上进行向量量化、基于 Transformer 的预训练，是否能产生可跨越多种 EEG 任务的通用表示？
• RQ2 与端到端方法相比，离散表示学习是否能提高 EEG 模型的可解释性？
• RQ3 EEGFormer 在 TUH 之外的数据集（如新生儿癫痫检测）上的迁移能力有多强？

主要发现

Model	Pretrain	Metric	TUAB	TUAR	TUSL	TUSZ	Neonate
EEGNet	✗	AUROC	0.841 ± .011	0.752 ± .006	0.635 ± .015	0.820 ± .030	0.793 ± .019
EEGNet	✗	AUPRC	0.832 ± .011	0.433 ± .025	0.351 ± .006	0.470 ± .017	0.499 ± .044
TCN	✗	AUROC	0.841 ± .004	0.687 ± .011	0.545 ± .009	0.817 ± .004	0.731 ± .020
TCN	✗	AUPRC	0.831 ± .002	0.408 ± .009	0.344 ± .001	0.383 ± .010	0.398 ± .025
EEG-GNN	✗	AUROC	0.840 ± .005	0.837 ± .022	0.721 ± .009	0.780 ± .006	0.760 ± .010
EEG-GNN	✗	AUPRC	0.832 ± .004	0.488 ± .015	0.381 ± .004	0.388 ± .023	0.419 ± .021
GraphS4mer	✗	AUROC	0.864 ± .006	0.833 ± .006	0.632 ± .017	0.822 ± .034	0.719 ± .007
GraphS4mer	✗	AUPRC	0.862 ± .008	0.461 ± .024	0.359 ± .001	0.491 ± .001	0.374 ± .013
BrainBERT	✓	AUROC	0.853 ± .002	0.753 ± .012	0.588 ± .013	0.814 ± .009	0.734 ± .019
BrainBERT	✓	AUPRC	0.846 ± .003	0.350 ± .014	0.352 ± .003	0.386 ± .018	0.398 ± .027
EEGFormer s	✓	AUROC	0.862 ± .007	0.847 ± .013	0.683 ± .018	0.875 ± .004	0.842 ± .008
EEGFormer s	✓	AUPRC	0.862 ± .005	0.488 ± .012	0.397 ± .011	0.553 ± .014	0.578 ± .023
EEGFormer b	✓	AUROC	0.865 ± .001	0.847 ± .014	0.713 ± .010	0.878 ± .006	0.842 ± .014
EEGFormer b	✓	AUPRC	0.867 ± .002	0.483 ± .026	0.393 ± .003	0.560 ± .010	0.568 ± .036
EEGFormer l	✓	AUROC	0.876 ± .003	0.852 ± .004	0.679 ± .013	0.883 ± .005	0.833 ± .017
EEGFormer l	✓	AUPRC	0.872 ± .001	0.483 ± .014	0.389 ± .003	0.556 ± .008	0.544 ± .026

• EEGFormer 的变体在 TUH 下游任务的 AUROC 和 AUPRC 指标上优于若干基线。
• 预训练的轮数对下游性能有正向影响。
• 使用 EEGFormer 进行线性探针已经能与部分有监督基线相媲美，端到端微调可取得最佳结果。
• 离散码本实现了可解释性，对学习到的索引进行 n-gram 分析即可在不微调的情况下定位癫痫模式。
• 在 Neonate 和 TUSZ 数据集上，与基线相比，EEGFormer 在 AUPRC 和 AUROC 上取得显著提升。

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