作者 | 科技猛兽 编辑 | 极市平台
点击下方卡片，关注“自动驾驶之心”公众号
戳我-> 领取自动驾驶近15个方向学习路线
>>点击进入→自动驾驶之心『Transformer』技术交流群
本文只做学术分享，如有侵权，联系删文
导读
本文介绍了一种新型的注意力机制Agent Attention，它结合了Softmax Attention和Linear Attention的优点。Agent Attention通过引入Agent token来平衡Query token与Key-value对之间的关系，提高了Transformer模型的效率和性能。
本文目录
1 Agent Attention：集成 Softmax 和 Linear 注意力机制
(来自清华，黄高老师团队)
1 Agent Attention 论文解读
1.1 Agent Attention 集成 Softmax Attention 和 Linear Attention 的优势
1.2 Softmax Attention 和 Linear Attention 计算范式
1.3 Agent Transformer
1.4 感知任务实验结果
1.5 生成任务实验结果
1.6 大感受野和高分辨率
太长不看版
注意力机制 (Attention module) 是 Transformers 中的关键组成部分。虽然全局的注意力机制具有很高的表征能力，但其计算成本较大，限制了其在各种场景下的适用性。本文提出一种新的注意力范式 Agent Attention, 目的在计算效率和表征能力之间取得良好的平衡。具体而言, Agent Attention 表示为四元组 , 在传统的注意力模块中引入了一组额外的 Agent token 。Agent token 首先充当 Query token 的代理来聚合来自 和 的信息, 然后将信息广播回 。鉴于 Agent token 的数量可以设计为远小于 Query token 的数量, 代理注意力明显比 Softmax 注意力更有效, 同时保留了全局上下文建模能力。
有趣的是，本文展示了 Agent attention 等效于 Linear attention 的广义形式。因此，代理注意力无缝集成了强大的 Softmax attention 和高效的 Linear attention。
作者通过大量实验表明，Agent attention 在各种视觉任务中证明了有效性，包括图像分类、目标检测、语义分割和图像生成。而且，代理注意力在高分辨率场景中表现出显着的性能，这得益于其线性注意力性质。例如，当应用于 Stable Diffusion 时，Agent attention 会加速生成并显着提高图像生成质量，且无需任何额外的训练。
1 Agent Attention：集成 Softmax 和 Linear 注意力机制
论文名称：Agent Attention: On the Integration of Softmax and Linear Attention (Arxiv 2023.12)
论文地址：
https://arxiv.org/pdf/2312.08874
代码链接：
https://github.com/LeapLabTHU/Agent-Attention
1.1 Agent Attention 集成 Softmax Attention 和 Linear Attention 的优势
将 Transformer 和 Self attention 引入视觉领域会带来巨大的挑战。现代 Transformer 模型通常采用 Softmax attention，计算每个 Query 和 Key 之间的相似度，导致计算复杂度随 token 数量呈二次方关系。为了解决这个问题，现有的工作通过设计高效的注意力机制来降低计算复杂度。比如，Swin Transformer[1]减少了感受野，将 Self-Attention 的计算限制在局部窗口上。PVT[2]采用稀疏 Attention，通过减少 Key 和 Value 的数量来减轻计算负担。尽管这些方法很有效，但它们不可避免地损害了对远程关系进行建模的能力，并且仍然不如全局 Self-Attention。
本文将一组额外的 token 引入注意力三元组 , 产生一个四元组注意力范式 , 称为Agent Attention。如图1所示, Agent Attention 由两个传统的 Softmax Attention 操作组成。第一个 Softmax attention 应用于三元组 , 其中代理 token 作为 Query 来聚合来自 的信息, 注意矩阵在 和 之间计算。第二个 Softmax attention 在三元组 上执行, 其中 是上一步的结果, 形成 Agent Attention 的最终输出。新引入的 token 可以被视为 Query token 的 "代理", Query token 不再需要与原始 和 直接通信。因此将token 称为代理 token。
由于全局 Self-Attention 的内在冗余，Agent token 的数量可以设计为远小于 Query token 的数量。作者发现简单地汇集原始 Query token 来当做 Agent token 的效果就非常好。这个做法可以将 Softmax Attention 的二次复杂度降低到线性复杂度，同时保留了全局上下文建模能力。有趣的是，如图 1 所示，Agent Attention 可以看作是广义的线性注意力。换句话说，Agent Attention 集成了 Softmax 和线性注意力，并享受二者的优势。
1.2 Softmax Attention 和 Linear Attention 计算范式
假设输入为 , 为 个 tokens。Self-Attention 中每个 head 可以表述为：
式中, 为投影矩阵, 为模块的维度, 为 head 的维度, 为相似度函数。
当使用 时, 1式就变成 Softmax Attention, 其在现代视觉 Transformer 中非常成功。但是, Softmax Attention 强制计算所有 Query 和 Key 之间的相似度,会带来 的计算复杂度。因此, 使用具有全局感受野的 Softmax Attention 会导致很大的计算复杂度。
Linear Attention 把映射函数分别应用于 和 , 即 。这样就可以根据矩阵乘法的结合律将计算顺序从 改为 。这样可以使计算复杂度降低到 。
然而, 设计有效的映射函数 并不容易。简单的 ReLU 导致性能显著下降, 而更复杂的设计 或矩阵分解 会引入额外的计算开销。一般来说, 当前的线性注意力方法仍然不如 Softmax Attention, 限制了它们的实际应用。
1.3 Agent Transformer
Softmax 和 Linear 注意力要么计算复杂度过高，要么模型表达能力不足。以前的研究通常将这两种注意力范式视为不同的方法，并尝试降低 Softmax Attention 的计算成本或提高 Linear Attention 的性能。Agent Attention 集成了 Softmax Attention 和 Linear Attention，同时享受线性复杂度和高表现力的好处。
首先将 Softmax 和 Linear Attention 缩写为：
其中, 表示 Query、Key 和 Value 矩阵， 表示 Softmax 函数。那么 Agent Attention 可以写成:
上式等价于：
式中，为 Agent tokens。
Agent Attention 由两个 Softmax Attention 操作组成, 即 Agent 聚合和 Agent 广播。具体而言, Agent token 首先被视为 Query 并在 和 之间执行注意力计算。然后使用 作为第二个 Attention 计算中的 Key 和 Value, 使用 Query 矩阵 , 将全局信息从 Agent 特征广播到每个 Query 标记并获得最终输出 。通过这种方式, 避免了 和 之间的成对相似度的计算, 同时通过 Agent token 保留每个 Query Key 对之间的信息交换。
Agent token 本质上用作 的代理, 聚合来自 和 的全局信息, 然后将其广播回 。实际上, 将 Agent token 的数量 设置为一个较小的超参数值, 在保持全局上下文建模能力的同时,实现了 的计算复杂度, 相对于输入特征的数量 呈线性。
有趣的是, Agent Attention 实际上集成了 Softmax Attention和高效的 Linear Attention, 通过两个 Softmax 操作建立了广义 Linear Attention 范式, 等价映射函数为
在实践中，Agent token 可以通过不同的方法获得，比如简单地设置为一组可学习的参数或通过池化从输入特征中提取。本文作者采用简单的池化策略来获得 Agent token，已经表现得非常好。
Agent Attention 继承了 Softmax 和 Linear Attention 的优势。实际使用中作者进一步做了两个改进来最大化代理注意力的潜力，即 Agent Bias：
式中, 为 Agent Bias。它用空间信息增强了普通 Agent 的注意力，帮助不同的 Agent tokens 专注于不同的区域。
多样性恢复模块
尽管 Agent Attention 受益于低计算复杂度和高模型表达能力，但它也受到特征多样性不足的影响。作为补救措施，作者遵循[5]的做法并采用深度卷积 (DWC) 模块来保留特征多样性。
在这些设计的基础上，作者提出了一种新的 Agent Attention 模块，其可以表述为：
式中, 。
Agent Attention 模块的优势
1) 高效计算和强表达能力： 之前的工作通常将 Softmax Attention 和 Linear Attention 视为两种不同的注意力范式，旨在解决它们各自的局限性。作为这两种注意力形式的无缝集成，Agent Attention 继承了两者的优点，同时具有较低的计算复杂度和高模型表达能力。
2) 大感受野： Agent Attention 可以在保持相同数量的计算量的同时采用较大的感受野。得益于线性复杂度，Agent Attention 可以在保持线性计算量的同时享受大甚至全局感受野的优势。
1.4 感知任务实验结果
ImageNet-1K 实验结果
如图3所示，在各种模型中将 Softmax Attention 替换为 Agent Attention 可以显着提高性能。例如，Agent-PVT-S 在仅使用 30% 的参数和 40% 的 FLOPs 时超过了 PVT-L。Agent-Swin-T/S 在保持相似 FLOPs 的同时比 SwinT/S 高出 1.3% 和 0.7%。这些结果明确证明 Agent Attention 方法具有优势，且能够适应不同的架构。
作者通过在各种设备上部署模型来进一步进行实际速度测量。图4说明了本文模型在 CPU 上实现了 1.7 到 2.1 倍的推理速度，同时提高了性能。在 RTX3090 GPU 和 A100 GPU 上，本文模型也实现了 1.4 倍到 1.7 倍的推理速度。
COCO 目标检测实验结果
作者将本文模型应用于 RetinaNet、Mask R-CNN 和 Cascade Mask R-CNN 框架来评估本文方法的性能。使用具有不同检测头的 1x 和 3x schedules 进行了一系列实验。如图5所示，本文模型在所有配置中都表现出一致的增强。Agent-PVT 优于 PVT 模型，box AP 从 +3.9 增加到 +4.7，而 Agent-Swin 超过 Swin 模型高达 +1.5 box AP。这些实质性的改进可以归因于大感受野，证明了 Agent Attention 在高分辨率场景的有效性。
ADE20K 语义分割实验结果
作者将本文模型应用于2个分割模型，SemanticFPN 和 UperNet。结果如图6所示。值得注意的是，Agent-PVT-T 和 Agent-Swin-T 比 PVT-T 和 Swin-T 高 +3.61 和 +2.17 的 mIoU。结果表明本文模型与各种分割 Backbone 兼容，并且都实现了改进。
1.5 生成任务实验结果
扩散模型的出现使得生成高分辨率和高质量的图像成为可能。然而，当前的扩散模型主要使用具有全局感受野的 Softmax Attention，导致计算成本大，且生成速度慢。作者将 Agent Attention 应用于 Stable Diffusion[6]，希望提高模型的生成速度。经过简单的调整，使用 Agent Attention (称为 AgentSD) 的 Stable Diffusion 的生成速度展示出显着改进，并且在没有任何额外训练的情况下产生了更好的图像质量。
作者实际上将 Agent Attention 应用于 ToMeSD 模型[7]。ToMeSD 在 Stable Diffusion 的注意力计算之前减少了 token 的数量，提高了生成速度。尽管如此，合并后的 token 数量仍然很大，导致持续的复杂度和延迟。因此，作者将 ToMeSD 模型中的 Softmax Attention 替换为 Agent Attention，以进一步提高速度。作者通过实验发现，通过 token merging[8]生成 Agent token 时，Agent Attention 可以直接应用于 Stable Diffusion 和 ToMeSD 模型，而不需要任何额外的训练。除此之外，作者通过在早期扩散生成步骤中应用 Agent Attention 并保持后面的步骤不变来获得显着提升。
作者定量比较了 AgentSD 与 Stable Diffusion 和 ToMeSD。如图7所示，ToMeSD 在保持相似图像质量的同时加速了 Stable Diffusion。AgentSD 不仅进一步加速了 ToMeSD，而且显著提高了图像生成质量。具体而言，在保持卓越的图像生成质量的同时，与 Stable Diffusion 和 ToMeSD 相比，AgentSD 的生成速度提高了 1.84 倍和 1.69 倍。在等效的生成速度下，与 ToMeSD 相比，AgentSD 生成 FID 分数低 0.9。
作者在图8中展示了一些可视化结果。与 Stable Diffusion 和 ToMeSD 相比，AgentSD 显著地减少了歧义和生成错误。例如，在第1列中，Stable Diffusion 和 ToMeSD 产生一条腿和两个尾巴的鸟类，而 AgentSD 的样本没有表现出这个问题。在第3列中，当提供 "A high quality photo of a mitten" 的 prompt 时，Stable Diffusion 和 ToMeSD 错误地生成猫，而 AgentSD 产生了正确的图像。
用于微调的 AgentSD
作者将代理注意力应用于基于 SD 的 Dreambooth[9]，以验证其在微调下的性能。当微调时，Agent Attention 可以集成到所有扩散生成步骤中，与原始 Dreambooth 相比，生成速度提高了 2.2 倍。
1.6 大感受野和高分辨率
现代视觉 Transformer 通常将 Self-Attention 的计算限制在局部窗口上，以降低计算复杂度，比如 Swin。如下图9所示，作者逐渐扩展 Swin 的窗口大小从 7^27^2 到 56^256^2 。显然，随着感受野的扩展，模型的性能不断提高。这表明，虽然窗口注意力范式是有效的，但它不可避免地损害了 Self-Attention 的远程建模能力，仍然不如全局注意力机制。由于 Agent Attention 的线性复杂度，可以从全局感受野中受益，同时仍然保持相同的计算复杂度。
受 Softmax attention 的二次复杂度的限制，当前的视觉 Transformer 通常通过增加模型深度和宽度来扩大。作者也尝试了 EfficientNet 中提出的提升输入分辨率的方法，结果如图10所示。
与 DeiT-B 相比，Agent-DeiT-B 实现了 0.2 的精度增益，而 448^2448^2 分辨率下的 Agent-DeiT-S 仅用四分之一的参数达到了 83.1% 的精度。作者在缩放 Agent-PVT-M 和 Agent-Swin-S 时观察到了类似的趋势，在图11中，作者逐渐增加 Agent-Swin-S、Swin-S 和 Swin-B 的分辨率。在高分辨率场景中，本文模型性能始终比较优越。
1.7 与其他线性注意力机制的对比
作者使用 DeiT-T 和 Swin-T 将本文的 Agent Attention 与其他 Linear Attention 方法进行比较。如图12所示，各种 Linear Attention 方法替换 DeiT-T 和 Swin-T 所采用的 Softmax Attention 通常会导致显著的性能下降。值得注意的是，本文模型优于其他方法以及 Softmax 基线模型。
Agent token 数量的消融实验结果
模型的计算复杂度可以通过改变 Agent token 的数量来加以调节。如图13所示，可以观察到减少浅层中的 agent token 数量对性能没有不利的影响。但是，在更深层中减少 agent token 的数量导致性能下降。
参考
^Swin Transformer: Hierarchical Vision Transformer using Shifted Windows
^Pyramid Vision Transformer: A Versatile Backbone for Dense Prediction without Convolutions
^Rethinking Attention with Performers
^SOFT: Softmax-free Transformer with Linear Complexity
^FLatten Transformer: Vision Transformer using Focused Linear Attention
^High-Resolution Image Synthesis with Latent Diffusion Models
^Token Merging for Fast Stable Diffusion
^Token Merging: Your ViT But Faster
^DreamBooth: Fine Tuning Text-to-Image Diffusion Models for Subject-Driven Generation
投稿作者为『自动驾驶之心知识星球』特邀嘉宾，欢迎加入交流！
① 全网独家视频课程
BEV感知、BEV模型部署、BEV目标跟踪、毫米波雷达视觉融合、多传感器标定、多传感器融合、多模态3D目标检测、车道线检测、轨迹预测、在线高精地图、世界模型、点云3D目标检测、目标跟踪、Occupancy、cuda与TensorRT模型部署、大模型与自动驾驶、Nerf、语义分割、自动驾驶仿真、传感器部署、决策规划、轨迹预测等多个方向学习视频（扫码即可学习）
② 国内首个自动驾驶学习社区
国内最大最专业，近3000人的交流社区，已得到大多数自动驾驶公司的认可！涉及30+自动驾驶技术栈学习路线，从0到一带你入门自动驾驶感知（2D/3D检测、语义分割、车道线、BEV感知、Occupancy、多传感器融合、多传感器标定、目标跟踪）、自动驾驶定位建图（SLAM、高精地图、局部在线地图）、自动驾驶规划控制/轨迹预测等领域技术方案、大模型、端到端等，更有行业动态和岗位发布！欢迎扫描下方二维码，加入自动驾驶之心知识星球，这是一个真正有干货的地方，与领域大佬交流入门、学习、工作、跳槽上的各类难题，日常分享论文+代码+视频
③【自动驾驶之心】技术交流群
自动驾驶之心是首个自动驾驶开发者社区，聚焦感知、定位、融合、规控、标定、端到端、仿真、产品经理、自动驾驶开发、自动标注与数据闭环多个方向，目前近60+技术交流群，欢迎加入！扫码添加汽车人助理微信邀请入群，备注：学校/公司+方向+昵称（快速入群方式）
④【自动驾驶之心】全平台矩阵
扫一扫
1452
被折叠的 条评论
为什么被折叠?
到【灌水乐园】发言
