Transformers Assemble（PART IV）

上一篇Transformers Assemble（PART III）重点在transformer位置信息的改进，这一集挑选了几篇都带有「Sparse」的标签，主要关注点在于transformer结构的复杂度问题。先来看看都有哪些：

Sparse Transformer from OpenAI，NAACL 2019
Adaptively Sparse Transformers，EMNLP2019
Explcit Sparse Transformer from PKU

Generating Long Sequences with Sparse Transformers

来自OpenAI的工作，同样关注于原始Transformer的计算复杂度问题，尤其是在面对长序列输入的情况。为此，论文中将full attention进行分解，通过多个sparse attention来代替，在不牺牲性能的情况下降低复杂度至 $\sqrt{n})$

)。

下图是三种不同的注意力形式，其中上半部分表示一个6x6图像的像素之间相互attend，下半部分表示对应的connectivity matrix。

（a）原始Transformer的full attention；
（b）Strided Attention：这种方式主要应用于图像或者音频，每一个位置通过attend其对应的行和列来获取信息，两个head的具体表示为：第一个head用于attend该位置前面的 $l$ 个位置，第二个head用于attend间隔 $l$ 的位置（如果输入是图像 $l$ 为图像的宽，则attend对应的列）：
$t=max⁡(0,i−l)A_{i}^{(1)}=\{t, t+1, \ldots, i\} \text { for } t=\max (0, i-l)$ $Ai(2)={j:(i−j)modl=0}A_{i}^{(2)}=\{j:(i-j) \bmod l=0\}$ （c）Fixed Attention：这种方式主要应用于像文本之类没有周期性的数据，首先将文本分成固定长度的块，然后第一个head处理该块中该位置之前的所有元素，第二个head处理每个块的最后一部分的固定大小的部分。

Other Tricks

上面就是Attention主要的改进，文中还涉及了一些其他的tricks。

pre-activation residual block
来自Identity mappings in deep residual networks可以使Transformer的训练更加容易
$(Hk−1)y=softmax⁡(norm⁡(HN)Wout)\begin{aligned} H_{0} &=\text { embed }\left(X, W_{e}\right) \\ H_{k} &=H_{k-1}+\text { resblock }\left(H_{k-1}\right) \\ y=& \operatorname{softmax}\left(\operatorname{norm}\left(H_{N}\right) W_{\text {out}}\right) \end{aligned}$ 其中resblock(H)的计算如下 $(H)=a(H)+b(H)\begin{aligned} a(H) &=\text { dropout }(\text { attention }(\operatorname{norm}(H))) \\ b(H) &=\operatorname{dropout}(\mathrm{ff}(\operatorname{norm}(H+a(H)))) \\ & \text { resblock }(H)=a(H)+b(H) \end{aligned}$
Gradient check-pointing
Efficient block-sparse attention kernels
Mixed-precision training

reference

Code Here
openai sparse-transformer blog

Adaptively Sparse Transformers

这篇论文也是对vanilla Transformer的改进，提出了Adaptively Sparse Transformers （AST），优化的两个关键就在其名字中：

Sparse： 通过替换Softmax函数为 $α−entmax\alpha-entmax$ 达到稀疏注意力；
Adaptively： 每个attention head都是模型可自动学习的；

作者指出与先前的sparse transformer（就是上面两个~）研究不同的是，他们的这一方法可以关注在不连续的输入集合，如下图：

Sparse Attention

softmax函数所有结果都不为0，并且最终所有元素之和为1，这样的特性决定了相对重要的部分的权值会“缩水”。这一方向的研究很多，作者选用了最近提出的alpha-entmax：
$(z):=argmax⁡p∈Δdp⊤z+HαT(p)\alpha \text { -entmax }(\boldsymbol{z}):=\underset{\boldsymbol{p} \in \Delta^{d}}{\operatorname{argmax}} \boldsymbol{p}^{\top} \boldsymbol{z}+\mathrm{H}_{\boldsymbol{\alpha}}^{\mathrm{T}}(\boldsymbol{p})$

$HαT(p):={1α(α−1)∑j(pj−pjα),α≠1−∑jpjlog⁡pj,α=1\mathrm{H}_{\alpha}^{\mathrm{T}}(\boldsymbol{p}):=\left\{\begin{array}{ll} {\frac{1}{\alpha(\alpha-1)} \sum_{j}\left(p_{j}-p_{j}^{\alpha}\right),} & {\alpha \neq 1} \\ {-\sum_{j} p_{j} \log p_{j},} & {\alpha=1} \end{array}\right.$

AST

对于Transformer类模型的功能至关重要的是，不同的head会捕获不同的语言现象，这让我们想到对于不同的head，使用不同的 $α\alpha$ 值，使其自适应地让一些head稀疏化，一些head更接近softmax。利用上面的 $α−entmax\alpha-entmax$ 替换原始的softmax函数后，将 $α\alpha$ 看成是与其他网络参数意义的可学习参数，通过随机梯度进行优化。但是通过梯度方法对其自动优化并不容易，然后作者在下面就开始一系列数学推导。。。

TODO

我实在是看不动了。。。
后面的实验和分析也非常有意思的。。。
大家记得看，我先溜了。。。
PS. 在油管上发现了作者的分享视频，放在reference里

Reference

Code Here
作者分享视频

Explicit Sparse Transformer: Concentrated Attention Through Explicit Selection

Motivation和上一篇论文一样，如下图，对于文本I thanked him with all my heart, and I asked him, 'why are you helping me?'，vanilla Transformer（蓝色标记）会对所有元素都有注意，而噪音的注意力会对效果产生影响；新提出的显式稀疏注意力机制（橙色标记）只会关注文本的 $t o p k$ 个attention score最大的元素，从而移除无关信息。

具体实现也非常简单易于实现，且不会增加额外的内存和计算开销。

沿用vanilla transformer的attention计算公式得到attention score，
$P=QKTdP=\frac{Q K^{\mathrm{T}}}{\sqrt{d}}$
假定分值越大的元素其相关性越大，计算Masking矩阵。找出 $P$ 中每行的 $k$ 个最大元素，记录其位置，并得到一个threshold vector， $t=[t1,t2,⋯,tlQ]t=\left[t_{1}, t_{2}, \cdots, t_{l_{Q}}\right]$
将Making矩阵应用到原始 $P$ 矩阵上，
$i)\mathcal{M}(P, k)_{i j}=\left\{\begin{array}{ll} {P_{i j}} & {\text { if } P_{i j} \geq t_{i}(k \text { -th largest value of row } i)} \\ {-\infty} & {\text { if } P_{i j}<t_{i}(k \text { -th largest value of row } i)} \end{array}\right.$ 反向传播时，
$i)\begin{aligned} &\frac{\partial M_{i j}}{\partial P_{k l}}=0 \quad(i \neq k \text { or } j \neq l)\\ &\frac{\partial M_{i j}}{\partial P_{i j}}=\left\{\begin{array}{ll} {1} & {\text { if } P_{i j} \geq t_{i}(k-\text { th largest value of row } i)} \\ {0} & {\text { if } P_{i j}<t_{i}(k-\text { th largest value of row } i)} \end{array}\right. \end{aligned}$
归一化，
$A=softmax⁡(M(P,k))A=\operatorname{softmax}(\mathcal{M}(P, k))$ 反向传播时，
$i)\begin{aligned} \frac{\partial A_{i j}}{\partial P_{k l}} &=\sum_{m=1}^{l_{Q}} \sum_{n=1}^{l_{K}} \frac{\partial A_{i j}}{\partial M_{m n}} \frac{\partial M_{m n}}{\partial P_{k l}} \\ &=\frac{\partial A_{i j}}{\partial M_{k l}} \frac{\partial M_{k l}}{\partial P_{k l}} \\ &=\left\{\begin{array}{cl} {\frac{\partial A_{i j}}{\partial M_{k l}}} & {\text { if } P_{i j} \geq t_{i}(k \text { -th largest value of row } i)} \\ {0} & {\text { if } P_{i j}<t_{i}(k \text { -th largest value of row } i)} \end{array}\right. \end{aligned}$
输出表示，
$C = A V$

整体流程如下最右图所示，

另外，参数 $k$ 的选择至关重要，当 $k$ 取与序列长度一致时即为vanilla transformer。作者在NMT实验中发现当 $k = 8$ 时效果最好。

reference

Code Here
Open Review