A Separable Self-attention Inspired by the State Space Model for Computer Vision

📄 arXiv: 2501.02040v2 📥 PDF

作者: Juntao Zhang, Shaogeng Liu, Kun Bian, You Zhou, Pei Zhang, Jianning Liu, Jun Zhou, Bingyan Liu

分类: cs.CV, cs.AI

发布日期: 2025-01-03 (更新: 2025-05-20)

🔗 代码/项目: GITHUB

💡 一句话要点

受状态空间模型启发，提出可分离自注意力机制，用于计算机视觉任务。

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

关键词: 可分离自注意力 状态空间模型 计算机视觉 图像分类 密集预测

📋 核心要点

  1. 现有Vision Transformer计算复杂度高，而状态空间模型在视觉任务中表现出潜力，但其与自注意力机制的联系有待探索。
  2. 论文提出一种新的可分离自注意力方法，借鉴Mamba的设计理念，构建了名为VMINet的新型网络架构。
  3. 实验表明，VMINet在图像分类和高分辨率密集预测任务上取得了与现有方法具有竞争力的结果。

📝 摘要（中文）

Mamba是一种高效的状态空间模型(SSM)，具有线性计算复杂度。尽管SSM不适合处理非因果数据，但Vision Mamba (ViM)方法在图像分类和目标检测等任务中仍然表现出良好的性能。最近的研究表明，状态空间模型和注意力变体之间存在丰富的理论联系。我们提出了一种新的可分离自注意力方法，首次将Mamba的一些优秀设计理念引入到可分离自注意力中。为了确保与ViM进行公平的比较，我们引入了VMINet，这是一个简单而强大的原型架构，仅通过堆叠我们的新型注意力模块和最基本的下采样层构建。值得注意的是，VMINet与传统的Transformer架构有显著的不同。我们的实验表明，VMINet在图像分类和高分辨率密集预测任务上取得了有竞争力的结果。代码可在https://github.com/yws-wxs/VMINet获取。

🔬 方法详解

问题定义：现有Vision Transformer模型计算复杂度高，尤其是在处理高分辨率图像时。虽然Vision Mamba (ViM)等基于状态空间模型的方法在视觉任务中展现出潜力，但如何将SSM的优势融入到更广泛使用的自注意力机制中，仍然是一个待解决的问题。现有方法通常基于Transformer架构，计算量大，效率较低。

核心思路：论文的核心思路是将Mamba状态空间模型中的一些关键设计理念，例如选择机制和线性复杂度，引入到可分离自注意力机制中。通过这种方式，可以构建一种既能捕捉全局上下文信息，又能保持计算效率的视觉模型。

技术框架：VMINet架构主要由堆叠的可分离自注意力模块和下采样层构成。输入图像首先经过一系列下采样层进行特征提取，然后通过多个可分离自注意力模块进行特征增强。最后，根据具体任务，可以连接分类器或密集预测头。整个架构避免了传统Transformer的复杂结构。

关键创新：论文的关键创新在于首次将Mamba的设计理念引入到可分离自注意力机制中，从而构建了一种新型的视觉模型。这种方法不仅降低了计算复杂度，还保留了捕捉全局上下文信息的能力。与传统的Transformer架构相比，VMINet更加简洁高效。

关键设计：VMINet的关键设计包括：1) 使用可分离卷积来降低自注意力的计算复杂度；2) 引入Mamba中的选择机制，动态地选择重要的特征；3) 采用简单的下采样层进行特征提取，避免使用复杂的卷积模块；4) 通过堆叠多个可分离自注意力模块来增强特征表达能力。具体的参数设置和损失函数根据不同的任务进行调整。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

VMINet在图像分类和高分辨率密集预测任务上取得了有竞争力的结果。具体而言，在ImageNet图像分类任务上，VMINet达到了与现有ViM模型相当的精度，同时保持了较低的计算复杂度。在高分辨率密集预测任务上，VMINet也展现出了良好的性能，证明了其在高分辨率图像处理方面的潜力。

🎯 应用场景

该研究成果可应用于图像分类、目标检测、语义分割等多种计算机视觉任务。由于其计算效率高，尤其适用于资源受限的设备或需要处理高分辨率图像的场景，例如移动设备上的图像处理、自动驾驶中的环境感知等。未来，该方法有望进一步扩展到视频理解、三维视觉等领域。

📄 摘要（原文）

Mamba is an efficient State Space Model (SSM) with linear computational complexity. Although SSMs are not suitable for handling non-causal data, Vision Mamba (ViM) methods still demonstrate good performance in tasks such as image classification and object detection. Recent studies have shown that there is a rich theoretical connection between state space models and attention variants. We propose a novel separable self attention method, for the first time introducing some excellent design concepts of Mamba into separable self-attention. To ensure a fair comparison with ViMs, we introduce VMINet, a simple yet powerful prototype architecture, constructed solely by stacking our novel attention modules with the most basic down-sampling layers. Notably, VMINet differs significantly from the conventional Transformer architecture. Our experiments demonstrate that VMINet has achieved competitive results on image classification and high-resolution dense prediction tasks.Code is available at: https://github.com/yws-wxs/VMINet.