Overview
你能打开这篇文章,相信对LayerNorm(LN)、BatchNorm(BN)多少是有些了解,它们分布在神经网络中,对上一层输出的激活值做归一化(normalize),这样做的好处是可以在一定程度上避免梯度消失问题。
搞懂深度网络初始化中讲到:激活值的方差会逐层递减,这导致反向传播中的梯度也逐层递减。因此有人就会想到,那能不能在模型中加入一些归一化操作,让激活值始终维持在均值为0,方差为1呢?
实践表明,这些normalize层对提高模型精度是有帮助的,很多模型甚至用它们替代了原来的激活函数。归一化层除了LN和BN,还有IN(InstanceNorm)和GN(GroupNorm),它们的区别只在于在哪个维度上做normalize。对于不同的模型,它们的表现也各不相同。
例如BN,它沿着BatchSize这个维度做normalize,在CV模型的表现就很好,但到了NLP模型,由于BatchSize一般都较小,如果还是用BN,那效果就不好了,反之,由于channel(HiddenSize)维度很大,用LN的效果会很好,因此LN就成了现在NLP模型的标配。
LayerNorm
这里的normalize指的是正态分布的标准化,如图示,相比统计学上的计算公式,它多了3个变量,其中是很小的常量,如1e-7,为的是防止分母为0,和是可训练的参数,用来调整曲线的高矮胖瘦和顶点坐标。
Nivdia的Apex包已经为开发者提供了现阶段最快速的LN函数:fused_layer_norm(),它的性能比你自己用python实现的要快N倍。
Fused Layer Norm
LN的计算除了mean()、pow()和sqrt()这三个显性函数外,每个加减乘除背后都会调用一个相应的函数。对于异构计算来说,每次函数调用都伴随着data在host memory和device memory、device memory和registers之间的反复传输,而内存访问从来都是提高性能的瓶颈。
fused_layer_norm()之所以快,就是把原本需要调用多个函数的计算融合到一个函数中,这样不仅对内存带宽的要求要少很多,而且还能从全局来优化计算流程,如并行计算等。
Implementation
计算均值和方差是LN的主要工作量。在GPU编程中,求均值是一个reduce问题,相关的代码实例网上有很多,这里就不过多介绍。重点说下方差的算法,虽然通常它是通过均值来求的(python代码也这样做的),但这里用更快速的并行算法:Welford's online algorithm,它把计算分成warp内计算,warp间合并:
Conclusion
在模型中加入normalize层往往能收到奇效,BN/LN/GN/IN,它们的主要区别只是在哪个维度上做normalize而已。当你要用到LN时,应使用Apex的fused_layer_norm。不仅仅是LN,Apex提供了很多的融合算子,如FusedMLP、FusedAdam,在建模时你应该首选这些layer/算子。
欢迎关注和点赞,你的鼓励将是我创作的动力
欢迎转发至朋友圈,公众号转载请后台留言申请授权~