【ICML 2024】重参数化BatchNorm，可直接替换现有归一化方法！

 一、论文信息 

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

中文题目：SLAB：具有简化线性注意力和渐进式重参数化批量归一化的高效变换器

论文链接：https://arxiv.org/pdf/2405.11582

官方github：https://github.com/xinghaochen/SLAB

所属机构：华为诺亚方舟实验室

核心速览：本文提出了一种名为SLAB的高效Transformer模型，通过简化线性注意力模块（SLA）和渐进式重参数化批量归一化（PRepBN）来解决Transformer在资源受限设备上部署的高计算成本问题。

二、论文概要 

图1. 在ImageNet上不同方法的比较。

1. 研究背景:

• 研究问题：Transformer架构在自然语言处理和计算机视觉任务中已成为基础架构，但其高计算成本使得在资源受限的设备上部署变得具有挑战性。研究者们需要找到一种方法，能够在保持性能的同时降低Transformer的计算成本。

• 研究难点：Transformer的归一化层（如LayerNorm）和注意力模块在推理时计算成本高，且直接替换LayerNorm为更高效的BatchNorm会导致训练不稳定和性能下降。此外，线性注意力虽然计算效率高，但要在效率和准确性之间取得良好平衡仍然具有挑战性。

• 文献综述：文章提到了Vision Transformer（ViT）的引入显著扩展了Transformer架构在计算机视觉任务中的应用。同时，为了提高Transformer的效率，研究者们尝试了多种方法，包括限制自注意力机制中token交互的范围、采用线性注意力机制以及替换LayerNorm为BatchNorm等。这些方法虽然在一定程度上提高了Transformer的效率，但仍然存在性能下降和计算成本高的问题。

2. 本文贡献:

• 重参数化的BatchNorm（RepBN）：提出了一种新的BatchNorm重新参数化公式，通过引入一个可学习的参数η，使得RepBN在训练结束后可以重新参数化为BatchNorm的形式。这种方法旨在提高训练的稳定性并增强模型性能。

 三、创新方法 

图2. 提出的渐进式重参数化批量归一化（Progressive Re-parameterized BatchNorm）的整体框架。(a) 在训练过程中，逐步用RepBN替换LayerNorm，RepBN是批量归一化的一个新的重参数化公式，用以进一步提升性能。(b) 在推理过程中，可以得到γ=0，因此，Transformer块转变为基于RepBN的架构，该架构可以进一步重参数化为批量归一化，并与线性层合并。

图2放大：RepBN是BatchNorm的一种新的重参数化公式，用以进一步提高性能。

RepBN是用于加速Transformer模型推理的归一化方法，其核心思想是将BatchNorm与线性层合并，以减少推理时的计算开销：

1. RepBN的定义：RepBN通过引入一个可学习的参数η，与BatchNorm的输出相结合，形成新的归一化公式。具体来说，RepBN的计算公式为：

其中，BN(X)代表标准的BatchNorm操作，η是一个与X同维度的可学习参数，通过端到端的方式联合训练。

2. RepBN的优势：RepBN结合了BatchNorm在训练时的稳定性和线性层在推理时的高效性，从而在不牺牲模型性能的前提下，提高了Transformer模型的推理速度。

 四、实验分析 

1. 图像分类结果：提出的渐进式替换BatchNorm（PRepBN）在不同骨干网络上均显示出与LayerNorm相当或更优的性能。例如，PRepBN在DeiT-S模型上实现了80.2%的top-1准确率，比BN+FFNBN方法高出1.4%。此外，SLAB-SwinT模型在ImageNet-1K上取得了83.6%的top-1准确率，比FlattenSwin模型低2.4ms的延迟。

2. 目标检测和实例分割结果：在COCO数据集上，使用PRepBN的模型在保持相似准确率的同时，显著降低了延迟。例如，PVT-S-PRepBN模型在对象检测任务上的延迟从64ms降低到59.6ms。

3. 语言建模结果：在Wikitext-103数据集上，PRepBN在保持与LayerNorm相似困惑度的同时，显著降低了每令牌的延迟。在LLaMA-350M模型上，PRepBN提高了吞吐量并略微提高了平均准确率。

五、代码 

https://github.com/AIFengheshu/Plug-play-modules
2025年全网最全即插即用模块，全部免费！适用于图像分类、目标检测、实例分割、语义分割、单目标跟踪（SOT）、多目标跟踪（MOT）、RGBT、图像去噪、去雨、去雾、去模糊、超分等计算机视觉（CV）和图像处理任务，持续更新中......

欢迎转发、点赞、收藏~