简介

本推文主要介绍了由清华大学黄高老师团队发表在ECCV 2024上的一篇论文《Agent Attention: On the Integration of Softmax and Linear Attention》，文中提出了一种新型的代理注意力（Agent Attention）。近年来，Transformer在计算机视觉领域中具有重要地位。然而，传统的基于Softmax的自注意力具有O(N2)的二次复杂度，导致其难以扩展到高分辨率的视觉任务中。基于此，作者提出了代理注意力，通过引入代理tokens来聚合全局语义信息，从而将传统的二次复杂度降低为线性复杂度。实验结果表明，Agent Attention在多个计算机视觉任务中表现出色。

本推文由李杨撰写，审校为陆新颖和朱旺。

论文链接：https://arxiv.org/abs/2312.08874。

代码链接：https://github.com/LeapLabTHU/Agent-Attention。

一、研究背景

Transformer模型在计算机视觉领域因其强大的全局上下文建模能力取得了显著进展，但其核心组件——基于Softmax的自注意力机制，计算复杂度为O(N2)，在高分辨率任务中面临巨大计算负担，限制了应用。为解决这一问题，提出了如Swin Transformer和PVT等方法，通过局部窗口或稀疏注意降低计算量，但这也降低了对全局关系的建模能力。

本文提出了Agent Attention，通过引入代理（Agent）tokens扩展注意力机制为四元组(Q, A, K, V)。代理tokens先聚合键值对(K, V)的信息，再将其广播回查询tokens Q，从而将计算复杂度从二次降低为线性复杂度，兼顾效率和全局建模能力。实验表明，Agent Attention在图像分类、目标检测、语义分割和图像生成等任务中表现出色，尤其在高分辨率场景下显著提升了性能，如在Stable Diffusion模型中不仅加速了生成，还提高了质量，无需额外训练。

图1 Softmax Attention、Linear Attention和Agent Attention之间的差异

图1展示了Softmax Attention、Linear Attention和Agent Attention之间的差异

    · Softmax Attention计算所有查询和键的相似度，复杂度为O(N2)，具有较高的表达能力但计算量大。

    · Linear Attention通过对查询和键分别应用映射函数ϕ(·)，改变计算顺序，从而降低复杂度为O(N)，但表达能力有所减弱。

    · Agent Attention引入了代理tokens，用于聚合和广播全局信息，以实现Softmax和Linear Attention的优点结合，既保持了较高的表达能力，又显著降低了计算复杂度。

二、方法

图2 Agent Attention的方法细节

首先，Agent Attention的计算过程分为两步：代理信息聚合（Agent Aggregation）和代理信息广播（Agent Broadcast）。在代理信息聚合阶段，代理tokens A作为查询，从键值对(K, V)中聚合全局信息，得到代理特征VA；在代理信息广播阶段，代理tokens A将聚合的信息重新广播回查询tokens Q，完成全局信息的传递，从而实现最终的输出。

其次，图2还展示了模块组成，包括代理tokens的生成和深度可分离卷积（DWC）模块。代理tokens通常通过池化或其他降维方法从输入特征中生成，以降低计算复杂度，而DWC模块用于保持特征的多样性，确保在聚合和广播过程中不会丢失重要信息。

三、实验结果

图3 在ImageNet-1K上的结果

图3展示了Agent Attention在不同视觉Transformer模型（DeiT、PVT、Swin、CSwin）上的性能表现。每个子图分别对比了使用Agent Attention前后的准确率与计算复杂度（FLOPs）关系。

DeiT：Agent-DeiT在相同计算复杂度下显著提高了准确率。

PVT：Agent-PVT在低计算量下实现了更高的精度，优于其他基线模型。

Swin：Agent-Swin显著提高了准确率，优于其他变体，如Swin T/S/B等。

CSwin：Agent-CSwin在不同计算复杂度下均取得了较高的准确率，与其他基线相比具有明显优势。

整体来看，图3展示了使用Agent Attention能够在同等或更少计算量的情况下显著提高模型的准确率。

图4 利用Stable Diffusion生成的样本

图4展示了Stable Diffusion、ToMeSD和AgentSD生成的样本对比，分别生成了金丝雀、桌子、手套和青椒的图像。

可以看到，AgentSD生成的图像在细节和质量上显著优于其他两种方法，尤其在对象的清晰度和细节保留上表现更为出色。具体来说，Stable Diffusion生成的图像有时会缺乏清晰的边界和细节；ToMeSD虽然有所改善，但仍存在一定的细节模糊问题。而AgentSD通过引入Agent Attention，显著提高了生成质量，使得生成的对象更符合高质量图像的要求，特别是在复杂结构和纹理表现上更加真实。

图5 不同模型在使用Agent Attention时的效果对比

图5展示了不同生成模型在使用Agent Attention时的效果对比。包括输入图像、生成的提示（Prompt）、Dreambooth、仅在生成阶段使用Agent Attention、以及在微调和生成阶段均使用Agent Attention。

从图5可以看出，仅在生成阶段使用Agent Attention的效果在细节和准确性上有所提升，但相比于在微调和生成阶段均使用Agent Attention，生成的图像在一致性和细节保留方面略有不足。而在微调和生成阶段都使用Agent Attention的情况下，生成的图像在视觉效果上最为逼真和一致。

四、总结

作者提出了一种新的注意力机制——Agent Attention，通过引入代理tokens，有效地降低了Transformer模型中的计算复杂度，同时保持了强大的全局建模能力。实验结果表明，Agent Attention在多个视觉任务中均表现优异，尤其在高分辨率场景下能够显著提高计算效率和模型性能。此外，Agent Attention在图像生成任务（如Stable Diffusion）中的应用也验证了其在提升生成速度和图像质量方面的潜力。整体来看，Agent Attention为Transformer模型在高效性和表达能力之间取得了良好的平衡，是一种值得进一步研究和推广的注意力机制。