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 authors propose a Transformer-based MIL method for WSI classification.   The method is designed to avoid the IID assumption inherent to standard MIL inference methods (e.g. mean, max pooling). Secondly, they propose a positional encoding module for patches in WSIs based on CNNs to incorporate spatial information into the prediction, along with semantic information learned in the MSA layers. In the ablation study, they demonstrated that PPEG model performed better than conventional sin-cos positional encoding. They used three WSI classification datasets from open sources with previously benchmarked results. The authors were able to achieve SOTA results on each task.