论文提出了包括渐进重参数化批归一化和简化线性注意力在内的新策略，以获取高效的
Transformer
架构。在训练过程中逐步将
LayerNorm
替换为重参数化批归一化，以实现无损准确率，同时在推理阶段利用
BatchNorm
的高效优势。此外，论文设计了一种简化的线性注意力机制，其在计算成本较低的情况下达到了与其他线性注意力方法可比的性能。
来源：晓飞的算法工程笔记公众号

论文: SLAB: Efficient Transformers with Simplified Linear Attention and Progressive Re-parameterized Batch Normalization

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

给定输入为
\(N\)
个令牌的特征
\(X \in \mathbb{R}^{N \times C}\)
，其中
\(C\)
是特征维度，
Transformer
块的一般架构可以写成：

\[\begin{equation}
\begin{split}
X = X + \mathrm{Attn}(\mathrm{Norm}(X)), \\
X = X + \mathrm{MLP}(\mathrm{Norm}(X)),
\end{split}
\end{equation}
\]

其中，
\(\mathrm{Attn}(\cdot)\)
计算注意力分数，
\(\mathrm{MLP}(\cdot)\)
表示多层感知机，
\(\mathrm{Norm}(\cdot)\)
是归一化函数。在
Transformer
块的默认配置中，
\(\mathrm{Norm}(\cdot)\)
通常是一个
LayerNorm
操作，
\(\mathrm{Attn}(\cdot)\)
是基于
softmax
的注意力机制

注意力在
Transformer
中扮演着重要角色。将查询、键和值矩阵表示为
\(Q, K, V \in \mathbb{R}^{N \times C}\)
，
softmax
注意力首先计算查询和键之间的成对相似性。成对相似性计算导致与查询和键的数量
\(N\)
相关的二次计算复杂度
\(O(N^2C)\)
，使得
Transformer
在处理具有长序列输入的任务时计算成本昂贵。线性注意力旨在解耦
softmax
函数，通过适当的近似方法或者用其他核函数先计算
\(K^T V\)
，计算复杂度变为
\(O(NC^2)\)
，与查询和键的数量
\(N\)
线性相关。

然而，
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
公式如下：

\[\begin{equation}
\mathrm{RepBN}(X) = \mathrm{BN}(X) + \eta X,
\end{equation}
\]

其中，
\(\eta\)
是一个可学习的参数，以端到端的方式联合训练。一旦训练完成，
RepBN
可以重新参数化为
BatchNorm
的一种规范形式。

根据引理
4.1
，
RepBN
输出的分布由
\(\alpha+\eta\sigma\)
和
\(\beta+\eta\mu\)
控制，分别对应于方差和均值。
RepBN
可以借助
\(\sigma\)
和
\(\mu\)
来恢复分布。

同时，当
\(\alpha=0, \beta=0\)
时，相当于跳过了
BatchNorm
。当
\(\eta=0\)
时，
RepBN
则退化为纯粹的
BatchNorm
。

Progressive LN
\(\rightarrow\)
RepBN

为了促进基于纯粹
BN
的
Transformer
模型的训练，论文建议在训练过程中逐步过渡从
LN
到
RepBN
，即

\[\begin{equation}
\mathrm{PRepBN}(X) = \gamma\mathrm{LN}(X) + (1 - \gamma)\mathrm{RepBN}(X),
\end{equation}
\]

其中，
\(\gamma\)
是一个超参数，用于控制不同归一化层的输出。通常，在训练初期
LN
主导架构时，
\(\gamma=1\)
；在训练结束时，为了确保过渡到基于纯粹
BN
的
Transformer
，
\(\gamma=0\)
。我们采用了一个简单而有效的衰减策略来调整
\(\gamma\)
的值：

\[\begin{equation}
\gamma = \dfrac{T - T_{cur}}{T}, \gamma \in [0, 1],
\end{equation}
\]

其中，
\(T\)
表示使用
LayerNorm
进行训练的总步数，
\(T_{cur}\)
表示当前的训练步数。这种渐进策略有助于减轻训练纯粹基于
BN
的
Transformer
的难度，从而在各种任务上实现强大的性能表现。

还有一些其他衰减策略可以逐渐减小
\(\gamma\)
的值，例如余弦衰减和阶梯衰减。从实验来看，线性策略是比较有效且简单的一种方法。

Simplified Linear Attention

注意力模块是
Transformer
网络中最重要的部分，通常表述为：

\[\begin{equation}
\begin{split}
&Q=XW_{Q}, K=XW_{K}, V=XW_{V},\\
&O_{i} = \sum_{j=1}^{N}\dfrac{\mathrm{Sim}(Q_{i}, K_{j})}{\sum_{j}\mathrm{Sim}(Q_{i}, K_{j})}V_{j},
\end{split}
\end{equation}
\]

其中，
\(W_Q, W_K, W_V \in \mathbb{R}^{C \times C}\)
将输入的标记投影到查询（
query
）、键（
key
）和值（
value
）张量。
\(\mathrm{Sim}(\cdot, \cdot)\)
表示相似性函数。对于注意力的原始形式，相似性函数是

\[\begin{equation}
\mathrm{Sim_{softmax}}(Q_i , K_j) = \exp (\frac{Q_iK_j^{T}}{\sqrt{C}}),
\end{equation}
\]

这种基于
softmax
的注意力导致了较高的计算复杂度。近年来，有几种方法研究了使用线性注意力来避免
softmax
计算，从而提高
Transformer
的效率。然而，这些方法仍然存在相当复杂的设计，并且计算效率不够高。因此，论文提出了一种简化的线性注意力（
SLA
）：

\[\begin{equation}
\begin{split}
&{\rm Sim}_{SLA}\left(Q_{i},K_{j}\right)=\mathrm{ReLU}\left(Q_{i}\right){\mathrm{ReLU}\left(K_{j}\right)}^T,\\
&\tilde {\rm O}_{i} = \sum_{j=1}^{N}\dfrac{\mathrm{Sim}_{SLA}(Q_{i}, K_{j})}{\sum_{j}\mathrm{Sim}_{SLA}(Q_{i}, K_{j})}V_{j},\\
&\!{\rm O}_{SLA}\!=\tilde {\rm O}+\!{\rm DWC}(V),
\end{split}
\end{equation}
\]

其中，
\(DWC(\cdot)\)
表示深度可分离卷积（
depth-wise convolution
）。这是一种简单而高效的线性注意力方法，因为它通过先计算
\(K^T V\)
，享受了解耦的计算顺序，从而显著减少了复杂度。此外，该方法只使用了
ReLU
函数和深度可分离卷积，这两种操作在大多数硬件上都具有良好的计算效率。

这里的整体逻辑跟
FLatten Transformer
基本一样，只是将其提出聚焦函数替换为
ReLU
函数。这里的效率提升通过摘除
softmax
计算从而达到先计算
\(K^T V\)
实现的（公式7做下乘法结合律），
ReLU
（也有保证内积为正数的作用）和
DWC
是补充计算顺序改变带来的性能损失。