SLAB:采用简化线性注意力与渐进式重参数化批归一化的高效 Transformer
SLAB: Efficient Transformers with Simplified Linear Attention and Progressive Re-parameterized Batch Normalization
AI总结:
针对高效 Transformer 的归一化与注意力瓶颈,提出渐进式重参数化 BatchNorm(PRepBN)和简化线性注意力(SLA),在视觉与语言任务上实现更低延迟和相当或更优性能。
摘要:
Transformer 已成为自然语言和计算机视觉任务的基础架构。然而,其高昂的计算成本使其在资源受限设备上的部署极具挑战。本文研究高效 Transformer 的计算瓶颈模块,即归一化层和注意力模块。LayerNorm 常用于 Transformer 架构,但由于推理阶段的统计量计算,其在计算上并不友好。然而,在 Transformer 中用更高效的 BatchNorm 替换 LayerNorm 往往会导致性能下降和训练崩溃。为解决该问题,我们提出一种名为 PRepBN 的新方法,在训练过程中以重参数化 BatchNorm 渐进式替换 LayerNorm。此外,我们提出简化线性注意力(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。
英文摘要:
Transformers have become foundational architectures for both natural language and computer vision tasks. However, the high computational cost makes it quite challenging to deploy on resource-constraint devices. This paper investigates the computational bottleneck modules of efficient transformer, i.e., normalization layers and attention modules. LayerNorm is commonly used in transformer architectures but is not computational friendly due to statistic calculation during inference. However, replacing LayerNorm with more efficient BatchNorm in transformer often leads to inferior performance and collapse in training. To address this problem, we propose a novel method named PRepBN to progressively replace LayerNorm with re-parameterized BatchNorm in training. Moreover, we propose a simplified linear attention (SLA) module that is simple yet effective to achieve strong performance. Extensive experiments on image classification as well as object detection demonstrate the effectiveness of our proposed method. For example, our SLAB-Swin obtains $83.6\%$ top-1 accuracy on ImageNet-1K with $16.2$ms latency, which is $2.4$ms less than that of Flatten-Swin with $0.1\%$ higher accuracy. We also evaluated our method for language modeling task and obtain comparable performance and lower latency.Codes are publicly available at https://github.com/xinghaochen/SLAB and https://github.com/mindspore-lab/models/tree/master/research/huawei-noah/SLAB.
机构、总结、中文摘要均由AI分析而来,可能存在错误,请以论文正文为准。