PaperlayerView exact revision on arXiv

Source-linked AI summary

MobileBERT: a Compact Task-Agnostic BERT for Resource-Limited Devices

Zhiqing Sun, Hongkun Yu, Xiaodan Song, Renjie Liu, Yiming Yang, Denny Zhou

arXiv:2004.02984v2cs.CLcs.LG

TL;DR

Large pre-trained NLP models are difficult to deploy on resource-limited mobile devices because of their size and latency. MobileBERT addresses this with a deep, narrow bottleneck architecture trained by transferring knowledge from an inverted-bottleneck BERT_LARGE teacher. It is 4.3× smaller and 5.5× faster than BERTBASE while achieving competitive benchmark results, though it still degrades substantially relative to its teacher.

  • Problem

    BERT’s heavy model size and high latency make it impractical for resource-limited mobile devices, while task-agnostic compact pre-trained alternatives were not yet available.

  • Method

    MobileBERT uses a deep, narrow bottleneck architecture and is trained through knowledge transfer from an inverted-bottleneck BERTLARGE teacher.

  • Results

    MobileBERT is 4.3× smaller and 5.5× faster than BERTBASE while achieving competitive results on GLUE and SQuAD benchmarks.

  • Takeaways & Limitations

    MobileBERT supports deploying various NLP applications on mobile devices while retaining task-agnostic fine-tuning and competitive benchmark performance.

  • Takeaways & Limitations

    MobileBERT degrades greatly compared with its IB-BERT teacher, leaving substantial room for improvement.

Abstract

from arXiv · show

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 GLUEscore o 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).

1 Introduction

MobileBERT addresses the difficulty of deploying large BERT models on mobile devices by providing a task-agnostic compact model that retains depth while reducing width. It uses bottleneck structures, an inverted-bottleneck BERT_LARGE teacher, and knowledge transfer to achieve competitive benchmark performance with lower size and latency.

  • BERT’s heavy model size and high latency make it impractical to deploy on resource-limited mobile devices.
  • Task-agnostic compact pre-trained BERT models were missing, whereas existing compression methods focused on task-specific distillation.
  • Straightforward narrower or shallower BERT variants incur significant accuracy loss because shallow networks lack representation power and narrow deep networks are difficult to train.
  • MobileBERT remains as deep as BERTLARGE while making each layer narrower through bottleneck structures and balancing self-attention with feed-forward networks.
  • MobileBERT is trained through knowledge transfer from an inverted-bottleneck incorporated BERTLARGE teacher called IB-BERT.
  • 4.3× smaller and 5.5× faster than BERTBASE, MobileBERT achieves a GLUE score of 77.7 with 62 ms latency on a Pixel 4 phone.On SQuAD v1.1/v2.0, it achieves dev F1 scores of 90.3/80.2, respectively.

2 Related Work

Prior work compressed BERT through task-specific distillation, reduced depth, or other teacher-student strategies. MobileBERT instead targets task-agnostic width reduction and uses pre-training-stage knowledge transfer without requiring a fine-tuned teacher or downstream data augmentation.

  • Earlier efforts distilled BERT into compact models for task-specific knowledge, sequence labeling, multi-task learning, or extremely small recurrent models.
  • Concurrent approaches used layer-wise distillation, often reducing BERT depth during pre-training and sometimes during fine-tuning.
  • MobileBERT uses knowledge transfer only during pre-training and does not require a fine-tuned teacher or downstream data augmentation.
  • Unlike prior depth-reduction approaches, MobileBERT compresses BERT by reducing width, which prior work reports as more effective.

3 MobileBERT

MobileBERT uses a deep, narrow bottleneck architecture with an inverted-bottleneck BERT_LARGE teacher and multiple knowledge-transfer objectives. The design also rebalances attention and feed-forward capacity and reduces inference costs through simpler operations.

  • Architecture: MobileBERT keeps BERT_LARGE’s depth while narrowing each building block with bottlenecks, using a hidden dimension of 128 and linear transformations to 512.The architecture uses two linear transformations per block to adjust input and output dimensions.
  • Teacher and transfer: The IB-BERTLARGE teacher is trained to convergence before transferring knowledge to MobileBERT, whose matching 512 feature-map size enables layer-wise comparison.Direct training from scratch performs worse according to the paper.
  • Architecture: Stacked feed-forward networks restore the balance between MHA and FFN disrupted by bottlenecks; MobileBERT uses one MHA followed by four FFNs per layer.MHA supports attention across subspaces, whereas FFN increases model non-linearity.
  • Operational optimizations: Operational optimizations replace layer normalization with NoNorm and GELU with ReLU to reduce latency, while the embedding dimension is reduced to 128 and projected to 512 by a kernel-size-3 convolution.NoNorm is an element-wise linear transformation parameterized by γ and β.
  • Training objectives: Layer-wise transfer combines feature-map matching with attention-map matching, while pre-training distillation combines MLM, NSP, and MLM knowledge-distillation losses.Feature-map transfer uses mean squared error, and attention transfer minimizes KL divergence between per-head attention distributions.
  • Training strategies: The paper evaluates auxiliary, joint, and progressive knowledge-transfer strategies, including staged training that separates layer-wise transfer from pre-training distillation and progressively trains layers.Progressive transfer divides training into L stages and can tune lower layers with a smaller learning rate.

4 Experiments

Experiments show that MobileBERT achieves strong benchmark performance with substantially lower model size and latency, supported by bottleneck-based architecture choices and knowledge-transfer strategies.

  • Architecture Search: Reducing the IB-BERT teacher’s inter-block hidden size to 512 does not damage performance, whereas shrinking its intra-block hidden size makes performance dramatically worse.
  • Architecture Search: MobileBERT performance peaks when MHA-to-FFN parameter ratios range from 0.4∼0.6, motivating a 128-dimensional, four-FFN student architecture.
  • Results on GLUE: 77.7 GLUE score is only 0.6 lower than BERTBASE, while MobileBERT outperforms other compressed BERT models with smaller or similar model sizes.
  • 4.3× smaller and 5.5× faster than BERTBASE, MobileBERT remains competitive on well-known NLP benchmarks.
  • Operational Optimizations: Quantization compresses MobileBERT by 4× with nearly no performance degradation, indicating additional compression potential.
  • Training Strategies: Progressive knowledge transfer consistently outperforms auxiliary and joint strategies, while Feature Map Transfer contributes most to performance improvement.
  • Training Objectives: The IB-BERTLARGE teacher matches the original teacher’s power, but MobileBERT degrades greatly relative to its teacher, leaving room for improvement.

5 Conclusion

MobileBERT is presented as a task-agnostic compact BERT variant that approaches BERTBASE performance while reducing model size and latency for mobile deployment.

  • MobileBERT is a task-agnostic compact variant of BERT.
  • MobileBERT is comparable with BERTBASE while being much smaller and faster.
  • MobileBERT can enable various NLP applications to be deployed on mobile devices.
  • The paper identifies deep-and-thin design, bottleneck structures, and progressive knowledge transfer as key findings.

Appendix for “MobileBERT: a Compact Task-Agnostic BERT for Resource-Limited Devices”

This related-work discussion traces knowledge transfer from output distributions to intermediate representations, attention maps, and sequential teacher-to-student transfer.

  • Knowledge distillation trains a smaller student network to mimic a larger teacher's class-distribution output.
  • FitNets extend transfer by matching intermediate hidden layers to train narrow and deep networks.
  • Prior work also transfers knowledge from top hidden neurons, sequentially from deep teachers to shallow students, and through attention maps or hidden-state similarities.

B Extra Related Work on Compact Architecture Design

Compact architecture research has largely targeted mobile CNNs, while many lightweight convolutional operations do not directly transfer to Transformer or BERT models.

  • Recent efficient CNN research has focused on mobile vision applications.
  • Lightweight operations such as depth-wise convolution cannot be directly applied to Transformer or BERT architectures.
  • The NLP literature includes group LSTMs, which adapt group-convolution ideas to recurrent neural networks.

C Visualization of Attention Distributions

The visualization compares attention distributions across selected layers and models, showing how attention transfer relates student and teacher attention patterns.

  • Figure 3 visualizes attention distributions in the first and twelfth layers of selected models from the ablation study.
  • Attention transfer helps the student mimic the teacher's attention distributions closely.
  • MobileBERT(bare)+PD+FMT produces attention distributions that are a reordering of those from the attention-transfer model and teacher.
  • The authors connect this pattern to multi-head attention's role in BERT non-linearity and to feature-map alignment.

D Extra Experimental Settings

The experiments match BERT’s preprocessing and use a fixed loss-balancing coefficient for pre-training distillation. Several alternative training techniques are left for future work.

  • 15% of WordPiece tokens are randomly masked, with next sentence prediction retained to match BERT’s preprocessing.
  • Span prediction and removing next sentence prediction are identified as potential future improvements.
  • α = 0.5 balances the original masked language modeling loss and the distillation loss during pre-training distillation.

E Architecture of MobileBERTTINY

MobileBERTTINY reduces multi-head attention parameters by using reduced intra-block feature maps as its query, key, and value inputs, although this may reduce model capacity.

  • MobileBERTTINY uses reduced intra-block feature maps as the key, query, and value inputs to multi-head attention.This replaces hidden states from inter-block feature maps as the MHA inputs.
  • The lighter multi-head attention structure effectively reduces parameters in its MHA modules.
  • The parameter reduction might harm MobileBERTTINY’s model capacity.

F GLUE Dataset

The GLUE benchmark covers linguistic acceptability, sentiment, paraphrase, similarity, question-answering NLI, textual entailment, and multi-genre natural language inference tasks. The experiments skip WNLI because few prior works outperform majority-class prediction.

  • Task descriptions: GLUE includes CoLA, which evaluates whether examples are grammatical English sentences using Matthews correlation coefficient.CoLA consists of acceptability judgments from books and journal articles on linguistic theory.
  • Task descriptions: SST-2 predicts sentiment in movie-review sentences and is evaluated by accuracy.
  • Task descriptions: MRPC and QQP assess whether sentence or question pairs are semantically equivalent using accuracy and F1 score.MRPC uses sentence pairs from online news sources, while QQP uses question pairs from Quora.
  • Task descriptions: STS-B predicts human similarity scores from 1 to 5 for sentence pairs using Pearson and Spearman correlations.
  • Task descriptions: MNLI predicts entailment, contradiction, or neutrality between premise-hypothesis pairs, with accuracy reported for matched and mismatched sections.
  • Task descriptions: QNLI determines whether a context sentence contains the answer to a question and is evaluated by test accuracy.The dataset is converted from SQuAD.
  • Task descriptions: RTE predicts textual entailment for sentence pairs using accuracy.The datasets come from annual textual entailment challenges.
  • Task descriptions: WNLI is skipped because few previous works outperform majority-class prediction on its pronoun-referent reading-comprehension task.
Loading 2004.02984v2…