2DMamba让图像模型高效处理超大分辨率图像,适合病理切片等场景。
2DMamba: Efficient State Space Model for Image Representation with Applications on Giga-Pixel Whole Slide Image Classification
- 将Mamba的线性复杂度引入2D图像,保留空间结构并提升计算效率
- 在10个病理图像数据集上,最高提升5.52%生存预测准确率
- 适用于超大图分类,尤其适合病理切片和遥感影像分析
高效建模大尺度二维上下文对数字病理学与遥感等领域至关重要。基于Transformer的模型虽具高并行性,但因序列过长导致二次复杂度问题。近期提出的Mamba通过选择性状态空间模型(SSM)实现线性复杂度与高并行性,有效处理一维长序列。然而,将其直接扩展至具有固有二维结构的视觉任务时,会因一维序列处理限制造成空间信息失真。现有二维SSM虽能建模二维结构,却因缺乏高效并行算法而计算极慢。本文提出2DMamba,一种新型二维选择性SSM框架,将图像的二维空间结构融入Mamba,并设计硬件感知的优化算子,兼顾空间连续性与计算效率。我们在病理全切片图像(WSI)与自然图像上验证了方法的通用性。在10个公开的WSI分类与生存分析数据集上,2DMamba使AUC最高提升2.48%,F1分数提升3.11%,准确率提升2.47%,C-index提升5.52%。此外,与VMamba结合用于自然图像时,在ADE20k语义分割上mIoU提升0.5至0.7,在ImageNet-1K分类上准确率提升0.2%。代码已开源:https://github.com/AtlasAnalyticsLab/2DMamba。
原文摘要 · Abstract (English)
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.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。