基于注意力和Transformer的阿尔兹海默症分类

汪悦恺，王文伟，孟慧茹

(武汉大学 电子信息学院，湖北 武汉 430072)

0 引言

痴呆是以认知功能受损为核心的综合征[1]。阿尔兹海默症(Alzheimer’s Disease,AD)是最常见的病因,占痴呆患者的60%～80%[2],在中国尤为普遍[3]。作为一种不可逆的神经退行性疾病,随着人均寿命延长和人口老龄化加剧,痴呆给患者和家属甚至医疗体系带来巨大负担。临床通过影像、量表和病史诊断AD,但这些方法效率低、可重复性差,依赖医生的经验和能力,且脑结构改变或损伤早于临床症状的出现[4],人工诊断容易漏诊和误诊。因此设计高效、便捷的方法诊断AD对于疾病预防和改良治疗方案具有重要的价值。

磁共振成像(Magnetic Resonance Imaging,MRI)利用磁共振原理无创地从人体中获得电磁信号,三维重建人体图像,清晰显示受检部位。MRI分辨率高且具有较强的组织对比能力,有助于鉴别AD早期特征改变,是AD诊断的有效工具[5]。随着计算机及大数据的发展,将人工智能技术与MRI医学成像技术相结合进行辅助诊断成为当今研究的热点。

在老年痴呆计算机辅助诊断方面,传统方法主要基于机器学习,首先根据临床经验提取感兴趣区域,再手工设计特征,步骤繁琐且与模型相独立。深度学习能够自动学习特征表示,避免手工设计,特征提取和模型构建相统一,面向任务性更好,在医学图像分析中逐渐取代传统方法。

现有深度学习方法往往提取病理区域或图像块输入[6-8],这需要医学领域知识且每个受试者提取的部位是一样的,没有考虑患者间差异,同时切块或切片导致大脑信息不完整。采用全脑输入能保留尽可能多的信息,Fulton等[9]对全脑MRI进行二维切片训练残差网络ResNet-50来分类AD和正常对照。考虑到医学图像的空间信息,切片会损失一个维度的信息,Wang等[10]使用三维卷积适应医学图像的空间性,实现痴呆的分类。为应对从全脑图像中提取细微变化的困难,注意力机制被引入AD诊断模型。3D att-ResNet[11]引入残差连接加深模型层数,并采用注意力模块使网络聚焦重要区域,提供一种辅助诊断方法。混合网络HybNet[12]使用第一阶段网络最后一层卷积图指导第二阶段基于全图像的全局分支和基于图像块的局部分支,在AD分类上实现较好的结果,验证了临床上海马、杏仁核等病理位置。Etminani等[13]采用端到端的三维卷积处理PET图像,发现后扣带皮层的辨别性作用。由于大脑结构复杂而病变细微,端到端网络处理全脑MRI仍然具有挑战性,上述方法为提高性能大多依赖医学知识,或在单尺度上提特征,难以获得多尺度信息。

由于卷积核感受野的局部性,卷积神经网络(CNN)很难建模全局依赖关系。全局自注意架构Transformer[14]首次在自然语言处理(Natural Language Processing,NLP)领域中提出,被设计用来建模长程依赖关系,成为NLP领域的首选模型。受Transformer在自然语言处理上的成功激励,Vision Transformer(ViT)[15]提出对原始图片分块,展平成序列输入Transformer编码,接入一个全连接层执行图像分类任务,将Transformer引入图像处理领域。

考虑到Transformer具有比CNN强大的全局建模能力,但是局部特征提取能力不如CNN,对于医学图像来说,病灶的大小和位置具有随机性和不确定性,彼此间也可能具有联系,因此3D MRI局部和全局的空间信息都有重要意义。因此本文提出一种端到端的结合3D CNN和Transformer的网络用于AD诊断;另外,考虑到特征图的语义信息能够反映网络在分类时候的注意力,设计了一个多尺度注意力模块,融合个体特异性的多尺度特征作为疾病相关的先验知识指导Transformer全局建模,进一步提高分类性能,并可视化热力图。

1 本文模型及方法

1.1 整体结构

提出的方法整体结构如图1所示,由CNN局部特征提取、多尺度注意力模块和全局特征提取组成。卷积部分引入密集连接的思想将前一层卷积的输入连接到下一层卷积的输入,在不增加模型复杂度的情况下加强对特征的复用,使网络易于训练。得到的多尺度局部特征作为疾病相关的注意力图,与输入加权突出疾病相关位置,指导全局模块建模,既能识别特定于受试者的疾病相关辨别性位置,也能有效提取局部特征和全局特征。全局特征提取模块对全局信息整合、有效建模全局长程依赖。

图1 方法总体结构Fig.1 Overall structure of the method

首先,输入的MRI图像X∈RC×H×W×D三维空间大小为H×W×D,通道数为C。直接采用ViT形式将MRI切块并展平建立全局关系无法充分提取局部特征,因此首先利用堆叠卷积层和池化层的CNN提取MRI图像的局部和深层语义信息,得到一系列不同层次的特征图;直接下采样的原始输入数据经过基于多尺度特征加权融合思想的注意力图编码后,由一个3×3×3卷积层进行嵌入维度转换,并将特征图像素按空间展平后送入堆叠的Transformer编码器在全局范围内建立长程依赖关系,最终由全连接层和softmax对特征进行分类,实现AD的自动诊断。

1.2 CNN局部特征提取和多尺度注意力模块

CNN核心是卷积操作,CNN局部特征提取模块的作用是学习图像的局部特征。二维卷积广泛应用于自然图像中,对于医学图像特别是MRI影像具有三维的空间信息,切片与切片之间也有关于病灶的结构信息,使用2D卷积会丢失层与层间的信息。为了充分利用3D MRI的空间信息,采用三维卷积,通过堆积多个卷积层和池化层得到输入数据由高到低不同分辨率的特征图,自动提取到由浅层到深层的语义特征。局部特征提取模块的具体结构如图2所示,包括6个卷积核大小为3×3×3的卷积(Conv)层、3个步长为2的池化层,池化层将特征图的大小减半以提高卷积核的感受野。在每一个特征图维度大小上,将前一层卷积的输入连接到下一个卷积层的输入,加强模型对特征的复用,并采用1×1×1卷积降低特征通道数防止参数增加过多。所有Conv层后面跟批量归一化(BN)层和校正线性单元(ReLU)激活。

图2 3D CNN特征提取模块Fig.2 3D CNN feature extraction module

图3 多尺度注意力模块Fig.3 Multi-scale attention module

1.3 Transformer全局模块

(1)

图4 3D Transformer模块Fig.4 3D Transformer module

(2)

2 实验与结果分析

2.1 实验数据集

本文所使用的AD分类数据均来源于开源的阿尔兹海默神经影像计划(The Alzheimer’s Disease Neuroimaging Initiative,ADNI)数据集(http:∥adni.loni.usc.edu/),其宗旨是提供开源的临床统计、影像、生物标志物和基因等信息来帮助研究AD。从中随机选择了160名受试者。所有受试者均接受了基本临床评估,包括人口统计信息、体格检查、神经心理评估和影像检查,再由专业的医生进行诊断,具体划分为80名AD患者、80名正常对照(NC),如表1所示。影像数据由160个受试者的3T类型T1加权的MRI扫描组成。

表1 160名受试者的信息

此外,为了验证本文方法在医学图像分类任务上的适用性和鲁棒性,在MedMNIST[16-17]提供的3个公开的3D医学图像分类数据集AdrenalMNIST3D、No-duleMNIST3D和VesselMNIST3D上进行了实验。数据集具有相同的大小 28 pixel×28 pixel×28 pixel,分辨率均为1 mm×1 mm×1 mm。AdrenalMNIST3D数据集是3D 肾上腺掩码图像集,由内分泌专家标记为正常或肾上腺肿块2类,包含519例训练、141例测试图像。NoduleMNIST3D数据集是3D CT扫描的肺结节数据,根据恶化程度分为阴性和阳性2类,包含591例训练、128例测试图像。VesselMNIST3D数据集是3D全脑血管图像,分为正常血管段和动脉瘤段2类,包含302例训练、91例测试图像。

2.2 数据预处理

使用的MRI预处理流程包括格式转换、前连合(AC)-后连合(PC)校正、配准到Colin27模板、去颅骨和裁剪。具体来说,首先使用MRIcron软件将磁共振设备获取的DICOM格式的原始MRI数据转换为NIFIT格式,并丢弃了颈脖部位的冗余数据,只保留了有意义的头部数据。转换后的数据具有160 pixel×256 pixel×256 pixel、170 pixel×256 pixel×256 pixel和160 pixel×192 pixel×192 pixel等不同尺寸,空间分辨率也不完全相同,图5(a)展示了一个NIFIT格式的原始数据样本。立体定向脑图谱是以前连合与后连合的连线(AC-PC)为基准的,不同人脑的AC-PC没有显著性差异,故接着使用基于Matlab的SPM12工具包进行AC-PC校正,将原点校正在AC点附近使不同人脑图像具有可比性。然后使用SPM12将所有MRI配准到Colin27模板上,去除全局线性误差,并将三维尺寸统一到181 pixel×217 pixel×181 pixel,空间分辨率统一到1 mm×1 mm×1 mm,配准后的图像如图5(b)所示。颅骨对卷积网络提取AD相关的特征没有意义且会带来大量的计算量,因此使用基于SPM12的CAT12包剥离颅骨,去颅骨的数据如图5(c)所示。最后,为了尽可能减少不具有信息的背景,求出每个受试者MRI图像的非零大脑区域的外接立方体长宽高并将最大值作为最终保留的大小,最终所有MRI被裁剪到160 pixel×192 pixel×148 pixel。

图5 一例原始脑MRI及相应关键预处理后图像示例Fig.5 An example of raw brain MRI and thecorresponding preprocessed images

2.3 实验设置

在AD诊断任务上,本研究对160名受试者的160张T1加权的3D MRI进行了分类(即AD和NC)。数据预处理过程见2.2节,处理后的数据为160 pixel×192 pixel×148 pixel,空间分辨率为1 mm×1 mm×1 mm。所有实验都通过Python编写的基于TensorFlow的Keras环境实现,模型的训练和测试均在一块NVIDIA GeForce GTX TITAN X 12 GB的GPU上完成。160例实验数据包含80例NC和80例AD,比例接近1∶1,防止正负样本数量差距大导致模型拟合结果偏向数据量多的类别。

采用Adam[18]优化算法最小化交叉熵损失函数对模型进行训练,交叉熵损失函数定义如下:

(3)

由于训练样本少、3D的数据维度高,对多个卷积层应用0.3的丢弃概率随机忽略部分神经元减少神经网络过拟合。批量大小设置为4,迭代次数epoch设置为200。

2.4 评估指标

为了评估模型对AD分类的性能,引入了二分类混淆矩阵,如表2所示。真阳性(True Positive,TP)表示实际为AD判断为AD的受试者个数;假阴性(False Negative,FN)表示实际为AD判断为NC的受试者个数;假阳性(False Positive,FP)表示实际为NC判断为AD的受试者个数;真阴性(True Ne-gative,TN)表示实际为NC判断为NC的受试者个数。

表2 二分类混淆矩阵

混淆矩阵能够较为全面地反映模型的性能,为了对模型的分类性能进行定量比较,实验采用了3种混淆矩阵衍生的分类性能指标,分别是准确性(Accuracy,ACC)、灵敏性(Sensitivity,SEN)和特异性(Specificity,SPE)。ACC是最常见、最直观的评价指标,由式(4)计算,表示模型正确分类的AD和NC的数量占样本总数量的比例,ACC越大模型正确分类的样本数量越多。SEN由式(5)计算,反映了模型对AD病人的识别能力,SEN越高漏诊的AD越少。SPE表达了模型对正常人的识别能力,SPE越高NC被误诊为AD的比例越低,由式(6)计算。

(4)

(5)

(6)

2.5 分类效果对比实验

首先在AD数据集上进行实验并与几种深度学习AD诊断方法3D-CNN[13]、3D att-ResNet[11]和HybNet[12]进行了对比。根据文献中的网络结构在Keras环境中复现了3D-CNN和3D att-ResNet的代码,批量大小设置为4,训练100个轮次,学习率0.000 1,使用Adam优化器对交叉熵损失训练。对于HybNet方法,采用公开的源码,第一阶段全卷积网络批量大小设置为2,训练150个轮次;第二阶段的全局分支和局部分支的批量大小设置为2,训练100个轮次;最终的混合网络批量大小设置为5,训练100个轮次,均应用0.3的丢弃概率,初始学习率0.000 1,采用Adam优化器。

AD和NC两类数据类别比接近1∶1,训练包含120张MRI,测试集包括40张,实验结果如表3所示。本文所提方法准确率达到了95.00%,敏感性95.00%,特异性95.00%,这些指标与现有的深度学习痴呆症诊断方法具有可比性或者更好。实验结果证明了结合CNN的局部细节特征和Transformer长程依赖关系的有效性,同时充分利用多尺度信息作为疾病的注意力引导。由结果可以看出所提模型经过训练后在AD诊断分类任务上能获得较好的性能,同时模型训练参数量较少,具有一定的优势,是一种轻量级的端到端方法,快捷高效,具有临床应用价值。

表3 AD vs NC的对比实验结果

图6展示了本文多尺度注意力模块产生的3例AD受试者和3例NC受试者的疾病注意力图,其中注意力图从3个不同的切面以二维形式展示,突出显示的区域表明在网络进行AD分类任务时受到更多的关注。从图中可以看出,AD受试者的注意力图的权重在海马、后扣带回皮质等区域累积,具有明显差异,与AD患者相比,NC受试者的注意力图在大脑空间上没有显著区别。本文的注意力模块在识别痴呆患者时突出强调了额叶、扣带回皮层、顶叶皮层和海马等区域,而这些区域在临床上是痴呆常见的辨识性区域。此外,可以观察到本文的注意力图在不同受试者间强调的位置在病理区域上具有相似性,但个体间不完全相同,能反映不同人脑结构、发病区域和病情程度带来的个体差异。结果表明,面向任务进行训练的端到端网络中注意力模块产生的疾病相关注意力图能够作为网络分类AD的指导。

图6 注意力图示例Fig.6 Examples of attention map

为了验证所提方法在3D医学图像分类上的适用性和鲁棒性,分别在MedMNIST提供的3个公开的3D医学图像分类数据集AdrenalMNIST3D、No-duleMNIST3D和VesselMNIST3D上进行了实验,并采用相同的分类性能指标ACC、SEN和SPE。在肾上腺肿块分类数据集AdrenalMNIST3D上的实验结果显示如表4所示。可以看出,ACC相比于其他3种方法分别提升了16.29%、9.2%、12.75%。肺结节分类数据集NoduleMNIST3D上的结果如表5所示。可以看出,ACC相比于其他3种方法分别提升了33.37%、16.07%和9.94%。脑动脉瘤分类数据集VesselMNIST3D上的结果如表6所示。可以看出,ACC相比于其他3种方法分别提升了15.38%、11.32%和7.7%。所提方法在3个公开3D医学分类数据集上相比与其他3种对比方法,ACC、SEN和SPE指标均有所提升,说明结合CNN的局部性和Transformer的长程依赖建模特性的方法不仅在AD分类上能够达到较好效果,同时在其他的三维医学图像分类任务上具有有效性。

表4 AdrenalMNIST3D的对比实验结果

表5 NoduleMNIST3D的对比实验结果

表6 VesselMNIST3D的对比实验结果

2.6 消融实验

为了验证所提方法设计的有效性并得到最优的结构,在AD分类数据集上进行了消融实验:① 探索了Transformer全局建模模块和多尺度特征图注意力编码模块的有效性;② 探讨了Transformer的输入词向量嵌入维度大小和Transformer层数对模型的影响。

所提方法由卷积特征提取、多尺度特征加权注意力编码和Transformer全局建模3个部分组成。表7展示了组成模块消融实验结果,通过分别添加Transformer全局部分和多尺度注意力模块验证有效性。由结果可以看到,仅使用不断整合卷积窗口内局部特征再由全连接层分类的纯卷积网络,ACC、SEN和SPE分别为77.16%、70.96%和83.07%;添加Transformer模块对卷积得到的细节特征图逐像素间建立全局关系,相比于纯CNN,结合局部和全局特征使,3个分类指标分别提高13.39%、20.97%和6.16%,验证了添加Transformer建立长程依赖关系在医学图像分类中的有效性。在此基础上加入多层特征加权融合的注意力编码模块,分类性能进一步提升,验证了所设计的深层和浅层特征融合作为AD相关的位置先验信息与输入加权作为Transformer输入的有效性,促进上下文信息提取和表达。

表7 组成模块的消融实验

表8、表9展示了Transformer嵌入维度和层数2个模型参数对分类性能影响的消融实验结果。从结果中可以观察到,表8固定Transformer层数为2,随着嵌入维度的增加,特征表达能力得到提升,分类性能先有所提高,但维度过高的时候分类性能反而下降,在嵌入维度为128时达到最高的准确率。导致准确率发生这种变化的可能原因是,对于使用的数据样本数量不是很多但是每个样本的数据量很大的情况下,过高的模型特征表示度会导致模型复杂度增高从而不利于训练,容易出现过拟合,采用128的嵌入维度达到最好的效果。表9在固定嵌入维度为128时增加Transformer的层数,模型的分类性能先提高后下降,在层数为2时达到最好的效果,说明过大或过小的模型大小都不能达到最好的效果,模型过小没有充分提取输入的特征容易出现欠拟合,过大则会导致模型参数量过多,医学图像分类经常面临样本数量小但数据复杂度高的情况,需要根据实际情况合理选择模型的参数。

表8 嵌入维度的消融实验

表9 Transformer层数的消融实验

3 结束语

处于不同阶段的AD患者大脑呈现出不同程度的萎缩,但具体萎缩区域和程度不明确,不同的病人也可能不固定位置出现脑出血、脑梗死和白质弥漫等,导致人工诊断方法准确率偏低耗费精力,AD的计算机辅助诊断非常有意义。本文提出一种基于多尺度特征融合的结合3D CNN和Transformer的深度学习方法用于准确而快速利用3D MRI对AD进行自动分类,用ACC、SEN和SPE指标对模型分类性能进行评价,并在ADNI分类数据集和公开的3D医学分类数据集上验证了所提方法的效果,在AD诊断上达到95%的准确率,高于几种常规深度学习方法,在其他3个医学数据集上也取得较好的结果。本文既利用了3D CNN的局部性也利用了Transformer的全局性,能够有效提取MRI的特征信息。针对AD的发病机制尚不明确的问题,使用了多尺度特征加权融合的注意力模块,既利用了多层语义信息又结合了精确的位置信息,进一步提升了分类性能,突出强调个体特异性的疾病相关位置。