论文阅读笔记——Generating Long Sequences with Sparse Transformers

Sparse Transformer 论文
解决了 Transformer 在长序列建模时的计算开销和内存过大的问题。
可视化了一个 128 层自注意力在 CIFAR-10 的数据集上学习到的注意力模式，发现：1）稀疏性普遍存在：大多数层在多数数据点上表现出稀疏注意力；2）例外：部分层想要捕捉全局依赖关系。Transformer 的注意力机制呈现了和卷积模型类似的归纳偏置，即浅层的网络倾向于提取纹理信息，深层的网络倾向于提取语义信息。

分解自注意力(Factorized self-attention)

Local 自注意力只关注自身相邻的，其余设为 0，类似于卷积；Atrous 自注意力是跳着计算，类似膨胀卷积；一种简单思路是交替使用 Local 自注意力和 Atrous 自注意力。但 OpenAI 并没有这么做，而是将二者合为一。

由于 Transformer 的最复杂的计算是 $Q{K}^T$，稀疏注意力是让设置好的像素点参与注意力的计算。由此，引入了连接模式的变量 S = { S 1 , … … , S n } S=\{S_1,……,S_n\} S={S1​,……,Sn​}。其中 $S_i$ 是在预测第 i 个时间片的索引，是一个由 0 和 1 组成的二维矩阵。
Attend ⁡ ( X , S ) = ( a ( x i , S i ) ) i ∈ { 1 , … , n } ( 2 ) a ( x i , S i ) = softmax ⁡ ( ( W q x i ) K S i T d ) V S i ( 3 ) K S i = ( W k x j ) j ∈ S i V S i = ( W v x j ) j ∈ S i ( 4 ) \begin{aligned} \operatorname{Attend}(X, S) = \left(a(\mathbf{x}_i, S_i)\right)_{i \in \{1, \ldots, n\}} \quad (2) \\a(\mathbf{x}_i, S_i) = \operatorname{softmax}\left(\frac{(W_q \mathbf{x}_i) K_{S_i}^T}{\sqrt{d}}\right) V_{S_i} \quad (3) \\K_{S_i} = \left(W_k \mathbf{x}_j\right)_{j \in S_i} \quad V_{S_i} = \left(W_v \mathbf{x}_j\right)_{j \in S_i} \quad (4) \end{aligned} Attend(X,S)=(a(xi​,Si​))i∈{1,…,n}​(2)a(xi​,Si​)=softmax(d ​(Wq​xi​)KSi​T​​)VSi​​(3)KSi​​=(Wk​xj​)j∈Si​​VSi​​=(Wv​xj​)j∈Si​​(4)​
其中 $W_q,W_k,W_v$ 是计算 Query，Key，Value 三个向量的权值矩阵。稀疏 Transformer 通过让链接模式作用到 $K^T$ 上，从而降低 $Q{K}^T$ 的复杂度

跨步注意力（Stride Attention） 由两种形式的连接模式组成。假设步长 $l$，行注意力是当前时间片的前 $l$ 个时间片的值为 1，其余为 0；列注意力是每隔 $l$ 个时间片段值为 1， 其余为 0。行注意力和列注意力的表达式如下，复杂度均为 $O(\sqrt{n})$：
A i ( 1 ) = { t , t + 1 , t + 2 , … … , i } , w h e r e t = m a x ( 0 , i − l ) A i ( 2 ) = { j : ( i − j ) m o d    l = 0 } \begin{aligned} A_i^{(1)}=\{t,t+1,t+2,……,i\},where\quad t = max(0,i-l) \\A_i^{(2)}=\{j:(i-j)\mod l =0\} \end{aligned} Ai(1)​={t,t+1,t+2,……,i},wheret=max(0,i−l)Ai(2)​={j:(i−j)modl=0}​
固定注意力（Fixed Attention） 也有行注意力和列注意力组成：
A i ( 1 ) = { j : ( [ j / l ] = [ i / l ] ) } A i ( 2 ) = { j : j m o d    l ∈ { t , t + 1 , … … , l } } \begin{aligned} A_i^{(1)}=\{j:([j/l]=[i/l])\} \\A_i^{(2)}=\{j:j\mod l \in\{t,t+1,……,l\}\} \end{aligned} Ai(1)​={j:([j/l]=[i/l])}Ai(2)​={j:jmodl∈{t,t+1,……,l}}​
将以上注意力核融入网络中：

• 每个残差块使用不同的注意力类型：$attention(X)=W_p\cdot attend(X,A^{(r\mathrm{mod}p)})$ 其中 r 是当前残差块的缩影，p 是注意力核的类别数；
• 每个注意力头计算所有类型注意力核，合并他们的结果：$attention(X)=W_p\cdot attend(X,{\cup }_{m=1}^p{A}^{(m)}$
• 对于多头注意力，每个头选择一个注意力核，合并结果： a t t e n t i o n ( X ) = W p ( a t t e n d ( X , A ) ( i ) ) i ∈ { 1 , … … , n h } attention(X)=W_p(attend(X,A)^{(i)})_{i\in\{1,……,n_h\}} attention(X)=Wp​(attend(X,A)(i))i∈{1,……,nh​}​ 其中 $n_h$ 组不同注意力核并行计算，然后在特征维度拼接。

多层 Transformer 训练

作者使用了在 ResNet v2 中提出的激活前置的残差模块，一个 $N$ 层的网络可以表示为：
$$\begin{array}{r}{H}_0=embed(X,W_e)\\ H_k=H_{k-1}+resblock(H_{k-1})\\ y=softmax(norm(H_N)W_{out})\end{array}$$
其中 embed 是可学习的嵌入层：$embed(X,W_e)=\left({\mathbf{x}}_i{W}_e+{\sum }_{j=1}^{n_{emb}}{\mathbf{o}}_i^{(j)}{W}_j\right)$ 其中 $n_{emb}$ 的值为 $d_{data}$ 或 $d_{attn}$，${\mathbf{x}}_i$ 是序列中第 i 个元素的 one-hot 编码，${\mathbf{o}}_i^{(j)}$ 是 ${\mathbf{x}}_i$ 在第 $j$ 维特征上的 one-hot 编码。
resblock(h) 由一个注意力模块和一个前馈神经网络组成：
$$\begin{array}{c}a(H)=\text{dropout( attention }(\mathrm{norm}(H)))\\ b(H)=\mathrm{dropout}(\mathrm{ff}(\mathrm{norm}(H+a(H))))\\ \mathrm{resblock}(H)=a(H)+b(H)\end{array}$$

梯度检查点

一个以时间换空间的一个策略，在反向传播的过程中，不是保存所有节点的参数值，而是只保留部分关键节点的值，然后通过这些关键节点反向推出其他节点的值。这样虽然引入了额外的节点参数的计算工作，但是大大节约了显存，从而使得训练更长的序列成为可能。

实验结果