基于注意力機制多任務的肺結節癌變風險判斷①

王廣涵,程遠志,史 操,許燦輝

(青島科技大學 信息科學技術學院,青島 266061)

肺癌是全球范圍內發病率和死亡率增長最快的惡性腫瘤之一,病死率在惡性腫瘤中居首位[1].高病死率的原因之一就是肺癌患者早期癥狀不明顯,然而等到發病明顯時,再檢查就已經是中晚期,難以有效的治療.Orlacchio 等人[2]指出,如果可以早期診斷,進行肺癌切除手術后,患者5年生存率可達到40%～80%.因此若能在肺癌早期階段進行正確診斷和治療,將顯著改善患者預后,降低肺癌病死率,早期肺癌沒有明確轉移,可以得到幾乎根治的效果.

肺癌的確診需要進行穿刺或者活檢,但是這兩項檢查需要進行手術操作,對醫生要求比較高,且耗時耗費,一般只有通過影像檢查確定癌癥風險和惡性結節位置之后才會視情況進行病理診斷確認.計算機斷層掃描技術(computed tomography,CT)是目前臨床上肺癌診斷最常用的影像檢查手段.

想要盡早的發現癌癥風險,則需要經常性的進行檢查,低劑量螺旋CT 掃描(LDCT)在顯著減低輻射劑量的同時可保證較好的圖像質量,使得肺癌篩查也可以成為日常體檢的項目之一,也就能及時進行早期肺癌診斷[3].

其實際應用價值,使得國內外的9 部肺癌篩查指南/共識都推薦高危人群定時采用LDCT 進行篩查[4],但是目前LDCT 肺癌篩查還未大規模普及應用是因為尚存在兩個問題:放射診斷醫師的培養數量遠跟不上需要分析的CT 影像的增長量,且人工分析的工作量也十分巨大,時間和經濟上的成本居高不下;放射診斷醫師人工觀察CT 圖像給出的結果也還存在假陰性和假陽性而導致的漏診、延誤.因此一種能夠根據CT 圖像自動得出肺結節癌變良惡性風險的計算機輔助診斷系統(CAD)就顯得十分必要.

近年來的研究中,根據算法可以分為兩類:基于影像語義特征的方法和基于深度學習的方法.

前一種方法是通過機器學習算法來提取肺結節明確的影像組學特征,如紋理特征、灰度特征、形態學特征等,然后根據這些語義特征來得出肺結節的良惡性的分類結果.比較有代表性的有:Suzuki 等人[5]提出的訓練多個神經網絡,采用3 個灰度特征、2 個邊緣特征、1 個形狀特征和臨床信息用于良惡性肺結節的分類方法;張佳嘉等人[6]也是用同一數據集訓練出多個分支網絡,先提取出醫生診斷時所用的肺結節語義特征信息,然后用于完成肺結節良惡性診斷任務.這種方法的優點在于能夠給出病理診斷的依據,可解釋性上更強一些.

后一種方法是直接應用深度學習等方法,不去關心放射科醫師的診斷依據,充分利用CT 圖像的全部特征信息,直接得出良惡性判斷的結果.比較有代表性的有:Ardila 等人[7]提出了一種端到端的肺結節檢測和肺癌風險預測模型,利用3D CNN 網絡,可以對同一病患的多張不同時期的CT 影像進行對比處理,檢測出肺結節區域后結合CT 全局信息給出一個惡性腫瘤得分,取得了很好的效果;Hua 等人[8]提出了兩種深度學習架構,來對肺結節的良惡性直接進行分類,規避了對語義特征進行細化的需要,也取得了不錯的分類性能.

在對比過兩種方法的優缺點之后,本文提出了一種基于多頭注意力機制的多任務學習方法,既采用深度學習網絡來充分提取CT 圖像特征,又利用硬參數共享的多任務模式來得出結節良惡性以及相關語義特征,避免了兩階段判斷網絡受到第一步分析識別各語義特征任務性能影響的情況,同時還能利用關聯語義特征提升對于良惡性判斷的性能.

方法流程如圖1所示,從原CT 圖像中根據檢測出的結節坐標,將需要判斷的結節截取出來,然后進行圖像處理輸入到基于多頭注意力機制的多任務學習網絡中,進而得出良惡性程度等語義屬性的判斷結果.

圖1 整體流程圖

1 多任務框架

多任務學習(multitask learning,MTL)的概念最早由Caruana[9]在 1997年提出,他構造了一個前饋神經網絡,提出了一種通過訓練單一的多層的感知器來執行多個任務的機制,在輔助醫療診斷和自動駕駛等領域得到了成功應用.

選擇多任務學習是因為單任務只關注一個目標而可能忽略了更多有用的信息.具體來說,這些信息來自于相關任務的訓練信號.通過在相關任務之間共享表示,我們可以使我們的模型能夠更好地對原始任務進行泛化[10].多任務學習的目的是在一并解決多個相關任務的同時實現內在特征的共享.已經證明,這種共享可以提高部分或所有任務的性能[11].

目前基于深度神經網絡的多任務學習中常用的模式有兩種:硬參數共享與軟參數共享,相對軟參數共享模式來講,硬參數共享模式不需要對每個任務都構建和訓練單獨的模型,且可以降低過擬合的風險[12],因此本文選擇硬參數共享的多任務模式.

1.1 肺結節各屬性特征

LIDC-IDRI[13]是由美國國家癌癥研究所發起收集的肺結節公開數據集,用于高危人群早期肺癌的檢測診斷.除了影像數據還包含了4 位胸部放射科醫師的診斷結果XML 文件.醫師針對每個結節給出9 個醫學語義特征(對比度、分葉征、毛刺征、球形度、邊緣、紋理、內部結構、鈣化程度和惡性程度)的具體分級(1-5 級,鈣化程度為1-6 級),等級越高,語義特征越明顯.

LUNA16[14]是該數據集的子集,刪除了LIDC-IDRI中切片厚度大于2.5 mm 和肺結節小于3 mM的CT 影像.相對于原數據集,LUNA16 數據集能夠更好的對計算機輔助診斷系統的性能進行評估.其中包含888 張肺部LDCT 掃描數據,以及1 186 個肺結節的坐標位置和直徑大小.

結合LIDC-IDRI 數據集中4 位胸部放射科醫師的診斷結果,就可以得到LUNA16 數據集中全部肺結節的語義特征,我們取4 位醫師的診斷結果的算術平均作為最終各語義特征的分級數值.

1.2 各屬性特征與癌變風險相關關系

在構造多任務模型之前,我們要先探討其他語義特征與結節良惡性之間的相關關系,以此來估計添加到多任務學習中的特征是否會對良惡性判斷產生正向影響,對此我們采用余弦相似度來衡量每一個語義特征Y與結節良惡性等級X的關系,計算公式如下:

計算結果如圖2所示,通過其他8 個語義特征和結節大小與良惡性的相似系數,可以發現:除內部結構這一特征外,其他語義特征與結節良惡性之間都存在較強的相關性,因此我們選擇這些語義特征作為多任務學習中的輔助任務.

圖2 與良惡性相關關系

1.3 多任務損失函數權重優化方法--動態加權平均

對于多任務的loss,最簡單的方式是直接將這兩個任務的loss 直接相加,得到整體的loss.這種loss 計算方式的不合理之處是顯而易見的,不同任務loss 的量級很有可能不一樣,loss 直接相加的方式有可能會導致多任務的學習被某個任務所主導或學偏.當模型傾向于去擬合某個任務時,其他任務的效果往往可能受到負面影響,效果會相對變差.因此需要對每個任務的loss 進行加權[15].

通過上一步得出的8 個語義特征以及良惡性,每個結節有9 個標簽來對應9 個特征的平均等級評分,每個子任務都有其損失函數,如式(2)所示,將這些子任務的損失函數加權求和便是整個網絡的損失函數:

其中,Li(X,Yi)是以X為輸入,Yi為輸出的第i個子任務的損失函數,具體到語義特征屬性的等級預測任務,可以表示為式(3):

其中,n為結節數量,是第i個結節的真實值,是網絡輸出的預測值.

多任務學習中,不同子任務的收斂速度,訓練難度都是不同的,不能讓簡單任務主導整個訓練,導致各個子任務的表現差距過大,對于大多數多任務網絡而言,訓練過程中的最大難題是為每一個子任務找到合適的權重,讓每個子任務的重要性得到平衡,解決這些問題的辦法有幾類,有代表性的有:梯度歸一化(GradNorm)[16],使不同的任務loss 量級接近;動態任務優先級(dynamic task prioritization)[17],利用不確定性賦權值,根據任務難易程度進行賦權值等.

本文采用的權重設計機制為動態加權平均(dynamic weight average)[18],通過考慮每個任務損失的變化率來調整任務權重.受梯度歸一化方法的啟發,動態加權平均方法對每個子任務首先計算前一個epoch 對應損失的比值,然后除以一個固定的值T進行exp 映射后,計算各個損失所占比.如式(4)所示,首先計算一個epoch后的損失變化:

然后,將ωi(t-1)帶入式(5),得到對應子任務i的權重:

其中,T的大小代表了任務間的松散程度,如果該值足夠大,那么 λ便會趨向于1,代表各任務的權重相同.得到各任務權重之后,就按照整體損失函數對共享網絡結構進行優化調整.

2 基于多頭注意力機制的多任務學習框架

Transformer 是Vaswani 等人[19]于2017年提出的一種模型架構.它開創性的思想,顛覆了以往序列建模和RNN 劃等號的思路,之后被廣泛應用于自然語言處理的各個領域,BERT、GPT 等模型也都是基于Transformer 的模型,在自然語言處理各個任務中都取得了突破性的成果.

因其在自然語言處理領域的巨大成功,開始有研究者嘗試將Transformer 引入到計算機視覺領域,其中由Dosovitskiy 等人[20]提出的Vision Transformer(ViT)模型很好的保留了Transformer 的原始框架,在圖像分類任務上可以獲得與當前最優卷積網絡相媲美的結果.

并且因為Transformer 中所應用的多頭注意力機制,將模型分為多個頭,形成多個子空間,可以讓模型去關注不同方面的特征/信息,正適用于多任務學習中對不同的子任務關注不同的特性,因此本文采用ViT 模型作為多任務學習框架的共享特征提取層.

2.1 多任務學習分類框架

本文的多任務學習方法框架,如圖3所示.

圖3 模型框架圖

(1)根據檢測出的結節坐標,從CT 原圖上取出小塊,然后采用了一種保留3 維信息的CT 圖像轉換ViT 切片數據的方法,將3 維數據通過切片拼接轉為符合特征提取網絡輸入所需的二維圖片形式,來完成第一部分的圖像處理任務.

(2)將訓練集數據輸入基于ViT 模型的多頭注意力機制特征提取網絡中,提取出的特征向量作為后續不同的分類任務網絡的輸入.

ViT 模型為了將一個標準Transformer 直接應用到圖像上,盡可能少的修改,選擇將圖像分割成小塊,并將這些塊轉化為線性嵌入序列,然后在圖像patch 的嵌入中加入位置嵌入,通過不同的策略在全局范圍內保留空間/位置信息.之后因為是分類任務,只需要輸入Transformer 的Encoder 層就可以了,而無需像自然語言處理的翻譯任務中還需要加Decoder 層.

本文中,通過第一步圖像處理的方法,將3 維CT圖像的每一個截面切片作為一個patch,進行線性序列化后,嵌入位置序列信息,作為Transformer 的輸入,更能夠保留和獲取空間/位置信息.

(3)采用多個分類任務共享特征提取網絡的硬參數共享模式,將ViT 模型提取出的特征向量分別輸入不同的MLP 網絡,對結節的多個屬性進行良惡性的判斷.使用動態加權平均(DWA)的損失函數權重計算優化方法,保證多個分類任務的loss 同步降低,反向調整特征提取網絡以獲取跟任務更為相關的特征,動態整體優化各分類網絡的準確性.模型在共享層就要學到一個通用的嵌入式表達使得每個任務都表現較好,從而降低過擬合的風險,以達到通過多任務學習,提升該模型的泛化效果.

2.2 保留3 維信息的CT 圖像轉換ViT 切片數據的方法

本文選用基于ViT 模型的多頭注意力機制特征提取網絡,因為Transformer 模型的特性,加載在更大樣本上預訓練好的模型參數,然后再根據具體的任務遷移到目標數據集進行調整,可以獲得更好的結果,并且所需的計算資源大大減少.然而ViT 模型現有的預訓練參數都是用二維圖像數據訓練得出的,因此想要應用于3 維的CT 數據上,需要對CT 數據進行變形調整.

受到ViT 對二維圖像進行切片取patch 方法的啟發,本文采用的方法是,對3 維的CT 數據進行截面切片,形成的切片大小可以對應ViT 模型的patch 大小,然后將一整套切片圖片拼接成一副二維圖像,就可以應用現有的ViT 預訓練模型進行遷移學習.這樣即最大程度上保留了CT 數據的3 維信息,還可以保證ViT模型進行裁剪圖像時,取到的每一片patch 都是原CT圖像的一個完整的截面圖像.

圖4 為拼接后的二維圖像以及經過ViT 的多頭注意力機制處理后的注意力熱力圖,可以看到注意力集中于結節區域.

圖4 注意力熱力圖

3 實驗分析

本文對LUNA16 數據集上進行五折交叉驗證實驗,結節各特征屬性值按照第1 節所述的,取4 個醫師給定值的算數平均,用兩種方式分類:(1)評級小于3 的歸為負樣本,評級大于3 的歸為負樣本,評級等于3(4 名醫師全部認為該結節評級不確定)的歸為不確定樣本,進行三分類;(2)評級小于2.75 的歸為負樣本,評級大于3.25 的歸為正樣本,排除掉最不確定的情況進行二分類.

分別以上述兩種分類方式進行單任務和多任務的實驗,三分類的結節良惡性判斷準確率遠小于二分類的準確率,此結果表明人類醫師無法準確判斷識別的結節,使用本文提出的判斷語義特征的多任務學習框架同樣無法識別判斷.

下面僅就二分類的實驗結果進行具體說明.

3.1 單任務網絡實驗

首先用最基礎的MVC 模型,通過第1 節選定的9 個語義特征的值來判斷結節良惡性類別,準確率達到了93%,符合我們對于語義特征與良惡性具有相關關系的分析.

我們使用ViT 模型分別對每一個語義特征單獨訓練,部分結果如表1所示.

表1 單任務網絡分類結果

可以發現單任務對各語義特征判斷的效果并不是很好,如果用各個單任務網絡得出的屬性值再來對良惡性進行分類,則相當于對原輸入數據增加20%-30%的噪聲干擾,最終的準確率甚至會低于單任務直接對良惡性判斷的結果.

3.2 多任務網絡實驗

如表2所示,在使用多任務網絡對全部語義特征一起訓練時,因其內在特征的共享,使得各任務或多或少的有了性能上的提升.

表2 多任務網絡分類結果

雖然對于結節良惡性分類的準確率最高只有87.3%,未達到SOTA,但是本文提出的框架除良惡性判斷之外,還可以一步同時輸出結節的9 個語義特征屬性值,提供了良惡性判斷的依據與標準,更具有實際應用價值.

4 結論與展望

在本文中,我們提出了一種基于注意力機制的多任務學習框架,該方法可以同時對CT 影像中肺結節的良惡性等9 類語義特征進行判斷.我們在LUNA16 數據集上的1 186 個CT 結節進行了試驗評估,實驗結果表明該方法對肺結節癌變風險的似然性預測是有效的.雖然準確率沒有達到目前3D 預測模型的精度,但是本文的方法同時給出了其他語義特征的預測值,能夠對良惡性的判斷,下一步需要進行的研究工作是:(1)交叉驗證每一項語義特征對良惡性結果的影響;(2)進一步提升包括良惡性在內的各特征屬性的預測準確率.