TinyBERT shrinks BERT-base to 14.5M parameters and runs 9.4x faster

Researchers at Huawei Noah’s Ark Lab and Huazhong University of Science and Technology published TinyBERT on September 23, 2019, a compressed version of BERT-base trained by knowledge distillation, where a small student model learns to reproduce the behaviour of a larger teacher. The four-layer student carries 14.5M parameters against the teacher’s 109M, and the authors report it at 7.5x smaller and 9.4x faster on inference while reaching more than 96.8% of the teacher’s score on the GLUE language understanding benchmark. A forward pass costs 1.2B FLOPs against 22.5B for BERT-base.

Rather than fitting only the logits the teacher outputs, the authors define losses on the embedding layer output and, for every Transformer layer, on the hidden states and on the attention matrices those layers produce, following a finding by Kevin Clark and co-authors that BERT’s attention weights capture substantial linguistic knowledge. Dropping the attention term costs the most in their ablation, taking a four-task development average from 75.6 to 71.0, against 72.9 without the hidden states and 73.5 without the prediction layer logits. The distillation also runs twice, first as a general stage against the unfine-tuned BERT on a large text corpus, which produces a general TinyBERT that initialises the student, then as a task-specific stage against the fine-tuned BERT on a task dataset expanded by data augmentation. Removing either stage costs more still, with that same average falling to 72.5 without the general stage and 68.5 without the task-specific one.

On the GLUE test set the authors report the 14.5M-parameter model at an average of 77.0 against 79.5 for BERT-base, with 82.5 on MNLI-m, 92.6 on SST-2, 87.7 on QNLI, 71.3 on QQP, 86.4 on MRPC, 80.4 on STS-B and 66.6 on RTE. The widest gap is CoLA, the linguistic acceptability task, at 44.1 against 52.8. The authors measured the speedups on a single NVIDIA K80 GPU, not on a phone. A six-layer variant at 67M parameters averages 79.4, which the authors call on par with the teacher.

Huawei released the checkpoints on Hugging Face, where the general four-layer model is published as huawei-noah/TinyBERT_General_4L_312D alongside a six-layer version, with the training code in the company’s Pretrained-Language-Model repository on GitHub. The paper later appeared in Findings of EMNLP 2020.