本文目录

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：

式中,