【AI学习】Mamba学习(二十):Mamba是一种线性注意力

论文《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模型,同时享有可并行计算和快速推理速度

文章的主要结论:

  1. Mamba与线性注意力Transformer的关系:文章揭示了Mamba模型与线性注意力Transformer之间惊人的相似性,并将Mamba重新表述为线性注意力Transformer的一个变体,具有六个主要区别:输入门、遗忘门、快捷方式、无注意力标准化、单头和修改的块设计。

  2. 关键设计因素分析:通过细致分析每个设计因素的利弊,并实证评估其对视觉任务中模型性能的影响,发现遗忘门和块设计是Mamba成功的核心因素,而其他四个设计因素则不那么关键。

  3. 提出Mamba-Like Linear Attention (MLLA)模型:基于上述发现,文章提出了一个Mamba-Like Linear Attention (MLLA)模型,该模型将两个关键设计的优点结合到线性注意力中。结果表明,MLLA模型在图像分类和高分辨率密集预测任务方面都优于各种视觉Mamba模型,同时具有可并行计算和快速推理速度。

  4. 遗忘门的替代方案:文章还探讨了遗忘门的本质,并验证了在视觉任务中可以用适当的位置编码替代遗忘门,以保持并行化的计算和快速推理。

这些结论提供了对Mamba模型成功因素的深入理解,并展示了如何通过结合Mamba的关键设计优点来提升线性注意力模型的性能。

**Mamba块设计的宏观架构:**现代线性注意力Transformer模型通常采用如图3(a)所示的块设计,包括线性注意力子块和MLP(多层感知机)子块。与之相比,Mamba通过结合H3和Gated Attention两种基本设计,修改了块设计,如图3(b)所示。Mamba的改进块设计整合了多种操作,如选择性状态空间模型、深度卷积、线性映射、激活函数、门控机制等,比传统的Transformer块设计更有效。
在这里插入图片描述

### 关于线性注意力机制在 Mamba 中的实现和应用 Mamba一种状态空间模型,其设计旨在通过降低计算复杂度来提高处理大规模数据集尤其是高分辨率图像的能力。不同于传统的自注意力机制,线性注意力机制能够显著减少内存占用并加快推理速度。 #### 线性注意力机制概述 传统Transformer中的自注意力层通常涉及矩阵乘法操作,这使得时间复杂度达到O(),对于长序列而言非常低效。而线性注意力方法则采用了一种更高效的策略,它利用特定函数将键值对映射到较低维度的空间内再执行加权求和运算,从而实现了近似线性的复杂度 O(n)[^2]。 #### 在 Mamba 中的应用 具体来说,在 Mamba 架构里引入了线性变换后的查询向量 q'、键 k' 及值 v' 来代替原始版本: \[ \text{Attention}(Q, K, V ) = softmax\left(\frac{{q'}^{T}k'}{\sqrt {d}}\right)v'\] 这里 d 表示特征维度大小;上述公式经过简化后可以表示成如下形式: \[ \text{Linformer Attention}(Q,K,V)=VW_{v}\sigma (K W_{k}^{T}) Q W_{q}\] 其中 \(W_k\) ,\(W_v\) 和 \(W_q\) 分别代表应用于键(key), 值(value)以及询问(query)上的权重参数矩阵[^1]。 这种改进不仅保持了原有模型的强大表达能力,还极大地提高了训练效率与资源利用率。此外,由于减少了不必要的冗余计算,该技术特别适合部署于边缘设备或其他受限环境中运行复杂的AI应用程序。 ```python import torch.nn as nn class LinearAttention(nn.Module): def __init__(self, dim, heads=8, dropout=0.1): super().__init__() self.heads = heads self.scale = dim ** (-0.5) self.to_qkv = nn.Linear(dim, dim * 3, bias=False) self.to_out = nn.Sequential( nn.Linear(dim, dim), nn.Dropout(dropout) ) def forward(self, x): b, n, _, h = *x.shape, self.heads qkv = self.to_qkv(x).chunk(3, dim=-1) q, k, v = map(lambda t: rearrange(t, 'b n (h d) -> b h n d', h=h), qkv) dots = einsum('b h i d, b h j d -> b h i j', q, k) * self.scale attn = dots.softmax(dim=-1) out = einsum('b h i j, b h j d -> b h i d', attn, v) out = rearrange(out, 'b h n d -> b n (h d)') return self.to_out(out) ```
评论
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

当前余额3.43前往充值 >
需支付:10.00
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包

打赏作者

bylander

你的鼓励将是我创作的最大动力

¥1 ¥2 ¥4 ¥6 ¥10 ¥20
扫码支付:¥1
获取中
扫码支付

您的余额不足,请更换扫码支付或充值

打赏作者

实付
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

1.余额是钱包充值的虚拟货币,按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载,可以购买VIP、付费专栏及课程。

余额充值