EdgeBERT：极限压缩，比ALBERT再轻13倍！树莓派上跑BERT的日子要来了？

2021 年 1 月 14 日 机器之心

机器之心转载

来源：夕小瑶的卖萌屋

作者：Sheryc_王苏

这个世界上有两种极具难度的工程：第一种是把很平常的东西做到最大，例如把语言模型扩大成能够写诗写文写代码的 GPT-3；而另一种恰恰相反，是把很平常的东西做到最小。对于 NLPer 来说，这种 “小工程” 最迫在眉睫的施展对象非 BERT 莫属。

从 18 年那个 109M 参数的 BERT，到 52M 参数的蒸馏后的 DistilBERT，再到 14.5M 参数的蒸馏更多层的 TinyBERT，最后到 12M 参数的层级共享的 ALBERT，曾经那个在集群上加载参数都费劲的 BERT 现在甚至已经可以跑在手机平台上了。当我们为 BERT 的轻量化欢呼雀跃之时，有这样一群人站了出来——只是手机端可不够！他们的理想，是让 BERT 跑在物联网设备上，跑在低功耗芯片上，跑在我们能触及的每一个电子器件上！

这样一群来自哈佛 / 塔夫茨 / HuggingFace / 康奈尔的软件和硬件极客们，此刻已披上了法袍，化身为为 BERT 极限瘦身的炼金术士，向着这个看似不可能的目标添加着许多让人意想不到的配方…

论文题目：EdgeBERT: Optimizing On-Chip Inference for Multi-Task NLP
论文链接：https://arxiv.org/pdf/2011.14203.pdf

配方基底：ALBERT

出处：ALBERT: A Lite BERT for Self-supervised Learning of Language Representations（ICLR'20）
链接：https://arxiv.org/pdf/1909.11942.pdf

EdgeBERT 是在 ALBERT 的基础上进行优化的。

ICLR'20 上谷歌提出的 ALBERT 是目前最佳的 BERT 压缩方案。相比于过去利用知识蒸馏从原始 BERT 模型得到压缩模型（例如 DistilBERT [1]、TinyBERT [2]）和利用浮点数量化得到压缩模型（例如 Q8BERT [3]），ALBERT 选择直接抛弃 BERT 的预训练参数，只继承 BERT 的设计思想。正所谓不破不立，继承 BERT 灵魂的 ALBERT 仅用 12M 参数就取得了和其他 BERT 变体相当的性能。

ALBERT 对 BERT 的设计采取了以下三点改进：

嵌入层分解：BERT 中，WordPiece 的嵌入维度和网络中隐藏层维度一致。作者提出，嵌入层编码的是上下文无关信息，而隐藏层则在此基础上增加了上下文信息，所以理应具有更高的维数；同时，若嵌入层和隐藏层维度一致，则在增大隐藏层维数时会同时大幅增加嵌入层参数量。ALBERT 因此将嵌入层进行矩阵分解，引入一个额外的嵌入层 E。设 WordPiece 词汇表规模为V，嵌入层维度为 E，隐藏层维度为 H，则嵌入层参数量可由 O（V x H）降低为O（V x E + E x H）。

参数共享：BERT 中，每个 Transformer 层参数均不同。作者提出将 Transformer 层的所有参数进行层间共享，从而将参数量压缩为仅有一层 Transformer 的量级。

上下句预测任务→句序预测任务：BERT 中，除语言模型的 MLM 任务外还进行了上下句预测任务，判断句 2 是否为句 1 的下一句，然而该任务被 RoBERTa 和 XLNET 等模型证实效果一般。作者提出将其替换为句序预测任务，判断句 2 和句 1 之间的句子顺序来学习文本一致性。

ALBERT 的设计相当成功，成为了压缩 BERT 的经典范例，而为了做到最极限压缩 BERT，以 ALBERT 作为起点确实是个不错的主意。ALBERT 已经如此强大，EdgeBERT 又能压到什么程度？作者一上来就用一张在 QQP 上的内存占用 / 运算量 / 性能比较图吊足了读者的胃口。（注意：纵坐标的内存占用量是对数刻度！）

本文对于 ALBERT 的利用除了作为初始化参数外，还在对下游任务做 fine-tune 时利用已经 fine-tune 好的 ALBERT 作为 teacher 进行知识蒸馏，来进一步提升模型性能。

初级配方：算法优化

1. 基于熵的提前退出机制

出处：DeeBERT: Dynamic Early Exiting for Accelerating BERT Inference（ACL'20）
链接：https://arxiv.org/pdf/2004.12993.pdf

ALBERT 虽好，但 Transformer 太深了，算起来太慢，让他变浅一点怎么样？

ACL'20 的 DeeBERT 恰好就提出了一种动态的提前退出机制（Early Exit）。这一机制的设计是希望让简单的文本经过较少的运算，而复杂的文本经过更多的运算。

实现上，DeeBERT 向 n 层的 BERT 模型添加了（n-1）个 “出口层” 分类器（Early Exit Off-Ramps）。出口层分类器 f_i 被放置在第 i 和 i+1 层 Transformer 之间，作为判断第 i 层 Transformer 的信息是否足以进行推断的标志。进行推断时，从最底层的分类器开始逐层计算出口层分类器的熵，当某一层的出口层分类器熵小于某个阈值时，则将该出口层分类器结果作为模型结果，省去后续层的计算。

上图表现的是不同的熵阈值在 MNLI、QQP、SST-2、QNLI 四个数据集上的平均退出层数、理论运行时间节省和相应的准确度。添加提前退出机制后，在 Acc 损失 1 个百分点时，能够在这四个数据集上分别降低 30%、45%、54% 和 36% 的理论运行时间；在 Acc 损失 5 个百分点时，能将在这四个数据集上的理论运行时间的降低进一步降低至 44%、62%、78% 和 53%。

2. 动态注意力范围

出处：Adaptive Attention Span in Transformers（ACL'19）
链接：https://arxiv.org/pdf/1905.07799.pdf

ALBERT 虽好，但 Attention 范围太广了，算起来太慢，让他变窄一点怎么样？

ACL'19 的 Adaptive Attention 提出的动态注意力范围正是试图通过这种方法减少注意力计算的。在 Transformer 的多头自注意力机制中，不同 head 对应的注意力范围并不一致，而让每一个 head 都对所有 token 进行注意力运算无疑增添了无用开销。为此，Adaptive Attention 为每一个 head 添加一个不同的 mask，使得每个 token 只能对周边的 token 计算注意力，从而降低矩阵运算的开销。

具体地，mask 函数

基于两个 token 之间的距离为注意力机制的权重计算添加了 soft masking。注意力机制中的权重

变为：

其中

为控制 soft 程度的超参数，

为序列截止到 token

的长度（原文采用了 Transformer Decoder 结构学习语言模型，故每个 token 只能于自己之前的 token 计算注意力。在 EdgeBERT 中没有提及公式，不过根据模型图的结构来看，分母应修改为对整个序列求和）。mask 函数中的

为 mask 的边界，此边界值会跟随注意力的 head 相关参数和当前输入序列变化：对于注意力机制中的每一个 head

，有

，其中

、

可训练，

为 sigmoid 函数。

EdgeBERT 甚至对 Adaptive Attention 又进一步做了简化：

连算都不用算了，直接给每一个 head 赋一个可学习的

，连输入序列都不考虑了，多出来的参数只有 12 个

（因为有 12 个 head）。那么，这样做的结果如何呢？作者将所有序列都 pad/trunc 到 128 长度，经过实验，得到了一个惊人的结果：

表中展示的是经过优化后各个 head 的

值，和模型在 MNLI/QQP/SST-2/QNLI 四个任务上的准确度。在一大半 head 几乎完全被 mask 掉（

）之后，模型居然只在这几个任务上掉了 0.5 甚至 0.05 的准确度！而这一方法也为模型带来了最高

的计算量降低。

3. 一阶网络剪枝

出处：Movement Pruning: Adaptive Sparsity by Fine-Tuning（NeurIPS'20）
链接：https://arxiv.org/pdf/2005.07683.pdf

ALBERT 虽好，但参数存起来占用的内存太长了，开销太大，让他变短一点怎么样？

这里的网络剪枝方式使用到了 NeurIPS'20 的一篇针对模型 Fine-tune 过程的剪枝算法。该论文的作者提出，传统的零阶网络剪枝（即给模型里的参数绝对值设定一个阈值，高于它的保留，低于它的置零）的方法并不适用于迁移学习场景，因为该场景下模型参数主要受原始模型影响，却需要在目标任务上进行 fine-tune 和测试，所以直接根据模型参数本身剪枝可能会损失源任务或目标任务的知识。与此相反，作者提出一种基于 Fine-tune 过程中的一阶导数进行剪枝的 Movement Pruning：尽可能保留 fine-tune 过程中更加偏离 0 的参数。

具体地：对于模型参数

，为其赋予同样 size 的重要性分数

，则剪枝 mask

。

前向传播过程中，神经网络利用加 mask 的参数计算输出的各分量：

。

反向传播过程中，利用 Straight-Through Estimator[4]的思想，将

省略近似得到损失函数

对重要性分数

的梯度：

对模型参数，有：

将上述两个式子代换后，省略的 mask 矩阵后可得：

根据梯度下降，当

时，重要性

增大，此时

与

异号。这表示，只有当在反向传播时为正的参数变得更大或为负的参数变得更小时才会得到更大的重要性分数，避免被剪枝。

4. 零阶网络剪枝

出处：Deep Compression: Compressing Deep Neural Networks with Pruning, Trained Quantization and Huffman Coding（ICLR'16）
链接：https://arxiv.org/pdf/1510.00149.pdf

变短是变短了，但感觉这剪得还不够好啊，再换一种算法让它更短一点怎么样？

这种方法的做法非常简单：给模型里的参数设定一个绝对值阈值，绝对值高于它的保留，绝对值低于它的置零。由于方法实在太过简单，不用公式也能很容易理解的吧

一阶和零阶网络剪枝的效果对比如上图所示（MvP：一阶网络剪枝，MaP：零阶网络剪枝）。在参数稀疏程度更高时，一阶剪枝的效果更好，其他情况下是简单的零阶剪枝更有效。同时，研究还发现，当剪掉 95% 的嵌入层参数时，模型竟然在 4 个任务上都保持了至少 95% 的准确度。

5. 动态浮点数量化

出处：AdaptivFloat: A Floating-point based Data Type for Resilient Deep Learning Inference（arXiv Preprint）
链接：https://arxiv.org/pdf/1909.13271.pdf

诶，怎么还有？网络的运算和存储过程从深度、宽度和长度进行了全方位的优化，模型还能再轻？

的确，到此为止，一般的炼丹师已经看着自己三维裁剪过的模型感到成就感满满了，但这距离让 BERT 跑遍所有设备的目标还差得远。以下的部分就要开始深入一般 NLP 工程师见不到的第四个维度——硬件维度了。在深入硬件优化之前，先来道偏软件的开胃菜，看看如何在浮点数量化上下手优化存储吧！

我们平时在考虑利用浮点数特性进行计算加速时，最先考虑的是使用 FP16 混合精度，虽然有效，但毕竟损失了信息，效果也多多少少会受影响。既想保留精度又想加速训练减小存储，那就只有深入底层，修改浮点数的表示方法了！

这正是 AdaptivFloat 的初衷：设计一种更加适合深度学习场景的浮点数数据类型。不过，要讲明白 AdaptivFloat 数据类型需要涉及到一些与机器学习无关的知识。

根据 IEEE 754 二进制浮点数标准，一个浮点数的二进制表示由三个域组成：符号位（Sign，S）、指数偏移值（Exponent bias，E）和分数值（Fraction，或 Mantissa，F）。由此，一个数可以表示为

。

这时聪明的你可能发现有什么不对劲：按照表示无符号整形的方法取指数偏移值只能取出正数啊！2 的负次幂怎么办！这正是为什么称其为“指数偏移值”：它并不代表实际上 2 的指数，而是在其基础上需要加一个常数作为 2 的指数：

。

我们常用的浮点数保证了

的选取能够使得

在数轴两侧分布几乎均等（例如在 32 位浮点数 FP32 中，指数范围为至），但这样的数作为机器学习模型的参数显然有些不太合适：为了增加小数的精度，我们甚至要允许

这样显然不会出现的数也能表示，这真的不是在浪费内存？

AdaptivFloat 的最关键动机正在于此：根据模型参数动态修改

。所谓动态体现在每个 Tensor 都能得到量身定做的

。方法也很简单，找到 Tensor 里最大的一个数，让它能被指数范围恰好覆盖到就好。不过说来简单，为了实现这一方法需要配套地对现有浮点数表示方法进行许多其他修改，感兴趣的话可以去看看 AdaptivFloat 原文，此外 IEEE 754 标准 [5] 同样也可以作为参考哦~

上图的结果中，Bit Width 是浮点数的总位数，后五行在模型中使用了 AdaptivFloat 并将指数偏移值限制为 3 个 bit。不比不知道，谁能想到修改个量化方式居然直接用 8bit 在四个数据集上都干翻了 FP32？！不仅节省了 3/4 的内存，还取得了更好的性能，或许这就是吹毛求疵挑战极限的极客们的浪漫？

高级配方：存储介质选择

光软件优化可不够！如果软件优化就够了的话，还买什么 SSD，换什么 GPU（不）

EdgeBERT 的目的，是尽可能降低边缘计算场景中使用 BERT 的推理延迟和耗能。为了最大限度地降低推理延迟，需要为网络中不同的组成部分根据其增删改查的需求选取符合最大效能的存储介质。

BERT 类模型的一大特点，在于它们都是预训练模型：这类模型并非开箱即用，而是需要在目标任务上 fine-tune 后才能使用。这使得这类模型天生地存在着两类存储需求：

嵌入层：保存了 Embedding 向量。EdgeBERT 在进行下游任务 fine-tune 时一般不对嵌入层进行修改。这类参数相当于只读参数，只对快速读取有较高要求，同时希望能够在掉电时依然保持原有数据来降低数据读写开销，因此适用耗能低、读取速度快的 eNVM（Embedded Non-Volatile Memory，嵌入式非挥发性记忆体）。本文选取的是基于 MLC 的 ReRAM，一种低功耗、高速度的 RAM。

其他参数：这些参数需要在 fine-tune 时进行改变。此处使用的是 SRAM（与计算机内存的 DRAM 不同，SRAM 更贵但功耗更低、带宽更高，常被用于制造 cache 或寄存器）

嵌入层用了 ReRAM 究竟能带来多大的影响？上图结果表明，仅仅是改变了嵌入层的硬件介质就能带来约的推理延迟降低，以及低至约的能耗降低！这对于边缘计算场景下简直是质变了！（为何 ReRAM 只有读，但 DRAM 那边却要算 DRAM 读 + SRAM 读 / 写呢？因为此处的 ReRAM 是特殊设计的只读结构，并且可以直接读入处理器进行运算。与此相反，DRAM，即电脑里一般使用的内存，需要经过基于 SRAM 的处理器 cache，所以读写开销需要加上这部分读写开销。）

合并结果

好了，所有的基础配方一个一个单独使用的结果已经出来了！那么，把它们全都加在一起能产生什么样的结果呢？

这张图展示了完全体 EdgeBERT 在四个数据集上的性能、运算量和内存占用。其中所有红点的实验配置参照上表（即 TABLE IV）。

在性能（准确度）相比 ALBERT 下降 1 个百分点时，EdgeBERT 能取得的内存降低和的推理速度；下降 5 个百分点时甚至能取得的推理速度。
Embedding 经过裁剪后仅保留了 40%，使得存储进 eNVM 的嵌入层参数仅 1.73MB。
QQP 的 Transformer 参数被 mask 掉 80%，MNLI、SST-2、QNLI 的 Transformer 参数被 mask 掉 60% 后，性能可以仅下降 1 个百分点。

究极配方：硬件加速器

这是什么东西？来，给你展示一下谷歌给树莓派定制的 TPU 加速器 Coral：

EdgeBERT 专属的硬件加速器想来应该也是差不多的样子。

这一部分完全不是王苏的菜… 给各位上一张 EdgeBERT 加速器的硬件结构图：

感兴趣的各位可以去参照原文进行学习_(:з」∠)_

这个加速器有什么用呢？它是基于 EdgeBERT 的运算特点量身定做的加速器，能够把 fine-tune 好的 EdgeBERT 完整地装进去进行运算。至于运算效果，它们修改了模型图中 VMAC 序列（即进行矩阵运算的单元序列）长度，与 NVIDIA 的移动端 TX2 mGPU 进行了推理时间和耗能的对比：

本文中提出的硬件加速器能够为 EdgeBERT 带来相比于 baseline 硬件加速器的能耗降低，相比于英伟达 TX2 移动端 GPU 甚至能够带来的能耗降低！耗电大户 BERT 家族终于也有能被说 “省电” 的一天了！

总结

压缩 BERT 是一项研究，但极限压缩 BERT 则是一项不易完成的工程：无论是对 Transformer 模型的全方位裁剪，还是对硬件存储介质读写性能与容错的取舍，亦或是对专属硬件加速器的设计，单独拿出一项已足够艰难，将它们合在一起就不仅可能互相冲突，甚至还可能产生逆向优化。这篇文章通过大量的实验，测试了已有的几种优化方法在边缘计算场景下的性能，比较了不同优化方法之间的差别，分析了所有优化方法进行组合后的影响及效果，并进一步提出了专属的硬件结构，实现了对目前已有的最轻量 BERT 变体的即插即用。对于需要长待机、低功耗、短延迟的场景，例如智能家居或是其他需要 NLP 技术加持的物联网设备，或许我们真的能在不远的将来看到实体的类似 EdgeBERT 加速器的解决方案出现。

虽然我们对于可能带来更大变革的模型结构依然处在探索当中，但从当下实用的角度而言，用基于 Lottery Ticket Hypothesis[6]的 BERT 优化方法寻找一个更优的类 BERT 子结构依然是一个不错的课题，至少它能让更多人、更多时候、更多场景能够用上性能强大的预训练模型。本文中提到的这些优化方法是不是也给爱思考的你带来了什么启发呢？

作者介绍

北航高等理工学院 CS 专业的市优秀毕业生，蒙特利尔大学 / MILA 准 Ph.D.，资深 ACG 宅，目前作为实习生在腾讯天衍实验室进行 NLP 研究。虽主攻 NLP，却对一切向更完善的智能迈进的系统和方向充满好奇。如果有一天 N 宝能真正理解我的文字，这个世界应该会被卖萌占领吧。（还没发过东西的）知乎 ID：Sheryc

参考链接：

[1] Sanh et al. DistilBERT, a Distilled Version of Bert: Smaller, Faster, Cheaper and Lighter. In NeurIPS'19 EMC2 Workshop. https://arxiv.org/pdf/1910.01108.pdf

[2] Jiao et al. TinyBERT: Distilling BERT for Natural Language Understanding. In Findings of EMNLP'20. https://arxiv.org/pdf/1909.10351.pdf

[3] Zafrir et al. Q8BERT: Quantized 8Bit BERT. In NeurIPS'19 EMC2 Workshop. https://arxiv.org/pdf/1910.06188.pdf

[4] Bengio et al. Estimating or Propagating Gradients Through Stochastic Neurons for Conditional Computation. arXiv Preprint. https://arxiv.org/pdf/1308.3432.pdf

[5] IEEE 754 - Wikipedia. https://zh.wikipedia.org/wiki/IEEE_754

[6] Frankle et al. The Lottery Ticket Hypothesis: Finding Sparse, Trainable Neural Networks. In ICLR'19. https://arxiv.org/pdf/1803.03635.pdf