只需一个损失函数、一个超参数即可压缩BERT，MSRA提出模型压缩新方法

论文链接：https://arxiv.org/pdf/2002.02925.pdf

这篇论文提出了一种新型模型压缩方法，能够通过逐步模块替换（progressive module replacing）有效地压缩 BERT。该方法首先将原版 BERT 分割成多个模块，并构建更加紧凑的替代模块；然后，用替代模块随机替换原始模块，训练替代模块来模仿原始模块的行为。在训练过程中，研究者逐步增加模块的替换概率，从而实现原始模型与紧凑模型之间的更深层次交互，使得训练过程流畅进行。

与之前用于 BERT 压缩的知识蒸馏方法相比，该方法仅利用一个损失函数和一个超参数，将开发者从调参这一繁琐过程中解放出来。该方法在 GLUE 基准上的性能优于现有的知识蒸馏方法，为模型压缩开启了新方向。

模型压缩方法知多少

随着深度学习的流行，很多巨大的神经模型诞生，并在多个领域中取得当前最优性能。尤其是在自然语言处理（NLP）领域中，预训练和调参已经成为其中大多数任务的新规范。基于 Transformer 的预训练模型在自然语言理解（NLU）和自然语言生成（NLG）领域中成为主流。这些模型从其「过参数化」属性中获益，它们往往包含数百万甚至数十亿个参数，这就使得此类模型计算成本高昂，且从内存消耗和高延迟的角度来看其计算是低效的。这一缺陷极大地阻碍了此类模型在生产环境中的应用。

为了解决该问题，研究者提出很多神经网络压缩技术。一般而言，这些技术可分为三类：量化、权重剪枝和知识蒸馏（KD）。其中，KD 能够压缩预训练语言模型，并因此得到了极大关注。KD 利用大型教师模型「教」紧凑的学生模型模仿教师的行为，从而将教师模型中嵌入的知识迁移到较小的模型中。但是，学生模型的性能状况取决于设计良好的蒸馏损失函数，正是这个函数使得学生模型模仿教师行为。近期关于 KD 的研究甚至利用更复杂的模型特定蒸馏损失函数，以实现更好的性能。

模型压缩新思路——Theseus 压缩

与显式地利用蒸馏损失函数来最小化教师模型与学生模型距离的 KD 不同，该研究提出一种新型模型压缩方法。研究者受到著名哲学思想实验「忒修斯之船」的启发（如果船上的木头逐渐被替换，直到所有的木头都不是原来的木头，那这艘船还是原来的那艘船吗？），提出了 Theseus Compression for BERT (BERT-of-Theseus)，该方法逐步将 BERT 的原始模块替换成参数更少的替代模块。研究者将原始模型叫做「前辈」（predecessor），将压缩后的模型叫做「接替者」（successor），分别对应 KD 中的教师和学生。

该方法的工作流程如下图 1 所示：首先为每个前辈模块（即前辈模型中的模块）指定一个替代（接替者）模块；然后在训练阶段中以一定概率用替代模块随机替换对应的前辈模块，并按照新旧模块组合的方式继续训练；模型收敛后，将所有接替者模块组合成接替者模型，进而执行推断。这样，就可以将大型前辈模型压缩成紧凑的接替者模型了。

图 1：BERT-of-Theseus 的工作流程。

Theseus 压缩与 KD 的思路有些类似，都鼓励压缩模型模仿原始模型的行为，但 Theseus 压缩有很多独特的优势。

首先，Theseus 压缩在压缩过程中仅使用任务特定的损失函数。而基于 KD 的方法除了使用任务特定的损失函数外，还把一或多个蒸馏损失函数作为优化目标。Theseus 方法在整个压缩过程中仅使用一个损失函数，从而联结不同的阶段，使压缩以完全端到端的形式进行。此外，为不同任务和数据集选择多个损失函数并平衡每个损失函数的权重，通常是一件耗时费力的事。

其次，与近期研究 [15] 不同，Theseus 压缩不使用 Transformer 特定特征进行压缩，这就为压缩广泛模型提供了可能性。

第三，与 KD 仅使用原始模型执行推断不同，该方法允许前辈模型与压缩后的接替者模型共同训练，从而实现更深层次的梯度级交互，并简化训练过程。此外，混合了前辈模块和接替者模块的不同模块组合添加了额外的正则化项（类似于 Dropout）。该方法还基于课程学习（Curriculum Learning）方法来驱动模块替换，将模块替换概率从低到高逐渐增加，从而实现优异的 BERT 压缩性能。

该研究的贡献如下：

• 提出新方法 Theseus 压缩。该方法仅使用一个损失函数和一个超参数，为模型压缩提供了新方向。

• 利用该方法压缩得到的 BERT 模型运算速度是之前的 1.94 倍，并且保留了原始模型超过 98% 的性能，优于其它基于 KD 的压缩基线。

BERT-of-Theseus

接下来，我们来看模块替换和课程学习方法。BERT-of-Theseus 方法的工作流程参见图 1。

在该示例中，研究者将一个 6 层的前辈模型 P = {prd_1, .., prd_3} 压缩成 3 层的接替者模型 S = {scc_1, .., scc_3}。prd_i 和 scc_i 分别包含两个层和一个层。(a) 在模块替换训练阶段，以概率 p 将每个前辈模块 prd_i 替换成对应的接替者模块 scc_i。(b) 在接替者微调和推断阶段，把所有接替者模块 scc_1..3 组合到一起，以执行计算。

图 2：恒定模块替换率和替换调度器的替换曲线。该图用不同灰度表示 Theseus 压缩的两个阶段：1）模块压缩，2）接替者微调。

实验

现在，我们来看 Theseus Compression for BERT 的实验结果。研究者对比了 BERT-of-Theseus 和其它压缩方法，并通过进一步实验来分析结果。

基线

如表 1 所示，研究者对比了新方法与现有方法的层数、参数量、损失函数、外部数据使用和模型无关性。

表 1：不同 BERT 压缩方法的对比。「CE」和「MSE」分别表示交叉熵和均方差，「KD」表示知识蒸馏的损失函数，「CETASK」和「CEMLM」分别表示在下游任务上和在遮蔽语言建模预训练任务中计算得到的交叉熵。其它损失函数参见相关论文。

实验结果

表 2 展示了模型在 GLUE 开发集上的实验结果。研究者将其预测结果提交至 GLUE 测试服务器，得到了官方排行榜结果，详见表 3。

表 2：在 GLUE 开发集上的实验结果。每个数据集下的数字表示数据集中的训练样本数量。

表 3：在 GLUE 服务器测试集上的实验结果。每个数据集下的数字表示数据集中的训练样本数量。

通用模型

作者也提供了一个在 MNLI 上压缩的 6 层继承者模型权重，可供直接微调使用，在 GLUE 上也取得了超过 DistillBERT 的表现。借助 transformers 库可以轻松使用三行代码加载这个模型权重：

from transformers import AutoTokenizer, AutoModel tokenizer = AutoTokenizer.frompretrained("canwenxu/BERT-of-Theseus-MNLI") model = AutoModel.frompretrained("canwenxu/BERT-of-Theseus-MNLI")

模型压缩后，研究者在其它句子分类任务上微调接替者模型，并与 DistillBERT 的性能进行对比（参见表 4）。这一通用模型在 MRPC 上实现了同等性能，且在其它句子级任务中的性能显著优于 DistillBERT。

表 4：该研究训练出的通用模型在 GLUE-dev 上的实验结果。