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

TL;DR

MambaMIM introduces a novel pre-training framework for the Mamba model using token interpolation and hybrid masking, significantly improving long-range dependency modeling in 3D medical image segmentation.

Contribution

It proposes MambaMIM, a new masked image modeling method with token interpolation and hybrid masking, enhancing Mamba's ability to capture long-range dependencies in medical imaging.

Findings

Achieved state-of-the-art segmentation performance on multiple benchmarks.

Demonstrated effective learning of causal relationships in state space sequences.

Enhanced Mamba architecture with improved multi-scale and long-range representations.

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.