Agent Attention：Softmax与线性注意力的融合研究

摘要

https://arxiv.org/pdf/2312.08874.pdf
在Transformer中，注意力模块是其关键组件。虽然全局注意力机制提供了高度的表达能力，但其过高的计算成本限制了其在各种场景下的应用。本文提出了一种新颖的注意力范式，称为Agent Attention，以在计算效率和表示能力之间取得良好的平衡。具体来说，Agent Attention由四个部分组成(Q, A, K, V)，它在常规注意力模块中引入了一组额外的agent tokens A。这些agent tokens首先充当查询tokens Q的代理，从K和V中聚合信息，然后将信息广播回Q。由于agent tokens的数量可以设计得比查询tokens的数量少得多，因此agent attention的计算效率显著高于广泛采用的Softmax attention，同时保留了全局上下文建模能力。有趣的是，我们发现所提出的agent attention相当于线性注意力的一种广义形式。因此，agent attention将强大的Softmax attention和高效的线性attention无缝地集成在一起。大量的实验证明了agent attention在各种视觉Transformer和各种视觉任务中的有效性，包括图像分类、目标检测、语义分割和图像生成。值得注意的是，由于其线性注意力的性质，agent attention在高清场景中表现出色。例如，当应用于Stable Diffusion时，我们的agent attention加速了生成过程，并显著提高了图像生成质量，而无需任何额外的训练。代码可在https://github.com/LeapLabTHU/Agent-Attention上获取。