Rethinking Embedding Coupling in Pre-trained Language Models

1 INTRODUCTION . The performance of models in natural language processing ( NLP ) has dramatically improved in recent years , mainly driven by advances in transfer learning from large amounts of unlabeled data ( Howard & Ruder , 2018 ; Devlin et al. , 2019 ) . The most successful paradigm consists of pre-training a large Transformer ( Vaswani et al. , 2017 ) model with a self-supervised loss and fine-tuning it on data of a downstream task ( Ruder et al. , 2019 ) . Despite its empirical success , inefficiencies have been observed related to the training duration ( Liu et al. , 2019b ) , pre-training objective ( Clark et al. , 2020b ) , and training data ( Conneau et al. , 2020a ) , among others . In this paper , we reconsider a modeling assumption that may have a similarly pervasive practical impact : the coupling of input and output embeddings1 in state-of-the-art pre-trained language models . State-of-the-art pre-trained language models ( Devlin et al. , 2019 ; Liu et al. , 2019b ) and their multilingual counterparts ( Devlin et al. , 2019 ; Conneau et al. , 2020a ) have inherited the practice of embedding coupling from their language model predecessors ( Press & Wolf , 2017 ; Inan et al. , 2017 ) . However , in contrast to their language model counterparts , embedding coupling in encoder-only pre-trained models such as Devlin et al . ( 2019 ) is only useful during pre-training since output embeddings are generally discarded after fine-tuning.2 In addition , given the willingness of researchers to exchange additional compute during pre-training for improved downstream performance ( Raffel ∗equal contribution †Work done as a member of the Google AI Residency Program . 1Output embedding is sometimes referred to as “ output weights ” , i.e. , the weight matrix in the output projection in a language model . 2We focus on encoder-only models , and do not consider encoder-decoder models like T5 ( Raffel et al. , 2020 ) where none of the embedding matrices are discarded after pre-training . Output embeddings may also be et al. , 2020 ; Brown et al. , 2020 ) and the fact that pre-trained models are often used for inference millions of times ( Wolf et al. , 2019 ) , pre-training-specific parameter savings are less important overall . On the other hand , tying input and output embeddings constrains the model to use the same dimensionality for both embeddings . This restriction limits the researcher ’ s flexibility in parameterizing the model and can lead to allocating too much capacity to the input embeddings , which may be wasteful . This is a problem particularly for multilingual models , which require large vocabularies with high-dimensional embeddings that make up between 47–71 % of the entire parameter budget ( Table 1 ) , suggesting an inefficient parameter allocation . In this paper , we systematically study the impact of embedding coupling on state-of-the-art pretrained language models , focusing on multilingual models . First , we observe that while naı̈vely decoupling the input and output embedding parameters does not consistently improve downstream evaluation metrics , decoupling their shapes comes with a host of benefits . In particular , it allows us to independently modify the input and output embedding dimensions . We show that the input embedding dimension can be safely reduced without affecting downstream performance . Since the output embedding is discarded after pre-training , we can increase its dimension , which improves fine-tuning accuracy and outperforms other capacity expansion strategies . By reinvesting saved parameters to the width and depth of the Transformer layers , we furthermore achieve significantly improved performance over a strong mBERT ( Devlin et al. , 2019 ) baseline on multilingual tasks from the XTREME benchmark ( Hu et al. , 2020 ) . Finally , we combine our techniques in a Rebalanced mBERT ( RemBERT ) model that outperforms XLM-R ( Conneau et al. , 2020a ) , the state-of-the-art cross-lingual model while having been pre-trained on 3.5× fewer tokens and 10 more languages . We thoroughly investigate reasons for the benefits of embedding decoupling . We observe that an increased output embedding size enables a model to improve on the pre-training task , which correlates with downstream performance . We also find that it leads to Transformers that are more transferable across tasks and languages—particularly for the upper-most layers . Overall , larger output embeddings prevent the model ’ s last layers from over-specializing to the pre-training task ( Zhang et al. , 2020 ; Tamkin et al. , 2020 ) , which enables training of more general Transformer models . 2 RELATED WORK . Embedding coupling Sharing input and output embeddings in neural language models was proposed to improve perplexity and motivated based on embedding similarity ( Press & Wolf , 2017 ) as well as by theoretically showing that the output probability space can be constrained to a subspace governed by the embedding matrix for a restricted case ( Inan et al. , 2017 ) . Embedding coupling is also common in neural machine translation models where it reduces model complexity ( Firat et al. , 2016 ) and saves memory ( Johnson et al. , 2017 ) , in recent state-of-the-art language models ( Melis et al. , 2020 ) , as well as all pre-trained models we are aware of ( Devlin et al. , 2019 ; Liu et al. , 2019b ) . Transferability of representations Representations of large pre-trained models in computer vision and NLP have been observed to transition from general to task-specific from the first to the useful for domain-adaptive pre-training ( Howard & Ruder , 2018 ; Gururangan et al. , 2020 ) , probing ( Elazar & Goldberg , 2019 ) , and tasks that can be cast in the pre-training objective ( Amrami & Goldberg , 2019 ) . last layer ( Yosinski et al. , 2014 ; Howard & Ruder , 2018 ; Liu et al. , 2019a ) . In Transformer models , the last few layers have been shown to become specialized to the MLM task and—as a result—less transferable ( Zhang et al. , 2020 ; Tamkin et al. , 2020 ) . Multilingual models Recent multilingual models are pre-trained on data covering around 100 languages using a subword vocabulary shared across all languages ( Devlin et al. , 2019 ; Pires et al. , 2019 ; Conneau et al. , 2020a ) . In order to achieve reasonable performance for most languages , these models need to allocate sufficient capacity for each language , known as the curse of multilinguality ( Conneau et al. , 2020a ; Pfeiffer et al. , 2020 ) . As a result , such multilingual models have large vocabularies with large embedding sizes to ensure that tokens in all languages are adequately represented . Efficient models Most work on more efficient pre-trained models focuses on pruning or distillation ( Hinton et al. , 2015 ) . Pruning approaches remove parts of the model , typically attention heads ( Michel et al. , 2019 ; Voita et al. , 2019 ) while distillation approaches distill a large pre-trained model into a smaller one ( Sun et al. , 2020 ) . Distillation can be seen as an alternative form of allocating pre-training capacity via a large teacher model . However , distilling a pre-trained model is expensive ( Sanh et al. , 2019 ) and requires overcoming architecture differences and balancing training data and loss terms ( Mukherjee & Awadallah , 2020 ) . Our proposed methods are simpler and complementary to distillation as they can improve the pre-training of compact student models ( Turc et al. , 2019 ) . 3 EXPERIMENTAL METHODOLOGY . Efficiency of models has been measured along different dimensions , from the number of floating point operations ( Schwartz et al. , 2019 ) to their runtime ( Zhou et al. , 2020 ) . We follow previous work ( Sun et al. , 2020 ) and compare models in terms of their number of parameters during finetuning ( see Appendix A.1 for further justification of this setting ) . For completeness , we generally report the number of pre-training ( PT ) and fine-tuning ( FT ) parameters . Baseline Our baseline has the same architecture as multilingual BERT ( mBERT ; Devlin et al. , 2019 ) . It consists of 12 Transformer layers with a hidden size H of 768 . Input and output embeddings are coupled and have the same dimensionality E as the hidden size , i.e . Eout = Ein = H . The total number of parameters during pre-training and fine-tuning is 177M ( see Appendix A.2 for further details ) . We train variants of this model that differ in certain hyper-parameters but otherwise are trained under the same conditions to ensure a fair comparison . Tasks For our experiments , we employ tasks from the XTREME benchmark ( Hu et al. , 2020 ) that require fine-tuning , including the XNLI ( Conneau et al. , 2018 ) , NER ( Pan et al. , 2017 ) , PAWS-X ( Yang et al. , 2019 ) , XQuAD ( Artetxe et al. , 2020 ) , MLQA ( Lewis et al. , 2020 ) , and TyDiQA-GoldP ( Clark et al. , 2020a ) datasets . We provide details for them in Appendix A.4 . We average results across three fine-tuning runs and evaluate on the dev sets unless otherwise stated . 4 EMBEDDING DECOUPLING REVISITED . Naı̈ve decoupling Embeddings make up a large fraction of the parameter budget in state-of-theart multilingual models ( see Table 1 ) . We now study the effect of embedding decoupling on such models . In Table 2 , we show the impact of decoupling the input and output embeddings in our baseline model ( §3 ) with coupled embeddings . Naı̈vely decoupling the output embedding matrix slightly improves the performance as evidenced by a 0.4 increase on average . However , the gain is not uniformly observed in all tasks . Overall , these results suggest that decoupling the embedding matrices naı̈vely while keeping the dimensionality fixed does not greatly affect the performance of the model . What is more important , however , is that decoupling the input and output embeddings decouples the shapes , endowing significant modeling flexibility , which we investigate in the following . Input vs output embeddings Decoupling input and output embeddings allows us to flexibly change the dimensionality of both matrices and to determine which one is more important for good transfer performance of the model . To this end , we compare the performance of a model with Ein = 768 , Eout = 128 to that of a model with Ein = 128 , Eout = 7683 ( the remaining hyperparameters are the same as the baseline in §3 ) . During fine-tuning , the latter model has 43 % fewer parameters . We show the results in Table 3 . Surprisingly , the model pre-trained with a larger output embedding size is competitive with the comparison method on average despite having 77M fewer parameters during fine-tuning.4 Reducing the input embedding dimension saves a significant number of parameters at a noticeably smaller cost to accuracy than reducing the output embedding size . In light of this , the parameter allocation of multilingual models ( see Table 1 ) seems particularly inefficient . For a multilingual model with coupled embeddings , reducing the input embedding dimension to save parameters as proposed by Lan et al . ( 2020 ) is very detrimental to performance ( see Appendix A.5 for details ) . The results in this section indicate that the output embedding plays an important role in the transferability of pre-trained representations . For multilingual models in particular , a small input embedding dimension frees up a significant number of parameters at a small cost to performance . In the next section , we study how to improve the performance of a model by resizing embeddings and layers .

Summary: This work studied the impact of embedding coupling in pre-trained language models, by taking a multilingual model as backbone. The major finding is that decoupling the input and output embedding shapes can bring benefits, and the output embedding plays an important role in the transferability of pre-trained representation. A rebalanced mBERT is designed by combing and scaling up the investigated techniques, achieving strong results on the XTREME benchmark.