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: The authors of this paper try to distill large pretrained language models into a bidirectional CMOW/CBOW-Hybrid model. The proposed architecture was designed to output per-token representations, so as to perform knowledge distillation on the masked language modeling pretraining task. A two-sequence encoding scheme (DiffCat) was designed for the downstream task on sentence pairs. Empirical results show that the proposed method can achieve comparable results to ELMo while using only half of the parameters and providing three times faster inference speed.