将注意力机制与状态空间模型融合,提升视觉识别性能。
A2Mamba: Attention-augmented State Space Models for Visual Recognition
- 设计多尺度注意力增强的状态空间模型,实现跨层动态交互
- ImageNet-1K上达86.1%准确率,优于所有同类架构
- 参数量少40%却在检测与分割任务中表现更优,适合高效部署
Transformer与Mamba最初用于自然语言处理,启发了视觉识别的骨干网络设计。近期研究将局部注意力Transformer与Mamba结合,以捕捉局部细节与全局上下文。然而,这些方法仅采用简单堆叠,缺乏两者间的交互机制,深层融合仍是开放问题。本文提出A2Mamba,一种新型Transformer-Mamba混合架构,引入多尺度注意力增强的状态空间模型(MASS),通过多尺度注意力图对状态空间模型的隐藏状态进行空间聚合,强化二维空间依赖性并提升动态建模能力。A2Mamba在多项视觉任务中超越现有卷积网络、Transformer及Mamba基线模型。例如,A2Mamba-L在ImageNet-1K上达到86.1%的顶级准确率;在语义分割中,A2Mamba-B相较CAFormer-S36提升2.5% mIoU且效率更高;在级联掩码R-CNN的检测与实例分割任务中,A2Mamba-S分别超过MambaVision-B 1.2%/0.9% AP^b/AP^m,同时参数量减少40%。代码已开源:https://github.com/LMMMEng/A2Mamba。
原文摘要 · Abstract (English)
Transformers and Mamba, initially invented for natural language processing, have inspired backbone architectures for visual recognition. Recent studies integrated Local Attention Transformers with Mamba to capture both local details and global contexts. Despite competitive performance, these methods are limited to simple stacking of Transformer and Mamba layers without any interaction mechanism between them. Thus, deep integration between Transformer and Mamba layers remains an open problem. We address this problem by proposing A2Mamba, a powerful Transformer-Mamba hybrid network architecture, featuring a new token mixer termed Multi-scale Attention-augmented State Space Model (MASS), where multi-scale attention maps are integrated into an attention-augmented SSM (A2SSM). A key step of A2SSM performs a variant of cross-attention by spatially aggregating the SSM's hidden states using the multi-scale attention maps, which enhances spatial dependencies pertaining to a two-dimensional space while improving the dynamic modeling capabilities of SSMs. Our A2Mamba outperforms all previous ConvNet-, Transformer-, and Mamba-based architectures in visual recognition tasks. For instance, A2Mamba-L achieves an impressive 86.1% top-1 accuracy on ImageNet-1K. In semantic segmentation, A2Mamba-B exceeds CAFormer-S36 by 2.5% in mIoU, while exhibiting higher efficiency. In object detection and instance segmentation with Cascade Mask R-CNN, A2Mamba-S surpasses MambaVision-B by 1.2%/0.9% in AP^b/AP^m, while having 40% less parameters. Code is publicly available at https://github.com/LMMMEng/A2Mamba.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。