Multi-head Attention-based Deep Multiple Instance Learning

Authors: hassan  keshvarikhojasteh
josien p w pluim
mitko  veta
Venue: MICCAI_org_2024_Workshop_COMPAYL
Published: 2024-06-10 11:23:24

Abstract: This paper introduces MAD-MIL, a Multi-head Attention-based Deep Multiple Instance
Learning model, designed for weakly supervised Whole Slide Images (WSIs) classification in digital pathology. Inspired by the multi-head attention mechanism of the Transformer, MAD-MIL simplifies model complexity while achieving competitive results against advanced models like CLAM and DS-MIL. Evaluated on the MNIST-BAGS and public datasets, including TUPAC16, TCGA BRCA, TCGA LUNG, and TCGA KIDNEY, MAD-MIL consistently outperforms ABMIL. This demonstrates enhanced information
diversity, interpretability, and efficiency in slide representation. The model’s effectiveness, coupled with fewer trainable parameters and lower computational complexity makes it a promising solution for automated pathology workflows. Our code is available at https://github.com/tueimage/MAD-MIL.

Main paper: 
 introduction: 
 Digital pathology leverages digital imaging technology to create high-resolution images of pathology slides, which can then be viewed, analyzed, and shared electronically. The adoption of digital pathology offers several advantages, including increased efficiency, improved collaboration, and the potential for more accurate and reproducible diagnoses Williams et al. (2017) . Deep learning has emerged as a powerful tool in digital pathology, offering advanced image analysis capabilities. Convolutional Neural Networks (CNNs) Krizhevsky et al. (2012) and recently Vision Transformers (ViTs) Dosovitskiy et al. (2020) excel in learning complex patterns from large datasets and have been applied to various aspects of digital pathology Coudray et al. (2018) ; Litjens et al. (2017) ; Campanella et al. (2019) . Despite the promising potential of deep learning in digital pathology, its adoption encounters various challenges Esteva et al. (2019) ; Holzinger et al. (2017) ; Madabhushi and Lee (2016) . Challenges include the need for large and diverse annotated datasets, model generalization across different institutions and populations, and ethical considerations. Multiple Instance Learning (MIL) has been widely utilized in digital pathology, to address challenges associated with weakly labeled annotated datasets. In MIL, data is organized into bags, and each bag contains multiple instances, with the label of the bag determined by the presence of at least one positive instance. Attention-based deep Multiple Instance Learning (ABMIL) Ilse et al. (2018) introduces an attention mechanism that allows the model to selectively focus One of the key components of the Transformer architecture Vaswani et al. (2017) is the multi-head self-attention mechanism, which runs the attention mechanism in parallel across multiple heads to capture diverse aspects of relationships. Although recent models advanced the state of the art for weakly supervised WSIs classification, they increased the complexity of the ABMIL model with respect to architecture design and number of trainable parameters. Aiming for simpler models, we propose the Multi-head ABMIL (MAD-MIL) model for weakly supervised WSIs classification. It is inspired by the multi-head self-attention mechanism not only to reduce the number of trainable parameters but also to incorporate various facets of the input WSI using multiple attention heads. Although there is a multi-head version of ABMIL available in the official repository, our method differs from that implementation. In ABMIL, attention weights are computed using full-dimensional input features, and a larger classifier is employed for the larger slide representation. Consequently, this approach increases the number of parameters and the overall complexity of ABMIL. While similar ideas have been explored in previous works such as Li et al. (2019) for multi-tagging WSIs and ACMIL for weakly supervised WSI classification, our model distinguishes itself from Li et al. (2019) in two key ways. First, we adhere to the original architecture proposed in ABMIL to maintain simplicity in the developed model. Second, the MAD-MIL is specifically designed for the task of weakly supervised WSIs classification while the model proposed by Li et al. (2019) was designed for multi-tagging WSIs using the Multi-Tag Attention Module to construct tag-related global slide representations for final prediction. Regarding ACMIL, several attention weights are computed for input features alongside the introduction of two supplementary regularization terms: semantic and diversity regularizations. In contrast to ACMIL, our approach within ABMIL does not entail the extraction of attention weights utilizing full-dimensional input features. This choice is made to avoid the addition of extra parameters. Furthermore, we abstain from introducing additional regularization terms, thereby preserving the simplicity of the ABMIL framework. We evaluate MAD-MIL for weakly supervised classification both in controlled settings, using variants of the MNIST-BAGS dataset, and on four publicly available clinical datasets, namely, TUPAC16, TCGA BRCA, TCGA LUNG, and TCGA KIDNEY. 
problem formulation: 
 If we consider the input WSI as I, then our goal is to predict its slide-level label Y . As it is infeasible to feed the I directly to the network due to computational constraints, it is tiled into N small instances I = {p 1 , p 2 , . . . , p N } to create a bag of instances, with unknown labels. To predict the slide label, in the first step the instances' features are extracted using a pre-trained network and compressed by a fully connected layer{f 1 , f 2 , . . . , f N } = F (pre - trained{p 1 , p 2 , . . . , p N }), ∀n f n ∈ R 1×D .In the second step, the features are aggregated with an attention-based function to obtain the slide embeddingZ = H(f 1 , f 2 , . . . , f N ), Z ∈ R 1×D .Finally, the slide-label is produced with a fully connected layer Ŷ = L(Z). Figure 1 shows the overall framework of the approach.In this method, two types of bags (positive and negative) are considered for binary WSIs classification. Positive bags contain at least one key-instance, while there are none in negative bags. Then, two types of attention functions are proposed, which we only present the gated version for brevity Ilse et al. (2018) : a n = exp{w ⊤ (tanh(Vf n ⊤ ) ⊙ sigm(Uf n ⊤ ))} N l=1 exp{w ⊤ (tanh(Vf l ⊤ ) ⊙ sigm(Uf l ⊤ ))} (1)With U, V ∈ R d×D and w ∈ R d×1 learnable parameters, ⊙ an element-wise multiplication and sigm() the sigmoid function. The attention weights a n ideally assign higher values to key-instances and lower values to non-key ones, providing interpretability of the model.Empirical evidence suggests that the multi-head self-attention mechanism in the Transformer outperforms its single-head counterpart Vaswani et al. (2017) . Additionally, by leveraging the attention weights from each head, we can generate diverse heatmaps that enhance the interpretability of the network. Since our goal is to employ the attention function from ABMIL, which differs significantly from the scaled dot-product self-attention module in the Transformer, a direct adaptation of the multi-head idea is impractical. Therefore, we propose to divide the input features into M equal parts, i.e. ∀n, m f n,m ∈ R 1×⌈ D M ⌉ ; n = {1 : N }, m = {1 : M }. Then, we feed each part to distinct attention modules given by:a n,m = exp{w m ⊤ (tanh(V m f n,m ⊤ ) ⊙ sigm(U m f n,m ⊤ ))} N l=1 exp{w m ⊤ (tanh(V m f l,m ⊤ ) ⊙ sigm(U m f l,m ⊤ ))} (2)In the next step, the features of each part are aggregated and ultimately, slide representation is calculated as the concatenation of the aggregated features: EQUATIONThe slide label can be predicted same as before Ý = L( Ẑ). Our developed method (MAD-MIL) has fewer trainable parameters, as we use smaller attention modules and it can return various attention heatmaps. However, one more hyperparameter M (number-ofheads) should be tuned. Figure 2 presents ABMIL and MAD-MIL models. 
dataset: 
 We used four public WSI datasets. The TUPAC16 dataset Veta et al. (2019) comprises 821 H&E slides from the Cancer Genome Atlas (TCGA) project, annotated with breast tumor proliferation labels. We categorize the labels into two classes: low-grade (slides labeled as grade 1) and high-grade (slides labeled as grade 2 or 3). The TCGA Breast Invasive Carcinoma Cancer (TCGA BRCA) dataset includes 1038 WSIs for two-types of breast cancer; Invasive Ductal Carcinoma (IDC) versus Invasive Lobular Carcinoma (ILC) with slide-level labels. The TCGA Lung Cancer dataset is a public dataset for Non-Small Cell Lung Carcinoma (NSCLC) subtyping. The dataset contains 1046 slides for Lung Adenocarcinoma (LUAD) and Lung Squamous Cell Carcinoma (LUSC) with only slidelevel labels. The TCGA Kidney Cancer includes 918 slides for three types of kidney cancer, namely Clear Cell Renal Cell Carcinoma (CCRCC), Papillary Renal Cell Carcinoma (PRCC), and Chromophobe Renal Cell Carcinoma (CHRCC).We also adopt the approach proposed in Ilse et al. (2018) to construct a dataset by utilizing images from the MNIST image dataset. In practical applications, particularly in weakly supervised WSI classification, a fundamental assumption of conventional MIL is challenged, as negative bags may also encompass some key-instances Li (2015). To address this, we introduce soft-bags and designate the digit 8 as the key-instance and create positive and negative bags with varying numbers of the key-instance.For the MNIST-BAGS dataset, we create 50 bags for training, 100 for evaluation, and 900 for testing, every bag containing 20 images. In our experimental setup, we explore varying percentages of the key-instance within positive and negative bags (p-pos, p-neg.) Specifically, in our analysis, we consider the (0.4, 0.2), (0.6, 0.4), and (0.8, 0.6) pairs. These pairs correspond to positive bags with values 0.4, 0.6, 0.8 and negative bags with values 0.2, 0.4, 0.6. We directly input MNIST images into the model, (input dim = 28 × 28 = 784 pixels), and then they are compressed to a size of D = 128. Throughout these experiments, we utilize the Adam optimizer and train models for 20 epochs and identify optimal learning rate and weight decay values based on the validation loss. We repeat the experiments with ten different seeds and report the mean AUC and F1 values to maintain consistency with prior methodologies. For the preprocessing of the TUPAC16 and TCGA datasets, we adhere to the steps outlined in Lu et al. (2021) , incorporating the extraction of features from non-overlapping 256×256 patches at 20× magnification using the ImageNet pretrained ResNet50. We feed extracted features into the model, where the feature dimensionality is input dim = 1024. These features are then compressed to a size of D = 512. We repeat the experiments five times to evaluate our framework on the TUPAC16 dataset, as there is an official test split. We report the average performance of the test set for the five experiments. However, for the TCGA datasets we employ the same 10-folder cross-validation setup as adopted in Chen et al. ( 2022), and present the mean and standard variance values of AUC and F1 scores. The training of the models for these datasets involves using the Adam optimizer for 50 epochs. A learning rate of 1e-4 is employed, and we tune the weight-decay value based on the validation loss.For all experiments we determine the optimal number of heads (M) according to the validation loss value. Regarding the Mean Pool and Max Pool models, we ensure a fair comparison by implementing the bag embedding-based variant as opposed to the instancelevel. Specifically, after the feature compression step, we calculate the slide representation using either the max or mean operation and subsequently predict the slide label. In what follows we use brief notation to indicate the number of heads for the MAD-MIL: for instance, MAD-MIL/3 means the MAD-MIL model with three heads.Table 1 reports the outcomes obtained from the MNIST-BAGS. Across all three configurations, the MAD-MIL model consistently exhibits superior performance compared to AB-MIL, despite having a lower number of trainable parameters and computational complexity. Notably, the most substantial enhancements are observed in the last setting (p-pos= 0.8, p-neg= 0.6), where AUC and F1 metrics experience notable increases of 4.3% and 12.7%, respectively. The results for the weakly-supervised WSI classification are presented in Tables 2, 3, 4 and 5. Notably, the MAD-MIL model outperforms ABMIL across all datasets. Interestingly, it achieves comparable performance to sophisticated networks like CLAM and DS-MIL. It is worth noting that while the Mean-Pool and Max-Pool models exhibit similar performance to other models, their lack of interpretability hinders their practical application.We present attention heatmaps in Figure 4 to assess the interpretability of the methods applied to a LUAD slide from the TCGA LUNG dataset. Notably, all three methods (ABMIL, CLAM-MB, and MAD-MIL/8) focus on similar regions within the input WSI, demonstrating negligible differences. However, MAD-MIL distinguishes itself by generating additional heatmaps. This enhancement could potentially contribute to the overall interpretability of the network. 
conclusion: 
 In 5 Supplementary Material 
qualitative results on mnist-bags: 
 We present the attention weights of the ABMIL and MAD-MIL models on the MNIST-BAGS dataset in Figure 4 to delve deeper into their interpretability. Unlike ABMIL, which distributes attention across both key and non-key instances, certain heads of the MAD-MIL model focus solely on either key instances (e.g., head-2 and head-4) or non-key instances (e.g., head-1 and head-6).Figure 5 presents the AUC, ACC, and Val Loss of the proposed method with respect to various number of heads on the datasets. These curves indicate that the multi-head idea is beneficial to both AUC and ACC metrics. However, ACC is more sensitive to the number of heads, with a significant drop observed for certain head counts. Conversely, AUC tends to consistently improve with an increase in the number of heads. Looking ahead, MAD-MIL opens avenues for further research and exploration. Future work could involve adapting the model for specific pathology tasks or refining its applicability across diverse clinical settings.  
 
 Appendix:  



Reviews: review: title: title: A lighter multi-head attention-based MIL framework for pathology application
main_review: main_review: The authors present a new lighter (in terms of number of parameters) multi-head attention-based MIL algorithm and evaluate it on several pathology datasets (TUPAC16, and several TCGA datasets) as well as bags created from MNIST data. The paper is well written and the results indicate that the proposed method works well for pathology applications. However, the discussion could be expanded as currently many results are presented in the tables and not properly discussed. For example, ACMIL method (which also uses multi-head approach) performs on part or better than the proposed method, hence this could have been discussed in more detail and suggestions made why the reader could still prefer their proposed method. 

The experiments on MNIST data with varying fraction of important instances in positive and negative bags is very interesting. It would be interesting to replicate this experiment with pathology data to see if the results are similar. Of course, this would require to have detailed labels which can be expensive to do. Finally, some details of implementation are missing such as what network was used to extract features, i.e. what data was it trained on.
title: title: Next step of ABMIL with some efficiency gains.
main_review: main_review: The paper outlines the multihead attention formulation of the WSI MIL problem. Compared to ACMIL, the efficiency gains appear modest. I would love to see a more real-world metric that captures the performance / compute (FLOPs) tradeoff that demonstrates the clear superiority of the proposed method.

As the paper stands, the multihead formulation makes sense. The metrics look good. The new method is a good reference for people looking to develop attention based MIL models.
title: title: Review of paper 1
main_review: main_review: The paper extends ABMIL within the context of the slide-level classification task in pathology, introducing multi-head attention to map and aggregate patch-level embeddings into slide-level embeddings. The paper extensively validates the approach on real pathology slide datasets from TUPAC16, TCGA-BRCA, TCGA-Lung, and TCGA-Kidney, demonstrating the superiority of using multi-head attention. However, there are some details that need addressing.

1. Which pre-trained network was used as the feature extractor? This should be specified.

2. Figure 5 visualizes the impact of the number of heads on performance. To facilitate comparison with the baseline (a single head), the result for "1" head should be added to the figure. Additionally, why are the experiment head counts not consistent across different experiments (e.g., some range from 2-5, while others from 2-10)? This should be standardized.

3. What do p-pos and p-neg represent? This should be clarified.

4. Moreover, the article appears to group based on the order of channels. I would be curious to know how much performance fluctuation occurs if the channel order is shuffled.