PlainMamba: Improving Non-Hierarchical Mamba in Visual Recognition
作者: Chenhongyi Yang, Zehui Chen, Miguel Espinosa, Linus Ericsson, Zhenyu Wang, Jiaming Liu, Elliot J. Crowley
分类: cs.CV, cs.LG
发布日期: 2024-03-26 (更新: 2024-08-15)
备注: Accepted to BMVC 2024
🔗 代码/项目: GITHUB
💡 一句话要点
提出PlainMamba以提升视觉识别中的非层次化模型性能
🎯 匹配领域: 支柱二：RL算法与架构 (RL & Architecture)
关键词: 视觉识别 状态空间模型 特征学习 高分辨率图像 深度学习 计算效率 方向感知
📋 核心要点
- 现有的视觉识别模型在处理二维图像时，往往无法有效捕捉空间关系和特征连续性。
- PlainMamba通过连续的二维扫描和方向感知更新，增强了模型对图像特征的学习能力。
- 在多项视觉识别任务中，PlainMamba表现出优越的性能，尤其在高分辨率输入下计算效率显著提高。
📝 摘要（中文）
我们提出了PlainMamba：一种简单的非层次化状态空间模型（SSM），旨在用于一般视觉识别。近期的Mamba模型展示了SSMs在序列数据上的竞争力，并初步尝试将其应用于图像。本文进一步调整了Mamba的选择性扫描过程，以增强其从二维图像中学习特征的能力，具体通过（i）连续的二维扫描过程提高空间连续性，确保扫描序列中标记的相邻性，以及（ii）方向感知更新，使模型能够通过编码方向信息来识别标记的空间关系。我们的架构易于使用和扩展，由相同的PlainMamba模块堆叠而成，所有层的宽度保持不变。通过去除特殊标记的需求，架构进一步简化。我们在多种视觉识别任务上评估PlainMamba，取得了比以往非层次化模型更好的性能，并与层次化替代方案具有竞争力。特别是在需要高分辨率输入的任务中，PlainMamba在保持高性能的同时计算需求显著降低。
🔬 方法详解
问题定义：本文旨在解决现有视觉识别模型在处理二维图像时对空间关系和特征连续性捕捉不足的问题。现有方法在特征学习上存在局限性，难以充分利用图像的空间信息。
核心思路：PlainMamba的核心思路是通过引入连续的二维扫描过程和方向感知更新机制，提升模型对图像特征的学习能力。这种设计旨在确保标记的相邻性和空间关系的有效识别。
技术框架：PlainMamba的整体架构由多个相同的PlainMamba模块堆叠而成，形成一个宽度恒定的模型。该架构简化了设计，去除了对特殊标记的需求，使得模型易于使用和扩展。
关键创新：PlainMamba的主要创新在于其选择性扫描过程的适应性调整，使其能够在视觉领域中有效学习特征。这一方法与传统的层次化模型相比，提供了更灵活的特征学习方式。
关键设计：在参数设置上，PlainMamba采用了固定宽度的模块堆叠结构，损失函数和网络结构经过优化，以确保在高分辨率输入下的计算效率和性能表现。
🖼️ 关键图片
📊 实验亮点
在多项视觉识别任务中，PlainMamba相较于以往的非层次化模型实现了显著的性能提升，尤其在高分辨率输入下，计算需求降低了许多，同时保持了高水平的识别准确率。这表明PlainMamba在视觉识别领域的有效性和竞争力。
🎯 应用场景
PlainMamba在视觉识别领域具有广泛的应用潜力，尤其适用于需要处理高分辨率图像的任务，如医学影像分析、自动驾驶视觉系统和安防监控等。其高效的计算性能和优越的识别能力将推动相关领域的技术进步和实际应用。
📄 摘要（原文）
We present PlainMamba: a simple non-hierarchical state space model (SSM) designed for general visual recognition. The recent Mamba model has shown how SSMs can be highly competitive with other architectures on sequential data and initial attempts have been made to apply it to images. In this paper, we further adapt the selective scanning process of Mamba to the visual domain, enhancing its ability to learn features from two-dimensional images by (i) a continuous 2D scanning process that improves spatial continuity by ensuring adjacency of tokens in the scanning sequence, and (ii) direction-aware updating which enables the model to discern the spatial relations of tokens by encoding directional information. Our architecture is designed to be easy to use and easy to scale, formed by stacking identical PlainMamba blocks, resulting in a model with constant width throughout all layers. The architecture is further simplified by removing the need for special tokens. We evaluate PlainMamba on a variety of visual recognition tasks, achieving performance gains over previous non-hierarchical models and is competitive with hierarchical alternatives. For tasks requiring high-resolution inputs, in particular, PlainMamba requires much less computing while maintaining high performance. Code and models are available at: https://github.com/ChenhongyiYang/PlainMamba .