Transformers are RNNs: Fast Autoregressive Transformers with Linear Attention

Linear Transformer: 使用线性注意力实现快速自回归的Transformer.

paper：Transformers are RNNs: Fast Autoregressive Transformers with Linear Attention
arXiv：link

作者提出了一种通过“线性化”降低自注意力机制的计算复杂度的方法，并构造了一种自回归的Transformer结构，能够更快地实现长句子生成等任务。

标准的Attention首先将输入序列 $X=[x_1,...,x_n] \in \Bbb{R}^{n×d}$ ( $n$ 个维度为 $d$ 的特征向量，通常 $n>d$ )转换成查询矩阵 $Q$ ,键矩阵 $K$ ,值矩阵 $V$ ：

Q = X W_{q} \in R^{n \times d}, W_{q} \in R^{d \times d}

$Q = XW_q \in \Bbb{R}^{n×d}, \quad W_q \in \Bbb{R}^{d×d}$

K = X W_{k} \in R^{n \times d}, W_{k} \in R^{d \times d}

$K = XW_k \in \Bbb{R}^{n×d}, \quad W_k \in \Bbb{R}^{d×d}$

V = X W_{v} \in R^{n \times d}, W_{v} \in R^{d \times d}

$V = XW_v \in \Bbb{R}^{n×d}, \quad W_v \in \Bbb{R}^{d×d}$

并通过下式计算自注意力，对于第 $i$ 个输入 $x_i$ ，其输出计算为:

(softmax (\frac{Q K^{T}}{\sqrt{d}}) V)_{i} = \frac{\sum_{j = 1}^{n} e^{\frac{q_{i}^{T} k_{j}}{\sqrt{d}}} v_{j}}{\sum_{j = 1}^{n} e^{\frac{q_{i}^{T} k_{j}}{\sqrt{d}}}}

$(\text{softmax}(\frac{QK^T}{\sqrt{d}})V)_i=\frac{\sum_{j=1}^{n}e^{\frac{q_i^Tk_j}{\sqrt{d}}}v_j}{\sum_{j=1}^{n}e^{\frac{q_i^Tk_j}{\sqrt{d}}}}$

上式计算中矩阵乘法 $QK^T$ 会引入 $O(n^2)$ 计算复杂度。

一般地，引入相似度函数 $\text{sim}(\cdot,\cdot)≥0$ ，则Attention也可表示为一般形式：

Attention (Q, K, V)_{i} = \frac{\sum_{j = 1}^{n} sim (q_{i}, k_{j}) v_{j}}{\sum_{j = 1}^{n} sim (q_{i}, k_{j})}

$\text{Attention}(Q,K,V)_i=\frac{\sum_{j=1}^{n}\text{sim}(q_i,k_j)v_j}{\sum_{j=1}^{n}\text{sim}(q_i,k_j)}$

注意到标准的Attention计算相当于选择了相似度函数：

sim (q_{i}, k_{j}) = e^{\frac{q_{i}^{T} k_{j}}{\sqrt{d}}}

$\text{sim}(q_i,k_j) = e^{\frac{q_i^Tk_j}{\sqrt{d}}}$

若把相似度函数看作核函数，即 $\text{sim}(q_i,k_j)=\phi(q_i)^T\phi(k_j)$ ，则有：

Attention (Q, K, V)_{i} = \frac{\sum_{j = 1}^{n} ϕ (q_{i})^{T} ϕ (k_{j}) v_{j}}{\sum_{j = 1}^{n} ϕ (q_{i})^{T} ϕ (k_{j})} = \frac{ϕ (q_{i})^{T} \sum_{j = 1}^{n} ϕ (k_{j}) v_{j}^{T}}{ϕ (q_{i})^{T} \sum_{j = 1}^{n} ϕ (k_{j})}

$\text{Attention}(Q,K,V)_i=\frac{\sum_{j=1}^{n}\phi(q_i)^T\phi(k_j)v_j}{\sum_{j=1}^{n}\phi(q_i)^T\phi(k_j)}=\frac{\phi(q_i)^T\sum_{j=1}^{n}\phi(k_j)v_j^T}{\phi(q_i)^T\sum_{j=1}^{n}\phi(k_j)}$

注意到通过上述转换，将计算复杂度从 $O(n^2)$ 降为 $O(n)$ ，即循环内从两次乘法减少为一次乘法。本文选择的 $\phi$ 如下：

ϕ (x) = elu (x) + 1, elu (x) = {\begin{cases} x, & x > 0 \\ e^{x} - 1, & x \leq 0 \end{cases}

$\phi(x)=\text{elu}(x)+1 , \quad \text{elu}(x)=\begin{cases} x, & x>0 \\ e^x-1, & x≤0 \end{cases}$

使用上述线性注意力(linear attention)构建Transformer时需要注意，由于Transformer采用语言模型的训练策略，因此需要引入mask，即mask掉未来的输入信息。实践中只需将求和 $\sum_{j=1}^{n}$ 替换为 $\sum_{j=1}^{i}$ ：