MobileBERT模型简单介绍

目录

一、概要

二、深入扩展

2.1 知识蒸馏方法

2.2 渐进式知识迁移

---

一、概要

MobileBERT 可以看作一个“瘦身”后的BERT-large模型，其使用了瓶颈结构（Bottleneck Structure），并且在自注意力和前馈神经网络的设计上也有一定的改进。MobileBERT能够达到教师模型 BERT-base 99.2% 的性能效果（以GLUE数据集 为测试基准），推理速度快 5.5 倍，参数量降低至23.2%。

MobileBERT 在 BERT 的基础上进行了若干结构改进，例如去掉层归一化、使用ReLU 激活函数等，这里不再展开介绍，感兴趣的读者可以阅读文献进一步了解模型结构细节。下面重点介绍MobileBERT采用的知识蒸馏方法。

二、深入扩展

2.1 知识蒸馏方法

MobileBERT 的损失函数由四部分组成：有监督MLM损失、有监督NSP损失、隐含层蒸馏损失和注意力蒸馏损失。

式中，0≤α≤1表示调节损失函数权重的超参数，在MobileBERT中取α=0.5。

其中，有监督MLM损失和有监督NSP损失与原版BERT的实现是一样的。隐含层蒸馏损失则与TinyBERT一致，计算教师模型和学生模型各层隐含层输出之间的均方误差损失。需要注意的是，由于MobileBERT（学生模型）与教师模型的层数一致（均为12层），这里不需要设计映射函数，只需要将教师模型和学生模型的每一层进行一一对应即可。注意力蒸馏损失也与TinyBERT类似，但在MobileBERT中使用的是基于KL散度的方法，而不是TinyBERT中的均方误差损失。

式中，K 表示注意力头数；KL（·）表示KL散度函数。与 MSE 损失函数不同的是，KL 散度并不是对称的，这一点需要特别注意。MobileBERT 模型的整体结构如下图所示。

2.2 渐进式知识迁移

MobileBERT使用了一种 渐进式知识迁移 （Progressive Knowledge Transfer）策略。下图给出了一个示例，其中教师模型是3层 Transformer 结构，每种颜色的浅色版本表示参数冻结，即参数不参与训练。

可以看到在渐进式知识迁移中，词向量层和最终分类输出层的权重是直接从教师模型拷贝至学生模型的，始终不参与参数更新。而对于中间的Transformer层，采用了渐进的方式逐步训练。首先，学生模型开始学习教师模型的第一层。接下来，学生模型继续学习教师模型的第二层，而此时学生模型的第一层权重是不参与更新的。依此类推，当学生模型学习教师模型的第i层时，学生模型中所有小于 i 层的权重均不参与更新。论文作者通过实验证明这种渐进式知识迁移方法显著优于其他直接蒸馏方法，感兴趣的读者可以阅读文献了解更多细节。