论文地址：https://arxiv.org/pdf/2506.15976
代码地址：https://github.com/CiaoHe/bi-mamba

---

关注UP CV缝合怪，分享最计算机视觉新即插即用模块，并提供配套的论文资料与代码。
https://space.bilibili.com/473764881

摘要

Mamba 是一种状态空间模型 (SSM)，它通过将递归重铸为并行选择性扫描来加速训练，最近已成为一种线性缩放、高效的自注意力替代方案。由于其单向性，Mamba 中的每个状态都只包含其先前状态的信息，而对之后的状态视而不见。当前，基于 Mamba 的计算机视觉方法通常通过在 Mamba 的全局前向扫描基础上增加全局后向扫描来克服这一限制，形成双向扫描以恢复全部感受野。然而，此操作会使计算负荷加倍，从而削弱 Mamba 原本的效率优势。

为了消除这些额外的扫描，本研究引入了局部双向 SSM 模块 LBMamba，该模块在前向选择性扫描内嵌入轻量级局部后向扫描，并在每线程寄存器中完全执行。在本研究在 LBMamba 的基础上提出了可扩展的视觉骨干网络 LBVim，它每两层交替扫描方向，无需额外的后向扫描即可恢复全局感受野。本研究在自然图像和全幻灯片图像 (WSI) 上验证了该方法的多功能性。

结果表明，LBVim 始终提供优越的性能-吞吐量权衡。也就是说，在相同的吞吐量下，LBVim 在 ImageNet-1K 分类数据集上的 top-1 准确率提高了 0.8% 到 1.6%，在 ADE20K 语义分割数据集上的 mIoU 提高了 0.6% 到 2.7%，在 COCO 检测数据集上的 APb 和 APm 分别提高了 0.9% 和 1.1%。本研究还将 LBMamba 集成到最先进的病理学多实例学习 (MIL) 方法 MambaMIL 中，该方法使用单向扫描。在 3 个公开的 WSI 分类数据集上进行的实验表明，本研究的方法实现了高达 3.06% 的 AUC、3.39% 的 F1 和 1.67% 的准确率的相对改进。

引言

本研究关注基于Mamba的状态空间模型（SSM）在计算机视觉中的应用。SSM因其线性时间和内存复杂度而成为序列建模中自注意力机制的替代方案。然而，传统SSM的训练和推理速度仍然受到限制，因为它们无法充分利用现代GPU的并行性。Mamba通过将状态更新与隐藏到输出的卷积解耦，并将计算重新表述为在现代GPU上高效运行的并行选择性扫描来克服这个问题。因此，Mamba在长距离任务上达到了与Transformer相当的精度，同时展现出更好的分辨率缩放特性，使其成为研究和生产系统的理想选择。它最初是为自然语言处理引入的，后来被应用于计算机视觉。基于Mamba的选择性扫描内核构建的视觉模型可显著提高GPU速度和节省内存，同时始终优于基于Transformer的基线模型。

然而，标准的基于Mamba的计算机视觉模型通常需要从不同方向多次扫描图像才能提高性能。这种多次扫描有两个主要原因：第一，Mamba将图像视为扁平化的一维序列，因此单一的从左到右的扫描只能捕获行方向的上下文信息。为了恢复垂直方向的依赖关系，视觉流水线通常会对图像块的列方向进行额外的扫描，从而产生两次正交扫描，共同逼近二维空间关系。尽管一些专门的二维Mamba/SSM方法被提出以更直接地解决这种结构不匹配问题，但多次扫描的需求依然存在。第二个原因是SSM的单向性: 位置t的潜在状态仅以先前位置1到t为条件。因此，模型无法获取位置t之后的信息，这通常会导致视觉任务的性能欠佳。

目前常见的解决方案是添加反向扫描（从右到左或从下到上）以恢复对未来标记的访问，从而产生双向扫描机制。虽然这种策略重建了完整的感受野，但每次额外的扫描都会大致使计算量加倍，从而削弱了Mamba最初提供的效率优势。为了解决全局双向扫描带来的额外计算开销问题，本研究提出了LBMamba，一种局部双向SSM，以及用于视觉任务的LBVim框架。LBMamba在正向选择性扫描内部嵌入了轻量级的局部反向扫描，并在每个线程的寄存器中完全执行，从而无需进行全局反向扫描。基于LBMamba构建的LBVim视觉主干网络每两层交替扫描方向，无需额外的反向扫描即可恢复全局感受野。

论文创新点

本研究提出了一个名为LBMamba的局部双向SSM模块，旨在提高Mamba模型在计算机视觉任务中的效率和性能。其创新点如下：

1. ⏪ 引入了局部反向扫描机制： ⏪

   • 不同于传统的全局双向扫描，LBMamba在Mamba的前向选择性扫描过程中嵌入了轻量级的局部反向扫描。
   • 这种局部反向扫描在子序列内进行，无需进行全局反向扫描，从而显著提高了计算效率。

2. ⚡️ 基于硬件感知的线程级双向扫描算子： ⚡️

   • 本研究设计了一个快速的、硬件感知的线程级双向扫描算子。
   • 该算子完全在线程私有寄存器中执行局部反向扫描，避免了额外的内存读取和线程间通信，最大限度地减少了计算开销。

3. 🚀 消除了全局反向扫描的额外开销： 🚀

   • 通过将局部反向扫描整合到前向扫描过程中，LBMamba消除了全局反向扫描带来的额外计算负担和内存访问，从而显著提高了模型的吞吐量。
   • 实验结果表明，LBMamba的运行时间仅增加了2%，却避免了全局反向扫描带来的运行时间翻倍。

4. 🔄 基于序列反转操作恢复全局感受野： 🔄

   • 为了弥补局部反向扫描在全局感受野方面的不足，本研究提出在每个编码器层后反转特征序列。
   • 这种交替的扫描方向策略，使得每个token在每两个编码器层后都能获得全局感受野，从而增强了模型对长距离上下文信息的建模能力。

5. ✅ 在不同视觉任务和数据集上的有效性验证： ✅

   • 本研究在自然图像和病理图像等不同类型的图像数据上进行了广泛的实验，涵盖了图像分类、语义分割、目标检测和多实例学习等多种计算机视觉任务。
   • 实验结果表明，LBMamba在多种任务和数据集上均取得了显著的性能提升，证明了其有效性和泛化能力。
   • 例如，在ImageNet-1K分类数据集上，相同吞吐量下，LBVim的top-1准确率比基线模型提高了0.8%到1.6%。在ADE20K语义分割数据集上，mIoU提高了0.6%到2.7%。

论文实验