脑电大模型--LaBraM

LaBraM是上海交通大学吕宝粮老师团队提出的第一个脑电领域的大模型，将NLP领域的掩码预训练任务引入到脑电领域，得到的大模型在多个不同的下游任务上都得到了性能验证。


一、论文的主要贡献：

1、解决单一数据集数据量小：

         本文使用了来自20个不同数据中超过2500小时的不同人物和格式的EEG数据，分别来自不同的公开数据集及该团队自建的数据集。

2、对大规模未标记数据的利用/不同数据的统一：

         仿照NLP/CV中预训练的模式，作者在脑电中也引入了改进的MASK机制，让模型能够学习上下文信息/特征信息。

3、有效的通用特征学习：

         作者利用neural Transformer模型从时间/空间域学习不同通道和长度的EEG特征，可以适用于更广泛的下游任务。同时引入了一个神经码本（neural codebook）的设计，获得了一个更紧凑且通用的EEG信号表示。同时该设计的引入也使得模型可以采用预训练方法来进行特征学习。

4、丰富的下游任务验证：

        作者在四个不同的具有代表性的下游任务上验证了LaBraM模型，超过了目前SOTA的方法。

二、论文的主要算法模块

（1）Neural Transformer

1、分割Patch：

       首先按照固定的时间窗和通道宽度将EEG数据划分为patch。首先在时间维度可使用有重叠的分割，将EEG数据进行第一次分割。其次在小片段内部利用一定的窗口长度将EEG数据再次分割（无重叠），同时分割开每个通道的数据。

2、时域编码器：

       根据之前的研究，作者指出利用卷积对EEG数据进行特征提取十分有效，因此在编码之前引入了一个由卷积、Group Norm和GELU组成的时域编码器对每一个patch进行编码。

3、时域和空间域的embedding：

      初始化一个可学习的时域的embedding表和一个可学习的空间域的embedding表，时域embedding表的大小>=patch的序列长度，空域embedding表的大小与通道数相当。把时域和空域的embedding信息加入到时域编码器的输出。

4、Transformer Encoder：

      这一步的模块采用的是Transformer的Encoder，为了使得训练更加稳定和高效，对于Q、K添加了Layer Norm操作。

5、下游任务：

        对于输出的Output Embedding，作者将其应用到下游任务上时首先通过一个平均池化的操作，然后使用特定于任务的预测头。

（2）neural tokenizer training

作者想模仿NLP中自回归的预测任务，但是直接预测EEG数据无法收敛，因此作者借鉴自VQ-VAE是量化的想法，将EEG信号表示为频率上离散的标记，借此对网络进行训练。

1、Neural Codebook：

      codebook的维度为K*D，D为embedding的维度，K为离散的embedding的维度。将Neural Transformer输出的embdding进行量化表示，然后在codebook中计算L2距离将EEG信号替换为离散表示。

2、傅立叶频谱预测：

       将离散表示输入到Neural Decoder中，让模型预测该EEG片段的幅度谱和相位谱，利用MSE误差来进行训练。以此更新codebook中的表示。

（3）预训练任务

对于编码器的输出，以一定的比例随机mask掉一些Patch，然后通过时空Embedding和Transformer Encoder模块进行编码，最后利用一个线性的预测头给出mask掉的patch的输出。

三、实验结果