来源：知乎—张凯
地址：https://zhuanlan.zhihu.com/p/549552888
01
背景
自Transformer提出以来，其在NLP和CV领域都取得了很好的性能，以及展现了其模型大一统的潜力。面向工业上的强烈应用需求，我们组最近一直在探索和研究Transformer如何在工业上部署和落地，也欢迎更多的人能够加入进来一起讨论和探究。近期我们组关于transformer归一化算子（LayerNorm）改进的工作《Unified Normalization for Accelerating and Stabilizing Transformers》被ACM MM 2022接受，和大家一起分享讨论下。
论文地址：
https://arxiv.org/pdf/2208.01313.pdf
目前代码正在整理和审核，后续会进行开源，代码地址：
https://github.com/hikvision-research/Unified-Normalization
02
研究动机
Transformer这两年在CV任务上取得了广泛的应用和验证，其在工业界有着很强的部署和落地需求。对比CNN中的采用的BN归一化方式（BN是非常高效的，在部署中可以吸收到卷积中），Transformer中采用LayerNorm的归一化方式，对部署是非常不友好的。主要原因有两点：1）LN需要在线计算过程，具体地，在线计算均值和方差；2）LN的方差计算需要开方操作，这在某些部署平台上是非常低效的，甚至是不支持的。那么直接在Transformer中将LN替换成BN怎么样呢？实验性的答案是性能会变差甚至训练会直接崩溃。其实这个问题很多同学都发现了，也都进行了讨论，transformer 为什么使用 layer normalization，而不是其他的归一化方法？。深究其中的原因，我们是将其归结为transformer在训练过程中的激活值及其梯度的统计值异常问题（后面会详细讨论）。对此的话，我们是提出了一种统一的归一化方法Unified Normalization（UN），首先，它可以像BN一样是非常高效，能够被合并到相邻的线性操作中。其次，利用数据的先验统计信息，相较于LN，它可以做到性能几乎相同或者很小的掉点。我们在机器翻译、图像分类、目标检测和图像分割等任务上都验证了其有效性，并且在GeForce RTX 3090上推理，可以带来31%的加速和18%内存节省。
03
研究方法
首先我们对归一化的方法做了一个归类：
1）online methods，这类方法的特点在于推理过程中需要在线计算的过程，其统计的维度一般是token维或者channel维度，例如LN【2】，IN【3】，DTN【4】等。这里要注意的是，在CNN中的LN，往往是对CHW维度做归一化的，而在Transformer中的LN，往往是对C维度做归一化的。
2）offline methods，对应的，这类方法的特点是无需在线计算过程，其统计维度一般是batch维度，例如BN【5】，MABN【6】（致敬下孙剑老师，在各个方向都留下了很好的工作）和我们的方法UN。
这里关于PowerNorm（PN）【7】，我们做了一个区分：1）PN with layerscale，这个是原论文最终采用的方式，前置了一个在线计算的layerscale层。2）PN without layerscale，这里在我们文中标注为PN*。
这里顺便提下曹跃老师的一篇工作《Leveraging Batch Normalization for Vision Transformers》【8】，他们也探讨了BN在transformer中的应用问题。因为其在应用中需要对结构进行调整（FFNBN），这个在swin中是较为稳定的。但在其他结构中如平坦的结构DeiT中会有掉点，在原始的transformer中，例如机器翻译任务中则是大幅掉点。本文我们主要研究在通用transformer中的归一化方法。
为了更好地获取数据集的先验统计信息，类似于PN和MABN，我们将平滑策略引入到了激活值的均值和方差统计中，对应地，在反向的梯度中也引入了平滑策略。在这种情况下，我们发现了两个现象：
1）激活值统计值的差异要比梯度统计值的差异要剧烈的多，并且随着训练的增加和深度的加深而更加突出。（可能很多做大模型的同学都知道，大模型增大后会无法收敛，其中一个原因是激活值太大，向上溢出了。这个其实在Swin v2中也提到了激活值增大的问题。）
我们定义了一个指标PNAC，来衡量这种变化差异的程度。其物理含义是，其统计值符合正态校验的占比，其值越低，说明其差异越大。详情可见原论文。可以看到，随着训练过程的增加和网络深度的加深，其激活值的差异会变得更加剧烈。
针对这一点，我们是在激活值统计值的平滑和梯度统计值的平滑采用了不同的策略。注意到由于在前向中引入了对统计值的平滑，所以反向中准确的计算梯度代价是非常大的，这里我们采用了对梯度的平滑。详情可以见论文。整体算法流程如下：
2）Transformer训练过程中的outlier会破坏训练的稳定性。而一旦前向激活值出现这种outlier，其统计值的不准确，会导致梯度的估计会出现较大偏差（前文所述梯度是没法准确计算的，是平滑估计得到的），从而影响训练的稳定性。
对应地，我们设计了一个outlier的自适应判断和处理的方法：
在outlier发生时，我们直接采用当前batch的统计值进行前向和反向的计算，保证梯度传播的正确性。
我们也理论证明了在该自适应判断和处理的策略下，其误差被大大减少。
我们可视化了LN模型，UN模型和BN模型之间的特征相似性，可以看到，LN和UN之间的特征相似性更高，这也间接说明了UN利用了先验统计信息后，其表现更接近于LN。
04
实验结果
我们主要验证了机器翻译、图像分类、目标检测、图像分割任务。
1）机器翻译任务
首先可以看到，归一化是非常重要的，没有归一化的话，在IWSLT14上无法收敛，在WMT14上是掉了7个多点。其次是BN在两个任务上掉点都很大，所以在NLP任务中默认采用的都是LN。同样可以看到，一些LN的简化算法，例如RMSNorm也几乎没有性能损失，但是其依然需要在线计算过程。最后是UN在两个数据集上性能波动在 个点以内。
2）图像分类任务
我们验证了两类ViT结构，以塔式结构为代表的SwinT和以平坦结构代表的T2T-ViT。结果上，在Swin-T上掉了0.3个点，在T2T-ViT上掉了0.6个点。但是在下游任务上，例如CIFAR100上，其性能相比LN高了0.5个点，这可能是UN更可以享受到ImageNet上带来的先验数据统计信息，这个在DTN中也有类似的讨论。
3）目标检测和图像分割任务
UN相比于LN有0.3-0.4的掉点，而相比于其他的归一化方法例如PN和BN则有1个点的优势。
4）加速效果
我们实测了在GPU上的内存占用和耗时，在分类任务上，UN相比于LN有31%的加速和18%的内存节省。
5）其他消融实验详情可以见原论文。
05
总结讨论
我们探索了transformer中归一化算子在部署中的高效优化问题。经过分析，其核心问题是transformer训练中的激活值及梯度的统计值异常问题，继而我们提出了UN，其在各种transformer模型和任务中都进行了验证，均取得了很好的结果，其在速度和内存上也有明显的优势。
最后再说一点，transformer在工业上的部署不仅仅是归一化算子的问题，需要大家一起努力探索，希望能够早日看到transformer的在工业上的推广和推动行业的进步。
参考文献
【1】Qiming Yang，Kai Zhang，Chaoxiang Lan，Zhi Yang，Zheyang Li，Wenming Tan，Jun Xiao， Shiliang Pu. Unified Normalization for Accelerating and Stabilizing Transformers. ACM MM 2022.
【2】Jimmy Lei Ba, Jamie Ryan Kiros, and Geoffrey E. Hinton. 2016. Layer Normalization. arXiv preprint arXiv:1607.06450 (2016).
【3】Dmitry Ulyanov, Andrea Vedaldi, and Victor Lempitsky. 2016. Instance normalization: The missing ingredient for fast stylization. arXiv preprint arXiv:1607.08022 (2016).
【4】Wenqi Shao, Yixiao Ge, Zhaoyang Zhang, Xuyuan Xu, XiaogangWang, Ying Shan, and Ping Luo. Dynamic Token Normalization Improves Vision Transformer. arXiv preprint arXiv:2112.02624 (2021).
【5】Sergey Ioffe and Christian Szegedy. 2015. Batch normalization: Accelerating deep network training by reducing internal covariate shift. In International conference on machine learning. PMLR, 448–456.
【6】Junjie Yan, Ruosi Wan, Xiangyu Zhang, Wei Zhang, Yichen Wei, and Jian Sun. 2020. Towards stabilizing batch statistics in backward propagation of batch normalization. arXiv preprint arXiv:2001.06838 (2020).
【7】Sheng Shen, Zhewei Yao, Amir Gholami, Michael Mahoney, and Kurt Keutzer. 2020. Powernorm: Rethinking batch normalization in transformers. In International Conference on Machine Learning. PMLR, 8741–8751.
【8】 Zhuliang Yao，Yue Cao， Yutong Lin， Ze Liu， Zheng Zhang， Han Hu. Leveraging Batch Normalization for Vision Transformers，ICCVW，2021
猜您喜欢：
附下载 |《TensorFlow 2.0 深度学习算法实战》
4405
被折叠的 条评论
为什么被折叠?
到【灌水乐园】发言
