本文目录

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 的数量导致性能下降。

参考

1. ^Swin Transformer: Hierarchical Vision Transformer using Shifted Windows
2. ^Pyramid Vision Transformer: A Versatile Backbone for Dense Prediction without Convolutions
3. ^Rethinking Attention with Performers
4. ^SOFT: Softmax-free Transformer with Linear Complexity
5. ^FLatten Transformer: Vision Transformer using Focused Linear Attention
6. ^High-Resolution Image Synthesis with Latent Diffusion Models
7. ^Token Merging for Fast Stable Diffusion
8. ^Token Merging: Your ViT But Faster
9. ^DreamBooth: Fine Tuning Text-to-Image Diffusion Models for Subject-Driven Generation

公众号后台回复“数据集”获取100+深度学习各方向资源整理