Task-oriented Embedding Counts: Heuristic Clustering-driven Feature Fine-tuning for Whole Slide Image Classification

In general, MIL approaches are categorized into instance-based and embedding-based methods. Instance-based methods focus on acquiring labels for individual instances and aggregating them for bag-level classification. Bag-level prediction is aggregated by the probability of all patches with Mean or Max-pooling [7]. As for embedding-based methods, they embed all patch-level features to form the feature set of bags, directly predicting the bag’s features. Due to their superior performance and enhanced interpretability, these methods are widely used in WSI classification.

In the realm of embedding-based MIL, a notable advancement is the attention-based MIL (ABMIL) introduced by [31]. This pioneering method leverages gated attention mechanisms to assign learnable weights to features at the patch level, marking a significant milestone in the evolution of MIL techniques. Concurrently, [41] have enhanced performance by employing attention mechanisms for the selective sampling of positive and negative instances, coupled with cluster constraints to mine patch-level information and integrate bag-level predictions. Furthermore, [34] identifies key instances and their relations with other instances for bag-level embedding, this approach also underscores the potential of self-supervised methods in enhancing model performance. [74] replaced the self-attention mechanism with a multi-head attention mechanism and a masking strategy is applied to patches with high attention to mitigate issues of excessive attention concentration and model overfitting.

However, several MIL models consider only the semantic information and neglect the spatial relationships between patches. This results in the loss of valuable positional information and overlooks the significance of context in pathological diagnosis. This oversight has been addressed by [34, 54, 37]. [34] have utilized multi-scale information to construct pyramid-structured data, thereby implicitly incorporating spatial information. Expanding upon this, [54] have proposed a model based on the transformer architecture to map the interrelations between patches, while [37] have employed backdoor adjustment for intervention training to mitigate biases introduced by contextual priors. Another method [46, 75, 2, 36] to represent the spatial relationships between different patches is through graph theory, where [75] have integrated MIL with deep graph convolutional networks, demonstrating superior performance in predicting lymph node metastasis. Similarly, [2] have applied graph neural networks to learn the relationships between patches, using graph pooling to infer patches of higher relevance automatically.

Other researches concentrate on the representation of bags, [4] and [33] identifying prominent patches to form bag embeddings. [73] introduced the concept of pseudo bags, developing a two-layer distillation framework that forms diverse bags through distinct distillation methods.

These approaches have to decouple the training of the feature encoder and the MIL aggregator to avoid the memory bottleneck [71], resulting in performance degradation.

Self-supervised learning, an emergent training paradigm within representation learning [72, 28, 17, 27, 18, 11, 10]. Unlike supervised learning, which is constrained by the availability of labeled data, self-supervised learning can learn from a vast amount of unlabeled data [16, 42]. In the domain of computer vision, self-supervised learning is capable of generating universal visual features that can even surpass the performance of models trained on labeled data, even on highly competitive benchmarks like ImageNet-1K [62, 28]. It is predominantly employed in computational pathology to construct patch embeddings. Given its capacity to derive domain-specific patch embeddings without the need for annotated data, leveraging extensive pathological data for SSL to obtain pre-trained weights emerges as an effective strategy.

The proliferation of patch-level self-supervised training paradigms within computational pathology has been noteworthy [34, 22, 3]. For instance, contrastive learning facilitates the mapping of patches and their augmented versions [28] (typically through pathology-relevant transformations such as random rotation, cropping, and stain jitter) to analogous embeddings. By utilizing pretext tasks, such as discerning correspondences between global and local-level details [15] or reconstructing randomly masked areas within patches [27], as well as leveraging similarities between patches from the same whole slide image or those sharing labels [65]. Self-supervised training methods, which learn similar spatial semantic information, are increasingly used for pre-trained feature encoders. Despite their ability to extract valuable information without annotations, SSL approaches require massive data and computational resources. As a result, a common approach is to use publicly available pre-trained weights on pathological datasets. Yet, the effectiveness of this strategy is limited due to variations in feature spaces across different tasks.

Early research has explored end-to-end training of MIL methods to update features. One approach centered around optimizing the training workflow, [12] leveraged unified memory architectures, allowing GPUs direct access to host memory for streamlined end-to-end model training. [47] capitalized on the locality of convolution operations in CNN to diminish memory demands, [48] also introduced an enhanced streaming CNN technique, tailored to optimize whole slide images for end-to-end training. These strategies have not achieved widespread adoption due to their computationally demanding nature. An alternative perspective focuses directly on WSIs. [55] developed an end-to-end training framework utilizing patch clustering, and [61] applied unsupervised techniques to compress WSIs. While, in theory, end-to-end training models are posited to deliver optimal outcomes, their propensity for overfitting hampers their generalizability.

Fine-tuning is indispensable for acquiring task-oriented features while it is unavailable for direct use in WSI classification due to memory bottlenecks [71]. A pragmatic strategy involves updating the feature encoder with a subset of samples. This task-oriented method iteratively transfers supervised information from one round to the next to update feature representations. For example, [67, 66] sampled patches from WSIs during training. [51] proposed a weakly supervised knowledge distillation concept to update the feature encoder using soft labels while discarding easily classified samples to construct challenging pseudo bags. [71] amalgamated bayesian and collaborative learning principles, applying a quality-aware strategy to update the feature encoder with pseudo labels. [8] engaged in iterative sampling of patches from WSIs to train the feature pyramid encoder, and [60] employed masking on patches with high attention scores, compelling teacher-student models to learn more discriminative features. Although these methods successfully updated the feature encoder, the information used to update the feature encoder may be noisy, and the mask method to implicitly search for hard negative samples may cause the model to lose some key information from high attention patches.

In a dataset of whole slide images with bag-level labels, denoted as $\mathcal{D}=\{(X_{i},Y_{i})\}_{i=1}^{|\mathcal{D}|}$, where $X_{i}$ represents a slide and $Y_{i}={1,2,\cdots,n}$ represents its corresponding label. Our goal is to establish a mapping from the $\mathcal{X}$ domain to the $\mathcal{Y}$ domain. However, learning direct mapping is extremely challenging due to the gigapixel nature of WSI. Thus, each slide is commonly segmented into patches $X_{i}=\{x_{i,k}\}_{k=1}^{N}$ for multiple instance learning, where $x_{i,k}$ denotes $k$-th patch cut from $i$-th slide, the number $N$ of patches in a slide is often different from others, and patch’s corresponding label $y_{i,k}$ remains unknown during training. In the binary MIL classification problem, the relationship between patch-level labels $y_{i,k}$ and bag-level labels $Y_{i}$ typically follows the following pattern:

where $+1$ represents the positive sample and $-1$ represents the negative sample.

Multiple instance learning involves a three-step modeling process: (1) Instances are transformed into low-dimensional embeddings through a feature encoder:

(2) All instance embeddings are aggregated into a bag-level representation:

where the attention mechanism is often used for aggregation:

where $a_{i,k}=\sigma(h_{i,k})$ represents the attention value of patch $x_{i,k}$. In the case of ABMIL, a gated attention mechanism is adopted as follows:

(3) Predict the aggregated bag-level representation with an MLP as follows:

Task-oriented feature fine-tuning methods are usually based on the idea of distillation which is typically iterative. Take the first round of iteration as an example: (1) extract features with an encoder initialized with pre-trained weights $\omega_{(0)}$ which are trained usually on ImageNet-1K as Eq. 2; (2) train the aggregator $g_{\varphi}$ and MLP $\Gamma{\phi}$ following Eq. 3 and Eq. 6; (3) distill the patch-level information from the $g_{\varphi_{(0)}^{*}}$ and $\Gamma_{\phi_{(0)}^{*}}$; (4) fine-tune the feature encoder $f_{\omega}$ by the patch-level information and update the features with the new encoder $f_{\omega_{(1)}^{*}}$. Subsequently, the aforementioned steps can be repeatedly executed until reaching an optimal outcome or satisfying the termination criteria.

Taking the framework of task-oriented feature fine-tuning methods, a common distillation practice for feature encoder updating is the pseudo label strategy. In this context, a patch-level pseudo label $z_{i,k}$ is assigned to each selected instance $x_{i,k}$. This supervised information is employed to update the feature encoder $f_{\omega}$, making the selection of pseudo labels crucial. As learning too many noisy samples could devastate the fine-tuning process, we first introduce a pseudo label cleaning module. Furthermore, overly simplistic samples distilled may fail to provide the feature encoder with ample information, resulting in overfitting issues. Hence, we introduce a hard negative sample mining module, designed to furnish the training process with adversarial information, thereby enhancing the generalization ability.

Then we will introduce each component of the heuristic clustering-driven feature fine-tuning method (HC-FT) in Fig. 2.

Through the two aforementioned modules, we can obtain a patch-level dataset $\mathcal{D}^{*}=\{(x_{i,k},z_{i,k})\}$ annotated with relatively clean pseudo labels for feature fine-tuning. Notably, the patch-level classification task typically encompasses $2n-1$ classes, where initial $n$ classes consist of positive samples from each category, and the remaining $n-1$ classes represent the corresponding hard negative samples, excluding the non-tumor class.

Our experiments employ two publicly available whole slide image classification datasets, namely CAMELYON16 [26] and BRACS [5].

CAMELYON16 dataset is a binary classification dataset focused on the detection of early-stage breast cancer lymph node metastases. It consists of 399 hematoxylin and eosin (H&E) stained WSIs, sourced from two distinct medical institutions, with 270 allocated for training and 129 for testing. Additionally, this dataset provides detailed pixel-level manual annotations of metastatic tumor areas, thereby facilitating in-depth patch-level analysis and enhanced visualization capabilities. We conducted a three-fold cross-validation, where 270 training samples were divided into a training and a validation set at a ratio of 8:2. All slides were cut into non-overlapping patches of $256\times 256$ pixels at a $20\times$ magnification.

BReAst Carcinoma Subtyping (BRACS) dataset emerges as a promotion for the automated detection and classification of breast tumors, also constituted by H&E stained histopathological images. The dataset encompasses 547 whole slide images. Each WSI has been meticulously annotated according to consensus by a panel of three board-certified pathologists, categorizing different types of lesions. Specifically, this dataset includes three types of lesions: benign, malignant, and atypical, further divided into seven subtypes. In this paper, we focus on the three-class classification. We adhered to the official division and employed distinct seeds across three experimental trials. All slides were cut into non-overlapping patches of $256\times 256$ pixels at a $5\times$ magnification.

In our experiments, we employ several evaluation metrics in bag-level classification, including accuracy (ACC), the area under the receiver operating characteristic curve (AUC), and the harmonic mean of precision and recall (F1 score). For the multi-class dataset BRACS, AUC is computed using a one-versus-all mode. In addition, we also conducted a patch-level evaluation with another two evaluation metrics: the free-response receiver operating characteristic (FROC) [43] and the competition performance metric (CPM) score [45]. FROC is defined as the plot of sensitivity versus the average number of false positive patches per image (FPI) and CPM is defined as the average sensitivity at seven predefined false positive rates which seven values are 0.125, 0.25, 0.5, 1, 2, 4, and 8 in FROC curve and was calculated as:

All experiments were conducted on one NVIDIA RTX 3090 GPU which possesses 24 GB memory. Our methodology was implemented within the PyTorch v2.0 deep learning framework. In our experiments, we applied the resnet50 [29] backbone as the feature encoder, and Adam optimizer was utilized for training. For the MIL training, the batch size was set to 1, with an initial and minimum learning rate of $1\times 10^{-3}$ and $1\times 10^{-4}$, respectively. For the encoder fine-tuning, the batch size was set to 64, with a fixed learning rate of $5\times 10^{-5}$. The maximum number of epochs for training both the MIL training and the encoder fine-tuning was established at 200. The early stopping strategy was implemented to prevent overfitting, which was monitored through the loss on the validation set.

In general, we evaluate the performance of different features under certain MIL models and the performance of certain features under different MIL models.

As our method aims at enhancing feature representation, we first evaluate features generated by different approaches: (1) Random initialization. (2) Pre-trained on natural images, Imagenet-1K [29]. (3) The weight from employing self-supervised methods on extensive WSI datasets, including Barlow Twins [72], MoCo V2 [17], and SWAV [10]. (4) Task-oriented feature fine-tuning methods, including WSI-FT [35] and BCL [71].

Additionally, as a multi-instance learning approach for WSI classification, we also evaluate other relevant MIL methods: (1) The classic ABMIL [31]. (2) Four variants of ABMIL, including single-attention-branch CLAM_SB [41], multi-attention-branch CLAM_MB [41], non-local attention pooling DSMIL [34] and muti-head attention ACMIL [74]. (3) Transformer-based MIL, TransMIL [54]. (6) Double-tier feature distillation MIL, DTFD [73]. (7) End-to-end feature encoder fine-tuning, WSI-FT [35]. (8) Bayesian collaborative learning-based MIL, BCL [71]. All results of these methods are conducted using their official codes under the same settings.

To validate the efficiency of the embedding generated by our approach, we conducted a comparative analysis of performance results obtained with various feature extraction techniques, as reported in Tab. 1. Initially, features derived from a randomly initialized feature encoder exhibited limited classification ability. Through supervised learning with ImageNet-1K, significant enhancements were observed in the representational capacity of embeddings, especially in the CAMELYON16 dataset. Self-supervision based methods facilitated fine-tuning of the feature encoder using extensive pathological images. Embeddings generated through Barlow Twins, MoCo V2, and SWAV failed to demonstrate superior performance in the CAMELYON16 dataset but exhibited improved performance in the BRACS dataset. This phenomenon could be attributed to the dataset’s inherent characteristics. The utilization of self-supervised weights may inadvertently complicate features, resulting in different MIL classification outcomes according to the category distinctions in different datasets. In task-oriented feature fine-tuning methods, BCL outperformed most of the other methods as it distills efficient patch-level information from MIL models for encoder fine-tuning. Notably, our method consistently outperformed other feature fine-tuning approaches in both datasets. This robust performance underscores the effectiveness and versatility of our approach. Furthermore, our results remained optimal even when applied to different MIL models, highlighting the generalizability of our approach.

To further validate the WSI classification performance of our approach, we evaluated results across various MIL methods, as summarized in Tab. 2. Our analysis reveals that our approach consistently outperforms others on both datasets. Most of the other methods primarily focus on the improvement of the aggregator, yet they fail to achieve comparable performance to ours. This can be attributed to their reliance on features extracted by weights pre-trained on natural images, which is not optimal for WSI classification tasks. WSI-FT, an end-to-end training framework with partial patches, tends to exhibit poor generalization of features. Although BCL offers an excellent foundational framework, the presence of noise in the pseudo labels may potentially lead to model overfitting.

Both as iterative methods, we demonstrated the performance difference between BCL and our approach across different rounds of iteration on the validation and test sets of CAMELYON16 and BRACS datasets. Our observations in Fig. 3 revealed that our method demonstrates a more rapid convergence compared to BCL on the CAMELYON16 dataset, accompanied by a notable improvement in effectiveness. On the BRACS dataset, although BCL’s performance exhibits fluctuations, our method consistently converges swiftly to superior outcomes.

Given the detailed annotations for tumor regions provided by the CAMELYON16 dataset, we evaluated the patch-level classification ability of various MIL methods for comparison. All MIL methods were assessed using their bag-level classification classifiers, while both BCL and our method included an additional patch-level classifier. As presented in Tab. 3, the performance results revealed that most methods exhibited high AUC scores alongside low F1 scores, which primarily stemmed from the inherent characteristics of the MIL principle. These methods tend to focus on a few instances within a bag, as a bag is labeled positive even if only one instance within it is positive. Consequently, this often leads to only a few patches being classified with very high confidence. In contrast, our method consistently achieved optimal results in both ACC and F1 scores. This can be attributed to our approach of selecting hard negative samples to mitigate overfitting. The performance of our bag-level classifier implicitly demonstrated that our method effectively addressed the tendency of MIL to excessively focus on a minority of patches. Furthermore, we provide free-response receiver operating characteristic curves using different MIL models, as depicted in Fig. 4. These curves illustrated that our method outperformed others across all FPIs, indicating the ability to detect tumor slides at a low false positive cost.

To provide an intuitive demonstration of the effectiveness of our method, we visualized representative results in Fig. 5. Specifically, we selected a macro-metastasis (’test$\_$001’) and a micro-metastasis (’test$\_$010’) WSIs from the CAMELYON16 dataset for visualization, generating classification heatmaps using attention scores derived from the MIL aggregator. For ’test$\_$001’, our heatmap accurately identifies the macro-metastatic tumor region in the slide, along with some micro-metastatic tumor regions in the surrounding area (e.g., the lower-left corner). Upon closer analysis of regions with low attention in the heatmap, we observed that patches containing a large number of fat cells (e.g., patch ($b$)) were assigned lower attention. However, it is noteworthy that not all patches containing fat cells were dismissed as normal; for instance, in patches where only a small portion contains fat cells (e.g., patch ($a$)), our model still identifies them as tumors, indicating its ability to distinguish between fat cells and tumor cells. Additionally, regions with low attention also contained some lymph cells, which our model treats as hard negative samples in the dataset (e.g., patch ($c$)), leading to their lower attention assignment. For ’test$\_$010’, our model successfully identifies the small tumor region, in accord with the annotation. Even in regions with micrometastatic tumor cells (e.g., patch ($f$)), our model demonstrates accurate identification. It’s noted that since the patches are $256\times 256$ in size, our heatmap slightly extends beyond the annotated tumor region.

To validate the effectiveness of key components within our method, we conducted ablation experiments on both the CAMELYON16 and BRACS datasets. As is depicted in Tab. 4, scrutinized three crucial components: potential negative sample mining, hard negative sample searching, and positive sample cleaning. For the potential negative sample mining component, compared to using only positive samples, a noticeable improvement was observed on the CAMELYON16 dataset. This can be attributed to the dataset’s simplicity, which facilitated more accurate category center selection, thus enhancing the mining module’s functionality. Conversely, on the BRACS dataset, the initial suboptimal MIL model generates less accurate category centers, thus consequently fewer useful negative samples are found for fine-tuning. In scenarios without hard negative sample searching, relying solely on high-confidence patches as hard negative samples showed mediocre performance on the CAMELYON16 dataset. In contrast, on the BRACS dataset, significant improvements were observed, as some negative samples were misclassified, thereby providing opportunities for enhancement. For experiments focusing solely on positive sample cleaning, both datasets exhibited slight improvements, indicating the presence of noise in the original positive samples. Notably, our method, although divided into three parts, operates synergistically: the negative sample mining module enables the model to identify missing information from low-confidence patches and maximize the utilization of all patches; the hard negative sample searching module aids the model in learning more adversarial information; and the positive sample cleaning module ensures the purity of positive samples, preventing the model from overfitting to noisy samples.

In this study, two critical hyper-parameters, namely the initial number of patches $K_{0}$ and the number of clusters $C$ are used. We conducted extensive experiments to set the best values for these two hyper-parameters. Limit our comparison to the outcomes following the initial iteration, as a selection of hyper-parameters. As depicted in Fig. 6, the curves all exhibit a trend of rising first and then declining. Therefore, we choose the parameters corresponding to the peak values as hyper-parameters. Specifically, we set $K_{0}$ to 10 on both datasets. As for the number of clusters $C$, since it largely depends on the dataset itself, we set it to $5$ for the CAMELYON16 dataset and $10$ for the BRACS dataset respectively.