该论文发表于International Conference on Learning Representations 2024 (CCF A),题目为《Large Brain Model for Learning Generic Representations with Tremendous EEG Data in BCI》。
上海交通大学的姜卫邦为论文的第一作者,上海零唯一思科技有限公司赵黎明博士和上海交通大学计算机科学与工程系吕宝粮教授为本文共同通讯作者。
论文链接:
https://openreview.net/forum?id=QzTpTRVtrP
论文概要
当前基于EEG(脑电图)的深度学习模型通常是为了特定的数据集和应用而设计的,这种做法限制了模型的泛化能力,从而削弱了其感知能力和普适性。我们期望脑电大模型(LEM)能够超越不同任务类型脑电数据集的限制,通过无监督预训练获得对脑电信号的通用感知能力,然后再针对不同的下游任务对模型进行微调。然而,与文本数据相比,脑电图数据集通常规模较小,格式多样,可能面临电极数量不一致、数据样本长度不等、任务设计多样性以及信噪比较低等问题。
为了应对这些挑战,我们提出了一个用于脑电信号通用表征的大模型(LaBraM)。LaBraM通过实现跨数据集学习,将脑电图信号分割成脑电图通道贴片。利用向量量化神经谱预测技术训练出语义丰富的神经分词器,将连续的原始脑电图通道切片编码为紧凑的神经代码,然后通过预测原始神经代码来预训练神经Transformer模型。LaBraM在大约2,500小时的数据上进行了训练,并在多种不同类型的下游任务上进行了实验。在异常检测、事件类型分类、情绪识别和步态预测等任务中,LaBraM的表现均优于所有比较的当前最佳方法(SOTA)。
研究背景
目前基于脑电(EEG)信号的深度学习模型通常是为特定数据集和脑机接口(BCI)应用而设计的,这限制了模型的规模,从而削弱了脑电信号的表征能力和普适性。近来,大型语言模型(LLM)在文本处理方面取得了前所未有的成功,这些颠覆性的研究成果促使人们开展了大型脑电模型(LEM)的探索。我们希望大型脑电模型能突破不同任务类型脑电数据集的限制,通过无监督预训练获得脑电信号的通用表征能力,然后在不同的下游任务上对模型进行微调。
然而,与文本数据相比,通常脑电数据集的规模非常小,格式也千差万别。开发大型脑电模型面临着下列几方面的挑战: 缺乏足够的脑电数据、 EEG信号采集的多样化配置、缺乏有效的脑电表征学习范式。
研究方法
本文的目标是设计一种通用的大型脑电模型,称为LaBraM,如图1所示。该模型可有效处理不同通道和长度的各种脑电数据。通过对大量脑电数据进行无监督训练,作者设想该模型将具备通用的脑电表征能力,使其能够快速适应各种脑电下游任务。为了训练LaBraM,作者从20个公开的脑电数据集中收集了超过2500个小时的各种任务和格式的脑电数据。
图1 LaBraM的整体架构。首先,所有输入的脑电信号将通过一个固定长度的时间窗口分割成脑电信号片段,然后对每个片段应用时间编码器提取时间特征。然后,将时间和空间嵌入添加到片段特征中,以携带时间和空间信息。最后,将嵌入序列按片段顺序传入Transformer编码器,以获得最终输出。
本文通过将原始脑电信号分割成脑电信号通道片段,解决了不同电极数量和时间长度的问题。同时,本文采用了向量量化神经频谱预测技术来训练语义丰富的神经标记器,以生成神经词汇。具体而言,标记器是通过预测原始信号的傅立叶频谱来训练的。在预训练阶段,部分脑电片段会被掩蔽,而神经Transformer的目标是基于可见片段预测被掩蔽的标记。作者预训练了三个不同参数规模的模型,分别为580万、4600万和3.69亿参数,这在BCI领域是迄今为止规模最大的模型。随后,作者在四种不同类型的下游任务上对这些模型进行了微调,这些任务包括分类和回归。
本文引入了神经Transformer,这是一种用于解码脑电信号的通用架构,能够处理任意通道数和时间长度的输入脑电信号。实现这一目标的关键操作是将脑电信号分割成块,这一灵感来源于图像中的片段嵌入技术。由于脑电信号在时间域具有高分辨率,因此在通过自注意力机制进行片段交互之前,提取时间特征至关重要。作者采用了由多个时域卷积块组成的时域编码器,将每个脑电片段编码成片段嵌入。这些时域卷积块由一维卷积层、组归一化层和GELU激活函数组成。为了使模型能够感知片段嵌入的时间和空间信息,作者初始化了一个时间嵌入列表和一个空间嵌入列表。对于任意片段,作者将其对应的时间嵌入和空间嵌入加到片段嵌入上。最终,嵌入序列将直接输入到Transformer编码器中。
图2 神经标记符号训练和LaBraM预训练概述。上部分:训练神经标记器,通过重建傅立叶频谱将脑电信号离散为离散的神经标记。下部分:在预训练过程中,部分脑电片段会被掩蔽,而目标则是从可见片段中预测被掩蔽的标记。
在通过掩码和预测对LaBraM进行预训练之前,需要将脑电标记转化为离散的标记。作者提出了矢量量化神经频谱预测,它是通过预测傅立叶频谱来训练的,如图2上部分所示。其关键部分是将脑电样本编码成片段表示的神经标记器和从神经嵌入解码傅立叶频谱的神经解码器。为了让LaBraM利用大量脑电数据学习通用表征,作者提出了掩码脑电建模。整个过程如图二下部分所示。给定一个脑电样本,时序编码器首先将其转换为片段嵌入。随机生成一个掩码,被掩蔽的EEG片段将被加入时间和空间嵌入,然后输入变换器编码器,来预测掩蔽的EEG片段。作者进一步提出了一种对称掩码策略,以提高训练效率。
研究结果
作者主要在两个下游任务数据集上验证LaBraM:TUAB和TUEV。并设计了三种不同的LaBraM配置:LaBraM-Base、LaBraM-Large和LaBraM-Huge。LaBraM-Base的参数数为5.8M,LaBraM-Large为46M,LaBraM-Huge为369M。
图3 预训练损失曲线和掩码脑电建模准确率曲线。
图3比较了Base模型、Large模型和Huge模型的预训练总损失和掩码脑电建模精度的收敛曲线。从图3可以看到,具有更多参数的大型模型可以收敛到更小的损失和更高的精度。值得注意的是,Huge模型的损失似乎有明显的下降趋势,而如果对其进行更长时间的训练,准确率则会趋于提高。这一观察结果表明,扩大模型规模有可能获得更好的性能。
表1和表2列出了 TUAB 和 TUEV 中最佳的基线结果以及LaBraM的结果。很明显,LaBraM-Base模型在这两项任务的各种评估指标上都优于所有基线模型。特别是在更具挑战性的 TUEV 多类分类任务中,该模型取得了显著的性能提升。在该模型中,作者观察到随着模型参数数量的增加,LaBraM-Huge模型的表现最好,其次是 LaBraM-Large模型,最后是LaBraM-Base模型。作者将这一良好表现归功于预训练数据量和模型参数的增加。作者推断,只要有足够多的脑电数据,大规模脑电模型就能学习到更通用的脑电表征,从而提高脑电信号在各种下游任务上的性能。
图4 比较模型在TUAB和TUEV数据集上的性能,是否纳入预训练过程。
在预训练过程中,作者希望模型能够学习到不针对任何特定任务的通用脑电表征。虽然在预训练过程中没有使用标签数据,但为了消除预训练数据对下游任务的影响,作者比较了是否将下游任务数据集纳入预训练过程的结果。值得注意的是,TUAB和TUEV的记录与预训练数据集的记录是不相交的。如图4所示,是否将下游任务数据集纳入模型的预训练过程,对模型在下游任务上的性能影响不大。这表明该模型具有学习通用脑电表征的能力,并为将来收集更多脑电数据提供了指导。换句话说,不需要在预训练过程中花费大量精力标注脑电数据。
图5 随着预训练数据的增加,Base模型、Large模型和Huge模型在 TUAB 和 TUEV 数据集上的性能比较。
虽然该研究已经收集了约2,500小时的脑电数据,但与自然语言处理和图像处理的样本量相比,仍然相对较小。通过调整预训练数据的大小来应对训练不同大小的 LaBraMs对数据大小的需求。如图5所示,训练时间为 500 小时的基本模型的性能超过了训练时间为2500小时的模型在TUAB上的性能,同时接近训练时间为 2500 小时的模型在 TUEV 上性能的 90%以上。对于Large模型来说,性能一般会随着数据量的增加而提高,尽管在 1000 小时后增速会放缓。相比之下,随着数据量的不断扩大,Huge模型的性能呈现出明显的上升趋势。因此,作者相信随着数据集的进一步扩大,LaBraM可以获得更好的性能。预训练大型脑电模型需要多少脑电数据,无疑是该领域值得探讨的重要问题。然而,2500 小时至少不是这个问题的答案。观察结果基本上遵循了缩放定律(scaling law),由此推断,当数据量达到至少一万小时的数量级时,Huge模型将继续保持较好的性能。
结论
本文提出了一种基于大量EEG数据的通用表征大模型(Large Brain Model,LaBraM),通过对超过 2,500 小时的各种脑电图数据进行无监督预训练来学习通用表征。LaBraM 能够处理各种EEG数据集,这是因为它将原始脑电信号类比成图片,然后分割成通道片段(patch),并使用了矢量量化神经频谱预测,在预训练期间生成丰富的语义标记符号。此外,神经Transformer还能对脑电信号的时间和空间特征进行有效的表征学习,使其适用于脑电图分析中的各种下游任务。LaBraM 在多个下游任务中得到了验证,包括异常检测、事件类型分类、情绪识别和步态预测,并且其表现均优于所有 SOTA 方法。
撰稿:陆文涛
审核:潘家辉