本文目录

1 SLAB：使用简化线性注意力机制和渐进重参数化 BN 的视觉Transformer
(来自北京华为诺亚方舟实验室)
1 SLAB 论文解读
1.1 背景：提升 Transformer 架构的效率
1.2 渐进重参数化 BatchNorm
1.3 简化的线性注意力机制
1.4 实验结果

太长不看版

Transformer 已经成为自然语言处理和计算机视觉任务的基础架构。然而，Transformer 架构的高计算成本使得其部署在资源受限的设备上是一件具有挑战性的事。本文研究了高效 Transformer 的计算瓶颈，即 Normalization 层和 Attention 模块。LayerNorm 通常用于 Transformer 架构，但由于推理过程中的统计计算，对计算并不友好。然而，用更高效的 BatchNorm 替换 LayerNorm 层通常会导致训练性能较差甚至是崩溃。为了解决这个问题，本文提出了一种名为 PRepBN 的新方法，在训练中用重参数化的 BatchNorm 逐步替换 LayerNorm。此外，本文提出了一个简化的线性注意力 (Simplified Linear Attention, SLA) 模块，该模块简单但高效地实现了强大的性能。作者在图像分类和目标检测实验中验证了本文方法的有效性。比如，SLAB-Swin 在 ImageNet-1K 上实现了 83.6% 的 top-1 精度，延时为 16.2ms，比 Flatten-Swin 低 2.4ms，精度提高了 0.1%。作者还评估了文本任务的建模，并获得了相当的性能和更低的延时。

本文做了哪些具体的工作

1. 探索用 BatchNorm 替换 LayerNorm 来加速 Transformer 的推理过程。BatchNorm 导致较低的推理延迟，但可能导致训练崩溃和性能较差，而 LayerNorm 可以稳定训练，但在推理过程中具有额外的计算成本。为此，本文提出一种渐进策略，通过使用超参数来控制两个层的比例，将 LayerNorm 逐渐替换为 BatchNorm。作者还提出了一种新的 BatchNorm (RepBN) 重参数化方法，以提高训练稳定性和整体性能。
2. 作者提出了一个简化的线性注意 (Simplified Linear Attention, SLA) 模块，该模块利用 ReLU 作为核函数，并结合深度卷积进行局部特征增强。所提出的注意力机制比以前的线性注意力更有效，但仍然获得了相当的性能。
3. 本文的渐进式重参数化 BatchNorm 在图像分类和目标检测任务上表现出了强大的性能，以较低的推理延时获得了相当的精度。

1 SLAB：使用简化线性注意力机制和渐进重参数化 BN 的视觉Transformer

论文名称：SLAB: Efficient Transformers with Simplified Linear Attention and Progressive Re-parameterized Batch Normalization (ICML 2024)

论文地址：

http://arxiv.org/pdf/2405.11582

代码链接：

http://github.com/xinghaochen/SLAB

1.1 背景：提升 Transformer 架构的效率

Transformer 迅速成为语言模型中最突出的模型架构。随着视觉 Transformer (ViT) 的引入，它的影响得到了显着扩展，说明了基于 Transformer 架构的有效性和多功能性。与跨不同视觉任务的卷积神经网络 (CNN) 相比，Transformer 架构已经证明了它们实现具有竞争力的性能基准的能力。由于其强大的性能，Transformer 已成为深度学习的主流架构。然而，Transformer 架构的计算需求带来了重大挑战，这主要是由于其 Attention 机制的二次计算复杂度和 LayerNorm 层需要进行在线的统计计算。

为此，人们致力于提升 Transformer 这种架构的效率。一些方法试图通过限制 Self-Attention 机制中的 token 交互范围来降低计算复杂度，比如 PVT[1]中对 Key 和 Value 进行下采样。同时，线性注意力作为一种替代策略，通过将注意力机制分解为线性，来提高计算的效率[2][3]。但如何获得效率和准确性之间良好平衡的仍然是一项具有挑战性的任务。此外，有一些探索将 Transformer 中的 LayerNorm (LN) 替换为 BatchNorm (BN)，其动机是 LayerNorm (LN) 在推理过程中会产生额外的计算开销。[4]建议在前馈网络中的两个线性层之间添加一个 BatchNorm 层来稳定训练。然而，基于 LayerNorm 和基于 BatchNorm 的 Transformer 之间仍然存在性能差距。

1.2 渐进重参数化 BatchNorm

LayerNorm 需要在训练和推理中进行统计计算，从而显著阻碍了 Transformer 的运行速度。相比之下，BatchNorm 可以在推理过程中与线性层合并，更适合高效的架构。

但是，在 Transformer 中直接使用 BatchNorm 会带来不满意的性能[5]。所以，本文作者建议在训练过程中逐步将 LayerNorm 替换为 BatchNorm，并提出一种新的 BatchNorm 重参数化房还是来进一步提升性能，如下图2所示。

重参数化 BatchNorm (Re-parameterized RepBN)

本文提出的 RepBN 可以表示为：

其中， 是一个可学习的参数，以端到端的方式联合训练。一旦训练完成，RepBN 就可以重新参数化为一个新的 BN。

Lemma 4.1： 将均值 、标准差 、缩放和移位参数 和 的 层表示为 。可以将式1中的 RepBN 重参数化为：

证明

基于上述的 Lemma 4.1, RepBN 输出的分布由 和 控制, 对应于方差和均值。RepBN 可以在 和 的帮助下恢复分布。

同时, 当 时, 相当于跳过 BatchNorm。当 时, RepBN 被转换为纯 BatchNorm。

渐进式 LN → RepBN。为了便于训练纯基于 BN 的 Transformer，作者建议在训练过程中逐步将 LN 转移到RepBN，即：

其中, 是用于控制不同归一化层输出的超参数。通常, 当 主导架构时, 在训练开始时 , 训练时 , 以确保它过渡到纯基于 BN 的 Transformer。作者对 使用了一种简单而有效的衰减策略:

其中, 的训练的总步骤, 是当前步骤。这种渐进式策略简化了基于 BN 的 Transformer 的训练难度，提升了很多任务的性能。

还有一些诸如余弦衰减或者阶跃衰减的方式可以衰减 的值。根据经验，作者发现线性策略是更有效的。

1.3 简化的线性注意力机制

注意力机制是 Transformer 中最重要的部分，一般可以表述为：

其中， 表示相似度函数。对于原始形式的注意力，相似度函数为：

基于 Softmax 的注意力导致计算复杂度高。最近的几种方法[3]研究了使用线性注意力来消除 Softmax 计算，从而提高 Transformer 的性能。然而，这些方法仍然存在相当复杂的设计，而且计算效率不够。本文提出了一种简化的线性注意力 (SLA)，其公式如下：

其中, 表示深度卷积。SLA 一个简单而有效的线性注意力, 因为它还通过首先计算 , 大大降低计算复杂度。此外, 只使用了 ReLU 函数和深度卷积, 并且这两种操作在大多数硬件中都是友好的。

为了证明本文的方法仍然保持特征多样性，可视化了注意图的影响。把一个 DeiT-T 模型应用渐进重参数化BatchNorm 和简化线性注意 (SLA) 的策略，如图3所示。可以发现，本文提出的方法仍然保持了较高的秩，证明了它在捕获注意信息方面的良好能力。

1.4 实验结果

ImageNet-1K 图像分类实验结果

在 ImageNet-1K 图像分类实验中，PRepBN 的线性衰减步骤 在不同的 Backbone 上略有不同。因为 droppath 带来的方差便宜，作者冻结模型参数，并在训练结束时专门更新重参数化 BatchNorm 的统计数据10个epoch。作者遵循 Flatten Transformer [3]的宏观架构设计和训练设置。所有报告的吞吐量/延迟结果都是在单个 V100 GPU 上获得的。对于分类任务，作者测量 FLOPs 以及 224×224 图像分辨率的吞吐量/延迟。

如下图4所示是使用本文的 PRepBN 归一化层在不同 Backbone 上的结果。本文提出的 PRepBN 与 LayerNorm 相比表现出相当或更好的性能。作者进一步将本文的方法与 BN+FFNBN[5]进行了比较，BN+FFNBN 也希望用 BatchNorm 训练 Transformer 模型。可以看出，本文的 PRepBN 在不同的 Backbones 上实现了一致的改进。例如，PRepBN 在 DeiT-S 模型上实现了 80.2% 的 top-1 精度，比 BN+FFNBN 方法好 1.4%。对于 Swin Transformer，本文的 PRepBN 在 Swin-T、Swin-S 和 Swin-B 模型上比 BN+FFNBN 带来了 +0.5%、+0.4% 和 +0.5% 的精度提升。

如图5所示，作者在基于 DeiT、基于 PVT 和基于 Swin 的模型上对我们的方法进行了比较分析。很明显，配备 PRepBN 的 Transformer 模型在保持相似精度级别的同时实现了更高的吞吐量。

图6显示了本文 SLAB Transformer 的性能，它由作者提出的渐进式重参数化 BatchNorm 和简化的线性注意力组成。作者将本文的模型与 Flatten Transformer 进行比较，后者利用聚焦的线性注意力来提高效率。在包括 DeiT[6]、PVT、CSwin[7]和 Swin[8]在内的各种架构上的实验表明，SLAB Transformer 比 Flatten Transformer 获得了更好的性能。更具体地说，SLAB-SwinT 模型在 ImageNet-1K 上获得了 83.6% 的 top-1 精度，延迟为 16.2 毫秒，比 FlattenSwin 低 2.4 毫秒，精度提高了 0.1%。本文模型的计算效率更高，主要是由于硬件友好的归一化层和简化的线性注意模块。

目标检测实验结果

作者使用 Mask R-CNN 来评估本文的方法在 COCO 数据集上用于目标检测和实例分割任务的有效性。Mask R-CNN 中使用的主干在 ImageNet-1K 上进行了预训练。所有模型都训练了 1× 时间表，即 12 个 epoch。延迟在 V100 GPU 上以 1 的批大小测量，平均值为 100 轮。

作者将本文提出的 PRepBN 与标准 LayerNorm 在各种 Backbone 上进行比较，包括 Swin 和 PVT 用于目标检测任务。结果如图7所示。结果表明，本文方法实现了与配备 LayerNorm 的原始模型相当的性能。利用离线归一化，本文提出的基于 PRepBN 的模型实现了更低的推理延时。例如，当 PVT-S Backbone 配备本文的 PRepBN 时，Mask R-CNN 的延迟从 64ms 减少到 59.6ms，目标检测和实例分割的准确性相似。图8还显示了 mAP 和延迟之间的权衡的更清晰的可视化，这表明我们提出的方法在目标检测上实现了更好的整体性能。

语言建模实验结果

作者还基于自适应输入[9]评估本文方法在语言建模任务的结果。作者在 Wikitext-103[10]数据集上训练模型，该数据集包含超过 100 万个 token。我们将每个 GPU 的 token 数设置为 4096，并在 8 个 GPU 上训练。每个样本的 token 数量限制为 512。作者还将本文方法应用于 LLaMA-350M 模型，遵循与先前工作[11]类似的架构和训练设置。

如图9所示，本文的 PRepBN 与配备 LayerNorm 的模型实现了类似的困惑度，而延迟每个令牌从 13.9ms 减少到 12.9ms。

此外，作者将 PRepBN 应用于更现代的 LLMs，例如 LLaMA，它采用 LayerNorm 的变体来消除平均值的计算，即 RMSNorm。如图10所示，作者的方法在 V100 GPU 上成功地将吞吐量从每秒 44.0 提高到 50.4 个令牌，并获得了略好的平均精度。这些结果证明了 PRepBN 在语言建模任务上的有效性。