Extremely Small BERT Models from Mixed-Vocabulary Training

TL;DR

This paper introduces a novel mixed-vocabulary distillation method to create extremely small BERT models, significantly reducing size while maintaining performance for resource-constrained NLP applications.

Contribution

It proposes a new distillation approach that aligns teacher and student embeddings with reduced vocabulary, enabling effective compression of BERT-LARGE into a much smaller, task-agnostic model.

Findings

Achieves an order of magnitude smaller BERT models

Maintains competitive accuracy on language benchmarks

Offers improved size-accuracy trade-offs

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.