GAU-α：尝鲜体验快好省的下一代Attention

2022-04-27 06:04

©PaperWeekly 原创 · 作者 | 苏剑林

单位 | 追一科技

研究方向 | NLP、神经网络

在FLASH：可能是近来最有意思的高效Transformer设计中，我们介绍了 GAU（Gated Attention Unit，门控线性单元），在这里笔者愿意称之为“目前最有潜力的下一代 Attention 设计”，因为它真正达到了“更快（速度）、更好（效果）、更省（显存）”的特点。

然而，有些读者在自己的测试中得到了相反的结果，比如收敛更慢、效果更差等，这与笔者的测试结果大相径庭。本文就来分享一下笔者自己的训练经验，并且放出一个尝鲜版“GAU-α”供大家测试。

开源地址：

https://github.com/ZhuiyiTechnology/GAU-alpha

GAU-α

首先介绍一下开源出来的“GAU-α”在 CLUE 任务上的成绩单：

所有的模型都是 Base 版，上表显示的是 CLUE 任务上验证集上的结果，大家的运行方式和比较都是公平的，作为一个相对比较来说是合理的。另外，这里的 RoFormerV2* 并非RoFormerV2：自然语言理解的极限探索中的多任务版本，而是仅仅进行了 MLM 预训练的版本（该版本没开源），这样对比是因为 GAU-α 也仅仅进行了 MLM 预训练。

从表中可以看出，除了 WSC 这个数据量极少的“异类”外，GAU-α 在多数任务上都有优势，并且除了 WSC 外的平均成绩是最好的。其中，RoFormerV2* 与 GAU-α 的比较是最为公平的，因为它们的训练脚本、训练数据、整体结构都是一样的，唯一不同就是 GAU-α 是将 RoFormerV2* 中的 Attention+FFN 组合换成了两层 GAU，两者对比充分显示出了 GAU 设计“更好”的特点。

此外，我们在RoFormerV2：自然语言理解的极限探索介绍过 RoFormerV2 对结构进行了简化，从而获得更快的速度，具有同样整体结构的 GAU-α 也是如此，所以 GAU-α 的速度是比表中的 BERT、RoBERTa、RoFormer 都要快的，但平均效果却更胜一筹。更进一步的测试显示，当序列长度超过 512 时，GAU-α 的速度开始超过同样精简过的 RoFormerV2，并且显存占用更低，越长则对 GAU-α 更有利。

训练

现在介绍一下模型的训练细节，完整的代码已经开源到 Github 中，如有疑惑可以对照着代码来读。

模型架构：GAU-α 就是将 RoFormerV2 的 Attention+FFN 换成了两层 GAU，在之前的文章中我们比较过两层 GAU 的计算量和参数量大致相当于 Attention+FFN 组合，所以这样的替换是合理的；RoFormerV2 的特点是保留了 Post Norm 结构，去掉了所有的 Bias 项，并且 Layer Norm 换成了 RMS Norm 的最简单变体，在 GAU-α 中也是如此。

归一化：在听说Attention与Softmax更配哦～中我们讨论过 Attention 的归一化问题，GAU-α 的 Attention 归一化选取了其中笔者自行提出的具有较好外推能力的熵不变性 Softmax（在 bert4keras 中暂称为 softmax_plus）。

训练方式：在初始化方面笔者按照训练1000层的Transformer究竟有什么困难？进行了调整，因此无须 Wamrup 就可以直接训练，优化器用的是 LAMB，学习率分段线性衰减；预训练任务用的是全词 MLM，分词工具用百度的 LAC，这些跟 RoFormerV2 都是对齐的。好像值得一提的也就这么多了，确实没进行多大的改变。除了在归一化方式上花了点时间进行测试，其他方面也没多费时间，直接训练就得到了不错的效果。