EMNLP 2020 | 忒修斯之船启发下的知识蒸馏新思路

随着 深度学习 的流行，很多大型 神经网络 模型诞生，并在多个领域中取得当前最优的性能。尤其是在 自然语言处理 （NLP）领域中，预训练和调参已经成为其中大多数任务的新范式。基于 Transformer 的预训练模型在 自然语言理解 （NLU）和 自然语言生成 （NLG）领域中成为主流。尽管这些模型从“过 参数 化”的特性中获益，但它们往往包含数百万甚至数十亿个 参数 ，这就使得此类模型的计算成本高昂，且从内存消耗和高延迟的角度来看计算低效。这一缺陷极大地阻碍了此类模型在生产环境中的应用。

为了解决该问题，研究人员提出了很多 神经网络 压缩技术。一般而言，这些技术可以分为三类： 量化 、 权重 剪枝 和 知识蒸馏 （Knowledge Distillation）。其中，由于 知识蒸馏 能够压缩预训练 语言模型 ，所以得到了极大关注。 知识蒸馏 利用大型教师模型“教”紧凑的学生模型模仿教师的行为，从而将教师模型中嵌入的知识迁移到较小的模型中。但是，学生模型的性能状况取决于设计良好的蒸馏 损失函数 ，正是这个函数使得学生模型可以模仿教师的行为。近期关于 知识蒸馏 的研究甚至利用更复杂的模型特定蒸馏 损失函数 ，以实现更好的性能。

近日，来自 微软亚洲研究院 自然语言计算组的研究员们提出了一种与显式地利用蒸馏 损失函数 来最小化教师模型与学生模型距离的 知识蒸馏 不同的模型压缩新方法。受到著名哲学思想实验“忒修斯之船”的启发（即如果船上的木头逐渐被替换，直到所有的木头都不是原来的木头，那这艘船还是原来的那艘船吗？），研究员们在 EMNLP 2020 上发表了 Theseus Compression for BERT (BERT-of-Theseus)，该方法逐步将 BERT 的原始模块替换成 参数 更少的替代模块（点击文末阅读原文，了解论文详情）。研究员们将原始模型叫做“前辈”（predecessor），将压缩后的模型叫做“接替者”（successor），分别对应 知识蒸馏 中的教师和学生。

该方法的工作流程如下图所示。首先为每个前辈模块指定一个接替者模块，然后在训练阶段中以一定的概率（如抛硬币）决定是否用替代模块随机替换对应的前辈模块，并按照新旧模块组合的方式继续训练。在模型 收敛 后，将所有接替者模块组合成接替者模型，进而执行推断。这样就可以将大型前辈模型压缩成紧凑的接替者模型了。

举例来说，假设现在有两支篮球队每支各五人，一支是经验老道的全明星球队，另一支则是年轻球员组成的青训队。为了提高青训队的水平，所以随机选派青训队员去替换掉全明星队中的球员，然后让这个混合的球队不断地练习、比赛。通过向前辈学习经验，新加入成员的实力会有所提升，也能学会和其他队员的配合，逐渐的这个混合球队就拥有了接近全明星球队的实力。之后重复这个过程，直到青训队员都被充分训练，最终青训队员也能自己组成一支实力突出的球队。相比之下，如果没有“老司机”来带一带，青训队无论如何训练，水平也不会达到全明星队的实力。

事实上，Theseus 压缩与 知识蒸馏 的思路有些类似，都是鼓励压缩模型模仿原始模型的行为，但 Theseus 压缩有很多独特的优势。

首先，Theseus 压缩在压缩过程中仅使用任务特定的 损失函数 。而基于 知识蒸馏 的方法除了使用任务特定的 损失函数 外，还需加入繁琐的蒸馏 损失函数 作为优化目标。

其次，与近期研究 TinyBERT 等不同，Theseus 压缩不使用 Transformer 特定特征进行压缩，这就为压缩广泛模型提供了可能性。与 知识蒸馏 仅使用原始模型执行推断不同，该方法允许前辈模型与压缩后的接替者模型共同训练，从而实现更深层次的梯度级交互，并简化训练过程。

此外，混合了前辈模块和接替者模块的不同模块组合还添加了额外的 正则化 项（类似于 Dropout ）。该方法基于课程学习（Curriculum Learning）方法来驱动模块替换，将模块替换概率从低到高逐渐增加，从而实现优异的 BERT 压缩性能。利用Theseus 压缩方法压缩得到的 BERT 模型运算速度是之前的1.94 倍，并且保留了原始模型超过 98% 的性能，优于其它基于 知识蒸馏 的压缩的基线方法。

通过在预训练 语言模型 BERT 上的成功实验， 微软亚洲研究院 的研究员们希望可以为模型压缩打开一种全新的思路，并希望看到这一方法在 计算机视觉 等领域的更广泛应用。