EMNLP2020 | 模型压缩系列：BERT-of-Theseus(一种基于模块替换的模型压缩方法)
当古希腊神话遇到BERT，于是有了BERT-of-Theseus
背景
论文标题：
BERT-of-Theseus: Compressing BERT by Progressive Module Replacing
论文作者：
Canwen Xu, Wangchunshu Zhou, Tao Ge, Furu Wei, Ming Zhou
机构：
武汉大学、北京航空航天大学、微软亚洲研究院
论文地址：
https://arxiv.org/abs/2002.02925
收录会议：EMNLP 2020
论文代码：
https://github.com/JetRunner/BERT-of-Theseus
摘要
文章提出一种新型模型压缩方法(Theseus)，该方法通过逐步更换模块从而实现BERT模型压缩。该方法先将原始的BERT分为数个模块并为其创建对应的压缩模块(即替代模块)，再随机地用压缩模块替换掉原始的模块，并训练这些压缩模块使其能够效仿原始的模块。在训练过程中逐步增加替换的概率，从而实现整体模块(即整个模型)替换。
举个栗子，有7个菜逼想要学习武当七侠的真武七截阵，先派一个菜逼，替换掉武当七侠中的一位，比如莫声谷，然后让这6神+1菜的组合不断的练习。经过一段时间，在6位大神的指(毒)导(打)下，新加入的菜逼实力会提升，这个6神+1菜的组合拥有接近武当七侠真武七截阵的真正实力。再陆续派出其他菜逼，重复上述过程，直到所有菜逼都接受了都被充(社)分(会)训(毒)练(打)，那么最终7个菜逼也能施展真武七截阵。所以，名师出高徒，组团抢大龙。
基于Theseus这种方法在原始模块和压缩模块之间带来了更深层次的交互，并平滑了训练过程。基于Theseus的方法在GLUE上优于知识蒸馏，且相比于BERT的知识蒸馏方案，基于theseus的压缩方法只需要一个损失函数和一个超参数，这真真是炼丹师和调参侠们的福音。
模型
Theseus之船
既然文章的模型名为 BERT-of-Theseus ，那先说说Theseus(忒修斯)之船的典故。忒修斯(Theseus)是传说中的雅典国王。克里特岛国王米诺斯在战争中打败过雅典。他要求雅典人每九年（亦传每年）奉祭七名少年七名少女给怪物米诺陶洛斯(米诺斯的儿子，半人半牛,被囚禁于米诺斯迷宫)。轮到第三次奉祭时，忒修斯自告奋勇要去杀死那个怪物。他和父亲埃勾斯(海神波塞冬的儿子)约定，当他的船回航时若成功则挂白帆。在克里特，米诺斯的女儿阿里阿德涅爱上了忒修斯，她给了他一个线团，以便他在迷宫中标记退路。忒修斯杀死了米诺陶并带领其他雅典人逃离了迷宫，他还带走了阿里阿德涅，但在回航路上将她丢在在纳克索斯岛，所幸酒神狄俄倪索斯怜悯她，娶她为妻．阿里阿德涅诅咒了忒修斯，他忘记挂上白帆，埃勾斯以为忒修斯失败身亡，悲痛之下堕下海岸而死，从此这片海域就被称为爱琴海。另一种说法是忒修斯被狄俄倪索斯胁迫，将阿里阿德涅遗弃，沮丧之余忘记了挂白帆。人们为了纪念他的英雄壮举，忒修斯回航所乘的船据说在雅典保留了很多年，随着时光流逝，那艘船逐渐破旧，木头朽毁后不断修补替换，人们依次更换了船上的甲板，以至最后更换了它的每一个构件。这时人们禁不住发出疑问：更换了全部构件的忒修斯之船还是原来那艘船吗？后来人们常把部分替换后主体是否仍然保留的哲学命题称为忒修斯之船后来，常把其所有部分被替换后原主体是否仍然存留的哲学问题称为“忒修斯之船”，其挑战性在于：如何理解和刻画跨越时间或空间的个体的同一性？文归正传，继续介绍BERT-of-Theseus的其他细节。
模块替换
模型压缩致力于在降低模型尺度和计算量的同时最大程度地保持原有精度。常见的模型压缩方法有剪枝、量化和知识蒸馏。
剪枝是用二分类器剪除模型中多余的权重，从而使其变为参数较少的小模型；量化不改变模型结构，但将模型换一种数值格式，比如通常训练模型用的是 float32 类型，而如果换成半精度float16 就不仅能提速还能节约显存，有时会进一步转换成 8 位整数甚至 2 位整数（二值化）；知识蒸馏则是训练student模型来模仿teacher模型，除了学习输出之外，有的还是设计模型的中间层结果、Attention 矩阵等部分的学习，这就涉及到多项损失函数的设计。BERT-of-Theseus模型压缩的核心思想：循序渐进地用更少参数的模块替换BERT中的各个原始模块。原始模块和用以替换的压缩模块分别称为 predecessor(前辈) 和 successor(传人)， successor要学习效仿predecessor，这类似于知识蒸馏中的teacher和student。但与知识蒸馏所不同的是，知识蒸馏一方面需要确切地定义teacher和student之间的相似损失函数，且最终性能严重依赖于所设计的损失函数。此外，该损失函数还需要联合特定任务的损失。Theseus压缩只面向下游任务阶段，而不涉及预训练阶段，所以只需要特定下游任务的损失函数。换句话说，Theseus压缩是面向微调阶段的模型压缩方案。受Dropout启发，文章中提出一种模块替换策略：
(1)将原始模型划分 predecessor，并为每个predecessor模块指定一个successor模块。比如，对于BERT压缩，为每2层 Transformer 作为一个 predecessor，为其指定1层的Transformer作为 successor。如此，则12层的BERT划分出6个predecessor模块，压缩后的BERT则由6个successor模块组成。
(2)第 i+1 个模块，通过一个伯努利分布，采样一个随机变量 r i + 1 r_{i+1} ri+1，则该变量值为1的概率为 p p p，值为0的概率为 1 − p 1-p 1−p。
(3)第 i+1 个模块的输出：
y i + 1 = r i + 1 ∗ s c c i ( y i ) + ( 1 − r i + 1 ) ∗ p r d i ( y i ) y_{i+1}=r_{i+1} * s c c_{i}\left(\mathbf{y}_{i}\right)+\left(1-r_{i+1}\right) * p r d_{i}\left(\mathbf{y}_{i}\right) yi+1=ri+1∗scci(yi)+(1−ri+1)∗prdi(yi)
通过这种方式，predecessor 模块和 successor 模块在训练过程协同工作。由于上述混合模型中引入了一个伯努利分布，所以在模型中引入了额外的噪声可以作为训练 successor 中的正则化操作(类似Dropout)。
(4)这个训练过程跟微调阶段很类似，都是面向下游特定任务。另外，这里需要关注一个训练细节：在梯度反向传播的时候， predecessor 的权重参数都freeze不参与梯度计算，只有 successor 模块的权重参数会参与梯度更新。
BERT-of-Theseus的具体架构如Figure 1所示：
successor 的微调和推理
在上述训练阶段后，将所有 successor 进行联合训练，即所谓的替换后的微调阶段。这是因为上述的模块替换训练阶段，各个 successor 模块都是孤立地训练，所以需要将所有 successor 模块整合到一起进行微调。
替换概率策略(curriculum策略)
对于整个模型压缩方法来说，撇去原始NLP任务中的超参数外，只有替换概率p(即替换率)这一超参数需要关注。一般来说，使用一个常量0.5就可以达到比较不错的效果。但是文章还提出了一个线性策略来为每个训练step设置不同的替换概率，即文章中提到的curriculum学习策略，其具体公式如下：
p d = min ( 1 , θ ( t ) ) = min ( 1 , k t + b ) p_{d}=\min (1, \theta(t))=\min (1, k t+b) pd=min(1,θ(t))=min(1,kt+b)
其中，t表示当前训练的step，k表示每个step的增长率，b表示初始的替换概率。通过这种策略，统一了此前谈及的两个训练阶段成为一个端到端的从易到难的学习过程。在初始阶段，暴露较多的 predecessor，让这些predecessor参与训练，能够提升整个模型的质量，得到更小的loss，这能够平滑整个训练过程平滑而不会过于震荡；而在训练后期，由于模型整体学习都比较好了，模型可以逐渐摆脱对 predecessor 的依赖，并让更多的successor参与后续的训练，进而使整个模型能够平稳过度到 successor 微调阶段。此外，在替换训练的初试阶段，当 θ ( t ) < 1 \theta(t)<1 θ(t)<1，训练过程平均被替换的模块数量的期望值是 n* p_d，所有n个 successor 的期望学习率：
l r ′ = ( n p d / n ) l r = ( k t + b ) l r l r^{\prime}=\left(n p_{d} / n\right) l r=(k t+b) l r lr′=(npd/n)lr=(kt+b)lr
lr表示整个训练过程设置的固定学习率，lr’表示所有successor模块的等效学习率。因此，使用这种替换策略相当于在训练阶段间接得给学习率应用了一个warm-up策略，间接得提升了模型训练的效率。
实验结果
实验过程使用得数据集是 GLUE。由于基于 Theseus 的模型压缩方案并不是针对预训练阶段而是面向特定任务(即常说的微调阶段)，所以只需要用到GLUE中的train set即可进行模型压缩。相比于 DistillBERT 的模型压缩需要耗费720个小时的GPU，基于 Theseus方案凭借不到20个小时的GPU，喜提低碳环保节能小先锋称号。不同模型压缩方法的直观对比如 Table 1所示：
对比的方面有：层数、参数量、损失函数、使用外部数据的使用情况等。需要注意的是，DistillBERT是预训练阶段蒸馏，所以不能直接比较。另外，由于TinyBERT的尺寸设置不同、不仅使用2阶段蒸馏还使用了数据扩增，所以在后续试验阶段的对比并不考虑TinyBERT。如果从损失函数、是否使用外部数据和是否局限特定模型来看，基于 Theseus 的方案无疑是最佳的。
各个模型在GLUE dev、test数据集上的评测结果如 Table 2和 Table 3所示：
可以看出，BERT-of-Theseus能够在GLUE的dev set 和 test set上分别能够保持BERT-base的98:4%、98:3%性能。其中 Fine-tuning 模型是用6-layer(层数与BERT-of-Theseus一样) 的BERT在GLUE上fine-tune的结果。在GLUE的每个任务上，BERT-of-Theseus都优于Fine-tuning 模型，由于使用相同损失函数，所以这充分表明BERT-of-Theseus中所用到的方法确实可以将知识从 predecessor 有效地转移到 successor。
此外，有学者发现在 MNLI 数据集上微调过的模型能够成功迁移到其他句子分类任务。因此，文章也将基于 MNLI 压缩后的模型作为一个通用的压缩版BERT，以用于下游任务。再用其他句子分类任务来微调压缩后的模型，该模型与 DistillBERT 在GLUE dev数据集上的对比如 Table 4所示：
可以看出，通用型的BERT-of-Theseus在MRPC上与DistillBERT的得分相同，但在其他句子级任务上碾压式地将DistillBERT摁在地上摩擦。
分析
替换模块的影响
由于不同的 Transformer 层在特征抽取上所起的作用不同，所以文章还进一步探讨了不同的替换模块所带来的影响，实验结果如 Table 5所示：
Table 5：在GLUE dev set实验替换不同模块所带来的影响
在QNLI、MNLI和QQP这3个数据集上各个模块替换后平均性能的下降情况如 Figure 3所示
可以看出，替换后两个模块对整体性能的影响很小，而替换第一个模块会严重损害性能。这是因为语言特征的抽取主要通过前几层。因此，被降低的表征能力成为后续层的瓶颈。相反，高质量的低层级特征有助于后续的网络层，因此缩小的模块大小对最终结果的影响有限。
替换率的影响
不同替换率在MRPC 和 RTE数据集的影响如Figure 4所示：
可以看出，不同的替换率对等效学习率的影响并不是绩效差异的主要因素。一般来说，BERT-of-Theseus对不同的替换率不是很敏感。在0.5和0.7之间的替换率可以使所有GLUE任务的性能令人满意。但是，如果替换率太小(例如，p = 0.1)，所有任务的性能都会显著下降。另一方面，不同任务的最佳替换率是不同的。
替换策略的影响
固定替换率(0.5、0.7和0.9这三个值)、curriculum替换策略和对应的反curriculum策略(即用1-p_d作为替换率)的对比结果如 Table 6所示：
可以看出，用curriculum替换策略压缩的模型始终优于用固定替换率压缩的模型。相反，使用反curriculum替换策略压缩后的模型性能明显下降，进一步验证了curriculum替换策略的有效性。
总结
文章提出一种新的名为 Theseus(提修斯)的模型压缩方法，并将该方法应用于BERT模型压缩，实验结果表明该方法优于其他知识蒸馏方法，更加可喜的是该方法只有一个超参数和一个损失函数，无需额外的数据。文章所提出的 theseus 模型压缩方法在模型压缩领域开创了一个新的方向。基于theseus的模型压缩方法有一个显著的不足：继承人模块替换前辈模块时需要二者在输入和输出维度上都各自相同。纵使有上述限制，并不妨碍基于theseus的方案降低网络深度。有学者研究表明，在固定参数容量下，隐层维度尺寸对于计算量的影响低于网络深度。此外，现有很多替换模块，比如 ShuffleNet单元对于ResBlock、Reformer Layer对于Transformer Layer ，这些替换模块可以直接用作 theseus 中的 successor。最后，基于theseus的方案可以使用前馈神经网络在不同大小的隐空间之间做特征映射。该方案在BERT上能够取得可喜的结果，那么未来将会尝试将其扩展到其他神经网络模型。基于theseus的模型压缩并不依赖于任何与特定模型相关的特征，因此有潜力将其应用到其他大模型(比如ResNet)的压缩。后续的研究方向是将基于theseus的模型压缩方法应用于卷积神经网络和图神经网络。