Source-linked AI summary
TransMIL: Transformer based Correlated Multiple Instance Learning for Whole Slide Image Classification
Zhuchen Shao, Hao Bian, Yang Chen, Yifeng Wang, Jian Zhang, Xiangyang Ji, Yongbing Zhang
TL;DR
WSI MIL methods commonly assume independently distributed instances, overlooking correlations that matter for pathology diagnosis. The paper introduces correlated MIL with a convergence proof and develops TransMIL to model morphological and spatial information. Across three computational pathology problems, TransMIL achieves better performance and faster convergence than state-of-the-art methods while supporting balanced or unbalanced binary and multiple classification.
Problem
Existing WSI MIL methods generally assume independent and identically distributed instances, neglecting correlations among regions used in pathology diagnosis.
Method
The paper proposes correlated MIL with a convergence proof and develops TransMIL using Transformer self-attention to model morphological and spatial information.
Results
TransMIL achieves better performance and faster convergence than state-of-the-art methods across three computational pathology problems, with best results across the reported classification settings.
Takeaways & Limitations
TransMIL supports balanced or unbalanced binary and multiple classification with visualization and interpretability.
Takeaways & Limitations
The correlated-MIL theory treats bag instances as random variables under either correlation or i.i.d. assumptions.
Abstract
from arXiv · showhide
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.
1 Introduction
WSI diagnosis is difficult because slides are enormous and lack pixel-level annotations, so MIL provides weakly supervised analysis. Existing MIL methods assume independent instances, whereas pathology decisions also use contextual and cross-region correlations; the paper proposes correlated MIL and TransMIL to address this gap.
- 1 Introduction: WSI diagnosis faces huge image sizes and a lack of pixel-level annotations, motivating weakly supervised MIL.
- 1 Introduction: Existing MIL methods assume instances are independent and identically distributed, although pathologists use contextual and inter-region correlations.
- 1 Introduction: Transformer self-attention models pairwise correlations between sequence tokens, unlike bypass attention mechanisms based on independent instances.
- 1 Introduction: Traditional Transformer sequences are computationally limited and generally handle only shorter sequences.
- 1 Introduction: The paper introduces a correlated MIL framework with a convergence proof and generic three-step algorithm.
- 1 Introduction: TransMIL uses a Transformer-based MIL design to explore morphological and spatial information between instances.
2 Related Work
Prior WSI MIL work includes instance-level and embedding-level approaches, while attention-based methods assign importance but generally do not model correlations between instances. TransMIL applies self-attention to comprehensively consider within-bag instance correlations for WSI classification.
- 2 Related Work: Instance-level MIL assigns pseudo-labels from bag labels and aggregates top-k instances, but requires many WSIs because few instances contribute to training.
- 2 Related Work: Embedding-level MIL maps every slide patch to a fixed-length embedding before aggregating feature embeddings.
- 2 Related Work: Attention methods weight feature channels, spatial distributions, time series, or MIL instances, but do not consider correlations between different instances.
- 2 Related Work: TransMIL is presented as the first Transformer-based WSI classification method to comprehensively model correlations among instances within a bag.
3 Method
The method formulates correlated MIL for bags whose instances may be dependent and ordered, then implements it with TransMIL to model morphological, spatial, and long-range correlations. Its TPT module combines Transformer aggregation, pyramid position encoding, and linear-complexity approximate self-attention for long WSI sequences.
- Correlated Multiple Instance Learning: Correlated MIL represents bags with dependent and ordered instances rather than assuming independent and identically distributed instances.
- Correlated Multiple Instance Learning: Theorem 1 and its inference provide an arbitrary approximation form for the scoring function S(X), while Theorem 2 gives the advantage of correlated MIL.
- Correlated Multiple Instance Learning: The generic algorithm extracts morphological and spatial information, aggregates it with a Pooling Matrix P, and transforms the result into a bag-level prediction.
- TransMIL: TransMIL maps each bag through Transformer space to label space, using Transformer layers for morphological aggregation and PPEG for spatial information encoding.
- TransMIL: Nystrom-approximated self-attention reduces computational complexity from O(n^2) to O(n), enabling TPT to process bags containing thousands of tokens.
- TransMIL: PPEG reshapes patch tokens into 2-D image space, applies different-sized convolution kernels, fuses spatial information, and reconnects the flattened tokens with the class token.
4 Experiments and Results
Experiments evaluate TransMIL on three WSI classification datasets, including binary and multiclass settings, and assess PPEG, positional encoding, interpretability, and convergence. TransMIL performs strongly across these evaluations while using spatial and morphological information.
- Datasets and evaluation: Experiments cover CAMELYON16, TCGA-NSCLC, and TCGA-RCC across binary and multiclass WSI classification tasks.The evaluation includes metastasis detection, lung cancer subtype classification, and three renal cell carcinoma subtypes.
- WSI classification results: TransMIL outperformed competing methods on CAMELYON16, with AUC at least 5% higher than ABMIL, PT-MTA, and CLAM.The dataset contains many negative regions because tumor occupies less than 10% of each positive slide on average.
- WSI classification results: TransMIL achieved 1.40% higher AUC and 2.16% higher accuracy than the second-best method on TCGA-NSCLC.Positive slides contain more than 80% tumor area on average in this dataset.
- WSI classification results: TransMIL achieved the best accuracy and AUC on the unbalanced, multiclass TCGA-RCC dataset.MIL-RNN was excluded because it did not consider multiclass classification.
- Effects of PPEG: PPEG-based conditional position encoding improved classification over no positional encoding and was more effective than sinusoidal encoding.Using different convolution kernel sizes in one layer adds multilevel positional encoding and more context to each token.
- Effects of conditional position encoding: Conditional position information improved AUC by up to 0.9% on CAMELYON16 and 0.61% on TCGA-NSCLC when input order was disrupted.The comparison evaluates sequential versus randomly disordered input sequences.
- Interpretability and attention visualization: TransMIL produced interpretable attention heatmaps whose highlighted regions were highly consistent with fine cancer annotations.The heatmaps identify regions of interest and important morphology used for diagnosis.
- Convergence: TransMIL required approximately two to three times fewer training epochs than traditional MIL, ABMIL, DSMIL, and CLAM.The paper reports better convergence and validation AUC than other MIL methods.
5 Conclusion
The paper presents correlated MIL and TransMIL for weakly supervised WSI classification, combining contextual, morphological, and spatial information. Across three public datasets, TransMIL supports varied classification settings and outperforms state-of-the-art MIL algorithms in AUC and accuracy, while higher magnification remains computationally challenging.
- TransMIL explores morphological and spatial information within a correlated MIL framework for weakly supervised WSI classification.The framework is designed to reflect contextual and inter-area correlations considered in pathology diagnosis.
- TransMIL applies to unbalanced and balanced datasets as well as binary and multiple classification tasks, with visualization and interpretability.The paper reports applicability across these settings and describes the network as easy to train.
- TransMIL outperforms state-of-the-art MIL algorithms in both AUC and accuracy across three public datasets.
- Higher-magnification WSIs produce longer sequences that increase computational and memory requirements.The paper identifies this as a challenge for future work after experiments at ×20 magnification.
A Appendix A
Appendix A develops the theoretical basis for correlated MIL by relating continuous set-function approximation to entropy under correlated instances. It concludes that correlation reduces uncertainty relative to the i.i.d. assumption.
- A Hausdorff-continuous set function S(X) can be arbitrarily approximated by a function constructed through an invertible map and auxiliary functions.
- The proof represents instances as jointly distributed random variables and decomposes joint entropy into conditional entropy terms.
- Under the correlation assumption, the information source has smaller entropy than under the i.i.d. assumption, reducing uncertainty and providing more useful information.
B Appendix B
Appendix B defines Transformer-based MIL as a mapping from bags to Transformer representations and then labels. Its design uses self-attention for morphological aggregation, positional encoding for spatial information, and class-token or global-pooling classification.
- Transformer-based MIL learns the mapping X → T → Y from bag space to Transformer space and then to label space.Each bag contains multiple instances and a corresponding label.
- The Transformer input prepends a class token to instance embeddings and adds positional encoding.
- Multi-head self-attention aggregates information across tokens through concatenated attention heads and an output projection.Layer normalization is applied before each multi-head self-attention block.
- The label mapping can use the class token or global average pooling, but WSI application must handle hundreds to thousands of instances per bag.