NeurIPS 2024 | EEGPT:通过预训练Transformer实现通用且可靠的脑电图信号表示

本文的标题为《EEGPT: Pretrained Transformer for Universal and Reliable Representation of EEG Signals》,作者团队为哈尔滨工业大学计算机学院李海峰教授团队。该论文提交至NeurIPS 2024并录用,探讨了通用且可靠的脑电图(EEG)信号特征表示的预训练Transformer模型EEGPT(EEG Pretrained Transformer)。

论文链接:https://openreview.net/forum?id=lvS2b8CjG5

论文概要

EEGPT是一种新提出的、具有1000万参数的预训练Transformer模型,旨在为EEG信号的特征提取提供通用而可靠的表示。本文提出用于通用 EEG 特征提取的EEGPT 模型,设计了双自监督学习方法,结合时空表示对齐和掩码重建,提高特征质量和模型收敛性,为 EEG 信号处理提供有效的创新解决方案。与其他自监督学习方法相比,EEGPT通过引入时空表示对齐技术,构建了一个自监督任务,专注于高信噪比和丰富语义信息的EEG表示。这种方法有效缓解了从低信噪比的原始信号提取特征时常见的质量问题。此外,EEGPT的分层结构能够分别处理空间和时间信息,从而降低计算复杂性,提高对脑机接口(BCI)应用的灵活性和适应性。实验结果验证了EEGPT在多种下游任务中的有效性和可扩展性,达到了最先进的性能。

研究背景

EEG动态反映大脑功能状态,广泛应用于情绪识别、运动想象分类等多个领域。目前EEG研究面临数据层面的多样性、模型层面的专业化和预训练及迁移学习方法的局限性,面临的挑战包括低信噪比、高个体间变异性以及EEG信号的任务依赖性变化。这些问题使得从EEG信号中提取稳健的通用表示变得复杂,当前EEG研究多针对特定任务或数据集设计,缺乏通用性和跨任务能力。

自监督学习在自然语言处理、计算机视觉和语音分析等领域展现了显著优势,许多最先进的模型通过自监督学习在大数据集上进行预训练,并针对特定应用进行微调,从而有效降低了对大量标记数据的需求。尽管近年来在EEG分析中应用自监督学习技术取得了良好效果,但现有的掩码自编码器技术难以捕捉EEG数据的时序依赖,限制了迁移学习的效果,尤其是在处理低SNR信号时的特征提取质量。

研究方法

EEGPT提出了一种基于时空一致性的双自监督EEG通用表示方法,主要包括以下几个关键技术要点:

  • 时空表示对齐:通过对齐预测特征与动量编码器的输出,增强编码器提取稳健特征的能力。此方法确保编码器输出包含高质量的全局特征,显著减少计算复杂性。
  • 掩码重构:使用掩码自监督学习任务,EEGPT不仅重构原始波形,还将预测的EEG信号特征与完整的EEG信号特征进行对齐,提高模型的表示质量和鲁棒性。
  • 局部时空嵌入:通过将EEG信号划分为均匀大小的块进行处理,结合通道嵌入信息,提升模型对不同EEG采集设备的兼容性。这种方法有助于提取宏观特征,忽略较小尺度的噪声信号。
  • 线性探测方法:在下游任务中,采用线性探测法,冻结预训练模型的参数,仅调整附加模块的参数,以减少过拟合风险并优化计算资源的使用。

EEGPT的结构如图1所示。模型将EEG信号(M个通道和T个时间点)分为若干个patch,并通过局部时空嵌入将每个patch作为token嵌入,然后分别分为有无掩码的patch使用时空表示对齐和掩码重建。最后,线性探测方法用于应用在下游任务中。

1. 时空表示对齐

掩码自动编码器是一种去噪自动编码器来学习特征,通过将被随机patch掩码遮挡的信号输入到编码器中,解码器预测掩码patch的原始嵌入,如公式1所示:

其中⊙表示逐元素乘积,是掩码块,and 分别是编码器和解码器,是学习表示,是相似性度量。通过最小化损失函数,模型学习输入信号的最佳表示。,在类BERT模型中编码器和解码器没有分割,没有明确表示,本文在EEGPT模型中添加了时空表示对齐分支来显式表达,这也就是双自监督方法,能够鼓励编码表示具有更大程度的语义,提高编码质量和泛化能力,如公式2所示:

时空表示对齐方法将预测特征和动量编码器的输出对齐,增强了编码器提取鲁棒特征的能力,确保编码器的输出包含高质量的全局特征。本文采用编码器和预测器来解耦空间和时间特征,能够降低计算复杂度。

编码器

编码器整合了来自掩码patch的空间信息,公式3描述了编码器(ENC)如何处理时间j处的所有掩码,并产生相应的输出特征。

预测器

如公式4所示,预测器(PRED)利用来自编码器掩码部分的特征,结合时间位置信息,来预测完整的编码特征。本文采用了旋转位置嵌入方法来生成,引入相对位置和时间信息。为了生成属于M的预测特征,使用可学习的向量查询作为查询标记,在自监督训练中鼓励编码器提取token之间相关性信息。

动量编码器

动量编码器与编码器的结构相同,每次训练迭代后,编码器的参数以的系数累积到动量编码器中,如公式5所示。

并使用基于均方误差(MSE)的对齐损失来实现时空表示对齐,LN表示层归一化,有助于减轻极值和协变量偏移的影响,如公式6所示。

2. 掩码重构

基于掩码的重建方法将重建器生成的重建patch与原始的patch进行对齐。

重建器

如公式7所示,重建器(REC)利用编码器编码的M部分的特征和预测器预测的M部分的特征,以及时间位置,生成重建的patch 。

掩码重建时使用基于均方误差的重建损失来实现的,如公式8所示。

最后通过将和相加构建完整的预训练损失。

3. 局部时空嵌入

局部时空嵌入方法首先在时空维度上对EEG信号进行分块和嵌入编码,然后将其输入编码器,如图2所示。首先对EEG信号进行分割,在时空维度上划分为大小相等的块,如公式10所示。

其中d表示patch的时间长度, 是时间patch的数量。此后结合嵌入信息线性地嵌入patch,构造一个集合其中包含所有可学习的通道嵌入向量和从通道名称到通道嵌入向量,使得模型能够适应多个数据集,提高通道适应性。最后得到的token由patch线性嵌入后得到,如公式11所示。

基于自监督学习任务,提取的patch特征是互相可预测的,可以忽略较小尺度的噪声信息。该方法旨在提取跨越更大尺度的宏观特征,这些特征更容易识别且认为是有意义的特征。

4. 线性探测方法

在下游任务中,使用预训练的编码器并连接额外的模块来进行分类。如图3所示,该方法冻结预训练模型中的参数,仅更改附加线性模块中的参数。该模块包括用于对齐EEG和模型之间通道的自适应空间滤波器(1*1卷积),以及将特征映射到网络最后一层的输出logits的线性层,连接多分类头进行多分类任务。

实验设计

本文进行了多项实验,以评估EEGPT在不同EEG任务中的表现。实验数据集包括多种EEG信号来源,如PhysioMI、HGD和M3CV等,每个数据集涵盖了不同的任务和主题。实验评估了EEGPT在运动想象分类、事件相关电位检测和睡眠阶段检测等任务上的性能。

1. 数据集

预训练数据集包括:

  • PhysioMI:涉及运动想象(MI)和运动执行(ME)任务,共109名受试者。
  • HGD:仅涉及运动想象任务,共14名受试者。
  • TSU:稳态视觉诱发范式(SSVEP)数据集,共35名受试者。
  • SEED:情绪三分类数据集,共15名受试者。
  • M3CV:多任务、多受试者数据集,涵盖多种EEG任务,共106名受试者。

下游任务数据集包括:

  • BCIC-2A/B:与运动想象相关的EEG信号,分别包含10名受试者。
  • Sleep-EDFx:睡眠阶段检测数据集,共197名受试者。
  • KaggleERN:使用错误相关负波范式,共26名受试者。
  • PhysioP300:使用事件相关电位P300范式,共9名受试者。
  • TUAB:异常检测数据集,共2383名受试者。
  • TUEV:事件类型分类数据集,共288名受试者。

2. 实验设置

模型设置

EEGPT模型的实现采用了vision transformer(VIT)结构,设置了S个可学习的摘要标记(类似于分类标签[CLS]标记),使用58个电极,如图4所示。采样率为256Hz,输入信号的时间长度为1024。每个patch的时间长度为64,对应250ms的时间窗口。训练时对50%的时间维度和80%的通道维度进行掩码遮掩。

预训练策略

对于每个预训练数据集随机抽取10%的样本作为验证集。如表6所示,训练了8个不同编码长度、模型层数、摘要标记的变体。训练过程中,使用AdamW优化器,采用OneCycle学习率策略(初始学习率为2.5e-4,最大学习率为5e-4,最小学习率为3.13e-5),训练200轮次,在8张Nvidia 3090 GPU上进行64和16位混合精度训练。

评估策略

对于TUAB和TUEV数据集的分割,遵循于BIOT模型的策略。其余数据集使用留一法(Leave-One-Subject-Out, LOSO)进行验证,其中Sleep-EDFx遵循6:2:2的分割比例并使用十折交叉验证,KaggleERN使用四折交叉验证,并对下游任务使用线性探测方法。例如在睡眠阶段检测任务中,使用4层transformer编码器模型作为分类器,每0.25s输出一次模型结果,以处理30s的长任务。为保证实验可靠性,将每个实验重复三次并计算标准偏差。

基线指标

对于TUAB和TUEV数据集,使用来自BIOT模型的相同基线。在其它任务中,使用预训练的BENDR、BIOT和LaBraM作为基线。以下指标用于比较:平衡准确率(Balanced Accuracy, BAC)、AUC-ROC曲线下面积(AUROC)、加权F1、科恩卡帕系数(Cohen’s Kappa),其中AUROC用于二分类任务,加权F1仅用于多分类任务。

结果分析

实验结果表明,EEGPT在多个数据集上的表现均优于现有基线模型。结果如表2、表3所示,在TUAB数据集上,EEGPT的性能与BIOT模型相当,在TUEV数据集上,EEGPT比BIOT模型提高了9.5%的平衡准确率,加权F1提高了6.9%。

同时也在其他下游任务数据集上与BENDR、LaBraM进行了比较实验,结果如表4所示,EEGPT的性能也显著优于其他模型,证明了其在多任务EEG信号处理中的有效性。在所有任务上,EEGPT模型与BENDR、BIOT和LaBraM相比都取得了具有竞争力的结果。这表明EEGPT在时空维度上学习了一致的表征特征,使得该模型能够更广泛地应用于多种范式任务,实现更好的分类性能。

结论

EEGPT通过创新的双自监督学习方法和时空表示对齐技术,显著提升了EEG信号的特征提取能力。该模型是具有1000多万个参数的自监督EEG预训练Transformer模型,能够用于学习通用的EEG特征,通过采用双自监督方法进行预训练(包括时空表示对齐和基于掩码的重建)。时空表示对齐将是否使用掩码的特征patch进行对齐,提高了EEG表征的质量。基于掩码的重建是利用EEG信号所表现出的空间和时间一致性,在两个维度上提取互补特征,通过层次结构首先从短期EEG信号中提取稳定的空间表示,再捕获长期EEG信号之间的时间相关性。实验结果验证了其在运动想象、睡眠阶段检测和ERP类型分类等多种下游任务中的优越表现,展示了其广泛的应用前景和潜力,为生物信号处理和人工智能应用提供了新的解决方案。未来的研究可以进一步探索EEGPT在更复杂的脑机接口和神经科学研究中的潜在应用。

附录:近2年脑电信号大模型的相关论文对比

(需要的话可以在下方对应论文,加上我们的往期推文链接)

《Large Brain Model for Learning Generic Representations with Tremendous EEG Data in BCI》:受大语言模型的启发,提出了一个名为大脑大模型(LaBraM)的统一 EEG 基础模型,旨在从大量EEG数据中学习通用特征表示。模型通过分割 EEG 信号为通道补丁和向量量化神经频谱预测,能够有效处理不同电极配置和信号长度。在来自约 20 个数据集的大约 2500 小时的各种类型的 EEG 信号上进行预训练,能处理不同电极配置和长度的 EEG 数据。

《NEUROLM: A UNIVERSAL MULTI-TASK FOUNDATION MODEL FOR BRIDGING THE GAP BETWEEN LANGUAGE AND EEG SIGNALS》:提出多任务基础模型NeuroLM,旨在整合语言和EEG信号,以提高情感识别和认知状态分析的准确性,将 EEG 信号视为外语,融入大语言模型框架。通过文本对齐的神经标记器编码EEG信号、多通道自回归预训练和多任务指令调整,实现多任务学习和推理。

《FoME: A Foundation Model for EEG using Adaptive Temporal-Lateral Attention Scaling》:提出 FoME (脑电图基础模型),开创性地采用时间-频率融合嵌入技术和自适应时间-横向注意力缩放(ATLAS)机制处理EEG信号,捕捉其复杂的时空动态,适应不同数据集,并结合CNN和Transformer增强了对局部和全局特征的学习能力,促进稳健的多通道建模。

《Brant - X: A Unified Physiological Signal Alignment Framework》:提出Brant - X模型,用于对 EEG 与其他生理信号之间的相关性进行建模,解决了多种生理信号间对齐的问题。利用 EEG 基础模型,通过两级对齐策略(补丁级和序列级)分别处理局部和整体特征,并引入特征提取和对齐算法,有效对齐 EEG 与其他信号(EOG、ECG、EMG)的语义,以及适应多任务学习的方法同时处理多个生理信号任务以提高性能。

《EEGPT: Pretrained Transformer for Universal and Reliable Representation of EEG Signals》:提出用于通用 EEG 特征提取的EEGPT 模型,设计了双自监督学习方法,结合时空表示对齐和掩码重建,提高特征质量和模型收敛性,为 EEG 信号处理提供有效的创新解决方案。

《EEGPT: Unleashing the Potential of EEG Generalist Foundation Model by Autoregressive Pre-training》:提出自回归EEG 预训练模型,通过引入电极建模策略、开发自回归预训练方法以及实施带有可学习电极图网络的多任务迁移学习范式,捕捉了EEG信号中固有的顺序和时间依赖关系。

相关链接

论文链接:https://openreview.net/forum?id=lvS2b8CjG5

PPT汇报链接:

https://nips.cc/media/neurips-2024/Slides/93793.pdf

撰稿人:陈宗楠

审稿人:潘家辉


登录用户可以查看和发表评论, 请前往  登录 或  注册
SCHOLAT.com 学者网
免责声明 | 关于我们 | 用户反馈
联系我们: