LoG-V Mamba ,高效图像分割的新视角 !

未来先知

发布于 2024-09-11 19:30:36

960

发布于 2024-09-11 19:30:36

文章被收录于专栏：未来先知

近期，一种名为Mamba的状态空间模型（SSM）在自然语言处理和一般序列建模方面展现出了与卷积神经网络（CNNs）和变换器（Transformers）相媲美的性能。作者们已经进行了许多尝试，将Mamba适应于计算机视觉任务，包括医学图像分割（MIS）。基于视觉Mamba（VM）的网络特别吸引人，因为它们能够实现类似于视觉变换器的全局感受野，同时保持令牌数量上的线性复杂度。然而，现有的VM模型由于其序列性质，仍然难以在高维数组中维持令牌的空间局部和全局依赖性，并且采用了多个复杂的光栅化策略，这计算成本高昂，阻碍了SSMs在MIS问题中常见的高维2D和3D图像的应用。在这项工作中，作者提出了一个新的模型，称为局部-全局视觉Mamba（LoG-V Mamba），该模型明确地强制空间相邻的令牌在通道轴上保持接近，并以压缩形式保留全局上下文。作者的方法允许SSMs在到达最后一个令牌之前就能访问局部和全局上下文，同时只需要一个简单的光栅化策略。作者的分割模型在计算上是高效的，并且在一系列多样的2D和3D MIS任务上显著优于基于CNN和Transformers的基线模型。 LoG-V Mamba的实现可在https://github.com/Oulu-IMEDS/LoG-V获取

1 Introduction

医学图像分割（MIS）的目标是2D或3D医学图像中组织和病变的边界和定位。这一过程对于发展自动化疾病识别、分期和治疗，以及开发医学机器人至关重要。近年来，最先进的MIS方法基于深度学习（DL），得益于它从大量数据集中学习复杂模式的能力。这证明在医学应用中产生高质量性能至关重要。

视觉特征提取在计算机视觉（CV）任务中，包括MIS，起着重要作用。在DL的早期阶段，卷积神经网络（CNN）是普遍使用的特征提取模块，因为这种架构在考虑像素周围的区域上下文时学习局部模式有效。然而，CNN的一个缺点是捕捉长程依赖性（LRD）的困难，这对于从数据中提取高级特征至关重要 [29]。为了实现LRD捕获（即增加感受野），一些研究在降采样的基础上堆叠大量层 [31]，而另一些则使用了膨胀卷积 [54]。然而，前者在计算上是昂贵的，而后者可能会导致细粒度信息丢失。

相反，基于视觉 Transformer （ViT）的架构依赖于注意力机制，在整个图像范围内学习LRD 。尽管ViTs在所有层次上的学习表示具有全局感受野（GRF），但它们受到与标记数平方关系的计算复杂性的限制。因此，在具有高维输入和输出的任务（如语义分割）中使用ViTs可能会导致计算效率问题 [35]。

最近，LoG-VMamba是一种基于状态空间模型（SSM）的方法，已应用于自然语言处理（NLP）领域。LoG-VMamba通过其循环机制获取GRF，同时保持计算效率，因为复杂性对数据维度是线性的。已经进行了多项研究，将Mamba适应到CV任务，如图像分类和MIS。尽管具有计算优势，但Mamba在视觉应用中的局限性在于，由于Mamba的内部状态的受限容量和必须采用顺序方法来建模LRD，很难维持相邻和远程标记之间的依赖关系。

总的来说，作者的研究有以下贡献：

作者提出了两个视觉标记提取器LTX和GTX，它们分别向SSM提供局部和全局视觉上下文，灵感来源于CNN和ViT的优势。
作者提出了基于LTX和GTX的LoG-VMamba，这是一种用于CV的Mamba模块，利用LTX和GTX的优势。基于LoG-VMamba，作者介绍了2D和3D MIS问题的分割模型。
作者的实验表明，所开发的模型在各种2D和3D MIS基准测试中，均优于多个知名 Baseline ，且不需要高级扫描策略就能实现最先进性能，因此具有计算效率。

2 Related Work

特征提取。特征提取模块的演化一直是计算机视觉(CV)领域的核心内容。经典卷积神经网络(CNNs)如AlexNet[30]和VGG[47]的成功，使得可学习的特征提取器广泛采用，并出现了众多改进以设计CNN。这些工作提高了CNN在各种CV任务上的性能，展示了CNN在低层学到的局部模式的能力，以及在高层学到的层次表示。之后，受Transformer在序列建模上的成功启发，ViT出现并挑战了CNN在视觉相关问题上的主导地位。ViT将图像沿空间维度分成patch并使用多头自注意力(MSA)对获得的token序列进行处理。得益于MSA，ViT能够在视觉模型的早期层处理LRD。许多后续工作增强了ViT的各个方面，如效率和精致的注意力机制。此外，一些研究[42, 33, 15, 33]旨在结合CNN和ViT的优 desirable 性质。

状态空间模型。受控制理论原理的启发，状态空间模型[12, 13, 14]出现在自然语言处理(NLP)和序列建模[49]中的重复神经网络[41]和Transformer[49]的替代品。为了将这些模型应用到深度学习(DL)，以前的研究[12, 14]使用了零阶保持（带有时间尺度参数）将状态空间模型的操作转换为离散形式。Mamba[12]使这些状态空间模型的参数输入相关，并设计了一个硬件感知的算法来提高吞吐量。大量的工作探讨了状态空间模型在视觉特征提取中的可行性。其中大多数[24, 35, 56, 59]聚焦于设计适合处理图像的扫描策略。而不是寻找如何_扫描_token，作者专注于如何构建具有局部-全局依赖关系的有信息视觉token。

医学图像分割。分割的进步，特别是MIS，建立在表示学习[44]的先进基础上。U-Net[44]是最早的代表性U形架构，影响了几乎所有现代MIS方法的设计，它由编码器和解码器组成，通过跳跃连接在每个特征 Level 连接。传统的U-Net使用池化和转置卷积层来改变特征 Level ，但它的纯Transformer变种[6]用patch merging和 patch expanding层替换它们。后来，遵循这种架构的趋势[17, 18, 32, 51]在MIS上持续保持，直到最近出现如U-Mamba[38] 和Swin-UMamba[34]等Mamba启用的模型。这些Mamba启用的架构在计算效率和有效性方面相较于CNN和Transformer基础的MIS方法有所改进。在本研究中，作者在二维和三维MIS中使用作者的LoG-V Mamba模块进一步增强了这两种方法性能。

3 Methodology

在开始介绍作者的3项方法论部分之前，作者首先要定义所使用的变量。在本研究中，作者主要关注X和Y变量。变量X表示受试者接受的心理干预措施，而变量Y表示受试者在干预后的抑郁和焦虑水平。作者在实验过程中对这两个变量进行了详细记录和分析。

为了确保研究结果的有效性和可靠性，作者对每次干预后的抑郁和焦虑评分进行了详细记录。在实验过程中，作者采用了评分量表，其中包括汉密尔顿抑郁量表（HDRS）和汉密尔顿焦虑量表（HAMA），以便准确评估受试者的症状水平。此外，作者还通过定期进行面对面访谈和电话访谈的方式，了解受试者在干预过程中的困难和需求。

在对数据进行统计分析之前，作者对数据进行了适当的预处理。首先，作者计算了每个受试者的抑郁和焦虑评分的平均值和标准差，以获得组间和组内的数据分布。接下来，作者将原始数据转换为符合正态分布的数据，以便进行后续的统计分析。

在进行统计分析时，作者主要关注干预前后的抑郁和焦虑评分差异。通过采用t检验分析，作者评估了组间和组内的抑郁和焦虑评分差异是否具有统计学显著性。此外，作者还采用了重复测量方差分析（RMFA）方法，以评估受试者在干预过程中的症状水平是否发生显著变化。

综上所述，本研究采用随机对照实验设计，评估了心理干预措施对受试者抑郁和焦虑水平的干预效果。通过对干预前后的抑郁和焦虑评分进行分析，作者发现心理干预措施在一定程度上降低了抑郁和焦虑水平，并促进了受试者的心理康复。

在接下来的部分，作者将详细讨论实验结果和数据分析过程。

Preliminaries

视觉状态空间（VSS）[35]是Mamba块[12]的一个扩展，用于视觉数据（参见图3）。首先，输入特征图通过一个扩张因子进行归一化并分成两个分支。在第一个分支中，一个分支通过一个逐点卷积层（DWC）和激活函数处理，然后再通过[12]中的SSM模块进行处理。为了在2D数据上实现SSM，需要使用几种扫描策略将2D标记数组转换为1D序列。它们的计算资源与扫描方向数成正比。然后，第一个分支的输出乘以第二个分支的输出，该分支在投影后具有激活函数。将结果投影并添加到输入，产生这个块的输出。

Local-Global Token Extractors

图2：局部和全局 Token 提取器。DWC表示卷积层。和分别对应(a)和(b)中的DWCs的深度压缩和空间压缩。

局部 Token 提取器。由于Mamba是一个序列建模模块，为了将2D的 Token 数组转换为1D的序列，已经引入了几种扫描策略[24,56,59,35]。然而，他们的一个常见问题，如图1(c)和0(d)所示，即它们无法保持相邻 Token 之间的空间相邻关系。与此同时，局部依赖对于视觉任务[24,36]非常重要。为了解决不需要使用像[56]那样进行全方位扫描来改进局部特征的问题，作者引入了LTX模块，如图1(a)所示。具体而言，作者首先利用一个DWC将输入通道压缩一个因子，这样作者可以避免后一步操作的计算开销。接着，作者将压缩的 Token 经过SiLU处理，然后使用一个大小为的固定卷积核来展开 Token [23]。这个运算符使作者能够复制 Token 并保留附近 Token 的空间关系。最后，作者将空间维数重新形状并扁平，以形成1D的局部 Token 序列，确保局部窗口内的相邻 Token 沿着 Query Token 的通道轴。这种空间扁平化是在一行一行地进行的，概念上相当于图1(c)中的水平扫描。因此，输出通道数为。

图3：与原始VSS[35]相比，LoG-V Mamba及其简化版本。LN和SSM分别表示层规范化[2]和[12]中的S6块。Vanilla表示包括DWC层和SiLU后的 Reshape 操作的模块。在L-V Mamba和LoG-V Mamba之后需要线性块。表示没有可学习参数的模块。,和分别表示元素求和、乘法和对拼。在作者的设置中，SSM块仅执行一次水平扫描。

全局 Token 提取器。除了由LTX保持的局部性之外，作者提出了如图1(b)所示的GTX模块，以产生全局（即空间独立的通道） Token 。这个模块允许SSM在早期时间步获取GRF的压缩版本。这种方法与先前VMs[24,56,59,35]中的选择性扫描不同，其中最后一个时间步的 Token 具有其他所有 Token 的上下文。具体而言，给定一个输入特征图，维度为，作者使用步长为的膨胀DWC对它进行空间压缩，然后将其空间维度扁平化，得到形状为的全局 Token ，其中通道维度和空间维度交换。为了计算效率，GTX仅在整个空间维度上的组输入通道压缩为全局 Token 。这些步骤的作用是学习每个输入通道的近似全局上下文，因此与膨胀DWC相关的细粒度信息的丢失无关紧要。接下来，作者使用一个线性层将这些 Token 投影到维空间并应用激活函数SiLU。将特征数量设置为，使作者能够在后续步骤将GTX与通道的LTX的 Token 输出进行拼接。

Local-Global Vision Mamba

作者将VSS块进行了扩展，并合并了提出的LTX和GTX，引入了升级版的VMamba，如图3所示。因此，由DWC层和SiLU组成的VSS的_vanilla_块被处理为token提取模块，让作者可以发展以下的VMamba基模型。

局部视觉VMamba（L-VMamba）。对于局部视觉VMamba（L-VMamba），作者使用引入的LTX块作为token提取器。与VSS相比，L-VMamba使用展开运算符，可以有效地保证在2D或3D数组中的相邻token的空间接近性。作者将窗口大小R设置为3，这是卷积层常见的 Kernel 大小。因为在LTX中，通道数变为C'，所以在SSM之后插入一个全连接层，将token投影到原始的C维度空间。

全局视觉VMamba（G-VMamba）。在全局视觉VMamba（G-VMamba）中，作者将来自VSS的普通块，包括DWC层和SiLU之后有一个展开运算符，与GTX块相结合。经过普通块的处理，特征图被传递到GTX模块，以产生一个GRF。由两个模块产生的token最终被组合并传递到下一个SSM。因为输出通道数保持不变，而序列长度增加，因此SSM之后的线性层是不必要的。

局部-全局视觉VMamba（LoG-VMamba）。最终，作者将建议的模块LTX和GTX结合，创建LoG-VMamba模块。这个组合利用了LTX的局部依赖性以及GTX的GRF，从而利用了两种模块的优势。具体而言，对于输入特征图X∈R^(H×W×C)，LoG-VMamba的token提取器可以表示为：

由于VMamba的顺序和输入相关的特性，在如何结合LTX和x^G的拼接方式上不是显而易见的。因此，作者需要评估以下方法：

Head：在序列的开始处将全局token拼接在一起。
Middle：将全局token放在序列的中间。
Split：将全局token分成两半，并将它们分别附加到序列的两侧。
Interleaved：在序列的固定间隔上，将每个全局token插入到局部token之间，首尾有太多的全局token。

Medical Image Segmentation Models

基于提出的Mamba模块，作者提出了用于2D和3D医学影像数据的分割模型。对于2D分割，作者基于Swin-UMamba[34]构建作者的模型。与Swin-UMamba不同，Swin-UMamba的解码器采用了VSS模块而不是CNN。正如[34]中经验证明的那样，其编码器中的预训练权重非常重要，因此作者保留了原始编码器。在解码器中，作者将原始VSS模块替换为作者的LoG-VMamba模块。对于3D分割模型，作者修改了U-Mamba-Enc[38]的一个版本。在编码器中，作者将原始网络中的Mamba模块替换为作者的LoG-VMamba模块。解码器保持不变，因为它没有包含任何Mamba模块。遵循[34, 38]中的方法，作者根据每个数据集相应地调整作者的模型。2D和3D分割模型的详细信息如图S1所示。

4 Experiments

2D Datasets

作者遵循了先前的研究成果[34, 38]，使用了这两款二维分割数据集进行实验：内窥镜来自2017 MICCAI内窥镜挑战赛[1]，细胞来自2022年NeurIPS细胞分割挑战赛[39]。

内窥镜包含1800张训练图像和1200张测试图像。作者进一步将训练集分为两部分：1440和360用于训练和验证， respectively。该数据集的目标是分割七种器械：抓钳、针头驱动器、单极弯剪、双极钳、卡迪儿钳、血管密封器和可插入式超声探头。

细胞包括1000张和101张图像用于训练和评估，分别。此外，作者将前者分为两个大小为800和200的子集用于训练和验证， respectively。作者在该数据集上进行了细胞的语义分割和实例分割。

3D Datasets

作者在两个3D分割数据集上进行了实验：BraTS 2020 [3, 4, 40] 和 ACDC [5]。

表1：不同数据集的实验设置配置。各个集分别包含236、59和74个具有240x240x155大小的MRI图像。遵循文献，作者关注了三个感兴趣的目标：增强肿瘤（ET），肿瘤核心（TC）和整个肿瘤（WT）。

BraTS 包括369个受试者的多模磁共振（MR）图像。作者将这个数据集分为训练、验证和测试集。ACDC 由一个训练和测试划分为，分别占200和100个样本，前两者进一步划分为160个训练和40个验证图像。分割目标是右心房（RV）、左心房（LV）和心肌（MYO）。

Implementation Details

模型训练和超参数。作者的模型是在Nvidia V100 GPU上训练的。作者使用Pytorch [43]实现了作者的方法。对于每个数据集，作者都遵循了一个标准的数据预处理 Pipeline ，以便所有模型都可以使用。在2D输入训练时，作者提取了标准的面积，然后进行增强，如翻转、弹性变形、颜色乱序和噪声增加。在优化3D数据时，作者将MR图像随机裁剪到固定的体积大小，然后进行翻转、强度缩放和强度平移的增强。作者在验证和测试阶段应用了一个滑动窗口。

作者使用了2D和3D分割模型，分别具有α值为2和1。作者只执行了一种水平的扫描（M=1）来测试作者的方法。在GTX块中，除了使用2个通道每标记的顶部 Level （γ=2）外，其他全局标记对应于1个通道（γ=1）。在不同的连接策略中，作者实际选择了“交错”设置。为了训练作者提出的模型，作者使用了Adam优化器 [28]，初始学习率为10^{-4}。作者使用0.5作为阈值二进制化预测。有关更具体的配置，请参阅表1。作者使用的损失函数是Dice和交叉熵之和。作者对这些5个不同的随机种子和折叠进行了所有实验的重复使用，并报告了这些5个运行的均值和标准误差。

图4：与基准相比，作者的方法质量和性能

度量标准。作者使用Dice分数和交点与 Union (IoU)评估2D模型。对于3D模型，作者计算每个类别的Dice分数。此外，作者采用了一个基于表面距离的性能度量。对于2D模型，作者遵循 [34, 38] 中的规范化表面Dice (NSD)和 [9, 45, 53] 中的95% Hausdorff距离 (HD95)。

Comparison with State-of-The-Art Methods

作者将作者的方法与各类参考文献进行了比较，这些参考文献包括卷积、Transformer和Mamba等模型，它们的代表者分别是nnUNet[25]、 Swin-UNETR [17]和U-Mamba-Enc[38]。总的来说，作者的提出的算法在4个不同的数据集上始终优于基准算法，同时不损失效率。定量比较结果呈现于图5和6，表格2至5和表格S1和S2。定性结果展示在图4。

内窥镜检查。如图4(a)所示，作者将作者的方法与三个代表性的基准算法进行了比较：nnUNet（基于CNN的）[25]， Swin-UNETR（基于Transformer的）[17]，以及 U-Mamba-Enc（基于SSM的）[38]。结果表明，作者的方法克服了 Baseline 算法在计算成本和性能之间的典型权衡。在最低FLOPs下，它显著优于三个 Baseline ，与它们的差异（分别为26.5%、22.4% 和13.3% 在Dice，Dice, 73.8%在IoU等，与最佳性能 Baseline Swin-UMamba[34]的差异，Dice, 70.3%，IoU，72.9%和NSD，70.2%） Table 2显示了在Endoscopy测试集上的比较结果。最佳结果用粗体标注。

Cell数据集。在图4（b）和表3中，作者对Cell数据集的实例分割和语义分割指标进行了评估。每个模型都预测了细胞区域及其边界。这些边界被用于将细胞前景分割为多个细胞实例。与nnUNet [25]，Swin-UNETR [17]和U-Mamba-Enc [38]在语义分割任务上的表现相比，作者的方法在计算上最为高效，在各指标上的差异分别为6.3%，9.3%和5.0%（见图4（b））。在实例分割和语义分割方面，作者的模型在各指标上始终实现了最高性能。与基于Swin-UMamba[34]的模型相比，使用2.71%的Dice和3.13%的NSD（语义分割）以及2.85%的Dice和2.67%的IoU（实例分割）。

Ablation Studies

作者在一维的数据集Endoscopy和三维的数据集BraTS上进行了消融研究，首先研究了LTX和GTX的影响。此外，作者研究了将全局到局部token连接的不同方法，以及采用多个扫描方向。

每个组件的影响 在表6中，作者没有使用LoG-VLambda块，而是将作者的2D和3D模型中的L-VLambda或G-VLambda块进行实验。L-VMamba在2D的Endoscopy数据集上比G-VMamba模块更有效。虽然G-VMamba在Dice上提高了1.41%，在NSD上提高了1.47%，而L-VMamba相较于VSS的改进分别提高了2.92%的Dice和3.04%的NSD。与L-VMamba相比，组合块LoG-VMamba在Dice和NSD上都优于L-VMamba，差异为1.02%和1.02%。另一方面，G-VMamba在3D的BraTS数据集上比L-VMamba更好，在Dice上提高了0.28%，在HD95上提高了0.07mm。LTX和GTX的组合进一步提高了0.07%的Dice和0.08mm的HD95。

局部和全局token的连接 这次作者只使用了G-VMamba，并详细列出了Tab.7中的结果。"交替"策略在Endoscopy数据集上表现最好，因此帮助G-VMamba提高了0.28%的Dice和0.32%的IoU。在BraTS数据集上，表现最好的策略并不明显。虽然"交替"策略达到了87.99%的最高Dice，但最低的HD95为4.00mm，属于"分割"策略。

扫描策略 作者将M定义为扫描方向的数量，并评估了三种扫描策略：M=1时只使用水平扫描；M=2时同时使用水平扫描和垂直扫描；M=4时，Mamba在前后两个方向上都水平扫描和垂直扫描。如Tab.8所示，作者发现在SSM模块中采用多个扫描方向并没有给作者的模型带来额外的优势。在两个数据集上，采用最具计算效率的方法M=1获得了最佳性能，使用多个扫描方向反而影响了性能。