LoG-VMamba:局部-全局视觉Mamba医学图像分割

LoG-VMamba : Local-Global Vision Mamba for Medical Image Segmentation

2408.14415 (arxiv.org)

Mamba, a State Space Model (SSM), has recently shown competitive performance to Convolutional Neural Networks (CNNs) and Transformers in Natural Language Processing and general sequence modeling. Various attempts have been made to adapt Mamba to Computer Vision tasks, including medical image segmentation (MIS). Vision Mamba (VM)-based networks are particularly attractive due to their ability to achieve global receptive fields, similar to Vision Transformers, while also maintaining linear complexity in the number of tokens. However, the existing VM models still struggle to maintain both spatially local and global dependencies of tokens in high dimensional arrays due to their sequential nature. Employing multiple and/or complicated scanning strategies is computationally costly, which hinders applications of SSMs to high-dimensional 2D and 3D images that are common in MIS problems. In this work, we propose Local-Global Vision Mamba, LoG-VMamba, that explicitly enforces spatially adjacent tokens to remain nearby on the channel axis, and retains the global context in a compressed form. Our method allows the SSMs to access the local and global contexts even before reaching the last token while requiring only a simple scanning strategy. Our segmentation models are computationally efficient and substantially outperform both CNN and Transformers-based baselines on a diverse set of 2D and 3D MIS tasks.

Mamba，作为一种状态空间模型（SSM），最近在自然语言处理和一般序列建模方面展示了与卷积神经网络（CNNs）和Transformer相媲美的竞争性能。人们已经进行了各种尝试，将Mamba适应于计算机视觉任务，包括医学图像分割（MIS）。

基于Vision Mamba（VM）的网络因其能够像Vision Transformer一样实现全局感受野，同时保持标记数量上的线性复杂度，而显得尤为吸引人。然而，由于现有VM模型的顺序性质，它们在高维数组中仍难以同时保持标记的空间局部和全局依赖性。采用多个和/或复杂的扫描策略计算成本高昂，这阻碍了SSM在MIS问题中常见的高维2D和3D图像上的应用。

本文提出了局部-全局视觉Mamba（LoG-VMamba），它明确地将空间相邻的标记保持在通道轴上的附近位置，并以压缩形式保留全局上下文。本文的方法允许SSM在到达最后一个标记之前就能访问局部和全局上下文，同时只需要一个简单的扫描策略。

本文的分割模型在计算上非常高效，并在各种2D和3D MIS任务上显著优于基于CNN和Transformer的基线方法。

Introduction

......

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

与卷积神经网络（CNNs）相比，基于视觉Transformer（ViT）的架构依赖于注意力机制来学习整个图像中的长距离依赖性（LRD）。尽管在所有学习表示层上都具有全局感受野（GRF），但ViT在处理与标记数量相关的二次复杂度时受到了阻碍。在ViT的情况下，这些标记是将图像分割成小块后，输入网络之前的图像块。因此，在如语义分割等具有高维输入和输出的任务中使用ViT可能会产生计算效率问题。

最近，在自然语言处理（NLP）领域引入了一种基于状态空间模型（SSM）的方法——Mamba。Mamba通过其循环机制实现了全局感受野的获取，同时由于其复杂度与数据维度的线性关系而保持了计算效率。已有几项研究尝试将Mamba适应于计算机视觉（CV）任务，如图像分类和医学图像分割（MIS）。尽管具有计算上的优势，但在视觉应用中，由于Mamba状态的容量受限以及建模LRD需要顺序方法，Mamba在维持相邻和远距离标记之间的依赖性方面存在困难。图1展示了与卷积和ViT相比，基于Vision Mamba（VM）的方法中两种常见扫描协议[34,35,52]的弱点。一些研究试图通过使用复杂的扫描策略（如多方向或全方向扫描）来维持相邻标记的局部性，但这会显著增加计算复杂度。

本文认为如果有一个适合图像的标记器，那么开发更复杂的扫描策略是不必要的。因此，本文提出了局部-全局视觉Mamba（LoG-VMamba），其核心组件是局部标记提取器（LTX）和全局标记提取器（GTX）。LTX在高维数组中保持相邻标记的局部性，明确确保空间上相邻的标记在通道轴上保持接近。另一方面，GTX跨所有空间维度压缩特征，为SSM模块提供在达到最终时间步之前的全局感受野的压缩版本。这两个组件的综合效果使LoG-VMamba与先前的基于VM的方法显著不同，因为本文消除了对复杂扫描策略的需求。

Local-Global Token Extractors

局部标记提取器（LTX）

由于Mamba是一个序列建模模块，因此已经引入了多种扫描策略[24,35,56,59]来将二维的标记数组转换为一维序列。然而，如图1c和1d所示，它们的共同缺陷在于无法保持相邻标记之间的空间邻近性。同时，局部依赖性对于视觉任务来说非常重要。为了改进局部特征的建模，同时避免像[56]中那样采用全方向扫描，本文引入了LTX模块，其图形化表示如图2a所示。

首先使用深度可分离卷积（DWC）将输入通道C按因子S进行压缩，从而避免后续操作的计算开销。将压缩后的标记通过SiLU激活函数后，使用一个固定大小为R×R的卷积核来展开标记。这个操作允许复制标记并保留相邻标记的空间关系。最后，重新调整并展平空间维度，以形成局部标记的一维序列，确保局部窗口内的相邻标记位于查询标记的通道轴上。这种空间展平是按行进行的，概念上相当于图1c中的水平扫描。因此，输出通道的数量为 。

全局标记提取器（GTX）

除了LTX保持的局部性外，本文还提出了GTX模块（如图2b所示），以生成全局（即空间独立但通道相关）的标记。该模块允许SSM在早期时间步访问全局感受野的压缩版本。这种方法与先前VM中的选择性扫描不同，后者只有最后一个时间步的标记具有所有其他标记的上下文。具体来说，给定一个维度为H×W×C'的输入特征图，本文使用步长为K×K的扩张深度可分离卷积（DWC）对其进行空间压缩，然后展平其空间维度以生成形状为的全局标记，其中通道维度和空间维度被转置。为了计算效率，GTX仅将所有空间维度上的每组γ输入通道压缩为全局标记。这些步骤的作用是学习每个输入通道中的全局上下文近似值，因此与扩张DWC相关的细粒度细节丢失问题并不显著。随后，本文使用一个线性层将这些标记投影到C'维空间，并应用SiLU激活函数。选择C'作为特征数量，以便在后续步骤中将GTX的输出与LTX的C'通道空间标记进行拼接。

Local-Global Vision Mamba

本文推广了VSS块，并融入了提出的LTX和GTX，以引入VMamba的升级版，如图3所示。因此，VSS中由DWC层和SiLU组成的原始块被视为标记提取模块，从而引导开发了以下基于Mamba的模型。

局部视觉Mamba（L-VMamba）

对于局部视觉Mamba（L-VMamba），采用引入的LTX块作为标记提取器。与VSS相比，L-VMamba包含了展开操作，以有效保证二维或三维数组中相邻标记的空间邻近性。将窗口大小R设置为3，这是卷积层中常见的核大小。由于LTX中通道数变为了C'，因此在SSM之后插入了一个线性层，以将标记投影回原始的C维空间。

全局视觉Mamba（G-VMamba）

在全局视觉Mamba（G-VMamba）中，将VSS中的原始块（包括DWC层、SiLU和随后的展平操作）与GTX块相结合。在通过原始块处理后，特征图通过GTX模块生成具有全局感受野（GRF）的标记。两个模块产生的标记最终结合并传递给下一个SSM。由于输出中的通道数保持不变，而序列长度增加，因此SSM之后的线性层是不必要的。

局部-全局视觉Mamba（LoG-VMamba）

最终，将提出的LTX和GTX模块结合起来，创建了LoG-VMamba模块。这种组合利用了LTX的局部依赖性和GTX的全局感受野，从而发挥了两者的优势。具体来说，给定输入特征图x ∈ R^H×W×C，LoG-VMamba的标记提取器可以表示为：

其中，。与L-VMamba类似，SSM模块之后需要一个线性层，以将输出xLG​映射到C维空间。

由于Mamba的序列性和输入依赖性，如何将x^L与x^G结合起来以利用这两种类型的上下文并不简单。因此，评估了以下方法：

• 头部（Head）：将全局标记拼接在序列的开头。
• 中部（Middle）：将全局标记放置在序列的中间。
• 分割（Split）：将全局标记分为两半，并附加到序列的两端。
• 交错（Interleaved）：在序列的固定间隔中插入每个全局标记，并将多余的全局标记放在开头。