ICML 2024 | 无需LayerNorm简化Attention，精度无损推理效率大幅提升

2024-06-20 06:06

©作者 | 陈醒濠

Transformer 已经成为了语言和视觉任务中常用的基础架构之一。然而，由于 Transformer 结构高计算开销的影响，其在端侧等资源受限设备中的应用依然面临很大的挑战。我们对 Transformer 结构中的标准化层和注意力机制两个模块的优化策略进行了深入探索，从而构建一个高效的 Transformer 结构。其中，LayerNorm 作为 Transformer 结构中常用的标准化层，但模型推理时仍需计算数据的统计值，导致了推理的低效。

我们提出了渐进式的 LayerNorm 替换策略，并对标准的 BatchNorm 进行了改进以更好地取代 LayerNorm 层。同时，我们采用了一种简单高效的线性注意力模块（Simplified Linear Attention），来获得更强的模型性能。我们将这两种策略的结合简称为 SLAB。

我们在图像分类、目标检测以及语言任务上都进行了大量的实验，获得了很好的效果。例如，我们的 SLAB-Swin-S 在 ImageNet1k 数据集上获得了 83.6% 的分类精度，相对 Flatten-Swin-S 在精度提升 0.1% 的情况下，时延减少了 2.4ms。

论文标题：

SLAB: Efficient Transformers with Simplified Linear Attention and Progressive Re-parameterized Batch Normalization

论文链接：

https://arxiv.org/abs/2405.11582

PyTorch链接：

https://github.com/xinghaochen/SLAB

Mindspore代码：

https://github.com/mindspore-lab/models/tree/master/research/huawei-noah/SLAB

方法

1.1 渐进式重参数化BatchNorm

LN 作为 Transformer 中常用的标准化层结构，由于其在训练和推理两阶段均存在均值和方差的计算，影响了 Transformer 的执行速度。与之相对，BN 仅在训练阶段存在均值和方差的计算，且在推理阶段可与相邻的线性层融合，可以去除标准化层对模型推理速度的影响。但是，在 Transformer 结构中将 LN 简单替换为 BN 训练会导致模型精度下降以及训练崩溃等问题。为解决这个问题，我们对 BN 进行了优化，并提出了渐进式重参数化批标准化策略。

首先，重参数化批标准化的定义如下：

式中，是一个可学习参数。其中， RepBN 可以通过调节 BN 的权值和偏移量，是特定层 BN 操作被跳过；当为 0 时，RepBN 等效为纯 BN 结构。同时，RepBN 能重参数化为 BN 的表现形式，并实现与相邻线性层的融合。

其次，为增强 BN 在 Transformer 结构中的训练稳定性，我们引入了渐进式替换策略。其表示形式如下：

式中，是一个超参数，用于控制 LN 和 RepBN 的输出比例。在训练开始阶段，一般设置为 1，此时 LN 在模型中发挥主导作用；在训练结束阶段，将衰减至 0，此时模型将转变为纯 BN 组成的结构。在实际应用中，我们采用了简单的线性替换策略，的值输出如下：

其中，为训练中包含 LN 的总训练步数，为模型当前的训练步数。相对于其他衰减策略，我们发现线性策略更为简单且高效。因此，后续实验中我们均采用了线性衰减的策略。

1.2 简单线性注意力

Attention 是 Transformer 网络中重要的模块之一。为进一步压缩模型计算量，我们引入了线性注意力模块。在该模块中，我们仅使用了硬件亲和的 ReLU 算子作为相似度函数，并增加了一个深度可分离模块增强局部特征提取。该简单线性注意力模块（simplified linear attention, SLA）形式如下：

式中，DWC 表示深度可分离卷积。

实验结果
2.1 分类任务

我们在 ImageNet1k 数据集上进行了实验，实验结果证明在多个 backbone 上，我们的 PRepBN 均获得了与 LN 相当甚至更好的性能。从实验结果看，相当基于 LN 的模型，PRepBN 模型的分类精度有 0.1%~1.4% 的提升。而基于我们 SLAB 的模型，能在精度与 Flatten Transformer 相当的情况下，减少模型的推理的时延。

2.2 检测任务

此外，我们验证了不同 backbone 在 COCO 数据集上的效果。从实验结果可以看出，我们的方法实现了与原 Backbone 模型相当的性能，但拥有更低的模型推理时延。

2.3 语言任务

我们基于 Adaptive inputs 方法在 Wikitext-103 数据集上评测了 PRepBN 在语言任务的能力。同时，我们也将 PRepBN 应用在了 LlaMA-350M 模型中，并评测了模型在下游任务的性能。从实验结果可以看出，我们的 PRepBN 方法在语言任务上也表现出了不错的性能，精度无损的情况下将 LLaMA-350M 速度从 44 tokens/s 提升到了 50.4 tokens/s。

总结

我们对 Transformer 结构中的标准化层和注意力机制两个模块的优化策略进行了深入探索，提出了渐进式的 LayerNorm 替换策略，同时采用一种简单高效的线性注意力模块，来获得更加高效的 Transformer 模型架构。这个方法在图像分类、目标检测以及语言任务上进行了大量的实验验证，在精度无损的情况下，大幅提升 Transformer 的推理效率。

更多阅读