MambaMIM:利用状态空间 Token 插值预训练 Mamba 及其在医学图像分割中的应用
MambaMIM: Pre-training Mamba with State Space Token Interpolation and its Application to Medical Image Segmentation
AI总结:
本文提出 MambaMIM 预训练框架,通过状态空间 token 插值和 3D 混合掩码增强 Mamba 的长程因果建模,并在 CT 预训练与八个分割基准上验证其先进性能。
摘要:
近来,状态空间模型 Mamba 展现出高效的长序列建模能力,尤其适用于处理 3D 医学成像中的长序列视觉任务。然而,现有生成式自监督学习方法尚未充分释放 Mamba 在处理长程依赖方面的潜力,因为它们忽视了掩码建模中状态空间序列固有的因果特性。为应对这一挑战,我们提出了一个名为 MambaMIM 的通用预训练框架,这是一种基于新型 TOKen-Interpolation 策略(TOKI)的掩码图像建模方法,面向选择性结构状态空间序列,用于学习掩码序列内状态空间的因果关系。此外,MambaMIM 引入自底向上的 3D 混合掩码策略,以保持不同架构之间的掩码一致性,并可用于任何单一或混合 Mamba 架构,以增强其多尺度和长程表示能力。我们在包含 6.8K 张 CT 扫描的大规模数据集上预训练 MambaMIM,并在八个公开医学分割基准上评估其性能。大量下游实验揭示了使用 Mamba 进行医学图像预训练的可行性和先进性。特别是,当我们将 MambaMIM 应用于一个混合 MedNeXt 和 Vision Mamba 的定制架构时,持续获得最先进的分割性能。代码可在 https://github.com/FengheTan9/MambaMIM 获取。
英文摘要:
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.
机构、总结、中文摘要均由AI分析而来,可能存在错误,请以论文正文为准。