Transformers Assemble（PART I ）

写在前面

经过之前一段时间的 NLP Big Bang，现在相对比较平静了，Transformer 派已经占据了绝对的主导地位，在各类应用中表现出色。看标题大家也可以猜个差不多，整理了一系列自《Attention is all you need》之后的对 Vanilla Transformer 的改进论文，和大家一起梳理整个发展过程。这篇是第一趴，都来自ICLR。OK，来看看今天的 Transformers：

Bi-BloSAN from UTS，ICLR2018
Universal Transformers from UVA&Google，ICLR2019
Reformer from Google，ICLR2020

BI-DIRECTIONAL BLOCK SELF-ATTENTION FOR FASTAND MEMORY-EFFICIENT SEQUENCE MODELING

这篇论文首先分析了目前几大类特征抽取器CNN、RNN、Attention的优缺点，针对其不足提出了一种对self-attention的改进，双向分块自注意力机制（bidirectional block self-attention (Bi-BloSA)）。

Masked Block Self-Attention

其最主要的组件是掩码分块自注意力机制（masked block self-attention (mBloSA)），基本思想是将序列分成几个等长的block（必要时进行padding），对每个单独的block内应用self-attention（intra-block SAN）捕获局部特征，然后对所有block输出再应用self-attention（inter-block SAN）捕获全局特征。这样，每个attention只需要处理较短的序列，与传统的attention相比可以节省大量的内存。最后通过Context fusion模块将块内SAN、块间SAN与原始输入结合生成最终上下文表示。整体框架如下图所示：

Masked Self-Attention

在传统的self-attention中时序信息是很难编码进去的，通过对某些位置进行掩码可以实现位置信息的融合。
$f(xi,xj)=c⋅tanh⁡([W(1)xi+W(2)xj+b(1)]/c)+Mij1f\left(x_{i}, x_{j}\right)=c \cdot \tanh \left(\left[W^{(1)} x_{i}+W^{(2)} x_{j}+b^{(1)}\right] / c\right)+M_{i j} 1$
而建模双向信息，只需要将上式最后一项修改即可得到前向和后向SAN
$M_{i j}^{f w}=\left\{\begin{array}{ll} {0,} & {i<j} \\ {-\infty,} & {\text { otherwise }} \end{array}\right.$
$M_{i j}^{b w}=\left\{\begin{array}{ll} {0,} & {i>j} \\ {-\infty,} & {\text { otherwise }} \end{array}\right.$

Bi-directional Block Self-Aattention Network

最终将上述模块整合后得到的网络框架如下图。

实验工作也非常充分，在九大公开数据集的多类NP任务中取得SOTA，同时模型的计算和内存更高效。

Reference

Code Here

UNIVERSAL TRANSFORMERS

上一篇论文针对的是Transformer内存占用大、对长序列输入不友好的缺陷，除此之外，这篇论文指出其还存在着以下几个问题

针对以上存在的不足，论文提出了一种Transformer的改进模型：Universal Transformer（UT），如下图所示。乍一看与Transformer非常相似（不同的地方在下图用红框框起来会在下文介绍），下面我们来仔细看看UT是如何解决上述缺陷的。

Parallel-in-Time Recurrence Mechanism

原始的Transformer堆叠层数是固定的（6层或12层），为了结合RNN的优势，对每一个token设置了一个Transition Function（参数共享），并且可以一次一次循环迭代计算（Transformer是对attention层后直接输入一个Dense Layer）。如下图，横坐标是输入序列的token，纵坐标是循环迭代次数Depth。每一次迭代是在整个序列上进行的，不同于RNN需要从左往右依次计算，大大提高了计算效率。当迭代次数Depth人为固定时，就变成了原始Transformer。

假设输入为长度为 $m$ 的序列，第 $t$ 次循环后有：
$(At))H^{t}=\text{LAYER NORM }\left(A^{t}+\text { TRANSITION }\left(A^{t}\right)\right)$ $(Ht−1+Pt))A^{t}=\text { LAYERNORM }\left(\left(H^{t-1}+P^{t}\right)+\text { MULTIHEADSELFATTENTION }\left(H^{t-1}+P^{t}\right)\right)$ 其中Transition Function可以是深度可分离卷积也可以是全连接。

Coodinate Embeddings

在上面公式中有一个 $P$ ，指的是两种embedding，（position，time）：

position embedding和原始Transformer的一样；
timestep embedding，其实本质上和position embedding一样，只不过这里是关注于循环步数 $t$ ，因为不同的循环次数学习的是不同层面的知识，这一信息可以提供给模型学习了几次，什么时候该停止；
$Pi,2jt=sin⁡(i/100002j/d)+sin⁡(t/100002j/d)Pi,2j+1t=cos⁡(i/100002j/d)+cos⁡(t/1000002j/d)\begin{aligned} P_{i, 2 j}^{t} &=\sin \left(i / 10000^{2 j / d}\right)+\sin \left(t / 10000^{2 j / d}\right) \\ P_{i, 2 j+1}^{t} &=\cos \left(i / 10000^{2 j / d}\right)+\cos \left(t / 100000^{2 j / d}\right) \end{aligned}$

Dynamic Halting

正如前面所说的，输入中的某些部分需要更多次的计算模型才可以理解。例如最常见的I went to the bank to deposit money里bank一词就需要更多次的运算。在UT中引入了Adaptive Computation Time (ACT)来实现，ACT可以基于模型在每步循环中预测的标量halting probability，动态调制处理每个输入token所需的计算步骤数（即响应时间ponder time）。UT将动态ACT暂停机制分别应用于每个位置，一旦特定的循环块停止，它将其状态复制到下一个步骤，直到所有块都停止，或者直到达到最大步数为止，编码器的最终输出是以此方式产生的最后一层表示。Dynamic Halting策略流程可以概括为：

在第 $t$ 步循环时，有第 $t$ 步的halting probability、第 $t - 1$ 步的token state以及提前设置的超参停止阈值；
利用UT计算当前token state；
利用全连接层和sigmoid激活函数计算ponder time，ponder time表示模型对每个token还需要多少次运算的预测；
判断每个token当前是否需要停止，规则为：
当 (halting probability + pondering) > halt threshold 时，停止；
当 (halting probability + pondering) ≤ halt threshold 时，继续运算；
对于那些仍然在运算的token，更新停止概率：halting probability += pondering；
当所有位置都停止或者模型达到循环最大步数时，结束。