LocalMamba: Visual State Space Model with Windowed Selective Scan
作者: Tao Huang, Xiaohuan Pei, Shan You, Fei Wang, Chen Qian, Chang Xu
分类: cs.CV, cs.AI
发布日期: 2024-03-14
🔗 代码/项目: GITHUB
💡 一句话要点
提出LocalMamba以优化视觉状态空间模型的扫描策略
🎯 匹配领域: 支柱二：RL算法与架构 (RL & Architecture)
关键词: 视觉状态空间模型 局部扫描策略 动态选择 图像分类 目标检测 图像分割 计算机视觉
📋 核心要点
- 现有的视觉Mamba方法在处理图像时未能有效保留局部2D依赖，导致性能未能超越传统CNN和ViT。
- 本文提出了一种局部扫描策略，通过将图像划分为窗口来捕捉局部依赖，并引入动态搜索方法优化每层的扫描选择。
- 实验结果显示，本文模型在ImageNet上相较于Vim-Ti提升了3.1%的准确率，且计算复杂度相同。
📝 摘要（中文）
近年来，状态空间模型（如Mamba）在长序列建模方面取得了显著进展，但在视觉任务中的应用尚未超越传统卷积神经网络（CNN）和视觉变换器（ViT）。本文认为，提升视觉Mamba（ViM）性能的关键在于优化序列建模的扫描方向。传统的ViM方法通过扁平化空间标记，忽视了局部2D依赖关系，导致相邻标记之间的距离过长。我们提出了一种新的局部扫描策略，将图像划分为不同窗口，有效捕捉局部依赖，同时保持全局视角。此外，考虑到不同网络层对扫描模式的偏好差异，我们提出了一种动态方法，独立搜索每层的最佳扫描选择，显著提升了性能。大量实验表明，我们的方法在有效捕捉图像表示方面优于现有方法。
🔬 方法详解
问题定义：本文旨在解决传统视觉Mamba方法在图像处理时未能有效捕捉局部依赖的问题，导致性能未能超越CNN和ViT。
核心思路：我们提出了一种局部扫描策略，将图像划分为多个窗口，以更好地捕捉局部依赖，同时引入动态搜索方法为每层网络独立优化扫描选择。
技术框架：整体架构包括局部扫描模块和动态选择模块。局部扫描模块负责将图像分割为窗口并进行局部特征提取，动态选择模块则根据每层的特征需求优化扫描模式。
关键创新：本研究的创新点在于引入局部扫描策略和动态搜索方法，使得每层网络能够根据特定需求选择最优的扫描模式，从而有效提升模型性能。
关键设计：在参数设置上，我们对窗口大小和扫描步幅进行了优化，损失函数采用了交叉熵损失，网络结构则基于现有的ViM架构进行改进，确保在相同计算复杂度下实现性能提升。
🖼️ 关键图片
📊 实验亮点
在实验中，LocalMamba模型在ImageNet数据集上相较于Vim-Ti提升了3.1%的准确率，且计算复杂度保持在1.5G FLOPs，显示出其在图像表示捕捉方面的显著优势。
🎯 应用场景
该研究的潜在应用领域包括计算机视觉中的图像分类、目标检测和图像分割等任务。通过优化视觉状态空间模型的扫描策略，LocalMamba能够在保持高效计算的同时，提升图像表示的准确性，具有广泛的实际价值和未来影响。
📄 摘要（原文）
Recent advancements in state space models, notably Mamba, have demonstrated significant progress in modeling long sequences for tasks like language understanding. Yet, their application in vision tasks has not markedly surpassed the performance of traditional Convolutional Neural Networks (CNNs) and Vision Transformers (ViTs). This paper posits that the key to enhancing Vision Mamba (ViM) lies in optimizing scan directions for sequence modeling. Traditional ViM approaches, which flatten spatial tokens, overlook the preservation of local 2D dependencies, thereby elongating the distance between adjacent tokens. We introduce a novel local scanning strategy that divides images into distinct windows, effectively capturing local dependencies while maintaining a global perspective. Additionally, acknowledging the varying preferences for scan patterns across different network layers, we propose a dynamic method to independently search for the optimal scan choices for each layer, substantially improving performance. Extensive experiments across both plain and hierarchical models underscore our approach's superiority in effectively capturing image representations. For example, our model significantly outperforms Vim-Ti by 3.1% on ImageNet with the same 1.5G FLOPs. Code is available at: https://github.com/hunto/LocalMamba.