Pre-RMSNorm与Pre-CRMSNorm Transformer：等价且高效的Pre-LN Transformer
机器学习
2023-10-27 v2 人工智能
神经与进化计算
摘要
Transformer在机器学习应用中取得了巨大成功。归一化技术，如层归一化（LayerNorm, LN）与均方根归一化（RMSNorm），在加速并稳定Transformer训练中起着关键作用。LayerNorm对输入向量重新中心化并重新缩放，而RMSNorm仅依据向量的RMS值进行重新缩放。尽管计算更高效，RMSNorm可能损害Transformer的表示能力。目前对于偏好何种归一化技术尚无共识，部分模型采用LayerNorm而另一些使用RMSNorm，尤其在近期大语言模型中。将一种归一化的Transformer转换为另一类型颇具挑战。在两类归一化持续争议之际，我们提出统一两种主流Transformer架构（Pre-LN与Pre-RMSNorm Transformer）的方案。通过去除Pre-LN Transformer主分支中固有的冗余均值信息，我们可将LayerNorm化简为RMSNorm，实现更高效率。我们进一步基于零均值向量的无损压缩提出压缩均方根归一化（CRMSNorm）与Pre-CRMSNorm Transformer。我们正式确立了Pre-LN、Pre-RMSNorm与Pre-CRMSNorm Transformer变体在训练与推理上的等价性。这意味着Pre-LN Transformer几乎可零代价替换为Pre-(C)RMSNorm对应结构，提供相同算术功能并免费获得效率提升。实验表明，我们可将Pre-LN Transformer的训练与推理时间减少1%–10%。
引用
@article{arxiv.2305.14858,
title = {Pre-RMSNorm and Pre-CRMSNorm Transformers: Equivalent and Efficient Pre-LN Transformers},
author = {Zixuan Jiang and Jiaqi Gu and Hanqing Zhu and David Z. Pan},
journal= {arXiv preprint arXiv:2305.14858},
year = {2023}
}
备注
NeurIPS 2023 spotlight. Code is available at https://github.com/ZixuanJiang/pre-rmsnorm-transformer