LLM时代的Normalization

常用Normalization这篇文章中,我们介绍了Batch NormalizationLayer NormalizationInstance NormalizationGroup Normalization等常用的Normalization手段。在大型语言模型(Large Language Model)时代,研究者们对Layer Normalization做了一些改进。

Norm方法

LayerNorm

LayerNorm使用均值和方差来重新调整数据分布,其公式如下:

RMSNorm

RMSNorm全称是Root Mean Square Layer Normalization。RMSNorm通过实验证明re-center操作并不重要,移除了LayerNorm中的均值项,可以看作LayerNorm在均值为0时的一个特例。它的计算公式如下:

由于不需要计算均值,使用RMSNorm提高了训练速度,作者认为其相比LayerNorm可以在各个模型上减少约7%-64%的计算时间。同时一些采用了RMSNorm的模型也展示了其在性能上的优越性。

DeepNorm

为了进一步稳定深度Transformer的训练,微软推出了DeepNorm,其不仅用作标准化,还作为残差连接。据论文所述,有了DeepNorm的帮助我们可以训练1000层的Transformer模型,同时保持稳定性和高性能。其计算公式如下:

Norm位置

Post Norm

原始的Transformer论文使用的就是Post Norm,其使用方式如下:

Post Norm结构的最终效果是要好于Pre Norm的,只不过Post Norm要达到自己的最优效果,需要加Warmup等训练技巧,如果和Pre Norm用一样的训练配置则效果不如Pre Norm

优点:

  • Post Norm效果上限高于Pre Norm

缺点:

  • Post Norm在深层的梯度范式逐渐增大,导致使用Post Norm的深层Transformer容易出现训练不稳定的问题

Pre Norm

同一设置下,Pre Norm结构往往更容易训练,但是最终效果通常不如Post Norm。其使用方式如下:

Pre Norm结构无形地增加了模型的宽度而降低了模型的深度,而在深度学习中深度通常比宽度更重要。也就是说Pre Norm的深度有“水分”,一个$L$层的Pre Norm模型其实际等效层数不如$L$层的Post Norm模型,层数少就会导致效果变差。

优点:

  • 相比于Post NormPre Norm在深层的梯度范式近似相等,所以使用Pre Norm的深层Transformer训练更加稳定

缺点:

  • 相比于Post NormPre Norm的模型效果略差

Sandwich Norm

Sandwich Norm是结合了Pre NormPost Norm的归一化方式,在一定程度上平衡了它们各自的优势,其使用方式如下:

该方式在Pre Norm的基础上,额外插入了一个Layer Norm

优点:

  • 有效控制每一层的激活值,避免它们过大,模型能够更好地学习数据特征

缺点:

  • 训练不稳定,可能会导致训练崩溃

经典开源模型使用的Normalization

Model Normalization
GPT3 Pre Layer Norm
Llama Pre RMS Norm
baichuan Pre RMS Norm
ChatGLM-6B Post Deep Norm
ChatGLM2-6B Post RMS Norm

【参考文献】

  1. 苏剑林-为什么Pre Norm的效果不如Post Norm
  2. Pre-trained Model Summary