[NeurIPS 2024] 视觉重建 论文解读:Visual Decoding and Reconstruction via EEG Embeddings with Guided Diffusion
视觉重建论文
Visual Decoding and Reconstruction via EEG Embeddings with Guided Diffusion
老规矩开头先来一段“废话”,讲讲最近的感悟,主要就是最近读论文,发现自己有点囫囵吞枣的感觉,看的很粗糙,很多地方都还没有弄懂。实验的结果也不看。一味的追求速度,我感觉这样不是很好,最近也听老师讲 快就是慢 慢就是快。 我们应该好好沉淀 走好每一步。
一、研究目的
通过非侵入性脑记录(例如功能磁共振成像、脑磁图或脑电图)解码/重建人类看到的视觉世界
二、研究背景
对比学习和生成模型在解码任务(例如图像分类和检索)和生成任务(例如图像重建)中极大地推进了基于功能磁共振成像的视觉解码。
多模态对齐的最新进展使得 MEG/EEG 视觉解码成为可能
** Question One : 什么是对比学习 ?**
对比学习(Contrastive Learning)是一种无监督学习的方法,它的核心思想是通过正负样本对的对比来学习数据的表示。在机器学习中,对比学习通常用于学习数据的嵌入表示,使得相似的数据点在嵌入空间中更接近,而不相似的数据点更远离。
对比学习的基本步骤通常包括:
- 数据增强:对原始数据进行变换,生成正样本对(相似的数据)和负样本对(不相似的数据)。
- 特征提取:使用神经网络或其他模型提取数据的特征表示。
- 对比损失:定义一个损失函数,如InfoNCE损失,来度量正样本对之间的相似度应该高于负样本对之间的相似度。
- 优化:通过梯度下降或其他优化算法来最小化损失函数,从而学习到一个好的特征表示。
具体内容可以看看另外一个博主的博客深入浅出讲解对比学习
** Question: 原文中出现了一个词叫做 zero-shot framework? 这是什么意思?**
“zero-shot framework”指的是一种机器学习或深度学习的框架,旨在在没有特定训练数据的情况下进行任务执行。具体来说,zero-shot学习通常意味着模型可以在没有直接接触过某些类别或任务的情况下,理解和处理这些类别或任务。
在电生理学(EEG)信号的视觉重建中,这种框架可能意味着:
无监督学习:模型能够利用已有的知识来推断新类别的特征,而不需要为每个类别收集训练数据。
泛化能力:模型在未见过的任务或数据上表现良好,能够将从训练数据中学到的知识迁移到新场景。
高效性:能够减少对大量标注数据的需求,适用于数据稀缺或难以获取的情境。
使用 EEG/MEG 数据进行神经解码:以前的研究表明 TemporalSpatial 模块在表示神经数据方面的有效性 [39, 26]。例如,EEGNet 和 ShallowNet [39] 等轻量级卷积神经网络在小型 EEG 和 MEG 数据集中取得了相当大的性能。使用对比学习,已经表明仅使用卷积神经网络和投影层可以在神经数据集上产生令人满意的结果 [6]。最近,Benchetrit 等人提出了一种基于 MEG 的实时视觉感知重建方法 [4]。Song 等人提出了一种使用 ShallowNet Temporal-Spatial 卷积模块的 EEG 编码器,该模块具有一个大型卷积内核,具有一些用于 EEG 嵌入的参数,因此在基于 EEG 的视觉解码方面具有良好的性能 [46]。
以前研究的局限性: 以前的 EEG 研究主要面向理解人脑中的视觉感知,而不是最大限度地提高 EEG 解码性能。因此,视觉解码性能远非最佳。
三、研究内容(主要贡献)
本文提出了一个基于端到端脑电图的视觉重建框架,由一个定制的大脑编码器组成,称为自适应思维映射器(ATM),它将来自不同通道的EEG信号投影到与clip嵌入共享的子空间中,以及两阶段多管道 EEG 到图像生成策略。
-
在第一阶段,EEG被嵌入以对齐高层次的clip嵌入,然后利用先前的扩散模型将EEG嵌入细化为图像先验,同时也从EEG解码出一幅模糊图像,以保持低层次特征。
-
在第二阶段,我们将高层次的clip嵌入、模糊图像和来自EEG潜变量的标题输入到一个预训练的扩散模型中。
四、技术路线
4.1 模型的架构图
4.2 任务描述
设 T T T 代表数据时间窗口的长度, C C C 代表EEG通道的数量, N N N 代表数据样本的总数。我们的目标是从大脑活动数据 E ∈ R N × C × T E ∈ R^{N×C×T} E∈RN×C×T 中派生出EEG嵌入 Z E = f ( E ) ∈ R N × F Z_E = f(E) ∈ \mathbb{R}^{N×F} ZE=f(E)∈RN×F 其中 f f f 是EEG编码器, F F F 是嵌入的投影维度。同时,我们使用CLIP模型从图像I中提取图像嵌入 Z I ∈ R N × F Z_I ∈ R^{N×F} ZI∈RN×F 我们的目标是有效地将EEG表示与图像表示对齐
如图 所示。在训练阶段,EEG 编码器是使用对比学习框架使用脑电图和图像对进行训练。在推理阶段,来自经过训练的脑电图投影仪的脑电图嵌入可用于各种零样本任务,包括基于脑电图的图像分类、检索和图像重建。(这段话对应的图片上边1/3的内容)
4.3 EEG嵌入表示模型:ATM(Adaptive Thinking Mapper)
ATM大致流程
我们提出了脑电编码器ATM,以学习高质量的EEG Embedding。ATM在结构设计上考虑了EEG通道的空间位置(位置编码)和信号的时空特性(通道注意力和时空卷积)。具体而言,ATM结合了Transformer编码器和时空卷积模块,从EEG信号中提取脑电嵌入。首先,将每个通道视为一个patch,并使用一维线性层进行token嵌入。接着,这些EEG token嵌入经过位置编码和自注意力模块处理,生成高维嵌入表示。随后,时空卷积模块(STConv)高效提取EEG特征,并支持灵活替换不同类型的时空卷积组件,以适应各种EEG和MEG数据集。最后,通过MLP投影层确保训练的稳定性,得到最终的EEG嵌入。
怎么理解Shared Token、Subject Tokens和Sequence Tokens呢? 它们又是怎么样发挥作用的?
In addition to entering the original series, we provide an identification input for a known subject and can specifically use this token for downstream tasks. For unknown subjects, we use shared tokens or average all tokens equally directly into the MLP projector.除了输入原始序列外,我们还为已知对象提供了一个识别输入,并可专门将该token用于下游任务。对于未知对象,我们使用shared tokens或将所有token平均直接输入 MLP 投影仪。
Shared Token:在代码中,Shared Token可能不是显式地作为一个单独的变量存在,但它隐含在模型的初始化参数中,作为模型的一部分被初始化,并在处理输入数据时被应用到每个序列上。
Subject Tokens:在代码中,subject_ids参数用于识别每个受试者,然后通过self.subject_embedding(subject_ids)获取对应的受试者嵌入。这个嵌入随后与输入的序列x在序列维度上进行拼接,形成新的输入。
Sequence Tokens:在代码中,x代表输入的EEG信号序列,它通过模型的一系列处理(如时间嵌入和位置嵌入的添加、掩码的应用等)最终被转换成嵌入向量。
我觉得这里可以参考一下BERT用于分类任务的时候,会在输入的最前面加入一个[CLS]的token(但是其实这个位置不一定要在最前面,因为Bert的基本单元还是Transformer,所以无论在哪,[CLS]都能学到全局的信息),这个Token在经过BERT后,被认为是掌握着这个句子的全局信息的这么一个东西,所以在随后直接把[CLS]这个token经过BERT以后的这么一个Embedding作为Linear Classifier的输入直接进行文本的分类操作。
后续论文中其实有讨论说,这个ATM这个Encoder在不同受试者上的表现,说不同被试在这个EEG数据上面还是有差异的。shared Token应该就是在不指定subject的情况下,作为和Bert中的[CLS]差不多的作用。
这里其实还有一个疑问 就是在BERT这个模型中,不止有[CLS]这一个特殊的token,还有一个就是[SEP]这个特殊的token,这个token的意义是用来断句的,就是用来分开一个输入batch中的不同句子,但是这个和[CLS]应该是一样的操作流程,按照上面的说法[SEP]也可以用来作为分类的输入,但是这个只用来断句,所以可能二者在初始值上的设定会有所区别,让[SEP]不参与Transformer的一系列操作。
Embedding Layer的组成
共享的值嵌入(Shared Value Embedding):
这是模型中的一个关键部分,用于将输入的EEG信号转换成高维空间中的点,从而提取出有用的特征。
在联合训练模式下,模型会根据每个受试者的ID选择相应的值嵌入,这意味着模型可以为每个受试者学习特定的表示。
在独立训练模式下,模型使用一个共享的值嵌入,适用于所有受试者。
时间嵌入(Temporal Embedding)和位置嵌入(Position Embedding):
这些嵌入用于向模型提供关于EEG信号时间动态和空间位置的信息。
时间嵌入帮助模型理解信号随时间的变化,而位置嵌入则提供关于信号在序列中位置的信息。
受试者嵌入(Subject Embedding):
这是模型中一个可选的部分,用于捕捉不同受试者之间的差异。
通过受试者ID获取相应的嵌入,并将其与输入的EEG信号拼接,从而为模型提供受试者特定的信息。
掩码嵌入(Mask Embedding):
用于处理输入序列中的缺失或无关部分,确保模型专注于重要的信号。
掩码嵌入通过将掩码应用到输入上,将不需要的部分替换为特殊的掩码token。
Question Two: TS Conv的主要内容是什么?
这个是代码中关于TS Conv模块的描述。不知道这里的时空卷积和EEGformer中的时空卷积的区别和联系在哪? 这里的TS卷积层有两种方式,后面的实验数据对两个模块的替换做了相应的分析。时空卷积模块是用少量参数表示脑电图数据的有效方法
4.4 Image Embedding
在本研究中,我们使用 CLIP 进行图像嵌入,用 EEG 嵌入表示为 Z I ∈ R N × 1024 Z_I ∈ \mathbb{R}^{N×1024} ZI∈RN×1024,而不是 Z I ∈ R N × 257 × 768 Z_I ∈ \mathbb{R}^{N×257×768} ZI∈RN×257×768。
clip模型的基本原理
CLIP(Contrastive Language-Image Pre-training)模型是由 OpenAI 开发的一种多模态神经网络,它通过对齐大量图像和文本数据来学习视觉表示,使得模型能够理解和生成与文本描述相匹配的图像。CLIP 的核心思想是利用对比学习(contrastive learning),通过比较正样本对(匹配的图像和文本)和负样本对(不匹配的图像和文本),来学习图像和文本的联合表示。
总结一下,就是将text和image分别通过Bert和Vit,分别获取文本的Embedding和图像的Embedding,将配对的(上图对角线)上的两个Embedding的内积最大化
具体的原理可以参考一位博主的一篇博客CLIP对比学习 语言-图像预训练模型
这里需要更正一下我之前的一个错误的地方,就是我原来以为clip模型做的事情是让 text embedding 和 image embedding 变成一个共享的embedding.但是事实却不是这样的,clip模型只是训练了两个encoder,使得两个encoder之间的输出能有一个很好的映射关系(就可以通过文本看看那个image最适合,或者通过image 看看哪个text最适合),并不能通过text embedding得到image embedding。要实现这个过程那就得重新自己训练一个模型。
4.5 EEG Guidance Image Generation
在本研究中,我们设计了一个两阶段的图像生成流程,目的是利用脑电图(EEG)记录生成相应的视觉刺激图像。具体过程如下:
第一阶段 - 先验扩散:
- 利用脑电图编码器(ATM)将图像转换为EEG嵌入(zE)。
- 训练一个以EEG嵌入为条件的扩散模型(εprior),用于学习CLIP嵌入(zI)的分布。
- 该扩散模型根据EEG嵌入生成对应的CLIP嵌入,作为第二阶段生成图像的先验。(作为第二阶段的先验 很疑惑什么意思)
第二阶段 - CLIP引导的图像生成:
- 使用预训练的SDXL模型和IP-Adapter模型来构建生成器,根据CLIP嵌入(zI)生成图像(I)。
- 引入img2img技术来引入图像的低级特征。
Question: GIT model img2text、Ip adapter、VAE的内容都没有详细说明?
VAE
VAE的原理还是比较复杂(公式比较多,后续了解后我会在这做更新)
IP adapter
IP adapter的论文解读可以参考我的这篇博客可控图像生成 论文解读 IP-Adapter: Text Compatible Image Prompt Adapter for Text-to-Image Diffusion Models,这里可以简单的介绍一下,IP adapter你可以把它理解为Stable Diffusion的一个插件,他的作用就是为Stable Diffusion多添加了一个图像prompt(原生Stable Diffusion实际上只有Text Prompt),他的实现过程就是讲图像经过clip图像编码器后再经过线性映射和一个LN网络后直接加入到一个Cross Attention中,让他的输出和原生Stable Diffusion的Cross Attention的输出相加代替原来的输出。
Question: EEG Embedding code和 VAE的latent code中间用MAE LOSSS连起来是什么意思?
我们希望通过与 VAE 的潜在空间对齐,从 EEG 中恢复轮廓、姿势、方向和其他像素级信息的基本信息。应该就是利用EEG-Image配对数据,将EEG通过ATM获得Embedding,将图像经过VAE的image Encoder,再利用MSE损失函数对EEG Embedding和VAE Latent进行对齐。
我们对Low-level pipeline进行了 200 次历时训练,尝试了潜均方误差(MSE)损失、对比学习损失和变异自动编码器(VAE)图像重构损失,以便将从投影层和上采样 CNN 获得的 4 × 64 × 64 EEG 潜像与 VAE 潜像对齐。然而,重构损失或对比学习损失比仅应用潜空间损失的性能更差,而且需要的 GPU 内存也更多。
Question: 怎么理解先验知识
Question: 其实不理解的地方还是那个用很多个Unet结构去把EEG转为Clip的image Embedding具体是如何做到的
在图 3 的左侧,我们通过 EEG 编码器 ATM 获得了每个图像的 EEG 嵌入 z E z_E zE。现在我们的目标是使用这些 EEG 嵌入来生成相应的图像。图像、脑电图嵌入和图像嵌入的联合分布可以表示为 p ( I , z E , z I ) = p ( z I ∣ z E ) p ( I ∣ z I ) p(I,z_E,z_I)=p(z_I|z_E)p(I|z_I) p(I,zE,zI)=p(zI∣zE)p(I∣zI),分别对应于先验扩散和 CLIP 引导生成。在第一阶段,我们首先关注先前的扩散阶段。受 DALL-E 2 [37] 和 Mind’s Eyes [42] 的启发,我们训练了一个以 EEG 嵌入 Z ^ E \hat{Z}_E Z^E 为条件的扩散模型,以学习 CLIP 嵌入 p ( z I ∣ z E ) . p(z_I|z_E). p(zI∣zE). 的分布。 在这个阶段,我们构建了一个轻量级的 U-Net: ϵ prior ( z I t , t , z E ) , \epsilon_\text{prior}(z_I^t,t,z_E), ϵprior(zIt,t,zE),,其中 z I t z^{t}_I zIt 表示扩散时间步长 t t t 处的含噪 CLIP 嵌入。我们使用 EEG 和 CLIP 嵌入来训练先验扩散模型。通过这个扩散模型,我们可以从 EEG 嵌入中生成相应的 CLIP 嵌入 z I z_I zI,作为第二阶段的先验。
这里可以去阅读一下Hierarchical Text-Conditional Image Generation with CLIP Latents这篇论文(就是DALL E2或者说unclip这个模型),这里的这个结构应该就是借鉴了DALL E2的这个思想。原文也是利用clip的text-image配对数据去训练了这么一个模型来获取image embedding,这样做的原因是这样显式地生成图像表征可提高图像多样性,同时将逼真度和标题相似性的损失降到最低。
4.6 Loss Function
作者对损失函数采用双重方法,服务于不同的目标。对于分类和检索任务,我们仅利用 CLIP 损失,该损失函数有助于将 EEG 数据 E 与相应的图像数据 I 对齐。对于生成任务,除了 CLIP 损失之外,我们还添加了均方误差(MSE)损失以促进回归中的一致性学习。因此,模型的整体损失函数是这两种不同损失类型的组合,表示为:
L
o
s
s
=
λ
⋅
L
C
L
I
P
+
(
1
−
λ
)
⋅
L
M
S
E
Loss = λ · L_{CLIP} + (1 − λ) · L_{MSE}
Loss=λ⋅LCLIP+(1−λ)⋅LMSE
这里,
λ
λ
λ是一个超参数,用于平衡每种损失类型的贡献。
五、实验结果
5.1 Training and Computational Considerations
在THINGS-EEG数据集上,我们进行了实验以验证ATM在处理电生理数据的能力,并在THINGS-MEG数据集上测试了其在MEG数据模态上的表现。实验在NVIDIA RTX 4090 GPU上进行,使用Adam优化器训练跨被试和同一被试模型,分别在496,200个和66,160个样本上,初始学习率为0.003,批次大小分别为16和1024,初始温度参数为0.07。我们在每个训练周期后对zero-shot测试数据集进行评估,并保持所有模型超参数一致。我们比较了不同编码器在受试者内和跨受试者测试集上的性能(详见附录H)。
初始温度参数是什么?
5.2 EEG Decoding Performance
我们的方法获得用于分类任务的脑电图嵌入。我们输出与文本嵌入的余弦相似度最高的 EEG 类别(如下如所示)。C显示了受试者中不同方法的平均准确性,并表明我们的方法优于其他方法。有关基于 EEG 的图像分类的更多详细信息,请参阅附录 B。
这张图描绘的是ATM将EEG转为EEG embedding后用于下游图像以及文本检索任务中的效果,可以发现在图像检索(通过计算 EEG 和 CLIP Embedding的余弦相似度)上,几乎都能够把相应的刺激图像检索出来(大部分在预测的前五名之中),图cd分别展示了不同模型在分类任务和检索任务上的平均准确度,可以发现本文提出的ATM EEG Encoder在这些下游任务中都有着不俗的表现(第一)。
其实这里还有一个问题 为什么图像分类用的是EEG和text的embedding的余弦相似度? 而图像检索用的是EEG和clip的嵌入? 而且这里的CLIP的嵌入指的是CLIP的text的嵌入还是image的嵌入?
这里图像检索用的是clip的image Embedding [This implies that by changing the text embeddings of image labels to image embeddings, we can transition the task from classification to image retrieval.]
ATM 消融研究系统地解构和分析了脑电图投影仪的每一层。作者对 ATM 中的每个组件(即 MLP 投影仪、时空卷积模块和通道注意块)进行了消融研究。我们指定了两种不同的卷积架构,ShallowNet (ATM-S) 和 EEGNetV4 (ATM-E),作为我们的卷积主干。附录 B.3 显示了在不同实验配置下获得的结果。(那按结果来说就是ATM-S更好,就是用ShallowNet的效果会好一些)
5.3 Image Generation Performance
图 6a 显示了在 EEG 嵌入的指导下生成图像的过程,并评估了生成图像的质量。为了评估生成性能,我们进行了一项图像检索任务。具体来说,我们提取生成图像的 CLIP 嵌入,并比较所有图像的 CLIP 嵌入之间的相似性,以检索生成的图像。
图 6b 显示了分布的相似性。图 6c 显示了生成的样本。生成的图像与看到的图像具有很高的语义相似性,并且在低级视觉特征方面具有良好的多样性,这可以通过引导量表超参数进行操作(图 6d)。我们还在表 1 中报告了来自不同数据集的各种指标的 EEG 、 MEG 和 fMRI 的解码和重建性能。
这个表我看的就不是很懂了
我们的研究使用各种指标来评估我们从大脑数据(EEG、MEG、fMRI)中重建视觉刺激的能力(表 1 和表 7)。这些指标包括 PixCorr(像素相关性,真实值和重建之间)、SSIM(结构相似性指数指标)[53]、SwAV(SwAV-ResNet50,指平均相关距离)[5]和使用神经网络的双向识别(AlexNet(2/5)、Inception、CLIP。这里 AlexNet(2/5) 是 AlexNet 的第 2 层和第 5 层特征层),用于低级和高级图像特征。在这里,双向识别可以看起来像是 [33] 中描述的双向检索任务。在表 1 中,我们的结果表明,在 THINGS 数据集上,我们可以在使用 ATM 的脑电图重建上实现优于 MEG 的性能。表 6 显示了不同数据集(fMRI、MEG、EEG)在视觉刺激任务上的解码性能,我们甚至实现了与 fMRI 和 MEG 相同或更好的性能。我们的结果表明,合适的神经表示在下游任务中起着决定性的作用。
5.4Temporal Analysis
为了研究不同脑电图时间窗对视觉解码的影响,作者计算了滑动时间窗和生长时间窗的平均top-1分类精度。滑动时间窗涵盖从视觉刺激开始到时间点t的整个时期,而生长时间窗则包括时间点t之前100ms的数据。结果显示,在视觉刺激后500ms内,准确率达到约30%的上限,之后不再提高,且MEG解码表现出类似的趋势。我们发现不同类别图像的最佳重建时间窗口存在差异.
5.5 Spatial Analysis
为了进一步探讨不同大脑区域对视觉解码的贡献,作者将THING-EEG数据中的EEG电极分为额叶、颞叶、中心、顶叶和枕叶五个区域。消融实验显示,使用所有大脑区域的信息效果最佳,枕叶区域在检索精度和重建性能上表现优异,而顶叶和颞叶区域包含一定的语义信息,额叶和中央区域的贡献则较小。
5.6 Representational analysis
如图 12 所示,我们展示了潜在空间中的表征相似性矩阵和可视化。为了研究从 EEG 获得的表征与图像表征之间的关系,我们进行了一个表征相似性矩阵。我们专注于受试者 8,它表现出最高的检索准确性。通过将聚类算法应用于与测试集中 200 张图像对应的图像嵌入,我们观察到不同的类别内聚类。我们根据图像和文本嵌入生成了相似性矩阵,然后将其与 EEG 表示进行比较。如图 12 所示,在带有图像的表征相似性矩阵中可以观察到明显的类别内聚类,而这种现象在带有文本的表征相似性矩阵中不存在。
这几张图似乎看不太懂
5.7 Concept analysis
我们采用了Wei等[55]提出的概念嵌入编码器,它将原始图像的剪辑嵌入编码为42维向量,每个维度代表一个不同的概念。利用 ATM 直接投影测试集中 200 个类别的 EEG 数据,我们获得了 EEG 嵌入,然后将其与冻结权重一起馈送到概念编码器中,产生 200 个概念嵌入,如图 13 所示。对概念层面的表征相似性的分析表明,从我们的 EEG 投影仪派生的 EEG 嵌入有效地与概念空间保持一致。这确保了高度对齐的语义一致性,为从 EEG 数据重建图像提供了令人信服的证据。
六、实验结果图像及未来可能研究方向
THINGS-EEG 数据集上零镜头检索的总体准确性。我们显示了在批量大小 = 1024 的情况下,受试者内和跨受试者检索任务性能 (Ave ± Std.%)。我们比较了来自不同脑电图嵌入方法的 2 路、4 路、10 路、Top-1 和 Top-5 的 200 路精度。我们的 ATM 优于所有其他 ATM。
如何评价生成图像的准确性?
根据相应的脑电图从 10 个受试者生成 10 张不同的图像。随后,将对生成的每个图像执行图像检索。计算前 1 名和前 5 名的准确度。它有助于评估生成的图像与其原始对应图像之间的语义一致性。
作者展示了他们生成的图像中效果比较好、中等和较差的图像。
EEG 数据中的视觉信息主要包含在 200-400ms 范围内,与之前的 EEG 研究一致。
有趣的现象和未来的方向:首先,跨对象和对象内设置之间存在不可忽视的性能差异。这种表现差距源于脑电图信号的受试者间差异,可能归因于个体大脑的异质性、个体之间视觉感知的差异,甚至脑电图记录期间噪声分布的变化。因此,需要在脑电图编码器上做出更多努力,例如更灵活的神经网络架构或更好的预训练模型权重初始化。迁移学习和元学习也是未来值得探索的方向。此外,在预训练大型脑电图模型时,如何统一不同脑电图数据集的各种电极蒙太奇是一个挑战。脑电图源定位将传感器级脑电图信号转换为标准脑源空间 ,可能是一个潜在的解决方案。
七、感想
开放原子开发者工作坊旨在鼓励更多人参与开源活动,与志同道合的开发者们相互交流开发经验、分享开发心得、获取前沿技术趋势。工作坊有多种形式的开发者活动,如meetup、训练营等,主打技术交流,干货满满,真诚地邀请各位开发者共同参与!
更多推荐
所有评论(0)