2DMamba: Efficient State Space Model for Image Representation with Applications on Giga-Pixel Whole Slide Image Classification

TL;DR

2DMamba addresses the challenge of modeling large 2D contexts in gigapixel Whole Slide Images by introducing a 2D selective State Space Model that preserves spatial continuity while enabling parallel computation through a hardware-aware 2D selective scan. The approach integrates a MIL-friendly 2D-Mamba block (2DMambaMIL) with a learnable padding token for non-tissue regions and a 2D hardware-optimized scanning operator, achieving linear memory access and substantial speed/memory benefits. Empirically, 2DMambaMIL outperforms state-of-the-art MIL and 1D Mamba on WSI classification and survival across 10 datasets, while 2DVMamba improves ADE20K and ImageNet-1K results; the hardware-aware scan also delivers strong natural-image performance and reduced GPU memory usage. The work demonstrates a scalable, spatially aware 2D long-context modeling paradigm with practical hardware efficiency, and code is released for broader adoption.

Abstract

Efficiently modeling large 2D contexts is essential for various fields including Giga-Pixel Whole Slide Imaging (WSI) and remote sensing. Transformer-based models offer high parallelism but face challenges due to their quadratic complexity for handling long sequences. Recently, Mamba introduced a selective State Space Model (SSM) with linear complexity and high parallelism, enabling effective and efficient modeling of wide context in 1D sequences. However, extending Mamba to vision tasks, which inherently involve 2D structures, results in spatial discrepancies due to the limitations of 1D sequence processing. On the other hand, current 2D SSMs inherently model 2D structures but they suffer from prohibitively slow computation due to the lack of efficient parallel algorithms. In this work, we propose 2DMamba, a novel 2D selective SSM framework that incorporates the 2D spatial structure of images into Mamba, with a highly optimized hardware-aware operator, adopting both spatial continuity and computational efficiency. We validate the versatility of our approach on both WSIs and natural images. Extensive experiments on 10 public datasets for WSI classification and survival analysis show that 2DMamba improves up to 2.48% in AUC, 3.11% in F1 score, 2.47% in accuracy and 5.52% in C-index. Additionally, integrating our method with VMamba for natural imaging yields 0.5 to 0.7 improvements in mIoU on the ADE20k semantic segmentation dataset, and 0.2% accuracy improvement on ImageNet-1K classification dataset. Our code is available at https://github.com/AtlasAnalyticsLab/2DMamba.

Figures (12)

• Figure 1: Left: Conventional MIL Bagging of patches adopts no spatial context. Middle: 1D Mamba-based methods flatten a WSI into a 1D sequence and lose the 2D structure. The adjacent blue and orange patches are far away in the sequence. We call this "spatial discrepancy". Right: 2DMamba processes a WSI in a 2D manner, preserving 2D structure. The blue and orange patches maintains adjacent in the sequence. We call it "spatial continuity".
• Figure 2: Left: The overall architecture of 2DMambaMIL for WSI representation. An input WSI is first tiled into patches and these patches are embedded by a feature extractor into a 2D features map. Non-tissue regions are padded with the learnable token to maintain the 2D spatial relationships. The 2D feature map is then fed to $U$ layers of 2D-Mamba blocks, where the key difference, compared with vanilla Mamba block, is our 2D selective scan module. Right: Our 2D selective scan algorithm. It performs parallel horizontal scan and parallel vertical scan for each state dimension $d$ independently. Parameter $C$ then aggregates $N$ state dimensions into a single dimension output $y$.
• Figure 3: Our hardware-aware 2D selective scan operator with efficient caching mechanism and high parallelism. Orange color represents operations on SRAM and green color represents those on HBM. (a) The 1D Mamba scan operator intakes a flattened sequence on HBM. It tiles the input into sub-sequences. Each sub-sequence is loaded from HBM to SRAM, scanned and reduced across $N$ intermediate dimensions, and then written back to HBM. The total memory access complexity is $\mathcal{O}(L)$. (b) A Naive 2D scan operator tiles the 2D feature map by rows and columns, and performs 1D Mamba scans on each row, column, and on $N$ independent state dimensions. This will explicitly instantiate $N$ intermediate feature maps on HBM, resulting in a memory access complexity of $\mathcal{O}(NL)$. (c) Our 2D scan operator tiles the feature map into 2D grids and scans each grid in 2 directions. Intermediate features are reduced inside each tile, only the aggregated result is stored back to HBM. The memory complexity is $\mathcal{O}(L)$. (d) GPU memory hierarchy: SRAMs are small but fast; HBMs have large capacities but are slow. (e) NVIDIA's CUB BlockScan only supports 1D sequences, with sizes of multiples of $32$. Scanning a two-row grid requires two sequential kernel launches and padding elements. (f) Our SegmentedBlockScan enables scanning multiple rows and columns in parallel. It reduces the amount of memory transactions and padding data.
• Figure 4: The attention visualization of 2DMambaMIL and four other methods on a TCGA-KIRC sample for survival analysis. Tumor regions are outlined in green. AB-MIL and SRMambaMIL primarily focus on non-tumor areas, while CLAM also shows substantial attention to non-tumor regions. In contrast, both 2DMambaMIL and MambaMIL focus predominantly on tumor regions. Compared with MambaMIL, attention of 2DMambaMIL shows a more heterogeneous fashion, focusing more on critical regions related to survival (red arrows) while paying less attention to less related ones (violet arrows).
• Figure 5: Comparison of Effective Receptive Fields (ERF) between Swin-T, VMamba-T and 2DVMamba-T. Pixels with higher intensity indicate larger responses regarding the central pixel.