LayerNorm核心技术

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/算子。

欢迎关注和点赞，你的鼓励将是我创作的动力

欢迎转发至朋友圈，公众号转载请后台留言申请授权～