全新注意力范式!清华黄高团队提出Agent Attention:无缝集成Softmax和Linear的注意力机制

AI最新资讯4周前发布 tree
24 0 0

今日应用


今日话题


全新注意力范式!清华黄高团队提出Agent Attention:无缝集成Softmax和Linear的注意力机制
全新注意力范式!清华黄高团队提出Agent Attention:无缝集成Softmax和Linear的注意力机制
 

重点标签 Agent AttentionSoftmax AttentionLinear AttentionTransformer模型视觉任务

文章摘要


Agent Attention的核心思想是在传统的Softmax AttentionLinear Attention的基础上,引入一组额外的Agent token,使得Agent token作为Query token的代理,先聚合来自Key和Value的信息,然后将信息广播回Query token。这样,Agent Attention既保留了Softmax Attention的全局上下文建模能力,又具有Linear Attention的线性复杂度优势。

在Softmax Attention中,计算每个Query和Key之间的相似度会带来二次方的计算复杂度。而Linear Attention通过将映射函数应用于Key和Value,降低了计算复杂度。Agent Attention则通过两个Softmax Attention操作,实现了Agent token的聚合和广播,避免了成对相似度的计算,同时保留了信息交换。

Agent Attention在多种视觉任务中表现出色,包括图像分类、目标检测、语义分割和图像生成。特别是在高分辨率场景中,Agent Attention的性能得到了显著提升。例如,在Stable Diffusion模型中,应用Agent Attention可以加速生成并提高图像生成质量,无需额外训练。

此外,Agent Attention还具有大感受野和高分辨率的优势。通过逐渐扩展Swin Transformer的窗口大小,模型性能不断提高,表明Agent Attention可以从全局感受野中受益,同时保持线性计算量。在与其他线性注意力机制的对比中,Agent Attention也显示出了优越的性能。

总之,Agent Attention作为一种新型的注意力机制,通过集成Softmax Attention和Linear Attention的优势,提高了Transformer模型的效率和性能,在各种视觉任务中展现出了广泛的应用前景。[/Agent Attention/Softmax Attention/Linear Attention/Transformer模型/视觉任务]

极市导读:本文介绍了一种新型的注意力机制Agent Attention,它结合了Softmax AttentionLinear Attention的优点,通过引入Agent token来平衡Query token与Key-value对之间的关系,提高了Transformer模型的效率和性能。

1 Agent Attention:集成 Softmax 和 Linear 注意力机制
本文提出的Agent Attention是在传统的注意力模块中引入了一组额外的Agent token。Agent token首先充当Query token的代理来聚合来自Key和Value的信息,然后将信息广播回Query token。由于Agent token的数量可以设计为远小于Query token的数量,代理注意力明显比Softmax注意力更有效,同时保留了全局上下文建模能力。

1.1 Agent Attention 集成 Softmax Attention 和 Linear Attention 的优势
将Transformer和Self attention引入视觉领域会带来巨大的挑战。现代Transformer模型通常采用Softmax attention,计算每个Query和Key之间的相似度,导致计算复杂度随token数量呈二次方关系。为了解决这个问题,现有的工作通过设计高效的注意力机制来降低计算复杂度。本文将一组额外的token引入注意力三元组,产生一个四元组注意力范式,称为Agent Attention。

1.2 Softmax Attention 和 Linear Attention 计算范式
假设输入为X,X由N个tokens组成。Self-Attention中每个head可以表述为:其中W^Q、W^K、W^V为投影矩阵,d为模型的维度,d_k为head的维度,f为相似度函数。当使用f为softmax时,公式就变成Softmax Attention,其在现代视觉Transformer中非常成功。但是,Softmax Attention强制计算所有Query和Key之间的相似度,会带来O(N^2)的计算复杂度。Linear Attention把映射函数分别应用于Key和Value,这样就可以根据矩阵乘法的结合律将计算顺序从O(N^2)改为O(N)。

1.3 Agent Transformer
Softmax和Linear注意力要么计算复杂度过高,要么模型表达能力不足。Agent Attention集成了Softmax Attention和Linear Attention,同时享受线性复杂度和高表现力的好处。Agent Attention由两个Softmax Attention操作组成,即Agent聚合和Agent广播。具体而言,Agent token首先被视为Query并在Key和Value之间执行注意力计算。然后使用Agent token作为第二个Attention计算中的Key和Value,使用Query矩阵,将全局信息从Agent特征广播到每个Query标记并获得最终输出。

1.4 感知任务实验结果
在ImageNet-1K实验中,将Softmax Attention替换为Agent Attention可以显着提高性能。例如,Agent-PVT-S在仅使用30%的参数和40%的FLOPs时超过了PVT-L。在COCO目标检测实验中,本文模型在所有配置中都表现出一致的增强。[Agent Attention/Softmax Attention/Linear Attention/Transformer模型/视觉任务] 本文介绍了一种新型的注意力机制Agent Attention,它结合了Softmax Attention和Linear Attention的优点,通过引入Agent token来平衡Query token与Key-value对之间的关系,提高了Transformer模型的效率和性能。Agent Attention在多种视觉任务中表现出色,包括图像分类、目标检测、语义分割和图像生成。特别是在高分辨率场景中,Agent Attention的性能得到了显著提升。此外,Agent Attention还具有大感受野和高分辨率的优势,在与其他线性注意力机制的对比中显示出了优越的性能。

摘要:

本文提出了一种新型的注意力机制——Agent Attention,它成功地结合了Softmax AttentionLinear Attention的优势。通过引入额外的Agent token,Agent Attention在保持全局上下文建模能力的同时,实现了计算复杂度的显著降低。这种机制特别适用于Transformer模型,并且在多种视觉任务中展现出了卓越的性能。

Agent Attention的核心原理是在传统的注意力模块中加入Agent token,这些token作为Query token的代理,首先聚合来自Key和Value的信息,然后将这些信息广播回Query token。这种方法不仅减少了计算量,而且保留了Softmax Attention的全局建模能力。

在实验部分,作者展示了Agent Attention在多种视觉任务中的应用效果。在ImageNet-1K数据集上,使用Agent Attention替换传统的Softmax Attention后,模型性能得到了显著提升。例如,Agent-PVT-S在参数和计算量减少的情况下,性能超过了PVT-L。此外,在COCO目标检测任务中,Agent Attention同样展现出了其优势,无论是在1x还是3x的训练计划下,都取得了一致的性能提升。

特别值得一提的是,Agent Attention在高分辨率场景下的性能表现。在Stable Diffusion模型中,应用Agent Attention可以加速图像生成过程,并在无需额外训练的情况下提高生成质量。这得益于Agent Attention的线性复杂度特性,使其在处理高分辨率图像时仍能保持高效的计算性能。

此外,Agent Attention还具有大感受野的优势。通过逐渐扩展Swin Transformer的窗口大小,模型性能得到了持续提升,这表明Agent Attention能够从全局感受野中受益,同时保持线性计算量。在与其他线性注意力机制的对比中,Agent Attention的性能也显示出了明显的优势。

总之,Agent Attention作为一种新型的注意力机制,通过集成Softmax Attention和Linear Attention的优势,显著提高了Transformer模型的效率和性能,在各种视觉任务中展现出了广泛的应用前景。

文章来源


原文地址: 点我阅读全文
原文作者: 极市平台

© 版权声明

相关文章

暂无评论

暂无评论...