基于3D注意力卷积与自监督学习的脑疾病分类方法

2024-03-10冀俊忠雷名龙

北京工业大学学报 2024年3期

冀俊忠, 于乐, 雷名龙

(1.北京工业大学北京人工智能研究院, 北京 100124;2.北京工业大学信息学部多媒体与智能软件技术北京市重点实验室, 北京 100124)

近年来,脑研究引起了人们的广泛关注。功能磁共振成像(functional magnetic resonance imaging,fMRI)是神经影像学中一种重要的成像方法,它能够安全无创地记录和呈现大脑活动,且具有良好的空间与时间分辨率,因而被广泛地应用于脑功能研究与脑疾病诊断等领域[1]。功能磁共振成像的原理是利用磁共振造影来测量神经元活动所引发的血液动力学变化,从而产生多个时刻的三维脑成像数据。基于fMRI数据的脑疾病分类是脑科学领域中重要的研究方向,其目的是通过数据分析的手段,有效地判断受试者是否患有相应的脑疾病。由于fMRI数据能够反映大脑中的活动状态与模式,分类fMRI数据不仅有助于理解孤独症谱系障碍[2]、注意缺陷多动障碍[3]和阿尔兹海默病[4]等脑疾病的发病机理,而且有望揭示脑疾病的生物标志物[5],为临床的辅助诊断提供重要的决策支持。

目前,fMRI数据的分类方法主要包括2种[6],分别为基于脑区的分类方法和基于体素的分类方法。其中,基于脑区的分类方法采用大脑划分的思想,通常针对fMRI三维图像中的每2个脑区来构建功能连接体[7-9]。但是,脑区级别的分类将体素点特征转化为脑区特征,忽略了大部分体素点的特征信息和体素点之间的空间位置关系[10],这可能会丢失全脑体素点信息和整个大脑的细粒度功能连接信息。而基于体素的分类方法[11-12]可以考虑到全脑体素的信息以及体素点之间的空间位置关系,在当前的分类任务中表现出了更好的效果。

由于基于体素的分类方法需要考虑细粒度的空间信息,因而对分类模型的复杂度提出了更高的要求。近年来,深度神经网络在许多fMRI数据分类任务中取得了很好的效果[13-14],其强大的数据表征能力能够较好地建模体素级的特征,产生了许多适用于不同数据形式的神经网络模型。其中,最基本的方法就是直接对体素数据进行分类。例如,早期模型将全脑体素矢量化为一维向量,利用全连接神经网络(fully connected neural network,FCNN)等传统深度学习方法[15-16]进行特征学习与分类。由于矢量化破坏了体素数据的空间依赖关系[17],后续方法[17-18]将卷积神经网络(convolutional neural network,CNN)扩展到三维空间中,利用三维卷积神经网络(3-dimensional convolutional neural network,3DCNN )对三维体素信息进行分类,从而较好地保留了空间信息。与直接对体素进行分类的方法不同,一些方法[19-20]从fMRI体素的角度出发,通过对大脑中的局部空间关系与整体空间关系进行编码进一步构建了基于体素的全脑功能连接映射图。该类方法构建的功能连接具有三维特性,因此,也常采用3DCNN模型来实现特征提取与脑疾病分类。然而,常规的3DCNN难以有效利用具有更强关系的体素,对空间信息的提取能力仍然有待提高[21]。

综上,常规FCNN与CNN模型难以捕获fMRI体素之间的空间信息。3DCNN模型为分类具有空间信息的fMRI数据提供了新的解决方法。但是,基于3DCNN的体素级脑疾病分类方法目前存在以下不足。首先,由于三维空间中体素-体素之间以及体素-脑区之间关联性的重要程度存在差异,现有3DCNN模型在提取空间特征的过程中未考虑这种差异性,导致模型难以准确地关注到fMRI数据中对当下脑疾病分类最为有效的空间信息;其次,基于体素的fMRI数据具有较为复杂的空间依赖关系,仅利用标签作为监督信息无法充分挖掘数据内丰富的空间信息。此外,通过提升模型的复杂度来增加特征提取能力也容易导致深度模型的过拟合,从而影响分类模型的性能。

为了解决上述问题,本文提出一种基于3D注意力卷积与自监督学习融合的脑疾病分类方法。旨在充分挖掘fMRI数据的三维空间特征以实现更有效的分类。首先,针对三维空间中不同体素的差异性,本文将注意力机制引入3DCNN网络,通过区分三维空间中不同位置的重要性来得到更适用于当下脑疾病分类任务的特征;其次,为了进一步挖掘体素的空间依赖关系,本文提出一种基于辅助任务的自监督学习方法,以鼓励卷积神经网络充分挖掘fMRI数据的三维空间特征;最后,通过辅助任务与目标分类任务的联合训练实现多任务学习,进而在有限的训练数据下提高目标分类任务的性能。在2个数据集上的实验结果表明,该方法能够有效地提升脑疾病分类的准确率。

1 基于3D注意力卷积与自监督学习的脑疾病分类方法

本文提出了基于3D 注意力卷积与自监督学习的脑疾病分类方法,整体框架如图1所示。该模型主要由3个部分组成:第1部分为如图1(a)所示的基于3D注意力卷积的特征提取模块,用于关注体素数据中不同空间位置的重要性差异。第2部分为如图1(b)所示的基于自监督学习的辅助任务模块,通过反卷积来预测三维空间中的体素信息,以鼓励神经网络进一步挖掘体素的空间依赖关系。第3部分为如图1(c)所示的目标分类任务模块,用于3D注意力CNN的输出特征通过分类模块映射为最终的分类预测结果。目标分类任务与自监督辅助任务采用联合训练的方式进行参数优化,分别以基于体素的全脑功能连接映射图和掩蔽后的结果作为输入,通过共享的3D注意力CNN获取隐层特征后,再分别输入到2个任务中得到最终损失。

图1 基于3D注意力卷积与自监督学习的脑疾病分类方法框架Fig.1 Architecture of brain disease classification based on 3D attention CNN and self-supervised learning

1.1 基于3D注意力卷积的特征提取

本文将基于体素的全脑功能连接映射图作为输入,从而能够保留大部分的体素信息,并进一步考虑体素与脑区之间的功能连接关系。输入数据的构建过程主要包含体素-脑区相关性计算与三维空间信息还原2个部分。

体素-脑区相关性计算方法如下。三维脑成像数据可以表示为R∈RT×D1×D2×D3,其中T表示时间序列长度,D1、D2、D3分别表示脑成像数据的3个维度,体素数为D1×D2×D3=D。为了计算体素与脑区之间的相关性,本文选用CC200大脑模板[22]将fMRI数据的三维图像划分为M个脑区。基于所得脑区计算每个脑区中所有体素点的平均时间序列U∈RM×T,同时提取每个体素的时间序列并将其重新组织为一个二维矩阵V∈RD×T。给定时间序列u=[u1,u2,…,uT],v=[v1,v2,…,vT],u与v之间的皮尔森相关系数(Pearson correlation coefficients,PCC)可以表示为

(1)

基于体素与脑区之间的相关性,本文进一步将其还原成具有空间特性的三维结构。具体地,将相关系数矩阵P按脑区展开成立方体,立方体中的每个元素代表当前体素与脑区之间的相关性,按原来的空间位置关系进行还原,能够得到M个维度为D1×D2×D3=D的全脑体素功能连接映射图,可将其表示为X∈RM×D1×D2×D3。由带有时间维度的三维全脑体素fMRI数据生成多通道的三维体素级功能连接数据,既完整地保留了原始体素信息,又提供了细粒度的大脑功能连接信息。

为了提取基于体素的全脑功能连接映射图中的空间信息,本文采用CNN作为特征提取器,探索体素数据中的局部关联性。由于fMRI体素数据存在于三维空间中,传统的卷积神经网络无法准确地描述数据中的三维空间关系,基于此,3DCNN模型具有三维卷积核,能够有效地在三维空间上挖掘fMRI数据的局部特征。

(2)

进而,带有激活函数的3DCNN层可以定义为

H(l+1)=σ(K(l)⊗H(l)+B(l))

(3)

为了考虑fMRI数据中不同体素的重要性差异,本文在一般3DCNN模型的基础上引入了注意力机制,进而为不同的体素分配不同的权重。具体来说,本文利用体素间的空间关系来生成空间注意系数,空间注意力系数可以关注到空间中含有关键信息的位置,这里利用Woo等[23]提出的方法来生成空间注意力系数,其计算方法可以表示为

(4)

虽然注意力卷积层能够提升三维卷积的特征提取能力。但是,多个卷积层的堆积可能会产生梯度消失的现象。因此,本文进一步设计了具有残差结构的3D注意力卷积层,保留初始输入特征在三维卷积中的作用,增加特征的判别性。带有残差结构的3D注意力卷积层可以定义为

H(l+1)=C3d(H(l))+C3d(H(l))⊙A(l)

(5)

给定了具有残差结构的3D注意力卷积层后,通过层间的传递方式,可以把L层的3D注意力CNN用Ff(·)表示。那么,对于输入数据X,第L层3D注意力卷积的输出可以用Z来表示,即

Z=Ff(X)

(6)

1.2 基于自监督学习的辅助任务

由于人脑功能高度复杂,神经影像数据通常样本少、维度高,并且存在大量噪声,故仅通过3DCNN模型难以充分挖掘fMRI数据中丰富的空间信息。研究表明,从额外的辅助任务中学习通常会得到鲁棒的特征表示。因此,本文基于三维空间信息的特性提出了一种自监督辅助任务(见图1(b)),用于辅助目标任务的特征学习过程。

自监督任务的重点在于设计适合数据与目标任务的辅助任务。对于基于体素的全脑功能连接映射图数据,其空间结构含有丰富的位置信息。为此,本文基于三维空间连续性假设为目标任务设计了一个自监督辅助任务——预测缺失体素,该任务的本质是对体素进行重构。与自编码器等重构数据本身的任务不同,本文的辅助任务采用缺失的数据预测完整的原始数据,该过程有助于发现数据中的模式,提升模型的特征学习能力。在训练辅助任务的过程中,模型需要从原始数据中深度挖掘三维空间信息,该过程对目标分类任务有很大的帮助。此外,这种方法也可以看作是在训练数据中添加噪声以避免模型产生过拟合。

辅助任务将随机掩蔽后的全脑功能连接映射图X∈RM×D1×D2×D3作为输入。随机掩蔽过程需要依据掩蔽的比率α∈[0,1]来确定每个体素是否被掩蔽,然后生成一个与X维度相同的掩码张量M∈[0,1]M×D1×D2×D3来随机遮挡内容,最终依靠掩码张量来生成残缺的数据∈RM×D1×D2×D3。该过程可以表示为原始数据X与掩码M之间的运算,即

=X⊙M

(7)

=Ff()

(8)

可以把辅助任务中使用的反卷积神经网络定义为Fs(·),那么自监督辅助任务的输出为

′=Fs()

(9)

自监督辅助任务的损失函数为

(10)

1.3 联合优化框架

本文采用联合训练的方式,通过自监督辅助任务与目标分类任务的联合优化,构建了一个多任务的学习框架。

在该框架中,分类任务如图1(c)所示,将3D卷积神经网络的输出特征Z展平拼接成一维向量,再输入到分类模块中,本文的分类模块由全连接神经网络构成,首先利用多层全连接网络来进行特征整合并实现高级抽象特征的挖掘,最后一层使用Softmax函数来得到最终的分类结果。可以将分类模块用函数Fc(·)来表示,那么目标分类任务最后的输出为

Y′=Fc(Z)

(11)

目标分类任务的损失函数为

(12)

式中:N代表训练集的样本数;C代表类别标签的个数;Yi,j和Y′i,j分别代表样本的标签与相应的预测值。

因此,本文所提方法总的目标函数可以定义为

L=Lmain+λLss1

(13)

式中λ为控制损失权重的系数,用于调整自监督辅助任务对分类任务的影响。在反向更新参数时,自监督辅助任务与目标分类任务2个部分均可对3D注意力卷积神经网络的参数产生影响,从而使得网络能够学习到更多的空间信息。通过这种联合训练的方法,可以使3D注意力卷积网络具有更强大的表达能力,从而提升分类任务的效果。

2 实验结果及分析

2.1 实验数据与预处理

本文采用孤独症脑影像数据ABIDE-Ⅰ和ABIDE-Ⅱ进行实验。经过质量评估,分别选择871和518名被试构建数据集。数据预处理分别采用连接组分析(configurable pipeline for the analysis of connectomes, CPAC)和数据处理助手 DPARSF(data processing assistant for resting-state,FMRI)完成。

2.2 实验设置和评价指标

实验在处理器为AMD Ryzen 9 3950X、显卡NVIDIA 3090、RAM为128 GB、操作系统为Ubuntu 18.04.1的环境下,利用Python编写代码并实现。

实验的具体参数为:D1=30,D2=36,D3=30,M=200,T=77、115、145、151、175、195、205、235、245、295、315(数据集中不同采集站点的时间序列长度有所不同)。3D注意力CNN包含3个卷积层,其特征数量分别为128、32、32,卷积核大小为3×3×3,在第2层3D注意力卷积后使用池化核大小为2×2×2进行池化;计算空间注意力的3D卷积采用一个卷积层,卷积核大小为5×5×5。辅助任务中的反卷积也相应包含3个反卷积层,其特征数量分别为32、128、200,反卷积核大小为3×3×3,在第1层反卷积后使用池化核大小为2×2×2进行反池化。全连接神经网络为2层,维度分别为1 024和2。控制重构任务与分类任务的损失权重系数为100。模型参数量为5.8×107,训练时间为1 380 s。训练过程中使用Adam梯度下降算法最小化损失函数,每次迭代的批大小为24,学习率为1×10-4。实验以8∶1∶1的比例随机划分为训练集、验证集和测试集,采用十折交叉验证对实验性能进行评估,采用广泛应用的准确率(accuracy,ACC)、灵敏度(sensitivity,SEN)、特异度(specificity,SPE)、正预测率(positive predictive value,PPV)、负预测率(negative predictive value,NPV)和综合指标(f1-measure,F1)来评价模型的分类性能。

2.3 参数分析

为了分析不同掩蔽比例α对实验结果的影响,本文在ABIDE-Ⅰ数据集上取不同的α进行实验,除α外,其余超参数设置相同,表1给出了不同掩蔽比例对分类精度的影响。

表1 ABIDE-Ⅰ上不同掩蔽比例的实验结果

可以看出,当α=0.3时分类准确率达到最高。当α=0.1时,对掩蔽后3D数据进行恢复的自监督辅助任务的难度较小,因此神经网络不能很好地挖掘fMRI数据的空间信息。而随着α的增加,自监督辅助任务的难度逐渐增大,分类性能也逐步提升。当α>0.3时,α的增加使得分类性能出现降低的趋势,其原因是过大的掩蔽比例会导致模型难以恢复原始的数据,进而影响分类性能。因此,后续实验将比例设为0.3。

2.4 消融实验

为评估模型中3D注意力卷积(attention convolution,ATT)和基于自监督学习的辅助任务(self-supervised learning,SSL)2个部分的有效性,本文在ABIDE-Ⅰ与ABIDE-Ⅱ数据集上进行了相关的消融实验,BASE表示同时移除注意力机制与自监督辅助任务的模型,BASE-SSL表示只移除注意力机制的模型,BASE-ATT表示只移除自监督辅助任务的模型,3D-SACNN表示完整的模型。结果如表2 所示。

表2 在ABIDE-Ⅰ与ABIDE-Ⅱ上的消融实验结果

可以观察到,ATT和SSL都可以提高分类精度。这一结果证明注意力机制可以通过学习特征权重,从而关注到更重要的空间特征,而自监督辅助任务能够更好地挖掘数据内丰富的空间信息,使得模型学习到更完备的fMRI数据特征,从而提升分类效果。此外,同时包含注意力机制和自监督辅助任务的模型取得了比2个消融实验模型更好的分类性能,证明了所提模型的有效性。

2.5 对比实验

为了验证所提方法的有效性,本文分别在ABIDE-Ⅰ和ABIDE-Ⅱ数据集上选择了8种基于fMRI数据的脑疾病分类方法进行了对比实验。其中,LASSO、 Ridge和SVM为基于传统机器学习的方法,而SSAE[7]、BrainNetCNN[8]、CKEW[9]和3DCNN[19]为基于深度学习的方法。SAE、BrainNetCNN和CKEW方法将脑区级的功能连接作为输入,其余方法将体素级的连接数据作为输入,结果如表3所示。

表3 8种方法在ABIDE-Ⅰ与ABIDE-Ⅱ上的分类性能对比

本文所提方法3D-SACNN在多数指标上达到最优,分类效果优于其他方法。在这些方法中,将体素级数据作为输入的方法优于将脑区级数据作为输入的方法,原因在于体素级数据包含细粒度的信息,采用浅层的SVM即可得到较好的实验结果。在以体素级数据作为输入的方法中,基于深度学习的方法显著优于基于传统机器学习的方法。以ABIDE-Ⅰ数据集为例,基于传统的机器学习方法由于模型结构较为简单,因此难以有效提取fMRI数据中有意义的高阶信息,其准确率均低于68.5%。与之相比,深度学习的方法以其强大的拟合能力在处理高维fMRI数据方面展现出了巨大优势,其准确率均高于68.5%。此外,与其他几种深度学习方法相比,本文提出的3D-SACNN方法明显优于BrainNetCNN、SAE、CKEW和3DCNN,在ABIDE-Ⅰ和ABIDE-Ⅱ这2个数据集上ACC指标分别达到了70.36%和67.46%,其原因主要在于2个方面:1) 3D注意力卷积神经网络可以挖掘到fMRI数据中具有差异的空间信息,从而关注到更重要的空间特征。2) 自监督学习可以进一步辅助3D注意力神经网络学习到更丰富的空间特征表示,从而进一步提升分类准确率。

2.6 异常脑区分析

本节采用显着图[25]的可视化方法来分析本文模型得到的能够判断被试是否患有疾病的重要的特征,这些特征即为模型预测出的可能与疾病相关的异常脑区。通过这些脑区与当前生物学研究的一些发现进行对比,来进一步证明分类结果的可解释性。

显着图的可视化方法使用训练后神经网络的反向传播来计算输出相对于输入数据的梯度。因此,可以进一步计算不同脑区内所有体素的平均梯度值,从而计算脑区级的显著性。具体来说,给定一个输入数据P∈RM×D,对于输出oc,输入P中各元素对该输出的重要程度W∈RM×D,可以通过输出对输入的梯度绝对值来计算,即

(14)

可得到输入数据对分类的贡献程度值,计算出的取值大小能够反映该特征对分类结果贡献程度的高低。为了获得脑区水平的显著性s∈RM,可对W在D维取平均值,即

(15)

式中avg表示平均操作。

通过以上方法,能获得所有脑区对分类的贡献程度值,本文选择出前10个贡献程度最高的脑区进行可视化展示,如图2所示,这些脑区即为由本文所提方法进行分析而得到的前10个可能与脑疾病相关的重要脑区。其中一些脑区位于默认模式网络(default mode network,DMN)中,如左颞中回(left middle temporal gyrus,MTG.L)、右颞中回(right middle temporal gyrus,MTG.R)、右颞上回(right superior temporal gyrus,STG.R)、右楔前叶(right precuneus,PCUN.R)和左内侧额上回(left superior frontal gyrus, medial,SFGmed.L)。右丘脑(right thalamus,THA.R)是突显网络(salience network,SN)中的脑区,相关研究中发现,这些脑区可能与孤独高度相关[26]。左背外侧额上回(left superior frontal gyrus, dorsolateral,SFGdor.L)位于执行控制网络(executive control network,ECN),孤独症组在执行心理旋转任务时,SFGdor.L的活动较少[27]。右岛盖部额下回(right inferior frontal gyrus, opercular part,IFGoper.R)和右三角部额下回(right inferior frontal gyrus, triangular part,IFGtriang.R)位于镜像神经元系统(mirror neuron system,MNS), 都属于额下回,研究发现孤独症儿童在模仿和观察情绪表达时,在额下回没有表现出镜像神经元活动[28]。此外,前额皮质和颞皮质的高级大脑功能直到青年末期和成年早期才完全成熟[29],孤独症患者的这些区域在成熟过程中受到损害,从而导致社交、语言等功能的障碍,MTG.L和MTG.R是颞皮质的一部分,SFGdor.L、IFGoper.R和IFGtriang.R位于前额叶。此外,右额中回(right middle frontal gyrus,MFG.R)也被认为是与孤独症最相关的一些重要脑区,这些脑区的异常与孤独症关系密切。