最近了解到一种称为 "BERT-of-Theseus" 的 BERT 模型压缩方法,源自论文《BERT-of-Theseus: Compressing BERT by Progressive Module Replacing》。这是一种以 "可替换性" 为出发点所构建的模型压缩方案,相比常规的剪枝、蒸馏等手段,它整个流程显得更为优雅、简洁
模型压缩
简单来说,模型压缩就是 " 简化大模型,得到推理速度更快的小模型 "。当然,一般来说模型压缩是有一定牺牲的,比如最后的评测指标会有一定的下降,毕竟" 更好又更快 " 的免费午餐是很少的,所以选择模型压缩的前提是能允许一定的精度损失。其次,模型压缩的梯度通常只体现在预测(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 训练过程示意图(右)
在 BERT-of-Theseus 的整个流程中,Predecessor 的权重都被固定住,6 层的 Predecessor 被分为 3 个模块,与 Successor 的 3 层模型一一对应。训练的时候,随机用 Successor 层替换掉 Predecessor 的对应模块,然后直接用下游任务的优化目标进行微调(只训练 Successor 的层)。训练充分后,再把整个 Successor 单独分离出来,继续在下游任务中微调一会儿,直到验证集指标不再上升
在实现的时候,事实上是类似 Dropout 的过程,同时执行 Predecessor 和 Successor 模型,并将两者对应模块的输出之一置零,然后求和、送如下一层中,即
$$ \begin{equation}\begin{aligned} &\varepsilon^{(l)}\sim U(\{0, 1\})\\ &x^{(l)} = x_p^{(l)} \times \varepsilon^{(l)} + x_s^{(l)} \times \left(1 - \varepsilon^{(l)}\right)\\ &x_p^{(l+1)} = F_p^{(l+1)}\left(x^{(l)}\right)\\ &x_s^{(l+1)} = F_s^{(l+1)}\left(x^{(l)}\right) \end{aligned}\tag{1}\end{equation} $$
由于 $\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 数据集中的实验结果:
$$ \begin {array}{c|c|c} \hline & \text {直接微调} & \text {BERT-of-Theseus}\\ \hline \begin {array}{c}\text {层数} \\ \text {效果}\end {array} & \begin {array}{ccc}\text {完整 12 层} & \text {前 6 层} & \text {前 3 层} \\ 60.11\% & 58.99\% & 57.96\%\end {array} & \begin {array}{cc}\text {6 层} & \text {3 层} \\ 59.61\% & 59.36\% \end {array}\\ \hline \end {array} $$
可以看到,相比直接拿前几层微调,BERT-of-Theseus 确实能带来一定的性能提升。对于随机置零方案,除了均等概率选择 0/1 外,原论文还尝试了其他策略,有轻微提升,但会引入额外超参
另外,对于蒸馏来说,如果 Succesor 跟 Predecessor 有同样的结果(同模型蒸馏),那么通常来说 Successor 的最终性能比 Predecessor 还要好些,BERT-of-Theseus 有没有这一特点呢?苏剑林大佬的实验发现结论好似否定的,也就是同模型情况下 BERT-of-Theseus 训练出来的 Successor 并没有 Predecessor 好,所以看来 BERT-of-Theseus 虽好,但也不能完全取代蒸馏