[Paper Review] Attention-Challenging Multiple Instance Learning for Whole Slide Image Classification

Yunlong Zhang, Honglin Li+4 more|arXiv (Cornell University)|Nov 13, 2023

Advanced Image and Video Retrieval Techniques4 citations

TL;DR

This paper proposes Attention-Challenging Multiple Instance Learning (ACMIL), a novel MIL framework for whole slide image (WSI) classification that combats overfitting by forcing attention mechanisms to focus on more challenging instances. ACMIL integrates Multiple Branch Attention (MBA) to capture diverse discriminative patterns and Stochastic Top-K Instance Masking (STKIM) to suppress salient instances, achieving state-of-the-art performance across three WSI datasets with improved generalization and robustness.

ABSTRACT

In the application of Multiple Instance Learning (MIL) methods for Whole Slide Image (WSI) classification, attention mechanisms often focus on a subset of discriminative instances, which are closely linked to overfitting. To mitigate overfitting, we present Attention-Challenging MIL (ACMIL). ACMIL combines two techniques based on separate analyses for attention value concentration. Firstly, UMAP of instance features reveals various patterns among discriminative instances, with existing attention mechanisms capturing only some of them. To remedy this, we introduce Multiple Branch Attention (MBA) to capture more discriminative instances using multiple attention branches. Secondly, the examination of the cumulative value of Top-K attention scores indicates that a tiny number of instances dominate the majority of attention. In response, we present Stochastic Top-K Instance Masking (STKIM), which masks out a portion of instances with Top-K attention values and allocates their attention values to the remaining instances. The extensive experimental results on three WSI datasets with two pre-trained backbones reveal that our ACMIL outperforms state-of-the-art methods. Additionally, through heatmap visualization and UMAP visualization, this paper extensively illustrates ACMIL's effectiveness in suppressing attention value concentration and overcoming the overfitting challenge. The source code is available at \url{https://github.com/dazhangyu123/ACMIL}.

Motivation & Objective

• Address the persistent challenge of overfitting in Multiple Instance Learning (MIL) for Whole Slide Image (WSI) classification due to limited data, high resolution, and class imbalance.
• Investigate the root cause of overfitting in existing MIL models, particularly the over-concentration of attention on a few discriminative instances as revealed by heatmap visualization.
• Develop a method that encourages attention mechanisms to focus on more diverse and challenging instances rather than relying on a small subset of salient patches.
• Improve model generalization and robustness in WSI analysis by explicitly promoting attention diversity and suppressing over-reliance on top-attention instances.

Proposed method

• Introduce Multiple Branch Attention (MBA), which employs multiple parallel attention branches to capture discriminative instances with distinct patterns, enhancing feature diversity.
• Propose Stochastic Top-K Instance Masking (STKIM), a data augmentation-like technique that randomly masks the top-K attention-weighted instances and redistributes their attention scores to remaining instances to prevent over-concentration.
• Combine MBA and STKIM within a unified ACMIL framework to force attention mechanisms to attend to more challenging and less obvious instances, improving generalization.
• Incorporate a diversity loss $\mathcal{L}_d$ in MBA to explicitly encourage each attention branch to learn distinct discriminative knowledge, preventing redundancy.
• Utilize UMAP visualization and heatmap analysis to validate the model's improved attention distribution and feature clustering behavior.
• Train and evaluate ACMIL on three WSI datasets (Camelyon16, BRACS, LBC) using two backbones (ResNet18 and ViT-S/16), with both ImageNet and SSL pre-training.

Experimental results

Research questions

• RQ1 To what extent does attention over-concentration on a few discriminative instances contribute to overfitting in existing MIL models for WSI classification?
• RQ2 Can forcing attention to focus on less salient, more challenging instances improve model generalization and performance on WSI datasets?
• RQ3 How effective is Stochastic Top-K Instance Masking (STKIM) in disrupting attention dominance and promoting attention diversity?
• RQ4 Does Multiple Branch Attention (MBA) with a diversity loss $\mathcal{L}_d$ lead to more robust and generalized feature learning compared to single-branch attention?
• RQ5 Can heatmap and UMAP visualizations confirm that ACMIL learns more discriminative and better-separated features than prior SOTA methods?

Key findings

• ACMIL significantly outperforms state-of-the-art MIL methods on three WSI datasets, achieving a 5.9% improvement in F1-score on the LBC dataset with ResNet18 and ImageNet pre-training.
• The inclusion of the diversity loss $\mathcal{L}_d$ in MBA leads to a 5.3–8.0% performance gain across datasets, confirming its critical role in promoting attention diversity.
• STKIM at test time slightly degrades performance, indicating it is not necessary during inference and functions best as a training-time regularizer, similar to cutout.
• UMAP visualization shows ACMIL learns more discriminative and better-separated features, with a V-measure of 0.316 on the LBC dataset, compared to 0.224 for ABMIL.
• Heatmap visualization confirms that ACMIL reduces attention concentration on a few instances, distributing attention more broadly across patches, which correlates with reduced overfitting.
• ACMIL achieves state-of-the-art performance across all metrics (F1-score, AUC) on Camelyon16, BRACS, and LBC, with consistent improvements over ABMIL and other SOTA baselines.