如何修剪BERT达到加速目的？理论与实现

作者 | RASA

编译 | NewBeeNLP

在先前讨论了各种加速模型（例如BERT） ^[1] 的方法之后，在这篇博客文章中，我们根据经验评估了修剪方法。通过本文你可以：

• 了解有关使用BERT进行 权重和神经元修剪 的实现方法；

• 看看我们可以使BERT变得 多快多小 ；

• 了解修剪神经元可以得出的结论-例如， 有时BERT根本不需要attention

Introduction

在之前的博客文章中，我们讨论了为什么需要加速大型神经语言模型，流行的模型压缩方法以及量化模型能达到多大的效果。在此博客文章中，我们将更多地讨论修剪BERT的动手经验。特别是，我们使用权重修剪和神经元修剪，即从BERT中删除单个权重连接或整个神经元。请注意，我们将此视为研究项目，因此，过程中发现的『趋势』比『绝对数值』更重要，绝对数在不同的硬件和数据集之间也可能有所不同。

我们在修剪什么？

为了更好地了解我们可以压缩的BERT组件，让我们提醒一下构成BERT体系结构的权重矩阵和矩阵乘法（在本例中为“基本”变体）。

请注意，如何仅用三个矩阵（）巧妙地实现多头自我注意力，而不是每个注意头使用三个矩阵。与self-attention模块的权重矩阵相比，和大了4倍。为了显着加速BERT，我们至少需要缩小这两个权重矩阵。在实际的实现中，我们修剪和。我们忽略，因为它的维数取决于特定任务中目标类别的数量，并且相对于其他权重矩阵而言，它通常较小。

Implementation

所有代码都可以在Rasa github库的分支 ^[2] 中找到。

权重修剪

我们的实现主要基于TensorFlow权重修剪接口 tf.contrib.model_pruning （最近在TensorFlow的 Model Optimization Toolkit ）。随附google-research/state_of_sparsity/的代码库 ^[3] 提供了一个很好的例子，可以将流行的tensor2tensor库中的权重修剪应用于transformer组件。这对于使用BERT的用户非常有用，因为BERT正是使用这些组件构建的。本质上，我们只需要用更改后的版本替换BERT中使用的所有线性层（ tf.layers.dense ），即可支持权重矩阵和偏差向量的修剪。

神经元修剪

为了实现神经元修剪，我们改编了的上述权重修剪代码。主要变化是：

• 「使修剪作用于整个神经元，而不是单个权重。」为了理解这是什么意思，我们将一个简单的线性层可视化，并尝试抑制一个神经元的激活：

该图说明了如何使神经元失活，从而mask掉bias向量相应输入以及权重矩阵中的相应列。

• 「使用重要性函数来确定要失活的神经元。」我们使用了先前博客中讨论的方法，这要求：

• 累积多个小批量的激活和相对应的训练loss梯度；

• 将激活与梯度相乘，并对每个线性层分别进行L2归一化；

• 使乘积最接近零的神经元失活。

「物理移除神经元」

要真正加快BERT的速度，仅仅按照上述方法掩盖实现中的权重矩阵元素还不够。我们实际上需要删除它们以减少矩阵乘法的数量。但是，实时调整权重矩阵的大小的做法太慢。这将需要在每个修剪步骤之后保存调整后的矩阵，使用更新后的矩阵尺寸重新构建计算图并将权重加载到图中。为此，我们提出了一个更有效的两阶段程序：

• 反复屏蔽越来越多的神经元（通过更新修剪蒙版），

• 在耗时的迭代部分之后，实际删除被掩盖的神经元，并保存缩小的权重矩阵和偏差向量以进行推断。

处理任意大小的激活

修剪并缩小尺寸的dense层会产生尺寸较小的激活。为了防止尺寸不匹配，我们将激活放大回原始大小，这样就无需调整使用激活的下一层。如图所示，放大过程是通过 tf.scatter_nd 实现的，并且会参考learned pruning mask：

即使这种膨胀增加了一些计算开销，并且也必须存储掩码，但我们稍后将证明较少的矩阵乘法的积极影响更大。

cross-pruning

你可能会认为：放大的激活包含很多零，下一层肯定不能从中受益！确实，放大只是为了简化尺寸处理。在BERT中有很多矩阵转置，逐元素张量加法和reshape。正确选择尺寸可能会成为噩梦。但是，在极少数情况下，一层的激活直接馈入下一层时，我们可以轻松地忽略膨胀。我们调整下一层的权重矩阵，删除其中的某些行。我们称为权重矩阵的交叉修剪（cross-pruning）：

在BERT中，仅在的情况下才可以进行交叉修剪，但由于是模型的2个最大权重矩阵之一，因此可以大大减少矩阵乘法的数量。

Experiments

我们的主要目标是在保持准确性的同时加快BERT的推理速度。所有的实验都是在真实的对话数据集上进行的，该数据集是根据人们与Rasa演示机器人Sara的对话建立的。我们现有的轻量级意图分类器只需几毫秒即可处理一条消息，并在数据集的测试部分上获得测试集的宏平均F1（或简称为“ F1”）得分为0.86。我们使用BERT成功达到了约0.92的F1，但是处理一条消息大约需要143ms。

在我们的实验中，我们修剪了经过微调的BERT，因此从非常好的F1开始，然后看它在修剪后如何变化。如果我们可以大大加快推理速度，并且仍远高于F1 = 0.86的基线值，那么我们可以得出结论，提速BERT是可行的方法。

权重修剪

我们知道这不会加快BERT的速度，但是可以为我们提供有关如何稀疏建立模型的提示。我们采用了最简单的方法，一次完成所有修剪。该模型不允许任何进一步的训练步，因为这可能有助于其在修剪后“恢复”。如果模型仍然运行良好，我们几乎可以肯定，随着恢复时间的增加，这种情况只会变得更好。我们尝试修剪到不同的目标稀疏度并展示F1：

结果令人信服；我们只需一步就可以移除多达50％的配重连接，而不会显着降低精度。因此，我们希望在使用神经元修剪时也可以删除至少50％的所有神经元。

神经元修剪

实际上，在这里我们实际上期望实现加速，正如之前其他人所实现的那样（请参阅此博客文章）。整个模型被修剪到所预设的稀疏性，这意味着要对每个权重矩阵进行不同程度的修剪，以反映矩阵的重要性。请注意，在通过“稀疏”修剪神经元时，我们指的是权重矩阵的神经元（列）稀疏性。

「setup」

由于这很耗时，因此修剪是使用Google Cloud Platform VM中的GPU进行的（即使如此，它也需要3个小时才能完成40个epoch）。推理速度是在MacBook Pro上使用CPU进行测量的。我们测量了实际的推理时间，即调用TensorFlow的 session.run() 的时间。

「pruning VS recovery」

自然地，神经元修剪需要对数据进行一些扫描以累积激活和梯度。我们称这些为 修剪epoch ，在此期间我们在每个修剪步骤中更新 pruning masks 。在此修剪期之后是恢复期，在此期间不再进行修剪，该模型有机会进一步调整其权重，也许已从过于激进的修剪中恢复。n那问题是，我们应该多快修剪一次模型，以及用多少时间来恢复？我们决定总共使用40个epochs，将其中的前几个epoch分配给修剪，其余的则分配给恢复。我们每个epoch执行两次修剪步骤，并且稀疏度在修剪期间线性增加，直至达到目标稀疏度。

「initial exploration」

首先，我们需要选择一个良好的修剪策略，即确定修剪模型的速度以及恢复模型的恢复时间。在40个epoch的总预算和60％的目标稀疏性的情况下，我们尝试了不同的修剪时间长度（如下图，从修剪到恢复的过渡由虚线标记）： 修剪策略很重要，但是过于急切（仅1个epoch）修剪会对模型造成太大损害。更有趣的是，修剪太慢也是有害的！我们假设，随着修剪的进行，剩余的神经元会以可能降低其权重质量的方式对其输入的变化做出反应。通过观察到非常慢的修剪（18或25个epoch），模型恢复需要更长的时间。我们决定在所有其他实验中使用10个修剪epoch+ 30个恢复epoch的设置。

「我们可以修剪多大程度」

看到60％稀疏的BERT仍可以达到接近0.90的F1，我们想知道是否可以继续进行修剪。我们尝试修剪到不同的稀疏度，并报告F1以及推理速度：

显然，稀疏性和加速度是正相关的。60％稀疏模型仍然可以达到F1〜0.895（相对降低2.8％），并且快28％！不幸的是，稀疏度超过60％会对F1造成太大伤害。请注意，即使是90％稀疏的BERT，它也比我们的轻量级分类器慢大约20倍，显然，BERT太大了。为什么50％稀疏模型几乎没有加速？这是由于放大激活导致的计算开销，我们测得约为15ms。还要注意的是，在GPU（矩阵乘法非常快但 tf.sactter_nd 尚未优化）上，除非你将BERT修剪到极高的目标稀疏度（例如80％），否则此开销超过了速度的提高。

「模型大小呢」

自然地，神经元修剪也会使BERT变小。原始的完整BERT约为406MB，而60％稀疏版本只有197MB。不幸的是，嵌入层很大，甚至90％稀疏的BERT仍约有115MB。即便如此，这表明如果只想缩小大型模型而几乎不损失准确性，神经元修剪可以为你提供帮助。

神经元修剪分析

我们让模型决定哪些组件需要修剪得多而哪些需要修剪的少。那么，这会导致什么呢？60%稀疏的BERT的component-specific和layer-specific稀疏性如下所示：

显然，特定于组件的稀疏度在不同层之间变化不大，但是通常，最后2-3层似乎不太重要，因为它们会被修剪得更多。有趣的是，可以非常多地修剪和，这告诉我们可以使用更少的参数来计算注意力权重。attention输出（）和layer输出（）似乎非常重要，因为它们只被修剪了一点。

「我们需要attention吗」

看到和被严重地修剪了，我们尝试了一个简单的事情：修剪这两个到100％的稀疏度，并保留所有其他参数不被修剪。所得模型达到F1 = 0.897。在这种情况下，self-attention组件仅通过对所有时间步长上的值向量求平均的方式来生成上下文向量，而不实用注意力。结果告诉我们，在这个特定的任务和数据集上，BERT不需要self-attention，这让我们感到惊奇。这个结果很好，但是如果我们可以缩小更大的，则BERT可以得到更大的加速。因此，我们尝试了类似的实验，在该实验中，我们将该权重矩阵修剪为100％的稀疏度。我们将其与消除和进行比较：

显然，中间层不能被这么多地修剪---它可能比self-attention更重要。一个有趣的现象是两种情况下修剪过程中的F1急剧下降（对于self-attention，发生在稀疏度约为90％；而对于的稀疏度约为50％）。我们假设每个组件都可以具有其特有的最大稀疏性，超过该稀疏性将导致模型“shock”。因此，在设置某个组件的目标稀疏度时，最好查看一下在模型级别设置了目标稀疏度的 “organic” pruning 中达到的稀疏度。只需简单聆听你的模型即可 :-)

「神经元修剪可以优化吗」

我们看到修剪epoch的数量很重要。恢复情况如何？那么除了60％以外的稀疏率呢？我们显示了一系列稀疏性的F1（虚线显示了恢复时期的开始）：

恢复时期显然很重要，这使我们从非常低的F1值中恢复过来。但是，至少对于已使用的修剪策略，修剪的模型越多，恢复的速度就越慢。稀疏度超过60％时，无法完全恢复，稀疏度超过80％时，它似乎完全失去恢复能力。尽管如此，即使是70％的稀疏模型也可以通过花费更多时间或使用其他修剪策略来恢复。

「接下来呢」

即使使用神经元修剪可以成功加快BERT的速度，我们还是决定不继续执行此想法。我们认为，对于对话NLU所需的分类而言，BERT太大（也许也太强大）。与我们现有的分类器相比，即使100％稀疏版本仍然非常慢。进一步研究的一个有趣方向仍然是知识蒸馏。最近，Hugging Face团队发布了DistilBERT ^[4] , 设法将BERT分解为自身的2倍小版本（仍然是一个大模型，具有66M参数！），实现了60％的加速。在论文『Distilling Task-Specific Knowledge from BERT into Simple Neural Networks ^[5] 』中，作者将BERT的知识分散到参数小于1M的小型Bi-LSTM中，但以准确性稍差为代价。

Conclusion

我们已经成功调整了BERT的权重修剪和神经元修剪功能。我们在sara对话数据集的扩展版本上观察到，多达60％的神经元可以删去，同时保留了测试集宏平均F1，这也产生了28％的相对推理加速和51％的模型尺寸减小。不幸的是，它仍然使BERT比我们现有的分类器慢得多。我们认为，对于我们的任务和数据集，BERT可能只是不必要地又大又慢---即使没有attention weights，也可以通过获得良好的结果来证实这一直觉。

我们基于BERT对Rasa意图分类器的原始改编 ^[6] ，开源了我们的实现 ^[7] 。我们的代码支持权重和神经元修剪，强烈建议你尝试一下。我们也很乐意在社区论坛中 ^[8] 与你讨论有关模型压缩的任何内容。

本文翻译自RASA博客，可以点击阅读原文直达。

https://blog.rasa.com/compressing-bert-for-faster-prediction-2/

本文参考资料