ICML 2024 | 无需LayerNorm简化Attention,精度无损推理效率大幅提升
©作者 | 陈醒濠
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
https://github.com/xinghaochen/SLAB
https://github.com/mindspore-lab/models/tree/master/research/huawei-noah/SLAB
方法
1.2 简单线性注意力
式中,DWC 表示深度可分离卷积。
实验结果
2.1 分类任务
我们在 ImageNet1k 数据集上进行了实验,实验结果证明在多个 backbone 上,我们的 PRepBN 均获得了与 LN 相当甚至更好的性能。从实验结果看,相当基于 LN 的模型,PRepBN 模型的分类精度有 0.1%~1.4% 的提升。而基于我们 SLAB 的模型,能在精度与 Flatten Transformer 相当的情况下,减少模型的推理的时延。
2.3 语言任务
我们基于 Adaptive inputs 方法在 Wikitext-103 数据集上评测了 PRepBN 在语言任务的能力。同时,我们也将 PRepBN 应用在了 LlaMA-350M 模型中,并评测了模型在下游任务的性能。从实验结果可以看出,我们的 PRepBN 方法在语言任务上也表现出了不错的性能,精度无损的情况下将 LLaMA-350M 速度从 44 tokens/s 提升到了 50.4 tokens/s。
总结
更多阅读
#投 稿 通 道#
让你的文字被更多人看到
如何才能让更多的优质内容以更短路径到达读者群体,缩短读者寻找优质内容的成本呢?答案就是:你不认识的人。
总有一些你不认识的人,知道你想知道的东西。PaperWeekly 或许可以成为一座桥梁,促使不同背景、不同方向的学者和学术灵感相互碰撞,迸发出更多的可能性。
PaperWeekly 鼓励高校实验室或个人,在我们的平台上分享各类优质内容,可以是最新论文解读,也可以是学术热点剖析、科研心得或竞赛经验讲解等。我们的目的只有一个,让知识真正流动起来。
📝 稿件基本要求:
• 文章确系个人原创作品,未曾在公开渠道发表,如为其他平台已发表或待发表的文章,请明确标注
• 稿件建议以 markdown 格式撰写,文中配图以附件形式发送,要求图片清晰,无版权问题
• PaperWeekly 尊重原作者署名权,并将为每篇被采纳的原创首发稿件,提供业内具有竞争力稿酬,具体依据文章阅读量和文章质量阶梯制结算
📬 投稿通道:
• 投稿邮箱:[email protected]
• 来稿请备注即时联系方式(微信),以便我们在稿件选用的第一时间联系作者
• 您也可以直接添加小编微信(pwbot02)快速投稿,备注:姓名-投稿
△长按添加PaperWeekly小编
🔍
现在,在「知乎」也能找到我们了
进入知乎首页搜索「PaperWeekly」
点击「关注」订阅我们的专栏吧
微信扫码关注该文公众号作者