SLAB：简化线性注意力与渐进重参数化批量归一化的高效 Transformer
计算机视觉与模式识别
2024-06-18 v2 计算与语言
摘要
Transformer 已成为自然语言和计算机视觉任务的基础架构。然而，其高计算成本使其在资源受限设备上部署提出了很大挑战。本文调查了高效 Transformer 的计算瓶颈模块，即归一化层和注意力模块。LayerNorm 是 Transformer 架构中常用的模块，但由于在推理期间进行统计计算，不够计算友好。然而，用更高效的 BatchNorm 替代 LayerNorm 常常导致性能下降和训练崩溃。为解决此问题，我们提出了一种新方法，称为 PRepBN，通过逐步替换 LayerNorm 为重参数化 BatchNorm 来实现训练。此外，我们提出了一种简化线性注意力(SLA)模块，简单而有效，能够实现卓越的性能。在图像分类和目标检测方面的大量实验表明了我们方法的有效性。例如，我们的 SLAB-Swin 在 ImageNet-1K 上获得 83.6% 的 top-1 准确率，延迟为 16.2ms，比 Flatten-Swin 少 2.4ms，同时以 0.1% 更高的准确率。我们还对该方法在语言建模任务中进行了评估，获得了可比的性能和更低的延迟。代码公开于 https://github.com/xinghaochen/SLAB 和 https://github.com/mindspore-lab/models/tree/master/research/huawei-noah/SLAB。
关键词
引用
@article{arxiv.2405.11582,
title = {SLAB: Efficient Transformers with Simplified Linear Attention and Progressive Re-parameterized Batch Normalization},
author = {Jialong Guo and Xinghao Chen and Yehui Tang and Yunhe Wang},
journal= {arXiv preprint arXiv:2405.11582},
year = {2024}
}
备注
ICML 2024