Agent Attention On the Integration of Softmax and Linear Attention

二三两。

已于 2024-03-12 16:25:46 修改

阅读量2.3k

点赞数 35

分类专栏：论文解读文章标签：深度学习人工智能计算机视觉

于 2024-03-12 16:16:17 首次发布

本文链接：https://blog.csdn.net/qq_47182555/article/details/135884885

版权

论文解读专栏收录该内容

8 篇文章 0 订阅

订阅专栏

1 前言

本文要是对《Agent Attention On the Integration of Softmax and Linear Attention》这篇论文的一个解读与总结，原文链接[2312.08874] Agent Attention: On the Integration of Softmax and Linear Attention

本文提出一种新型Transformer模型，它结合了Linear Attention和Softmax Attention的优点，在计算效率和表示能力之间取得良好的平衡。具体来说，Agent Attention，表示为四元组 $\left ( Q,A,K,V \right )$ ，在传统的Attention模块中引入了一组额外的Agent token A。Agent token首先作为Query token Q的代理，从K和V中聚合信息，然后将信息广播回Q。代理令牌的数量可以设计得比查询令牌的数量小得多，从而减少了时间复杂度。本文证明了所提出的代理注意等同于线性注意的广义形式。

2 现有研究存在的问题

上图描述了Softmax Attention，Linear Attention，Agent Attention三种注意力机制的计算过程。

自Vision Transformer诞生以来，自关注在计算机视觉领域取得了显著的进步。然而，普遍存在的Softmax注意力的二次复杂度给将自注意应用于视觉任务带来了挑战。以下是一些常见的注意力机制的简单介绍，同时分析了它们各自的优缺点。

（1）Softmax注意力机制

做法：采用Softmax函数计算Q，K之间的注意力权重。

公式： $Softmax Attention= Softmax\left ( Q,K \right )V$ 。

优点：可以关注的到全局感受野。

缺点：时间复杂度为二次，复杂度高不适合处理长序列数据。

（2）线性注意力机制

做法：抛弃了Softmax函数，使用映射函数 $\Phi \left ( \cdot \right )$ 代替Softmax函数计算Q，K之间的注意力权重，改变计算次序从而降低时间复杂度。

公式： $Linear Attention= \Phi \left ( Q \right )\Phi \left ( K \right )V$ 。

优点：模型简单，计算效率高，时间复杂度为一次，适用性广。

缺点：可能无法捕捉到非线性的关系，缺乏灵活性同时对异常值敏感。

（3）窗口注意力机制

做法：每个时间步只关注序列中的局部窗口范围，而不是整个序列。

优点：减少计算复杂度，更好的捕捉局部信息。

缺点：丢失全局信息，对窗口大小的选择不易。

（4）稀疏注意力机制

做法：在计算注意力权重时引入稀疏性约束，使得每个时间步只能关注少数位置或者特征。

优点：提高效率，改善泛化能力。

缺点：有用信息丢失，稀疏性约束的选择需要谨慎。

3 Agent Attention

根据上面的分析，我们可以将改进注意力机制方法分为两种：

① 通过限制自注意力的整体感受野的大小，比如PVT、Swin Transformer。

② 通过采用不同的映射函数改变计算次序的线性注意力机制。

虽然已有的改进注意力机制的方法都有效，但是都存在一定的缺陷。在本片论文中，作者提出了一种新的注意力范式，而不是增强Softmax或线性注意力，将这两种注意力类型整合在一起，在各种任务中获得卓越的表现。

3.1 Agent注意模块的信息流

在上图中，描述的是Agent注意模块的信息流。该例子采用的是对Q进行池化操作以获得Agent token。

首先使用三个独立的线性变换层将输入转换成注意力查询，键，值矩阵 $Q,K,V\in \mathbb{R}^{N\times C}$ 将它们分别写成 $Q= \left [ q_{1} ,q_{2}\cdots q_{N}\right ]$ ， $K= \left [ k_{1} ,k_{2}\cdots k_{N}\right ]$ ， $V= \left [ v_{1} ,v_{2}\cdots v_{N}\right ]$ 。再定义我们的Agent token $A\in \mathbb{R}^{n\times C}$ 。

随后，利用 A 结合 K，V 采用Softmax注意力机制聚合来自 K,V 的信息，添加 Agent Bias 增加位置信息，得到 $V_{A}$ 。

最后 $Q,A,V_{A }$ 采用Softmax注意力机制， $Q$ 从 $V_{A}$ 中查询特征。此外，对 V 进行 DWC 来保持特征多样性。

3.2 Agent Attention公式表示

将Softmax注意力机制和Linear注意力机制简化为：

根据上方公式，Agent Attention可以简化为：

等价与

Agent Attention由两个 Softmax Attention 操作组成，即 Agent Aggregation 和 Agent Broadcast 。其中括号内的 Attn 为 Agent Aggregation ，聚合特征。括号外的 Attn 为 Agent Broadcast ，广播特征。

（1）首先将 Agent token A视为查询，并在 A、K和V 之间执行注意力计算，从所有值中聚合代理特征 $V_{A}$ 。

（2）随后，利用 A 作为键， $V_{A}$ 作为值，利用查询矩阵 Q 进行第二次关注计算，将智能体特征的全局信息广播到每个查询令牌上，得到最终输出 O 。

这样既避免了Q和K的成对相似度计算，又保留了每个查询键 Q 通过 A 进行的信息交换。新定义的 Agent token A 本质上充当 Q 的代理，聚合来自 K 和 V 的全局信息，并随后将其广播回 Q 。实际上，将 A 的数量 n 设置为一个小的超参数，在保持全局上下文建模能力的同时，实现了 $O\left ( Nnd \right )$ 的线性计算复杂度。

3.3 进一步改进

在实际应用中，作者进一步进行了两方面的改进，以最大限度地发挥代理注意力的潜力。

(1) Agent Bias：为了更好地利用位置信息，作者为 Agent Attention 精心设计了的一个Agent Bias。本文中采用三个偏差分量（列偏量，行偏量，块偏量）来构建 Agent Bias ，用空间信息增强了普通Agent Attention 的表达能力，允许不同的Agent token 专注与不同的区域。