LBMamba通过局部双向扫描提升视觉模型性能,不增加计算开销。
LBMamba: Locally Bi-directional Mamba
- 在前向扫描中嵌入轻量级局部反向扫描,避免额外计算。
- 在ImageNet等数据集上提升0.8%~1.6%准确率,效率不变。
- 适合追求高精度且对计算成本敏感的视觉任务研究者。
Mamba作为状态空间模型(SSM),通过将循环运算重构成并行扫描实现线性扩展,成为自注意力机制的高效替代方案。然而其单向特性导致每个状态仅能获取历史信息,无法感知后续状态。现有基于Mamba的视觉方法通常通过全局反向扫描弥补,但使计算量翻倍,削弱了原本的效率优势。为此,本文提出LBMamba,一种局部双向SSM模块,将轻量级局部反向扫描嵌入前向扫描,并在每线程寄存器中执行。基于此构建了LBVim主干网络,每两层交替扫描方向,实现全局感受野而无需额外反向扫描。在自然图像与全切片图像(WSI)上验证表明,该方法持续提供更优的性能-吞吐权衡:相同吞吐下,LBVim在ImageNet-1K上提升0.8%~1.6%准确率,ADE20K上提升0.6%~2.7% mIoU,COCO检测上提升0.9% APb和1.1% APm。同时,该方法使四种SOTA Mamba模型(VMamba、LocalVim、PlainMamba、Adventurer)准确率提升0.5%~3.4%。将LBMamba集成至无向病理多实例学习模型MambaMIL,在三个公开WSI分类数据集上实现最高3.06%的AUC提升、3.39%的F1提升和1.67%的准确率提升。代码已开源。
原文摘要 · Abstract (English)
Mamba, a State Space Model (SSM) that accelerates training by recasting recurrence as a parallel scan, has recently emerged as a linearly-scaling alternative to self-attention. Because of its unidirectional nature, each state in Mamba only has information of its previous states and is blind to states after. Current Mamba-based computer-vision methods typically overcome this by augmenting Mamba's global forward scan with a global backward scan, forming a bi-directional scan to restore a full receptive field. However, this operation doubles the computational load, eroding much of the efficiency advantage that originally Mamba have. To eliminate this extra scans, we introduce LBMamba, a locally bi-directional SSM block that embeds a lightweight locally backward scan inside the forward scan and executes it in per-thread registers. Building on LBMamba, we present LBVim, a backbone that alternates scan directions every two layers to recover a global receptive field without extra backward sweeps. We validate our approach on both natural images and whole slide images (WSIs) and show that it constantly offers a superior performance-throughput trade-off. Under the same throughput, LBVim achieves 0.8% to 1.6% higher top-1 accuracy on the ImageNet-1K classification dataset, 0.6% to 2.7% higher mIoU on the ADE20K semantic segmentation dataset, 0.9% higher APb and 1.1% higher APm on the COCO detection dataset. Our method also boosts the accuracy of four SOTA Mamba models, namely VMamba, LocalVim, PlainMamba and Adventurer, by 0.5% to 3.4%. We integrate LBMamba into the SOTA pathology multiple instance learning (MIL) model, MambaMIL, which is unidirectional. Experiments on 3 public WSI classification datasets show that our method achieves a relative improvement of up to 3.06% better AUC, 3.39% better F1, 1.67% better accuracy. Our code is available at https://github.com/cvlab-stonybrook/LBMamba.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。