浅谈BERT/Transformer模型的压缩与优化加速
©作者 | 姚益武
单位 | 阿里巴巴集团
研究方向 | 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 实现，支持 FP16 和 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 应用提供了坚实基础。
特别鸣谢
感谢 TCCI 天桥脑科学研究院对于 PaperWeekly 的支持。TCCI 关注大脑探知、大脑功能和大脑健康。
更多阅读
#投 稿 通 道#
让你的文字被更多人看到
如何才能让更多的优质内容以更短路径到达读者群体，缩短读者寻找优质内容的成本呢？答案就是：你不认识的人。
总有一些你不认识的人，知道你想知道的东西。PaperWeekly 或许可以成为一座桥梁，促使不同背景、不同方向的学者和学术灵感相互碰撞，迸发出更多的可能性。
PaperWeekly 鼓励高校实验室或个人，在我们的平台上分享各类优质内容，可以是最新论文解读，也可以是学术热点剖析、科研心得或竞赛经验讲解等。我们的目的只有一个，让知识真正流动起来。
稿件基本要求：
• 文章确系个人原创作品，未曾在公开渠道发表，如为其他平台已发表或待发表的文章，请明确标注
• 稿件建议以 markdown 格式撰写，文中配图以附件形式发送，要求图片清晰，无版权问题
• PaperWeekly 尊重原作者署名权，并将为每篇被采纳的原创首发稿件，提供业内具有竞争力稿酬，具体依据文章阅读量和文章质量阶梯制结算
投稿通道：
• 投稿邮箱：[email protected]
• 来稿请备注即时联系方式（微信），以便我们在稿件选用的第一时间联系作者
• 您也可以直接添加小编微信（pwbot02）快速投稿，备注：姓名-投稿
△长按添加PaperWeekly小编
现在，在「知乎」也能找到我们了
进入知乎首页搜索「PaperWeekly」
点击「关注」订阅我们的专栏吧
·
