EMNLP杰出论文 | 当注意力遇到RNN，五倍以上训练加速！

作者 | 高斯定理 欢迎关注知乎"高斯定理"和专栏"边走边看" 整理 | NewBeeNLP

今天分享EMNLP杰出论文（outstanding paper）—当注意力遇到RNN

• 论文：When Attention Meets Fast Recurrence: Training Language Models with Reduced Compute
• 地址：https://arxiv.org/pdf/2102.12459.pdf

由于众所周知的原因，当前训练语言模型的开销越来越大，非常不环保。所以作者在文章中发出了一个灵魂拷问：注意力机制是建模的唯一法宝吗？（Is attention all we need for modeling?）

作者在文中结合了SRU循环时序网络（recurrent network）和注意力机制，能够使用更短的训练时间达到相同的测试效果，测试效果以BPC（bits-per-character，字混淆度，越低越好）：

其中蓝线是Transformer-XL，黄线是SRU+Attention，绿线是SRU+一层Attention。提速了约5-9倍。除了训练速度更快外，作者比较了Attention的重要性，结果显示只用少数发放在合理位置的Attention就可以达到让人接受的性能（绿线）。

本文分别讲述方法和文中一些有趣的实验。码字不易，欢迎点赞、收藏与交流探讨！

方法

1.什么是SRU

SRU的计算公式

在公式1中U矩阵通过一次矩阵相乘运算得到U[t, 0]，U[t, 1]，U[t, 2]。在公式2中变量都是通过向量元素相乘得到，因此极大地减少了计算量。另一个值得注意的点，由于在公式2中进行元素相乘操作（element-wise multiplication），相关操作可以在GPU中并行，这也提升了GPU的利用效率，大幅提升了训练速度。

2.加入Attention（注意力机制）

作者将Attention放在得到U矩阵的过程中，即公式1被替换为

这里的Q、K、V的计算公式稍有不同，Q的特征维度被降维至原来的1/4或1/8，然后K和V由Q做一个等维度映射而来。（当然作者也说了可以直接由X得到降维后K、Q、V，数学上等价，效果当然也一样）。这里对QKV的降维其实也是减少计算量的一个操作。最后通过一个残差链接、layernorm和升维映射 W^o 输出U。

整个思想非常简单，示意图如下

其中a）是原始的SRU，MatMul将一个2048维度的X，映射到U（3*2048）。

b）是映射结构，先降维，再升维，其实减少了计算量。我们可以简单地比较a）和b），其中a）的矩阵运算量正比于2048*2048*3，b）的矩阵运算量正比于2048*512+512*2048*3=2048*512*4=2048*2048，约为a）的1/3。

上图c）则是加入了attention结构。

实验

作者分别在ENWIK8、WIKI-103、BILLION WORD、IWSLT’14 De→En进行了测试，使用的网络分别是Transformer-XL和10层SRU++，完全体现了更快更强的效果，如下图。

作者另一个非常有意思的实验是

实验a）我只用一个attention layer（只有一层是sru++，其他都是sru）；

实验b）一共使用2层sru++，固定最后一层使用sru++，其他每一层分别使用sru++进行测试。结果如下图所示。其中上部分对应于实验a），可以看到在第一层使用BPC最差，使用attention的层数越深，效果越好。下部分对应于实验b），在最后一层固定使用sru++的情况下，依次分别将其他层的sru替换成sru++，替换位置越深效果越好。该实验指明了如果要尽量少用attention，那么应该把attention放在网络较深的位置。

作者还比较了在注意力机制中降维的比例（相同模型大小下），如下表所示。结果显示维数从4480变至560最优（ratio：8/1）。

题外话

有一些小伙伴可能会说虽然提速了，但是最后收敛的效果可能差于Transformer模型。我觉得这个情况是有可能的。虽然个人实在是没有机会去训练那些特别大的模型，但是即使是目前在我自己的小任务上，等待模型完全收敛也是要天荒地老。（譬如设置个15个epoch的early stopping，一般情况下绝对会训到设置的最大训练迭代数。）因此在合理的迭代次数内比较不同模型的性能是有意义的，特别是GPU资源受限的条件下。