ICML 2024 | 简化Transformer，华为提出无需LayerNorm的Attention，精度无损推理效率大幅提升...

最新推荐文章于 2024-10-10 22:47:40 发布

小白学视觉

最新推荐文章于 2024-10-10 22:47:40 发布

阅读量208

点赞数

文章标签： transformer 深度学习人工智能

原文链接：https://mp.weixin.qq.com/s?__biz=MzU0NjgzMDIxMQ==&mid=2247623539&idx=2&sn=eaab8b5341e538b95bd4f0e9aa030b67&chksm=facf153d7532b209cd7d7b8ff73df99e42ffe44b8f86da01f180b0e3221759cf2a577deb9846&scene=126&sessionid=0

版权

点击上方“小白学视觉”，选择加"星标"或“置顶”

重磅干货，第一时间送达

©作者 | 陈醒濠转载自：PaperWeekly

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 的推理效率。

下载1：OpenCV-Contrib扩展模块中文版教程

在「小白学视觉」公众号后台回复：扩展模块中文教程，即可下载全网第一份OpenCV扩展模块教程中文版，涵盖扩展模块安装、SFM算法、立体视觉、目标跟踪、生物视觉、超分辨率处理等二十多章内容。


下载2：Python视觉实战项目52讲
在「小白学视觉」公众号后台回复：Python视觉实战项目，即可下载包括图像分割、口罩检测、车道线检测、车辆计数、添加眼线、车牌识别、字符识别、情绪检测、文本内容提取、面部识别等31个视觉实战项目，助力快速学校计算机视觉。


下载3：OpenCV实战项目20讲
在「小白学视觉」公众号后台回复：OpenCV实战项目20讲，即可下载含有20个基于OpenCV实现20个实战项目，实现OpenCV学习进阶。


交流群

欢迎加入公众号读者群一起和同行交流，目前有SLAM、三维视觉、传感器、自动驾驶、计算摄影、检测、分割、识别、医学影像、GAN、算法竞赛等微信群（以后会逐渐细分），请扫描下面微信号加群，备注：”昵称+学校/公司+研究方向“，例如：”张三 + 上海交大 + 视觉SLAM“。请按照格式备注，否则不予通过。添加成功后会根据研究方向邀请进入相关微信群。请勿在群内发送广告，否则会请出群，谢谢理解~