苏剑林作者追一科技学校NLP、神经网络研究方向

BERT-of-Theseus:基于模块替换的模型压缩方法

最近了解到一种称为“BERT-of-Theseus”的 BERT 模型压缩方法,来自论文 BERT-of-Theseus: Compressing BERT by Progressive Module Replacing。这是一种以“可替换性”为出发点所构建的模型压缩方案,相比常规的剪枝、蒸馏等手段,它整个流程显得更为优雅、简洁。

论文标题:BERT-of-Theseus: Compressing BERT by Progressive Module Replacing
论文链接:https://arxiv.org/abs/2002.02925

本文将对该方法做一个简要的介绍,给出一个基于 bert4keras [1] 的实现,并验证它的有效性。
▲ BERT-of-Theseus,原作配图
一、模型压缩

首先,我们简要介绍一下模型压缩。不过由于笔者并非专门做模型压缩的,也没有经过特别系统的调研,所以该介绍可能显得不专业,请读者理解。

1.1 基本概念

简单来说,模型压缩就是“简化大模型,得到推理速度更快的小模型”。当然,一般来说模型压缩是有一定牺牲的,比如最明显的是最后的评测指标会有一定的下降,毕竟“更好又更快”的免费午餐是很少的,所以选择模型压缩的前提是能允许一定的精度损失。

其次,模型压缩的提速通常只体现在预测阶段,换句话说,它通常需要花费更长的训练时间,所以如果你的瓶颈是训练时间,那么模型压缩也不适合你。

模型压缩要花费更长时间的原因是它需要“先训练大模型,再压缩为小模型”。读者可能会疑惑:为什么不直接训练一个小模型?答案是目前很多实验已经表明,先训练大模型再压缩,相比直接训练一个小模型,最后的精度通常会更高一些。

也就是说,在推理速度一样的情况,压缩得到的模型更优一些,相关探讨可以参考论文 Train Large, Then Compress: Rethinking Model Size for Efficient Training and Inference of Transformers  [2] ,另外知乎上也有讨论《为什么要压缩模型,而不是直接训练一个小的 CNN?》[3]。

1.2 常见手段

常见的模型压缩技术可以分为两大类:1)直接简化大模型得到小模型;2)借助大模型重新训练小模型。这两种手段的共同点是都先要训练出一个效果比较好的大模型,然后再做后续操作。

第一类的代表方法是剪枝(Pruning)量化(Quantization)

剪枝,顾名思义,就是试图删减掉原来大模型的一些组件,使其变为一个小模型,同时使得模型效果在可接受的范围内;

至于量化,指的是不改变原模型结构,但将模型换一种数值格式,同时也不严重降低效果,通常我们建立和训练模型用的是 float32 类型,而换成 float16 就能提速且省显存,如果能进一步转换成 8 位整数甚至 2 位整数(二值化),那么提速省显存的效果将会更加明显。

第二类的代表方法是蒸馏(Distillation)。蒸馏的基本想法是将大模型的输出当作小模型训练时的标签来用,以分类问题为例,实际的标签是 one hot 形式的,大模型的输出(比如 logits)则包含更丰富的信号,所以小模型能从中学习到更好的特征。

除了学习大模型的输出之外,很多时候为了更进一步提升效果,还需要小模型学习大模型的中间层结果、Attention 矩阵、相关矩阵等,所以一个好的蒸馏过程通常涉及到多项 loss,如何合理地设计这些 loss 以及调整这些 loss 的权重,是蒸馏领域的研究主题之一。

二、Theseus

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

BERT-of-Theseus 的命名源于思想实验“忒修斯之船”:如果忒修斯的船上的木头被逐渐替换,直到所有的木头都不是原来的木头,那这艘船还是原来的那艘船吗?

2.1 核心思想

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

举个实际的类比:

假设现在有 A、B 两支球队,每支各五人。A 球队属于明星球队,实力超群;B 球队则是新手球队,待训练。为了训练 B 球队,我们从 B 球队中选 1 人,替换掉 A 球队中的 1 人,然后让这个“4+1”的A球队不断的练习、比赛。经过一段时间,新加入的成员实体会提升,这个“4+1”的球队就拥有接近原始 A 球队的实力。

重复这个过程,直到 B 球队的人都被充分训练,那么最终 B 球队的人也能自己组成一支实力突出的球队。相比之下, 如果一开始就只有 B 球队,只是 B 球队的人自己训练、比赛,那么就算他们的实力逐渐提升,但由于没有实力超群的 A 球队帮助,其最终实力也不一定能突出。

2.2 流程细节

回到BERT的压缩,现在假设我们有一个 6 层的 BERT,我们直接用它在下游任务上微调,得到一个效果还不错的模型,我们称之为 Predecessor(前辈)

我们的目的是得到一个 3 层的 BERT,它在下游任务重效果接近 Predecessor,至少比直接拿 BERT 的前 3 层去微调要好(否则就白费力气了),这个小模型我们称为 Successor(后辈)。那么 BERT-of-Theseus 是怎么实现这一点的呢?如下图。

▲ BERT-of-Theseus训练过程示意图

▲ Predecessor和Successor模型示意图

在 BERT-of-Theseus 的整个流程中,Predecessor 的权重都被固定住。6 层的 Predecessor 被分为 3 个模块,跟 Successor 的 3 层模型一一对应,训练的时候,随机用 Successor 层替换掉 Predecessor 的对应模块,然后直接用下游任务的优化目标进行微调(只训练 Successor 的层)。

训练充分后,再把整个 Successor 单独分离出来,继续在下游任务中微调一会,直到验证集指标不再上升。


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


由于  非 0 即 1(不作调整,各自 0.5 概率随机选效果就挺好了),所以每个分支其实就相当于只有一个模块被选择到,因此上面右图就相当于下述模型结构。由此每次的置零都是随机的,因此训练足够多的步数后,Successor 的每个层都能被训练好。

2.3 方法分析

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

前面说到,蒸馏多数时候也需要匹配中间层输出,这时候要涉及到的训练目标就有很多了:下游任务 loss、中间层输出 loss、相关矩阵 loss、Attention 矩阵 loss、等等,想想要平衡这些 loss 就是一件头疼的事情。

相比之下,BERT-of-Theseus 直接通过替换这个操作,逼着 Successor 能有跟 Predecessor 类似的输出,而最终的训练目标就只有下游任务 loss,不可谓不简洁。

此外,BERT-of-Theseus 还有一个特别的优势:很多的蒸馏方法都得同时作用于预训练和微调阶段,效果才比较突出,而 BERT-of-Theseus 直接作用于下游任务的微调,就可以得到相媲美的效果。这个优势在算法上体现不出来,属于实验结论。

从形式上来看,BERT-of-Theseus 的随机替换思路有点像图像中的数据扩增方案 SamplePairing 和 mixup(参考《从 SamplePairing到mixup:神奇的正则项》[4] ),都是随机采样两个对象加权求和来增强原模型;也有点像 PGGAN [5] 的渐进式训练方案,都是通过对两个模型进行某种程度的混合,实现两个模型的过渡。

如果了解它们的读者,继而就能够对 BERT-of-Theseus 提出一些拓展或者说疑问: 一定要非 0 即 1 吗,任意  的随机数行不?或者说不随机,直接让  慢慢地从 1 变到 0 行不?这些想法都还没有经过充分实验,有兴趣的读者可以修改下述代码自行实验。

三、实验效果

原作者们开源了自己的 PyTroch 实现:

https://github.com/JetRunner/BERT-of-Theseus

邱震宇老兄也分享了自己的讲解 [6] 以及基于原版 BERT 的 Tensorflow 实现 qiufengyuyi/bert-of-theseus-tf [7]。当然,既然笔者决定写这篇介绍,那就肯定少不了基于 bert4keras 的 Keras 实现了:

https://github.com/bojone/bert-of-theseus

这大概是目前最简洁、最具可读性的 BERT-of-Theseus 实现了,没有之一。

原论文的效果大家就自己去看原论文了。笔者在几个文本分类任务上实验了一下,结果大同小异,跟邱兄的实验结论也比较一致。其中在 CLUE 的 iflytek 数据集中实验结果如下:

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

另外,对于蒸馏来说,如果 Successor 跟 Predecessor 有同样的结构(同模型蒸馏),那么通常来说 Successor 的最终性能比 Predecessor 还要好些,BERT-of-Theseus 有没有这一特点呢?

笔者也实验了一下该想法,发现结论是否定的,也就是同模型情况下 BERT-of-Theseus 训练出来的 Successor 并没有比 Predecessor 好,所以看来 BERT-of-Theseus 虽好,但也不能完全取代蒸馏。

四、文末小结

本文介绍并实验了一种称为“BERT-of-Theseus”的 BERT 模型压缩方法,该方法的特点是简洁明了,纯粹通过替换操作来让小模型去学习大模型的行为,使得能在只有一个 loss 的情况下就能达到当前最优的模型压缩效果。

参考文献

[1] https://github.com/bojone/bert4keras

[2] https://arxiv.org/abs/2002.11794

[3] https://www.zhihu.com/question/303922732

[4] https://kexue.fm/archives/5693

[5] https://arxiv.org/abs/1710.10196

[6] https://zhuanlan.zhihu.com/p/112787764

[7] https://github.com/qiufengyuyi/bert-of-theseus-tf

PaperWeekly
PaperWeekly

推荐、解读、讨论和报道人工智能前沿论文成果的学术平台。

理论BERT
相关数据
权重技术

线性模型中特征的系数,或深度网络中的边。训练线性模型的目标是确定每个特征的理想权重。如果权重为 0,则相应的特征对模型来说没有任何贡献。

二值化技术

二值化是将像素图像转换为二进制图像的过程。

Dropout技术

神经网络训练中防止过拟合的一种技术

剪枝技术

剪枝顾名思义,就是删去一些不重要的节点,来减小计算或搜索的复杂度。剪枝在很多算法中都有很好的应用,如:决策树,神经网络,搜索算法,数据库的设计等。在决策树和神经网络中,剪枝可以有效缓解过拟合问题并减小计算复杂度;在搜索算法中,可以减小搜索范围,提高搜索效率。

文本分类技术

该技术可被用于理解、组织和分类结构化或非结构化文本文档。文本挖掘所使用的模型有词袋(BOW)模型、语言模型(ngram)和主题模型。隐马尔可夫模型通常用于词性标注(POS)。其涵盖的主要任务有句法分析、情绪分析和垃圾信息检测。

验证集技术

验证数据集是用于调整分类器超参数(即模型结构)的一组数据集,它有时也被称为开发集(dev set)。

分类问题技术

分类问题是数据挖掘处理的一个重要组成部分,在机器学习领域,分类问题通常被认为属于监督式学习(supervised learning),也就是说,分类问题的目标是根据已知样本的某些特征,判断一个新的样本属于哪种已知的样本类。根据类别的数量还可以进一步将分类问题划分为二元分类(binary classification)和多元分类(multiclass classification)。

知乎机构

作为中文互联网综合性内容平台,知乎将AI广泛应用与社区,构建了人、内容之间的多元连接,提升了社区的运转效率和用户体验。知乎通过内容生产、分发,社区治理等领域的AI应用,也创造了独有的技术优势和社区AI创新样本。

https://www.zhihu.com
量化技术

深度学习中的量化是指,用低位宽数字的神经网络近似使用了浮点数的神经网络的过程。

推荐文章
暂无评论
暂无评论~