TinyBERT: 面向预训练语言模型的知识蒸馏方法
以BERT为代表的预训练语言模型显著提升了诸多自然语言处理(NLP)任务的效果,成为当前NLP领域主流的语义理解模型。但预训练语言模型参数量巨大,推理时间长,难以在资源受限的终端设备上进行部署。为了解决该痛点,华为诺亚方舟实验室语音语义团队提出了一种面向BERT模型的知识蒸馏算法。通过此算法压缩得到的小模型TinyBERT在只有原始BERT模型13%参数量的情况下,推断加速达到9倍,同时在自然语言处理标准评测GLUE上获得原始BERT模型96%的效果;另一方面,对比当前其他主流蒸馏方法,TinyBERT在只有28%的参数和31%的推断时间情况下,平均提升3个点以上。同时,TinyBERT结合诺亚自研的bolt( https://github.com/huawei-noah/bolt )在手机芯片上实现了进一步加速,目前在Kirin810上单句(长度为32)平均推理时间可以优化至7ms,满足手机端侧的推理需求,为预训练语言模型的端侧部署提供了一套完整的解决方案。TinyBERT的相关论文公开在: https://arxiv.org/abs/1909.10351 ,相关代码和论文发布在 https://github.com/huawei-noah/Pretrained-Language-Model 。
模型方法
面向Transformer的蒸馏
现有的预训练语言模型大多基于Transformer的结构,因此我们首先提出一种新的面向 transformer 的知识蒸馏算法。通过均匀选层的策略,新的蒸馏算法从老师模型中选择若干层与学生模型各层进行相对应的知识蒸馏(如下图a)。面向Transformer结构的蒸馏综合考虑了不同类型层,包括嵌入层embedding layer、转换层transformer layer和预测层prediction layer。其中transformer层的蒸馏尤其关键,其包含了基于注意力矩阵和基于隐状态的蒸馏(如下图b)。通过这种新的蒸馏算法,原始的老师BERT 模型中编码的大量知识可以有效地迁移到学生 TinyBERT 模型中去。另一方面,在学生模型构建上,TinyBERT在深度(层数)和宽度(隐状态维度)的选择相比之前蒸馏算法会更加灵活。
两阶段知识蒸馏
为了充分迁移老师模型的语义知识,TinyBERT学习包含通用性蒸馏(general distillation)和任务相关性蒸馏(task-specific distillation)两个阶段,分别在预训练学习阶段和任务学习阶段执行面向Transformer的蒸馏算法(如下图)。通用性蒸馏阶段使用预训练得到的BERT-base为老师模型,在大规模文本语料上进行知识蒸馏得到通用性的TinyBERT。此阶段的蒸馏让TinyBERT学习到丰富的语义知识,提升了小模型的泛化能力。任务相关蒸馏阶段使用在任务数据上微调好的BERT-base模型作为老师,在增强的训练数据集上进行面向transformer的知识蒸馏算法,获得任务相关的TinyBERT。在数据增强方面,我们提出了基于预训练语言模型BERT和词向量模型GloVe的词替换方法,让TinyBERT在更多的语料上学习,进一步提升其泛化能力。两阶段的蒸馏起到相互辅助的作用:通用性的蒸馏为任务相关性蒸馏提供了一个好的模型初始;通过学习任务相关的知识,任务相关性蒸馏可以进一步提升TinyBERT的效果。
实验结果
我们在自然语言处理标准评测集GLUE上验证了TinyBERT的模型效果,结果如下表。实验结果表明:(1)相比直接使用原始预训练方法得到BERT-small模型,TinyBERT在平均分数上有6.3点的提升;(2)TinyBERT显著超过了当前的基准模型DistillBERT和BERT-PKD等,并且在模型大小和推断时间上有一定的优势;(3)相比老师模型BERT-base,TinyBERT只使用了1/7的参数量,在各项任务上获得了有竞争力的结果,同时达到了9.4倍的模型推断加速。(4)TinyBERT具有与基于LSTM的基准模型Distilled BiLSTM-SOFT相匹配的模型大小和推断速度,但效果显著好。
为了进一步分析TinyBERT效果,我们测试了不同模型结构下TinyBERT的表现,实验结果如下表。结果表明:(1)增加层数和隐藏层维度都能提升TinyBERT效果,6层的模型基本可以匹配老师模型的结果;(2)4层TinyBERT模型在效果上可以超过6层基准模型BERT-PKD和DistilBERT,进一步说明了所提出方法的有效性。
同时,我们对不同蒸馏模块(GD:General Distillation、TD:Task-specific Distillation和DA:Data Augmentation)和不同类型蒸馏层(Embd: Embedding layer、Pred:Prediction layer和Trm:Transformer layer)分别进行了消融实验,结果如下表所示。实验结果表明:(1)本工作提出的通用性蒸馏、任务相关性蒸馏以及数据增强对TinyBERT学习都有显著的帮助,其中任务相关蒸馏和数据增强模块在四个数据集上有持平的影响,同时两者相比通用性蒸馏模块影响更大;(2)Transformer层的蒸馏是TinyBERT学习的关键,基于注意力矩阵的蒸馏相比隐藏层的蒸馏更加重要。
总结和展望
本工作提出了一种行之有效的面向预训练语言模型的知识蒸馏方法,为在终端设备上部署BERT应用打下了基础。TinyBERT自公开以来受到业界广泛关注,同时也在相关业务场景下成功落地。在未来,诺亚方舟实验室语音语义团队会持续预训练语言模型小型化的相关研究,进一步探索在不同应用场景下(如搜索、问答和生成等)有效的蒸馏方法;同时结合量化和剪枝等不同类型的模型压缩方法,得到推断更快、模型更小、精度更高且满足终端部署要求的语义理解模型。
- 点赞
- 收藏
- 关注作者
评论(0)