Mamba-Sea: A Mamba-based Framework with Global-to-Local Sequence Augmentation for Generalizable Medical Image Segmentation
作者: Zihan Cheng, Jintao Guo, Jian Zhang, Lei Qi, Luping Zhou, Yinghuan Shi, Yang Gao
分类: cs.CV
发布日期: 2025-04-24
备注: Accepted by IEEE TMI 2025. The code is available at https://github.com/orange-czh/Mamba-Sea
🔗 代码/项目: GITHUB
💡 一句话要点
Mamba-Sea:基于Mamba和全局-局部序列增强的医学图像分割通用框架
🎯 匹配领域: 支柱二:RL算法与架构 (RL & Architecture)
关键词: 医学图像分割 领域泛化 Mamba 序列模型 全局增强 局部增强 分布偏移
📋 核心要点
- 现有医学图像分割方法在领域泛化方面存在不足,难以应对不同数据集间的分布偏移。
- Mamba-Sea通过全局和局部序列增强,抑制模型对领域特定信息的过度学习,提升泛化能力。
- 实验表明,Mamba-Sea在Prostate数据集上取得了超过90%的Dice系数,超越了现有最佳方法。
📝 摘要(中文)
为了解决医学图像分割中的分布偏移问题,领域泛化(DG)成为一种有前景的设置,它在源域上训练模型,使其能够泛化到未见过的目标域。现有的DG方法主要基于CNN或ViT架构。最近,以Mamba为代表的先进状态空间模型在各种监督医学图像分割中显示出良好的效果。Mamba的成功主要归功于其捕获长程依赖关系的能力,同时保持与输入序列长度的线性复杂度,使其成为CNN和ViT的有希望的替代方案。受此启发,本文探索了Mamba架构在解决医学图像分割DG中分布偏移问题的潜力。具体来说,我们提出了一种新的基于Mamba的框架Mamba-Sea,它结合了全局到局部的序列增强,以提高模型在领域偏移问题下的泛化能力。我们的Mamba-Sea引入了一种全局增强机制,旨在模拟不同站点之间外观的潜在变化,从而抑制模型对领域特定信息的学习。在局部层面,我们提出了一种沿输入序列的序列式增强,通过建模和重采样与领域偏移相关的风格统计信息,来扰动随机连续子序列中token的风格。据我们所知,Mamba-Sea是第一个探索Mamba在医学图像分割中泛化的工作,提供了一种先进且有前景的基于Mamba的架构,对领域偏移具有很强的鲁棒性。值得注意的是,我们提出的方法首次在Prostate数据集上超过了90%的Dice系数,超过了之前88.61%的SOTA。
🔬 方法详解
问题定义:医学图像分割任务在实际应用中面临着数据分布偏移的问题,即模型在训练集(源域)上表现良好,但在测试集(目标域)上性能下降。现有的方法,如基于CNN或ViT的模型,在处理这种偏移时泛化能力有限,容易过拟合到源域的特定特征。
核心思路:Mamba-Sea的核心思路是通过全局和局部两个层面的序列增强,来模拟不同领域之间的差异,从而提高模型对未知目标域的泛化能力。全局增强旨在模拟不同站点间外观的潜在变化,局部增强则扰动输入序列中token的风格,使模型对风格变化更加鲁棒。
技术框架:Mamba-Sea是一个基于Mamba架构的医学图像分割框架。它主要包含以下模块:1) Mamba编码器:用于提取图像特征;2) 全局序列增强模块:模拟不同领域间的全局外观差异;3) 局部序列增强模块:扰动局部token的风格;4) 解码器:将提取的特征映射回分割结果。整体流程是:输入图像经过Mamba编码器提取特征,然后分别经过全局和局部序列增强,最后通过解码器得到分割结果。
关键创新:Mamba-Sea的关键创新在于将Mamba架构与全局-局部序列增强相结合,用于解决医学图像分割中的领域泛化问题。这是首次将Mamba应用于该任务,并设计了专门的增强策略来提高模型的泛化能力。与现有方法相比,Mamba-Sea能够更好地捕获长程依赖关系,并有效地抑制模型对领域特定信息的学习。
关键设计:全局增强模块通过模拟不同站点的外观变化来实现,具体方法未知。局部增强模块通过建模和重采样与领域偏移相关的风格统计信息,来扰动随机连续子序列中token的风格。损失函数和网络结构等其他技术细节在论文中可能有所描述,但摘要中未提及。
🖼️ 关键图片
📊 实验亮点
Mamba-Sea在Prostate数据集上取得了显著的性能提升,Dice系数超过90%,超越了之前SOTA的88.61%。这一结果表明,Mamba-Sea在领域泛化方面具有强大的能力,能够有效地应对医学图像分割中的分布偏移问题。
🎯 应用场景
Mamba-Sea可应用于各种医学图像分割任务,尤其是在数据分布存在差异的情况下,例如不同医院、不同扫描仪或不同疾病阶段的数据。该研究有助于提高医学图像分析的可靠性和准确性,辅助医生进行诊断和治疗,并降低对大量标注数据的依赖。
📄 摘要(原文)
To segment medical images with distribution shifts, domain generalization (DG) has emerged as a promising setting to train models on source domains that can generalize to unseen target domains. Existing DG methods are mainly based on CNN or ViT architectures. Recently, advanced state space models, represented by Mamba, have shown promising results in various supervised medical image segmentation. The success of Mamba is primarily owing to its ability to capture long-range dependencies while keeping linear complexity with input sequence length, making it a promising alternative to CNNs and ViTs. Inspired by the success, in the paper, we explore the potential of the Mamba architecture to address distribution shifts in DG for medical image segmentation. Specifically, 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 Mamba-Sea introduces a global augmentation mechanism designed to simulate potential variations in appearance across different sites, aiming to suppress the model's learning of domain-specific information. At the local level, we propose a sequence-wise augmentation along input sequences, which perturbs the style of tokens within random continuous sub-sequences by modeling and resampling style statistics associated with domain shifts. To our best knowledge, Mamba-Sea is the first work to explore the generalization of Mamba for medical image segmentation, providing an advanced and promising Mamba-based architecture with strong robustness to domain shifts. Remarkably, our proposed method is the first to surpass a Dice coefficient of 90% on the Prostate dataset, which exceeds previous SOTA of 88.61%. The code is available at https://github.com/orange-czh/Mamba-Sea.