摘掉Softmax

制约Attention性能的关键因素，其实是定义里边的Softmax！事实上，简单地推导一下就可以得到这个结论。 $Q K^{⊤}$ 这一步我们得到一个 $n \times n$ 的矩阵，就是这一步决定了Attention的复杂度是 $O (n^{2})$ ；如果没有Softmax，那么就是三个矩阵连乘 $Q K^{⊤} V$ ，而矩阵乘法是满足结合率的，所以我们可以先算 $K^{⊤} V$ ，得到一个 $d \times d$ 的矩阵，然后再用 $Q$ 左乘它，由于 $d ≪ n$ ，所以这样算大致的复杂度只是 $O (n)$ （就是 $Q$ 左乘那一步占主导）。

也就是说，去掉Softmax的Attention的复杂度可以降到最理想的线性级别 $O (n)$ ！这显然就是我们的终极追求：Linear Attention，复杂度为线性级别的Attention。所以，本文的主题就是探究摘掉Softmax后的线形Attention。

问题是，直接去掉Softmax还能算是Attention吗？它还能有标准的Attention的效果吗？

几个例子

根据式3 如果直接去掉Softmax，那么就是 $sim (q_{i}, k_{j}) = q_{i}^{⊤} k_{j}$ ，问题是内积无法保证非负性，所以这还不是一个合理的选择。下面我们简单介绍几种可取的方案。

值得指出的是，下面介绍的这几种Linear Attention，前两种来自CV领域，第三种是苏神自己构思的，所以都还没有在NLP任务上做过什么实验。

核函数形式

一个自然的想法是：如果 $q_{i}, k_{j}$ 的每个元素都是非负的，那么内积自然也就是非负的。为了完成这点，我们可以给 $q_{i}, k_{j}$ 各自加个激活函数 $ϕ, φ$ ，即

sim (q_{i}, k_{j}) = ϕ (q_{i})^{⊤} φ (k_{j})

其中 $ϕ (\cdot), φ (\cdot)$ 是值域非负的激活函数。在论文Transformers are RNNs中选择的是 $ϕ (x) = φ (x) = elu (x) + 1$ 。

非要讲故事的话，式4 可以联想到“核方法（kernal method）”，尤其是 $ϕ = φ$ 时 $ϕ$ 就相当于一个核函数，而 $⟨ ϕ (q_{i}), ϕ (k_{j})⟩$ 就是通过核函数所定义的内积。这方面的思考可以参考论文原文，此处不做过多延伸。

妙用Softmax

另一篇更早的文章《Efficient Attention: Attention with Linear Complexities》则给出了一个更有意思的选择。它留意到在 $Q K^{⊤}$ 中， $Q, K, \in R^{n \times d}$ ，如果“ $Q$ 在 $d$ 那一维是归一化的、并且 $K$ 在 $n$ 那一维是归一化的”，那么 $Q K^{⊤}$ 就是自动满足归一化了，所以它给出的选择是：

A tt e n t i o n (Q, K, V) = so f t ma x_{2} (Q) so f t ma x_{1} (K)^{⊤} V

其中 $so f t ma x_{1}$ 、 $so f t ma x_{2}$ 分别指在第一个（ $n$ ）、第二个维度（ $d$ ）进行Softmax运算。也就是说，这时候我们是各自给 $Q, K$ 加Softmax，而不是 $Q K^{⊤}$ 算完之后才加Softmax。

如果直接取 $ϕ (q_{i}) = so f t ma x (q_{i}), φ (k_{j}) = so f t ma x (k_{j})$ ，那么很显然这个形式也是式4 的一个特例。另外这个设计在CV中出现过不止一次，比如A^2-Nets也包含了同样的做法。

自己的构思

在这里，苏神给出自己的一种构思。这个构思的出发点不再是式4，而是源于我们对原始定义2 的近似。由泰勒展开我们有

e^{q_{i}^{⊤} k_{j}} \approx 1 + q_{i}^{⊤} k_{j}

如果 $q_{i}^{⊤} k_{j} \geq - 1$ ，那么就可以保证右端的非负性，从而可以让 $sim (q_{i}, k_{j}) = 1 + q_{i}^{⊤} k_{j}$ 。到这里读者可能已经想到了，想要保证 $q_{i}^{⊤} k_{j} \geq - 1$ ，只需要分别对 $q_{i}, k_{j}$ 做 $l_{2}$ 归一化。所以，笔者最终提出的方案就是：

sim (q_{i}, k_{j}) = 1 + (\frac{q _{i}}{∥ q _{i} ∥})^{⊤} (\frac{k _{j}}{∥ k _{j} ∥})

这不同于形式4，但理论上它更加接近原始的Scaled-Dot Attention。

Linformer

跟本文所介绍的Linear Attention很相似的一个工作是Facebook的Linformer，它依然保留原始的Scaled-Dot Attention形式，但在进行Attention之前，用两个 $m \times n$ 的矩阵 $E, F$ 分别对 $K, V$ 进行投影，即变为

A tt e n t i o n (Q, K, V) = so f t ma x (Q (E K)^{⊤}) F V

这样一来， $Q (E K)^{⊤}$ 就只是一个 $n \times m$ 的矩阵，而作者声称对于哪怕对于很大的序列长度n，m也可以保持为一个适中的常数，从而这种Attention也是线性的。

但是，笔者认为“对于超长序列m可以保持不变”这个结论是值得质疑的，对于长序列原论文只做了MLM任务，而很明显MLM并不那么需要长程依赖，所以这个实验没什么说服力。因此，Linformer是不是真的Linear，还有待商榷。

自回归生成

Linformer的另一个缺点是 $E K, F V$ 这两个运算直接把整个序列的信息给“糅合”起来了，所以它没法简单地把将来信息给Mask掉（Causal Masking），从而无法做语言模型、Seq2Seq等自回归生成任务，这也是刚才说的原作者只做了MLM任务的原因。相比之下，前面介绍的几种Linear Attention都能做到这一点。以式3 和式4 为例，如果要Mask掉未来信息，那么只需要把求和 $j = 1 \sum n$ 改为 $j = 1 \sum i$ ：

A tt e n t i o n (Q, K, V)_{i} = \frac{j = 1 \sum i ( ϕ ( q _{i} ) ^{⊤} φ ( k _{j} ) ) v _{j}}{j = 1 \sum i ϕ ( q _{i} ) ^{⊤} φ ( k _{j} )} = \frac{ϕ ( q _{i} ) ^{⊤} j = 1 \sum i φ ( k _{j} ) v _{j}^{⊤}}{ϕ ( q _{i} ) ^{⊤} j = 1 \sum i φ ( k _{j} )}

实现上式有两种方式：第一方式是设 $S_{i} = j = 1 \sum i φ (k_{j}) v_{j}^{⊤}$ 以及 $z_{i} = j = 1 \sum i φ (k_{j})$ ，我们有

A tt e n t i o n (Q, K, V)_{i} = \frac{ϕ ( q _{i} ) ^{⊤} S _{i}}{ϕ ( q _{i} ) ^{⊤} z _{i}}, S_{i} = S_{i - 1} + φ (k_{i}) v_{i}^{⊤} z_{i} = z_{i - 1} + φ (k_{i})

这说明这种Attention可以作为一个RNN模型用递归的方式实现，它的空间复杂度最低，但是要串行计算，适合预测解码时使用；

第二种是直接将 $φ (K), V \in R^{n \times d}$ 做外积，得到一个 $n \times d \times d$ 的矩阵，然后对 n 那一维执行 $cumsum$ 运算，这样就一次性得到 $S_{1}, S_{2}, \dots, S_{n}$ 了，它的速度最快，但空间占用最大，适合训练时使用，不过很多时候都有 $d^{2} ≫ n$ ，一般情况下训练时都很难承受这个空间复杂度，因此多数还是用RNN形式。

但是如果用这种最基本的循环来训练有问题。

如果使用autograd，那么，每一个时刻的hidden state $S_{t}$ 都会被保存下来给反向传播算梯度。保存中间hidden state在之前的RNN训练里面没啥问题，但是在这里却是一个大问题。为什么呢？因为传统RNN 每个时刻的hidden state的大小是 $O (d)$ ，而linear attention每个时刻的hidden state大小是 $O (d^{2})$ ，如果全部存下来需要 $O (L d^{2})$ ，那是万万不可接受的（你可以想象，输入的QKV才 $3 L d$ , 一般 $d ≫ 3$ , 那么你基本就显存爆炸寄了不用玩了）

Transformers are RNNs 里面提出了一个memory-efficient 训练的解决方案。

在这个方案里， $S_{t}$ 在前向传播的过程中是不会被保存的。由于只有 $q_{t}$ 的梯度依赖于 $S_{t}$ 。那我们索性在backward的时候重新算一遍 $S_{t}$ ，但这时的output算的不是 $o_{t}$ ，而是 $dq_{t} = S_{t} do_{t}$

而 $k_{t}, v_{t}$ 的梯度不依赖于 $S_{t}$ ，可以直接通过平常的BPTT来算出来。

小结

上面介绍了一些从结构上对Attention进行修改从而降低其计算复杂度的工作，其中最主要的idea是去掉标准Attention中的Softmax，就可以使得Attention的复杂度退化为理想的 $O (n)$ 级别（Linear Attention）。相比于其他类似的改进结构的工作，这种修改能在把复杂度降到 $O (n)$ 的同时，依然保留所有的“token-token“的注意力，同时还能保留用于做自回归生成的可能性。

Reference

🔖 https://spaces.ac.cn/archives/7546

#

!