MambaMIM: Pre-training Mamba with State Space Token Interpolation and its Application to Medical Image Segmentation

• URL: http://arxiv.org/abs/2408.08070v2
• Date: Fri, 18 Apr 2025 13:21:39 GMT
• Title: MambaMIM: Pre-training Mamba with State Space Token Interpolation and its Application to Medical Image Segmentation
• Authors: Fenghe Tang, Bingkun Nian, Yingtai Li, Zihang Jiang, Jie Yang, Wei Liu, S. Kevin Zhou,
• Abstract summary: We propose a general-purpose pre-training framework called MambaMIM.<n>MambaMIM learns causal relationships of state space within a masked sequence.<n>We pre-train MambaMIM on a large-scale dataset of 6.8K CT scans.
• Score: 23.67774523461722
• License: http://arxiv.org/licenses/nonexclusive-distrib/1.0/
• Abstract: Recently, the state space model Mamba has demonstrated efficient long-sequence modeling capabilities, particularly for addressing long-sequence visual tasks in 3D medical imaging. However, existing generative self-supervised learning methods have not yet fully unleashed Mamba's potential for handling long-range dependencies because they overlook the inherent causal properties of state space sequences in masked modeling. To address this challenge, we propose a general-purpose pre-training framework called MambaMIM, a masked image modeling method based on a novel TOKen-Interpolation strategy (TOKI) for the selective structure state space sequence, which learns causal relationships of state space within the masked sequence. Further, MambaMIM introduces a bottom-up 3D hybrid masking strategy to maintain a masking consistency across different architectures and can be used on any single or hybrid Mamba architecture to enhance its multi-scale and long-range representation capability. We pre-train MambaMIM on a large-scale dataset of 6.8K CT scans and evaluate its performance across eight public medical segmentation benchmarks. Extensive downstream experiments reveal the feasibility and advancement of using Mamba for medical image pre-training. In particular, when we apply the MambaMIM to a customized architecture that hybridizes MedNeXt and Vision Mamba, we consistently obtain the state-of-the-art segmentation performance. The code is available at: https://github.com/FengheTan9/MambaMIM.

Related papers

• Mamba-Sea: A Mamba-based Framework with Global-to-Local Sequence Augmentation for Generalizable Medical Image Segmentation [40.802307155824394]
  We propose a novel Mamba-based framework, Mamba-Sea, incorporating global-to-local sequence augmentation to improve the model's generalizability under domain shift issues. Our proposed method is the first to surpass a Dice coefficient of 90% on the Prostate dataset, which exceeds previous SOTA of 88.61%.
  arXiv  Detail & Related papers  (2025-04-24T12:57:25Z)
• DefMamba: Deformable Visual State Space Model [65.50381013020248]
  We propose a novel visual foundation model called DefMamba. By combining a deformable scanning(DS) strategy, this model significantly improves its ability to learn image structures and detects changes in object details. Numerous experiments have shown that DefMamba achieves state-of-the-art performance in various visual tasks.
  arXiv  Detail & Related papers  (2025-04-08T08:22:54Z)
• TransMamba: Fast Universal Architecture Adaption from Transformers to Mamba [88.31117598044725]
  We explore cross-architecture training to transfer the ready knowledge in existing Transformer models to alternative architecture Mamba, termed TransMamba. Our approach employs a two-stage strategy to expedite training new Mamba models, ensuring effectiveness in across uni-modal and cross-modal tasks. For cross-modal learning, we propose a cross-Mamba module that integrates language awareness into Mamba's visual features, enhancing the cross-modal interaction capabilities of Mamba architecture.
  arXiv  Detail & Related papers  (2025-02-21T01:22:01Z)
• Mamba-SEUNet: Mamba UNet for Monaural Speech Enhancement [54.427965535613886]
  Mamba, as a novel state-space model (SSM), has gained widespread application in natural language processing and computer vision. In this work, we introduce Mamba-SEUNet, an innovative architecture that integrates Mamba with U-Net for SE tasks.
  arXiv  Detail & Related papers  (2024-12-21T13:43:51Z)
• Mamba-CL: Optimizing Selective State Space Model in Null Space for Continual Learning [54.19222454702032]
  Continual Learning aims to equip AI models with the ability to learn a sequence of tasks over time, without forgetting previously learned knowledge. State Space Models (SSMs) have achieved notable success in computer vision. We introduce Mamba-CL, a framework that continuously fine-tunes the core SSMs of the large-scale Mamba foundation model.
  arXiv  Detail & Related papers  (2024-11-23T06:36:16Z)
• A Comprehensive Survey of Mamba Architectures for Medical Image Analysis: Classification, Segmentation, Restoration and Beyond [2.838321145442743]
  Mamba is an alternative to template-based deep learning approaches in medical image analysis. It has linear time complexity, which is a significant improvement over transformers. Mamba processes longer sequences without attention mechanisms, enabling faster inference and requiring less memory.
  arXiv  Detail & Related papers  (2024-10-03T10:23:03Z)
• MAP: Unleashing Hybrid Mamba-Transformer Vision Backbone's Potential with Masked Autoregressive Pretraining [23.37555991996508]
  We propose Masked Autoregressive Pretraining (MAP) to pretrain a hybrid Mamba-Transformer vision backbone network. We show that both the pure Mamba architecture and the hybrid Mamba-Transformer vision backbone network pretrained with MAP significantly outperform other pretraining strategies.
  arXiv  Detail & Related papers  (2024-10-01T17:05:08Z)
• MambaVision: A Hybrid Mamba-Transformer Vision Backbone [54.965143338206644]
  We propose a novel hybrid Mamba-Transformer backbone, denoted as MambaVision, which is specifically tailored for vision applications. Our core contribution includes redesigning the Mamba formulation to enhance its capability for efficient modeling of visual features. We conduct a comprehensive ablation study on the feasibility of integrating Vision Transformers (ViT) with Mamba.
  arXiv  Detail & Related papers  (2024-07-10T23:02:45Z)
• DeciMamba: Exploring the Length Extrapolation Potential of Mamba [89.07242846058023]
  We introduce DeciMamba, a context-extension method specifically designed for Mamba. We show that DeciMamba can extrapolate context lengths 25x longer than the ones seen during training, and does so without utilizing additional computational resources.
  arXiv  Detail & Related papers  (2024-06-20T17:40:18Z)
• MambaOut: Do We Really Need Mamba for Vision? [70.60495392198686]
  Mamba, an architecture with RNN-like token mixer of state space model (SSM), was recently introduced to address the quadratic complexity of the attention mechanism. This paper conceptually concludes that Mamba is ideally suited for tasks with long-sequence and autoregressive characteristics. We construct a series of models named MambaOut through stacking Mamba blocks while removing their core token mixer, SSM.
  arXiv  Detail & Related papers  (2024-05-13T17:59:56Z)
• Vision Mamba: A Comprehensive Survey and Taxonomy [11.025533218561284]
  State Space Model (SSM) is a mathematical model used to describe and analyze the behavior of dynamic systems. Based on the latest state-space models, Mamba merges time-varying parameters into SSMs and formulates a hardware-aware algorithm for efficient training and inference. Mamba is expected to become a new AI architecture that may outperform Transformer.
  arXiv  Detail & Related papers  (2024-05-07T15:30:14Z)
• Visual Mamba: A Survey and New Outlooks [33.90213491829634]
  Mamba, a recent selective structured state space model, excels in long sequence modeling. Since January 2024, Mamba has been actively applied to diverse computer vision tasks. This paper reviews visual Mamba approaches, analyzing over 200 papers.
  arXiv  Detail & Related papers  (2024-04-29T16:51:30Z)
• ReMamber: Referring Image Segmentation with Mamba Twister [51.291487576255435]
  ReMamber is a novel RIS architecture that integrates the power of Mamba with a multi-modal Mamba Twister block. The Mamba Twister explicitly models image-text interaction, and fuses textual and visual features through its unique channel and spatial twisting mechanism.
  arXiv  Detail & Related papers  (2024-03-26T16:27:37Z)
• MambaMIL: Enhancing Long Sequence Modeling with Sequence Reordering in Computational Pathology [10.933433327636918]
  Multiple Instance Learning (MIL) has emerged as a dominant paradigm to extract discriminative feature representations within Whole Slide Images (WSIs) in computational pathology. 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. Our proposed framework performs favorably against state-of-the-art MIL methods.
  arXiv  Detail & Related papers  (2024-03-11T15:17:25Z)
• MedMamba: Vision Mamba for Medical Image Classification [0.0]
  Vision transformers (ViTs) and convolutional neural networks (CNNs) have been extensively studied and widely used in medical image classification tasks. Recent studies have shown that state space models (SSMs) represented by Mamba can effectively model long-range dependencies. We propose MedMamba, the first Vision Mamba for generalized medical image classification.
  arXiv  Detail & Related papers  (2024-03-06T16:49:33Z)
• Mamba-UNet: UNet-Like Pure Visual Mamba for Medical Image Segmentation [21.1787366866505]
  We propose Mamba-UNet, a novel architecture that synergizes the U-Net in medical image segmentation with Mamba's capability. Mamba-UNet adopts a pure Visual Mamba (VMamba)-based encoder-decoder structure, infused with skip connections to preserve spatial information across different scales of the network.
  arXiv  Detail & Related papers  (2024-02-07T18:33:04Z)
• Swin-UMamba: Mamba-based UNet with ImageNet-based pretraining [85.08169822181685]
  This paper introduces a novel Mamba-based model, Swin-UMamba, designed specifically for medical image segmentation tasks. Swin-UMamba demonstrates superior performance with a large margin compared to CNNs, ViTs, and latest Mamba-based models.
  arXiv  Detail & Related papers  (2024-02-05T18:58:11Z)
• Is Mamba Capable of In-Context Learning? [63.682741783013306]
  State of the art foundation models such as GPT-4 perform surprisingly well at in-context learning (ICL) This work provides empirical evidence that Mamba, a newly proposed state space model, has similar ICL capabilities.
  arXiv  Detail & Related papers  (2024-02-05T16:39:12Z)

This list is automatically generated from the titles and abstracts of the papers in this site.

This site does not guarantee the quality of this site (including all information) and is not responsible for any consequences.