GAU-α:尝鲜体验快好省的下一代Attention
By 苏剑林 | 2022-04-22 | 45820位读者 |在《FLASH:可能是近来最有意思的高效Transformer设计》中,我们介绍了GAU(Gated Attention Unit,门控线性单元),在这里笔者愿意称之为“目前最有潜力的下一代Attention设计”,因为它真正达到了“更快(速度)、更好(效果)、更省(显存)”的特点。
然而,有些读者在自己的测试中得到了相反的结果,比如收敛更慢、效果更差等,这与笔者的测试结果大相径庭。本文就来分享一下笔者自己的训练经验,并且放出一个尝鲜版“GAU-α”供大家测试。
GAU-α #
首先介绍一下开源出来的“GAU-α”在CLUE任务上的成绩单:
$$\small{\begin{array}{c|ccccccccccc}
\hline
& \text{iflytek} & \text{tnews} & \text{afqmc} & \text{cmnli} & \text{ocnli} & \text{wsc} & \text{csl} & \text{cmrc2018} & \text{c3} & \text{chid} & \text{cluener}\\
\hline
\text{BERT} & 60.06 & 56.80 & 72.41 & 79.56 & 73.93 & 78.62 & 83.93 & 56.17 & 60.54 & 85.69 & 79.45 \\
\text{RoBERTa} & 60.64 & \textbf{58.06} & 74.05 & 81.24 & 76.00 & \textbf{87.50} & 84.50 & 56.54 & 67.66 & 86.71 & 79.47\\
\text{RoFormer} & 60.91 & 57.54 & 73.52 & 80.92 & \textbf{76.07} & 86.84 & 84.63 & 56.26 & 67.24 & 86.57 & 79.72\\
\text{RoFormerV2}^* & 60.87 & 56.54 & 72.75 & 80.34 & 75.36 & 80.92 & 84.67 & 57.91 & 64.62 & 85.09 & \textbf{81.08}\\
\hline
\text{GAU-}\alpha & \textbf{61.41} & 57.76 & \textbf{74.17} & \textbf{81.82} & 75.86 & 79.93 & \textbf{85.67} & \textbf{58.09} & \textbf{68.24} & \textbf{87.91} & 80.01\\
\hline
\end{array}}$$
所有的模型都是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都是对齐的。
好像值得一提的也就这么多了,确实没进行多大的改变。除了在归一化方式上花了点时间进行测试,其他方面也没多费时间,直接训练就得到了不错的效果。
小结 #
GAU是笔者认为的“目前最有潜力的下一代Attention设计”,本文分享了GAU的一些训练经验,并开源了一个尝鲜版“GAU-α”。
转载到请包括本文地址:https://kexue.fm/archives/9052
更详细的转载事宜请参考:《科学空间FAQ》
如果您还有什么疑惑或建议,欢迎在下方评论区继续讨论。
如果您觉得本文还不错,欢迎分享/打赏本文。打赏并非要从中获得收益,而是希望知道科学空间获得了多少读者的真心关注。当然,如果你无视它,也不会影响你的阅读。再次表示欢迎和感谢!
如果您需要引用本文,请参考:
苏剑林. (Apr. 22, 2022). 《GAU-α:尝鲜体验快好省的下一代Attention 》[Blog post]. Retrieved from https://kexue.fm/archives/9052
@online{kexuefm-9052,
title={GAU-α:尝鲜体验快好省的下一代Attention},
author={苏剑林},
year={2022},
month={Apr},
url={\url{https://kexue.fm/archives/9052}},
}
April 22nd, 2022
赞, 后面会有显存和速度相关的评测吗?
显存和速度之前都评测过了,在序列较短的情况下,GAU没什么优势,但是序列长度较长(超过512),GAU肯定是更省显存的并且更快。
April 25th, 2022
我做了个pytorch实现,https://github.com/Tongjilibo/GAU-alpha-pytorch,再请教苏神一个问题,就是我直接用bert4keras的GAU-alpha,mlm的效果一般
重新测试了下,长文本效果还可以,短文本效果一般~
感谢你的实现和测试。
August 31st, 2022
请教苏神,GAU这种模型假如想用作生成模型的话,decoder是用类似t5的decoder还是基于GAU复现一个GAU式的decoder呢?
你喜欢,可以用GAU做decoder,Github上最新版本的bert4keras已经内置cross attention版本的GAU。
September 25th, 2022
苏神,想请教一下,在给类transformer block加position encoding(包括rope)的时候,您认为哪种方案比较合理:
a.只在第一个block前加一次
b.每个block加不同的
c.每个block加相同的
d.每几层加一次(比如降采样长度发生变化的时候肯定是需要加的)
看什么编码吧,相对位置编码一般都是每个block都加。
September 26th, 2022
[...]GAU-α:尝鲜体验快好省的下一代Attention[...]
May 30th, 2023
[...]GAU-α:尝鲜体验快好省的下一代Attention[...]
August 4th, 2023
大佬对最近出现的TransNormerLLM怎么看owo。
跟RetNet很相似,而且对于大模型只报告了速度和显存,没有报告效果,所以没什么可看的。