GroupMamba: Efficient Group-Based Visual State Space Model

📄 arXiv: 2407.13772v2 📥 PDF

作者: Abdelrahman Shaker, Syed Talal Wasim, Salman Khan, Juergen Gall, Fahad Shahbaz Khan

分类: cs.CV

发布日期: 2024-07-18 (更新: 2025-03-28)

备注: Accepted at CVPR-2025

🔗 代码/项目: GITHUB

💡 一句话要点

提出GroupMamba，一种高效的基于分组的视觉状态空间模型，提升图像识别性能。

🎯 匹配领域: 支柱二：RL算法与架构 (RL & Architecture)

关键词: 状态空间模型 图像分类 目标检测 语义分割 蒸馏训练 分组卷积 通道调制 长程依赖

📋 核心要点

  1. 纯粹基于SSM的模型在计算机视觉任务中面临稳定性和性能挑战，尤其是在模型规模较大时。
  2. 论文提出调制分组Mamba层，通过分组处理和通道调制，提升模型效率和跨通道信息交流。
  3. 引入蒸馏训练方法稳定大型模型训练，实验表明该方法在多个视觉任务上取得了显著性能提升。

📝 摘要（中文）

本文提出了一种参数高效的调制分组Mamba层，旨在解决基于状态空间模型（SSM）的视觉模型在扩展性和稳定性方面的挑战。该层将输入通道分为四个组，并对每个组独立应用提出的基于SSM的高效视觉单选择扫描（VSSS）块，每个VSSS块扫描四个空间方向之一。调制分组Mamba层还将四个VSSS块封装到通道调制算子中，以改善跨通道通信。此外，引入了一种基于蒸馏的训练目标来稳定大型模型的训练，从而实现持续的性能提升。综合实验表明，所提出的方法在ImageNet-1K图像分类、MS-COCO目标检测和实例分割以及ADE20K语义分割任务上均优于现有方法。仅有23M参数的微型变体在ImageNet-1K上实现了83.3%的top-1准确率，参数效率比现有最佳同等规模的Mamba设计高26%。

🔬 方法详解

问题定义：现有基于状态空间模型（SSM）的视觉模型在扩展到大型模型时，面临训练不稳定和效率低下的问题。尤其是在计算机视觉任务中，如何有效地利用SSM捕获长程依赖关系，同时保持模型的稳定性和参数效率，是一个关键挑战。

核心思路：论文的核心思路是将输入通道分组，并对每个组独立应用高效的视觉单选择扫描（VSSS）块。通过分组处理，降低了计算复杂度，提高了模型效率。同时，引入通道调制算子，增强了跨通道的信息交流，提升了模型性能。此外，采用蒸馏训练方法，稳定了大型模型的训练过程。

技术框架：GroupMamba模型的核心是调制分组Mamba层。该层首先将输入通道分为四个组，然后对每个组应用一个VSSS块，每个VSSS块扫描四个空间方向之一。四个VSSS块的输出被送入通道调制算子，以增强跨通道通信。整个模型通过蒸馏训练进行优化，以提高稳定性和性能。

关键创新：论文的关键创新在于调制分组Mamba层和蒸馏训练方法。调制分组Mamba层通过分组处理和通道调制，实现了参数效率和性能的平衡。蒸馏训练方法则有效地稳定了大型模型的训练过程，避免了性能下降。

关键设计：VSSS块是基于SSM的扫描模块，其具体实现细节（例如状态转移矩阵、输入矩阵、输出矩阵等）对性能有重要影响。通道调制算子的具体形式（例如线性层、非线性激活函数等）也会影响跨通道通信的效果。蒸馏训练的目标函数设计，例如使用软标签或硬标签，以及蒸馏温度的设置，都会影响模型的最终性能。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

GroupMamba在ImageNet-1K图像分类、MS-COCO目标检测和实例分割以及ADE20K语义分割任务上均取得了优于现有方法的性能。特别是，仅有23M参数的微型变体在ImageNet-1K上实现了83.3%的top-1准确率，参数效率比现有最佳同等规模的Mamba设计高26%。这些结果表明GroupMamba在效率和性能方面具有显著优势。

🎯 应用场景

GroupMamba具有广泛的应用前景，包括图像分类、目标检测、语义分割等计算机视觉任务。其高效的架构和稳定的训练方法使其能够应用于资源受限的设备，例如移动设备和嵌入式系统。此外，该模型还可以扩展到其他领域，例如视频理解和自然语言处理。

📄 摘要（原文）

State-space models (SSMs) have recently shown promise in capturing long-range dependencies with subquadratic computational complexity, making them attractive for various applications. However, purely SSM-based models face critical challenges related to stability and achieving state-of-the-art performance in computer vision tasks. Our paper addresses the challenges of scaling SSM-based models for computer vision, particularly the instability and inefficiency of large model sizes. We introduce a parameter-efficient modulated group mamba layer that divides the input channels into four groups and applies our proposed SSM-based efficient Visual Single Selective Scanning (VSSS) block independently to each group, with each VSSS block scanning in one of the four spatial directions. The Modulated Group Mamba layer also wraps the four VSSS blocks into a channel modulation operator to improve cross-channel communication. Furthermore, we introduce a distillation-based training objective to stabilize the training of large models, leading to consistent performance gains. Our comprehensive experiments demonstrate the merits of the proposed contributions, leading to superior performance over existing methods for image classification on ImageNet-1K, object detection, instance segmentation on MS-COCO, and semantic segmentation on ADE20K. Our tiny variant with 23M parameters achieves state-of-the-art performance with a classification top-1 accuracy of 83.3% on ImageNet-1K, while being 26% efficient in terms of parameters, compared to the best existing Mamba design of same model size. Code and models are available at: https://github.com/Amshaker/GroupMamba.