[Paper Review] Dual-stream Multiple Instance Learning Network for Whole Slide Image Classification with Self-supervised Contrastive Learning

Bin Li, Yin Li+1 more|arXiv (Cornell University)|Nov 17, 2020

Image Retrieval and Classification Techniques53 references42 citations

TL;DR

The paper introduces DSMIL, a dual-stream MIL framework with a trainable distance-based aggregation and self-supervised contrastive learning to classify whole slide images and localize tumors without patch-level labels, leveraging multiscale features for improved accuracy.

ABSTRACT

We address the challenging problem of whole slide image (WSI) classification. WSIs have very high resolutions and usually lack localized annotations. WSI classification can be cast as a multiple instance learning (MIL) problem when only slide-level labels are available. We propose a MIL-based method for WSI classification and tumor detection that does not require localized annotations. Our method has three major components. First, we introduce a novel MIL aggregator that models the relations of the instances in a dual-stream architecture with trainable distance measurement. Second, since WSIs can produce large or unbalanced bags that hinder the training of MIL models, we propose to use self-supervised contrastive learning to extract good representations for MIL and alleviate the issue of prohibitive memory cost for large bags. Third, we adopt a pyramidal fusion mechanism for multiscale WSI features, and further improve the accuracy of classification and localization. Our model is evaluated on two representative WSI datasets. The classification accuracy of our model compares favorably to fully-supervised methods, with less than 2% accuracy gap across datasets. Our results also outperform all previous MIL-based methods. Additional benchmark results on standard MIL datasets further demonstrate the superior performance of our MIL aggregator on general MIL problems. GitHub repository: https://github.com/binli123/dsmil-wsi

Motivation & Objective

• Address weakly supervised whole slide image (WSI) classification using only slide-level labels.
• Mitigate MIL aggregation weaknesses in unbalanced bags and large bag sizes.
• Learn robust patch representations via self-supervised contrastive learning (SimCLR) for MIL.
• Develop a multiscale pyramidal fusion to exploit multiscale tissue features.
• Demonstrate state-of-the-art performance on WSI datasets and validate on general MIL benchmarks.

Proposed method

• Propose DSMIL with a dual-stream MIL aggregator: a max-pooling stream to identify the critical instance and a distance-based attention stream to compute a bag embedding.
• Use a trainable distance U(h_i, h_m) between instance embeddings and the critical instance to form attention weights.
• Compute bag embedding as a weighted sum of instance information vectors v_i using U(h_i, h_m) as weights, then score via a bag classifier; final score is the average of the two streams.
• Train the feature extractor f using self-supervised contrastive learning (SimCLR) to obtain robust patch representations for MIL.
• Employ a pyramidal multiscale feature fusion where features from 5x and 20x magnifications are concatenated to form multiscale embeddings, enabling locally-constrained attention.
• Leverage a self-supervised learning framework to alleviate memory and label scarcity in large bags and improve representation quality.

Experimental results

Research questions

• RQ1 Can a dual-stream MIL aggregator with a trainable distance measurement improve instance and bag-level discrimination in weakly labeled WSIs?
• RQ2 Does self-supervised contrastive learning yield superior patch representations for MIL in WSIs compared to end-to-end training or ImageNet pretraining?
• RQ3 Does multiscale (pyramidal) feature fusion improve WSI classification and localization performance over single-scale approaches?
• RQ4 How does DSMIL perform on standard MIL benchmarks beyond WSIs, and does its aggregation generalize across tasks?

Key findings

• DSMIL outperforms other recent MIL models on WSI classification and tumor localization, with around 2-3% higher accuracy on Camelyon16 and TCGA lung cancer datasets.
• With multiscale attention (DSMIL-LC), DSMIL narrows the accuracy gap to fully supervised methods to under 2% on Camelyon16 for classification.
• Self-supervised contrastive learning for patch representations yields substantial gains over end-to-end training, especially in unbalanced bag scenarios (e.g., Camelyon16).
• Multiscale attention outperforms single-scale and other multiscale fusion strategies, with DSMIL-LC achieving state-of-the-art results on Camelyon16 and TCGA datasets.
• On classical MIL datasets, DSMIL achieves competitive or superior results compared to recent MIL aggregators, demonstrating generalization of the dual-stream approach.
• Ablation studies confirm the effectiveness of the DSMIL aggregator, contrastive learning, and multiscale attention in improving MIL performance.