TransMIL: Transformer based Correlated Multiple Instance Learning for Whole Slide Image Classification

1 Introduction . The advent of whole slide image ( WSI ) scanners , which convert the tissue on the biopsy slide into a gigapixel image fully preserving the original tissue structure [ 1 ] , provides a good opportunity for the application of deep learning in the field of digital pathology [ 2 , 3 , 4 ] . However , the deep learning based biopsy diagnosis in WSI has to face a great challenges due to the huge size and the lack of pixel-level annotations [ 5 ] . To address this problem , multiple instance learning ( MIL ) is usually adopted to take diagnosis analysis as a weakly supervised learning problem . In deep learning based MIL , one straightforward idea is to perform pooling operation [ 6 , 7 ] on instance feature embeddings extracted by CNN . Ilse et al . [ 8 ] proposed an attention based aggregation operator , giving each instance additional contribution information through trainable attention weights . In addition , Li et al . [ 9 ] introduced non-local attention into the MIL problem . By calculating the similarity between the highest-score instance and the others , each instance is given different ∗Contributed equally : { shaozc0412 , h2495067728 , sky374263410 } @ gmail.com . †Corresponding author : ybzhang08 @ hit.edu.cn . 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . attention weight and the interpretable attention map can be obtained accordingly . There were also other pioneering works [ 10 , 11 , 12 , 13 , 14 ] in weakly supervised WSI diagnosis . However , all these methods are based on the assumption that all the instances in each bag are independent and identically distributed ( i.i.d. ) . While achieving some improvements in many tasks , this i.i.d . assumption was not entirely valid [ 15 ] in many cases . Actually , pathologists often consider both the contextual information around a single area and the correlation information between different areas when making a diagnostic decision . Therefore , it would be much desirable to consider the correlation between different instances in MIL diagnosis . At present , Transformer is widely used in many vision tasks [ 16 , 17 , 18 , 19 , 20 , 21 ] due to the strong ability of describing correlation between different segments in a sequence ( tokens ) as well as modelling long distance information . As shown in Figure 1 , different from bypass attention network in the existing MIL , the Transformer adopts self-attention mechanism , which can pay attention to the pairwise correlation between each token within a sequence . However , traditional Transformer sequences are limited by their computational complexity and can only tackle shorter sequences ( e.g. , less than 1000 ) [ 22 ] . Therefore , it is not suitable for large size images such as WSIs . To address these challenges mentioned above , we proposed a correlated MIL framework , including the convergence proof and a generic three-step algorithm . In addition , a Transformer based MIL ( TransMIL ) was devised to explore both morphological and spatial information between different instances . Great performance over various datasets demonstrate the validity of the proposed method . 2 Related Work . 2.1 Application of MIL in WSI classification . The application of MIL in WSIs can be divided into two categories . The first one is instance-level algorithms [ 7 , 23 , 24 , 25 , 26 ] , where a CNN is first trained by assigning each instance a pseudo-label based on the bag-level label , and then the top-k instances are selected for aggregation . However , this method requires a large number of WSIs , since only a small number of instances within each slide can actually participate in the training . The second category is embedding-level algorithms , where each patch in the entire slide is mapped to a fixed-length embedding , and then all feature embeddings are aggregated by an operator ( e.g. , max-pooling ) . To improve the performance , the MIL attention based method [ 8 , 10 , 11 , 12 , 13 ] assigns the contribution of each instance by introducing trainable parameters . In addition , the feature clustering methods [ 14 , 27 , 28 ] calculated the cluster centroids of all the feature embeddings and then the representative feature embeddings was employed to make the final prediction . Recently , non-local attention [ 9 ] was also adopted in MIL to pay more attention to the correlation between the highest-score instance and all the remaining instances . 2.2 Attention and Self-attention in Deep Learning . Attention was initially used to extract important information about sentences in machine translation [ 29 ] . Then the attention mechanism was gradually applied to computer vision tasks , including giving different weights to feature channels [ 30 ] or spatial distribution [ 31 ] , or giving different weights to time series in video analysis [ 32 ] . Recently , attention was also applied in MIL analysis [ 10 , 11 , 12 , 13 ] . However , all these methods did not consider the correlation between different instances . The most typical self-attention application was the Transformer based NLP framework proposed by Google [ 33 ] . Recently , Transformer was also applied in many computer vision tasks , including object detection [ 16 , 17 ] , segmentation [ 18 , 19 ] , image enhancement [ 20 , 21 ] and video processing [ 34 ] . In this paper , for the first time , we proposed a Transformer based WSI classification , where the correlations among different instances within the same bag are comprehensively considered . 3 Method . 3.1 Correlated Multiple Instance Learning . Problem formulation Take binary MIL classification as an example , we want to predict a target value Yi ∈ { 0 , 1 } , given a bag of instances { xi,1 , xi,2 , . . . , xi , n } with Xi , for i = 1 , . . . , b , that exhibit both dependency and ordering among each other . The instance-level labels { yi,1 , yi,2 , . . . , yi , n } are unknown , and the bag-level label is Yi , for i = 1 , . . . , b . A binary MIL classification can be defined as : Yi = { 0 , iff ∑ yi , j = 0 yi , j ∈ { 0 , 1 } , j = 1 . . . n 1 , otherwise ( 1 ) Ŷi = S ( Xi ) , ( 2 ) where S is a scoring function , Ŷi represents the prediction . b is the total number of bags , n is the number of instances in ith bag , and the number of n can vary for different bags . Compared to the MIL framework proposed by Ilse et al . [ 8 ] , we further introduce the correlation between different instances . Theorem 1 and Inference give an arbitrary approximation form of the scoring function S ( X ) , and Theorem 2 provides the advantage of correlated MIL . Theorem 1 . Suppose S : X → R is a continuous set function w.r.t Hausdorff distance dH ( · , · ) . ∀ε > 0 , for any invertible map P : X → Rn , ∃ function σ and g , such that for any X ∈ X : |S ( X ) − g ( P X∈X { σ ( x ) : x ∈ X } ) | < ε . ( 3 ) That is : a Hausdorff continuous function S ( X ) can be arbitrarily approximated by a function in the form g ( P X∈X { σ ( x ) : x ∈ X } ) . Proof . By the continuity of S , we take ∀ε > 0 , ∃ δε , so that |S ( X ) −S ( X′ ) | < ε for any X , X′ ∈ X , if dH ( X , X′ ) < δε . Define K = ⌈ 1δε ⌉ and define an auxiliary function : σ ( x ) = ⌊Kx⌋ K . Let X̃ = { σ ( x ) : x ∈ X } , then : |S ( X ) − S ( X̃ ) | < ε , ( 4 ) because dH ( X , X̃ ) < 1K ≤ δε . Let P : X → Rn be any invertible map , its inverse mapping is expressed as P−1 : Rn → X . Let g = S ( P−1 ) , then : S ( P−1 ( P X∈X ( { σ ( x ) : x ∈ X } ) ) ) = S ( P−1 ( P X̃∈X ( X̃ ) ) ) = S ( X̃ ) . ( 5 ) Because |S ( X ) − S ( X̃ ) | < ε and S ( X̃ ) = S ( P−1 ( P ( X̃ ) ) ) = g ( P ( X̃ ) ) , we have : |S ( X ) − g ( P X∈X { σ ( x ) : x ∈ X } ) | < ε . ( 6 ) This completes the proof . Inference Suppose S : X → R is a continuous set function w.r.t Hausdorff distance dH ( · , · ) . ∀ε > 0 , for any function f and any invertible map P : X → Rn , ∃ function h and g , such that for any X ∈ X : |S ( X ) − g ( P X∈X { f ( x ) + h ( x ) : x ∈ X } ) | < ε . ( 7 ) That is : a Hausdorff continuous function S ( X ) can be arbitrarily approximated by a function in the form g ( P X∈X { f ( x ) + h ( x ) : x ∈ X } ) . Proof . The proof is in the Appendix A. Theorem 2 . The Instances in the bag are represented by random variables Θ1 , Θ2 , . . . , Θn , the information entropy of the bag under the correlation assumption can be expressed as H ( Θ1 , Θ2 , . . . , Θn ) , and the information entropy of the bag under the i.i.d . ( independent and identical distribution ) assumption can be expressed as ∑n t=1 H ( Θt ) , then we have : H ( Θ1 , Θ2 , . . . , Θn ) = n∑ t=2 H ( Θt | Θ1 , . . . , Θt−1 ) +H ( Θ1 ) ≤ n∑ t=1 H ( Θt ) . ( 8 ) Proof . The proof is in the Appendix A. Theorem 2 proved the correlation assumption has smaller information entropy , which may reduce the uncertainty and bring more useful information for the MIL problem . Motivated by the Inference and Theorem 2 , a generic three-step method like Algorithm1 was developed . The main difference between the proposed algorithm and existing methods is shown in Figure 2 . Algorithm 1 : A generic three-step approach under the correlated MIL Input : The bag of instances Xi = { xi,1 , xi,2 . . . , xi , n } Output : Bag-level predicted label Ŷi 1 ) Extracting morphological and spatial information of all the instances by f and h , respectively ; Xf ← f ( Xi ) , Xh ← h ( Xi ) , Xfh ← Xf +Xh , where Xf , Xh , Xfh ∈ Rn×d ; 2 ) Aggregating the extracted information for all instances by Pooling Matrix P ; XP ← PXfh , where P ∈ Rn×n ; 3 ) Transforming XP to obtain the predicted bag-level label by g ; Ŷi ← g ( XP ) .

Summary: The paper presents a transformer based approach for WSI classification, correlation among instances are considered to address the iid assumed by previous methods. The methods were tested on three public datasets, i.e. Camelyon16, TCGA-NSCLC and TCGA-RCC datasets, and the results suggest that the performance for binary and multiple category classification is better than state of the art.