M obile BERT : a Compact Task-Agnostic BERT for Resource-Limited Devices (2020)

Abstract: Natural Language Processing (NLP) has recently achieved great success by using huge pre-trained models with hundreds of millions of parameters. However, these models suffer from heavy model sizes and high latency such that they cannot be deployed to resource-limited mobile devices. In this paper, we propose MobileBERT for compressing and accelerating the popular BERT model. Like the original BERT, MobileBERT is task-agnostic, that is, it can be generically applied to various downstream NLP tasks via simple fine-tuning. Basically, MobileBERT is a thin version of BERT_LARGE, while equipped with bottleneck structures and a carefully designed balance between self-attentions and feed-forward networks. To train MobileBERT, we first train a specially designed teacher model, an inverted-bottleneck incorporated BERT_LARGE model. Then, we conduct knowledge transfer from this teacher to MobileBERT. Empirical studies show that MobileBERT is 4.3x smaller and 5.5x faster than BERT_BASE while achieving competitive results on well-known benchmarks. On the natural language inference tasks of GLUE, MobileBERT achieves a GLUE score of 77.7 (0.6 lower than BERT_BASE), and 62 ms latency on a Pixel 4 phone. On the SQuAD v1.1/v2.0 question answering task, MobileBERT achieves a dev F1 score of 90.0/79.2 (1.5/2.1 higher than BERT_BASE).

## introduction the nlp community has witnessed a revolution of pre-training self-supervised models. these models usually have hundreds of millions of parameters @xcite @xcite @xcite . among these models, bert @xcite * this work was done when the first author was an intern at google brain. there have been some efforts that taskspecifically distill bert into compact models @xcite @xcite . to the best of our knowledge, there is not yet any work for building a taskagnostic lightweight pre-trained model, that is, a model that can be generically fine-tuned on different downstream nlp tasks as the original bert does. in this paper, we propose mobilebert to fill this gap. in practice, task-agnostic compression of bert is desirable. task-specific compression needs to first fine-tune the original large bert model into a task-specific teacher and then distill. such a process is much more complicated @xcite and costly than directly fine-tuning a task-agnostic compact model. at first glance, it may seem straightforward to obtain a task-agnostic compact bert. for example, one may just take a narrower or shallower version of bert, and train it until convergence by minimizing a convex combination of the prediction loss and distillation loss @xcite . unfortunately, empirical results show that such a straightforward approach results in significant accuracy loss @xcite . this may not be that surprising. it is well-known that shallow networks usually do not have enough representation power while narrow and deep networks are difficult to train. our mobilebert is designed to be as deep as bert large while each layer is made much narrower via adopting bottleneck structures and balancing between self-attentions and feed-forward networks (figure 1 ). to train mobilebert, a deep and thin model, we first train a specially designed teacher model, an inverted-bottleneck incorporated bert large model (ib-bert). then, we conduct knowledge transfer from ib-bert to mobilebert. a variety of knowledge transfer strategies are carefully investigated in our empirical studies. empirical evaluations foot_0 show that mobilebert is 4.3× smaller and 5.5× faster than bert base , while it can still achieve competitive results on well-known nlp benchmarks. on the natural language inference tasks of glue, mobilebert can achieve a glue score of 77.7, which is only 0.6 lower than bert base , with a latency of 62 ms on a pixel 4 phone. on the squad v1.1/v2.0 question answering task, mobileber obtains a dev f1 score of 90.3/80.2, which is even 1.5/2.1 higher than bert base . ## related work recently, compression of bert has attracted much attention. @xcite propose to pre-train the smaller bert models to improve task-specific knowledge distillation. tang et al. (2019) distill bert into an extremely small lstm model. tsai et al. (2019) distill a multilingual bert into smaller bert models on sequence labeling tasks. clark et al. (2019b) use several single-task bert models to teach a multi-task bert. liu et al. (2019a) distill knowledge from an ensemble of bert models into a single bert. concurrently to our work, @xcite distill bert into shallower students through knowledge distillation and an additional knowledge transfer of hidden states on multiple intermediate layers. @xcite propose tinybert, which also uses a layer-wise distillation strategy for bert but in both pre-training and fine-tuning stages. @xcite propose distilbert, which successfully halves the depth of bert model by knowledge distillation in the pre-training stage and an optional fine-tuning stage. in contrast to these existing literature, we only use knowledge transfer in the pre-training stage and do not require a fine-tuned teacher or data augmentation @xcite in the down-stream tasks. another key difference is that these previous work try to compress bert by reducing its depth, while we focus on compressing bert by reducing its width, which has been shown to be more effective @xcite . ## mobilebert in this section, we present the detailed architecture design of mobilebert and training strategies to efficiently train mobilebert. the specific model settings are summarized in ## experiments in this section, we first present our architecture search experiments which lead to the model settings in ## conclusion we have presented mobilebert which is a taskagnostic compact variant of bert. empirical results on popular nlp benchmarks show that mo-bilebert is comparable with bert base while being much smaller and faster. mobilebert can enable various nlp applications 7 to be easily deployed on mobile devices. in this paper, we show that 1) it is crucial to keep mobilebert deep and thin, 2) bottleneck/invertedbottleneck structures enable effective layer-wise knowledge transfer, and 3) progressive knowledge transfer can efficiently train mobilebert. we believe our findings are generic and can be applied to other model compression problems. l1 h1 l1 h2 l1 h3 l1 h4 l12 h1 l12 h2 l12 h3 l12 h4 mobilebert (bare) + pd + fmt + at ib-bert teacher mobilebert (bare) mobilebert (bare) + pd + fmt mobilebert (bare) + pd figure 3: the visualization of the attention distributions in some attention heads of the ib-bert teacher and different mobilebert models.