ACM MM 2022 | 统一归一化：加速Transformer工业部署的归一化方法

2022 年 8 月 22 日 PaperWeekly

©作者 | 张凯

研究方向 | 计算机视觉

自 Transformer 提出以来，其在 NLP 和 CV 领域都取得了很好的性能，以及展现了其模型大一统的潜力。面向工业上的强烈应用需求，我们组最近一直在探索和研究 Transformer 如何在工业上部署和落地，也欢迎更多的人能够加入进来一起讨论和探究。

近期我们组关于 transformer 归一化算子（LayerNorm）改进的工作《Unified Normalization for Accelerating and Stabilizing Transformers》被 ACM MM 2022 接受，和大家一起分享讨论下。

论文题目：

Unified Normalization for Accelerating and Stabilizing Transformers

论文链接：

https://arxiv.org/pdf/2208.01313.pdf

代码链接：

https://github.com/hikvision-research/Unified-Normalization

研究动机

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% 内存节省。

研究方法

首先我们对归一化的方法做了一个归类：

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。

实验结果

我们主要验证了机器翻译、图像分类、目标检测、图像分割任务。

3.1 机器翻译任务

首先可以看到，归一化是非常重要的，没有归一化的话，在 IWSLT14 上无法收敛，在 WMT14 上是掉了 7 个多点。其次是 BN 在两个任务上掉点都很大，所以在 NLP 任务中默认采用的都是 LN。同样可以看到，一些 LN 的简化算法，例如 RMSNorm 也几乎没有性能损失，但是其依然需要在线计算过程。最后是 UN 在两个数据集上性能波动在个点以内。

3.2 图像分类任务

我们验证了两类 ViT 结构，以塔式结构为代表的 SwinT 和以平坦结构代表的 T2T-ViT。结果上，在 Swin-T 上掉了 0.3 个点，在 T2T-ViT 上掉了 0.6 个点。但是在下游任务上，例如 CIFAR100 上，其性能相比 LN 高了 0.5 个点，这可能是 UN 更可以享受到 ImageNet 上带来的先验数据统计信息，这个在 DTN 中也有类似的讨论。

3.3 目标检测和图像分割任务

UN 相比于 LN 有 0.3-0.4 的掉点，而相比于其他的归一化方法例如 PN 和 BN 则有 1 个点的优势。

3.4 加速效果

我们实测了在 GPU 上的内存占用和耗时，在分类任务上，UN 相比于 LN 有 31% 的加速和 18% 的内存节省。

其他消融实验详情可以见原论文。

总结讨论

我们探索了 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

更多阅读