MobileBERT runs in 62 ms on a Pixel 4 with 25.3M parameters
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.