基于注意力和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的發病機制尚不明確的問題,使用了多尺度特征加權融合的注意力模塊,既利用了多層語義信息又結合了精確的位置信息,進一步提升了分類性能,突出強調個體特異性的疾病相關位置。