一种多层对多层的BERT蒸馏方法

⬆⬆⬆ 点击蓝字

关注我们

AI TIME欢迎每一位AI爱好者的加入！

李健铨，天津大学计算机学院2014级硕士生，神州泰岳AI研究院深度学习实验室总监。在自然语言处理方向合作发表论文6篇，申请发明专利55项，获得授权8项。负责研发的产品“泰岳语义工厂”获得第九届（2019）年吴文俊人工智能科技进步奖（企业技术创新工程项目）。

一、简介

在 NLP 领域，BERT一经出现就吸引了所有人的目光。此后根据BERT推出了 XLNET、RoBERTa、ALBERT、T5、Megatron-LM、Turning-NLP 等性能卓越的改进模型。但是这些模型体积愈发庞大，从 BERT-large 的 3.4 亿参数到 Turing-NLP 参数规模的 170 亿参数。单个样本计算一次的开销动辄上百毫秒，很难应用到实际生产中。BERT蒸馏技术应运而生，解决了模型大小和推理速度的问题，而且使用该技术的模型大小并没有明显的下降。

近一年来，知识蒸馏（Knowledge Distillation）作为一种常用的模型压缩方法逐渐成为BERT蒸馏的研究热点分支。最初针对BERT模型的蒸馏学习是利用教师和学生网络的输出logits计算蒸馏损失，然而这样不能学到Transformer 中间层的信息。随后的研究通过对教师 BERT 模型中所有Transformer层的first token，attention，hidden parameter输出等进行蒸馏，将大 BERT 的中间层信息迁移到小模型上，学习到的学生模型有了更好的表现。这种方法最近也成为了BERT蒸馏的主流方法。包括PKD-BERT到TinyBERT，MobileBERT，都是采用指定层对应学习的方式进行蒸馏。

然而为了压缩模型大小，学生模型的层次数量一般小于教师BERT模型的层次数量，因此中间层的学习不能做到层次一一对应。目前的BERT蒸馏方法均使用跨层映射的方法（Skip），其层次映射函数为, 其中 i,j为教师，学生模型的指定层次，N,M为教师，学生模型的层数。举例而言，如果教师模型有12层，对应的学生模型有4层，具体对应为 student 第 1 层 transformer 对应 teacher 第 3 层，第 2 层对应第 6 层，第 3 层对应第 9 层，第 4 层对应第 12 层。

图 1 跨层映射策略

该指定方法获得了不错的效果，但存在如下问题：该蒸馏过程中部分层次的信息必然被舍弃。不同任务可能需要学习不同层次的知识，这种强制指定不一定适应所有任务。

作者提出的方法解决了这两个问题，无需进行层次对应指定，且能够实现多层对多层的蒸馏学习。我们的主要贡献如下：

1. 我们提出了先进的多层对多层的BERT蒸馏方法，BERT的中间层以自适应和整体的方式进行学习。

2. 我们引入EMD方法衡量学生模型和教师模型之间的差异。

3. 我们提出了一种Cost Attention 机制，在蒸馏学习中建模不同隐含层的重要程度。

4. 在GLUE任务上进行的大量实验表明，BERT-EMD的性能优于最新的BERT蒸馏方法。

二、模型实现细节

图 2 模型

与TinyBERT类似，我们的方法同样包括logits，hidden parameter，embedding和attention的蒸馏。下面我们分别介绍：

1.对embedding的蒸馏：

词向量在NLP任务中至关重要，我们希望压缩词向量矩阵维度而不牺牲性能。为此，我们将教师模型和学生模型的词向量之间的均方误差（MSE）最小化：