[论文解读] Efficacy of Transformer Networks for Classification of Raw EEG Data

Gourav Siddhad, Anmol Gupta+2 more|arXiv (Cornell University)|Feb 8, 2022

EEG and Brain-Computer Interfaces被引用 11

一句话总结

本文提出了一种基于Transformer的端到端深度学习框架，用于对未经手工特征提取的原始预处理脑电信号进行分类。在公开的STEW数据集上，该方法在两级工作负荷任务中实现了95.28%的最先进准确率，在三级工作负荷任务中实现了88.72%的准确率；在本地的年龄与性别数据集上，分别实现了94.53%（性别）和87.79%（年龄）的准确率，证明了Transformer在脑电信号分类中的有效性。

ABSTRACT

With the unprecedented success of transformer networks in natural language processing (NLP), recently, they have been successfully adapted to areas like computer vision, generative adversarial networks (GAN), and reinforcement learning. Classifying electroencephalogram (EEG) data has been challenging and researchers have been overly dependent on pre-processing and hand-crafted feature extraction. Despite having achieved automated feature extraction in several other domains, deep learning has not yet been accomplished for EEG. In this paper, the efficacy of the transformer network for the classification of raw EEG data (cleaned and pre-processed) is explored. The performance of transformer networks was evaluated on a local (age and gender data) and a public dataset (STEW). First, a classifier using a transformer network is built to classify the age and gender of a person with raw resting-state EEG data. Second, the classifier is tuned for mental workload classification with open access raw multi-tasking mental workload EEG data (STEW). The network achieves an accuracy comparable to state-of-the-art accuracy on both the local (Age and Gender dataset; 94.53% (gender) and 87.79% (age)) and the public (STEW dataset; 95.28% (two workload levels) and 88.72% (three workload levels)) dataset. The accuracy values have been achieved using raw EEG data without feature extraction. Results indicate that the transformer-based deep learning models can successfully abate the need for heavy feature-extraction of EEG data for successful classification.

研究动机与目标

• 评估Transformer网络在无需手工特征提取的情况下，对原始脑电信号进行端到端分类的有效性。
• 通过利用Transformer的注意力机制，解决脑电信号分类中长期依赖于预处理和特征工程的问题。
• 证明Transformer在本地和公开脑电信号数据集上，能够实现与现有深度学习模型相当或更优的性能。
• 通过自动化的、数据驱动的特征学习，减轻脑电信号中伪迹污染和个体间差异带来的瓶颈。

提出的方法

• 将Transformer编码器架构直接应用于原始、清洗后的脑电信号，使用可学习的标记嵌入和位置编码来建模时间依赖性。
• 对脑电信号适配了位置编码，尽管未针对其生理和时间特性进行专门优化。
• 模型使用多头自注意力机制，捕捉脑电通道和时间点之间的长程依赖关系，实现联合的时空表征学习。
• 在学习到的表征之上添加分类器头，用于在两个数据集上执行二分类和多分类任务。
• 采用交叉熵损失函数进行端到端训练，并使用早停法，通过验证性能调整超参数。
• 消融研究将所提模型与基于注意力的基线模型进行比较，证实多头自注意力在脑电信号表征学习中的优势。

实验结果

研究问题

• RQ1 Transformer网络是否能有效分类未经预处理的原始脑电信号，且性能可与或优于现有最先进方法？
• RQ2 基于Transformer的模型在脑电信号分类任务中，与传统机器学习和深度学习模型相比，性能如何？
• RQ3 Transformer中的多头自注意力在多大程度上能够捕捉脑电信号中与分类相关的时空模式？
• RQ4 位置编码的选择是否会影响脑电信号上的性能？是否可针对神经生理信号特性进行优化？
• RQ5 所提出的框架能否在不同脑电信号分类范式（如年龄/性别预测和心理工作负荷估计）中实现泛化？

主要发现

• 在本地脑电数据集上，该模型在性别分类任务中达到94.53%的准确率，在年龄分类任务中达到87.79%的准确率，证明了无需特征工程即可实现优异性能。
• 在公开的STEW数据集上，该模型在两级工作负荷分类任务中达到95.28%的准确率，在三级工作负荷分类任务中达到88.72%的准确率，优于现有最先进方法。
• 该模型在精确率、召回率和F1分数等各项指标上均表现稳定，混淆矩阵和ROC曲线表明其泛化能力出色。
• 消融研究证实，多头自注意力显著优于单头注意力机制，凸显其在学习复杂脑电信号表征中的关键作用。
• 结果表明，Transformer能够从原始脑电信号中有效学习有意义的表征，降低对人工特征提取的依赖。
• 尽管结果表现强劲，但在年龄/性别数据集上的性能略低于STEW数据集，可能由于位置编码不够优化且缺乏领域特定特征的整合。