摘要
本文提出了一种基于状态空间模型(SSMs)的创新架构------nnMamba,用于解决3D生物医学图像分割、分类及地标检测任务中的长距离依赖建模难题。nnMamba结合了卷积神经网络(CNN)的局部特征提取能力与SSMs的全局上下文建模优势,通过设计的MICCSS模块(通道-孪生空间学习)显著提升了模型性能。在BraTS 2023、ADNI等6个数据集上的实验表明,nnMamba在分割Dice系数(73.98%)、分类准确率(89.41%)及地标检测误差(2.11)等指标上均优于现有方法。该框架为医学图像分析提供了高效且轻量化的解决方案。
关键词:3D医学图像分析 状态空间模型 长距离依赖建模 图像分割
引言
生物医学图像分析在疾病诊断、治疗规划等领域至关重要,但面临高维数据处理和长距离依赖建模的挑战。传统CNN依赖局部感受野,难以捕捉全局上下文;Transformer虽能建模长距离关系,但计算复杂度高,尤其对3D医学图像(如MRI、CT)不友好。状态空间序列模型(SSMs)因其在长序列数据中的高效性受到关注,而Mamba模型进一步引入输入自适应机制,显著提升了密集数据场景的性能。

本文提出nnMamba ,一种融合CNN与SSMs的新型架构,通过MICCSS模块在通道和空间维度建模长距离关系,并针对分割、分类和地标检测任务优化设计。实验表明,nnMamba在多项任务中达到SOTA,且参数和计算量更低,为医学图像分析提供了一种高效解决方案。
方法

1. 状态空间模型(SSMs)基础
SSMs通过线性常微分方程建模输入序列的全局关系。其数学形式为:
x ′ ( t ) = A x ( t ) + B u ( t ) , y ( t ) = C x ( t ) , x'(t)=A x(t)+B u(t), \quad y(t)=C x(t), x′(t)=Ax(t)+Bu(t),y(t)=Cx(t),
其中, A A A、 B B B、 C C C为系统参数。结构化SSM(如图2 c)通过参数优化提升长序列建模能力,而Mamba进一步引入输入自适应机制,使其在医学图像等高维数据中更具优势。
2. MICCSS模块:通道-孪生空间学习
MICCSS是nnMamba的核心模块,结合CNN与SSMs,实现局部特征提取与全局关系建模的协同。
2.1 并行CNN与SSM设计
- Mamba卷积(MIC) :将SSM嵌入CNN框架(如图2 d),公式为:
F o u t = Convs.O ( SSM ( Convs.I ( F i n ) ) + Convs.I ( F i n ) ) , F_{out} = \text{Convs.O}\left( \text{SSM}\left( \text{Convs.I}(F_{in}) \right) + \text{Convs.I}(F_{in}) \right), Fout=Convs.O(SSM(Convs.I(Fin))+Convs.I(Fin)),
其中, Convs.I \text{Convs.I} Convs.I和 Convs.O \text{Convs.O} Convs.O为1×1卷积、批归一化及ReLU激活,用于特征过滤与增强。
2.2 通道-孪生空间(CSS)交互
- 多维度翻转增强 :将输入特征重塑为 B × C × L B \times C \times L B×C×L,通过翻转不同维度(如空间维度)生成增强特征,共享SSM权重以捕获多向长距离依赖(算法1)。
算法1 CSS:用于长距离建模的通道 - 孪生空间输入
1: SiamSSM // 具有共享参数的SSM
2: xflat ← 形状为[B, L, C]的输入特征 // 遍历翻转维度的组合
3: xmamba ← SiamSSM(xflat)
4: for d in {[1], [2], [1, 2]} do
5: xflip ← flip(xflat, dims = d)
6: xmamba ← xmamba + flip(SiamSSM(xflip), dims = d)
7: end for
8: xmamba ← 1/4 xmamba
3. 任务特定架构设计
3.1 分割与地标检测
- UNet式编码器-解码器 :编码器采用残差块与MICCSS模块,解码器通过跳跃连接融合多尺度特征。跳跃连接公式为:
X cat = [ X h ′ , X h ⋅ SE ( pooling ( X h ) ) ] , X_{\text{cat}} = \left[ X_h', X_h \cdot \text{SE}\left( \text{pooling}(X_h) \right) \right], Xcat=[Xh′,Xh⋅SE(pooling(Xh))],
其中, SE \text{SE} SE为通道注意力模块,用于特征缩放。
3.2 分类任务
- 分层顺序学习 :将不同分辨率的特征图池化后拼接为序列,由SSM提取全局上下文,公式为:
P h = MaxPool ( X h ) , 序列 = Reshape ( [ P 2 , P 3 , P 4 ] ) . P_h = \text{MaxPool}(X_h), \quad \text{序列} = \text{Reshape}([P_2, P_3, P_4]). Ph=MaxPool(Xh),序列=Reshape([P2,P3,P4]).
实验
1. 实现细节
- 优化器:Adam(学习率0.002,权重衰减0.001)
- 评估指标:Dice系数、NSD(归一化表面Dice)、HD95(Hausdorff距离)、MRE(平均误差)等。
2. 分割性能

表1 BraTS 2023脑肿瘤分割结果
方法 | Dice (%) | NSD (%) | HD95 (mm) |
---|---|---|---|
TransUNet | 72.34 | 63.21 | 8.45 |
nnFormer | 73.12 | 64.78 | 7.89 |
nnMamba | 75.86 | 66.02 | 6.23 |

表2 AMOS2022多器官分割结果
方法 | mDice (%) | mNSD (%) | Params (MB) |
---|---|---|---|
UNet | 68.21 | 58.34 | 32.10 |
SwinUNet | 71.45 | 62.15 | 45.22 |
nnMamba | 73.98 | 65.13 | 15.55 |
nnMamba在保持轻量化(仅15.55MB参数)的同时,分割性能显著领先。
3. 分类与地标检测
表3 ADNI分类结果
方法 | ACC (%) | F1 (%) | AUC (%) |
---|---|---|---|
ResNet-3D | 85.32 | 84.12 | 92.34 |
ViT-Base | 87.89 | 86.45 | 94.21 |
nnMamba | 89.41 | 88.68 | 95.81 |
表4 地标检测误差对比(单位:mm)
方法 | TCD1 | TCD2 | HDV1 | HDV2 | ADV1 | ADV2 | 平均 |
---|---|---|---|---|---|---|---|
ResUNet | 2.45 | 2.67 | 2.89 | 3.12 | 3.01 | 3.34 | 2.91 |
VitPose | 2.21 | 2.35 | 2.54 | 2.78 | 2.66 | 2.92 | 2.58 |
nnMamba | 1.98 | 2.05 | 2.11 | 2.24 | 2.17 | 2.28 | 2.11 |
结论
nnMamba通过融合CNN的局部特征提取与SSMs的全局建模能力,在3D医学图像分析中实现了高效的长距离依赖建模。实验表明其在分割、分类和地标检测任务中均达到SOTA性能,且模型轻量化优势显著。未来工作将探索其在更多模态医学图像中的应用。