论文《Demystify Mamba in Vision: A Linear Attention Perspective 》从线性注意力的视角,对Mamba 进行了阐释。
论文名称:Demystify Mamba in Vision: A Linear Attention Perspective (Arxiv 2024.05)
论文地址:https://arxiv.org/pdf/2405.16605
https://ar5iv.labs.arxiv.org/html/2405.16605
理解Mamba的一个维度是从RNN出发,另一个维度是从线性注意力的角度。
文章《Mamba是一种线性注意力?清华大学揭开视觉Mamba的真面目!》对这篇论文有详细的阐释。
(文章链接:https://mp.weixin.qq.com/s/0yPNXAujJyRLIIgM9bxF3Q)
我这里摘录一些关键信息。
Mamba的形式转化
直接看文章的图2
图2展示了如何将SSM的公式转换为与线性注意力机制相似的形式,从而揭示SSM和线性注意力之间的联系。
线性注意力的形式转化
Softmax注意力的形式:
线性注意力的形式:
线性注意力转化为递归形式:
Mamba与线性注意力的对比
分别从图示和公式的形式,将manba和线性注意力进行对比:
除了这4个差异之外,重要的是要注意式12表示单头线性注意力。这表明选择性状态空间模型类似于单头线性注意,不包含多头设计。
总之,选择性 SSM 和线性注意力之间的相似性和差异可以概括为:
选择性状态空间模型类似于具有附加输入门、遗忘门和 Shortcut 的线性注意力,同时省略归一化和多头设计。
进一步的分析:
文章的摘要和结论
摘要:Mamba是一个具有线性计算复杂度的有效状态空间模型。它最近在处理各种视觉任务中的高分辨率输入方面展现出了令人印象深刻的效率。在本文中,我们揭示了强大的Mamba模型与线性注意力Transformer之间惊人的相似之处,后者在实践中通常表现不如传统的Transformer。通过探索有效的Mamba和表现不佳的线性注意力Transformer之间的相似性和差异,我们提供了全面的分析,以揭开Mamba成功的幕后关键因素。具体来说,我们在统一的公式框架内重新表述了选择性状态空间模型和线性注意力,将Mamba描述为线性注意力Transformer的一个变体,具有六个主要区别:输入门、遗忘门、快捷连接、无注意力归一化、单头和修改后的块设计。对于每个设计,我们都仔细分析了其优缺点,并通过实证评估了其对视觉任务中模型性能的影响。有趣的是,结果突出显示,遗忘门和块设计是Mamba成功的两个核心贡献因素,而其他四个设计则不那么关键。基于这些发现,我们提出了一个Mamba启发的线性注意力(MILA)模型,将这两个关键设计的优点融入线性注意力中。结果表明,该模型在图像分类和高分辨率密集预测任务中都超越了各种视觉Mamba模型,同时享有可并行计算和快速推理速度
文章的主要结论:
-
Mamba与线性注意力Transformer的关系:文章揭示了Mamba模型与线性注意力Transformer之间惊人的相似性,并将Mamba重新表述为线性注意力Transformer的一个变体,具有六个主要区别:输入门、遗忘门、快捷方式、无注意力标准化、单头和修改的块设计。
-
关键设计因素分析:通过细致分析每个设计因素的利弊,并实证评估其对视觉任务中模型性能的影响,发现遗忘门和块设计是Mamba成功的核心因素,而其他四个设计因素则不那么关键。
-
提出Mamba-Like Linear Attention (MLLA)模型:基于上述发现,文章提出了一个Mamba-Like Linear Attention (MLLA)模型,该模型将两个关键设计的优点结合到线性注意力中。结果表明,MLLA模型在图像分类和高分辨率密集预测任务方面都优于各种视觉Mamba模型,同时具有可并行计算和快速推理速度。
-
遗忘门的替代方案:文章还探讨了遗忘门的本质,并验证了在视觉任务中可以用适当的位置编码替代遗忘门,以保持并行化的计算和快速推理。
这些结论提供了对Mamba模型成功因素的深入理解,并展示了如何通过结合Mamba的关键设计优点来提升线性注意力模型的性能。
**Mamba块设计的宏观架构:**现代线性注意力Transformer模型通常采用如图3(a)所示的块设计,包括线性注意力子块和MLP(多层感知机)子块。与之相比,Mamba通过结合H3和Gated Attention两种基本设计,修改了块设计,如图3(b)所示。Mamba的改进块设计整合了多种操作,如选择性状态空间模型、深度卷积、线性映射、激活函数、门控机制等,比传统的Transformer块设计更有效。