[Paper Review] TransMIL: Transformer based Correlated Multiple Instance Learning for Whole Slide Image Classification

Zhuchen Shao, Hao Bian+5 more|arXiv (Cornell University)|Jun 2, 2021

AI in cancer detection501 citations

TL;DR

TransMIL introduces a Transformer-based correlated MIL framework for whole slide image classification, leveraging both morphological and spatial information to outperform state-of-the-art MIL methods on multiple datasets with faster convergence and better interpretability.

ABSTRACT

Multiple instance learning (MIL) is a powerful tool to solve the weakly supervised classification in whole slide image (WSI) based pathology diagnosis. However, the current MIL methods are usually based on independent and identical distribution hypothesis, thus neglect the correlation among different instances. To address this problem, we proposed a new framework, called correlated MIL, and provided a proof for convergence. Based on this framework, we devised a Transformer based MIL (TransMIL), which explored both morphological and spatial information. The proposed TransMIL can effectively deal with unbalanced/balanced and binary/multiple classification with great visualization and interpretability. We conducted various experiments for three different computational pathology problems and achieved better performance and faster convergence compared with state-of-the-art methods. The test AUC for the binary tumor classification can be up to 93.09% over CAMELYON16 dataset. And the AUC over the cancer subtypes classification can be up to 96.03% and 98.82% over TCGA-NSCLC dataset and TCGA-RCC dataset, respectively. Implementation is available at: https://github.com/szc19990412/TransMIL.

Motivation & Objective

• Address weakly supervised WSI classification by moving beyond i.i.d. instance assumptions to model correlations among instances.
• Propose a Transformer-based MIL (TransMIL) to capture both morphological and spatial relationships between patches.
• Provide a convergent, interpretable, and efficient three-step algorithm for correlated MIL in pathology.
• Demonstrate superior performance and faster convergence on multiple public histopathology datasets.

Proposed method

• Propose a generic three-step correlated MIL framework with a convergence proof.
• Develop TransMIL with a Transformer-based Temporal Patch Transformer (TPT) module to model long-range correlations among patch embeddings.
• Introduces Pyramid Position Encoding Generator (PPEG) for conditional, multi-scale spatial position encoding.
• Use a Nystr40m-based approximate self-attention in TPT to handle long sequences efficiently.
• Combine morphology (f), spatial (h), and their sum (fh) embeddings, followed by a learned pooling matrix P and a final mapping g to bag labels.
• Provide an end-to-end training setup using cross-entropy loss and Lookahead optimizer with ResNet50-based patch embeddings (1024-d, reduced to 512-d).

Experimental results

Research questions

• RQ1 Can modeling correlations among instances within a WSI improve MIL-based classification compared to i.i.d.-based MIL methods?
• RQ2 How can Transformer-based architectures be adapted to large, variable-length WSI patch sequences while preserving spatial context?
• RQ3 Does conditional, multi-scale position encoding (PPEG) improve diagnostic performance and interpretability in MIL for pathology?
• RQ4 Is TransMIL effective across binary and multi-class WSI classification tasks and across unbalanced datasets?
• RQ5 Do correlated MIL approaches converge faster than existing MIL methods?

Key findings

• TransMIL achieves state-of-the-art performance on three public datasets: CAMELYON16, TCGA-NSCLC, and TCGA-RCC.
• On CAMELYON16, TransMILACCURACY 0.8837 and AUC 0.9309, outperforming i.i.d.-based and some non-local methods.
• On TCGA-NSCLC, TransMIL achieves ACC 0.8835 and AUC 0.9603, higher than competing methods.
• On TCGA-RCC, TransMIL achieves ACC 0.9466 and AUC 0.9882, the best among reported methods.
• Ablation shows PPEG and conditional position encoding significantly improve performance over sinusoidal or no position encoding.
• TransMIL shows faster convergence, requiring 2–3x fewer training epochs than competing MIL methods.
• Visualizations of attention maps demonstrate high overlap with cancer regions, supporting interpretability.