nnMamba：基于状态空间模型的3D生物医学图像分割、分类和地标检测
摘要
本文提出了一种基于状态空间模型（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性能，且模型轻量化优势显著。未来工作将探索其在更多模态医学图像中的应用。
代码开源地址：https://github.com/lhaof/nnMamba