CV技术指南 | 其实Mamba是一种线性注意力?清华大学黄高团队揭秘开视觉Mamba的真实面目!

本文来源公众号“CV技术指南”,仅用于学术分享,侵权删,干货满满。

原文链接:其实Mamba是一种线性注意力?清华大学黄高团队揭秘开视觉Mamba的真实面目!

前言 

本文揭示了 Mamba 模型与 Linear Attention Transformer 之间有着惊人的相似之处。作者通过探索 Mamba 和 Linear Transformer 之间的相似性和差异,在本文中提供了一个全面的分析来揭示 Mamba 成功的关键因素。

0 本文目录

1 揭秘视觉 Mamba:一种线性注意力机制视角
(来自清华大学 (黄高团队),阿里巴巴)
1 MLLA 论文解读
1.1 探索 Mamba 和 Linear Attention Transformer 的关联
1.2 线性注意力与选择性状态空间模型简介
1.3 用统一的视角看待选择性状态空间模型与线性注意力
1.4 关于每个差异的分析
1.5 宏观架构设计
1.6 每种差异的影响和 MLLA 最终方案
1.7 实验结果

太长不看版

Mamba 是一种具有线性计算复杂度的状态空间模型。其最近在处理不同视觉任务的高分辨率输入中展示出很不错的效率。本文揭示了 Mamba 模型与 Linear Attention Transformer 之间有着惊人的相似之处。作者通过探索 Mamba 和 Linear Transformer 之间的相似性和差异,在本文中提供了一个全面的分析来揭示 Mamba 成功的关键因素。

具体而言,作者使用统一的公式重新表述了选择性状态空间模型和线性注意力,将 Mamba 重新表述为 Linear Attention Transformer 的变体。它们主要有6个区别:输入门 (input gate)、遗忘门 (forget gate)、快捷连接 (shortcut)、无注意力归一化、single-head 和修改后的 Block Design。对于每个设计,本文仔细分析了它的优缺点,并实证性地评估了其对视觉模型性能的影响。更有趣的是,遗忘门 (forget gate) 和修改后的 Block Design 是 Mamba 模型成功的核心贡献,而其他的四种设计不太关键。

基于这些发现,作者将这两个比较重要的设计融入 Linear Attention 中,并提出一种类似 Mamba 的线性注意力模型,其在图像分类和高分辨率密集预测任务上都优于视觉 Mamba 模型,同时享受并行化的计算和快速推理。

本文做了哪些具体的工作

  1. 揭示了 Mamba 与 Linear Attention Transformer 之间的关系:Mamba 和 Linear Attention Transformer 可以使用一个统一的框架表示。与传统的 Linear Attention 的范式相比,Mamba 有6种不同设计:输入门 (input gate)、遗忘门 (forget gate)、快捷连接 (shortcut)、无注意力的归一化、single-head 和经过修改的 Block Design

  2. 对上述的每一种特殊的设计进行了详细分析,并实证验证了遗忘门 (forget gate) 和 Block Design 很大程度上是 Mamba 性能优越的关键。此外,证明了遗忘门 (forget gate) 的循环计算可能不是视觉模型的理想选择。相反,适当的位置编码可以作为视觉任务中的遗忘门 (forget gate) ,同时保持并行化的计算和快速的推理。

### 关于线性注意力机制在 Mamba 中的实现和应用 Mamba一种状态空间模型,其设计旨在通过降低计算复杂度来提高处理大规模数据集尤其是高分辨率图像的能力。不同于传统的自注意力机制,线性注意力机制能够显著减少内存占用并加快推理速度。 #### 线性注意力机制概述 传统Transformer中的自注意力层通常涉及矩阵乘法操作,这使得时间复杂度达到O(n²),对于长序列而言非常低效。而线性注意力方法则采用了一种更高效的策略,它利用特定函数将键值对映射到较低维度的空间内再执行加权求和运算,从而实现了近似线性的复杂度 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
发出的红包
实付
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

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

余额充值