arXiv 2305.148582023-10-26cs.LGcs.AIcs.NE

Pre-RMSNorm 与 Pre-CRMSNorm Transformer：等价且高效的 Pre-LN Transformer

Pre-RMSNorm and Pre-CRMSNorm Transformers: Equivalent and Efficient Pre-LN Transformers

Google（谷歌）；Arizona State University（亚利桑那州立大学）；The University of Texas at Austin（德克萨斯大学奥斯汀分校）

Zixuan Jiang, Jiaqi Gu, Hanqing Zhu, David Z. Pan

备注：: NeurIPS 2023 spotlight. Code is available at https://github.com/ZixuanJiang/pre-rmsnorm-transformer

AI总结：

本文通过移除Pre-LN Transformer主分支中的冗余均值信息，将LayerNorm简化为RMSNorm，并提出压缩RMSNorm及Pre-CRMSNorm架构，证明其与Pre-LN在训练和推理中等价，可在保持功能不变的同时降低1%-10%的训练与推理时间。

摘要：

Transformer 在机器学习应用中取得了巨大成功。归一化技术，如层归一化（LayerNorm，LN）和均方根归一化（RMSNorm），在加速和稳定 Transformer 训练方面发挥着关键作用。LayerNorm 对输入向量进行重新居中（recenter）和重新缩放（rescale），而 RMSNorm 仅通过向量的 RMS 值对其进行重新缩放。尽管 RMSNorm 计算效率更高，但可能会损害 Transformer 的表示能力。目前对于首选的归一化技术尚无共识，因为一些模型采用 LayerNorm，而另一些模型则使用 RMSNorm，尤其是在近期的大型语言模型中。将采用一种归一化的 Transformer 转换为另一种类型具有挑战性。在两种归一化类型之间持续存在分歧的情况下，我们提出了一种统一两种主流 Transformer 架构——Pre-LN 和 Pre-RMSNorm Transformer——的解决方案。通过移除 Pre-LN Transformer 主分支中固有的冗余均值信息，我们可以将 LayerNorm 简化为 RMSNorm，从而实现更高的效率。我们进一步基于零均值向量的无损压缩，提出了压缩 RMSNorm（CRMSNorm）和 Pre-CRMSNorm Transformer。我们正式建立了 Pre-LN、Pre-RMSNorm 和 Pre-CRMSNorm Transformer 变体在训练和推理中的等价性。这意味着 Pre-LN Transformer 几乎可以零成本地替换为 Pre-(C)RMSNorm 对应版本，在提供相同算术功能的同时，还能免费获得效率提升。实验表明，我们可以将 Pre-LN Transformer 的训练和推理时间减少 1% - 10%。

英文摘要：

Transformers have achieved great success in machine learning applications. Normalization techniques, such as Layer Normalization (LayerNorm, LN) and Root Mean Square Normalization (RMSNorm), play a critical role in accelerating and stabilizing the training of Transformers. While LayerNorm recenters and rescales input vectors, RMSNorm only rescales the vectors by their RMS value. Despite being more computationally efficient, RMSNorm may compromise the representation ability of Transformers. There is currently no consensus regarding the preferred normalization technique, as some models employ LayerNorm while others utilize RMSNorm, especially in recent large language models. It is challenging to convert Transformers with one normalization to the other type. While there is an ongoing disagreement between the two normalization types, we propose a solution to unify two mainstream Transformer architectures, Pre-LN and Pre-RMSNorm Transformers. By removing the inherent redundant mean information in the main branch of Pre-LN Transformers, we can reduce LayerNorm to RMSNorm, achieving higher efficiency. We further propose the Compressed RMSNorm (CRMSNorm) and Pre-CRMSNorm Transformer based on a lossless compression of the zero-mean vectors. We formally establish the equivalence of Pre-LN, Pre-RMSNorm, and Pre-CRMSNorm Transformer variants in both training and inference. It implies that Pre-LN Transformers can be substituted with Pre-(C)RMSNorm counterparts at almost no cost, offering the same arithmetic functionality along with free efficiency improvement. Experiments demonstrate that we can reduce the training and inference time of Pre-LN Transformers by 1% - 10%.

机构、总结、中文摘要均由AI分析而来，可能存在错误，请以论文正文为准。