GraphMamba: Whole slide image classification meets graph-driven selective state space model

Abstract

Multi-instance learning (MIL) has demonstrated promising performance in whole slide image (WSI) analysis. However, existing transformer-based methods struggle back and forth between global representation capability and quadratic complexity, particularly when handling millions of instances. Recently, the selective state space model (Mamba) has emerged as a promising alternative for modeling long-range dependencies with linear complexity. Nonetheless, WSI remains challenging for Mamba due to its inability to capture complex local tissue and structural patterns, which is crucial for accurate tumor region recognition. To this end, we approach WSI classification from a graph-based perspective and present GraphMamba, a novel method that constructs multi-level graphs across instances. GraphMamba involves two key components: intra-group graph mamba (IGM) to grasp instance-level dependencies, and cross-group graph mamba (CGM) for exploring group-level relationships. In particular, before aggregating group features into a comprehensive bag representation, CGM utilizes a cross-group feature sampling scheme to extract the most informative features across groups, enabling compact and discriminative representations. Extensive experiments on four datasets demonstrate that GraphMamba outperforms state-of-the-art ACMIL method by 0.5%, 3.1%, 2.6%, and 3.0% in accuracy on the TCGA BRCA, TCGA Lung, TCGA ESCA, and BRACS datasets. The source code will be available at https://github.com/titizheng/GraphMamba.