【模型压缩系列】一：模型替换
bert问世以来，相关研究和改进如雨后春笋般涌现，预训练模型得到了飞速的发展，并在NLP很多任务中取得了统治地位。工业界也对此产生了浓厚的兴趣，各大公司纷纷试图在搜索、推荐、分类等领域中引入bert。然而，效果较好的深层bert模型复杂度较高，难以在具有较多低耗时场景的工业界直接使用，能不能使小模型有接近large级bert系模型的表现，同时又可以满足时间需求呢？模型压缩会给你答案。
一、背景
贝壳找房是一家科技驱动的品质居住服务平台，为了更好地服务消费者，我们今年推出了小贝助手2.0产品，包括咨询助手、培训助手、带看助手、房客源助手等来帮经纪人赋能提效。文本理解能力是这些助手所核心依赖的能力，但是随着业务场景越来越复杂，对文本理解的能力和实时性都提出了很高的要求，因此我们结合贝壳的语料特点做了预训练，为了保证线上的实时推理需求进一步做了模型压缩的工作。
二、技术介绍
一个复杂的深层的神经网络模型，参数量百万级别甚至更大，推理时间可想而知。直觉上，要想降低推理时间可以怎么做？
减少结点个数 降低权值精度
基于上述直观想法，有了两个研究方向：剪枝和量化。剪枝相对比较复杂，即使在学术界也研究较少，而量化在效果上往往不如人意。
根据各种计算去剪枝确实困难，可如果以层为单位直接舍弃，能极大地减少模型的计算量，这就是模型替换的出发点。
此外，在分类模型中，无论模型有多少个中间层，多少个参数，直接决定分类结果的是最后全连接层的输出结果(soft targets)，如果我们去训练一个小模型，使得小模型的logits逼近大模型的logits，那么我们的小模型同样会有不错的效果。这就是模型蒸馏的思想，由于用到了大模型的logits(知识)，因此又叫知识蒸馏。
在知识蒸馏过程中，我们通常称大模型为teacher模型，小模型为student模型，teacher可以在自己知识储备最丰富(训练完成)时指导student(蒸馏)，也可以在学习的过程中与student进行交互，共同学习，我们把后者称为协同训练。
综上所述，我们以直观的思路引出3种模型压缩方法:
模型替换 模型蒸馏 协同训练
我们会分别对其进行原理介绍，实验分享与总结。本篇主要介绍模型替换。
三、模型压缩之模型替换
如上文所述，模型替换算是剪枝的一种，只不过剪得比较彻底，直接剪掉了若干层。可见，小模型和大模型不仅都要是transformer-based模型，而且小模型的hidden_size和intermediate_size与大模型对应部分也必须是一致的。
同时，模型替换又很像蒸馏。其初始化需要用到大模型的知识，又通过隐层替换的形式，不断利用大模型的信息来训练替换后的隐层参数。正如其名字的含义，最终每个隐层都不同于原来的隐层，但又拥有接近原模型的表现。下面让我们来深入了解其原理吧。
3.1 原理
模型替换的思想来自2020年EMNLP会议的论文[1]，在模型压缩方向，各种知识蒸馏方法相互角逐，我们能看到这样一篇新颖的论文实属不易。
3.1.1 步骤
Stage1:
根据压缩比确定prd和scc，如上图所示，压缩比为2:1，大模型的1-2,3-4,5-6层分别对应小模型的1,2,3层，分别用prd1,prd2,prd3和scc1,scc2,scc3表示。prd(predecessor)代表可能被替换的模块，scc(successor)代表替换模块 假设小模型是n层，用大模型的前n层初始化小模型，feed forward过程中，每层都以概率p决定是否用scc层替换prd层，pred层的参数是固定不变的，通过下游任务学习scc层的参数
Stage2:
如上图b所示，在同一个下游任务中，直接在所有的scc层finetune
上述两个阶段执行完之后，就得到了模型压缩后的小模型。
3.1.2 描述
不同于许多蒸馏方法，需要在预训练阶段就开始实施，更不需要设计loss，整个过程都在下游任务中直接finetune，突出一个简洁。我们需要提前在下游任务中训练好一个大模型，并且设计好替换概率，接下来就可以执行上述两阶段操作了。
由于第一阶段每个scc层都是以一定概率替换prd模块，所有的scc层并没有直接全部在下游任务中训练，因此就有了第二阶段的scc层微调，后面实验表明，第二步是很有必要的，能够提升小模型的表现。
具体到替换部分，每个prd, scc相对应的模块都通过一个伯努利分布采样一个随机变量，假设当前是第个模块，对应的，随机变量，也就是替换概率，，其中采样的概率为，有的概率为1,的概率为0。
假设scc模块输出为，prd模块输出为，那么本模块的最终输出取决于上述随机变量的取值。
3.2 实验
我们分别在小贝的技能分类任务和智能话务意图分类任务上做实验。
小贝数据取自线上im对话，其中用户意图有几百类，为了提高识别的准确率，我们在意图上层定义了13个技能，每个技能对应一部分意图，把一个复杂的分类任务分解成了两步，每步都是相对简单的分类任务，下面的实验是对技能层分类的实践。
智能话务数据来自招聘场景中的语音对话，经asr转换成文本，我们定义了77个意图类别，不同于小贝数据，我们直接对意图做预测。
两个任务一个类别较少，文本较规整，另一个类别较多，文本存在噪声，我们通过实验验证模型替换在两种特性数据下的效果。所有实验结果都是运行5次取的最大值。
实验1
我们将压缩前后的大小模型分别称为 teacher 和 student,实验所用的大模型均由加入了公司内部语料的公开中文预训练模型继续训练而来。其参数如下所示。
teacher: layer=12, hidden_size=768, intermediate_size=3072, attention_heads=12
student: layer=4, hidden_size=768, intermediate_size=3072, attention_heads=12
小贝技能13分类结果
| 模型 | precision | recall | f1-score | support |
|---|---|---|---|---|
| teacher | 0.8815 | 0.8813 | 0.8811 | 10582 |
| 前4层finetune | 0.8627 | 0.8626 | 0.8622 | 10582 |
| student | 0.8676 | 0.8675 | 0.8669 | 10582 |
智能话务意图77分类
| 模型 | precision | recall | f1-score | support |
|---|---|---|---|---|
| teacher | 0.8421 | 0.8388 | 0.8388 | 4703 |
| 前4层finetune | 0.8323 | 0.8297 | 0.8290 | 4703 |
| student | 0.8401 | 0.8384 | 0.8376 | 4703 |
分析
实验中，前4层finetune的参数值同student模型一样，都是由teacher模型的前四层初始化而来。从结果来看，直接finetune比模型替换后的结果要差一些的。两个任务中，student结果精度损失不同，技能13分类损失达到1.4%，而意图77分类精度损失只有0.2%，说明类别较少，文本较规整的技能分类任务对模型复杂度敏感，而对于77分类任务，虽然更复杂的模型能够提升效果，但受任务本身难度制约，提升比较有限，使得teacher模型，前4层finetune以及student结果都较为接近。这说明在更难的任务上使用复杂度更高的模型以及模型压缩收益更高。
问题
模型替换有其固有的问题，其仅缩小了模型层数，每层的结点个数并没有减少，响应时间能满足要求吗？我们与线上正在使用的4层albert做对比。结果如下。
albert: layer=4, hidden_size=312, intermediate_size=1248, attention_heads=12
bert-of-theseus: layer=4, hidden_size=768, intermediate_size=3072, attention_heads=12
| 模型 | 响应时间 |
|---|---|
| albert | 28ms |
| bert-of-theseus | 88ms |
其中，响应时间表示测试的1万条数据中99%的请求在该时间内响应完成。
结果表明，仅减少层数并不足以使模型压缩到响应时间满足需求的程度，如果我们将teacher模型的中间结点也缩小呢？如果压缩后的小模型效果接近上面的表现，我们可以着手训练层数高，中间结点少的“瘦高”模型。效果究竟怎样呢？为此，我们做了实验二。
实验二
本实验为模型替换在“瘦高”模型上的探索，所谓“瘦高”模型，就是隐层和中间层结点较少，而层数比较高。
“瘦高”模型参数如下:
teacher: layer=12, hidden_size=256, intermediate_size=1024, attention_heads=4
对应的，压缩后的小模型参数:
student: layer=4, hidden_size=256, intermediate_size=1024, attention_heads=4
小贝技能13分类结果
| 模型 | precision | recall | f1-score | support |
|---|---|---|---|---|
| teacher | 0.8771(0.8815) | 0.8769(0.8813) | 0.8766(0.8811) | 10582 |
| 前4层finetune | 0.8584(0.8627) | 0.8569(0.8626) | 0.8572(0.8622) | 10582 |
| student | 0.8602(0.8676) | 0.8594(0.8675) | 0.8590(0.8669) | 10582 |
智能话务意图77分类
| 模型 | precision | recall | f1-score | support |
|---|---|---|---|---|
| teacher | 0.8356(0.8421) | 0.8316(0.8388) | 0.8310(0.8388) | 4703 |
| 前4层finetune | 0.8247(0.8323) | 0.8214(0.8297) | 0.8209(0.8290) | 4703 |
| student | 0.8272(0.8401) | 0.8229(0.8384) | 0.8231(0.8376) | 4703 |
其中，括号内为标准bert-base参数下的系列实验结果。
分析
最为直观的结果就是，“瘦高”模型的系列实验中，所有结果全面下降，我们知道技能分类任务上，student和teacher的gap很大，在本实验中，其差距并没有被缩小。在意图分类上，甚至标准bert-base压缩后的小模型表现都好于“瘦高”teacher模型结果，因为这个任务更难，更需要模型有更高的复杂度。看来中间结点的减少相对于层数的缩小代价更大，两个任务都表现出某种程度的欠拟合，尤其是层数又少，中间结点又少的student模型，比实验一的结果下降了很多。因此，“瘦高”模型不会成为训练的重点，模型替换并不能实际解决我们的问题。
四、进阶
从上面的实验可以看出，我们每次都会和小模型直接finetune做对比，这是因为在调研方法的时候看到网上对模型替换是有争议的，尤其在中文数据集上。在一些实验中，模型替换结果与小模型直接finetune结果相差无几。在我们的实验中，从概率论的角度，替换后的结果是显著好于直接finetune的，不过提升均在1%以内。如果你在实验中，模型替换未能好于直接finetune，不妨做如下尝试。
降低stage2学习率 引入数据增强
stage2，即替换后的小模型在下游任务finetune阶段，是进一步训练模型的有效手段，然而实验初期，我们发现经过stage2后，模型效果下降，从效果变化来看像是步长过大导致，降低学习率后，效果回归正常，表现稳定的优于stage1结果。
数据增强，在这里作为引导模型学习的手段，增强方式为引入大模型预测了伪标签的数据。这部分数据直接加到小模型finetune上效果提升并不显著，但在模型替换中效果明显。我认为，在漫长的finetune过程中这部分数据可能难以充分发挥效用；而在模型替换stage1后，模型已经有了不错的预测性，此时增强的数据可能能够更好地指导参数向大模型学习。我们实验初期应用数据增强能够有效提升表现，不过上述实验为了保持finetune与模型替换的一致性，两个实验均未使用数据增强。
此外，我们尝试了修改替换概率，效果变化并不明显。尝试了在替换过程中引入映射矩阵，使替换后的小模型隐层可以映射到任意大小，效果表现不佳。
五、总结
综上可得，模型替换有如下特点:
优势：
直接在下游任务上压缩，方法非常简洁，易于操作 结果相对稳定且效果不错
劣势：
模型中间层size不可变，比较依赖大模型的结构 不能蒸到cnn等异构模型
至于模型替换相较于其他模型压缩方法，效果怎样呢？我们没有直接对比过，因为二者的场景不同，研究正常large的模型替换是没有意义的(从响应时间考虑)，中间结点不能减少使得其需要“瘦高”模型的支持。而从本篇的实验也可以看到，作为teacher的大模型需要尽可能“高胖”。
本篇介绍了我们在实际业务中，对模型替换的探索，除此之外，在蒸馏和协同训练的实践中，我们也积累了一些经验，取得了一些进展，你是否对此感到好奇呢？后面我们没有继续在bert-base上研究，teacher模型统一使用了24层的large模型，敬请期待系列文章的后续吧。
六、参考文献
[1] Xu C, Zhou W, Ge T, et al. BERT-of-Theseus: Compressing BERT by Progressive Module Replacing[J]. arXiv, 2020: arXiv: 2002.02925.
[2] https://zhuanlan.zhihu.com/p/112787764.
[3] https://kexue.fm/archives/7575
[4] https://github.com/ymcui/Chinese-BERT-wwm
作者介绍
杨蕴凯，2019年6月毕业于武汉大学计算机学院，毕业后加入贝壳找房人工智能中心业务智能部，主要从事自然语言理解相关工作。