Cross-Architecture Distillation Using Bidirectional CMOW Embeddings

1 INTRODUCTION . Large pretrained language models ( Devlin et al. , 2019 ; Raffel et al. , 2020 ) ( PreLMs ) have emerged as de-facto standard methods for natural language processing ( Wang et al. , 2018 ; 2019 ) . The common strategy is to pretrain models on enormous amounts of unlabeled text before fine-tuning them for downstream tasks . However , the drawback of PreLMs is that the models are becoming larger and larger with up to several billions of parameters ( Brown et al. , 2020 ) . This comes with high environmental and economic costs ( Strubell et al. , 2019 ) and shifts the development and research in the hands of a few global players only ( Bommasani et al. , 2021 , pp . 10-12 ) . Even though a single pretrained model can be reused for multiple downstream tasks , the sheer model size is often prohibitive . The immense resource requirements prevent using those models in small-scale laboratories and on mobile devices , which is tied to privacy concerns ( Sanh et al. , 2020b ) . There is a need for more efficient models or compressed versions of large models to make AI research more inclusive and energy-friendly , while fostering deployment in applications . Reducing the size of PreLMs using knowledge distillation ( Hinton et al. , 2015 ) or model compression ( Bucila et al. , 2006 ) is an active area of research ( Sanh et al. , 2020a ; Jiao et al. , 2020 ; Sun et al. , 2020 ) . Both knowledge distillation and model compression can be described as teacher-student setups . The student is trained to imitate the predictions of the teacher while using less resources . Typically , a large PreLM takes the role of the teacher while the student is a smaller version of the same architecture . Sharing the same architecture between student and teacher enables using dedicated distillation techniques , e. g. , aligning the representations of intermediate layers ( Sanh et al. , 2020a ; Sun et al. , 2020 ) . However , using more efficient architectures as student has already shown promising results such as the task-specific distillation approaches by Tang et al . ( 2019 ) and Wasserblat et al . ( 2020 ) . In their works , the student models are LSTMs ( Hochreiter & Schmidhuber , 1997 ) or models based on a continuous bag-of-words representation ( CBOW ) ( Collobert & Weston , 2008 ; Mikolov et al. , 2013 ) . On the one hand , LSTMs are difficult to parallelize as they need at least O ( n ) sequential steps to encode a sequence of length n. On the other hand , CBOW-based models are not order-aware , i. e. , can not distinguish sentences with the same words but in different order ( “ cat eats mouse ” vs. “ mouse eats cat ” are treated equivalent ) . There are , however , more efficient models such as Mai et al . ( 2019 ) ’ s continual multiplication of words ( CMOW ) that do capture word order by representing each token as a matrix , instead of a vector as in CBOW . A sequence in CMOW is modeled by the non-commutative matrix multiplication , which makes the encoding of a sequence becomes dependent on the word order . We denote such models as matrix embeddings . The present work investigates how order-aware matrix embeddings can be used as student models in cross-architecture distillation from large PreLM teachers . Thus , we complement the existing body of works that focused predominantly on same-architecture distillation . All previous cross-architecture approaches are task-specific , whereas we also explore general distillation . We aim to understand to which extent order-aware embeddings are suited to capture the teacher signal of a large PreLM such as BERT . To this end , we extend Mai et al . ( 2019 ) ’ s CMOW/CBOW-Hybrid model , which is a hybrid variant unifying the strength of CBOW and CMOW , with a bidirectional representation of the sequences . Furthermore , we add the ability to emit per-token representations to facilitate using a modern masked language model objective ( Devlin et al. , 2019 ) . We investigate both general distillation , i. e. , the distillation is applied during pretraining on unlabeled text , as well as task-specific distillation , when an already fine-tuned PreLM is distilled on a per-task basis . We further introduce a two-sentence encoding scheme to CMOW such that it can deal with similarity and natural language inferencing tasks . This improves the performance by 32 % compared to a naive joint encoding . Our results show that large PreLMs can be distilled into efficient order-aware embedding models and achieve performance comparable to ELMo ( Peters et al. , 2018 ) on the GLUE benchmark . On certain tasks , embedding-based models even challenge other size-reduced BERT models such as DistillBERT . In summary , our contributions are : • We extend order-aware embedding models with bidirection and make them amenable for masked language model pretraining . • We explore using order-aware embedding models as student models in a cross-architecture distillation setup with BERT as teacher and compare general and task-specific distillation . • We introduce the first encoding scheme that enables CMOW/CBOW-Hybrid to deal with two-sentence tasks ( 32 % increase over the naive approach ) . • Our results show that the best distilled embedding models can be on-par with more expensive models such as ELMo or DistilBERT on certain tasks of the GLUE benchmark , while having a much higher encoding speed ( thrice as high as DistilBERT ) . Below , we introduce our embedding models , our cross-architecture distillation setup , and our twosequence encoding scheme . The experimental procedure is described in Section 3 . The results are reported in Section 4 and discussed in Section 5 , where we also relate our work to the literature . 2 METHODS . First , we introduce our bidirectional extension to the CMOW/CBOW-Hybrid model . Subsequently , we introduce our approach for cross-architecture distillation that we use during the pretraining and fine-tuning stages . Finally , we introduce a two-sentence encoding scheme for order-aware embedding models that is crucial for fine-tuning on downstream tasks with paired sentences . 2.1 EXTENDING THE ORDER-AWARE EMBEDDING MODELS . We extend the CMOW/CBOW-Hybrid embeddings of Mai et al . ( 2019 ) with bidirection and the ability to emit per-token representations as a foundation for cross-architecture distillation . CMOW/CBOWHybrid embeddings , our baseline model , are a combination of matrix embeddings and vector embeddings . Compared to vector-only embeddings , the word order can be captured because matrix multiplication is non-commutative . Given a sequence of n tokens , matrix-space embeddings Xj ∈ Rd×d for each different token j , and vector-space embeddings xj ∈ Rdvec , the CMOW/CBOWHybrid embedding of a sequence of length n is the multiplication of embedding matrices Xi concatenated ( symbol ·||· ) to the summation of the embedding vectors xi : H ( CMOW ) : = X ( CMOW ) 1 ·X ( CMOW ) 2 · · ·X ( CMOW ) n h ( CBOW ) : = ∑ 1≤j≤n x ( CBOW ) j h ( Hybrid ) : = flatten ( H ( CMOW ) ) || h ( CBOW ) where flatten collapses the matrix into a vector . The original work on CMOW ( Mai et al. , 2019 ) has extensively analyzed the CMOW and CBOW components individually and found that joint training is generally preferable . The CMOW and CBOW components can have different dimensionalities , as they are combined via concatentation . We initialize each matrix Xj as identity plus Gaussian noise Id +N ( 0 , σ2init ) with σinit = 0.01 . H ( CMOW ) 4 . Proposed Model : Bidirectional CMOW/CBOW-Hybrid Inspired by the success of bidirection in RNNs ( Schuster & Paliwal , 1997 ) , LSTMs ( Peters et al. , 2018 ) , and Transformers ( Devlin et al. , 2019 ) , we extend CMOW by a bidirectional component . Hence , we introduce a second set of matrixspace embeddings that are multiplied in reverse order . We then have one matrix embedding for the forward direction X ( fw ) ∈ Rnvocab×d×d and one for the backward direction X ( bw ) ∈ Rnvocab×d×d . Then , we concatenate forward and backward directions . Figure 1 illustrates bidirectional CMOW . Furthermore , we emit one representation per token position i , which allows training with a masked language model objective ( Devlin et al. , 2019 ) . Thus , we are able to make use of the BERT teacher signal for pretraining . Since we can reuse computations , O ( n ) matrix multiplications are sufficient to encode a sequence of length n. For these intermediate representations , we also modify the CBOW component in a way that it yields partial sums for the forward and backward directions . Formally , we compute the CMOW/CBOW-Hybrid representation as follows : H ( Bidi.CMOW ) i : = X ( fw ) 1 ·X ( fw ) 2 · · ·X ( fw ) i ||X ( bw ) n ·X ( bw ) n−1 · · ·X ( bw ) i h ( Bidi.CBOW ) i : = i∑ j=1 x ( CBOW ) j || n∑ j=i x ( CBOW ) j h ( Bidi.Hybrid ) i : = flatten ( H ( Bidi.CMOW ) i ) || h ( Bidi.CBOW ) i For fine-tuning on tasks with full sentences as input , e. g. , natural language inferencing , we do not need per-token representations . In this case , we compute the representation of a token sequence as : H ( Bidi.CMOW ) : = X ( fw ) 1 ·X ( fw ) 2 · · ·X ( fw ) n ||X ( bw ) n ·X ( bw ) n−1 · · ·X ( bw ) 1 h ( CBOW ) : = n∑ j=1 x ( CBOW ) j h ( Bidi.Hybrid ) : = flatten ( H ( Bidi.CMOW ) i ) || h ( CBOW ) i Note that the forward and backward directions of the embedding vectors h ( CBOW ) conflate to equivalent formulas when we encode entire sequences . Thus , we only need to include a single CBOW representation along with the two CMOW components that do yield different results for the forward and backward direction . At inference time , the model is parallelizable along the sequential dimension . For regularization , we apply a mild dropout ( p=0.1 ) on both the embeddings and their aggregated representations during pretraining . Then , we feed them into a linear masked language modeling head , see Figure 1 , or an MLP classification head to tackle the downstream tasks . We have also experimented with linear , LSTM , and CNN classifiers , which are described in Appendix A . We chose an MLP because it adds nonlinearity to the model without introducing further complexity . The MLP has led to the best average performance across all GLUE tasks ( we report the results in Appendix C ) . 2.2 CROSS-ARCHITECTURE DISTILLATION . A central question of our research is whether we can distill a large PreLM , e. g. , BERT , into more efficient , non-transformer architectures such as the proposed bidirectional CMOW/CBOW-Hybrid model . This requires a cross-architecture distillation approach , which we describe below . In general , the idea of knowledge distillation is to compress the knowledge of a large teacher model into a smaller student model ( Hinton et al. , 2015 ; Bucila et al. , 2006 ) . It involves a loss function L that is a combination of two loss terms , i. e. , L = α · Lhard + ( 1−α ) · Lsoft with weighting parameter α. Lhard denotes the cross-entropy loss with respect to the ground truth and Lsoft = Σiti · log ( si ) is the cross-entropy between student logits s and the teacher signal t. Optionally , the softmax within Lsoft is flattened by a temperature parameter T . We distinguish general distillation , where BERT ’ s teacher signal is only used during pretraining , and task-specific distillation , where the BERT teacher signal is used during fine-tuning for the downstream ( see related work ) . Considering our goal to design a cross-architecture distillation , the general distillation approach has the conceptual benefit that the teacher model is not needed for fine-tuning . Thus , the student model is capable of tackling downstream tasks without supervision of the large teacher . This has the benefit that one does not need to carry around the BERT model for adaption to every new downstream task . Above , we have introduced the ability to emit per-token representations with lightweight ( bidirectional ) CMOW/CBOW-Hybrid embedding models . This enables us now to use BERT ’ s teacher signal during pretraining together with a masked language modeling objective . With other words , this enables us to perform cross-architecture distillation with matrix embeddings . We consider three variants of cross-architecture distillation in our experiments : a ) When using general distillation , depicted in Figure 1 , BERT acts as a teacher during pretraining and the model is fine-tuned to downstream tasks on its own . b ) For task-specific distillation , BERT acts as a teacher during fine-tuning as shown in Figure 2 . For this case , we have the option of either b1 ) starting with pretrained embeddings ( from general distillation , i. e. , the a ) variant ) , or b2 ) starting from scratch with randomly initialized embeddings .

Summary: This paper presents a knowledge distillation method for BERT, where BERT-base acts as a teacher network and  CMOW is used as a student network. To get a relatively strong student network, the authors extend the vanilla CMOW to bidirectional CMOW. The authors also adopt a two-stage distillation method consisting of general distillation and task-specific distillation to further improve the performances. Finally, the authors evaluate the proposed method on GLUE dataset.