去掉Attention的Softmax，复杂度降为O(n)

众所周知，尽管基于Attention机制的Transformer类模型有着良好的并行性能，但它的空间和时间复杂度都是 $\mathcal{O}(n^2)$ 级别的， $n$ 是序列长度，因此当 $n$ 比较大时Transformer模型的计算量难以承受。近来，也有很多工做致力于下降Transformer模型的计算量，好比模型剪枝、量化、蒸馏等精简技术，又或者修改Attention结构，使得其复杂度能下降到 $\mathcal{O}(nlog⁡n)$ 甚至 $\mathcal{O}(n)$ php

论文《Transformers are RNNs: Fast Autoregressive Transformers with Linear Attention》当中提到一种线性化Attention（Linear Attention）的方法，由此引起了个人兴趣，继而阅读了一些相关博客，有一些不错的收获，最后将本身对线性化Attention的理解汇总在此文中html

Attention

当前最流行的Attention机制当属Scaled-Dot Attention，即python

\begin{aligned}Attention(\boldsymbol{Q},\boldsymbol{K},\boldsymbol{V}) = softmax\left(\boldsymbol{Q}\boldsymbol{K}^{\top}\right)\boldsymbol{V}\tag{1}\end{aligned}

这里的 $\boldsymbol{Q}\in \mathbb{R}^{n\times d_k}, \boldsymbol{K}\in \mathbb{R}^{m\times d_k}, \boldsymbol{V}\in \mathbb{R}^{m\times d_v}$ ，简单起见我就没显示的写出Attention的缩放因子 $\frac{1}{\sqrt{d}}$ 了。本文咱们主要关心Self Attention的场景，因此为了介绍上的方便，统一设 $\boldsymbol{Q},\boldsymbol{K},\boldsymbol{V}\in \mathbb{R}^{n\times d}$ markdown

摘掉Softmax

读者也许想不到，制约Attention性能的关键因素，实际上是定义里边的Softmax！事实上，简单地推导一下就能够获得这个结论。 $QK^T$ 这一步咱们获得一个 $n\times n$ 的矩阵，以后还要作一个Softmax网络

对一个 $1\times n$ 的行向量进行Softmax，时间复杂度是 $O(n)$ ，可是对一个 $n\times n$ 矩阵的每一行作一个Softmax，时间复杂度就是 $O(n^2)$ app

若是没有Softmax，那么Attention的公式就变为三个矩阵连乘 $\boldsymbol{QK^{\top}V}$ ，而矩阵乘法是知足结合率的，因此咱们能够先算 $\boldsymbol{K^{\top}V}$ ，获得一个 $d\times d$ 的矩阵（这一步的时间复杂度是 $O(d^2n)$ ），而后再用 $Q$ 左乘它（这一步的时间复杂度是 $O(d^2n)$ ），因为 $d \ll n$ ，因此这样算大体的时间复杂度只是 $O(n)$ ide

对于BERT base来讲， $d=64$ 而不是768，why？由于768其实是经过Multi-Head拼接获得的，而每一个head的 $d=64$ svg

也就是说，去掉Softmax的Attention复杂度能够降到最理想的线性级别 $\mathcal{O}(n)$ ！这显然就是咱们的终极追求：Linear Attention函数

通常的定义

问题是，直接去掉Softmax还能算是Attention吗？他还能有标准的Attention的效果吗？为了回答这个问题，咱们先将Scaled-Dot Attention的定义等价的改写为（本文的向量都是列向量）oop

\begin{aligned}Attention(\boldsymbol{Q},\boldsymbol{K},\boldsymbol{V})_i = \frac{\sum\limits_{j=1}^n e^{\boldsymbol{q}_i^{\top}\boldsymbol{k}_j}\boldsymbol{v}_j}{\sum\limits_{j=1}^n e^{\boldsymbol{q}_i^{\top}\boldsymbol{k}_j}}\tag{2}\end{aligned}

这里稍微解释下，首先咱们知道 $\boldsymbol{Q},\boldsymbol{K}\in \mathbb{R}^{n\times d}$ ，令 $\boldsymbol{M} = \boldsymbol{Q}\times \boldsymbol{K^{\top}}$ ，由矩阵乘法法则可知， $\boldsymbol{M}$ 的第一行是由 $\boldsymbol{Q}$ 的第一行乘以 $\boldsymbol{K^{\top}}$ 的全部列获得的

$Attention(\boldsymbol{Q},\boldsymbol{K},\boldsymbol{V})_i$ 表示最终输出结果矩阵的第 $i$ 行

$\boldsymbol{q}_i^{\top}$ 表示 $\boldsymbol{Q}\in \mathbb{R}^{n\times d}$ 矩阵的第 $i$ 行（行向量）

$\boldsymbol{k}_j$ 表示 $\boldsymbol{K^{\top}}\in \mathbb{R}^{d\times n}$ 矩阵的第 $j$ 列（列向量）

$\boldsymbol{v}_j$ 表示 $V^{\top}\in \mathbb{R}^{d\times n}$ 矩阵的的第 $j$ 列（列向量）

因此，Scaled-Dot Attention其实就是以 $e^{\boldsymbol{q}_i^{\top}\boldsymbol{k}_j}$ 为权重对 $\boldsymbol{v}_j$ 作加权平均。因此咱们能够提出一个Attention的通常化定义

\begin{aligned}Attention(\boldsymbol{Q},\boldsymbol{K},\boldsymbol{V})_i = \frac{\sum\limits_{j=1}^n \text{sim}(\boldsymbol{q}_i, \boldsymbol{k}_j)\boldsymbol{v}_j}{\sum\limits_{j=1}^n \text{sim}(\boldsymbol{q}_i, \boldsymbol{k}_j)}\tag{3}\end{aligned}

也就是把 $e^{\boldsymbol{q}_i^{\top}\boldsymbol{k}_j}$ 换成 $\boldsymbol{q}_i,\boldsymbol{k}_i$ 的通常函数 $\text{sim}(\boldsymbol{q}_i,\boldsymbol{k}_j)$ ，为了保留Attention类似的分布特性，咱们要求 $\text{sim}(\boldsymbol{q}_i, \boldsymbol{k}_j)\geq 0$ 恒成立。也就是说，咱们若是要定义新的Attention，必需要保留式(3)的形式，而且知足 $\text{sim}(\boldsymbol{q}_i, \boldsymbol{k}_j)\geq 0$

这种通常形式的Attention在CV中也被称为Non-Local网络，出自论文《Non-local Neural Networks》

几个例子

若是直接去掉Softmax，那么就是 $\text{sim}(\boldsymbol{q}_i, \boldsymbol{k}_j) = \boldsymbol{q}_i^{\top}\boldsymbol{k}_j$ ，问题是内积没法保证非负性，因此这还不是一个合理的选择。下面咱们介绍几种可取的方案

值得一提的是，下面介绍的这几种Linear Attention，前两种来自CV领域，第三种是苏剑林大佬构思的（除了下面的介绍外，还有EMANet等CV领域对Attention的改进工做）

核函数形式

一个天然的想法是：若是 $\boldsymbol{q}_i, \boldsymbol{k}_j$ 的每一个元素都是非负的，那么内积天然也是非负的。为了完成这点，咱们能够给 $\boldsymbol{q}_i, \boldsymbol{k}_j$ 各自加个激活函数 $\phi,\varphi$ ，即

\begin{aligned}\text{sim}(\boldsymbol{q}_i, \boldsymbol{k}_j) = \phi(\boldsymbol{q}_i)^{\top} \varphi(\boldsymbol{k}_j)\tag{4}\end{aligned}

其中 $\phi(\cdot), \varphi(\cdot)$ 是值域非负的激活函数。本文开头提到的论文《Transformers are RNNs》选择的是 $\phi(x)=\varphi(x)=\text{elu}(x)+1$ ，其中

\text{elu}(x)=\begin{cases}x& \text{if} \ x>0\\ \alpha (e^x-1) & \text{if}\ x<0\end{cases}

常见的 $\alpha$ 取值为 $[0.1, 0.3]$

非要讲故事的话，式(4)能够联想到"核方法"，尤为是 $\phi=\varphi$ 时， $\phi$ 就至关于一个核函数，而 $\langle \phi(\boldsymbol{q}_i), \phi(\boldsymbol{k}_j)\rangle$ 就是经过核函数所定义的内积。这方面的思考能够参考论文《Transformer dissection: An unified understanding for transformer’s attention via the lens of kernel》，此处不作过多延伸

妙用Softmax

另外一篇更早的文章《Efficient Attention: Attention with Linear Complexities》则给出了一个更有意思的选择。它留意到在 $\boldsymbol{QK^{\top}}$ 中， $\boldsymbol{Q},\boldsymbol{K}\in \mathbb{R}^{n\times d}$ ，若是“ $\boldsymbol{Q}$ 在 $d$ 那一维是归一化的，而且 $\boldsymbol{K}$ 在 $n$ 那一维是归一化的”，那么 $\boldsymbol{QK^{\top}}$ 就是自动知足归一化了，因此它给出的选择是

\begin{aligned}Attention(\boldsymbol{Q},\boldsymbol{K},\boldsymbol{V}) = softmax_2\left(\boldsymbol{Q}\right)softmax_1(\boldsymbol{K})^{\top}\boldsymbol{V}\tag{5}\end{aligned}

其中 $softmax_1$ 、 $softmax_2$ 分别表示在第一个 $(n)$ 、第二个维度 $(d)$ 进行Softmax运算。也就是说，这时候咱们是各自给 $\boldsymbol{Q},\boldsymbol{K}$ 加Softmax，而不是算完 $\boldsymbol{QK^{\top}}$ 以后再加Softmax

其实能够证实这个形式也是式(4)的一个特例，此时对应于 $\phi(\boldsymbol{q}_i)=softmax(\boldsymbol{q}_i),\varphi(\boldsymbol{k}_j)=e^{\boldsymbol{k}_j}$ ，读者能够自行推导一下

苏神的构思

在这里，苏神给出了一种构思。这个构思的出发点再也不是式(4)，而是源于咱们对原始定义(2)的泰勒展开。由泰勒展开咱们有

\begin{aligned}e^{\boldsymbol{q}_i^{\top}\boldsymbol{k}_j} \approx 1 + \boldsymbol{q}_i^{\top}\boldsymbol{k}_j\tag{6}\end{aligned}

若是 $\boldsymbol{q}_i^{\top}\boldsymbol{k}_j\geq -1$ ，那么就能够保证右端的非负性，从而可让 $\text{sim}(\boldsymbol{q}_i, \boldsymbol{k}_j)=1 + \boldsymbol{q}_i^{\top}\boldsymbol{k}_j$ 。到这里读者可能已经想到了，想要保证 $\boldsymbol{q}_i^{\top}\boldsymbol{k}_j\geq -1$ ，只须要分别对 $\boldsymbol{q}_i,\boldsymbol{k}_j$ 作 $l_2$ 归一化。因此，苏神最终提出的方案就是：

\begin{aligned}\text{sim}(\boldsymbol{q}_i, \boldsymbol{k}_j) = 1 + \left( \frac{\boldsymbol{q}_i}{\Vert \boldsymbol{q}_i\Vert}\right)^{\top}\left(\frac{\boldsymbol{k}_j}{\Vert \boldsymbol{k}_j\Vert}\right)\tag{7}\end{aligned}

若 $\boldsymbol{x}=[x_1,x_2,...,x_n]$ ，则 $\Vert x\Vert=\sqrt{x_1^2+x_2^2+···+x_n^2}$

这不一样于式(4)，但理论上它更加接近原始的Scaled-Dot Attention

实现

这里主要是针对苏神所提出的方法进行实现，可是因为笔者本人水平有限，所以最终实现的代码当中其实存在一些问题，主要是：

从测试结果来看，改进后的计算速度并无提高
没法作到求和为1

代码实现主要是针对BERT的PyTorch实现这篇文章的代码，更具体的说，其实仅修改了ScaledDotProductAttention这个函数，所以下面只放出这部分代码

class ScaledDotProductAttention(nn.Module):
    def __init__(self):
        super(ScaledDotProductAttention, self).__init__()

    def forward(self, Q, K, V, attn_mask):
        Q = F.normalize(Q, dim=3)
        K = F.normalize(K, dim=3)
        M = (torch.ones(Q.shape[0], Q.shape[1], Q.shape[2], K.shape[2]) + torch.matmul(Q, K.transpose(-1, -2))) # scores : [batch_size, n_heads, seq_len, seq_len]
        M_sum = torch.sum(M, dim=3)
        M = M / M_sum.unsqueeze(3).repeat(1, 1, 1, M.shape[3])
        attn = M.masked_fill(attn_mask, 0) # Fills elements of self tensor with value where mask is one.
        context = torch.matmul(attn, V)
        return context
复制代码

若是您有更好的实现方法，还望不吝赐教

Reference

线性Attention的探索：Attention必须有个Softmax吗？