BERT - EMD : Many-to-Many Layer Mapping for BERT Compression with Earth Mover’s Distance (2020)

Abstract: Pre-trained language models (e.g., BERT) have achieved significant success in various natural language processing (NLP) tasks. However, high storage and computational costs obstruct pre-trained language models to be effectively deployed on resource-constrained devices. In this paper, we propose a novel BERT distillation method based on many-to-many layer mapping, which allows each intermediate student layer to learn from any intermediate teacher layers. In this way, our model can learn from different teacher layers adaptively for different NLP tasks. In addition, we leverage Earth Mover’s Distance (EMD) to compute the minimum cumulative cost that must be paid to transform knowledge from teacher network to student network. EMD enables effective matching for the many-to-many layer mapping. Furthermore, we propose a cost attention mechanism to learn the layer weights used in EMD automatically, which is supposed to further improve the model’s performance and accelerate convergence time. Extensive experiments on GLUE benchmark demonstrate that our model achieves competitive performance compared to strong competitors in terms of both accuracy and model compression

## introduction in recent years, pre-trained language models, such as gpt @xcite , bert @xcite , xl-net @xcite , have been proposed and applied to many nlp tasks, yielding state-of-the-art performances. however, the promising results of the pre-trained language models come with the high costs of computation and memory in inference, which obstruct these pre-trained language models to be deployed on resource-constrained devices and real-time applications. for example, the original bert-base model, which achieved great success in many nlp tasks, has 12 layers and about 110 millions parameters. it is therefore critical to effectively accelerate inference time and reduce the computational workload while maintaining accuracy. this research issue has attracted increasing attention @xcite @xcite , of which knowledge distillation @xcite is considered to be able to provide a practical way. typically, knowledge distillation techniques train a compact and shallow student network under the guidance of a complicated larger teacher network with a teacher-student strategy @xcite . once trained, this compact student network can be directly deployed in real-life applications. so far, there have been several studies, such as distilbert @xcite , bert-pkd @xcite , tinybert @xcite , which attempt to compress the original bert into a lightweight student model without performance sacrifice based on knowledge distillation. for example, bert-pkd @xcite and tiny-bert @xcite are two representative bert compression approaches, which encourage the student model to extract knowledge from both the last layer and the intermediate layers of the teacher network. despite the effectiveness of previous studies, there are still several challenges for distilling comprehensive knowledge from the teacher model, which are not addressed well in prior works. first, existing compression methods learn one-to-one layer mapping, where each student layer is guided by only one specific teacher layer. for example, bert-pkd uses the 2, 4, 6, 8, 10 teacher layers to guide the 1 to 5 student layers, respectively. how-ever, these one-to-one layer mapping strategies are assigned based on empirical observations without theoretical guidance. second, as revealed in @xcite , different bert layers could learn different levels of linguistic knowledge. the oneto-one layer mapping strategy cannot learn an optimal, unified compressed model for different nlp tasks. in addition, most previous works do not consider the importance of each teacher layer and use the same layer weights among various tasks, which create a substantial barrier for generalizing the compressed model to different nlp tasks. therefore, an adaptive compression model should be designed to transfer knowledge from all teacher layers dynamically and effectively for different nlp tasks. to address the aforementioned issues, we propose a novel bert compression approach based on many-to-many layer mapping and earth mover's distance (emd) @xcite , called bert-emd. first, we design a many-to-many layer mapping strategy, where each intermediate student layer has the chance to learn from all the intermediate teacher layers. in this way, bert-emd can learn from different intermediate teacher layers adaptively for different nlp tasks, motivated by the intuition that different nlp tasks require different levels of linguistic knowledge contained in the intermediate layers of bert. second, to learn an optimal many-to-many layer mapping strategy, we leverage emd to compute the minimum cumulative cost that must be paid to transform knowledge from teacher network to student network. emd is a well-studied optimization problem and provides a suitable solution to transfer knowledge from the teacher network in a holistic fashion. we summarize our main contributions as follows. ( 1 ) we propose a novel many-to-many layer mapping strategy for compressing the intermediate layers of bert in an adaptive and holistic fashion. ## related work language models pre-trained on large-scale corpora can learn universal language representations, which have proven to be effective in many nlp tasks @xcite @xcite . early efforts mainly focus on learning good word embeddings, such as word2vec @xcite and glove @xcite . although these pre-trained embeddings can capture semantic meanings of words, they are context-free and fail to capture higher-level concepts in context, such as syntactic structures and polysemous disambiguation. subsequently, researchers have shifted attention to contextual word embeddings learning, such as elmo @xcite , ulmfit @xcite , gpt @xcite , bert @xcite , enrie @xcite , xl-net @xcite , roberta @xcite . for example, @xcite released the bert-base of 110 million parameters and bert-large of 330 million parameters, which achieved significantly better results than previous methods on glue tasks. however, along with high-performance, the pretrained language models (e.g., bert) usually have a large number of parameters, which require a high cost of computation and memory in inference. recently, many attempts have been made to reduce the computation overhead and model storage of pre-trained language models without performance sacrifice. existing compression techniques can be divided into three categories: low-rank matrix factorization @xcite , quantization @xcite , and knowledge distillation @xcite . next, we mainly review the related works that use knowledge distillation to compress the bert model. knowledge distillation using the teacher-student strategy learns a lightweight student network under the guidance of a large and complicated teacher network. @xcite awadallah (2019 ) distilled bert into an lstm network via both hard and soft distilling methods. sun et al. (2019) proposed the bert-pkd model to transfer the knowledge from both the final layer and the intermediate layers of the teacher network. jiao et al. (2019) proposed the tinybert model, which performed the transformer distillation at both pre-training and fine-tuning processes. xu et al. (2020) proposed the bert-of-theseus model to learn a compact student network by replacing the teacher layers with their substitutes. sun et al. (2020) introduced the mobilebert model, which has the same number of layers with the teacher network, but was much narrower via adopting bottleneck structures. wang et al. (2020) distilled the self-attention module of the last transformer layer of the teacher network. however, the aforementioned bert compression approaches struggle to find an optimal layer mapping between the teacher and student networks. each student layer merely learns from a single teacher layer, which may lose rich linguistic knowledge contained in the teacher network. different from previous methods, we propose a many-tomany layer mapping method for bert distillation, where each intermediate student layer can learn from any intermediate teacher layers adaptively. in addition, an earth mover's disepstance is applied to learn the optimal many-to-many layer mapping solution. ## methodology in this section, we propose a novel bert compression method based on many-to-many layer mapping and earth mover's distance (called bert-emd). in addition, we also propose a cost attention mechanism to learn the layer weights used in emd automatically. ## conclusion in this paper, we propose a novel bert compression method based on many-to-many layer mapping by earth mover's distance (emd). to our knowledge, bert-emd is the first work that allows each intermediate student layer to learn from any intermediate teacher layers adaptively. in addition, a cost attention mechanism is designed to further improve the model's performance and accelerate convergence time by learning the layer weights used in emd automatically. extensive experiments on glue tasks show that bert-emd can achieve competitive performances with the large bert-base model while significantly reducing the model size and inference time.