原文作者：Grigory Sapunov

翻译：Camel

原文标题：Speeding up BERT

原文链接：https://blog.inten.to/speeding-up-bert-5528e18bb4ea

雷锋网AI科技评论按：BERT 在 NLP 领域的地位正如 ResNet 在 CV 领域的地位一样，属于里程碑的进展。目前，BERT 已经成为 NLP 深度学习管道中的重要组成部分。

沉重的BERT

但 BERT 并不是完美无瑕的，它最大的问题就是：太大。

BERT-base 模型能够包含一亿个参数，较大的 BERT-large 甚至包含 3.4 亿个参数。显然，很难将这种规模的模型部署到资源有限的环境（例如移动设备或嵌入式系统）当中。

模型太大是其一，BERT 的训练和推理时间也太长了！

在基于 Pod 配置的 4 个 Cloud TPUs（总共 16 个 TPU 芯片）上对 BERT-base 进行训练，或者在 16 个 Cloud TPU（总共 64 个 TPU 芯片）上对 BERT-large 进行训练，每次预训练都需要至少 4 天的时间才能完成。

当然对于最终的用户而言，训练时间通常似乎并不是什么大问题，因为反正只需要训练一次就够了（实际上往往会训练多次，因为你可能要多次重新训练模型，才能获得满意的结果）。不过如果能够加速，当然也是好的，因为迭代越快，你就能越早地解决你的问题。

BERT 的推理时间可能取决于你用的是什么设备（什么型号、多少个），在许多情况下，它会极大地限制你要处理数据的数量、速度以及成本。对于某些实时的应用程序来讲，用 BERT 简直是让人望而却步。

对以上的问题，能够想到的最直接的方法就是优化神经网络。这在神经网络中并不陌生，其他领域（例如计算机视觉）以前也有相同的问题，并且已经开发了几种压缩和加速神经网络模型的方法。大致可以分为几类：

• 架构改进：将原有的架构改进为更快的架构，例如，将 RNN 替换为 Transformer 或 CNN；使用需要较少计算的层等。当然也可以采用其他优化，例如从学习率和策略、预热步数，较大的批处理大小等；

• 模型压缩：通常使用量化和修剪来完成，从而能够在架构不变（或者大部分架构不变）的情况下减少计算总量；

• 模型蒸馏：训练一个较小的模型，该模型能够复制原始模型的行为。

我们来看下 BERT 在这些策略上可以做什么。雷锋网

1、架构和优化改进

大规模分布式训练

加速 BERT 训练能够想到的第一件事（甚至是第零件事）是将其分发到更大的集群当中。虽然最初的 BERT 已经使用了多台机器进行训练，但还有更多针对 BERT 分布式训练的优化解决方案，例如阿里巴巴 [1] 和英伟达 [2] 的一些工作。

英伟达最近使用 NVIDIA DGX SuperPOD（具有 92 个 DGX-2H 节点，共有 1472 个 V100 GPU，理论上可以提供 190PFLOPS）刷新了 BERT 训练的记录，在 53 分钟内训练出了 BERT-large 模型 [3]。当然，这个训练也是非常昂贵的，除了英伟达，其他公司和个人很难去做这样的事情。

还有另外一种更为巧妙优化方案，这是一种被称为 LAMB 的新的分层自适应大批量优化技术 [4]，这种技术可以将 TPUv3 Pod 上的 BERT 训练时间从 3 天降低到 76 分钟（1024 个 TPUv3 芯片，可以提供超过 100PFLOPS），显然，训练这个模型也不会便宜。

架构

现在考虑更加架构性的以及更少硬件的解决方案。

有一种基于对自注意层行为的观察来训练 BERT 的渐进式堆叠方法 [5]，该方法的分布局部地集中在其位置和 start-of-sentence token 周围，且在浅层模型中的注意力分布与深层模型相似。为此，作者提出了将知识从浅层模型转移到深层模型的堆叠算法，并逐步应用堆栈来加快 BERT 训练的速度。凭借这种方式，作者获得模型的训练时间相比原始 BERT 的训练时间缩短了约 25%，归其原因则在于对于相同数量的步骤，训练一个小的模型需要的计算也相对较少。

还有一种方法是注意力矩阵的稀疏分解（也称 Sparse Transformer，这项研究是由 OpenAI 完成的)[6] 和 block attention[7]，它们都属于能够减少内存和计算总量的体系结构改进。

ALBERT

最后不得不提的是，已经提交给 ICLR 2020 的一个 BERT 的后代，被称为 ALBERT（A Lite BERT）[8]。

ALBERT 结合了两种参数降低的技术。

第一种是分解式嵌入参数化，将隐藏层的大小与词嵌入的大小分开。这种分隔使得在不显著增加词汇表嵌入参数大小的情况下能够更容易增加隐藏层的大小。

第二种是跨层参数共享。这种技术可以防止参数随着网络深度的增加而增大。

这两种技术可以显著降低 BERT 的参数数量，而不会严重影响性能，从而提高参数效率。

与 BERT-large 具有相似配置的 ALBERT，其参数能够降低 18 倍，训练速度提高了月 1.7 倍。

相比严格调校的 RoBERTa，ALBERT 则同样更胜一筹 [9]。

2、量化和修剪

量化 会降低模型权重的数值精度。通常情况下，使用 FP32（32 位浮点）来训练模型，然后可以将其量化为 FP16（16 位浮点），INT8（8 位整数），甚至将其量化为 INT4 或 INT1。于是模型尺寸将随之减小 2 倍、4 倍、8 倍或 32 倍。这称为 post-training quantization。

另一个选项是量化感知训练（也是较为困难和较为不成熟的）。这种方法的 FP16 训练现在已成为一种普遍的方法。而在 ICLR 2020 提交的论文中有一个比较有趣的工作 [10]，它使用 ResNet、GNMT、Transformer 的 8 位浮点表示获得了最先进的训练结果。

修剪 即从模型中删除一些不重要或不太重要的权重（有时会是神经元），从而产生稀疏的权重矩阵（或较小的图层）。甚至还有一些研究直接去除掉与 Transformer 的注意力头相对应的整个矩阵。

量化，可以使用 TensorFlow Lite（用于设备上推断的 TensorFlow 的部分）[11] 来执行。TensorFlow Lite 提供了在移动设备、嵌入式设备和 IoT 设备上转换和运行 TensorFlow 模型的工具，它支持训练后量化和量化感知训练。

另一个选择是使用英伟达的 TensorRT 框架 [12]。英伟达 TensorRT 是用于高性能深度学习推理的平台，它包括深度学习推理优化器，并且在运行时能够为深度学习推理应用程序提供低延迟和高吞吐量。

英伟达最近发布了具有最新优化功能的 TensorRT 6[13]，它可以在 T4 GPU 上仅用 5.8 毫秒对 BERT-large 进行推理，而在 V100 上甚至只需要 4.2 毫秒，对于 Titan RTX，速度可能会更快。

当批处理数为 1 时，对于 340M 个参数的 BERT-large 模型，仅需 5.84 毫秒；对于 110M 参数的 BERT-Base 模型，则只需 2.07 毫秒。若将批处理数设为 128，你在使用 BERT-large 时甚至可以达到 250 个句子/秒的处理速度。

更多的数据我们这里就不再一一列举了。

PyTorch 最近也宣布了在其 1.3 版本 [14] 中支持量化。尽管目前它还处于实验阶段，但我们已经可以使用它了，因为在其教程中提到他们已经能够将动态量化应用于将模型权重转换为 INT8 的 LSTM 语言模型 [15]。

此外，还有一个众所周知的 BERT 量化，称为 Q-BERT。这项工作来源于「Q-BERT: Hessian Based Ultra Low Precision Quantization of BERT」[16]。在这个研究中，作者甚至降低到 2-bit 的超低精度量化，但其性能相比于基线却没有显著下降（仅下降 2.3%），而对应的模型参数压缩率最高可以达 13 倍，嵌入表压缩率和激活的最高都为 4 倍。

3、蒸馏

另一个有趣的模型压缩方法是蒸馏，这是一种将大型「teacher」网络的知识转移到较小的「student」网络的技术，训练学生网络来模仿教师网络的行为。

Rich Caruana 及其合作者率先采用了这种策略。在他们先驱性的论文 [17] 中，他们提供了令人信服的证明：大型集成模型所获得的知识可以转移到单个小型的模型中。

Geoffrey Hinton 等人在他们的「Distilling the Knowledge in a Neural Network」{18} 论文中证明了这种技术可以应用于神经网络模型。

DistilBERT

从 Hinton 开始，蒸馏的方法逐渐被应用到了不同的神经网络中，当然你可能也听过 HuggingFace 提出的 DistilBERT，这是一种对 BERT 的蒸馏。这项工作出自论文「DistilBERT, a distilled version of BERT: smaller, faster, cheaper and lighter」[19]，目前已经提交到 NeurIPS 2019。

DistilBERT 是一种较小的语言模型，受 BERT 的监督而训练。在该模型中，作者删除了令牌类型嵌入和合并器（用于下一个句子分类任务），并保持体系架构其余部分不变，同时将层数减少了两倍。

您可以在 HuggingFace（以前叫做 pytorch-transformers 和 pytorch-pretrained-bert）的 translators python 软件包的帮助下，使用现成的 DistilBERT。该软件包的 2.0.0 版本支持 TensorFlow 2.0 / PyTorch 互操作。

DistilBERT 的作者还使用了最新 RoBERTa 论文中的一些训练技巧，这些技巧表明 BERT 的训练方式对其最终性能至关重要。

DistilBERT 与 BERT 相比具有出人意料的结果：作者保留了 95％以上的性能，但参数却减少了 40％。

比较 GLUE 基准测试的开发集

在推断时间方面，DistilBERT 比 BERT 快了 60％以上，比 ELMo + BiLSTM 快 120％。

推理速度

TinyBERT

就在几天前，出现了一种新的 BERT 蒸馏方法，来自华为诺亚方舟实验室的刘群团队提出了 TinyBERT[20]。

为了构建一个具有竞争力的 TinyBERT，作者首先提出了一种新的 Transformer 蒸馏方法，来蒸馏嵌入 BERT 的知识。

具体来说就是，他们设计了几个损失函数来适合 BERT 层的不同表示形式：

1、嵌入层的输出；

2、Transformer 层派生的隐藏状态和注意力矩阵;

3、预测层输出的 logits 输出。

论文中基于注意力的拟合则得益于最近的发现，即 BERT 学习的注意力权重可以捕获大量的语言知识，这意味着语言知识可以很好地从教师 BERT 转移到学生 TinyBERT。而在 BERT 的现有知识蒸馏的方法（如 Distilled BiLSTM_SOFT，BERT-PKD 和 DistilBERT）中却忽略了这一点。

在这项工作中，作者还提出了一种新颖的两阶段学习框架，包括通用蒸馏和特定任务蒸馏。在通用蒸馏阶段，未经微调的原始 BERT 充当教师模型，学生 TinyBERT 通过在通用领域对大型语料库执行通常的 Transformer 蒸馏来学习模仿教师的行为。他们获得了可以对各种下游任务进行微调的通用 TinyBERT。在特定任务蒸馏阶段，他们将数据进行扩充，来提供更多与任务相关的材料供教师-学生学习，然后在增强的数据上重新执行 Tranformer 蒸馏。

这个两阶段方法对于提升 TinyBERT 的性能和泛化能力是必不可少的。

TinyBERY 在实验上取得了非常的成绩，相对于 GLUE 数据集的 BERT-base，其性能并没有下降多少，而推理参数小了 7.5 倍，推理时间快了 9.4 倍。

TinyBERT 与其他基准的比较

我们期待他们能够将这种方法应用到 BERT-large 和 XLNet 等大型模型中，同样也期待他们开放源码。

其他蒸馏方法

除了 DistilBERT 和 TinyBERT 外，还有其他一些为大家所熟知的蒸馏方法。

（2019/03）「Distilling Task-Specific Knowledge from BERT into Simple Neural Networks」[21]

这篇论文将 BERT 蒸馏到单层 BiLSTM 中，取得了与 ELMo 可比的结果，同时使用的参数减少了大约 100 倍，推理时间减少了 15 倍。

BiLSTM_SOF 是 TBiLSTM 的蒸馏，后者是在 soft logit 目标上训练出来的。

（2019/08）「Patient Knowledge Distillation for BERT Model Compression」[22]

这篇论文提出了一种耐心知识蒸馏的方法，这是首次尝试使用教师的隐藏状态，而不仅仅是最后一层的输出。他们的学生模型从教师模型的多个中间层「耐心」地学习来获得更多知识。在他们的耐心蒸馏知识框架中，只训练学生模仿中间层的 [CLS] 令牌的表示形式。代码已公开 [23]。

（2019/09）「Extreme Language Model Compression with Optimal Subwords and Shared Projections」

这是最近提交到 ICLR 2020 的一篇论文，这篇论文专注于一种用于训练词汇量显著较小、嵌入和隐藏状态维度较低的学生模型的知识蒸馏技术。作者采用了双重训练机制，可以同时训练教师和学生模型，从而能够获得针对学生词汇的最佳词嵌入。该方法能够将 BERT-base 模型压缩 60 倍以上，而下游任务指标只有很小的下降，从而使得语言模型占用的空间只有不到 7MB。

TinyBERT 的结果似乎更好，但一个 7MB 的类 BERT 模型简直爽的不能再爽！

需要强调，以上所介绍的方法并不互相冲突，所以我们期待能够将这些方法协同起来使用来加速 BERT 模型（或其他相似模型）。

参考资料

[1]https://www.alibabacloud.com/blog/perseus-bert-industry-leading-bert-training-solution-optimized-for-performance_594717

[2]https://github.com/NVIDIA/Megatron-LM

[3]https://devblogs.nvidia.com/training-bert-with-gpus/

[4]https://arxiv.org/abs/1904.00962

[5]http://proceedings.mlr.press/v97/gong19a/gong19a.pdf

[6]https://arxiv.org/abs/1904.10509

[7]https://arxiv.org/abs/1804.00857

[8]https://openreview.net/pdf?id=H1eA7AEtvS

[9]https://blog.inten.to/papers-roberta-a-robustly-optimized-bert-pretraining-approach-7449bc5423e7

[10]https://openreview.net/forum?id=HJe88xBKPr

[11]https://www.tensorflow.org/lite

[12]https://developer.nvidia.com/tensorrt

[13]https://news.developer.nvidia.com/tensorrt6-breaks-bert-record/

[14]https://pytorch.org/blog/pytorch-1-dot-3-adds-mobile-privacy-quantization-and-named-tensors/

[15]https://pytorch.org/tutorials/advanced/dynamic_quantization_tutorial.html

[16]https://arxiv.org/abs/1909.05840

[17]https://www.cs.cornell.edu/~caruana/compression.kdd06.pdf

[18]https://arxiv.org/abs/1503.02531

[19]https://arxiv.org/abs/1910.01108

[20]https://arxiv.org/abs/1909.10351

[21]https://arxiv.org/abs/1903.12136

[22]https://arxiv.org/abs/1908.09355

[23]https://github.com/intersun/PKD-for-BERT-Model-Compression

[24]https://arxiv.org/abs/1909.11687

雷锋网(公众号：雷锋网)报道。