LoG-VMamba:局部-全局视觉Mamba医学图像分割
LoG-VMamba : Local-Global Vision Mamba for Medical Image Segmentation
Mamba, a State Space Model (SSM), has demonstrated competitive performance compared to Convolutional Neural Networks (CNNs) and Transformers in tasks involving Natural Language Processing and general sequence modeling. Several attempts have been made to adapt Mamba for Computer Vision applications, such as medical image segmentation (MIS). Vision Mamba (VM)-based networks are particularly attractive due to their ability to achieve global context windows, akin to Vision Transformers, while maintaining linear complexity relative to the number of tokens. However, despite the challenges posed by their sequential structure, existing VM models continue to struggle with preserving both local and global token dependencies in high-dimensional data. The computational cost associated with employing multiple or intricate scanning strategies often hinders the practical application of SSMs to high-dimensional 2D and 3D images encountered in MIS problems. In this work, we introduce Local-Global Vision Mamba (LoG-VMamba), a novel architecture that explicitly enforces spatial adjacency of tokens along the channel axis while maintaining a compressed form of global context. Our approach enables these models to access both local and global contextual information without waiting until the final token is processed. By implementing a simple scanning strategy, our segmentation models achieve computational efficiency while outperforming baseline methods based on CNNs and Transformers across diverse 2D and 3D medical imaging tasks.
Mamba被归类为一种状态空间模型（SSM），近年来，在自然语言处理领域以及一般的序列建模任务中展现了与其当前流行的卷积神经网络（CNNs）和Transformer架构不相上下的能力。
研究者们已经尝试将该方法应用于多种计算机视觉任务，并取得了显著成效。
基于 Vision Mamba（VM）的网络因其具备与 Vision Transformer 相似的全局感受野能力的同时又维持了标记数量上的线性复杂度而备受关注。然而现有 VM 模型受限于其固有的顺序特性 在处理高维数组时无法兼顾标记的空间局部性和整体关联性。为了提高效率依赖复杂的多模态融合机制来提高计算效率受限于其在处理高维图像时带来的计算开销
该文提出了一种名为局部-全局视觉Mamba（LoG-VMamba）的方法；通过将空间上相邻的标记放置于通道轴上的邻近位置来实现这一目标；从而得以以压缩的形式保持住全局语境信息；该方法使得SSM无需等到最后一个标记才能同时获取局部信息与整体背景知识；并且仅需采用一种简单的扫描机制即可实现这一目标。
该分割模型在计算效率方面表现出色，并且在2D与3D MIS相关任务中均远超基于CNN与Transformer的基准模型。
Introduction
......
涵盖Mis在内的计算机视觉(CV)领域中,视觉特征提取在该领域扮演着关键角色. 在深度学习初期阶段,卷积神经网络视为一种广泛采用的特征提取模块,因其能够捕捉像素局部区域及其上下文信息的能力强而备受推崇. 然而,在CNN架构中存在的一个不足是难以捕捉长距离依赖性(LRD),而这对于数据中的高级特征提取至关重要[29]. 为了增强感受野以捕捉长距离依赖性(LRD),一些研究采用了多层下采样的策略,另一种方法则是采用空洞卷积技术. 但这种方式带来了过高的计算成本.
与卷积神经网络（CNNs）不同的是，在视觉Transformer（ViT）架构中利用了注意力机制来理解和捕捉图像中的长距离依赖性（LRD）。尽管该架构在整个学习表示层上都实现了全局感受野（GRF），但它对处理与标记数量相关的二次复杂度存在局限性。在这种情况下，在ViT模型中这些标记是通过将图像划分为小块，并将其作为输入传递给网络前的步骤。因此，在涉及高维输入输出的任务中使用ViT可能会导致计算效率下降
近年来，在自然语言处理（NLP）领域中提出了一种名为Mamba的状态空间模型（SSM）。该模型通过其循环机制实现了对全局感受野的捕捉，并且由于计算复杂度与数据维度呈线性关系而保持了高效的计算性能。已有研究表明，将该方法应用于计算机视觉（CV）任务是可行的，并包括图像分类以及医学图像分割等任务。然而，在视觉领域的应用中，尽管在计算效率方面表现优异，但由于该模型状态容量受限以及长距离依赖建模需要采用顺序方法的原因，在维持相邻标记之间的依赖关系方面仍存在挑战。图1对比展示了基于Vision Mamba的方法在扫描协议方面的不足之处与卷积网络及ViT方法的差异。此外，部分研究者试图通过引入复杂的扫描策略来改善相邻标记的局部特性问题，但这种做法往往会导致计算复杂度显著提升。
本文主张：若存在适用于图像的标记器，则开发更为复杂的扫描策略并非必要之举。为此，本文提出了一种新型视觉Mamba算法——局部-全局视觉Mamba（LoG-VMamba）。该算法由两个关键组件构成：局部标记提取器（LTX）与全局标记提取器（GTX）。具体而言，在高维数组中通过保持相邻样本间的局部特性来实现目标；而基于通道轴的设计则能明确地保证空间上相邻的样本在通道轴上相互接近这一重要特性。相反地，在全局层面，则采用跨所有空间维度的方式对特征进行压缩处理；这一设计使得该算法能够在达到最终时间步之前为SSM模块提供一个能够在相同计算预算内实现更大范围感受野表现的优势版本；两者的结合不仅显著提升了算法性能；更重要的是使该方法得以突破传统框架对复杂扫描策略的需求限制
Local-Global Token Extractors
局部标记提取器（LTX）
基于Mamba作为一个序列建模模块的基础上[24,35,56,59]整合了一系列新的扫描策略来实现对二维标记数组的一维化映射。然而这些方案均存在一个共用缺点即无法有效维持相邻标记间的空间关系。此外对于视觉任务而言局部依赖性的保持具有重要意义。为了提升局部特征建模的效果同时规避[56]中那种全方位扫描的方式本文创新性地提出了LTX模块其直观展示如图2a所示。
通过深度可分离卷积（DWC）对输入通道C按照压缩因子S进行降维处理后，在激活函数作用下完成标记信息的提取。随后利用一个固定尺寸为R×R的卷积核对提取到的关键点信息进行还原处理。该方法能够实现对提取到的关键点信息进行复制，并有效保持相邻关键点之间的空间关系。最终对该局部区域内的关键点信息实施进一步优化，并将其转换为一维序列表示以便后续处理。这种空间展平过程是按行完成的，并且其本质等同于图1c所示的水平扫描方式。输出通道的数量为
。
全局标记提取器（GTX）
除了通过LTX实现的局部性之外，在此基础上我们提出了GTX模块（如图2b所示），它能够生成一种特殊的标记——这种标记具有全局特性但通道之间相互独立的空间属性。通过该模块设计，在早期的时间步中能够访问到压缩后的全局感受野信息。与基于VM的选择性扫描方法不同，在VM中只有最后一个时间步所提取的信息才具备与其他时间步信息相同的上下文关联度。具体而言，在输入特征图经过尺寸为H×W×C'的空间域上应用扩张深度可分离卷积操作后（其中卷积核大小为K×K），随后我们将得到的空间维度展平以获得形状为H'×W'×C''的结果
其中通道维度与空间维度进行了交换以实现全局标记的生成为了提高计算效率GTX将所有空间维度上的每组γ输入通道合并为一个全局标记这一过程旨在估计每个输入通道的空间位置信息并且发现其潜在的特征关系值得注意的是这些步骤的主要目标是学习每个输入通道中的全局上下文近似值而非引入细粒度的空间细节随后本文通过一个线性层将这些标记投影到C'维特征空间并采用SiLU激活函数进行激活以进一步强化特征表示
Local-Global Vision Mamba
本文扩展了VSS块，并将其与提出的LTX和GTX整合以采用VMamba的升级版（如图3所示）。因此，在VSS中由DWC层和SiLU组成的原始块被视为基于Mamba的标记提取模块。
局部视觉Mamba（L-VMamba）
对于局部视觉Mamba（L-VMamba），采用了引入的LTX块作为其标记提取器。与基于矢量稀疏表示的方法（VSS）相比，L-VMamba通过执行展开展操作来增强相邻标记在二维或三维数组中的空间邻近性保证。通常情况下，在卷积神经网络（CNN）架构中选择窗口尺寸为3是其典型设计选择之一。鉴于其通道数目已被调整为C'值，在SSM处理后添加了线性层来恢复原始C维空间的信息结构
全局视觉Mamba（G-VMamba）
基于全局视觉Mamba（G-VMamba）设计的模型中，将VSS中的原始块与GTX块结合，并包含DWC层、SiLU及后续展平操作。经由GTX模块输出带有全局感受野（GRF）的标记。鉴于输出通道数量维持不变但序列长度增加，在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） 在固定间隔内插入所有项并整理多余项至开头。