论文地址
：
https://arxiv.org/abs/2405.11582
论文代码
：
https://github.com/xinghaochen/SLAB
Introduction
transformer
架构最初引入用于自然语言处理任务，迅速成为语言模型领域的杰出模型。随着
Vision Transformer
(
ViT
)的引入，其影响力显著扩展，展示了基于
transformer
的架构的有效性和多样性。这些架构在与卷积神经网络（
CNNs
）相比，在各种视觉任务中表现出了竞争力的性能基准。由于其强大的性能，
transformer
已成为深度学习中的主流架构。然而，
transformer
架构的计算需求构成了一个重大挑战，这主要是由于其注意力机制的二次计算复杂性和
LayerNorm
组件在线统计计算的必要性。
许多工作致力于提升
transformer
架构的效率。有的方法试图通过限制自注意机制中
token
交互的范围来减少计算复杂度，例如降采样键和值矩阵、采用稀疏全局注意模式以及在较小的窗口内计算自注意力。与此同时，线性注意力作为一种替代策略出现，通过将注意力机制分解为线性计算成本来增强计算效率，然而，在效率和准确性之间取得良好平衡仍然是一个具有挑战性的任务。此外，由于
LayerNorm
在推理过程中额外的计算开销，一些探索尝试将
BatchNorm
（
BN
）替代
transformer
中的
LayerNorm
（
LN
），比如在前向网络的两个线性层之间添加一个
BatchNorm
层来稳定训练。然而，
LayerNorm
和
BatchNorm
的
transformer
之间仍存在性能差距。
论文的重点是通过深入研究计算效率低下的模块，即归一化层和注意力模块，来获取高效的
transformer
架构。首先，论文探索了用
BatchNorm
替换
LayerNorm
以加速
transformer
的推理过程。
BatchNorm
可以降低推理延迟，但可能导致训练崩溃和性能下降，而
LayerNorm
可以稳定训练，但在推理过程中会增加额外的计算成本。因此，论文提出了一种渐进策略，通过使用超参数控制两种归一化层的比例，逐步将
LayerNorm
替换为
BatchNorm
。最初，
transformer
架构由
LayerNorm
主导，随着训练的进行逐渐过渡到纯
BatchNorm
。这种策略有效地减轻了训练崩溃的风险，并且在推理过程中不再需要计算统计信息。除了渐进策略外，论文还提出了一种新的
BatchNorm
重新参数化公式（
RepBN
），以增强训练稳定性和整体性能。
此外，注意力机制的计算成本对于高效的
transformer
架构至关重要，之前的方法在效率和准确性之间难以取得良好的平衡。因此，论文提出了一种简化的线性注意力（
SLA
）模块，该模块利用
ReLU
作为核函数，结合深度可分卷积来进行局部特征增强。这种注意力机制比之前的线性注意力更高效，而且能达到可比较的性能水平。
论文在各种架构和多个基准测试上广泛评估了提出的方法。渐进重新参数化的
BatchNorm
在图像分类和物体检测任务中表现出强大的性能，以更低的推理延迟获得类似的准确性。此外，结合渐进
RepBN
和简化线性注意力模块的
SLAB transformer
在提高计算效率的同时，与
Flatten transformer
相比达到了竞争性的准确性。例如，
SLAB-Swin-S
在
ImageNet-1K
上达到了
83.6%
的
Top-1
准确率，推理延迟为
16.2
毫秒，比
Flatten-Swin-S
的准确率高出
0.1%
，延迟则减少了
2.4
毫秒。论文还对提出的方法在语言建模任务上进行了评估，获得了可比较的性能和更低的推理延迟。
Preliminaries
给定输入为
个令牌的特征
, 其中
是特征维度, Transformer 块的一般架构可以写成:
其中，
计算注意力分数，
表示多层感知机，
是归一化函数。在 Transformer 块的默认配置中，
通常是一个 LayerNorm 操作，
是基于 sof tmax 的注意力机制
注意力在 Transformer 中扮演着重要角色。将查询、键和值矩阵表示为
, softmax 注意力首先计算查询和键之间的成对相似性。成对相似性计算导致与查询和键的数量
相关的二次计算复杂度
, 使得 Transformer 在处理具有长序列输入的任务时计算成本昂贵。线性注意力旨在解耦 softmax 函数, 通过适当的近似方法或者用其他核函数先计算
, 计算复杂度变为
, 与查询和键的数量
线性相关。
然而，
LayerNorm
在推理过程中需要统计计算，因此占据了不可忽视的延迟部分。因此，论文探索利用
BatchNorm
来构建高效的
Transformer
模型，
BatchNorm
仅在训练过程中存在，并且可以与前置或顺序线性层合并。此外，注意力模块对于
Transformer
至关重要，而基于
softmax
的注意力机制由于其二次计算复杂度而在计算效率上存在问题。因此，论文提出了一种简单而高效的注意力形式，极大地减少了延迟，同时在各种视觉任务上保持了良好的性能。
Methods
论文专注于构建高效的
Transformer
模型，并提出了一系列策略，包括逐步替换
LayerNorm
（
LN
）为重新参数化的
BatchNorm
（
BN
）以及简化的线性注意力（
SLA
）模块。所提出的
SLAB Transformer
模型在与先前方法相比表现出了强大的性能，同时具备更高的计算效率。
Progressive Re-parameterized BatchNorm
LayerNorm
在训练和推理过程中都需要进行统计量计算，因此显著影响了
Transformer
的运行速度。相比之下，
BatchNorm
在推理过程中可以简单地与线性层合并，更适合于高效的架构设计。然而，直接在
Transformer
中使用
BatchNorm
会导致性能表现不佳。为此，论文提出在训练过程中逐步替换
LayerNorm
为
BatchNorm
，并且还提出了一种受
Repvgg
启发的新的
BatchNorm
重新参数化公式，以进一步提高性能，如图
2
所示。
-
Re-parameterized BatchNorm
RepBN
公式如下：
其中，
是一个可学习的参数，以端到端的方式联合训练。一旦训练完成，
RepBN
可以重新参数化为
BatchNorm
的一种规范形式。
根据引理 4.1, RepBN 输出的分布由
和
控制, 分别对应于方差和均值。R epBN 可以借助
和
来恢复分布。
同时, 当
时, 相当于跳过了 BatchNorm 。当
时, RepBN 则退化为纯粹的 BatchNorm 。
为了促进基于纯粹
BN
的
Transformer
模型的训练，论文建议在训练过程中逐步过渡从
LN
到
RepBN
，即
其中,
是一个超参数, 用于控制不同归一化层的输出。通常, 在训练初期 LN 主导架构时,
; 在训练结束时, 为了确保过渡到基于纯粹 BN 的 Transformer,
。我们采用了一个简单而有效的衰减策略来调整
的值:
其中,
表示使用 LayerNorm 进行训练的总步数,
表示当前的训练步数。这种渐进策略有助于减轻训练纯粹基于
BN
的
Transformer
的难度，从而在各种任务上实现强大的性能表现。 还有一些其他衰减策略可以逐渐减小 的值，例如余弦衰减和阶梯衰减。从实验来看，线性策略是比较有效且简单的一种方法。
Simplified Linear Attention
注意力模块是
Transformer
网络中最重要的部分，通常表述为：
其中,
