MambaClinix：3D医学图像分割新网络！分层门控卷积和基于Mamba的U-Net

点击下方卡片，关注“CVer”公众号

AI/CV重磅干货，第一时间送达

点击进入—>【Mamba和医学影像】交流群

添加微信号：CVer111，小助手会拉你进群！

扫描下方二维码，加入CVer学术星球！可以获得最新顶会/顶刊上的论文idea和CV从入门到精通资料，及最前沿应用！发论文/搞科研/涨薪，强烈推荐！

转载自：极市平台

导读

 

MambaClinix通过自适应分阶段设计将多层级CNN的高阶空间交互能力与SSM的全局依赖能力相结合，提供一个灵活的框架，可无缝集成到临床医学图像分割任务中，训练过程最大限度地减少人工干预，增强模型学习的自动化水平

深度学习技术，尤其是CNN和Transformers，极大推动了3D医学图像分割的发展。然而CNN受制于局部感受野，阻碍其在复杂临床场景中的应用。Transformers能够有效捕捉长距离关联特征，但计算量较大，训练和部署成本高昂。近来，基于状态空间模型(SSM)的Mamba架构被提出，它在维持线性计算复杂度的同时，能够有效建模长距离依赖关系。然而，Mamba在医学图像分割任务中也暴露一些缺陷，特别是在捕捉高敏感的局部特征方面，临床效果并不理想。

该研究提出一种MambaClinix模型，它将分层门控卷积网络(HGCN)与Mamba集成在一个自适应的分阶段(stage-wise)U型框架中。这种设计方式利于特征图的高阶空间交互，使模型能有效捕捉医学图像中的近端和远端关系。其中，该研究提出的HGCN网络，利用纯卷积结构来模仿Transformers注意力机制，成功促进了高阶特征的交互计算。此外，该研究引入了一种区域特定的Tversky损失函数，用于强调特定体素区域的学习效果，从而优化模型的决策过程。该研究最终在五个基准数据集上进行消融和比较实验，结果表明所提出的MambaClinix在保持低模型复杂度的同时获得了高分割精度。

代码：https://github.com/CYB08/MambaClinix-PyTorch
论文：https://arxiv.org/abs/2409.12533
数据及预训练模型：
https://drive.google.com/drive/folders/111n2yo68O3s7kZFjwo7840B-pdNWkAvG

MedAna3.0为团队自研的医学影像综合分析平台。该平台的核心分割模型即为MambaClinix. MedAna3.0视频简介如下：

Link：https://www.bilibili.com/video/BV1VYbFepEDk/?spm_id_from=333.999.0.0

1 Introduction

医学图像分割能够精确描绘解剖结构，提高临床诊断质量。在该领域，卷积神经网络(CNN)因其固有的平移不变性而成为主流模型。CNN这种能力在医学成像中至关重要，因为与疾病相关的ROI可能出现在不同的方向和位置。然而大量研究表明，CNN的局部感受野制约了它们捕捉全局特征的能力，给更精准的医学图像分割带来了挑战。在医学图像分析中，理解更大的解剖结构，解析临床图像不同部分之间的空间关系，对于分析复杂影像特征，识别大范围病灶至关重要。

近年来，Transformer模型由于其独特的自注意力机制，使它在捕捉医学图像长距离特征方面展现出了效果。这种自注意力机制促进了不同距离和复杂度的空间特征相互作用，分割效果超越了标准CNN，产生了令人印象深刻的结果。然而，Transformers计算复杂度较高且需大量样本用于训练，使得它们不太适合轻量级模型部署。尽管如此，Transformers中的自注意力机制为图像分割带来新的启发。一些研究表明，自注意力机制的有效性源于其能够促进图像内高阶特征交互。这种能力是通过动态调整注意力权重和自适应地聚焦于各个空间维度来实现的，这有助于全面分析医学影像中的全局特征关系。以此为动机，我们仿照高阶空间交互特点，开发了一种用于3D医学图像分割的新型分层门控卷积网络(HGCN)。HGCN采用纯卷积结构，以递归方式在多个维度上进行空间特征交互。然而，进一步的实验表明，随着网络的加深，提升HGCN的空间交互阶数以捕获长距离特征的成本效益将逐渐降低。此外，涉及高阶空间交互的计算需要大量的递归和门控卷积过程，这不仅需要大量的GPU资源，还可能增加过拟合的风险。

最近，状态空间模型(SSM)，尤其是Mamba，在有效捕获长距离依赖关系方面表现出明显的优势。与基于Transformer的模型相比，Mamba卓越的计算效率和硬件加速算法，使其非常适合处理3D医学图像等复杂数据。随后，一些研究探索了将CNN与Mamba相结合，旨在将CNN详细的局部特征提取与Mamba的全局特征表示能力相结合，提高医学图像分割的精度。

在该研究中，为了平衡模型在捕捉局部和全局特征方面的能力，我们提出了MambaClinix，一种自适应的分阶段（stage-wise）建模框架，它融合了HGCN和Mamba组块，针对医学图像分割进行了优化。在编码器的较低阶段，使用HGCN来扩展高阶空间交互。为了克服增加HGCN交互阶数带来的收益递减，在较高阶段加入了一个残差Mamba模块来取代HGCN，用于提取长距离依赖特征。HGCN增强了Mamba块的语义特征，丰富了输入质量。这种分阶段方法使模型能够更深入地理解医学图像中的近端和远端关系，这对于临床医学分割任务至关重要。MambaClinix继承了nnU-Net的自配置策略，允许自动调整网络结构以匹配特定的数据集特征。这种自适应配置可确保模型架构经过精细调整以满足不同数据集的独特需求。总体而言，MambaClinix通过自适应分阶段设计将多层级CNN的高阶空间交互能力与SSM的全局依赖能力相结合，提供一个灵活的框架，可无缝集成到临床医学图像分割任务中，训练过程最大限度地减少人工干预，增强模型学习的自动化水平。

2 Method

2.1 Framework

该研究提出的MambaClinix架构如图1所示，采用自适应stage-wise设计，组合不同功能组块在一个U形框架中。在编码器的较低阶段，设计了一个分层门控卷积网络(HGCN)块，如图2(b)所示，通过纯卷积结构搭建空间交互网络。在较高阶段，采用残差Mamba块，如图2(a)所示，取代了HGCN，用于增强对长距离依赖关系的捕捉能力。由此，所提出的MambaClinix将编码器分为两部分：第一部分利用HGCN实现高效的空间特征交互，第二部分将这些特征输入残差Mamba块，以确保全面理解更大的空间信息。该架构使网络能够平衡细节特征处理和全局特征集成，优化医学图像分割的深度和广度。MambaClinix采用nnU-Net的自配置方法，可自动调整网络结构以适应每个数据集的特征。此外，所提出的HGCN能够深度集成到这种自适应配置中，允许根据数据集特征动态调整HGCN空间交互阶数。这种灵活性确保HGCN可以自适应地增加阶数以平衡结构复杂性和计算效率。在模型训练期间还应用了区域特定的Tversky损失函数，对3D医学图像的每个子体素进行不同的损失加权。这种方法将模型的注意力引导到更具挑战性的分割区域，改善细节学习并提高整体分割性能。

2.2 Hierarchical Gated Convolutional Network (HGCN)

HGCN嵌入在框架中，用于在早期阶段动态捕获图像特征。其核心处理由高阶门控卷积(hgConv)层实现，该层采用门控卷积机制和递归过程促进高阶空间交互。hgConv通过将交互扩展到多个级别来绕过自注意力机制的复杂计算，增强了CNN捕获医学图像中多级空间依赖关系的能力。此外，hgConv可适应各种临床数据场景，与自适应框架无缝集成。在编码器的第s阶段，具有通道的输入特征首先通过标准残差网络(Res)，然后Stem层(Stem)将通道数调整为。Stem能够确保后续hgConv具有稳定的输入结构。经过层归一化(LN)后，数据通过hgConv处理以实现高阶空间交互。然后将生成的特征图映射为输出量，并与残差分量相结合。该过程可以表示为：




高阶门控卷积(hgConv) 。给定一个3D输入特征。函数φ涉及两个连续卷积过程，用于将原始通道转为通道扩展并进一步细化的子向量集。输入投影首先将通道数扩展2倍，然后经过深度卷积层，将通道进一步分解为多个子向量组合，该过程可以表示为：

其中，n表示hgConv的空间交互阶数。这种细分策略是通过将通道划分为逐渐减小的子部分来实现的，如下式所示。

然后，每个细分的子向量经历从j到j+1的变换过程，其公式如下：

其中 表示对应于的卷积函数。一个外部参数γ用于调节前一次的输出。此操作促进了不同层之间的信息交互，增强了模型捕捉医学图像中复杂特征的能力。

2.3 Mamba Block

残差Mamba组块利用残差连接和层归一化来增强原始Mamba的空间建模能力。它放置在编码器的较高阶段，以总结HGCN输出的全局特征，捕获医学图像中的长距离依赖关系。在第s个阶段，输入特征首先经过核大小为1×1×1的残差卷积，然后进行批量归一化(BN)和非线性激活函数σ。变换后的输出进行展平、转置、归一化，并通过具有两个并行分支的Mamba块进行处理。一个分支从层归一化特征上的线性函数(LN)开始，经过深度卷积(dwConv)和SiLU函数，再由状态空间模型(SSM)层进行特征增强。对于没有SSM层的另一个分支，它先经过线性函数，然后是SiLU函数。最终，使用Hadamard积将两个分支的特征组合起来，重新整形为原始形状，并通过残差连接重新集成到网络中。这个过程可以写成：

其中，MLP(·)是一个多层感知结构。该架构保留了图像特征流的连续性和完整性，有助于梯度在网络中传播。

2.4 Stage-wise Integration of HGCN and Mamba

随着网络的加深，增加HGCN的计算阶数以捕捉医学图像全局特征的收益逐渐减少。为了解决这个问题，我们在更高阶段引入了残差Mamba块来替代HGCN。这种替代旨在提取长距离依赖关系，同时保持计算效率。HGCN在早期阶段捕获的详细图像特征为Mamba块提供了丰富的信息，提高了其接收输入的质量。给定一个具有s个阶段的编码器，其中H和M分别代表HGCN和残差Mamba组块。编码器的stage-wise设计,可以表示为：

其中，s是总体阶段(stage)数，s能够根据特定数据集特征自适应调整。考虑到计算复杂性和GPU限制，HGCN块计算的空间交互阶数最小为2，最大为6。表示第s//2阶段的HGCN块，代表该阶段计算了(s//2+1)阶的空间交互；表示第s阶段的Mamba块。这种stage-wise设计展现了的可定制功能，是如何与自适应配置策略深度集成的。

2.5 Region-Specific Loss Function

在模型训练过程中，Dice损失函数会优化真阳性(TP)，并同等惩罚假阳性(FP)和假阴性(FN)。然而，目标器官和背景之间的图像差异通常会造成数据不平衡，使预测偏向背景，导致FN比FP增加更快。而在临床环境中，高召回率(减少FN)指标是需要重点优化的，以确保能够正确识别到每个可能的疾病区域，即使以牺牲一些FP为代价。例如，肺部CT扫描中遗漏一个小结节可能会延误肺癌治疗，而FP可以通过额外的人工检查来纠正。为了解决这个问题，Tversky损失函数被提出，其定义为：

然而，这种Tversky损失函数关注的是总体三维影像，对目标区域和背景区域的体素进行统一惩罚，却忽略了不同子区域之间的分割难度差异。为了克服这一限制，我们提出了一个区域特定的损失函数，专注于优化三维医学图像中的特定子区域。在训练过程中，该方法会动态调整不同区域的惩罚，以提高整体预测准确性，为更难分割的子区域分配更高的权重。因此，将区域特定的Tversky损失函数改写为：

其中，由N个体素组成的医学图像被划分为k个子体素，表示为。模型聚焦计算这些子体素的损失。并且区域特定损失的梯度仅取决于对这些子区域的评估，而不是整个图像。

3 Experiment and results

3.1 Datasets

为了获取反映多样化和真实临床场景的数据，我们在青岛大学附属医院山东省数字医学与计算机辅助手术重点实验室中，进行了一项回顾性研究。本研究从2021年至2023年间，随机采样研究数据，构建三个内部数据集（PCD、LungT和LiverT）。此外，我们还收集了两个公共数据集（ABD [Umamba] 和 BraTs2021）用于实验评估。数据集详细信息如表1所示。MambaClinix继承了nnU-Net的自配置功能，针对不同数据集的网络配置参数如表2所示。

(1)肺循环CECT数据集(PCD)。该内部数据集包含547个三维增强CT(CECT)图像，这些图像经过精确注释以分割肺循环系统，并带有支气管、肺动脉和肺静脉的标签。它来自2022年至2023年期间在青岛大学附属医院随机抽样的病例。每个图像都由三名经验丰富的高年资医生手工标注，随后由临床医生审查以确保准确性。该数据集已用于临床环境中的辅助诊断。

(2)肺肿瘤CECT数据集(LungT)。该内部数据集包括800个CECT图像，重点聚焦肺肿瘤。它是从2021-2023年期间在青岛大学附属医院随机抽样的病例中收集的。由三名经验丰富的高年资医生精准标注，随后由临床医生完善，以确保在临床诊断和手术规划中的实用性。

(3)肝肿瘤CT数据集(LiverT)。该数据集包含292个肝肿瘤样本，并分为两部分。第一部分是内部收集的161张CECT图像，这些图像来自2021年至2023年期间青岛大学附属医院随机选择的病例。这些样本由经验丰富的医生注释并由临床医生验证。第二部分包括来自公开的医学分割十项全能(MSD)挑战赛的131个样本，专门针对CT图像中肝肿瘤分割，每个图像都清晰标记以区分肝脏和肿瘤。

(4)腹部CT数据集(ABD)。该数据集源自 MICCAI 2022 FLARE挑战赛，是一个由100例CT影像组成的分割数据集。旨在分割13种不同的腹部器官。该数据集与U-Mamba设置一致，训练集包含来自MSD Pancreas数据集的50个CT影像，注释由AbdomenCT-1K提供。另外50个病例来自多医疗中心作为测试集。此设置可在不同的成像场景中对模型进行稳健评估。

(5)脑肿瘤MRI数据集(BraTS)。该数据集来自BraTS2021挑战，专为脑肿瘤分割而设计，包含1,251个3D脑磁共振成像(MRI)。每个样本具有四种成像模式（T1、T1ce、T2 和 Flair）。该数据集包含三种分割区域：整个肿瘤 (WT)、增强肿瘤 (ET) 和肿瘤核心 (TC)。

3.2 Implementation Setup

MambaClinix基于nnU-Net和UMamba框架开发，总体模型在PyTorch 2.0.1上实现。所有实验在四块NVIDIA GeForce RTX 4090 GPU 上进行。优化器采用Adam，初始化学习率为1e-4。使用PolyLRScheduler进行学习率调节。MambaClinix在PCD数据集上训练了300个epoch，在LungT数据集和LiverT数据集上训练了1000个epoch，在ABD数据集和BraTS数据集上训练了500个epoch。

3.3 Benchmark Results

为了评估MambaClinix的效果，本研究比较了医学图像分割领域的几种成熟模型。结果如表3和图3-5所示。所有比较模型根据其基础架构分为三组：基于CNN的模型（nnU-Net 、SegResNet）、基于Transformer的模型（UNETR、SwinUNETR）和基于Mamba的模型（U-Mamba、LightM-Unet和MambaClinix）。使用DSC系数和mIoU等指标评估性能。所有模型均使用nnU-Net框架实现，并进行一致的图像预处理以确保公平比较。

在LungT数据集上，MambaClinix的Dice得分为72.78%，鉴于肺组织的高度变异性和相似的密度，这一结果非常显著，表明其能够准确分割较小的、对比度较低的区域。在包含一系列具有复杂几何形状的腹部器官ABD数据集上，MambaClinix获得了最高的84.69%。值得注意的是，该数据集的训练集和测试集来自不同的数据源，凸显了该模型在全局特征提取和关键特征迁移方面的稳健性。LiverT数据集来自两个不同源的混合，MambaClinix的Dice得分为 71.37%，优于其他模型，并证明了其在处理具有严重结构异常的图像方面是有效的。虽然MambaClinix在PCD和BraTS数据集上没有取得最高分，但它仍然取得了极具竞争力的结果，展示了其在复杂临床场景中的稳健性。

图 3. 在a)PCD数据集和b)ABD数据集上的视觉比较，突出显示具有显著差异的区域。

图 4. 在a)LiverT数据集和b)BraTS数据集上进行视觉比较，突出显示具有显著差异的区域。

图 5. 结果可视化。a)LiverT数据集上肝肿瘤的3D可视化；b)LungT数据集的分割结果，包括肺肿瘤的空间定位。

3.4 Ablation Study

我们进行了消融研究，以评估MambaClinix架构中不同组块的功能性。实验是在BraTS、LiverT和LungT数据集上进行的，结果详见表4。用标准残差块替换MambaClinix中Mamba和HGCN块来建立baseline模型，用作后续消融实验的比较。不同组块被逐步集成或替换到基线配置中。这种方法允许系统地评估每个块对整体分割性能的贡献。更详细的消融比较实验，可以参考作者原文。

3.5 Efficiency of Region Specific Loss

该实验中，作者首先采用交叉熵(CE)和Dice相结合的复合损失函数。进一步将三维图像划分为子体素，应用特定区域的Dice损失来单独优化每个子区域的预测，称为基线(RS)。接下来，加入Tversky损失项(LTRS)，FP和FN的固定惩罚系数分别设置为0.3和0.7。最后，引入这些惩罚系数的自适应策略，形成MambaClinix的最终损失函数。结果列于表5中。更详细的消融比较实验，可以参考作者原文。

4 Conclusion

该研究提出了MambaClinix，这是一种新颖的U型架构医学图像分割模型，它将HGCN和Mamba模块集成在一个自适应的stage-wise框架内。该方法进一步结合了区域特定的损失函数来优化模型的决策能力。MambaClinix架构能够自适应配置自身网络结构以适应不同的数据集，同时应用区域特定损失的策略，使其能够应用于多种临床分割任务。跨多个数据集的实验结果表明，与基于Transformer的方法相比，MambaClinix提供了卓越的计算效率，同时分割性能优于其他最先进的分割模型。

Reference

[1] Chenyuan Bian, et al., MambaClinix: Hierarchical Gated Convolution and Mamba-Based U-Net for Enhanced 3D Medical Image Segmentation, 2024

何恺明在MIT授课的课件PPT下载

在CVer公众号后台回复：何恺明，即可下载本课程的所有566页课件PPT！赶紧学起来！

ECCV 2024 论文和代码下载

在CVer公众号后台回复：ECCV2024，即可下载ECCV 2024论文和代码开源的论文合集

CVPR 2024 论文和代码下载

在CVer公众号后台回复：CVPR2024，即可下载CVPR 2024论文和代码开源的论文合集

Mamba和医学影像交流群成立

扫描下方二维码，或者添加微信号：CVer111，即可添加CVer小助手微信，便可申请加入CVer-Mamba、医学影像微信交流群。另外其他垂直方向已涵盖：目标检测、图像分割、目标跟踪、人脸检测&识别、OCR、姿态估计、超分辨率、SLAM、医疗影像、Re-ID、GAN、NAS、深度估计、自动驾驶、强化学习、车道线检测、模型剪枝&压缩、去噪、去雾、去雨、风格迁移、遥感图像、行为识别、视频理解、图像融合、图像检索、论文投稿&交流、PyTorch、TensorFlow和Transformer、NeRF、3DGS、Mamba等。
一定要备注：研究方向+地点+学校/公司+昵称（如Mamba、医学影像+上海+上交+卡卡），根据格式备注，可更快被通过且邀请进群

▲扫码或加微信号: CVer111，进交流群
CVer计算机视觉（知识星球）来了！想要了解最新最快最好的CV/DL/AI论文速递、优质实战项目、AI行业前沿、从入门到精通学习教程等资料，欢迎扫描下方二维码，加入CVer计算机视觉（知识星球），已汇集上万人！

▲扫码加入星球学习

▲点击上方卡片，关注CVer公众号
整理不易，请赞和在看