该论文发表于NeurIPS 2023(CCF A),题目为《Learning Topology-Agnostic EEG Representations with Geometry-Aware Modeling》。
华南理工大学的柯毅为本文第一作者,微软亚洲研究院首席研究员李东胜为本文的通讯作者。
论文链接:https://papers.nips.cc/paper_files/paper/2023/file/a8c893712cb7858e49631fb03c941f8d-Paper-Conference.pdf
大规模的预训练在视觉和语言上的下游任务上显示出增强模型的巨大潜力。由于EEG的未标记数据非常丰富,开发类似的头皮脑电EEG预训练技术是合适的。同时,多样的采样通道选择以及固有的结构和空间信息为进一步改进现有的预训练策略带来了挑战和途径。为了打破不同EEG资源之间的界限,促进跨数据集的EEG预训练,本文提出将各种通道选择映射到一个统一的拓扑结构。本文进一步引入了MMM(Multi-level hierarchy learning,Multi-dimensional position encoding,Multi-stage mask strategy),一个具有多维位置编码的预训练框架,多级信道层次结构,建立在统一拓扑上的多阶段预训练策略,以获得拓扑无关性。MMM整体架构如图1。
图1 MMM整体架构
研究背景
由于EEG数据能够以非侵入和低成本的方式采集,但需要大量的人力对其进行标注,自监督预训练方法被大量应用于EEG片段的学习中。虽然有许多公开可用的EEG数据集用于预训练,但在它们之间进行迁移是非常具有挑战性的,因为这些EEG记录在蒙太奇(头皮上放置电极的数量和位置)和采样率方面可能会有很大的变化。EEG预训练的另一个挑战是对结构进行编码,对来自不同位置的EEG信号之间的关系进行建模。这些信号由放置在头皮上的电极采集,可以看作是一个具有空间信息的二维流形。这种空间信息对于专家提供精确的标签往往是至关重要的。综上所述,为了得到一个有能力和强大的EEG预训练模型,需要1 )将空间信息很好地编码到表示中,2 )使用具有不同传感器配置的EEG语料预训练模型。
研究方法
本文提出了一种基于掩码自动编码器的自监督学习框架MMM,其使用DE特征聚合时间信息,而MMM框架本身则用于探索EEG表示在空间结构方面的学习。该框架主要包括以下几个技术要点。
(1)不同传感器配置的统一拓扑
现有的来自不同机构的EEG数据集大多使用不同的传感器配置(32通道、62通道等)。而同时,最近的神经科学研究表明,大脑可以划分为数个大脑功能组织。据此,本文将大脑划分为17个区域,如图1 ( b )所示。因此,处理各种传感器配置的挑战可以看作是将具有不同传感器配置Ex的数据x∈X映射到统一拓扑Z = Encoder ( x )的表示中。
(2)MMM:基于统一拓扑的EEG预训练框架
本文介绍了一种基于统一拓扑结构的EEG预训练框架MMM,该框架遵循掩蔽自编码器( MAE ) 的训练模式(如图1(a)所示),即在预训练期间,模型首先用编码器将部分掩蔽的不规则输入token( DE特征)编码为统一表示,然后用解码器从统一表示中重建掩蔽token。
为了更好地支持EEG预训练,本文设计了3种MMM:1 )多维位置编码,它将几何信息注入到tokens中;2 )多级通道层次结构,即在原始通道集合中添加代表区域的额外tokens;3 )多阶段预训练,交替使用全局随机掩蔽和区域掩蔽两种不同的掩蔽策略,以有效地学习层次化表示。
(3)多维位置编码
传统的位置编码只考虑加入通道间的顺序信息,这丢失了通道间的空间信息。本文借助二维流形的思想,将多维位置编码定义为:
其中,( x , y)为来自头皮的二维坐标,d为传感器个数,10 - 10系统网格对应的( x , y)的d对构成一个空间。函数R将x,y投影到其原始空间中的秩值。i为编码中特征维度的索引,Ω为一个大的常数(一般设置为10000)。最终的编码是通过将x和y坐标的编码串接在一起生成的。
多维位置编码在两个地方使用。首先将来自不同通道的DE特征经过线性层处理后加入,从而为数据注入几何信息。其次,在获得表示后重新应用编码。对重掩码表示进行位置编码,告知解码器如何区分重掩码token。
(4)多级通道层次结构
本文的模型骨干部分训练和基于Transformer的编码器-解码器作为掩码自动编码器模式相似。但相较于传统的直接将各原始通道节点直接放入transformer不同,本文还引入了区域节点/区域标记s∈R17 × D,并将其追加到序列末尾,如图1(a)下方所示。同时为了避免区域节点之间的交互冗余,本文为为编/解码器分配了一个固定的注意力掩码,图1(d)为其中一个例子。通过掩蔽建模的方式,该训练方式迫使模型能对EEG的统一拓扑结构进行良好的建模。
(5)不同掩蔽方式
为了让MMM更好的学习到全局信号知识和区域信号知识,本文相应地引入两种掩蔽策略,并在预训练过程中轮流使用。其中全局随机策略掩蔽通道的掩蔽信道数量固定,区域掩蔽策略对同一区域内的通道全部掩蔽或全部不掩蔽,如图2所示。
图2 掩蔽策略
实验结果
(1)被试相关分类任务
本文在SEED和SEED - IV的被试相关情况下的不同设置下与既定的基线进行了比较,结果如表1所示。窗口大小是构建输入的特征长度。例如,4s × 5表示每个DE特征聚合4秒的原始EEG信号,五个连续的DE特征被用作模型输入。匕首符号代表在计算最终的平均准确率时,每个受试者排除3个中表现最差的会话。
表1 MMM方法与其他方法在SEED和SEED - IV上的实验结果
(2)不同数据集间的知识迁移
在这一部分中,本文研究了模型在具有相同传感器配置的数据集上的可迁移性。本文在原始数据集上预训练模型,包括SEED、SEED - IV以及它们的组合。随后,在目标数据集SEED和SEED - IV上评估这些预训练模型的微调性能。
表2 基于SEED和SEED-IV的数据迁移实验结果
(3)不同通道设置数据集间的知识迁移
在这一部分中,本文进一步研究了本文的方法在不同传感器配置下的可迁移性。本文使用Lite系列数据集预训练模型,然后在完整的SEED数据集上进行部分微调,即冻结编码器的参数,只调整最后一个MLP,如表3所示。MMM在SEED - Lite和SEED - Union - Lite数据集上预训练的模型均优于从头训练的模型。
表3 在SEED数据集上的部分微调结果
(4)从大规模脑电语料中进行知识迁移
模型在TUEG数据集上进行预训练,( - )表示从头开始训练得到的模型。然后,在SEED数据集上对模型进行微调和测试。鉴于之前的研究发现了本文的方法在不同数据集和蒙太奇中的有效性,接着将方法部署在TUEG数据集上。值得注意的是,尽管TUEG数据集收集了10 - 20系统中涉及21个通道的不同通道配置的数据,但其在预训练中的应用仍然是有益的,如表4所示。
表4 大规模预训练
(5)区域划分的影响
在这一部分,本文深入分析了MMM中不同的统一拓扑对SEED数据集的影响。本文将默认的17区域设置与几种变体进行了比较:
表5 不同区域划分选择的变体结果
结论
在这项研究中,本文提出了一种用于EEG预训练的创新框架MMM,它有效地导航了EEG数据中与各种传感器配置相关的复杂性,并很好地编码了EEG的空间信息。本文的模型采用了基于掩码自动编码器的自监督学习方法,使其能够从大量的跨数据集语料库中获益,并提高下游任务的性能。本文的框架在情绪识别任务中展示了最先进的性能,进一步的评估证实了其方法的有效性。
撰稿人:肖睿
审稿人:李景聪