BERT-of-Theseus

155****7220

最近了解到一种称为"BERT-of-Theseus"的BERT模型压缩方法，源自论文《BERT-of-Theseus: Compressing BERT by Progressive Module Replacing》。这是一种以"可替换性"为出发点所构建的模型压缩方案，相比常规的剪枝、蒸馏等手段，它整个流程显得更为优雅、简洁。本文将对该方法做一个简要的介绍，并给出一个简单的Pytorch实现

模型压缩

简单来说，模型压缩就是"简化大模型，得到推理速度更快的小模型"。当然，一般来说模型压缩是有一定牺牲的，比如最后的评测指标会有一定的下降，毕竟"更好又更快"的免费午餐是很少的，所以选择模型压缩的前提是能允许一定的精度损失。其次，模型压缩的梯度通常只体现在预测（Inference）阶段，换句话说，它通常需要花费更长的训练时间，所以如果你的瓶颈是训练时间，那么模型压缩也不适合你

模型压缩要花费更长时间的原因是它需要"先训练大模型，再压缩为小模型"。读者可能会疑惑：为什么不直接训练一个小模型？答案是目前很多实验已经表明，先训练大模型再压缩，相比直接训练一个小模型，最后的精度通常会更高一些。也就是说，在推理速度一样的情况下，压缩得到的模型更优一些，相关探讨可以参考论文《Train Large, Then Compress: Rethinking Model Size for Efficient Training and Inference of Transformers》，另外知乎上也有讨论《为什么要压缩模型，而不是直接训练一个小的CNN？》

常见手段

常见的模型压缩技术可以分为两大类：

直接简化大模型得到小模型
借助大模型重新训练小模型

这两种手段的共同点是都要先训练出一个效果比较好的大模型，然后再做后续操作

第一类的代表方法是剪枝（Pruning）和量化（Quantization）。剪枝，顾名思义，就是试图删减掉原来大模型的一些组件，使其变为一个小模型，同时使得模型效果在可接受的范围内；至于量化，指的是不改变原模型的结构，但将模型换一种数值格式，同时也不严重降低效果，通常我们建立和训练模型用的是float32类型，而换成float16类型就能提速且省显存，如果能进一步转换成8位整数甚至2位整数，那么提速省显存的效果将会更加明显

第二类的代表方法是蒸馏（Distillation）。蒸馏的基本想法是将大模型的输出当作小模型训练时的标签来用，以分类问题为例，实际的标签是one-hot形式的，大模型的输出（比如logits）则包含更丰富的信号，所以小模型能从中学习到更好的特征。除了学习大模型的输出之外，很多时候为了更近一步提升效果，还需要小模型学习大模型的中间层结果、Attention矩阵等，所以一个好的蒸馏过程通常涉及到多项loss，如何合理地涉及这些loss以及调整这些loss的权重，是蒸馏领域的研究主题之一。下图从左到右分别是剪枝、量化、蒸馏

Theseus

本文将要介绍的模型压缩方法称为"BERT-of-Theseus"，属于上面说的两大类模型压缩方法的第二类，也就是说它也是借助大模型来训练小模型，只不过它是基于模块的可替换性来设计的

BERT-of-Theseus的命名源于思想实验"忒修斯之船"：如果忒修斯船上的木板被逐渐替换，直到所有的模板都不再是原来的木板，那这艘船还是原来的那艘船吗？

核心思想

前面说到，用蒸馏做模型压缩时，往往不仅希望小模型的的输出跟大模型的输出对齐，还希望中间层结果也对齐。"对齐"意味着什么呢？意味着可替换！所以BERT-of-Theseus的思想就是：干嘛要煞费苦心地通过添加各种loss去实现可替换性呢？直接用小模型的模块去替换掉大模型的模块然后训练不就好了吗？

举个实际的例子：

假设现在有A、B两支球队，每支各五人。A球队属于王者球队，实力超群；B球队则是青铜球队，待训练。为了训练B球队，我们从B球队中选1人，替换掉A球队中的1人，然后让这个"4+1"的A球队不断的练习、比赛。经过一段时间，新加入的成员实力会提升，这个"4+1"的球队就拥有接近原始A球队的实力。重复这个过程，直到B球队的人都被充分训练，那么最终B球队的人也能自己组成一支实力突出的球队。相比之下，如果一开始就只有B球队，只是B球队的人自己训练、比赛，那么就算他们的实力逐渐提升，但由于没有实力超群的A球队帮助，其最终实力也不一定能突出

流程细节

回到BERT的压缩，现在假设有一个6层预训练好的BERT，我们直接用它在下游任务上微调，得到一个效果还不错的模型，称之为Predecessor（前辈）；现在我们的目的是得到一个3层的BERT，并且这个它在下游任务中的效果接近Predecessor，至少比直接拿BERT的前3层去微调要好（否则就白费力气了），这个小模型我们称为Successor（传承者）。那么BERT-of-Theseus是怎么实现这一点的呢？如下图（右）

在BERT-of-Theseus的整个流程中，Predecessor的权重都被固定住，6层的Predecessor被分为3个模块，与Successor的3层模型一一对应。训练的时候，随机用Successor层替换掉Predecessor的对应模块，然后直接用下游任务的优化目标进行微调（只训练Successor的层）。训练充分后，再把整个Successor单独分离出来，继续在下游任务中微调一会儿，直到验证集指标不再上升

在实现的时候，事实上是类似Dropout的过程，同时执行Predecessor和Successor模型，并将两者对应模块的输出之一置零，然后求和、送如下一层中，即

$ε(l)∼U(0,1)\varepsilon^{(l)}\sim U({0, 1})$
$x(l)=xp(l)×ε(l)+xs(l)×(1−ε(l))x^{(l)} = x_p^{(l)} \times \varepsilon^{(l)} + x_s^{(l)} \times \left(1 - \varepsilon^{(l)}\right)$
$xp(l+1)=Fp(l+1)(x(l))x_p^{(l+1)} = F_p^{(l+1)}\left(x^{(l)}\right)$
$xs(l+1)=Fs(l+1)(x(l))x_s^{(l+1)} = F_s^{(l+1)}\left(x^{(l)}\right)$

由于 $ε\varepsilon$ 非0即1（不做调整，各自0.5的概率随机效果就挺不错了），所以每个分支其实就相当于只有一个模块被选择到。由于每次的置零都是随机的，因此训练足够多的步数后，Successor的每层都能被训练好

方法分析

与蒸馏相比，BERT-of-Theseus有什么优势呢？首先，这既然能被发表出来，所以至少效果应该是不相上下的，所以我们就不去比较效果了，而是比较方法本身。很明显，BERT-of-Theseus的主要特点是：简洁

前面说到，蒸馏多数时候需要匹配中间层的输出，这时涉及到的训练目标就有很多了：下游任务loss、中间层输出loss、相关矩阵loss、Attention矩阵loss等等，想要平衡这些loss本身就是一件头疼的事情。相比之下，BERT-of-Theseus直接通过替换这个操作，逼着Successor能有跟Predecessor类似的输出，而最终的训练目标就只有下游任务loss，不可谓不简洁。此外，BERT-of-Theseus还有一个特别的优势：很多的蒸馏方法都得同时作用于预训练和微调阶段，效果才比较突出，而BERT-of-Theseus直接作用于下游任务的微调，就可以得到相媲美的效果。这个优势在算法上体现不出来，属于实验结论

$ε\varepsilon$ 一定要非0即1吗？任意0~1的随机数行不？或者说不随机，直接让 $ε\varepsilon$ 慢慢地从1变到0行不？这些想法都还没有经过充分实验，有兴趣的读者可以自行实验

实验效果

原作者们开源了自己的PyTorch实现 JetRunner/BERT-of-Theseus，知乎用户邱震宇也分享了自己的讲解以及基于原本BERT的Tensorflow实现qiufengyuyi/bert-of-theseus-tf。原论文的效果大家就自己去看原论文了，这里po出苏剑林大佬在CLUE的iflytek数据集中的实验结果：

可以看到，相比直接拿前几层微调，BERT-of-Theseus确实能带来一定的性能提升。对于随机置零方案，除了均等概率选择0/1外，原论文还尝试了其他策略，有轻微提升，但会引入额外超参

另外，对于蒸馏来说，如果Succesor跟Predecessor有同样的结果（同模型蒸馏），那么通常来说Successor的最终性能比Predecessor还要好些，BERT-of-Theseus有没有这一特点呢？苏剑林大佬的实验发现结论好似否定的，也就是同模型情况下BERT-of-Theseus训练出来的Successor并没有Predecessor好，所以看来BERT-of-Theseus虽好，但也不能完全取代蒸馏