MIL是弱監督分類問題的有力工具。然而,目前的MIL方法通常基于iid假設,忽略了不同實例之間的相關性。為了解決這個問題,作者提出了一個新的框架,稱為相關性MIL,并提供了收斂性的證明。基于此框架,還設計了一個基于Transformer的MIL (TransMIL)。TransMIL可以有效地處理不平衡/平衡和二元/多分類,具有良好的可視化和可解釋性。在CAMELYON16數據集上,二元腫瘤分類的測試AUC高達93.09%。在TCGANSCLC和TCGA-RCC數據集上,癌癥亞型分類的AUC分別高達96.03%和98.82%。
來自:TransMIL: Transformer based Correlated Multiple Instance Learning for Whole Slide Image Classification
目錄
- 背景概述
- 方法
- 如何應用Transformer到相關性MIL
背景概述
WSI將活檢切片上的組織轉換成完全保留原始組織結構的十億像素圖像。然而,WSI中基于深度學習的活檢診斷由于像素空間龐大導致缺乏像素級標注。為了解決這個問題,通常采用MIL將診斷分析作為一個弱監督學習問題。
在基于深度學習的MIL中,一個簡單的想法是對CNN提取的instance特征嵌入進行pooling操作。Ilse等人提出了一種基于注意力的聚合算子,通過可訓練的注意力為每個實例提供額外的貢獻信息。此外,Li將非局部注意力引入了MIL問題。通過計算得分最高的實例與其他實例之間的相似度,賦予每個實例不同的注意力權重,從而得到可解釋的注意力圖。
然而,所有這些方法都基于這樣的假設:每個bag中的所有實例都是獨立且同分布的。雖然在許多任務中取得了一些改進,但在許多情況下,這種假設并不完全有效。實際上,在做出診斷決定時,病理學家通常會同時考慮單個區域周圍的環境信息和不同區域之間的相關信息。因此,在MIL診斷中考慮不同instance之間的相關性是可取的。
目前,Transformer由于具有較強的描述序列中不同token之間的相關性以及對遠距離信息建模的能力,被廣泛應用于視覺任務中。如圖1所示,Transformer采用自注意力機制,可以關注序列內每個token之間的兩兩相關性。然而,傳統的Transformer受到其計算復雜性的限制,只能處理較短的序列(例如,小于1000或512)。因此,它不適合WSI等大尺寸圖像。
- 圖1:決策過程圖示。MIL注意力機制:遵循iid假設。自注意機制:屬于相關性MIL。
方法
以二元MIL為例,我們想要預測target value Y i ∈ { 0 , 1 } Y_{i}\in\left\{0,1\right\} Yi?∈{0,1},給定一個bag X i X_{i} Xi?(instance為 { x i , 1 , x i , 2 , . . . , x i , n } \left\{x_{i,1},x_{i,2},...,x_{i,n}\right\} {xi,1?,xi,2?,...,xi,n?}),其中 i = 1 , . . , b i=1,..,b i=1,..,b,實例級標簽是未知的: { y i , 1 , y i , 2 , . . . , y i , n } \left\{y_{i,1},y_{i,2},...,y_{i,n}\right\} {yi,1?,yi,2?,...,yi,n?},bag標簽是已知的,并且與實例標簽有聯系:
b b b是袋的總數, n n n是第 i i i個袋里的實例數, n n n的個數可以根據不同的袋而變化。
關于相關性MIL的優勢,文中給出了證明,但是在此略過。主要意思是:
- 考慮實例相關性可以具有更小的信息熵,從而減少不確定性,為MIL帶來更多有效信息。TransMIL與過去方法的主要區別如圖2
- 圖2:不同pooling矩陣 P P P的差異。假設從a中的WSI采樣5個instance, P ∈ R 5 × 5 P\in\R^{5\times 5} P∈R5×5是對應的pooling矩陣,其中對角線內的值表示instance自身的注意力權重,其余值表示不同instance之間的相關性。b,c,d 都忽略了相關信息,因此 P P P是對角矩陣。在b中,第一個實例是由Max-pooling算子選擇的,所以在對角線位置只有一個非零值。在c中,由于Mean-pooling運算符,對角線內的所有值都是相同的。在d中,由于引入的是bypass注意力,對角線內的值可能會發生變化。但其余位置為0(獨立同分布假設)。e服從相關性假設,因此在非對角線位置存在非零值,表示不同實例之間存在相關性。
對于MIL的pooling,這里有一個通用的三步法:
- 形態信息:morphological,空間信息:spatial
如何應用Transformer到相關性MIL
Transformer使用自注意力機制對序列中所有令牌之間的交互進行建模,位置信息的添加進一步增加了順序信息。因此,將Transformer引入相關性MIL問題是合理的,其中函數 h h h對實例之間的空間信息進行編碼,pooling矩陣 P P P使用自注意力進行信息聚合。為了說明這一點,進一步給出一個正式的定義。
給定一個bag集合 { X 1 , . . . , X b } \left\{X_{1},...,X_{b}\right\} {X1?,...,Xb?},每個bag對應一個標簽 Y i Y_{i} Yi?。目標是學習映射: X → T → Y \mathbb{X} \rightarrow \mathbb{T}\rightarrow \mathbb{Y} X→T→Y,即從bag空間,到Transformer空間,再到標簽空間。
為了更好地描述 X → T \mathbb{X} \rightarrow \mathbb{T} X→T的映射,作者設計了一個包含兩個Transformer層和一個位置編碼層的TPT模塊,其中Transformer層用于聚合形態信息,PPEG (Pyramid position encoding Generator)用于編碼空間信息。所提出的基于MIL (TransMIL)的Transformer的概述如圖3所示。
- 圖3:每個WSI被裁剪成patch(背景被丟棄),并被ResNet50嵌入到特征向量中。然后用TPT模塊對序列進行處理:1)序列的平方;2)序列相關性建模;3)條件位置編碼與局部信息融合;4)深度特征聚合;5) T → Y \mathbb{T}\rightarrow \mathbb{Y} T→Y的映射。
序列來自每個WSI中的特征嵌入。TPT模塊的處理步驟如算法2所示,其中MSA表示多頭自注意力,MLP表示多層感知機,LN表示 Layer Norm。