1、摘要
在医学图像分割领域,基于神经网络和基于变换器的模型都得到了广泛的研究。然而,神经网络在长距离建模能力方面表现出局限性,而变形金刚则受到其二次计算复杂度的阻碍。近年来,以Mamba为代表的状态空间模型(SSM)已成为一种有前途的方法。它们不仅在模拟长距离相互作用方面表现出色,而且保持了线性计算复杂度。本文利用状态空间模型,提出了一种用于医学图像分割的Ushape体系结构模型--Vision Mamba UNet(VM-UNet).具体地说,引入视觉状态空间(VSS)块作为捕获大量上下文信息的基础块,并构造非对称的编解码器结构。我们在ISIC 17、ISIC 18和Synapse数据集上进行了全面的实验,结果表明VM-UNet在医学图像分割任务中具有较好的性能。据我们所知,这是第一个基于纯SSM模型构建的医学图像分割模型.我们的目标是建立一个基准,并为未来开发更高效、更有效的基于SSM的分割系统提供有价值的见解。我们的代码可在https://github.com/JCruan519/VM-UNet上获得。
2、引言
自动医学图像分割技术帮助医生更快地进行病理诊断,从而提高患者护理的效率。最近,基于CNN和基于Transformer的模型在各种视觉任务中表现出了出色的性能,特别是在医学图像分割方面。UNet [27]作为基于CNN的模型的代表,以其结构简单、可扩展性强而著称,后续的许多改进都是基于这种U型架构[11,37,28,29,30]。TransUnet [10]是基于Transformer模型的先驱,是第一个在编码阶段使用Vision Transformer(ViT)[13]进行特征提取的模型,并在解码阶段使用CNN,展示了全局信息获取的重要能力。随后,Transmission [36]采用了ViT和CNN的并行架构,同时捕获局部和全局特征。此外,Swin-UNet [9]将Swin Transformer [21]与U形架构相结合,首次引入了纯基于Transformer的U形模型。
尽管如此,基于CNN的模型和基于Transformer的模型都有固有的局限性。基于CNN的模型受到其局部感受野的限制,大大阻碍了它们捕获远程信息的能力。这通常会导致提取不充分的特征,从而导致次优分割结果。尽管基于变换器的模型在全局建模方面表现出上级性能,但自注意力机制要求图像尺寸方面的二次复杂性,导致高计算负担[31,13],特别是对于需要密集预测的任务,如医学图像分割。这些模型中目前的缺点迫使我们开发一种新的架构,医学图像分割,能够捕捉强大的远程信息,并保持线性计算复杂度。
最近,状态空间模型(SSM)引起了研究人员的极大兴趣。建立在经典SSM研究的基础上[18],现代SSM(例如,Mamba [16])不仅建立了长距离依赖关系,而且还表现出与输入大小相关的线性复杂度。此外,基于SSM的模型在许多领域都得到了大量的研究,包括语言理解[17,16],一般视觉[38,20]等。特别是,U-Mamba [24]最近引入了一种新的SSM-CNN混合模型,标志着其首次应用于医学图像分割任务。SegMamba [35]在编码器部分采用了SSM,同时在解码器部分仍然使用CNN,这表明SSM-CNN混合模型用于3D脑肿瘤分割任务。虽然上述工作已经利用SSM的医学图像分割任务,纯SSM为基础的模型的性能还有待探讨。
受VMamba 在图像分类任务中的成功影响,本文首次介绍了Vision Mamba UNet(VM-UNet),这是一种纯粹的基于SSM的模型,旨在展示医学图像分割任务的潜力。具体来说,VM-UNet由三个主要部分组成:编码器、解码器和跳过连接。编码器由来自VMamba的VSS块组成,用于特征提取,沿着用于下采样的补丁合并操作。相反,解码器包括VSS块和补丁扩展操作以恢复分割结果的大小。对于跳过连接组件,为了突出最原始的纯SSM模型的分割性能,我们采用了最简单的加法运算形式。
在器官分割和皮肤病变分割任务上进行了全面的实验,以证明纯SSM模型在医学图像分割中的潜力。具体来说,我们对Synapse [19],ISIC 17 [8]和ISIC 18 [12]数据集进行了广泛的实验,结果表明VM-UNet可以实现有竞争力的性能。此外,重要的是要注意,VM-UNet代表了纯基于SSM的分割模型的最基本形式,因为它不包括任何专门设计的模块。
本文的主要贡献可以概括如下:
1)我们提出了VM-UNet,标志着第一次探索纯粹基于SSM的模型在医学图像分割中的潜在应用。
2)在三个数据集上进行了综合实验,结果表明VM-UNet具有相当的竞争力。
3)我们为医学图像分割任务中的纯SSM模型建立了一个基线,为开发更高效和更有效的基于SSM的分割方法铺平了道路。
3、s(4)和mamba的知识
在现代基于SSM的模型中,即,结构化状态空间序列模型(S4)和Mamba都依赖于一个经典的连续系统,该系统将一维输入函数或序列(表示为x(t)∈ R)通过中间隐式状态h(t)∈ RN映射到输出y(t)∈ R。上述过程可以表示为线性常微分方程(ODE):
其中A ∈ RN×N表示状态矩阵,而B ∈ RN×1和C ∈ RN×1表示投影参数。
S4和Mamba将这个连续系统离散化,使其更适合深度学习场景。具体地说,他们引入了一个时间尺度参数λ,并使用固定的离散化规则将A和B转换为离散参数A-和B-。通常,零阶保持(ZOH)被用作离散化规则,并且可以定义如下:
在离散化之后,基于SSM的模型可以以两种方式计算:线性递归或全局卷积,分别定义为等式3和4。
其中K- ∈ RL表示结构化卷积核,L表示输入序列x的长度。
4、模型-Vision Mamba UNet (VM-UNet)
如图1(a)所示,VM-UNet的整体架构。具体地,VM-UNet包括补丁嵌入层、编码器、解码器、最终投影层和跳过连接。与以前的方法[9]不同,我们没有采用对称结构,而是采用了非对称设计。
Patch Embedding层将输入图像x ∈ RH×W×3划分为大小为4 × 4的非重叠块,随后将图像的维度映射到C,C默认为96。这个过程的结果是嵌入图像x′ ∈ RH4 × W4 ×C。最后,我们使用Layer Normalization [7]对x '进行归一化,然后将其送入编码器进行特征提取。编码器由四个阶段组成,在前三个阶段的末尾应用patch merging操作,以减少输入特征的高度和宽度,同时增加通道的数量。我们在四个级上采用[2,2,2,2] VSS块,每个级的通道计数为[C,2C,4C,8C]。
类似地,解码器被组织成四个阶段。在最后三个阶段的开始,一个补丁扩展操作被用来减少特征通道的数量和增加的高度和宽度。在四个阶段中,我们使用[2,2,2,1] VSS块,每个阶段的通道计数为[8C,4C,2C,C]。在解码器之后,使用最终投影层来恢复特征的大小以匹配分割目标。具体而言,通过补丁扩展进行4倍上采样以恢复特征的高度和宽度,然后通过投影层恢复通道数量。对于跳过连接,采用简单的加法操作,没有花哨的东西,因此不会引入任何额外的参数。
5、VSS Block
从VMamaba [20]派生的VSS块是VM-UNet的核心模块,如图1(B)所示。在经历层规格化之后,输入被分成两个分支。在第一个分支中,输入通过线性层,之后是激活函数。在第二分支中,输入在被馈送到2D选择性扫描(SS 2D)模块中用于进一步特征提取之前经历通过线性层、深度可分离卷积和激活函数的处理。随后,使用层规范化对特征进行规范化,然后使用第一个分支的输出执行逐元素生成,以合并两个路径。最后,使用线性层混合这些特征,并将此结果与剩余连接组合以形成VSS块的输出。在本文中,默认使用SiLU 作为激活函数。
SS2D由三个部分组成:扫描扩展操作、S6块和扫描合并操作。如图2(a)所示,扫描扩展操作将输入图像沿着四个不同方向(左上至右下、右下至左上、右上至左下和左下至右上)展开为序列。然后,这些序列由S6模块进行特征提取处理,确保来自各个方向的信息被彻底扫描,从而捕获不同的特征。随后,如图2(B)所示,扫描合并操作将来自四个方向,将输出图像恢复为与输入相同的大小。从Mamba [16]衍生的S6块通过基于输入调整SSM的参数在S4 [17]之上引入了选择性机制。这使得模型能够区分和保留相关信息,同时过滤掉不相关的信息。S6块的伪码在算法1中给出。
6、损失函数
VM-UNet的引入旨在验证纯SSM模型在医学图像分割任务中的应用潜力。因此,我们专门利用最基本的二进制交叉熵和骰子损失(BceDice损失)以及交叉熵和骰子损失(CeDice损失)分别作为二进制和多类分割任务的损失函数,如等式5和6所示。
表示样本总数,C表示类别总数。
yi, ˆyi 分别表示真实标签和预测。
yi,c 是指示符,如果样本i属于类别c,则该指示符等于1,否则等于0。
ˆyi,c 是模型预测样本i属于类别c的概率。
|X|和 |Y| 分别表示the ground truth 和 prediction。
λ1, λ2 是损失函数的权重,默认为1。
7、结论与展望
结论:在本文中,我们第一次介绍了一个纯粹的基于SSM的医学图像分割模型,提出VM-UNet作为基线。为了充分利用基于SSM的模型的功能,我们使用VSS块构造VM-UNet,并使用预训练的VMamba-S初始化其权重。在皮肤病变和多器官分割数据集上的综合实验表明,单纯的SSM模型在医学图像分割中具有很强的竞争力,值得进一步深入研究。
未来的工作:1)设计更适合分割任务的模块,基于SSM的机制。2)VM-UNet的参数计数约为30 M,可通过手动设计或其他压缩策略简化SSM,从而增强其在真实医疗场景中的适用性。3)考虑到SSM在捕获长序列信息方面的优势,进一步研究在更高分辨率下的分割性能将是有价值的。4)探索SSM在其他医学成像任务中的应用,如检测、配准和重建等。