来自清华大学的研究者提出了一种新的注意力范式——署理注意力 (Agent Attention)。
近年来，视觉 Transformer 模型得到了极大的发展，相关工作在分类、联系、检测等视觉任务上都取得了很好的效果。然而，将 Transformer 模型运用于视觉领域并不是一件简单的事情。与自然语言不同，视觉图片中的特性数量更多。由于 Softmax 注意力是平方复杂度，直接进行全部自注意力的算计往往会带来过高的算计量。针对这一问题，先前的工作通常通过减少参与自注意力算计的特性数量的方法来降低算计量。例如，设计稀疏注意力机制（如 PVT）或将注意力的算计限制在局部窗口中（如 Swin Transformer）。尽管有效，这样的自注意力方法很容易受到算计模式的影响，同时也不可避免地牺牲了自注意力的全部建模能力。
与 Softmax 注意力不同，线性注意力将 Softmax 解耦为两个独立的函数，从而不妨将注意力的算计顺序从 (query・key)・value 调整为 query・(key・value)，使得总体的算计复杂度降低为线性。然而，目前的线性注意力方法效果明显逊于 Softmax 注意力，难以实际运用。
注意力模块是 Transformers 的关键组件。全部注意力机制具良好的模型表达能力，但过高的算计成本限制了其在各种场景中的运用。本文提出了一种新的注意力范式，署理注意力 (Agent Attention)，同时具有高效性和很强的模型表达能力。
论文链接：https://arxiv.org/abs/2312.08874
代码链接：https://github.com/LeapLabTHU/Agent-Attention
具体来说，署理注意力在传统的注意力三元组 (Q,K,V) 中引入了一组额外的署理向量 A，定义了一种新的四元注意力机制 (Q, A, K, V)。其中，署理向量 A 首先作为查询向量 Q 的署理，从 K 和 V 中聚合信息，然后将信息广播回 Q。由于署理向量的数量可以设计得比查询向量的数量小得多，署理注意力不妨以很低的算计成本实现全部信息的建模。
此外，本文证明署理注意力等价于一种线性注意力范式，实现了高本能 Softmax 注意力和高效线性注意力的自然融合。该方法在 ImageNet 上使 DeiT、PVT、Swin Transformer、CSwin Transformer 等模型架构取得了明显的本能提拔，不妨将模型在 CPU 端加速约 2.0 倍、在 GPU 端加速约 1.6 倍。运用于 Stable Diffusion 时，署理注意力不妨将模型生成速度提拔约 1.8 倍，并明显提高图像生成质量，且无需任何额外训练。
方法
在本文中，我们创新性地向注意力三元组 (Q,K,V) 引入了一组额外的署理向量 A，定义了一种四元的署理注意力范式 (Q, A, K, V)。如图 1 (c) 所示，在署理注意力中，我们不会直接算计 Q 和 K 之间两两的相似度，而是使用少量的署理向量 A 来收集 K 和 V 中的信息，进而呈递给 Q，以很低的算计成本实现全部信息的建模。从整体结构上看，署理注意力由两个常规 Softmax 注意力操作组成，并且等效为一种广义的线性注意力，实现了高本能 Softmax 注意力和高效线性注意力的自然融合，因而同时具有二者的优点，即：算计复杂度低且模型表达能力强。
图 1：Softmax 注意力、线性注意力与署理注意力机制对比
1. 署理注意力
图 2：署理注意力示意图
上图即为署理注意力的示意图，下面给出具体数学形式。为了书写方便，我们将 Softmax 注意力和线性注意力分别缩写为：
其中，Q,K,V 分别为 Query、Key、Value 矩阵，表示 Softmax 函数，为线性注意力中的映射函数。则署理注意力可以表示为：
另一个等效的表示为：
其中 A 为新定义的署理矩阵。
如公式 (3) 和示意图第一行所示，署理注意力由两个 Softmax 注意力操作组成，分别为署理特性聚合和广播。具体来说，我们首先将 A 作为 Query，在 A、K 和 V 之间进行注意力算计，从所有特性中汇聚信息，得到署理特性
。随后，我们将 A 作为 Key，
作为 Value，和 Q 进行第二次注意力算计，将署理特性中的全部信息广播回每一个特性，并获得最终输出 O。这样一来，我们避免了 Q 和 K 之间相似度的算计，而是通过署理向量实现了每个 query-key 之间的信息交换。可以看到，在这一算计范式中，少量的署理特性 A 充当了 Q 的 “署理人”—— 从 K 和 V 中收集信息并呈递给 Q，因而本文将这种注意力机制命名为署理注意力。实际运用中，我们将 A 的数量设置为一个小的超参数 n，从而以线性算计复杂度
实现了全部建模。
值得指出的是，如公式 (4) 和示意图第二行所示，署理注意力实际上将高本能的 Softmax 注意力和高效的线性注意力融合在了一起，通过使用两次 Softmax 注意力操作实现了广义线性注意力范式，其中等效映射函数定义为
。
实际运用中，署理向量可以通过不同的方法获得，例如设置为一组可学习参数，或通过池化等方式从输入特性中得到。我们也可以使用更加优越的方法来获得署理向量，例如 Deformable Points、Token Merging 等。本文中，我们采用简单的池化来获取署理向量。
2. 署理注意力模块
为了更好地发挥署理注意力的潜力，本文进一步做出了两方面的改进。一方面，我们定义了 Agent Bias 以促进不同的署理向量聚焦于图片中不同的位置，从而更好地利用位置信息。另一方面，作为一种广义的线性注意力，署理注意力也面临特性多样性不足的问题，因此我们采用一个轻量化的 DWC 作为多样性恢复模块。
在以上设计的基础上，本文提出了一种新的署理注意力模块，其结构如下图：
图 3：署理注意力模块
结合了 Softmax 注意力和线性注意力的优势，署理注意力模块具有以下特点：
(1) 算计复杂度低且模型表达能力强。之前的研究通常将 Softmax 注意力和线性注意力视为两种不同的注意力范式，试图解决各自的问题和局限。署理注意力优雅地融合了这两种注意力形式，从而自然地继承了它们的优点，同时享受低算计复杂性和高模型表达能力。
(2) 不妨采用更大的体验野。得益于线性算计复杂度，署理注意力可以自然地采用更大的体验野，而不会增加模型算计量。例如，可以将 Swin Transformer 的 window size 由 7^2 扩大为 56^2，即直接采用全部自注意力，而完全不引入额外算计量。
实验结果
1. 分类任务
署理注意力是一个通用的注意力模块，本文基于 DeiT、PVT、Swin Transformer、CSwin Transformer 等模型架构进行了实验。如下图所示，在 ImageNet 分类任务中，基于署理注意力构建的模型不妨取得明显的本能提拔。例如，Agent-Swin-S 可以取得超越 Swin-B 的本能，而其参数量和算计量不到后者的 60%。
图 4：ImageNet 图片分类结果
在实际推理速度方面，署理注意力也具有明显的优势。如下图所示，在 CPU/GPU 端，署理注意力模型不妨取得 2.0 倍 / 1.6 倍左右的加速，同时取得更好的本能。
图 5：实际测速结果
2. 检测和联系
在检测和联系任务中，相较于基础模型，Agent Transformer 也不妨取得十分明显的本能提拔，这在一定程度上得益于署理注意力的全部体验野。
图 6：COCO 物体检测与联系结果
图 7：ADE20K 语义联系结果
3.Agent Stable Diffusion
特别值得指出的是，署理注意力可以直接运用于 Stable Diffusion 模型，无需训练，即可加速生成并明显提拔图片生成质量。如下图所示，将署理注意力运用于 Stable Diffusion 模型，不妨将图片生成速度提拔约 1.8 倍，同时提拔图片的生成质量。
图 8：Stable Diffusion, ToMeSD 和 AgentSD 的定量化结果
下图中给出了生成图片的样例。可以看到，署理注意力不妨明显降低 Stable Diffusion 模型生成图片的歧义和错误，同时提拔生成速度和生成质量。
图 9：生成图片的样例
4. 高分辨率与大体验野
本文还探究了分辨率和体验野对模型本能的影响。如下图所示，我们基于 Agent-Swin-T 将窗口大小由 7^2 逐步扩大到 56^2。可以看到，随着体验野的扩大，模型本能稳步提拔。这说明尽管 Swin 的窗口划分是有效的，但它依然不可避免地损害了模型的全部建模能力。
图 10：体验野大小的影响
下图中，我们将图片分辨率由 256^2 逐步扩大到 384^2。可以看到，在高分辨率的场景下，署理注意力模型持续展现出明显的优势。
图 11：高分辨率场景
总结
本文的贡献主要在三个方面：
(1) 提出了一种新颖、自然、有效且高效的注意力范式 —— 署理注意力，它自然地融合了高本能的 Softmax 注意力和高效的线性注意力，以线性算计量实现有效的全部信息建模。
(2) 在分类、检测、联系等诸多任务中充分验证了署理注意力的优越性，特别是在高分辨率、长序列的场景下，这或为开发大尺度、细粒度、面向实际运用场景的视觉、语言大模型提供了新的方法。
(3) 创新性地以一种无需训练的方式将署理注意力运用于 Stable Diffusion 模型，明显提拔生成速度并提高图片质量，为扩散模型的加速和优化提供了有效的新研究思路。