Bert/Transformer模型压缩与优化加速

前言

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)；

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

Lite-Transformer

参考文献: https://arxiv.org/abs/2004.11886

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：

SAN-M

参考文献: 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)相加：

MiniLM (知识蒸馏)

参考文献: https://arxiv.org/abs/2002.10957

针对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：

  

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

Switch-Transformer

参考文献：https://arxiv.org/abs/2101.03961

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