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

相关推荐
daifgFuture8 小时前
Android 3D球形水平圆形旋转,旋转动态更换图片
android·3d
牧子川13 小时前
【论文解读】CVPR2023 PoseFormerV2:3D人体姿态估计(附论文地址)
3d·cvpr2023·poseformerv2
资深设备全生命周期管理13 小时前
优化版本,增加3D 视觉 查看前面的记录
3d
m0_7482507418 小时前
GPUCUDA 发展编年史:从 3D 渲染到 AI 大模型时代(上)
人工智能·3d
少林6591 天前
谷歌地图高清卫星地图2026中文版下载|谷歌地图3D卫星高清版 V7.3.6.9796 最新免费版下载 - 前端工具导航
3d·谷歌地图
LeonDL1681 天前
HALCON 深度学习训练 3D 图像的几种方式优缺点
人工智能·python·深度学习·3d·halcon·halcon训练3d图像·深度学习训练3d图像
xhload3d3 天前
图扑软件 | 带你体验 Low Poly 卡通三维世界
物联网·3d·智慧城市·html5·webgl·数字孪生·可视化·工业互联网·三维建模·工控·轻量化·中国风·卡通动画·写实风格·科技风·low poly
图扑数字孪生3 天前
基于 HT for Web 轻量化 3D 数字孪生数据中心解决方案
3d·数字孪生·三维可视化·数据中心·智慧机房
njsgcs3 天前
PolyGen:一个用于 3D 网格的自回归生成模型 论文阅读
3d
Angel Q.3 天前
PnP(Perspective-n-Point)算法 | 用于求解已知n个3D点及其对应2D投影点的相机位姿
数码相机·算法·3d·pnp