Attention Free Transformer(AFT)技术小结

type

status

date

slug

summary

1 Motivation

原本基于dot product self attention Transformer的时间复杂度和空间复杂度都很高。提出了一个新的AFT层来降低transformer的计算量。

2 Method

2.1 Multi-Head Attention回顾

首先回顾一下经典的Multi-Head Attention(MHA)，每一个head的计算如下

其中：。i是非线性函数，默认为。通常情况下。假定输入 , 经过转化后的输出。将所有head的结果拼接起来得到最后的输出。

单头Attention的时间复杂度计算：

的计算,此处有3个矩阵乘法，计算量为，时间复杂度为：

的计算,计算量为：，时间复杂度为：

scale 的计算量为：，时间复杂度为：

softmax的计算量为：，时间复杂度为：

最后加权乘法计算量为，时间复杂度为：

对于MHA，时间复杂度为

2.2 Attention Free Transofrmer(AFT)

2.2.1 AFT full

第一步和MHA一样，输入经过三个linear transfer得到，3个矩阵，维度为。AFT引入了一个新的可训练参数矩阵，论文将其称之为可学习的一对一位置偏置（learned pair-wise position biases）。

我们以为视角看每一步的具体流程。

SETP1: 求。从取的向量, 和做加法后以列方向进行。该步骤的计算复杂度为

STEP2: 求矩阵。将q_t用sigmoid变换后，点乘wighted(K)。该步骤的计算复杂度为

STEP3: 计算。该步骤的计算复杂度为

对式（4）稍做变形，可得论文中的计算公式

将所有的步骤串起来的流程如下。可以看到AFT其实也用到了attention的思想。但AFT中的Attention Score的计算并没有用到矩阵乘法，只用到了向量点乘。虽整体的计算复杂度仍然是，但计算量已有所下降。

2.2.1 AFT local

在许多情况下，局部性是一个很重要的归纳偏置（inductive bias），而标准的Transformer的计算中没有引入局部信息。因此，作者提出AFT-local。其形式与AFT-Full一致。区别在于，引入了下式限制

式中的就是定义的局部窗口大小（local window size）。它进一步降低了计算量。变换后的如下图所示（此时，黑色方块为0）。

2.2.2 AFT simple

AFT simple是AFT local当时的特殊形式。此时没有位置偏置。可将式5化简为，因为对不同的，都是相同的。AFT simple的时间复杂度为

详细的计算流程和复杂度计算可参考下图：

2.2.3 AFT conv

作者进一步将局部性的思想扩展到空间权重共享（如卷积），提出AFT-conv。具体来说，让的值仅依赖和的相对位置。为了考虑参数数量随着数增加而增长的情况，作者采用了一个设计选择，将的维度与head数绑定在一起（MHA的思路）。这使得AFT-conv可以采用深度可分离卷积、全局池化和element-wise操作的实现方式。

可以看到与AFT simple相比，AFT conv引入了head思想，并通过1维卷积的计算结果引入局部信息。其形式与式(6)相比分子分母中新增了，。（上标表示第个head）。此时的为conv1d的filter。

从ViT可视化attention map中可以看出（横轴为head，纵轴为layer）。原本的ViT（左边）的不同层，head的attention map的响应最大区域基本都是中心区域。而用了AFT-conv后，不同层、head的attention都有所不同，有助于模型捕获不同尺度的特征。

3 小结

本文提出了一种Dot Product Attention Free的Transformer，最多能将transofmer的时间复杂度从降低到（AFT-simple）。