RPMIL: Rethinking Uncertainty-Aware Probabilistic Multiple Instance Learning for Whole Slide Pathology Diagnosis

1 Introduction

Deep neural networks have made significant strides in recent years, benefiting various fields worldwide. In the field of computational pathology, the application of various deep learning methods has led to new insights into disease diagnosis. Whole slide image (WSI) pathology diagnosis is regarded as the gold standard for identifying or excluding tumors and remains a hot topic of research. Due to the gigapixel resolution of WSIs and the absence of fine-grained patch labels, current methods predominantly follow the multiple instance learning (MIL) paradigm [Tellez et al., 2019; Pinckaers et al., 2020; Zhang et al., 2022]. This approach divides a WSI (bag) into many small patches (instances), where the presence of even a single positive instance can lead to the bag as positive.

∗Corresponding author: Jing Zhao.

Existing research primarily focuses on enhancing instance feature representation, such as through self-supervised learning [Huang et al., 2023; Lu et al., 2024; Chen et al., 2024b] or instance feature re-embedding [Chikontwe et al., 2022; Tang et al., 2024]. More work is dedicated to the design of bag aggregators, using methods like modifying instance weight allocation [Li et al., 2021] or designing Transformer-based aggregators [Shao et al., 2021]. Recently, some novel studies have begun to address false relationships in MIL [Lin et al., 2023; Chen et al., 2024a], recalculating intervention expectations to replace traditional likelihood estimation. In practice, training data is often limited, typically consisting of only a few hundred or a few thousand bags. In a deterministic embedding space, a small number of bag features can easily fit the decision boundary, but this fitting is often fragile and prone to overfitting. The bag feature obtained in this manner are suboptimal. Although existing research has explored data augmentation through splitting or mixing existing bags to generate pseudo-bags [Zhang et al., 2022; Chen and Lu, 2023; Liu et al., 2024], the number of training samples is typically only expanded by a limited factor.

From the perspective of uncertainty estimation [Kendall and Gal, 2017], we rethink the process of MIL modeling. We observe that these MIL methods fall under point estimation, where each bag is mapped to a deterministic embedding. Conversely, uncertainty estimation maps input samples to an embedding distribution, better capturing the bag feature distribution and aiding in understanding the confidence of the estimation. Figure 1(a) illustrates the process of point estimation in MIL, while Figure 1(b) depicts the probabilistic MIL utilizing bag uncertainty. The essential difference between the two lies in whether the bag feature is represented as a single value or a probability distribution. Furthermore, we consider the instance features within a bag as a distribution and assume that the instance feature distribution can also aid in prediction. This assumption aligns with the logical process in actual clinical diagnosis, where pathologists identify tumor regions (instance) to assess whether WSI is a tumor [Xiong et al., 2023]. Consequently, we jointly utilize the bag distribution and the instance distribution to drive predictions, as shown in Figure 1(c). Two main issues need to be verified when applying uncertainty estimation in MIL: i) Can uncertainty estimation outperform the dominant point estimationMILmethods? Specifically, this involves comparing P(y|z = g(x)) and ∫z P(z|x)P(y|z) dz. ii) Can instance distribution and bag distribution jointly enhance classification results? This involves comparing likelihood P(y|z) and P(y|z, x).

In this paper, we rethink probabilistic modeling in MIL and propose RPMIL, an uncertainty-aware probabilistic MIL method for whole slide pathology diagnosis. RPMIL learns a probabilistic aggregator that consolidates instance features into a bag feature distribution rather than the fixed bag feature. Specifically, the probabilistic aggregator uses a variational autoencoder approach [Cemgil et al., 2020; Michelucci, 2022] to compress multiple instance features into two low-dimensional latent space features, representing the mean and variance of the bag feature distribution, respectively. By reparameterization sampling [Kingma et al., 2015], an arbitrary number of bag feature representations can be obtained. We use Kullback-Leibler (KL) divergence and mean squared error (MSE) loss to constrain the bag feature distribution, preventing distribution drift. Importantly, we propose combining instance and bag feature distribution for the final prediction. Through extensive experiments and analysis, we reach two significant conclusions: i) Uncertainty estimation based on bags surpasses point estimation methods in MIL. ii) Combining instance feature distribution and bag feature distribution for prediction significantly outperforms using a single distribution. Through ablation studies, we further find that as the number of sampling in the bag feature distribution increases, the results become more pronounced.

2 RelatedWork

2.1 Multiple Instance Learning forWSIs

There are two main methods for predicting bags: the instance-based method and the bag-based method. A typical instance-based approach trains a instance classifier to assign pseudo-labels to each instance based on bag labels [Qu et al., 2024b; Lin et al., 2024]. Due to the large number of instances, the top-k policy is often used to select instances [Xu et al., 2019; Chikontwe et al., 2020], necessitating a very large number of WSIs [Chen and Lu, 2023]. Given the fact that WSIs lack pixel-level labels and contain both cancerous and normal regions, bag-based approaches have become dominant. The rapid development of deep learning, particularly the exploration of attention mechanisms [Vaswani, 2017; Ilse et al., 2018], has led to the rise of bag-based methods that can directly predict the label of a bag by aggregating instance features into a bag feature [Li et al., 2021; Shao et al., 2021; Zhang et al., 2022]. The advantage of this approach is that it does not require the labels of numerous instances, so only the bag labels are needed for classification. Some studies also explore improving bag feature representation, employing multi-scale learning [Xiong et al., 2023; Qu et al., 2023] or fusing multimodal information [Qu et al., 2024a; Li et al., 2024; Shi et al., 2024] to achieve a better representation of the bag. The above MIL methods fall under point estimation and fail to capture the full spectrum of data variability. Therefore, we propose RPMIL, an uncertaintyaware probabilistic MIL method to learns a probabilistic aggregator.

2.2 Probabilistic Methods inMIL

Introducing a probabilistic model in multiple instance learning can adequately take into account uncertainty [Haußmann et al., 2017; Kendall and Gal, 2017]. DGMIL [Qu et al., 2022] is an instance-based MIL method that considers instance feature distribution while assigning pseudo-labels to instances and focuses on constructing instance-level feature distribution. Bayes-MIL [Cui et al., 2023] uses the uncertainty of the attention weight of each instance as a measure of the accuracy of guessing whether the instance is positive or negative. However, our RPMIL is a bag-based method to construct distributions for bag features and we drive the prediction by jointly utilizing instance feature distribution and bag feature distribution. AGP [Schmidt et al., 2023] is also a bag-based method that utilizes Gaussian processes to dynamically obtain attention scores for each instance. In contrast, our RPMIL learns a probabilistic aggregator based on a generative model, constructing the bag features as a distribution, which allows for sampling to capture uncertainty. We note that Causal inference methods such as IBMIL [Lin et al., 2023] and CaMIL [Chen et al., 2024a] have also made adjustments to the probability calculation formula by using interventional likelihood expectations. Yet, these causal methods still fall under point estimation, and they rely on sampling based on clustering. Conversely, our RPMIL constructs a normal distribution about bag features. We perform sampling via reparameterization and utilize MCMC [Rubinstein and Kroese, 2016] to approximate the desired quantity.

3 Methodology

3.1 Preliminary

The MIL methodology involves the following steps: first, instance features are obtained using a feature extractor with fixed weights. Then, these instance features are aggregated into a bag feature by assigning attention scores or through global integration. Finally, these bag features are used for classification. For a given dataset of WSIs, each WSI has a corresponding label yi ∈ {0, 1} for the binary classification task, which determines whether a certain cancer is present. Moreover, each WSI referred to as a bag is divided into many patches pi, i.e., instances, {pi,1, pi,2, . . . , pi,ni } and ni is the number of instances from a bag. The bag label yi is defined as follows:

where yi,j corresponds to the label of pi,j.

A challenging problem arises because the number of instances usually reaches thousands or more. We only have access to bag labels yi. Our goal is to predict the bag label when the instance labels are unknown. This whole process of MIL can be summarized as follows:

where f (·) is a pre-trained feature extractor used to obtain the instance features x, g(·) is an aggregator that converges all instance features x into a bag feature z, and h(·) is a classifier to get the predicative result ˆy. The optimization objective is to minimize the cross-entropy (CE) loss between the real label y and the prediction ˆy.

There is a dependency among all the variables due to their sequential relationship, as illustrated in Figure 1(a). The traditional MIL methods fall under point estimation, and the prediction probability can be defined as follows:

In this approach, a bag is mapped to a deterministic embedding, which is then used to predict its output. P(y|z) is the final prediction probability.

3.2 Uncertainty-Aware Probabilistic MIL

MIL methods based on point estimation fail to capture the full spectrum of data variability due to their reliance on a fixed embedding, especially when the number of trainable bags is limited. Conversely, uncertainty estimation is better equipped to capture data diversity, as it maps input samples to an embedded distribution. Thus, we propose an uncertainty-aware probabilistic MIL method for whole slide pathology diagnosis, as shown in Figure 2. Our method learns a probabilistic aggregator that aggregates instance features into a bag feature distribution P(z|x) rather than a fixed bag feature z. Then, we obtain z by sampling from the dynamic distribution space, allowing us to derive a range of different values that provide openings for the bag classifier. Here, the specific prediction probability equation is different from Equation (3) and is defined as follows:

Now, we obtain the probability by sampling z from the distribution P(z|x) and then performing marginal integration. In an implementation, we use Markov Chain Monte Carlo (MCMC) [Rubinstein and Kroese, 2016] to approximate the desired quantity.

Furthermore, we conceptualize the instance features within a bag as representing a distribution. We consider that jointly leveraging these distributions can enhance the final prediction, as shown in Figure 1(c). This approach mirrors the clinical diagnostic process, wherein pathologists use tumor areas

(instance) to assess WSI. We modify the probability P(y|x) in Equation (4) to obtain the P(y|x, z), and the final prediction is as follows:

For detailed formula derivations, please refer to the supplementary materials. To calculate P(x|z), we use a crossattention [Vaswani, 2017] module, where z is act as query and x is act as key. We fuse the joint distributions of instance and bag:

where N(·) is the layer normalization, z ∈ Rk×d , x ∈ Rn×d , Wq ∈ Rd×d , Wk ∈ Rd×d and Wv ∈ Rd×d are learnable parameters of query, key and value, respectively.

Eventually, we use the pooling operator and feed it to the classifier h(·) to get the final predicted labels:

3.3 Building Bag Distribution

In our probabilistic framework, it is necessary to construct a distribution of bag features. Specifically, we compress multiple instance features into a new distribution for classification purposes. Here, we use a variational autoencoder to compress the instance features into two low-dimensional latent space features, representing the mean and variance of the bag feature distribution, respectively. We use KL loss to regularise the distribution P(z|x). We assume that P(z|x) obeys a normal distribution, with the mean µ and variance σ2 obtained by the encoder. Given the assumption that Q(z) is a prior that follows a standard normal distribution. The KL loss is defined as follows:

The benefit of keeping P(z|x) close to a standard normal distribution is the benefit of ensuring that the feature space is more compact and distinct, which aids the classifier in making better distinctions. Now, we have reformulated Equation (2), and the bag feature distribution is represented as follows:

By reparameterization sampling, an arbitrary number of bag features z can be obtained. To prevent excessive differences between features sampled from uncertainty estimation and those obtained from point estimation, we use MSE to constrain the differences between the mean of the sampled features ¯z and the point-estimated feature z′ . The overall loss is summarised as follows:

where λ1, λ2 and λ3 are hyperparameters.

4 Experiments

4.1 Datasets

The effectiveness of the proposed method is examined on both the Camelyon16 [Bejnordi et al., 2017] and the TCGANSCLC datasets. The Camelyon16 dataset, from a competition organized by ISBI to classify breast cancer, is officially divided into a training set and a test set, with 270 WSIs in the training set and 129 in the test set.

The TCGA is a joint project of the National Cancer Institute and the National Human Genome Research Institute, including clinical data on a wide range of human cancers. We select two types of non-small cell lung cancer (NSCLC) WSIs for classification: lung adenocarcinoma (LUAD) and lung squamous carcinoma (LUSC). TCGA-NSCLC dataset contains a total of 1053 WSIs, with 541 LUAD from 478 patients and 512 LUSC from another 478 patients. We divide the dataset into training, validation, and test sets in a 6:1.5:2.5 ratio.

4.2 Preprocessing and Implement Details

Following CLAM [Lu et al., 2021], all WSIs are divided into non-overlapping patches at 20× magnification with a window size of 256×256 pixels. Our statistics show that the Camelyon16 has a total of approximately 3.617 million patches, with an average of about 9,066 patches per WSI, while the TCGA-NSCLC has a total of about 12.731 million patches, with an average of about 12,102 patches per WSI.

All patches are extracted using ResNet-50 pre-trained on ImageNet-1k [He et al., 2016], resulting in features with a dimension of 1024. Additionally, PLIP pre-trained on largescale pathology images is used to obtain 512-dimensional features [Huang et al., 2023]. We use the Adam optimizer with an initial learning rate of 1e-4 and a weight decay of 1e-5, adjusting the learning rate using a cosine annealing scheme. The minibatch is set to 1, the number of epochs is set to 100, and the hyperparameter λ1 is 1, λ2 and λ3 are both set to 0.5. We record accuracy (ACC), F1 score, and area under curve (AUC) as criteria for evaluating our models. All experiments are performed on an NVIDIA GeForce RTX 3090 and are repeated five times, with all metrics reported as averages.

4.3 Baseline Methods

We compare many classical multiple instance learning methods for WSIs classification, including ABMIL, CLAM, TransMIL, DGMIL, DTFD, MMIL and DGR-MIL. We use the vanilla attention-based method as part of our default bag probabilistic aggregator. ABMIL enables the aggregation of instance features by assigning different attention scores to each instance. CLAM employs a multi-branch attention network that trains the instance classifier to obtain higherscoring instances. TransMIL and MMIL are multiple instance learning networks based on the Transformer architecture that take into account information between instances. DTFD is a dual-layer framework which introduces pseudo-bags to increase the number of bags, and the pseudo-bags are still at the same resolution. DGMIL performs WSI classification from the perspective of data distribution, while it focuses on instances rather than bags. DGR-MIL takes into account the relationship between global vectors and instance embeddings, but diverse global vectors make training unstable.

4.4 Experimental Results

As shown in Table 1, our method achieves the best results on both the Camelyon16 and the TCGA-NSCLC datasets, in terms of all three evaluation metrics: ACC, F1-score, and AUC. Especially, our method has a significant improvement on the Camelyon16 dataset compared with DGMIL, a method for constructing instance feature distribution, 9.55% increase on ACC, 10.23% increase on F1-score, and 9.42% increase on AUC. This improvement can be attributed to the limitations of DGMIL, which relies on clustering techniques to construct instance distributions and requires positive instances to train the instance classifier. On the Camelyon16 dataset, however, the small cancer regions in the WSI positive samples result in a limited number of available positive instances, hindering DGMIL's performance. In contrast, our method handles such imbalances more effectively, demonstrating better performance under these challenging conditions. Notably, these methods usually achieve better results on the TCGA-NSCLC compared with the Camelyon16. In particular, the two constructed baselines, mean-pooling and max-pooling [Wang et al., 2018], also performe admirably on the latter. Specifically, the TCGA-NSCLC dataset contains a larger number of WSIs, and the distribution of positive and negative samples is more balanced compared to the Camelyon16 dataset. The average tumor region in positive samples of the TCGA-NSCLC dataset is significantly larger, often around 80%, whereas, in the Camelyon16 dataset, the tumor region in positive samples is much smaller, averaging only about 10% [Shao et al., 2021].

4.5 Main Analysis

In this subsection, we primarily address the two main questions posed in the introduction: i) Can uncertainty estimation outperform the dominant point estimation MIL methods? ii) Can instance distribution and bag distribution jointly enhance classification results?

Comparison between Point Estimation and Uncertainty Estimation (P(y|z = g(x)) in Equation (3) vs ∫z P(z|x)P(y|z) dz in Equation (4). In point estimation methods, multiple instance features are aggregated into a single feature, whereas uncertainty methods aggregate them into a probability distribution. To compare these two types of MIL methods, we conduct experiments with both parameterized aggregators (using vanilla attention) and non-parameterized aggregators (using mean-pooling and max-pooling). The results, as shown in Table 2, reveal that under parameterized aggregators, uncertainty methods significantly outperform point estimation methods. We acknowledge that a highperformance aggregator can enhance the performance of MIL models. However, our experiments reveal that even with a basic pooling method, MIL models can still achieve significant benefits from uncertainty estimation and the joint distribution of instance and bag features. This finding suggests that, despite the focus on aggregator design in MIL, uncertainty estimation methods should be given more attention in practical WSI applications.

Comparison between Two Predictive Distributions (P(y|z) in Equation (4) vs P(y|z, x) in Equation (5)). The limitation of point estimation methods lies in their sole focus on the aggregated, fixed bag feature representation, overlooking the inherent distribution of instance features within the bag. To better utilize the information from instance feature distribution, we propose a method that jointly utilizes instance distribution and bag feature distribution. In this process, we remove the cross-attention branch and simplify the modeling P(y|z). As shown in Table 2, the experimental results indicate that using only P(y|z) leads to a significant shift in the bag feature distribution. In contrast, our method, which jointly leverages instance distributions and bag feature distributions, significantly improves the model's predictive performance, highlighting the importance of considering instance distributions. The effectiveness of this approach aligns with clinical observations, as tumor patches can confirm the tumor WSI.

4.6 Ablation Results

In this subsection, we investigate the proposed method in more detail, showing the contribution of each component through a series of ablation experiments.

Effect of Sampling Number. In our study, the probabilistic model makes predictions by learning the distribution of bag features, and the sampling number significantly impacts the experimental results. In Figure 3, we test different sampling numbers (1, 10, 100, and 1,000) and find that a larger sampling size effectively improves model performance and stabilizes training. With a sampling size of 1,000, the model obtains more enriched bag representations, achieving the best performance. When sampling sizes are small, the sampled bag features often fail to represent the overall distribution, leading to a significant deviation of the sample mean from the actual center of the bag feature distribution.

Importance of the Variance. Due to the introduction of variance, the results of each iteration will fluctuate, and the degree of fluctuation is related to the magnitude of the variance. We roughly think that larger variance indicates more uncertainty in the bag features, making the model unable to adequately estimate the results of classifying such bag features. We verify the above thought by removing some test samples with a large variance in bag features. We compare this with randomly removing a portion of the sample. We observe that the results become more favorable in Figure 4. The ACC improves by 0.46%, 0.75% and 1.19%, the F1-score improves by 0.20%, 0.29% and 0.44% and the AUC improves by 0.99%, 1.21%, and 1.32% when 5%, 10%, and 15% of the test set with high variance are removed on the TCGANSCLC. As the number of removed samples increases, there is an improvement, indicating that variance can be an effective measure of uncertainty to some extent.

Impact of KL and MSE Loss. Compared to the single feature in point estimation, we learn an enriched feature distribution. To prevent significant deviations in the features, we introduce additional KL and MSE loss functions to constrain the bag feature distribution, in addition to using cross-entropy loss. From Table 3, it can be concluded that the absence of either two loss functions degrades the model performance, especially on the Camelyon16. The effect of the missing MSE loss is greater than the KL loss, which keeps the pre-encoding and post-decoding bag features from deviating significantly, while the KL loss contributes to reducing the variance of the classification results.

Combination of Probabilistic Aggregator and Bag Pooling. We compare the effectiveness of learning bag features using three different aggregators: mean-pooling, maxpooling, and vanilla attention-based aggregators. Additionally, there is another bag pooling operation before classification to pool the fusion bag features, and we compare two pooling: mean-pooling and max-pooling. In Table 4, the best combination is the attention-based probabilistic aggregator and the bag mean-pooling. The former is able to learn a fine bag distribution, and the latter enables bag features sampled from the distribution to be near the center of the distribution.

4.7 Visualization

In this subsection, we follow the visualization technique in CLAM to assess whether the attention mechanism in the distribution aggregator is appropriately focused on the positive tissue regions. This evaluation is crucial as it helps to provide further confidence in the reliability of our approach and holds significant promise for real-world clinical applications, where precision and interpretability are paramount for diagnosis. In Figure 5, both our model and the classical CLAM successfully identify the general regions of interest that are deemed positive. While CLAM assigns high attention to areas that are contaminated or irrelevant in the original WSI, our method, enhanced by the incorporation of uncertainty, refrains from considering these areas as positive. This distinction is vital because such contaminated regions, often arising due to issues like uneven staining or sample folding during WSI acquisition, can lead to misclassifications. By preventing these areas from being mistakenly labeled as positive, our approach maintains a higher level of accuracy in classification, demonstrating robustness in the presence of common challenges faced during sample collection and preparation.

5 Conclusion

In this paper, we rethink the process of MIL modeling from the perspective of uncertainty estimation. We propose RPMIL, an uncertainty-aware probabilistic MIL method tailored for WSI pathology diagnosis. Unlike traditional methods that rely on point estimates, RPMIL constructs a bag feature distribution, enabling more excellent results. Our experimental results demonstrate that uncertainty estimation surpasses conventional point estimation approaches in MIL, achieving SOTA performance. Moreover, combining instance feature distribution with bag feature distribution for prediction yields significantly better results than relying on a single distribution alone. We also observe that increasing the number of samples in the bag distribution leads to more pronounced improvements in classification performance. Additionally, the variance in uncertainty serves as a supplementary support for enhancing classification accuracy.

Acknowledgments

This work was supported by the National Natural Science Foundation of China under Project 62476089 and the Foundation of Key Laboratory of System Control and Information Processing, Ministry of Education.