Extremely Small BERT Models from Mixed-Vocabulary Training (2021)

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 BERT-LARGE 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.

## introduction recently, pre-trained context-aware language models like elmo @xcite , gpt @xcite , bert @xcite and xlnet @xcite have outperformed traditional word embedding models like word2vec @xcite and glove @xcite , and achieved strong results on a number of language understanding tasks. however, these models are typically too huge to host on mobile/edge devices, especially for real-time inference. recent work has explored, inter alia, knowledge distillation @xcite to train small-footprint student models by implicit transfer of knowledge from a teacher model. most distillation methods, however, need the student and teacher output spaces to be aligned. this complicates task-agnostic distillation of bert to asterisk (*) denotes equal contribution. research conducted when all authors were at google. smaller-vocabulary student bert models since the input vocabulary is also the output space for the masked language modeling (mlm) task used in bert. this in turn limits these distillation methods' ability to compress the input embedding matrix, that makes up a major proportion of model parameters e.g. the ∼30k input wordpiece embeddings of the bert base model make up over 21% of the model size. this proportion is even higher for most distilled bert models, owing to these distilled models typically having fewer layers than their teacher bert counterparts. we present a task and model-agnostic distillation approach for training small, reduced-vocabulary bert models running into a few megabytes. in our setup, the teacher and student models have incompatible vocabularies and tokenizations for the same sequence. we therefore align the student and teacher wordpiece embeddings by training the teacher on the mlm task with a mix of teacher-tokenized and student-tokenized words in a sequence, and then using these student embeddings to train smaller student models. using our method, we train compact 6 and 12-layer reducedvocabulary student models which achieve competitive performance in addition to high compression for benchmark datasets as well as a real-world application in language understanding for dialogue. ## related work work in nlp model compression falls broadly into four classes: matrix approximation, weight quantization, pruning/sharing, and knowledge distillation. the former two seek to map model parameters to low-rank approximations @xcite and lower-precision integers/floats @xcite @xcite respectively. in contrast, pruning aims to remove/share redundant model weights @xcite . more recently, dropout @xcite has been used to cut inference latency by early exit @xcite . knowledge distillation focuses on implicit transfer of knowledge as soft teacher predictions @xcite , attention distributions @xcite and intermediate outputs @xcite . approaches close to our work rely on similar methods @xcite , while others involve combinations of layer-wise transfer @xcite , taskspecific distillation @xcite , architecture search @xcite and layer dropout @xcite ; many of these are specific to the transformer layer @xcite . another highly relevant line of work focuses on reducing the size of the embedding matrix, either via factorization @xcite or vocabulary selection/pruning @xcite . ## proposed approach here, we discuss our rationale behind reducing the student vocabulary size and its challenges, followed by our mixed-vocabulary distillation approach. ## discussion impact of vocabulary size: we trained a model from scratch identical to bert base except with our 5k-wp student vocabulary. on the sst-2 and mnli-m dev sets, this model obtained 90.9% and 83.7% accuracy respectively -only 1.8% and 0.7% lower respectively compared to bert base . since embeddings account for a larger fraction of model parameters with fewer layers, we trained another model identical to our 6×256 model, but with a 30k-wp vocabulary and teacher label dis-tillation. this model showed small gains (0.1% / 0.5% accuracy on sst-2 / mnli-m dev) over our analogous distilled model, but with 30% more parameters solely due to the larger vocabulary. this suggests that a small wordpiece vocabulary may be almost as effective for sequence classification/tagging tasks, especially for smaller bert models and up to moderately long inputs. curiously, increasing the student vocabulary size to 7k or 10k did not lead to an increase in performance on glue. we surmise that this may be due to underfitting owing to the embeddings accounting for a larger proportion of the model parameters. alternative vocabulary pruning: probing other strategies for a small-vocabulary model, we used the above 6×256 30k-wp vanilla distilled model to obtain a smaller model by pruning the vocabulary to contain the intersection of the 30k and 5k vocabularies (total 4629 wps). this model is 1.2% smaller than our 4928-wp distilled model, but drops 0.8% / 0.7% on sst-2/mnli-m dev sets. furthermore, to exploit the high overlap in vocabularies, we tried running our distillation pipeline but with the embeddings for student tokens (after projecting up to the teacher dimension) also present in the teacher vocabulary tied to the teacher embeddings for those tokens. this model, however, dropped 0.7% / 0.5% on sst-2/mnli-m compared to our analogous 6×256 distilled model. we also tried pretraining bert large from scratch with the 5k vocabulary and doing vanilla distillation for a 6×256 student: this model dropped 1.2% / 0.7% for sst-2/mnli-m over our similar distilled model, indicating the efficacy of mixed-vocabulary training over vanilla distillation. ## conclusion we propose a novel approach to knowledge distillation for bert, focusing on using a significantly smaller vocabulary for the student bert models. our mixed-vocabulary training method encourages implicit alignment of the teacher and student word-piece embeddings. our highly-compressed 6 and 12-layer distilled student models are optimized for on-device use cases and demonstrate competitive performance on both benchmark datasets and practical tasks. our technique is unique in targeting the student vocabulary size, enabling easy combination with most bert distillation methods.