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

文 | Sheryc_王苏
本文首发于NLP宝藏公号【夕小瑶的卖萌屋】，疯狂暗示！

这个世界上有两种极具难度的工程：第一种是把很平常的东西做到最大，例如把语言模型扩大成能够写诗写文写代码的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

Arxiv访问慢的小伙伴也可以在【夕小瑶的卖萌屋】订阅号后台回复关键词【0105】下载论文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\times H)$ 降低为
$O(V\times E+E\times 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函数

$m_z(x)=\min[\max[\frac{1}{R}(R+z-x)],1]$ 基于两个token之间的距离为注意力机制的权重计算添加了soft masking。注意力机制中的权重

$\alpha$ 变为：

$\alpha_{tr}=\frac{m_z(t-r)\exp(s_{tr})}{\sum^{t-1}{q=t-S}m_z(t-q)\exp(s{tq})}$

其中

$R$ 为控制soft程度的超参数，

$S$ 为序列截止到token

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

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

$h$ ，有

$z_t=S\sigma(\textbf{v}^T_h\textbf{x}_t+b_h)$ ，其中

$\textbf{v}_h$ 、

$b_h$ 可训练，

$\sigma(\cdot)$ 为sigmoid函数。

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

$z$ 连算都不用算了，直接

给每一个head赋一个可学习的

$z_h$ ，连输入序列都不考虑了，多出来的参数

只有12个

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

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

$z$ 值，和模型在MNLI/QQP/SST-2/QNLI四个任务上的准确度。在一大半head

几乎完全被mask掉（

$z=0$ ）之后，模型居然只在这几个任务上

掉了0.5甚至0.05的准确度！而这一方法也为模型带来了最高

$1.22\times$ 的计算量降低。

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的参数。

具体地：对于模型参数

$\textbf{W}$ ，为其赋予同样size的重要性分数

$\textbf{S}$ ，则剪枝mask

$\textbf{M}=top_v(\textbf{S})$ 。

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

$a_i$ ：

$a_i=\sum^n_{k=1}W_{i,k}M_{i,k}x_k$ 。

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

$top_v(\cdot)$ 省略近似得到损失函数

$\mathcal{L}$ 对重要性分数

$S$ 的梯度：

$\frac{\partial\mathcal{L}}{\partial S_{i,j}}=\frac{\partial\mathcal{L}}{\partial a_i}\frac{\partial a_i}{\partial S_{i,j}}=\frac{\partial \mathcal{L}}{\partial a_i}W_{i,j}x_j$

对模型参数，有：

$\frac{\partial\mathcal{L}}{\partial W_{i,j}}=\frac{\partial \mathcal{L}}{\partial a_i}M_{i,j}x_j$

将上述两个式子代换后，省略

${0,1}$ 的mask矩阵

$M_{i,j}$ 后可得：

$\frac{\partial\mathcal{L}}{\partial S_{i,j}}=\frac{\partial \mathcal{L}}{\partial W_{i,j}}W_{i,j}$

根据梯度下降，当

$\frac{\partial \mathcal{L}}{\partial S_{i,j}}<0$ 时，重要性

$S_{i,j}$ 增大，此时

$\frac{\partial \mathcal{L}}{\partial W_{i,j}}$ 与

$W_{i,j}$ 异号。这表示，只有当在反向传播时

为正的参数变得更大或为负的参数变得更小时才会得到更大的重要性分数，避免被剪枝。

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）。由此，一个数可以表示为

$V=S\times2^{E_{real}}\times F$ 。

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

$E_{real}=E+E_{bias}$ 。

我们常用的浮点数保证了

$E_{bias}$ 的选取能够使得

$E_{real}$ 在

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

$2^{127}$ 这样显然不会出现的数也能表示，这真的不是在

浪费内存？

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

$E_{bias}$ 。所谓动态体现在

每个Tensor都能得到量身定做的

$E_{bias}$ 。方法也很简单，找到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究竟能带来多大的影响？上图结果表明，仅仅是改变了嵌入层的硬件介质就能带来约

$50\times$

的推理延迟降低，以及低至约

$66000\times$

的能耗降低！这对于边缘计算场景下简直是质变了！（为何ReRAM只有读，但DRAM那边却要算DRAM读+SRAM读/写呢？因为此处的ReRAM是特殊设计的只读结构，并且可以直接读入处理器进行运算。与此相反，DRAM，即电脑里一般使用的内存，需要经过基于SRAM的处理器cache，所以读写开销需要加上这部分读写开销。）

合并结果

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

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

在性能（准确度）相比ALBERT下降1个百分点时，EdgeBERT能取得
$13.4\times$
的内存降低和
$2.4\times$
的推理速度；下降5个百分点时甚至能取得
$3.53\times$
的推理速度。
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硬件加速器

$5.2\times$

的能耗降低，相比于英伟达TX2移动端GPU甚至能够带来

$163\times$

的能耗降低！耗电大户BERT家族终于也有能被说“省电”的一天了！

总结

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

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

参考文献

[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