Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Extremely Small BERT Models from Mixed-Vocabulary Training

Sanqiang Zhao, Raghav Gupta, Yang Song, Denny Zhou

TL;DR

This work tackles the memory bottleneck of BERT-like models by introducing a mixed-vocabulary distillation method that aligns embeddings between a large teacher vocabulary and a compact student vocabulary. It employs a two-stage training process: Stage I blends teacher and student tokenizations to align embeddings, Stage II trains the student solely with the reduced vocabulary. The resulting 6- and 12-layer, 5K-vocabulary models achieve competitive results on GLUE and SNIPS while offering substantial parameter reductions, making on-device NLP more feasible. The approach also sheds light on the trade-offs of vocabulary size, showing small vocabularies can maintain strong performance and be more efficient for edge deployments.

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.

Extremely Small BERT Models from Mixed-Vocabulary Training

TL;DR

This work tackles the memory bottleneck of BERT-like models by introducing a mixed-vocabulary distillation method that aligns embeddings between a large teacher vocabulary and a compact student vocabulary. It employs a two-stage training process: Stage I blends teacher and student tokenizations to align embeddings, Stage II trains the student solely with the reduced vocabulary. The resulting 6- and 12-layer, 5K-vocabulary models achieve competitive results on GLUE and SNIPS while offering substantial parameter reductions, making on-device NLP more feasible. The approach also sheds light on the trade-offs of vocabulary size, showing small vocabularies can maintain strong performance and be more efficient for edge deployments.

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.

Paper Structure

This paper contains 12 sections, 2 equations, 2 figures, 1 table.

Table of Contents

  1. Introduction
  2. Related Work
  3. Proposed Approach
  4. Student Vocabulary
  5. Mixed-vocabulary training
  6. Experiments
  7. Evaluation Tasks and Datasets
  8. Models and Baselines
  9. Training Details
  10. Results
  11. Discussion
  12. Conclusion

Figures (2)

  • Figure 1: Depiction of our mixed-vocabulary training approach. (Left) Stage I involving retrained teacher BERT with default config (e.g., 30K vocabulary, 768 hidden dim) and mixed-vocabulary input. (Right) Stage II involving student model with smaller vocabulary (5K) and hidden dims (e.g., 256) and embeddings initialized from stage I.
  • Figure 2: Test set accuracy of distilled models, teacher model and baselines on the GLUE test sets, along with other parameters. MNLI-m and MNLI-mm refer to the genre-matched and mismatched test sets. All models other than NoKD and our distilled models use a 30K-WordPiece vocabulary. The average is computed using F1 score for MRPC and accuracy for MNLI-m and SST-2.