All updates

Researchers at Carnegie Mellon University and Google Brain published MobileBERT on April 6, 2020, a compressed version of the BERT language model built for phones. It has 25.3M parameters against 109M for BERT-base, and the authors measured 62 ms per inference on a Pixel 4, which they report as 4.3 times smaller and 5.5 times faster than BERT-base.

The saving comes from the shape of each layer. MobileBERT keeps the 24 layers of the much larger BERT-large but makes every block narrow, so that a block works internally at a width of 128 while the representation flowing between blocks stays 512 wide, with a small linear layer at each end to shrink the input and widen the output again, an arrangement the paper calls a bottleneck. Narrowing the block leaves the attention module holding too large a share of the parameters, so each block stacks 4 feed-forward networks behind its single attention module to restore the usual balance. The authors also traced a large part of the remaining latency to layer normalisation and the gelu activation and replaced both with cheaper operations, which cut inference from 192 ms to 62 ms without changing the number of arithmetic operations.

A network that deep and thin is hard to train directly, so the team first trained a teacher and then copied its behaviour layer by layer. The teacher is BERT-large fitted with inverted bottlenecks, which widen inside the block but narrow the representation passing between blocks to the same 512 the student uses, so the two models’ layer outputs line up and can be compared one to one during the transfer. That transfer happens only during pre-training, which keeps the result task-agnostic, so one distilled model is fine-tuned separately for each downstream task and no task-specific teacher is needed.

MobileBERT scores 77.7 on the GLUE language understanding benchmark against 78.3 for BERT-base, and on the SQuAD question answering sets v1.1 and v2.0 it reaches dev F1 scores of 90.0 and 79.2, which the authors put 1.5 and 2.1 above BERT-base. The latency figures are the authors’ own runs, with the models exported to TensorFlow Lite and timed on a 4-thread Pixel 4 at a fixed sequence length of 128, where BERT-base took 342 ms. A smaller variant with 15.1M parameters runs in 40 ms and scores 75.8 on GLUE, and 8-bit quantisation leaves the reported accuracies almost unchanged. Code and pre-trained weights are published in Google Research’s repository.

Hugging Face published DistilBERT on October 2, 2019, a compressed version of BERT trained by knowledge distillation, where a small student model learns to reproduce the output of a larger teacher. The student carries 66M parameters against BERT-base’s 110M, and Victor Sanh and his co-authors report it retaining 97% of BERT’s score on the GLUE language understanding benchmark while running 60% faster.

The authors tested the model on a phone by building a question-answering app for an iPhone 7 Plus and comparing it against their own BERT-base question-answering model. Excluding the tokenisation step, they measured DistilBERT at 71% faster, and the whole model in that app weighs 207 MB, which they note could shrink further with quantisation. On a server CPU at batch size 1, the same pair takes 410 seconds against 668 for a full pass over the STS-B development set.

Three losses train the student at once. A distillation loss pushes it towards the teacher’s full output distribution rather than only the correct token, a masked language modelling loss keeps BERT’s own pre-training objective in place, and a cosine embedding loss aligns the directions of student and teacher hidden state vectors. The student halves BERT’s layer count, drops the token-type embeddings and the pooler, and starts from every second layer of the teacher, an initialisation the ablation study values at 3.69 points of GLUE macro score. Training ran on eight 16 GB V100 GPUs for about 90 hours on English Wikipedia and the Toronto Book Corpus.

Hugging Face released the trained weights together with the training code in its transformers library, and published the iOS demo separately as Core ML code. The paper was accepted at the fifth Workshop on Energy Efficient Machine Learning and Cognitive Computing, co-located with NeurIPS 2019.

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.