LayerNorm与RMSNorm:大语言模型中的归一化技术深度解析

深入对比LayerNorm与RMSNorm在大语言模型中的应用,解析其技术原理和优劣势,揭示RMSNorm在梯度稳定性和计算效率上的优势。

原文标题:大语言模型中的归一化技术:LayerNorm与RMSNorm的深入研究

原文作者:数据派THU

冷月清谈:

本文深入探讨了大语言模型中归一化技术的必要性,对比了LayerNorm和RMSNorm的技术原理和应用。归一化通过调整数据量纲,提高计算稳定性并优化学习效率。LayerNorm通过计算每个样本的均值和标准差进行归一化,尤其适用于序列数据和Transformer模型。RMSNorm是LayerNorm的变体,仅使用均方根进行归一化,省略了均值计算,降低了计算成本,同时保持模型性能和提供更稳定的学习过程。文章还从梯度传播的视角分析了两种归一化技术的差异,RMSNorm通过保持非零均值,有效缓解了梯度消失问题。最后,文章总结了两种归一化技术在现代深度学习模型中的重要作用,以及RMSNorm对LayerNorm的有效补充。

怜星夜思:

1、LayerNorm和RMSNorm在实际应用中,除了Transformer,还在哪些类型的模型或任务中表现出色?原因是什么?
2、文章提到RMSNorm通过保持非零均值来缓解梯度消失,那么是否意味着均值为零就一定不好?在什么情况下,零均值反而有优势?
3、文章中说RMSNorm计算效率更高,但仅仅是少了一次均值计算。在实际训练中,这种计算效率的提升能带来多大的性能改善?在哪些场景下,这种提升会更明显?

原文内容


来源:Deephub Imba
本文约3000字,建议阅读6分钟
本文将系统分析归一化技术的必要性,并详细阐述为何原始Transformer架构中的LayerNorm在LLama模型中被RMSNorm所替代的技术原理。


在LLama等大规模Transformer架构的语言模型中,归一化模块是构建网络稳定性的关键组件。本文将系统分析归一化技术的必要性,并详细阐述为何原始Transformer架构中的LayerNorm在LLama模型中被RMSNorm所替代的技术原理。



归一化技术的基础原理


归一化的核心定义


归一化(Normalization)是一种将数据映射到特定数值区间的数学变换技术,旨在提升计算稳定性并优化学习效率。归一化的本质在于调整数据的量纲规模而保持分布形态不变。

归一化前后数据分布

归一化前后按特征划分的数据分布

通过对比归一化前后的数据分布可以明确观察到,尽管数据点的相对位置关系保持恒定,但数值范围发生了显著变化。如上图所示,横轴上原本分布在30至70区间的数据经归一化后映射至0到1区间,然而数据的拓扑分布结构保持不变。这构成了归一化技术的基本理论框架。

归一化技术的技术必要性


归一化技术在深度学习中具有双重技术价值。首要方面是提高数值计算稳定性和优化收敛速度。当数据量级差异显著时,评估指标(如均方误差MSE)可能产生数量级上的波动。在MSE计算过程中,大值输入会导致梯度幅值过大,从而使优化过程不稳定。

以均方误差计算为例,考虑两组输入值a和b的对比情况:

  1. 未经归一化处理:
    • a = [100, 200, 300], b = [110, 190, 310]
    • MSE = (1/3) * [(100–110)² + (200–190)² + (300–310)²]
    • MSE = (1/3) * [100 + 100 + 100] = 100
  2. 经归一化处理:
    • 归一化后的值:a_normalized = [1.0, 2.0, 3.0], b_normalized = [1.1, 1.9, 3.1]
    • MSE = (1/3) * [(1.0–1.1)² + (2.0–1.9)² + (3.0–3.1)²]
    • MSE = (1/3) * [0.01 + 0.01 + 0.01] = 0.01

从计算结果可见,未归一化的数据由于数值较大,产生了较大的MSE值,这可能导致梯度幅值过大,使模型训练过程不稳定。而归一化处理后,MSE值显著降低,梯度幅值维持在合理区间,有助于参数更新过程的稳定性。

通过将数据规范化到一致的数值范围,归一化能够有效控制梯度幅值,实现更为平稳的参数优化过程。

归一化的另一技术价值在于潜在的内存效率提升。归一化通过缩减数据表示范围间接提高了内存利用效率。特别是在采用8位或16位浮点格式进行训练时,数据范围的收窄使得低精度表示更为精确。这种特性可通过量化技术或FP16计算进一步降低模型训练阶段的内存占用。

需要说明的是,归一化本身并不直接减少内存占用,它需要与量化或低精度计算技术协同应用才能实现有效的内存优化。

归一化的实证效果分析


通过实际数据集分析可验证归一化的技术效果。以加州住房数据集为例,采用归一化技术的模型(蓝线)展现出快速且稳定的学习曲线,而未归一化的模型(红线)则表现出学习不稳定性。这种差异在训练过程中尤为明显,未归一化情况下可能出现数值溢出,导致损失函数出现异常峰值。

归一化前后MSE比较——加州住房数据集

归一化前后准确率性能比较——CIFAR-10

在CIFAR-10图像分类任务中,类似的数值稳定性问题同样存在。应用归一化技术后,学习过程能够迅速达到稳定状态。因此,归一化已成为现代深度学习的基础技术组件,对于提升模型性能至关重要。

层归一化(Layer Normalization)



层归一化最初由Jimmy Lei Ba、Jamie Ryan Kiros和Geoffrey Hinton在2016年发表的论文"Layer Normalization"中提出。该研究提出LayerNorm作为批量归一化的替代方案,以解决其局限性,特别强调了其在序列数据和循环神经网络模型中的适用性。

层归一化通过计算每个样本的均值(μ)和标准差(σ)来实现数据归一化。由于LayerNorm是在每个样本基础上独立应用的,不依赖于批量大小,因此在处理序列数据的模型中表现尤为出色。Transformer模型作为序列数据处理的核心架构,广泛采用了LayerNorm技术。

在Transformer架构中,LayerNorm扮演着至关重要的角色。Transformer处理的序列数据由多个token组成,LayerNorm所归一化的"样本"实际上是每个token位置的嵌入向量或隐藏状态。

例如,对于句子"The cat sat on the mat",Transformer模型会为每个单词(token)生成特定的向量表示(嵌入)。LayerNorm在此场景下归一化的是该token位置的嵌入向量或隐藏状态。LayerNorm的核心机制是对每个token位置计算的嵌入或隐藏状态向量进行归一化,以维持特定范围内的数值稳定性。


其中:

  • x 表示每个样本的输入值
  • μ 表示平均值
  • σ 表示标准差
  • ϵ 是一个微小常数,用于确保数值稳定性,防止除零错误,通常设定为0.00001或0.000001

LayerNorm将数据归一化为均值为0,方差为1的分布。


Transformer输入示例和LayerNorm计算过程


LayerNorm应用结果

从上图可观察到,每个token的嵌入维度为5,均值和方差计算在嵌入向量的5维值空间内进行。值得注意的是,这一计算过程独立于其他数据点或token。另一关键特性是数据分布向均值0对齐,这一特性与RMSNorm形成明显对比。

RMS归一化(RMS Normalization)


RMS归一化(RMSNorm)是层归一化的变体,其特点是仅使用均方根(Root Mean Square)进行归一化,省略了均值计算环节。RMSNorm在降低计算成本的同时能够维持模型性能,同时提供更加稳定的学习过程。


RMS(x)的定义如下:


其中RMS代表特征值的均方根。


Transformer输入示例和RMSNorm计算过程


从结果可以观察到,RMSNorm不执行均值对齐操作。由于它没有执行 x-μ 运算,因此不存在向0对齐的过程。与传统LayerNorm需要两次数据扫描(一次计算均值,一次计算方差)不同,RMSNorm仅需一次处理即可完成RMS计算。这种简化方法在保持性能的同时有效降低了计算资源需求。

LayerNorm与RMSNorm的技术对比


均值对齐特性分析


与LayerNorm显著不同的是,RMSNorm保持非零均值分布,不执行均值对零的对齐操作。


通过对比LayerNorm和RMSNorm应用到前述Transformer输入示例后的均值分布结果,差异变得更为明显。尽管两种方法都将数据分布收敛到特定区域,LayerNorm呈现零均值特性,而RMSNorm则表现为非零均值分布。

从梯度传播视角的比较分析


RMSNorm通过保持非零均值而不强制均值对零的对齐,有效缓解了梯度消失问题。LayerNorm的零均值对齐可能导致梯度在反向传播过程中逐层衰减。

LayerNorm与RMSNorm:跨Epoch的梯度稳定性对比

上图基于简化网络模型分析了LayerNorm和RMSNorm的梯度传播效果。左侧图表展示了按训练epoch划分的LayerNorm和RMSNorm的平均梯度范数变化。梯度范数是衡量网络权重变化速率的指标,具体表示反向传播过程中各层权重梯度的幅值。较大的梯度范数值表示权重变化幅度较大,而过大或过小的值可能导致学习不稳定。

RMSNorm相比LayerNorm维持了更高的梯度范数,尤其在训练初期阶段。这一特性对于防止深度学习模型中的梯度消失问题至关重要。右侧图表显示了RMSNorm与LayerNorm之间的梯度比率,初期RMSNorm的梯度显著高于LayerNorm,随着学习进程的推进,这一差异逐渐减小。

与LayerNorm相比,RMSNorm在计算效率和训练稳定性方面表现卓越。网络结构越深,RMSNorm通过有效防止梯度消失问题实现更稳定的学习过程,同时能够在减少计算资源消耗的条件下达成相似的性能水平。

结论


本文深入分析了原始Transformer架构中采用的层归一化技术与LLama架构中实现的RMS归一化技术之间的技术差异与特性。LayerNorm通过调整数据的均值和方差确保数值稳定性,在各类序列模型中发挥关键作用。而RMSNorm作为一种无需利用均值即可稳定数据尺度的技术,能够有效解决深层网络中的梯度消失问题,同时保持较快的训练收敛速度。

这两种归一化技术在现代深度学习模型中均发挥着不可替代的作用,对于最大化提升模型性能(尤其是在Transformer架构中)至关重要。RMSNorm在计算效率和训练稳定性方面对LayerNorm形成了有效补充,为提升复杂深层模型性能做出了重要贡献。

编辑:王菁




关于我们

数据派THU作为数据科学类公众号,背靠清华大学大数据研究中心,分享前沿数据科学与大数据技术创新研究动态、持续传播数据科学知识,努力建设数据人才聚集平台、打造中国大数据最强集团军。




新浪微博:@数据派THU

微信视频号:数据派THU

今日头条:数据派THU

我偏学术一点,补充一下。LayerNorm由于其对batch size不敏感的特性,在一些小batch size或者单样本的训练场景下应用较多,比如few-shot learning。而RMSNorm在图神经网络(GNN)中也有应用,因为它计算简单,可以加速图数据的处理。具体原因可能和GNN对计算效率的要求更高有关。

我来个不严谨的类比!RMSNorm就像一个优化过的代码,虽然只是少了几行代码,但在高并发的情况下,就能节省大量的服务器资源。当然,实际提升多少,还得看你的代码质量和服务器配置。

这个问题很有意思!除了Transformer,LayerNorm在RNN(循环神经网络)及其变体中也很常见,因为它能有效处理序列长度不一致的问题。RMSNorm我了解的相对少一些,不过考虑到它计算效率高,可能在资源受限的边缘设备上部署的模型中更有优势,例如一些轻量级的语音识别或图像分类模型。大家觉得呢?

我补充一点工程上的考量。在资源受限的设备上,哪怕是很小的计算量减少,都能带来显著的性能提升。比如在移动设备上部署模型时,RMSNorm的优势就可能更加突出。此外,如果训练过程中需要频繁进行归一化操作,例如在GAN(生成对抗网络)中,那么RMSNorm的效率优势也会更加明显。

从信号处理的角度来看,零均值信号可以消除直流分量,突出交流信号,这在某些需要关注信号变化的场景下是有利的。例如,在图像处理中,零均值化可以去除图像的亮度偏置,使模型更关注纹理细节。所以,是否需要零均值,取决于具体的任务和数据特性。

抖个机灵!LayerNorm就像一个班里的尖子生,啥都会,适应性强。RMSNorm就像一个偏科生,在特定领域特别厉害,比如计算资源紧张的时候。选哪个,就看你的任务是啥啦!

这个问题很有深度啊!零均值并非绝对不好。在某些情况下,它可以简化计算,例如在BatchNorm中,零均值可以使激活函数的输入分布更稳定。而且,零均值可能也有助于一些特定类型的特征学习。不过在深层网络中,如果每层都强制零均值,可能会导致信息丢失,这就是RMSNorm要保持非零均值的原因吧?

这确实是个好问题!虽然RMSNorm只少了一次均值计算,但在大规模训练中,每一次计算的节省都会累积成显著的性能提升。尤其是在模型非常深、数据量非常大的情况下,这种提升会更加明显。此外,在GPU等硬件加速器上,减少计算步骤也能更好地利用并行性,从而进一步提高效率。

我来皮一下!零均值就像做人要中庸,但有时候太中庸就没个性了!在模型里也是一样,太追求零均值,可能就把一些有用的信息给抹平了。当然,如果你的数据本身就有很多噪音,那零均值可能就像一个降噪器,帮你把噪音滤掉。