[論文レビュー] 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.

研究の動機と目的

• 巨大な未ラベルデータから dataset-specific pretraining ではなく普遍的な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コーパス外（例：新生児発作検出）データセットへどれだけ転移するか。

主な発見

モデル	事前学習	指標	TUAB	TUAR	TUSL	TUSZ	新生児
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解析によって微調整なしで発作パターンを局在化できる。
• 新生児データセットおよびTUSZデータセットでは、EEGFormerはベースラインと比較してAUPRCおよびAUROCの顕著な向上を示す。

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