Research map / Architectures and efficiency / Attention and transformers
2DMamba: Efficient State Space Model for Image Representation with Applications on Giga-Pixel Whole Slide Image Classification
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Proposes 2DMamba, a 2D selective state space model with a hardware-aware 2D scan operator, applied to WSI MIL classification/survival and to VMamba for natural image tasks.
Problem
Mamba flattens 2D images into 1D sequences, causing spatial discrepancy (distant neighbors in sequence), while existing 2D SSMs lack efficient parallel algorithms and are too slow.
Approach
A 2D selective SSM that performs horizontal then vertical scans over a 2D feature map (so hidden state aggregates over the upper-left region with Manhattan-distance decay), implemented with a tiled, hardware-aware SegmentedBlockScan CUDA kernel that keeps O(L) memory access.
Key results
- On five WSI classification datasets, 2DMambaMIL reaches 0.7517 accuracy and 0.8964 AUC on BRACS, outperforming the best SSM-based MIL by up to 2.48% AUC and 3.11% F1.
- On five TCGA survival datasets, 2DMambaMIL attains the highest C-index on all, e.g. 0.6198 on LUAD (about 5.5% relative gain over the best baseline).
- 2DVMamba-T reaches 82.8% ImageNet-1K top-1 vs 82.6% for VMamba-T, and 48.6 vs 47.9 single-scale mIoU on ADE20K.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Use a learnable non-tissue padding token instead of fixed-zero padding (ablation) | 2DMambaMIL · PANDA | accuracy: 0.4868 -> 0.5075 | Helped |
| Use a learnable non-tissue padding token instead of fixed-zero padding (ablation) | 2DMambaMIL · PANDA | AUC: 0.8057 -> 0.8184 | Helped |
| Use a learnable non-tissue padding token instead of fixed-zero padding (ablation) | 2DMambaMIL · TCGA-BRCA | accuracy: 0.9313 -> 0.9458 | Helped |
| Use a learnable non-tissue padding token instead of fixed-zero padding (ablation) | 2DMambaMIL · TCGA-BRCA | AUC: 0.9722 -> 0.9782 | Helped |
| Replace 1D scanning with intrinsic 2D scanning (ablation) | MambaMIL · PANDA | accuracy: 0.4679 -> 0.5075 | Helped |
| Replace 1D scanning with intrinsic 2D scanning (ablation) | MambaMIL · PANDA | AUC: 0.7781 -> 0.8184 | Helped |
| Replace 1D scanning with intrinsic 2D scanning (ablation) | MambaMIL · TCGA-BRCA | accuracy: 0.9333 -> 0.9458 | Helped |
| Replace 1D scanning with intrinsic 2D scanning (ablation) | MambaMIL · TCGA-BRCA | AUC: 0.9657 -> 0.9782 | Helped |
| Use 2-direction scanning instead of 1D scanning (ablation) | MambaMIL · PANDA | accuracy: 0.4679 -> 0.4853 | Helped |
| Use 2-direction scanning instead of 1D scanning (ablation) | MambaMIL · PANDA | AUC: 0.7781 -> 0.7749 | Hurt |
| Use 4-direction raster scanning instead of 1D scanning (ablation) | MambaMIL · PANDA | accuracy: 0.4679 -> 0.4923 | Helped |
| Use 4-direction cross scanning instead of 1D scanning (ablation) | MambaMIL · PANDA | AUC: 0.7781 -> 0.8006 | Helped |
Surprising findings
- Multi-directional 1D scanning (4-direction cross scan) improves MambaMIL but still underperforms the 2D scan, supporting that flattening loses cell-to-cell structure.
- Naive 2D scan is far slower than the CUB 1D scan, while the proposed kernel roughly doubles its throughput over the naive approach but still lags CUB at large sizes.
Limitations
- Throughput remains below vanilla Mamba's CUB 1D scan at larger feature sizes (e.g. 200x200).
- Natural image gains are small (0.2% ImageNet top-1), and segmentation results are single-run comparisons with modest margins.
Experimental setup
Datasets: TCGA-BRCA BRACS PANDA DHMC TCGA-NSCLC TCGA-KIRC TCGA-KIRP TCGA-LUAD TCGA-STAD TCGA-UCEC ImageNet-1K ADE20K
Benchmarks: BRACS DHMC PANDA TCGA-NSCLC TCGA-BRCA TCGA-KIRC TCGA-KIRP TCGA-LUAD TCGA-STAD TCGA-UCEC ImageNet-1K ADE20K
Metrics: accuracy F1 AUC C-index top-1 accuracy mIoU throughput GPU memory FLOPs
Baselines: AB-MIL CLAM DSMIL DTFD-MIL TransMIL S4-MIL MambaMIL SRMambaMIL VMamba-T Swin-T LocalVMamba-T EfficientVMamba-B
Builds on: Mamba VMamba MambaMIL 2D-SSM CUB BlockScan
Closest papers by meaning
- Spatial-Mamba: Effective Visual State Space Models via Structure-Aware State Fusion (ICLR 2025 · similarity 0.91)
- Multi-Scale VMamba: Hierarchy in Hierarchy Visual State Space Model (NeurIPS 2024 · similarity 0.88)
- M3amba: Memory Mamba is All You Need for Whole Slide Image Classification (CVPR 2025 · similarity 0.88)
- QuadMamba: Learning Quadtree-based Selective Scan for Visual State Space Model (NeurIPS 2024 · similarity 0.88)
- Vision Mamba Mender (NeurIPS 2024 · similarity 0.88)
- EfficientVMamba: Atrous Selective Scan for Light Weight Visual Mamba (AAAI 2025 · similarity 0.88)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.