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

一种多层对多层的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）最小化：

其中，是学生模型和教师模型Transformer输入的词向量，是可学的线性映射。

2.对logits的蒸馏：

知识蒸馏中学生模型需要学习教师模型输出的类别概率。我们使用的教师输出概率分布和学生输出概率分布的交叉熵作为损失函数：

其中和分别对应教师模型和学生模型的logits。

3.使用Earth Mover's Distance算法建模学生模型到老师模型隐含层的整体差异：

我们使用Earth Mover's Distance（EMD），衡量老师与学生网络之间的整体差异，EMD方法允许多对多的层映射且无需指定对应层次。

我们将模型的层输出分为hidden state和attention，那么蒸馏学习的目标应使两个分布（老师和学生）接近。给定为老师模型所有层次的Attention输出，为学生模型所有层次的Attention输出。我们给定和作为学生模型和老师模型层次对应的权重。该权重初始为 1/N 和 1/M 。同样的有和作为不同层次的hidden state输出和对应权重，权重同样使用均值初始化。

我们同样使用MSE衡量不同层次之间的距离：

其中，代表学生模型第i层和老师模型第j层对应的attention距离和hidden state距离。分别对应学生和老师模型第i，j层的attention，分别对应学生和老师模型第i，j层的hidden state，是映射函数。

根据EMD算法，我们可以按照距离找到一个层次之间信息的转移量和,分别对应attention和hidden state。以hidden state为例，整体的转移损失可被表示为：

对应的约束条件：

上述优化是一个经过充分研究的运输问题(Hitchcock, 1941)，可以通过线性规划的的方法来解决(Rachev, 1985)。在得到最佳映射流后，我们就可以定义EMD距离：

通过减小EMD距离使老师模型和学生模型的表示更加接近，因此损失函数定义如下：

4.模型整体损失函数：

由上述部分，模型的整体损失函数如下：

其中 β 用来控制层间损失的比重。

EMD算法中包含一个权重因子，代表该层次的最大转移量，并按照均值初始化。我们认为不同任务对应的不同层次权重应有差异。因此我们设计了一套更新该权重的方法。下面以教师模型的权重更新为例，学生模型与之相同。

我们给出了一种按照Cost Attention 更新权重的方式，单位转移成本越小所分配的权重就越大。首先，我们给定每一层的转移成本为，分别对应attention和hidden state。其定义如下如下：

然后根据单位转移成本的反比更新权重：

其中是更新后的权重。

最后将权重使用softmax归一化并取均值：

其中为温度系数，其值越大权重分布越平滑。

三、实验结果

我们对比了BERT Small，DistillBERT，BERT-PKD，TinyBERT，可以看到4层和6层的BERT-EMD均好于这些压缩方法。在未引入数据增强的情况下，6层的BERT-EMD甚至MNLI，QQP，QNLI，MRPC，RTE，STS-b这几个数据集上超过了BERT-base，同时比BERT-base快约一倍。

根据图2的结果，我们有几个关键的观察结果。首先，在压缩Transformer时，不同的任务可能会更倾向于不同的教师层。矩阵的对角线位置对于MNLI任务几乎总是很重要的，它与传统跳层的映射策略呈现相似的趋势。但是，对于RTE任务，每个学生Transformer层都可以向任何教师Transformer层学习。先前的映射方法无法充分利用教师网络。另外从结果中也可以看出，BERT-EMD在RTE数据集上的提升比MNLI要大。

四、总结：

作者推出的BERT-EMD相较于以前工作蒸馏BERT的方式，有效地解决了人工指定蒸馏学习层次对应关系的弊端，并且在隐含层的学习中对每层的重要程度进行了建模。在GLUE Benckmark上，不使用数据增强技术，6层模型有着12层BERT相当的结果。

我们可以进一步思考和优化的点是如何衡量学生模型和老师模型隐含层表征的差异。直接通过MSE去建模老师模型隐含层和学生模型隐含层的表示的差异，并不是很合理。两者并不是同一个向量空间，另外对应向量位置的正负，大小可能对最后概率预测的影响都不可预估。

论文：.06133

代码：

整理：李健铨

排版：田雨晴

本周直播预告：

AI Time欢迎AI领域学者投稿，期待大家剖析学科历史发展和前沿技术。针对热门话题，我们将邀请专家一起论道。同时，我们也长期招募优质的撰稿人，顶级的平台需要顶级的你，请将简历等信息发至yun.he@aminer！

微信联系：AITIME_HY

AI Time是清华大学计算机系一群关注人工智能发展，并有思想情怀的青年学者们创办的圈子,旨在发扬科学思辨精神，邀请各界人士对人工智能理论、算法、场景、应用的本质问题进行探索，加强思想碰撞，打造一个知识分享的聚集地。

更多资讯请扫码关注

（点击“阅读原文”下载本次报告ppt）

（直播回放：）