Demystify Mamba in Vision: A Linear Attention Perspective
作者: Dongchen Han, Ziyi Wang, Zhuofan Xia, Yizeng Han, Yifan Pu, Chunjiang Ge, Jun Song, Shiji Song, Bo Zheng, Gao Huang
分类: cs.CV
发布日期: 2024-05-26 (更新: 2024-12-02)
备注: NeurIPS 2024
🔗 代码/项目: GITHUB
💡 一句话要点
揭示视觉Mamba的奥秘：线性注意力视角下的深度解析与改进
🎯 匹配领域: 支柱二：RL算法与架构 (RL & Architecture)
关键词: Mamba模型 线性注意力 Transformer 状态空间模型 视觉任务 高分辨率图像 模型优化
📋 核心要点
- 现有Transformer模型在处理高分辨率图像时计算复杂度高，而线性注意力Transformer性能通常不如传统Transformer。
- 论文将Mamba模型重新表述为线性注意力Transformer的变体，并分析了Mamba成功的关键设计因素。
- 实验结果表明，遗忘门和块设计是Mamba成功的关键，基于此提出的MILA模型在图像分类和高分辨率密集预测任务中优于现有Mamba模型。
📝 摘要（中文）
Mamba是一种高效的状态空间模型，具有线性计算复杂度。最近，它在处理各种视觉任务中的高分辨率输入时表现出令人印象深刻的效率。本文揭示了强大的Mamba模型与线性注意力Transformer之间存在惊人的相似之处，而线性注意力Transformer在实践中通常不如传统的Transformer。通过探索有效的Mamba和次优的线性注意力Transformer之间的异同，我们提供了全面的分析，以揭示Mamba成功的关键因素。具体来说，我们将选择性状态空间模型和线性注意力在一个统一的公式中重新表达，将Mamba重新定义为线性注意力Transformer的一种变体，具有六个主要区别：输入门、遗忘门、shortcut、无注意力归一化、单头和修改后的块设计。对于每个设计，我们都仔细分析了其优缺点，并从经验上评估了其对视觉任务中模型性能的影响。有趣的是，结果表明遗忘门和块设计是Mamba成功的核心贡献者，而其他四个设计则不太重要。基于这些发现，我们提出了一种Mamba启发式线性注意力（MILA）模型，通过将这两个关键设计的优点融入线性注意力中。由此产生的模型在图像分类和高分辨率密集预测任务中均优于各种视觉Mamba模型，同时具有可并行计算和快速推理速度。
🔬 方法详解
问题定义：Mamba模型在视觉任务中表现出色，尤其是在处理高分辨率图像时，但其成功的原因尚不完全清楚。线性注意力Transformer作为一种计算效率更高的替代方案，在实践中却往往不如传统Transformer。因此，本文旨在揭示Mamba成功的关键因素，并利用这些因素改进线性注意力模型。
核心思路：论文的核心思路是将Mamba模型与线性注意力Transformer进行对比分析，找出它们之间的异同，从而理解Mamba成功的关键设计。通过将Mamba重新表述为线性注意力Transformer的一种变体，可以更清晰地识别Mamba的关键组成部分，并评估它们对模型性能的影响。
技术框架：论文首先将选择性状态空间模型和线性注意力在一个统一的公式中重新表达。然后，将Mamba定义为线性注意力Transformer的一种变体，并识别出六个主要区别：输入门、遗忘门、shortcut、无注意力归一化、单头和修改后的块设计。接下来，对每个设计进行分析和评估，以确定其对模型性能的影响。最后，基于分析结果，提出了一种Mamba启发式线性注意力（MILA）模型。
关键创新：论文最重要的技术创新在于揭示了Mamba模型与线性注意力Transformer之间的联系，并确定了遗忘门和块设计是Mamba成功的关键因素。通过将这些关键设计融入线性注意力中，可以显著提高线性注意力模型的性能。
关键设计：论文的关键设计包括：1) 统一的公式，用于表达选择性状态空间模型和线性注意力；2) 六个主要区别的识别，用于区分Mamba和线性注意力Transformer；3) 遗忘门和块设计的分析和评估，用于确定其对模型性能的影响；4) MILA模型的提出，通过将遗忘门和块设计融入线性注意力中，提高了模型性能。
🖼️ 关键图片
📊 实验亮点
实验结果表明，遗忘门和块设计是Mamba成功的关键。基于此，论文提出的MILA模型在图像分类和高分辨率密集预测任务中均优于各种视觉Mamba模型。例如，在ImageNet分类任务上，MILA模型取得了与Mamba相当甚至更好的性能，同时保持了线性注意力的并行计算和快速推理速度。
🎯 应用场景
该研究成果可应用于各种需要处理高分辨率图像的视觉任务，例如图像分类、目标检测、语义分割等。通过利用Mamba的关键设计，可以提高线性注意力模型的性能，从而在计算资源有限的情况下实现更好的效果。此外，该研究还可以为设计更高效的视觉模型提供新的思路。
📄 摘要（原文）
Mamba is an effective state space model with linear computation complexity. It has recently shown impressive efficiency in dealing with high-resolution inputs across various vision tasks. In this paper, we reveal that the powerful Mamba model shares surprising similarities with linear attention Transformer, which typically underperform conventional Transformer in practice. By exploring the similarities and disparities between the effective Mamba and subpar linear attention Transformer, we provide comprehensive analyses to demystify the key factors behind Mamba's success. Specifically, we reformulate the selective state space model and linear attention within a unified formulation, rephrasing Mamba as a variant of linear attention Transformer with six major distinctions: input gate, forget gate, shortcut, no attention normalization, single-head, and modified block design. For each design, we meticulously analyze its pros and cons, and empirically evaluate its impact on model performance in vision tasks. Interestingly, the results highlight the forget gate and block design as the core contributors to Mamba's success, while the other four designs are less crucial. Based on these findings, we propose a Mamba-Inspired Linear Attention (MILA) model by incorporating the merits of these two key designs into linear attention. The resulting model outperforms various vision Mamba models in both image classification and high-resolution dense prediction tasks, while enjoying parallelizable computation and fast inference speed. Code is available at https://github.com/LeapLabTHU/MLLA.