浅谈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: