Attention Free Transformer(AFT)

paper: An Attention Free Transformer

date: 2021-05

org: Apple

1 Motivation

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

2 Method

2.1 Multi-Head Attention回顾

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

$$f_i(X)=\sigma (\frac{Q_i(K_i{)}^T}{\sqrt{d_k}})V_i,s.t.Q_i=X{W}_i^Q,K_i=X{W}_i^K,V_i=X{W}_i^V,\text{\hspace{1em}(1)}$$

其中： $W_i^Q\in R^{d\times d_k},W_i^K\in R^{d\times d_k},W_i^V\in R^{d\times d_{\upsilon }}$。$\sigma$是非线性函数，默认为$softmax$。通常情况下$d_v=d_k,h=\frac{d}{d_k}$。假定输入 $X\in {\mathbb{R}}^{T\times d}$, 经过$f_i$转化后的输出$f_i(X)\in {\mathbb{R}}^{T\times d_v}$。将所有head的结果拼接起来得到最后的输出${\mathbb{R}}^{T\times d}$。

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

• $QKV$ 的计算,此处有3个矩阵乘法，计算量为 $d\times d_k\times T\times 3$， 时间复杂度为：$\mathcal{O}(\frac{1}{h}T{d}^2)$
• $Q{K}^T$的计算,计算量为： $d_k\times T\times T$， 时间复杂度为：$\mathcal{O}(\frac{1}{h}{T}^{2}d)$
• scale 的计算量为： $T\times T$， 时间复杂度为：$\mathcal{O}(T^2)$
• softmax的计算量为： $T\times T$， 时间复杂度为：$\mathcal{O}(T^2)$
• 最后加权乘法计算量为 $d_k\times T\times T$，时间复杂度为：$\mathcal{O}(\frac{1}{h}{T}^{2}d)$

对于MHA，时间复杂度为$\mathcal{O}(T{d}^2)$

2.2 Attention Free Transofrmer(AFT)

2.2.1 AFT full

第一步和MHA一样，输入$X$经过三个linear transfer得到$QKV$，3个矩阵， 维度为${\mathbb{R}}^{T\times d}$。AFT引入了一个新的可训练参数矩阵$w\in {\mathbb{R}}^{T\times T}$，论文将其称之为可学习的一对一位置偏置（learned pair-wise position biases）。

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

SETP1: 求$\mathrm{weighted}(K^{(t)})$。从$w$ 取$t=t$的向量, 和$K$做点乘后以列方向进行$\mathrm{softmax}$。该步骤的计算复杂度为$\mathcal{O}(T\times d)$

$$\mathrm{Weighted}(K^{(t)})=\frac{\exp (K+w_t)}{{\sum }_{i=1}^T\exp (k_i+w_{ti})}\text{\hspace{1em}(2)}$$

STEP2: 求${\mathrm{Attention}}^{(t)}$矩阵。将q_t用sigmoid变换后，点乘wighted(K)。该步骤的计算复杂度为$\mathcal{O}(T\times d)$

$$\mathrm{Attentio}{\mathrm{n}}^{(\mathrm{t})}=\sigma (q_t)\odot \mathrm{Weighted}(K^{(t)})=\frac{\sigma (q_t)\odot \exp (K+w_t)}{{\sum }_{i=1}^T\exp (k_i+w_{ti})}\text{\hspace{1em}(3)}$$

STEP3: 计算$y_t$。该步骤的计算复杂度为$\mathcal{O}(T\times d)$

$$y_t=\sum _{i=1}^T({\mathrm{Attentio}{\mathrm{n}}^{(\mathrm{t})}}_i\odot v_i)=\sum _{i=1}^T\frac{\sigma (q_t)\odot \exp (k_i+w_t)}{{\sum }_{i=1}^T\exp (k_i+w_{ti})}\odot v_i\text{\hspace{1em}(4)}$$

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

$$y_t=\sigma (q_t)\odot \frac{{\sum }_{i=1}^T\exp (k_i+w_t)\odot v_i}{{\sum }_{i=1}^T\exp (k_i+w_{ti})}\text{\hspace{1em}(5)}$$

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

式（4）计算pipeline

式（5）计算pipeline

2.2.1 AFT local

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

$$w_{t,t^{\prime }}=\{\begin{array}{l}{w}_{t,t^{\prime }},\mathrm{if}|t-t^{\prime }|<s\\ 0,otherwise.\end{array}\text{\hspace{1em}(6)}$$

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

2.2.2 AFT simple

AFT simple是AFT local当$s=0$时的特殊形式。此时没有位置偏置。可将式5化简为，因为对不同的$t$， ${\sum }_{i=1}^T(\mathrm{softmax}(K)\odot V{)}_i$都是相同的。AFT simple的时间复杂度为$\mathcal{O}(Td)$

$$y_t=\sigma (q_t)\odot \frac{{\sum }_{i=1}^T\exp (k_i)\odot v_i}{{\sum }_{i=1}^T\exp (k_i)}=\sigma (q_t)\odot \sum _{i=1}^T(\mathrm{softmax}(K)\odot V{)}_i\text{\hspace{1em}(6)}$$

2.2.3 AFT conv

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

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

$$y_t^j={\sigma }_q(q_t^j)\odot \frac{conv1d(\exp (K^j)\odot V^j,\exp (w^j)-1)+{\sum }_{i=1}^T\exp (k_i^j)\odot v_i^j}{conv1d(\exp (K^j),\exp (w^j)-1)+{\sum }_{i=1}^T\exp (k_i^j)}\text{\hspace{1em}(7)}$$

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

3 小结

本文提出了一种Dot Product Attention Free的Transformer，最多能将transofmer的时间复杂度从$\mathcal{O}(T^{2}d)$降低到$\mathcal{O}(Td)$（AFT-simple）。