Extremely Small BERT Models from Mixed-Vocabulary Training
Sanqiang Zhao*1, Raghav Gupta*2, Yang Song3 Denny Zhou4 1University of Pittsburgh, Pittsburgh, PA 2Google Research, Mountain View, CA 3Kuaishou Technology, Beijing, China 4Google Brain, Mountain View, CA
Abstract
Pretrained language models like BERT have achieved good results on NLP tasks, but are impractical on resource-limited devices due to memory footprint. A large fraction of this footprint comes from the input embeddings with large input vocabulary and embedding dimensions. Existing knowledge distillation methods used for model compression cannot be directly applied to train student models with reduced vocabulary sizes. To this end, we propose a distillation method to align the teacher and student embeddings via mixed-vocabulary training. Our method compresses BERTLARGE to a task-agnostic model with smaller vocabulary and hidden dimensions, which is an order of magnitude smaller than other distilled BERT models and offers a better size-accuracy trade-off on language understanding benchmarks as well as a practical dialogue task.
中文速览
大型预训练语言模型（如BERT）因参数量庞大而难以部署在手机等资源受限设备上，其中词嵌入矩阵（input embedding matrix）就占据了相当大一块存储空间，而现有的知识蒸馏（knowledge distillation）方法又无法直接处理师生模型词表不同的情况。为此，研究者提出了一种"混合词表训练"（mixed-vocabulary training）方法：先让教师模型同时处理自身词表和学生词表分词的混合输入，借此将两套词表的嵌入向量对齐，再用对齐后的嵌入初始化学生模型并完成正常蒸馏。实验表明，使用约5000个词片（WordPiece）的学生词表，最终模型体积比其他蒸馏版BERT小一个数量级，在GLUE语言理解基准和口语对话任务上均保持了有竞争力的准确率，最小版本在手机上推理速度比BERT快57倍、大小不足10 MB。这项工作的意义在于，它填补了词表缩减与知识蒸馏之间的技术空白，为将BERT类模型真正落地到边缘设备提供了切实可行的压缩方案。
原文 arXiv:1909.11687;中英对照 + 大白话阅读 https://aha.fim.ai/paper/1909.11687v2