MambaMIM: 基于状态空间令牌插值的Mamba预训练及其在医学图像分割中的应用
计算机视觉与模式识别
2025-05-15 v2
摘要
最近，状态空间模型Mamba在高效长序列建模方面展现出能力，特别是在处理3D医学成像中的长序列视觉任务方面。然而，现有的生成式自监督学习方法尚未充分释放Mamba处理长程依赖的潜力，因为它们忽略了状态空间序列在掩码建模中的固有因果特性。为应对这一挑战，我们提出了一种通用预训练框架MambaMIM，这是一种基于新型TOKen-Interpolation（TOKI）策略的掩码图像建模方法，用于选择性结构状态空间序列，在掩码序列中学习状态空间的因果关系。此外，MambaMIM引入了一种自底向上的3D混合掩码策略，以在不同架构之间保持掩码一致性，并可应用于任何单一或混合Mamba架构，以增强其多尺度和长程表示能力。我们在一个包含6.8K张CT扫描的大规模数据集上预训练MambaMIM，并在八个公共医学分割基准上评估其性能。广泛的下游实验揭示了使用Mamba进行医学图像预训练的可行性和进步。特别是，当我们把MambaMIM应用于将MedNeXt和Vision Mamba混合的自定义架构时，我们持续获得了最先进的语义分割性能。代码已开源，地址为：https://github.com/FengheTan9/MambaMIM。
引用
@article{arxiv.2408.08070,
title = {MambaMIM: Pre-training Mamba with State Space Token Interpolation and its Application to Medical Image Segmentation},
author = {Fenghe Tang and Bingkun Nian and Yingtai Li and Zihang Jiang and Jie Yang and Wei Liu and S. Kevin Zhou},
journal= {arXiv preprint arXiv:2408.08070},
year = {2025}
}
备注
Accepted by Medical Image Analysis. Code: https://github.com/FengheTan9/MambaMIM