传统的注意力机制需要对 $Q$ 和 $K$ 中的每对元素进行计算以得到注意力评分，在计算时需要得到一个注意力矩阵，这导致其空间开销的增长速度是 $O(T^2)$ 的。AFT通过，将空间复杂度降低到了 $O(Td)$ ，其中 $d$ 为特征向量长度。并且通过忽略长距离向量间的注意力，AFT还可以将时间复杂度降低到 $O(Tsd)$ ，其中 $s$ 为注意力区间的长度，在图像处理中很有用。

AFT

AFT的过程如下：对于输入向量 $X\in R^{T\ast d}$ ,首先和传统自注意力机制一样，通过三个线性变换得到 $Q$ , $K$ , $V$ 三个量：

$$\begin{array}{r}Q=X{W}^Q\\ K=X{W}^K\\ V=X{W}^V\end{array}$$

然后，通过下面的公式来生成结果：

$$Y=f(X);Y_t={\sigma }_q(Q_t)\odot \frac{\sum _{t^{\prime }=1}^T\exp (K_{t^{\prime }}+w_{t,t^{\prime }})\odot V_{t^{\prime }}}{\sum _{t^{\prime }=1}^T\exp (K_{t^{\prime }}+w_{t,t^{\prime }})}$$

其中， $w_{t,t^{\prime }}$ 是一个可以学习的参数，作用和位置编码类似。
初看这个公式可能很难理解为什么AFT要这样做，我们将其按特征向量的维度展开，使用上标 $i$ 来表示特征向量的第 $i$ 维，那么公式可以写成下面的形式：

$$\begin{array}{rl}{Y}_t^i& =⟨a_t^i,V^i⟩\\ a_t^i& =\frac{{\sigma }_q(Q_t^i)\odot \exp (K_{t^{\prime }}^i+w_{t,t^{\prime }})}{\sum _{t^{\prime }=1}^T\exp (K_{t^{\prime }}^i+w_{t,t^{\prime }})}\end{array}$$

然后 $Y_t=\text{Concat}(Y_t^1,Y_t^2,...,Y_t^d)$ 。可以发现上面的式子还可以进一步写成下面的形式：

$$\begin{array}{rl}{Y}_t^i& =g(K^i,Q_t^i)V\\ Y_t& =\text{Concat}(Y_t^1,Y_t^2,...,Y_t^d)\end{array}$$

下面是MHA的公式：

$$\begin{array}{rl}{f}_i(X)& =\text{Score}(Q,K)V\\ f(X)& =\text{Concat}(f_1(X),f_2(X),...,f_n(X))\end{array}$$

可以看出AFT实际上类似一个逐特征通道进行的MHA。

AFT的计算复杂度

在公式 $Y_t^i=g(K^i,Q_t^i)V$ 中，计算一个 $Y_t^i$ 的时间复杂度为 $O(T^2)$ ，这是因为计算 $a_t^i$ 的复杂度是 $O(T)$ ，与 $V^i$ 相乘后就是 $O(T^2)$ 。因此总的时间复杂度就是 $O(T^{2}d)$ 。所以上面的原始AFT相当于Transformer在时间复杂度上没有改进。

但在空间复杂度上，Transformer需要计算一个 $T^2$ 大小的注意力矩阵，而在AFT中没有了这个操作，而是使用了类似 $Softmax$

AFT 变体

AFT-Local

论文作者通过实验，发现Transformer应用在图片上时更加注重局部性而忽略长程的注意力，因此作者提出了 AFT-local ，通过仅保留局部的注意力来降低计算的复杂度。

AFT-local 的原理是改造了 $w_{t,t^{\prime }}$:

$$w_{t,t^{\prime }}=\begin{array}{r}\{\begin{array}{ll}{w}_{t,t^{\prime }}& |t-t^{\prime }|<s\\ 0& otherwise\end{array}\end{array}$$

其中，$s$ 为注意力窗口的长度。在这种情况下，AFT的时间复杂度降为了 $O(Tsd)$

AFT-Free

当上面的 $s=1$ 时，就得到了AFT-Free,这是一个时间复杂度为线性的算法。