# MobileBERT: Rethinking Model Design for Effective Knowledge Distillation - Canonical: https://33rdsquare.com/mobilebert/ - Published: 2024-09-03 - Author: Jordan Brown - Categories: [Artificial Intelligence & Machine Learning & ChatGPT](https://33rdsquare.com/category/tech/ai/) --- In the past few years, we‘ve seen remarkable progress in natural language processing (NLP) with the emergence of giant pre-trained language models like BERT, GPT-3, and PaLM. These models, with hundreds of billions of parameters trained on massive amounts of text data, have achieved unprecedented performance across a wide range of language understanding and generation tasks. However, the immense size of these state-of-the-art models presents a significant challenge when it comes to deploying them in real-world applications, particularly on resource-constrained devices like mobile phones. Inference with such large models is extremely computationally intensive, making it impractical to run them on-device. This has motivated a line of research into compressing large language models into smaller, more efficient models that can be feasibly deployed while preserving as much of the original model‘s performance as possible. The primary approach that has emerged for model compression is _knowledge distillation_. In knowledge distillation, we train a compact "student" model to mimic the behavior of a large "teacher" model. The basic idea is that the student model can learn to make similar predictions as the teacher by minimizing the difference between their outputs or intermediate representations. While knowledge distillation has proven to be an effective way to compress large models, most work has treated model distillation as an afterthought – something to be done after a high-performing teacher model is trained. In a 2020 paper, researchers at Google proposed a novel approach called [MobileBERT](https://arxiv.org/abs/2004.02984) that turns this paradigm on its head. MobileBERT is an efficient NLP model specifically designed from the ground up to be distilled, putting knowledge distillation at the forefront of the model design process. The key insight behind MobileBERT is that if we know we‘re going to be distilling the model later, we can make specific architectural choices to make the distillation process more effective and enable stronger performance in the final distilled model. Let‘s dive into some of the key innovations that make MobileBERT unique. ## Bottlenecks: Aligning Teacher and Student Representations One of the core challenges in knowledge distillation is aligning the intermediate representations of the teacher and student models. In order to train the student to mimic the teacher, we need some way to compare their activations at different layers. A common approach is to add a mean-squared error or cosine distance loss between corresponding layers. The problem is that this requires the student and teacher activations to be the same size. Most prior work achieves this by having the student mirror the teacher‘s architecture exactly, just with smaller layer sizes. However, this couples the student‘s architecture to the teacher‘s and can be suboptimal. Ideally, we‘d like the flexibility to design the student architecture however we want. MobileBERT introduces the concept of _bottleneck_ layers to address this. The idea is to add linear projection layers at the input and output of each transformer block that project to a fixed smaller dimension. These bottlenecks effectively compress the representations passed between blocks. The key is that the bottleneck dimension is kept the same between the student and teacher models. This allows corresponding blocks to have aligned representations without constraining the internal block dimensions to match. The bottlenecks provide a point of alignment for distillation while giving freedom to customize the block internals. ## Stacked Feed-Forward Networks: Balancing Speed and Expressiveness The overall design philosophy of MobileBERT is to slim down the model as much as possible while maintaining performance. One area where MobileBERT makes a key departure from standard BERT models is in the feed-forward networks (FFNs) inside each transformer block. In BERT, the FFNs contain the majority of the model‘s parameters. They consist of two linear layers with a large inner dimension and a non-linear activation in between. The ratio of parameters between the FFN and multi-head attention components is typically around 2:1. To keep the model lean, MobileBERT uses much narrower FFN layers. However, this presents a problem: Cutting down the FFN size so drastically reduces the model‘s capacity and empirically harms performance. The solution MobileBERT employs is to use _stacked_ FFN layers – essentially adding more FFN layers with narrower inner dimensions to compensate. Ablation studies found that stacking 4 FFN layers provided the best balance of speed and performance. By playing with the FFN inner dimensions and number of stacked layers, MobileBERT is able to maintain a high parameter ratio between the FFNs and attention layers while keeping the total model size small. This allows MobileBERT to devote more of its limited parameter budget to the attention layers, which are critical for capturing context and long-range dependencies. ## Operational Optimizations: Maximizing Inference Efficiency Beyond architectural design choices, MobileBERT also incorporates a few lower-level optimizations to further improve inference speed, aimed at maximizing efficiency on mobile devices. One optimization is replacing the GELU activation used in BERT with the simpler ReLU activation. GELU is somewhat expensive to compute, while ReLU can be implemented extremely efficiently. Experiments showed that switching to ReLU gave a significant speedup with minimal impact on model quality. MobileBERT also simplifies the layer normalization operation used extensively throughout the model. Standard layer norm requires computing the mean and variance statistics of the input. MobileBERT replaces this with a simple elementwise affine transformation that only requires storing a single scale and bias parameter per neuron. Again, this gives a nontrivial speedup while being much more parameter-efficient. Together, these optimizations allow MobileBERT to achieve very impressive inference speeds. On a Pixel 4 smartphone, MobileBERT is able to process input sequences in as little as 62 ms – over 5x faster than BERT-base. ## Distillation Objectives and Procedures: Training a Student From Scratch The MobileBERT architecture is designed with knowledge distillation in mind from the start, using bottlenecks and other design choices to make the model easier to distill effectively. But how exactly is the model distilled? The researchers propose two main distillation objectives: 1. **Feature map transfer**: This is the standard approach of aligning the activations of corresponding layers between student and teacher. MobileBERT minimizes the mean-squared error between bottleneck activations. 2. **Attention transfer**: In addition to aligning activations, MobileBERT also tries to mimic the attention distributions of the teacher. At each corresponding attention head, the Kullback-Leibler divergence between student and teacher attention distributions is minimized. The intuition is that the student should learn to "attend" to the same inputs as the teacher. In terms of the actual distillation process, the paper explores a few different approaches. The simplest is _auxiliary distillation_, where the distillation losses are simply added on top of the standard masked language modeling and next sentence prediction losses, and everything is trained end-to-end. However, the authors found that _joint_ distillation worked better, where the model is first trained using only the distillation losses until convergence, and then fine-tuned with the standard objectives. They also experimented with a _progressive_ distillation approach where the model is distilled one layer at a time, freezing each layer‘s weights before moving on to the next. The idea is to prevent distillation errors from earlier layers from propagating to later layers. With these training procedures, MobileBERT is able to be distilled effectively from a large teacher like BERT-large in a single stage. Critically, MobileBERT can then be fine-tuned on downstream tasks directly, without involving the teacher at all. This is in contrast to prior approaches like TinyBERT that required fine-tuning the teacher on each task before distilling a student, which is much more expensive. ## Results: Impressive Performance and Efficiency So how well does MobileBERT actually perform? The results are quite impressive. Despite being dramatically smaller and faster, MobileBERT is able to achieve 97-100% of the performance of BERT-base on the GLUE benchmark suite. The full 25M parameter MobileBERT without the inference optimizations slightly outperforms BERT-base on average while being 4.3x smaller and 5.5x faster. With the ReLU and layer norm optimizations, MobileBERT does trade off a couple accuracy points for a large speedup, but is still very competitive with larger models. Perhaps even more remarkable are the results for the "tiny" variant of MobileBERT with only 15M parameters. MobileBERT-tiny outperforms the 60M parameter TinyBERT model across the board while being 4x smaller. This really highlights the impact of MobileBERT‘s architectural choices in enabling highly parameter-efficient models. It‘s worth emphasizing again that these MobileBERT models are achieving these results without any fine-tuning of the teacher model. They are distilled entirely from a pre-trained teacher in a task-agnostic way, then fine-tuned on their own. This is significant, as it means MobileBERT can serve as a "generic" compact model that can be applied to new tasks cheaply and efficiently, much like how we use pre-trained BERT models today. ## Conclusion and Future Directions MobileBERT represents an important conceptual shift in how we approach the design of efficient NLP models. Rather than compressing models as an afterthought, MobileBERT shows the benefits of designing a model to be distilled from the start. By using techniques like bottleneck layers and stacked FFNs, MobileBERT is able to achieve an unprecedented combination of performance and efficiency. The results make it clear that putting knowledge distillation at the forefront of model design is a promising direction. While most work continues to focus on scaling models ever larger, MobileBERT points to an alternate path of developing compact, lightweight models that punch above their weight. That said, knowledge distillation is not a panacea. Distilled models are still fundamentally limited by the performance of their teachers. To make continued progress, we will need to keep improving our training techniques and model architectures. MobileBERT provides a strong foundation to build on – it will be exciting to see how these ideas evolve and where they take us next. One thing that has become clear is that model size is no longer the main bottleneck for performance. MobileBERT shows that with the right design and training, a model 1/4 the size of BERT-base can match its performance. The limiting factor now appears to be the pre-training process itself – there remains a lot to be gained from more sample-efficient pre-training that can extract more learning signal from the same data. Approaches like [ELECTRA](https://arxiv.org/abs/2003.10555) that use more sophisticated pre-training tasks show a promising path forward here. As NLP models become more widely deployed in real-world applications, efficiency will only become more important. MobileBERT lights the way for a future of highly capable NLP that can run on the devices in our pockets. Accessible, efficient, and performant natural language interfaces have the potential to fundamentally change how we interact with technology – MobileBERT brings us one step closer to that future. --- Source: [MobileBERT: Rethinking Model Design for Effective Knowledge Distillation](https://33rdsquare.com/mobilebert/)