论文阅读_知识蒸馏_MobileBERT
英文题目：MobileBERT: a Compact Task-Agnostic BERT for Resource-Limited Devices
中文题目：MobileBERT：面向资源有限设备的任务无关的压缩模型
论文地址：https://arxiv.org/pdf/2004.02984
领域：自然语言处理，模型蒸馏
发表时间：2020
作者：Zhiqing Sun，卡内基梅隆大学，谷歌
出处：ACL
被引量：162
代码和数据：https://github.com/google-research/google-research/tree/master/mobilebert
阅读时间：22.06.16
读后感
使用了瓶颈结构，同时微调模型实现加速，最终实现了任务无关的蒸馏 BERT 模型。
介绍
大模型效果好，但受资源限制。文中提出了 MobileBERT 模型对 BERT 模型进行压缩和加速，它和原始 BERT 一样，是任务无关的，可以泛化到不同下游任务。MobileBERT 是对 BERT LARGE 的优化，使用瓶颈结构。具体实现分两步：先训练一个倒瓶颈的特殊的老师模型，然后再将知识转移到学生模型 MobileBERT。它比 BERT BASE 模型小 4.3 倍，快 5.5 倍，在 GLUE 数据集上达到与 BASE BASE 类似的效果。
文中提出的 MobileBERT 模型与 BERT LARGE 的深度相同，用瓶颈结构使每一层更窄，如图 -1 所示：
与之前的模型相比，文中模型只在预训练阶段使用知识转移，在下游任务中不需要微调教师模型和数据增强；另外，它不修改模型层数，只改变了每层的宽度。
方法
图 -1 展示了模型参数的具体大小：
瓶颈和逆瓶颈
从表 -1 中可以看到，每块维度仅为 128，另外，将全流程的层输入输出都调整为更窄的 512。
在训练过程中，首先训练老师模型 IB-BERT，再将知识转移到学生模型 MobileBERT。老师模型设计如图 -1(b) 所示，它是一个逆瓶颈结构，师生模型的 Feature map 大小均为 512，这样可以直接对比层输出。同时使用瓶颈和逆瓶颈使老师模型保持性能，学生模型足够紧凑。
多层前馈网络
注意力层和前馈网络功能不同：前者连接了不同空间数据，后者提升了模型的非线性表达能力。一般 BERT 的注意力层和前馈网络参数比例为 1:2，而由于 MobileBERT 模型比较窄，注意力参数相对更多，为解决此问题，设计了多个叠加的前馈层，如图 -1(c) 所示。
操作优化
通过延迟分析发现，归一化层和 gelu 激活函数占延迟的很大比例，因此对其进行优化。
- 去掉归一化层：用针对元素的线性变换替换 n 通道的隐藏层归一化。
- 使用 relu 代替 glue
嵌入分解
嵌入词表占了模型的很大比例，如表 -1 所示，MobileBERT 将嵌入词表压缩到 128 维，再用核为 3 的卷积层将其扩展到 512 维。
训练目标
转移目标包含：feature map 转移和注意力转移。
Feature map 转移
BERT 中每层的输出是下层的输入，因此，需要让学生的层输出尽量与老师结果一致。这里使用 MSE 作为损失函数。
公式中 l 为层索引，T 为序列长度，N 为 feature map 大小。
注意力转移
注意力提升了 BERT 模型的语言识别能力，因此也通过学习注意力层将知识更好地转移到 MobileBERT 模型。使用 KL 散度来计算注意力层的差异，作为损失函数：
预训练蒸馏
除了针对各 Transformer 层的转移，在预训练时，还使用了蒸馏的 loss 函数，最终损失函数由三部分组成：遮蔽的损失 (MLM)，下一句预测的损失 (NSP) 和针对 MLM 的蒸馏损失：
训练策略
下面讨论三种训练策略，如图 -2 所示：
辅助知识转移（AKT）
将中间层的知识转移作为知识蒸馏的辅助任务。损失函数是所有层的知识转移损失和预测蒸馏损失的线性组合。
联合知识转移（JKT）
将训练分成两步：先训练中间层知识转移，然后训练预训练的知识蒸馏。
递进知识转移（PKT）
下层的错误可能影响到上层的学习，因此设计了先训练下层，再训练上层的知识转移方法。一般在训练上层时冻结下层参数，也可以使用软化方法，训练上层时对下层参数使用较小的学习率。实验证明该方法效果最好。
实验
表 -2 展示了在 SQuAD 数据集上，使用不同参数大小的蒸馏结果：
图 -3 展示了堆叠前馈层的效果：
精调下游任务的方法与普通 BERT 一致，略有不同的是 MobileBERT 常需要更大的学习率和更多次迭代。主实验结果如表 -8 所示：
可以看到，由于是对 BERT LARGE 的蒸馏，MobileBERT 与 BERT BASE 模型效果类似，而大小比 BERT 小 4.3 倍，速度是 BERT 的 5.5 倍。
表 -8 对比了三种不同策略的训练方法，可以看到递进训练效果更好。
文中用的每个优化技术都在实验中做了效果比较，这里只罗列了最重要的部分，其它详见论文正文。