【文献阅读】UniMind：释放LLM的力量，实现统一的多任务大脑解码

UniMind：释放LLM的力量，实现统一的多任务大脑解码

原文标题：UniMind: Unleashing the Power of LLMs for Unified Multi-Task Brain Decoding

1 概述

作者：Weiheng Lu, Chunfeng Song, Jiamin Wu, Pengyu Zhu, Yuchen Zhou, Weijian Mai, Qihao Zheng, Wanli Ouyang; 来自上海人工智能实验室，北京大学，香港中文大学

期刊： 论文在 arXiv 上公开，提交日期为 2025 年 6 月 23 日

利用的思想：将EEG识别与LLM结合，解码EEG信号后向LLM提问。（NeuralLM）这种思想在多任务EEG理解中，与单任务方法相比，其性能不令人满意

动机：解决上述问题

创新点：NLC，TQS

意义/结果：第一个在多任务中实现与单任务解码模型相当甚至更优性能的模型

2 背景

背景的背景：

EEG在脑机接口（BCI）中起着至关重要的作用，它提供了一种非侵入性的、实时的大脑活动测量方法，也是研究大脑感知机制的有力工具。通过分析EEG信号中的神经模式，已经进行了许多研究来解码特定的大脑状态，在癫痫发作分类、睡眠阶段分类、运动图像识别、异常检测、情绪分析和急性压力检测等不同应用中显示出巨大的潜力。

  • 1: 由于不同任务中EEG信号格式的变化，许多深度学习模型仅限于特定任务的范式。虽然这些模型对预期任务有效，但它们很难泛化到新的任务上。
  • 2: 收集特定任务的EEG数据成本高昂，因此为每个任务构建足够的数据集是不切实际的。这些局限性凸显了对通用EEG模型的迫切需求，这些模型可以学习跨任务的可转移表示。
  • 作为回应，已经提出了BIOT[19]和LaBraM[20]等基础模型，通过大规模的自我监督预训练来获得EEG信号的普遍感知能力。然而，这些模型需要对每个下游任务进行额外的微调，以超越特定任务的模型，因此仍在努力实现泛化，并导致额外的计算开销。

背景：

为了应对上述挑战，最近提出的NeuroLM开发了一种多任务EEG基础模型，该模型使用对抗域级对齐和多通道自回归将大型语言模型（LLM）集成在一起，以桥接EEG和语言模态。

不足：虽然NeuroLM展示了LLM在多任务EEG理解中的潜力，但与单任务方法相比，其性能仍然不令人满意，在SEED[22]、TUEV[23]和SleepEDF[2]等数据集上的表现落后10%以上。为了释放LLM解码神经模式和统一不同EEG任务的潜力，需要解决两个关键挑战：

  • 神经信号和语言之间存在巨大的模态差距。与文本标记的丰富、结构化特征不同，EEG信号以高噪声、稀疏的语义信息和复杂的神经模式为特征，LLM难以直接感知。EEG和语言模态之间的对齐是LLM理解EEG信号和跨模态推理的先决条件。现有的方法使用对抗训练在域级别对齐EEG和文本，但由于忽略了语义级别的关系，模态差距仍未得到解决。核心问题仍然是如何进行有效的跨模态桥接，以缓解EEG和LLM之间的模态差距。
  • EEG解码任务的异质性。由于不同的配置和认知机制，EEG任务显示出显著的异质性，导致数据集之间的信号特征不同，如电极通道、试验持续时间和时空脑活动。鉴于此，纯粹与任务无关的混合任务调优可能会导致特定任务的解码性能下降。因此，在多任务环境中，迫切需要一种任务感知机制来实现通用的大脑解码。此外，任务无关的方法忽视了任务之间的神经相关性，限制了对大脑认知的理解。如何开发一种任务感知机制，在不牺牲单个任务性能的情况下，实现对EEG任务异质性鲁棒的统一多任务学习，仍然是一个悬而未决的挑战。

为了应对这些挑战，我们引入了UniMind，这是一种通用的大脑基础模型，通过释放LLM在理解大脑模式方面的潜力，实现统一的多任务大脑解码，如图1所示。

意义：据我们所知，这是第一个在一个统一模型中匹配单任务方法性能的多任务EEG解码模型。为了缩小模态差距，我们提出了一种神经语言连接器（NLC），它充当EEG编码器和LLM之间紧凑、可训练的桥梁，以语义上有意义的方式从稀疏的EEG数据中压缩出基本的大脑模式，供LLM解释。

怎么做的：具体来说，为了利用EEG信号的时空特性[18,25]，神经语言连接器采用双分支架构，具有可学习的查询令牌，通过交叉注意力分别聚合神经信号的时间动态和空间依赖性。聚合的特征将对齐并映射到冻结LLM的语义空间。通过这样做，NLC将时空神经模式转换为LLM的可解释特征，从而实现无缝的神经语言集成。为了促进跨异构EEG任务的有效多任务学习，我们提出了一种任务感知查询选择模块（TQS）来生成任务自适应查询令牌。TQS维护包含多个查询令牌的空间和时间查询池。路由器机制用于根据输入特征从查询池中动态查找与任务相关的查询。通过允许每个任务自适应地选择自己的查询，TQS促进了相关任务之间的知识共享和相互增强，同时减少了冲突任务的干扰。此外，通过任务选择机制，我们通过检查任务间的相关性，揭示了不同任务在大脑中的功能组织方式，揭示了大脑如何调节不同的认知功能。

我们的贡献可以概括如下：

  • 我们提出了UniMind，这是一种通用的大脑基础模型，通过整合EEG和语言之间的时空跨模态桥接策略以及任务感知机制，用于统一的多任务大脑解码。
  • 我们提出了一种双分支神经语言连接器，将EEG信号的时间和空间模式编码为LLM可解释的表示，以及一个为任务相关表示生成自适应查询的任务感知模块。
  • 实验表明，UniMind的平均性能比现有的最佳多任务解码模型高出12%，并且是第一个在各种任务中实现与单任务解码模型相当甚至更优性能的模型。
  • 通过查询可视化提供神经科学角度的见解，揭示共享的神经机制，并为多任务大脑解码中的知识共享提供经验支持。

3 Method

image

3.1 神经语言连接器 (Neuro-Language Connector, NLC)

目的：核心目标是弥合 EEG 信号 和 语言模型 之间的模态鸿沟。由于 EEG 信号具有稀疏、低信噪比和复杂时空模式的特点，NLC 的设计目的是从这些嘈杂的信号中提取并解释时空神经特征，使其与 LLM 的语义空间对齐。

  • EEG 编码器：
    • 原始 EEG 信号 \(X\in R^{C\times S}\) (C 个电极通道，S 个时间采样点) 会通过一个预训练的 EEG 编码器处理。这个编码器将不同通道和长度的 EEG 信号转换为标准的令牌序列。
    • 它使用一个时间编码器将每个 EEG 通道分割成不重叠的 patches，类似于 Vit 处理图像的方式。如果滑动窗口长度为 \(t\)，每个通道会产生 \(T=⌊\frac{S}{t}⌋\) 个补丁。
    • 信号被编码成一系列补丁令牌嵌入，并在前面添加一个“类令牌”（class token），得到 \(E\in R^{(C\times T+1)\times D_e}\)。这里的 \(D_e\) 是 EEG 嵌入的维度。
  • NLC 核心：交叉注意力机制
    • NLC 使用交叉注意力机制来聚合 EEG 嵌入 \(E\) 中的时空特征。它设计了两组可学习的查询（Query）令牌：
      • 时间查询 \(Q_t\in R^{n_{qt}\times D_e}\)： 捕捉每个 EEG 通道内的时变模式，即时间动态。
      • 空间查询 \(Q_t\in R^{n_{qs}\times D_e}\)： 分析每个时间步长的空间通道激活模式，即空间依赖性。
    • 交叉注意力计算：
      • EEG 嵌入 \(E\) 会被重塑，以便分别进行时间和空间注意力：\(E_t\in\mathbb{R}^{C\times T\times D_e}\) 用于时间注意力，而 \(E_s\in\mathbb{R}^{T\times C\times D_e}\) 用于空间注意力。
      • \(F_t=\text{CrossAttn}(Q_t,E_t)\in\mathbb{R}^{n_{qt}\times C\times D_e}\)：时间查询 \(Q_t\) 作为查询，重塑后的 \(E_t\) 作为键和值。
      • \(F_s=\text{CrossAttn}(Q_s,E_s)\in\mathbb{R}^{n_{qs}\times T\times D_e}\)：空间查询 \(Q_s\) 作为查询，重塑后的 \(E_s\) 作为键和值。
      • 这种交叉注意力机制使得可学习的查询令牌能够从嘈杂的 EEG 信号中发现和聚合关键的大脑模式。
  • 特征融合与投影：
    • 经过交叉注意力后，得到压缩的时间特征 \(F_t\) 和空间特征 \(F_s\)。
    • 这些特征会与初始的类令牌 \(F_{cls}\) 进行拼接，形成神经语义嵌入 \(F=\text{Concat}(F_t, F_s, F_{cls})\in\mathbb{R}^{(n_{qs}\times C+n_{qs}\times T+1)\times D_e}\)。
    • 最后，这个神经语义嵌入 \(F\) 被投影到与 LLM 兼容的嵌入空间，得到 \(F'\in\mathbb{R}^{(n_{qs}\times C+n_{qs}\times T+1)\times D_L}\)，其中 \(D_L\) 是 LLM 的隐藏层维度。这一步确保 LLM 能够理解这些神经信号。

3.2 任务感知查询选择模块 (Task-aware Query Selection Module, TQS)

目的：TQS 旨在解决多任务 EEG 解码中任务间信号特征的异质性问题。由于不同 EEG 任务的信号模式差异大，直接使用共享架构可能会导致性能下降。TQS 通过动态学习任务自适应查询令牌，实现特征选择和任务适应性。

  • 查询池与路由器：
    • TQS 维护一个查询池 \(P\)，其中包含一系列可学习的查询令牌 \(P=[P_1,P_2,..,P_{Nq}]\)，分为时间查询池 \(P_t\) 和空间查询池 \(P_s\)。
    • 引入一个查询路由器 \(R\)（同样分为 \(R_t\) 和 \(R_s\)）。
    • 查询路由器 R 根据输入的 EEG 嵌入 \(E\) 在查询池 \(P\) 中的查询令牌上生成路由分数 \(S\)。这个分数表示每个查询与当前输入 EEG 特征的匹配程度。
  • 任务自适应查询选择：
    • 基于路由分数 \(S\)，TQS 从查询池 \(P\) 中选择出任务自适应查询 \(q_a^i\)，这通过 TopK 选择实现，即选择与当前任务最相关的 \(K\) 个查询。
    • 除了动态选择的 \(q_a^i\) 外，还会引入一个额外的静态查询令牌 \(q_{st}^i\)，用于捕获每个任务的稳定特征。
    • 最终的查询令牌序列是动态选择的查询和静态查询的拼接：\(Q_i=[q_a^i;q_{st}^i]\in\mathbb{R}^{(n_i+1)\times D_e},i∈\{t，s\}\)。
  • 与 NLC 结合：
    • 这些经过 TQS 模块动态选择和生成的 \(Q_t\) 和 \(Q_s\) 查询，随后被馈送到 NLC 中。
    • NLC 使用这些任务自适应的查询来从 EEG 信号中提取与当前任务最相关的时间和空间特征。
    • 这种机制使得每个任务都能自适应地选择其所需的查询，从而促进相关任务之间的知识共享，同时减少不相关任务的干扰。

[!NOTE]
查询令牌 \(P=[P_1,P_2,..,P_{Nq}]\) 等价于自己定义的 CodeBook
静态查询 \(q_{st}\) 等价于 \(X^{cls}\)

[!note]

静态查询 \(q_{st}\) 捕获的是全局的、通用的、稳定的大脑模式

动态查询 \(q_{a}\) 捕获的是局部的、任务特异的、可变的大脑模式

3.3 使用 LLM 进行多任务解码

目的：该模块将 EEG 解码任务转化为 LLM 能够处理的条件语言建模任务。

  • 多任务指令调优数据集：
    • 为了训练 LLM，UniMind 构建了一个专门的指令调优数据集，涵盖多种 EEG 解码任务类型。
    • 对于每个 EEG 样本，随机选择一个预先设计的 prompt，并与任务标签配对。
  • LLM 训练：
    • LLM 的输入序列由三部分拼接而成：
      • 神经语义表示 \(F’\) (由 NLC 生成)
      • 一个特殊的分隔符令牌 [SEP]，没有含义仅分隔
      • 描述任务的指令提示 P
    • 即输入为：\(X=[F',[SEP],P]\)。（[材料，“问”，问题）
    • 模型通过预测相应的目标令牌序列 \(t_a\) 来生成正确的标签或描述。
    • 训练过程使用因果交叉熵损失进行优化，这是一种标准的回声序列生成任务损失函数，目标是最大化生成正确标签的概率：

\[\mathcal{L}=-\sum_{i=1}^L\log p(t_{a,i}|X,t_{a,<i}) \]

其中 \(t_{a,i}\)，是预测的令牌，\(t_{a,<i}\) 是前面的令牌。

4 Experiment

4.1 实验设置

  • 数据集：采用了五个任务领域的十个公开可用的EEG数据集：睡眠阶段分类（HMC[26]、SleepEDF[27]、SHHS[28]）、情绪识别（SEED[22]、SEED-IV[29]）、临床异常检测（TUAB[23]、TUEV[23]、TUVL[30]）、认知工作量分类（workload[31]）和运动图像分类（SHU[32]）。所有实验都遵循训练、验证和测试分割，以确保公平一致的评估。
  • 指标：Balanced Accuracy；Cohen’s Kappa；Weighted F1
  • 训练细节：我们使用LaBraM的EEG编码器，并采用InternLM2.57B作为语言模型。在NLC中，每个分支都采用一个8头交叉注意力层，然后是两层MLP。查询路由器R被实现为两层MLP，查询池大小被设置为16。静态查询对空间和时间维度都使用1个标记，而任务自适应查询使用2个空间标记和1个时间标记。训练时，LLM冷冻以保持其预先训练的语言理解能力，而其他组件是可训练的。
  • 比较方法：为了确保全面评估，我们将UniMind与非基础和基础模型基线进行了比较。
    • （1）非基础模型。我们包括最先进的深度学习模型，如SPaRCNet[7]、ContraWR[14]、CNN Transformer（CNN Trans）[15]、FFCL[16]和ST Transformer（ST Trans）[17]。这些模型通常采用专门为单个任务或数据集设计的监督学习策略。
    • （2）基础模型。我们考虑NeuroLM[21]、BIOT[19]和LaBraM[20]。NeuroLM有三种参数量表（B、L、XL），我们报告了每个数据集的最佳表现版本。为了进行公平比较，我们遵循NeuroLM对SEED、HMC、TUEV和TUSL的数据分割和评估指标，并直接报告NeuroLM中提供的已发表结果。对于剩余的数据集，我们使用自己的数据分割和评估指标重新实现了基线模型。

4.2 与SOTA方法的对比

  • 与多任务模型进行比较：首先将我们的方法与NeuroLM在相同的多任务设置中进行比较。UniMind在所有数据集上的表现始终远远优于NeuroLM，平均性能提高了12%。这一改进源于UniMind通过神经语言连接器进行的深度语义对齐，这增强了LLM解释神经信号以进行多任务解码的能力，释放了其神经推理能力。此外，任务感知查询选择模块提高了模型对异构解码任务的任务适应性。值得注意的是，我们的方法在具有挑战性的SEED-IV和TUEV数据集上取得了显著的性能，大大提高了它们最初较低的平衡精度另一个有趣的现象是，UniMind在只有1k个样本的小规模Workload数据集上，平衡精度大幅提高了15.22%。这些改进源于TQS从共享池中进行的任务自适应查询选择，这促进了任务之间的知识转移和相互增强。这种机制有助于有效的多任务训练，提高复杂和数据有限任务的性能。
  • 与多任务模型进行比较：UniMind能够在统一模型中匹配甚至超越领先的单任务基线。与特定任务的调整相比，有效的大脑语言对齐可以缩小性能差距。此外，任务感知查询选择器促进的知识共享允许每个数据集从语义相似的数据集中受益，这是单个任务模型所不具备的能力。

4.3 消融实验

  • 关键组件分析：这些改进源于任务感知机制，该机制能够跨任务进行自适应查询学习，促进相关任务之间的知识共享，同时减少跨模态桥接中的任务间干扰。
  • 不同类型查询的比较：这种改进主要是由于每个任务都能自适应地选择查询，从而更好地捕捉与任务相关的时间和空间特征。此外，将静态查询和任务自适应查询相结合，可以在所有数据集中获得最佳的整体性能。
  • 任务自适应查询数量分析：结果表明，增加查询数量并不总是能提高性能：太少的查询限制了模型捕获不同EEG模式的能力，而太多的查询可能会引入噪声，因为一些查询可能与EEG嵌入不一致。虽然最佳设置因数据集而异，但当ns=2和nt=1时，总体结果最佳。
  • 查询池大小分析：当池太小时，特定于任务的选择往往会重叠，这可能会导致任务之间的干扰。相反，过大的池会减少相似任务之间的查询重叠，削弱相关任务之间可能发生的相互增强。通过平衡多样性和共享表示，将池大小设置为16可以获得最佳的整体性能。

EEG信号的不重叠分割是不是也是为了防止选择过多重叠而影响任务之间的干扰？

  • 模型参数大小分析：图4调查了不同模型尺寸对性能的影响。我们将InternLM模型与大约0.5B、1.8B和7B参数进行比较。我们发现，具有更多参数的模型往往会在更具挑战性的数据集（如SEED和TUEV）上带来一致的改进，这可能是因为它们对复杂模式的建模能力更强。相比之下，TUSL的性能随着模型尺寸的增加而降低。这可能是由于训练数据量有限（245个样本），这使得较大的模型更容易过拟合。

4.4 任务级查询分析

  • 任务自适应查询选择的可视化：为了研究任务自适应查询选择模块的工作机制，我们使用t-SNE在EEG数据集中可视化来自查询路由器R的空间和信道路由得分S。查询分数反映了生成任务自适应查询时的查询选择模式。如图5（a）所示，神经路由分布为时间和空间路由器呈现了一个清晰的任务感知模式，来自不同数据集的样本形成了明显分离的集群。良好分离的分布表明，TQS能够在异构任务中动态选择任务自适应神经LLM对齐的查询。

    此外，我们分析了跨数据集的任务自适应空间查询的选择频率，以揭示每个任务的偏好，如图5（b）所示。结果表明，查询选择呈现出明显的任务依赖趋势。例如，情绪识别和临床事件检测任务（SEED、SEED-IV、TUAB、TUEV）倾向于选择查询P4和P10，而睡眠阶段分类任务（SHHS、HMC）则始终倾向于P5。这些发现表明，P中的查询有效地捕捉了大脑各区域不同的潜在神经特征。

  • 任务之间的神经相关性：为了捕捉任务之间潜在的神经相关性，我们测量EEG任务之间查询选择频率的相似性，以反映神经模式。这里我们分析空间查询标记。如图6所示，我们强调了两个关键观察结果：

    • （1）属于同一任务域的数据集/任务往往表现出更高的相似性得分。例如，TUAB和TUEV（临床）、SEED和SEED-IV（情绪）以及SHHS和HMC（睡眠）显示出非常一致的选择模式。这些发现表明，TQS可以帮助捕获来自类似任务的神经信号背后的一致神经激活模式。
    • （2）来自不同领域的某些任务表现出很强的相似性，表明认知功能之间存在共享的神经机制。例如，临床异常检测任务（TUAB、TUEV）和情绪识别任务（SEED、SEED-IV）显示出相似的空间模式，这是由于双侧时间区域的重叠活动[22,34]，特别是在电极T7和T8周围。同样，由于中枢和额叶区域的神经活动趋同，睡眠分期和认知工作量任务在空间查询使用方面表现出超过75%的相似性[35,36]。这些发现揭示了任务之间潜在的神经相关性，这与现有的神经科学研究一致，暗示了多任务学习共享表征的潜力。

  • 多任务训练中的相互增强：我们通过选择六对具有高相似性得分的数据集，进一步验证了多任务训练中的相互增强效应。如图7所示，联合训练始终比单独训练每个任务带来更好的表现。对于与癫痫相关的任务，TUAB的TUEV提高了5.64%，SEED-IV的SEED提高了5.54%，睡眠数据使工作负荷任务提高了2.66%，TUEV与SEED配对时提高了5.74%。结果表明，UniMind促进了有效的多任务训练，通过共享神经表示提高了不同解码任务的性能。这一改进突显了多任务基础模型在推进广义大脑解码方面的潜力。

5 结论

本文介绍了UniMind，一种用于多任务脑解码的通用EEG基础模型。它利用神经语言连接器将时空EEG特征与LLM对齐，并利用任务感知查询选择模块来适应不同的任务。在10个数据集上的实验表明，UniMind的表现优于之前的模型。我们相信，这项工作为未来在脑电领域研究LLM——大脑交互和多任务学习奠定了坚实的基础。

局限性和未来的工作。这项工作没有研究该模型对完全看不见的主题和新任务的泛化能力，这可能会限制其更广泛的应用。未来的工作将集中在零样本或少热点环境中进行跨学科实验，并整合更广泛的EEG任务，以增强模型的鲁棒性和可推广性。

posted @ 2025-07-22 19:57  NeoAxiomN  阅读(122)  评论(0)    收藏  举报