0. 论文信息
标题：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
机构：Tsinghua University、Alibaba Group
原文链接：https://arxiv.org/abs/2405.16605
代码链接：https://github.com/LeapLabTHU/MLLA
1. 摘要
Mamba是一种有效的状态空间模型，具有线性计算复杂度。最近，它在处理各种视觉任务的高分辨率输入方面表现出了令人印象深刻的效率。在本文中，我们揭示了强大的Mamba模型与线性注意力变压器有着惊人的相似之处，后者在实践中的表现通常不如传统变压器。通过探索有效的曼巴和亚线性注意力转移器之间的相似性和差异，我们提供了全面的分析来揭示曼巴成功背后的关键因素。具体来说，我们在一个统一的公式中重新表述了选择性状态空间模型和线性注意力，将Mamba重新表述为线性注意力转换器的一种变体，具有六个主要区别:输入门、遗忘门、快捷方式、无注意力标准化、单头和修改的块设计。对于每个设计，我们仔细分析其利弊，并根据经验评估其对视觉任务中模型性能的影响。有趣的是，结果强调了遗忘门和块设计是Mamba成功的核心因素，而其他四种设计则不那么重要。基于这些发现，我们通过将这两个关键设计的优点结合到线性注意中，提出了一个类似曼巴的线性注意(MLLA)模型。结果模型在图像分类和高分辨率密集预测任务方面都优于各种视觉Mamba模型，同时具有可并行计算和快速推理速度。
2. 引言
近年来，以Mamba为代表的状态空间模型迅速获得了广泛的研究兴趣。与流行的Transformer模型所具有的二次复杂度相比，基于状态空间的Mamba以线性复杂度提供了有效的序列建模。这一关键特性使得Mamba能够以可控的计算成本处理极长的序列，成为自然语言处理和视觉识别领域颇具前景的架构。
然而，Mamba并非首个实现线性复杂度全局建模的模型。早期的线性注意力（Linear Attention）被提出作为一种计算高效的替代方案，以取代广泛采用的Softmax注意力（即点积注意力）。具体而言，线性注意力用线性归一化替换了注意力操作中的非线性Softmax函数，从而改变了计算顺序，从(QK⊤)V变为Q(K⊤V)，从而将计算复杂度从O(N²)降低到O(N)。尽管其效率高，但先前的工作证明线性注意力存在表达能力不足的问题，这使其在实际应用中不太实用。令人惊讶的是，我们发现了高性能的Mamba与表现不佳的线性注意力Transformer之间的密切关系。因此，一个引人注目的研究问题出现了：是什么因素促成了Mamba的成功，并使其显著优于线性注意力Transformer？
在本文中，我们从线性注意力Transformer的角度，通过理论分析和实证研究揭示了Mamba的奥秘。具体而言，我们在统一的框架下重新编写了选择性状态空间模型和线性注意力的公式，将Mamba描述为线性注意力Transformer的一种变体，具有六个不同之处：输入门、遗忘门、捷径连接、无注意力归一化、单头以及修改后的块设计。为了解开导致Mamba有效的因素之谜，我们进行了视觉任务的实证研究，以评估每个特殊设计的影响。结果表明，遗忘门和块设计往往是Mamba优越性的两个核心贡献者。虽然块设计可以轻松采用，但遗忘门需要递归计算，这可能不太适合非自回归的视觉模型。因此，我们深入探究了遗忘门的本质，并验证了在视觉任务中，它可以用适当的位置编码来替代。基于我们的发现，我们将这两个核心贡献者或它们的替代品引入线性注意力Transformer，提出了我们的类Mamba线性注意力（MLLA）模型。实验结果表明，在图像分类和高分辨率密集预测任务中，MLLA实现了优于各种Mamba模型的结果，验证了线性注意力在具备两个核心设计优势的情况下可以超越Mamba。
3. 主要贡献
我们的主要贡献和收获如下：
• 我们揭示了Mamba与线性注意力Transformer之间的密切关系：Mamba和线性注意力Transformer可以在统一的框架下表述，与传统的线性注意力范式相比，Mamba展现了六个独特的设计：输入门、遗忘门、捷径连接、无注意力归一化、单头以及修改后的块设计。
• 我们对每个特殊设计进行了详细分析，并通过实证研究验证了遗忘门和块设计在很大程度上导致了Mamba的优越性。此外，我们还证明了遗忘门的递归计算可能不太适合视觉模型。相反，在视觉任务中，适当的位置编码可以替代遗忘门的作用，同时保持并行计算和快速推理速度。
• 我们开发了一系列名为MLLA的线性注意力视觉Transformer模型，它们继承了Mamba的核心优点，并且比原始的Mamba模型更适合视觉任务。
4. 方法
分别展示了Mamba中的选择性SSM和单头线性注意力的示意图。可以看出，选择性SSM类似于单头线性注意力，但增加了输入门Δi、遗忘门Aei和捷径D⊙xi，同时省略了归一化QiZi。
现代线性注意力Transformer模型通常采用图3(a)所示的块设计，该设计由线性注意力子块和全连接层（MLP）子块组成。相比之下，Mamba通过结合H3和门控注意力两种基本设计，对块设计进行了改进，从而得到了图3(b)所示的结构。改进的Mamba块集成了多个操作，如选择性SSM、深度卷积、线性映射、激活函数、门控机制等，且通常比传统的Transformer块设计更为有效。
5. 实验结果
如表3所示，由于结合了Mamba和线性注意力的优点，我们的MLLA模型在所有模型尺寸上均优于各种视觉Mamba模型。这些结果也验证了，凭借Mamba的两个关键设计的优势，性能较差的线性注意力Transformer可以超越高性能的Mamba。值得注意的是，我们经验性地观察到，与视觉Mamba模型相比，MLLA表现出更高的可扩展性，因为MLLA-B达到了85.3%的准确率，远超其他模型。此外，MLLA还优于各种卷积神经网络（CNN）和视觉Transformer设计。例如，MLLA的性能优于MambaOut，后者是最近的一项工作，移除了Mamba中的选择性SSM并采用了门控卷积架构。
推荐课程：
聊一聊经典三维点云方法，包括：点云拼接、聚类、表面重建、QT+VTK等
。
如表4所示，在COCO数据集上，MLLA模型也优于视觉Mamba模型，表明其在高分辨率密集预测任务中的有效性。MLLA以线性复杂度O(N)提供了有效的全局建模，并可实现并行计算，因此非常适合高分辨率图像建模场景。值得注意的是，MLLA以显著优势超越了MambaOut，这与MambaOut中的发现相一致。
6. 总结 & 未来工作
本文揭示了功能强大的Mamba模型与次优的线性注意力Transformer之间出人意料的紧密联系，为Mamba的优越性和成功提供了一些见解。我们将Mamba重新表述为线性注意力Transformer的一种变体，并识别出其六大主要特殊设计：输入门、遗忘门、捷径、无注意力归一化、单头注意力以及修改后的块设计。实验验证表明，遗忘门和块设计在很大程度上提升了性能，而其他差异则提供了微不足道的贡献或损害了模型性能。基于我们的发现，我们提出了类似Mamba的线性注意力（MLLA）模型，通过将这些两个关键设计的优点融入线性注意力中。MLLA在多个任务上超越了各种视觉Mamba模型，同时保持了并行计算和较高的推理速度。
对更多实验结果和文章细节感兴趣的读者，可以阅读一下论文原文~
本文仅做学术分享，如有侵权，请联系删文。
3D视觉交流群，成立啦！
目前我们已经建立了3D视觉方向多个社群，包括
2D计算机视觉
、
最前沿
、
工业3D视觉
、
SLAM
、
自动驾驶
、
三维重建
、
无人机
等方向，细分群包括：
工业3D视觉
：相机标定、立体匹配、三维点云、结构光、机械臂抓取、缺陷检测、6D位姿估计、相位偏折术、Halcon、摄影测量、阵列相机、光度立体视觉等。
SLAM
：视觉SLAM、激光SLAM、语义SLAM、滤波算法、多传感器融合、多传感器标定、动态SLAM、MOT SLAM、NeRF SLAM、机器人导航等。
自动驾驶：深度估计、Transformer、毫米波|激光雷达|视觉摄像头传感器、多传感器标定、多传感器融合、自动驾驶综合群等、3D目标检测、路径规划、轨迹预测、3D点云分割、模型部署、车道线检测、Occupancy、目标跟踪等。
三维重建
：3DGS、NeRF、多视图几何、OpenMVS、MVSNet、colmap、纹理贴图等
无人机
：四旋翼建模、无人机飞控等
2D计算机视觉
：图像分类/分割、目标/检测、医学影像、GAN、OCR、2D缺陷检测、遥感测绘、超分辨率、人脸检测、行为识别、模型量化剪枝、迁移学习、人体姿态估计等
最前沿
：具身智能、大模型、Mamba、扩散模型等
除了这些，还有
求职
、
硬件选型
、
视觉产品落地、产品、行业新闻
等交流群
添加小助理: dddvision，备注：
研究方向+学校/公司+昵称
（如
3D点云+清华+小草莓
）, 拉你入群。
▲长按扫码添加助理：cv3d008
3D视觉知识星球
「
3D视觉从入门到精通
」知识星球，已沉淀6年，星球内资料包括：
秘制视频课程近20门
（包括
