Source-linked AI summary
MambaMIL: Enhancing Long Sequence Modeling with Sequence Reordering in Computational Pathology
Shu Yang, Yihui Wang, Hao Chen
TL;DR
Existing MIL methods face limited contextual interaction, computational overhead, and overfitting challenges when modeling long sequences of WSI instances. MambaMIL integrates Mamba into MIL and adds Sequence Reordering Mamba to exploit multiple instance orderings. Across two tasks and nine datasets, it outperforms existing competitors, including by 2.6% and 2.7% in mean survival-prediction performance across seven datasets under two feature settings.
Problem
Existing MIL methods struggle to comprehensively model contextual relationships in long WSI instance sequences while avoiding computational overhead and overfitting.
Method
MambaMIL integrates Mamba into MIL, using SR-Mamba to model long sequences through distinct sequential orderings and compressed hidden states.
Results
MambaMIL outperforms the second-best method by 2.6% and 2.7% in mean performance across seven survival-prediction datasets under two feature settings, and outperforms competitors across nine benchmarks.
Takeaways & Limitations
Sequence reordering helps MambaMIL capture more discriminative features and supports strong performance across diverse computational-pathology benchmarks.
Abstract
from arXiv · showhide
Multiple Instance Learning (MIL) has emerged as a dominant paradigm to extract discriminative feature representations within Whole Slide Images (WSIs) in computational pathology. Despite driving notable progress, existing MIL approaches suffer from limitations in facilitating comprehensive and efficient interactions among instances, as well as challenges related to time-consuming computations and overfitting. In this paper, we incorporate the Selective Scan Space State Sequential Model (Mamba) in Multiple Instance Learning (MIL) for long sequence modeling with linear complexity, termed as MambaMIL. By inheriting the capability of vanilla Mamba, MambaMIL demonstrates the ability to comprehensively understand and perceive long sequences of instances. Furthermore, we propose the Sequence Reordering Mamba (SR-Mamba) aware of the order and distribution of instances, which exploits the inherent valuable information embedded within the long sequences. With the SR-Mamba as the core component, MambaMIL can effectively capture more discriminative features and mitigate the challenges associated with overfitting and high computational overhead. Extensive experiments on two public challenging tasks across nine diverse datasets demonstrate that our proposed framework performs favorably against state-of-the-art MIL methods. The code is released at https://github.com/isyangshu/MambaMIL.
1 Introduction
WSI analysis is challenging because of high resolution and limited pixel-level annotations, motivating MIL’s bag-of-instances formulation. Existing MIL methods struggle with contextual interactions and long-sequence modeling, leading to MambaMIL and its sequence-reordering extension.
- High-resolution WSIs and limited pixel-level annotations make deep-learning-based WSI analysis challenging.
- MIL represents each WSI as a bag of tissue-patch instances and classifies the bag according to whether at least one instance is positive.
- Attention-based MIL methods primarily model instance-level information under independent and identically distributed assumptions, neglecting contextual relationships among instances.
- Mamba provides linear-complexity long-sequence modeling with global receptive fields, but direct application to flattened visual patches restricts interactions to previously scanned positions.
- MambaMIL incorporates Mamba into MIL, while SR-Mamba models sequential and transpositional orderings to capture long-range dependencies among scattered positive instances.
- Experiments across two challenging tasks and nine datasets report superior performance against state-of-the-art MIL methods.
2 Method
MambaMIL integrates Mamba into MIL to model long WSI instance sequences efficiently, while SR-Mamba uses distinct sequence orderings to capture richer instance dependencies. Its pipeline extracts and projects patch features, processes them with stacked SR-Mamba modules, and aggregates the resulting representations.
- 2.1 Preliminaries: S4 models represent long sequences through latent state-space dynamics and support recurrent or convolutional computation for efficient inference or parallelizable training.The continuous parameters are discretized before recurrent or convolutional processing.
- 2.1 Preliminaries: Mamba adds input-dependent selection and hardware-aware parallel computation, selectively propagating or forgetting information along a sequence with linear-complexity modeling.Selection depends on the current token while retaining a global receptive field.
- 2.2 Overview of MambaMIL: MambaMIL maps WSI patches to instance features, reduces their dimension, processes them through stacked SR-Mamba modules, and aggregates the outputs into bag-level representations.The pipeline consists of Feature Extractor, Linear Projection, stacked SR-Mamba modules, and Aggregation.
- 2.3 Sequence Reordering Mamba: SR-Mamba uses parallel branches for original and reordered instance sequences, assigning each ordering a compressed hidden state to model scattered positive patches and long-range dependencies.The original ordering is processed directly, while the reordered branch exploits position-sensitive sequence modeling.
- 2.3 Sequence Reordering Mamba: Sequence Reordering partitions instances into segments, reshapes them into a 2-D feature map, samples across segments, and restores the resulting features to the original ordering.Sequences not divisible by the segment size are zero-padded before reordering.
- 2.3 Sequence Reordering Mamba: SR-Mamba combines the two modeled feature streams by element-wise addition, and stacked SR-Mamba modules provide long-range modeling with linear complexity.The design preserves sequential ordering while generating a globally reordered representation for feature re-embedding.
3 Experiments
Experiments evaluate MambaMIL across two downstream tasks and nine datasets, comparing it with pooling, attention-, Transformer-, and state-space-based MIL methods. Results show strong benchmark performance, improved cancer-subtyping AUC, and more stable validation behavior than TransMIL.
- Experimental Setup: Experiments cover survival prediction and cancer subtyping across nine public challenging datasets using two distinct feature settings.The feature sets come from ImageNet-pretrained ResNet-50 and PLIP pretrained on pathology image-text pairs.
- Survival Prediction: MambaMIL achieves the best performance on all seven survival-prediction benchmarks, exceeding the second-best method by 2.6% and 2.7% in mean performance across feature settings.These comparisons use two feature settings across seven TCGA cancer datasets.
- Cancer Subtyping: MambaMIL attains AUCs of 80.4% on BRACS and 95.9% on NSCLC for cancer subtyping.Relative to ABMIL, its AUC improves by 3.9% on BRACS and 2.1% on NSCLC.
- Ablation Study: SR-Mamba surpasses vanilla Mamba and Bi-Mamba on survival-prediction datasets under the same training settings.The comparison isolates the effect of sequence reordering among Mamba-block variants.
- Ablation Study: During training on the BRCAS validation set, TransMIL shows increasing validation loss with decreasing ACC and AUC, whereas MambaMIL remains stable.The comparison is reported throughout the training process.
4 Conclusion
The paper introduces MambaMIL, a Mamba-based MIL method centered on SR-Mamba for long-sequence modeling in computational pathology. Across nine benchmarks, it reports performance advantages under all evaluated metrics and identifies possible extensions to other computational-pathology modalities.
- Conclusion: MambaMIL applies the Mamba framework to MIL to address long-sequence modeling and overfitting in computational pathology.The paper presents this as the first application of Mamba in computational pathology.
- Conclusion: SR-Mamba enables MambaMIL to leverage intrinsic global information contained in long sequences of WSI instances.SR-Mamba is the specially designed core module of the approach.
- Conclusion: Across nine benchmarks, MambaMIL outperforms existing competitors under all reported metrics.The conclusion attributes the result to benefits from long-sequence modeling.
- Conclusion: The paper anticipates extending the approach to genomics, pathology reports, and clinical data for multimodal computational-pathology tasks.The stated target applications include diagnosis, prognosis, and therapeutic-response prediction.
Appendix
The appendix specifies the SR-Mamba block’s tensor inputs, reordered-sequence processing, state-space operations, gating, residual connection, and output, alongside hyper-parameter configurations.
- Algorithm 1 SR-Mamba Block Process: SR-Mamba receives an instance sequence Xl−1 with shape (B, M, D) and returns Xl with the same shape.B denotes batch size, M instance number, and D dimension.
- Algorithm 1 SR-Mamba Block Process: The block creates a reordered sequence from a linear projection of the input and processes the original and reordered sequences in separate branches.The algorithm iterates over the two orderings, os and rs.
- Algorithm 1 SR-Mamba Block Process: Each branch applies SiLU and convolutional processing, derives state-space parameters through linear projections, and produces an SSM output.The listed operations include LinearA, LinearB, LinearC, parameterized state construction, and SSM processing.
- Algorithm 1 SR-Mamba Block Process: The branch outputs are gated with SiLU, combined, projected linearly, and added to the input through a residual connection.The resulting tensor is returned as Xl.
- Hyper-parameter Configurations: The appendix includes a table of hyper-parameter configurations for the reported experiments.The supplied appendix passage identifies the configuration table but does not enumerate its values.