浅谈BERT/Transformer模型的压缩与优化加速

2021 年 12 月 31 日 PaperWeekly

©作者 | 姚益武

单位 | 阿里巴巴集团

研究方向 | AI算法与工程架构

前言

BERT/Transformer 结构及其变体，已成为自然语言处理（NLP）、语音识别（ASR）等领域的主流序列建模结构。并且，相比于卷积操作的局部特征提取能力、以及平移不变性，Self-Attention/Cross-Attention 的全局 Context 信息编码或长时依赖建模能力，能够使图像/视觉模型具备更强的相关性建模能力（更好的全局视野）、与动态建模能力（领域或样本适应性）。

因此在计算机视觉领域（CV），Vision Transformer 模型结构（如 ViT、Deit、Swin-Transformer、PVT、SegFormer 与 DETR 等）也日益流行、方兴未艾。然而，Transformer 模型的高计算复杂度与参数量，限制了模型的推理部署（包括服务端与移动端应用），其计算、存储与运行时内存等资源开销都成为限制因素（存在大量密集计算算子、与访存算子，如 BatchMatMul、LayerNorm、Transpose/Reshape 等）。例如，从标准 Transformer layer 的 Tensor 处理结构来看，MHA 的 Context 特征计算与特征长度的平方、Embedding 特征维度成正比：

Standard Self-Attention (X=Y) / Cross-Attention (X!=Y):

Standard FFN:

本文从以下几个维度举例说明，解析 Bert/Transformer 模型的加速方法，体现了从建模依赖、数据依赖到硬件依赖的优化层次：

模型结构精简化与知识蒸馏
模型量化（Quantization）
网络结构搜索（NAS: Network Architecture Search）
计算图优化（计算图等价变换）
推理优化引擎（Faster Transformer / Torch-TensorRT / AICompiler）

模型结构精简化与知识蒸馏

2.1 Lite-Transformer

论文标题：

Lite Transformer with Long-Short Range Attention

收录会议：

ICLR 2020

论文链接：

https://arxiv.org/abs/2004.11886

代码链接：

https://github.com/mit-han-lab/lite-transformer

Lite Transformer 是一种融合卷积与 Self-Attention 操作的、高效精简的 Transformer 结构，可应用于 NMT、ASR 等序列生成任务。其核心是长短距离注意力结构（LSRA：Long-Short Range Attention），能够有效融合卷积与 Attention 机制。

LSRA 将输入 Embedding 沿 Feature 维度 split 成两部分，其中一部分通过 GLU（Gate Linear Unit）、一维卷积，用以提取局部 Context 信息；而另一部分依靠 Self-attention，完成全局相关性信息编码。一维卷积的引入，能够减轻计算量与参数量。Lite Transformer 核心结构如下，首先将 FFN 的宽度摊平（flatten），其次引入 LSRA 以替换 Self-Attention：

2.2 SAN-M

论文标题：

SAN-M: Memory Equipped Self-Attention for End-to-End Speech Recognition

收录会议：

INTERSPEECH 2020

论文链接：

https://arxiv.org/abs/2006.01713

SAN-M 表示 Self-Attention 与 DFSMN 记忆单元的融合 ，是一种 Transformer ASR 模型。DFSMN 适合捕获局部信息，Self-Attention 模块具备较强的长时依赖建模能力，因此二者存在互补性。SAN-M 通过将两个模块的特性融合在一起，实现了优势互补。Biasic Sub-layer 表示包含了 SAN-M 的 Self-Attention 模块，DFSMN 添加在 values 后面，其输出与Multi-head Attention（MHA）相加：

2.3 MiniLM（知识蒸馏）

论文标题：

MiniLM: Deep Self-Attention Distillation for Task-Agnostic Compression of Pre-Trained Transformers

论文链接：

https://arxiv.org/abs/2002.10957

代码链接：

https://github.com/microsoft/unilm/tree/master/minilm

针对 NLP 任务，深度自注意力知识蒸馏（Deep Self-Attention Distillation），通过迁移 Teacher model 最后一层 Self-Attention 模块的 Attention score 信息与 Value relation 信息，可有效实现 Student model 的诱导训练。只迁移最后一层的知识，能够直接迁移语义信息，显得简单有效、训练速度更快；而且相比于层间特征迁移，不需要手动设计 Teacher-student 之间的层对应关系。Attention score 信息与 Value relation 信息的知识迁移如下：

Attention score transfer:

Value relation transfer:

若选择的 Self-attention layer 为 LSRA， 除了在 Multi-head Attention（MHA）分支迁移 Attention score 与 Value relation；在 CNN 分支需要迁移 Feature map 的信息，这里主要计算 AT loss：

式中表示学生网络里第 j 个网络层的特征输出，表示 Teacher network 里第 j 个 group 的特征输出。

2.4 Switch-Transformer

论文标题：

Switch Transformers: Scaling to Trillion Parameter Models with Simple and Efficient Sparsity

论文链接：

https://arxiv.org/abs/2101.03961

代码链接：

https://github.com/tensorflow/mesh/blob/master/mesh_tensorflow/transformer/moe.py

针对内容理解与生成、以及多模态特征表征等 AI 任务，基于 MoE（Mixture of Experts）单元的大模型的参数规模不断扩展（Switch-Transformer 是其中的典型代表之一），但大模型对算力的需求、被 MoE 的稀疏激活（Sparse activation）或动态路由（Dynamic routing）机制有效缓解。

MoE 单元属于典型的 Sample 驱动型结构设计，使得 MoE-based 模型具备良好的动态建模能力（动态卷积、Self-Attention 等亦如是），因此 MoE 能够非常好的适应多任务/多模态、甚至在线实时的场景应用需求。 Switch-Transformer 的 MoE 单元结构如下，具体执行过程包括动态 路由（Dynamic routing）、数据分派（Data dispatch）、专家计算（Expert computation）与结果合并（Result combine）：

模型量化（Quantization）

3.1 Transformer ASR 模型压缩

论文标题：

Extremely Low Footprint End-to-End ASR System for Smart Device

收录会议：

INTERSPEECH 2021

论文链接：

https://arxiv.org/abs/2104.05784

针对 Transformer ASR，文章提出了联合随机稀疏与 PTQ 量化（KL 量化、ADMM 与混合精度设置）的压缩策略，整体实现了 10 倍压缩，且绝对精度损失约 0.5%。总体流程如下：

模型稀疏化，更新 Weight 重要性，渐进式增加稀疏度：

KL 量化，计算 Activation 量化参数：

ADMM，优化 Weight 量化参数：

混合精度量化设置，减少量化误差：

此外，针对 Transformer 模型的量化，需要讨论全网络 INT8 量化与计算的意义：

全网络 INT8 量化：

同时减少计算密集算子与访存密集算子的开销；
实现模型压缩，INT8 模型是 FP32 模型的 1/4；

Fully INT8 Attention Path：

基本的全 INT8 通路：

Dense layer 的输入与输出均为 INT8 Tensor；
Reshape 与 Transpose 的操作对象为 INT8 Tensor，节省内存开销：
Dense+Reshape+Transpose、与 Transpose+Reshape+Dense 可以实现 Op fusion；

BatchMatMul、Softmax 的操作对象保留为 FP32 Tensor，确保模型预测精度；

Encoder的Self-Attention：

X=Y；

Decoder 的 Self-Attention：

X=Y，Batch_size=Batch_size * Beam_size，Seq_len=1；
k_dense 与 v_dense 的输出会添加到 Cache，按照 INT8 类型搬运数据，可节省访存开销：

3.2 PTQ for Vision Transformer

论文标题：

Post-Training Quantization for Vision Transformer

论文链接：

https://arxiv.org/abs/2106.14156

代码链接：

https://github.com/hahnyuan/PTQ4ViT

针对 Vision Transformer 的 PTQ 量化， 分别针对 FFN 与 Self-Attention，提出了 Similarity-aware Quantization 与 Ranking-aware Quantization：

Similarity-aware Quantization：基于优化方式确定 Weight 与 Activation 量化的最优量化参数，并以 Pearson 相似度作为 Target loss；

Ranking-aware Quantization：在相似度 Loss 基础之上，叠加 Attention map 的 Ranking loss 作为约束，以准确感知不同 Attention map 的重要性排序，确保 Self-attention 量化的有效性（优先保留重要 Attention map 的信息）；

3.3 PTQ4ViT

参考博客文章：

https://blog.csdn.net/nature553863/article/details/121564248

3.4 Q-Bert

论文标题：

Q-BERT: Hessian Based Ultra Low Precision Quantization of BERT

收录会议：

AAAI 2020

论文链接：

https://arxiv.org/abs/1909.05840

针对 NLP 任务（MNLI、SQuAD 等），Q-Bert 是首个实现 Bert 模型混合精度量化的 QAT 算法。Bert 模型是由 Embedding layer、Encoder layer 以及 Task-specific output layer 构成的前馈结构模型，且 Encoder layer 包含了 MHA 与 FFN 等子模块。

不同的 Bert layer 体现了不同的量化敏感度（Quantization Sensitivity），配以不同的量化比特数，直接决定了量化后模型的预测精度。通常 Bert 的 Embedding layer 与 Output layer 的量化敏感度相对较高，故 Embedding layer 会配以较高的量化比特数（如 8bits），Output layer 则保留为 FP32 浮点实现。另外，尽管不同 Encoder layer 的结构相一致，但由于对 Context 信息编码的贡献有所区别，因此表现出不同的量化敏感性。

量化敏感度与混合精度量化： Q-Bert 采用线性均匀量化方式（min-max 量化），并通过 Dense layer 的 Hessian 信息（二阶梯度）定义其量化敏感度：

式中 Hessian 矩阵特征值（Top eigenvalues）的均值，反映了针对同一个下游任务（相同训练数据分布），不同 Dense layer 的量化敏感度；Hessian 矩阵特征值的标准差，反映了针对不同的下游任务（不同训练数据分布），Dense layer 量化敏感度的占比。然后，基于不同 Dense layer 的量化敏感度排序，可确定 Bert 模型的混合精度量化配置，并进一步执行 QAT 微调训练。

Group-wise量化： 为了进一步提高量化后模型的精度，将 MHA 的 Q、K、V 与 O 权重矩阵按 Attention heads 分成不同的 Group，每个 Group 定义各自的量化范围（min-max 范围），从而更精细的保证不同 Group 的量化分辨率。而每个 Group 又可以根据实际需求、进一步划分为不同的 Sub-group，以满足更为苛刻的部署精度需求：

网络结构搜索（NAS）

4.1 AdaBert

论文标题：

AdaBERT: Task-Adaptive BERT Compression with Differentiable Neural Architecture Search

收录会议：

IJCAI 2020

论文链接：

https://arxiv.org/abs/2001.04246

针对 NLP 任务，AdaBert 通过可微分搜索、与知识蒸馏，实现卷积类型的 NLP 模型搜索：

搜索空间： 类似于 DARTS 的 Cell 设计，实现搜索空间的构造；区别于 Vanilla CNN，DARTS 的 Cell 单元结构体现出了非规整性，可能不利于实际推理引擎的加速；

搜索策略：基于 Gumbel Softmax 实现网络结构的随机采样（类似于 FBNet）、与可微分搜索，并且采样概率随着训练的进行、逐渐锐化，逼近 Argmax 采样；另外，基于 FLOPS 与 Model size 构造了 Efficiency-aware Loss，作为搜索训练的资源约束，实现 Hardware-aware 搜索；

知识蒸馏：多层次、任务相关的知识迁移；

4.2 DynaBert

论文标题：

DynaBERT: Dynamic BERT with Adaptive Width and Depth

收录会议：

NeurIPS 2020 (Spotlight)

论文链接：

https://arxiv.org/abs/2004.04037

代码链接：

https://github.com/huawei-noah/Pretrained-Language-Model/tree/master/DynaBERT

针对 NLP 任务，DynaBert 模型压缩策略可实现多个维度的压缩搜索（主要是 Width 与 Depth 维度）。DynaBert 是一种动态可伸缩性质的 One-shot NAS：首先基于预训练 Bert 模型（或其变体），进行多维度正则化训练；正则化训练之后，按照不同的资源约束需求，能够进行网络参数的裁剪、以获得子网络，并进行子网络的微调训练；从正则训练、到子网络微调，可协同知识蒸馏训练，进一步保持子网络的预测精度。主要步骤描述如下：

第一步： 基于 Neuron与Attention head 的重要性，执行参数重排（Weight Re-wiring）：

第二步： 多阶段微调训练，实现宽度、与深度方向的自适应正则化：

4.3 NAS-Bert

参考博客文章：

NAS-Bert——确保One-shot与Task-agnostic：https://blog.csdn.net/nature553863/article/details/120764285

4.4 Evolved Transformer

论文标题：

Lite Transformer with Long-Short Range Attention

收录会议：

ICLR 2020

论文链接：

https://arxiv.org/abs/2004.11886

代码链接：

https://github.com/mit-han-lab/lite-transformer

针对序列生成任务，基于 NAS 搜索获得的 Transformer 结构:

搜索空间： 包括两个 Stackable cell，分别包含在 Transformer encoder 与 Transformer decoder。每个 Cell 由 NAS-style block 组成, 可通过左右两个 Block 转换输入 Embedding、再聚合获得新的 Embedding，进一步输入到 Self-Attention 模块。
搜索策略： 基于 EA（Evolutional Aligorithm）的搜索策略；

网络结构如下，融合了一维卷积与 Attention 的特点：

4.5 HAT: Hardware-aware Transformer

论文标题：

HAT: Hardware-Aware Transformers for Efficient Natural Language Processing

收录会议：

ACL 2020

论文链接：

https://arxiv.org/abs/2005.14187

代码链接：

http://github.com/mit-han-lab/hardware-aware-transformers.git

对 Transformer 模型进行网络结构搜索时，通常会面临两个关键问题：

FLOPS 不能完全代表 Transformer 模型的计算速度/实际执行效率，即只能间接反映硬件平台特性；
不同硬件设备上，与计算硬件或计算库相适配的最优 Transformer 结构不尽相同;

针对序列生成任务，HAT（Hardware-aware Transformer）从搜索空间、搜索策略与搜索预测三方面加以分析：

搜索空间： Encoder-decoder Attention 的任意连接方式，以及 Transformer layer 内部结构（网络深度、Attention head 数目、Width 等）；

搜索策略： 训练了 SuperTransformer 作为超网络，并在超网络预训练之后，结合资源约束，通过进化搜索算法（EA: Evolutional Algorithm）寻找最优子网络；不同于权重共享型 NAS（如 FBNet、SPOS 等），HAT 是一种动态可伸缩类型的One-shot NAS；

搜索预测： 单独、离线训练了 MLP 回归模型作为 Cost model 或 Predicter，用于预测不同子网络结构（Architecture embedding）、在指定硬件平台上的执行速度（GPU 或 CPU latency），并且回归预测效果良好（以相关性系数作为Metric）。搜索预测的好处在于：一方面可直接、有效获取硬件平台特性，作为超网络预训练的资源约束；另一方面，在训练或搜索阶段，无需在硬件平台上测量推理延迟，提升搜索效率，从而有利于跨平台模型结构搜索；

4.6 AutoFormer

论文标题：

AutoFormer: Searching Transformers for Visual Recognition

收录会议：

ICCV 2021

论文链接：

https://arxiv.org/abs/2107.00651

代码链接：

https://github.com/microsoft/AutoML

针对 Vision Transformer 的结构搜索，AutoFormer 提出的 Weight entanglement，在不额外增加 Block choice 的前提下，通过通道宽度、网络深度、Attention head 数目等多个维度的调整，实现了 Vision Transformer 模型的动态可伸缩预训练与结构搜索。

Weight entanglement 的做法，类似于 BigNAS、FBNet-v2 的通道搜索，都不会额外增加通道维度的权重参数量。相比于手工设计的 CNN 模型（ResNet、ResNext、DenseNet）与 Vision Transformer 模型（ViT、Deit），AutoFormer 模型在相同资源开销条件下、能够获得最好的识别精度。

AutoFormer 总体思路如上图所示，是一种基于 Weight entanglement 的动态可伸缩搜索方法，其搜索维度包括 Attention heads、通道宽度与网络深度。下面从搜索空间、搜索策略与搜索效率这些维度加以分析：

搜索空间： Embedding dimension、Q-K-V dimension、Attention heads、MLP ratio 与 Network depth；根据不同的资源约束，分别设置 Supernet-tiny、Supernet-small 与 Supernet-base 三个基本的 Template model；

搜索策略： 基于 Weight entanglement 的 One-shot NAS，完成超网络预训练之后，通过进化算法执行子网络搜索。包含 l 个网络层的子网络结构可按下式表示，表示第 i 层的 Block 结构，表示相应的权重参数：

而每个 Block 结构，都是按照 Weight entanglement 原则从超网络采样获得，下式表示 n 个动态选择范围：

搜索效率：按照权重共享形式、完成超网络预训练，训练效率高、且收敛速度快；另外，由于没有引入额外的 Block choice，因此训练时的 Memory cost 较低。

计算图优化

5.1 LINM

LINM（Loop-invariant Node Motion） 是一种计算图等价变换技术，通过将 Transformer 模型涉及的自回归解码（ Auto-regressive Decoding）的重复计算逻辑（Encoder-decoder Attention 的 k/v 计算、cache gather 等）移至 While-loop 之外，在确保计算功能不变的情况下，实现计算效率的提升：

5.2 MatMul替换为卷积

如果推理引擎的卷积算子具备很高的计算效率，并且为了避免一些不必要的访存算子开销，可以按如下计算图等价变换，将 Dense layer 替换为 Conv1D layer：

5.3 QKV计算合并

MHA（Multi-head Attention）模块在计算 Attention map 与 Context 编码特征之前，需要基于输入 Tensor 计算 QKV 矩阵、并 Split 成多份 Heads。可通过计算图等价变换，将 QKV 计算合并到同一个 Dense layer，能够充分利用 NPU 或 GPU 的并发计算特性：

1B, N, C = x.shape
2qkv = self.qkv(x).reshape(B, N, 3, self.sample_num_heads, -1).permute(2, 0, 3, 1, 4)
3q, k, v = qkv[0], qkv[1], qkv[2]

推理优化引擎

6.1 Faster Transformer

GitHub链接：

https://github.com/NVIDIA/FasterTransformer

Faster Transformer 是 NVIDIA 针对标准 BERT/Transformer 做的推理优化库，其发布时间线：

● 1.0版本：

2019 年 7 月，开源了 FasterTransformer 1.0，针对 BERT 中的 Transformer Encoder 进行优化和加速；
面向 BERT 分类、自然语言理解场景；
底层由 CUDA 和 cuBLAS 实现，支持 FP1 6 和 FP32 计算，FP16 可充分利用 Volta 和 Turing 架构的 Tensor Core 计算单元；
提供 C++ API、TF Op 与 TensorRT Plugin 三种接口；
参考资料： https://zhuanlan.zhihu.com/p/73715272

● 2.0版本：

2020 年 2 月，新增对 Transformer decoder 的优化和加速，包括 decoder与 decoding 两种加速模式；
面向生成式场景，如 NMT、文本内容生成与 ASR 等；
底层由 CUDA 和 cuBLAS 实现，支持 FP16 和 FP32 计算模，FP16 可充分利用 Volta 和 Turing 架构的 Tensor Core 计算单元；
提供 C++ API、TF Op 与 TensorRT Plugin 三种接口；
参考资料：快无止尽！FasterTransformer 2.0让Decoder不再是你的性能瓶颈

● 2.1版本：

2020 年 6 月，引入 Effective Transformer 优化；通过 remove_padding 的支持，提高计算与访存效率；
并新增 PyTorch Op 接口；
参考资料： https://github.com/bytedance/effective_transformer

● 3.0版本：

2020 年 9 月，新增 BERT encoder 的 INT8 量化加速支持；
仅支持 Turing 架构 GPU；
同时支持 PTQ 与 QAT 方法，提供了 TF 量化工具；
相比于 FP16 计算，约 20~30% 加速，但存在精度损失风险。

● 3.1版本：

2020 年 12 月，新增对 PyTorch 使用 INT8 推理的支援；
在 Turing 以后的 GPU 上，FP16 的性能比 3.0 提升了 10% ~ 20%；
INT8 的性能比 3.0 最多提升了 70%。

● 4.0版本：

2021 年 4 月，新增对 GPT-3 等百亿/千亿级参数规模模型的多机多卡推理加速能力；
新增 FP16 fused MHA 算子，同时支持 Volta 与 Turing 架构的 GPU；
以及对解码端 Kernel 的优化，可以省略已完成语句的计算，节省计算资源；
参考资料：https://zhuanlan.zhihu.com/p/363517823；

● 支持矩阵：

Transformer 计算图表示，主要包含三个部分的表示：

Encoder layer；
Decoder layer；
Decoding logic；

6.2 Torch-TensorRT

参考博客文章：

深度学习模型压缩与优化加速（Model Compression and Acceleration Overview）

https://blog.csdn.net/nature553863/article/details/81083955

6.3 AICompiler

基于 AICompiler（AI 编译优化），可通过 Op fusion 方式（算子融合），提升 Transformer 模型的执行效率。Op fusion 具备的好处主要是：1）首先，将多个零碎算子合并成计算功能等价的一个大算子，可减少密集的 Kernel launch 开销（CUDA 核函数启动）；2）其次，由于相邻算子之间存在内存读写开销，Op fusion 可有效减少密集的内存访问开销；3）最终通过算子实现的自动寻优、与代码生成（CodeGen），可自动生成优化的可执行代码。

总结

深度学习领域的高性能计算， 总结为如下公式：

高性能计算 = 高效率算法 + 模型压缩 + 系统/硬件优化

可以单独的强调其中一种优化策略，或者实现多种选项的联合优化（例如 Hardware-aware NAS，压缩与编译联合优化等），以满足实际应用场景的部署需求。在工业界，包括阿里淘系 MNN、阿里云 PAI、华为昇腾计算、Open AI Lab 的 Tengine、百度 EasyDL 等平台，在该领域均构建了相对完善的软硬件技术栈，为开放、普惠的 AI 应用提供了坚实基础。