专栏名称: 极市平台
极市平台是由深圳极视角推出的专业的视觉算法开发与分发平台,为视觉开发者提供多领域实景训练数据库等开发工具和规模化销售渠道。本公众号将会分享视觉相关的技术资讯,行业动态,在线分享信息,线下活动等。 网站: http://cvmart.net/
51好读  ›  专栏  ›  极市平台

青大附院数字医学实验室提出MambaClinix,医学影像分割任务超越传统SOTA!

极市平台  · 公众号  ·  · 2024-09-25 22:00

正文

↑ 点击 蓝字 关注极市平台
编辑丨极市平台

极市导读

MambaClinix通过自适应分阶段设计将多层级CNN的高阶空间交互能力与SSM的全局依赖能力相结合,提供一个灵活的框架,可无缝集成到临床医学图像分割任务中,训练过程最大限度地减少人工干预,增强模型学习的自动化水平。 >> 加入极市CV技术交流群,走在计算机视觉的最前沿

深度学习技术,尤其是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 个子体素,表示为







请到「今天看啥」查看全文