Efficient Contextual Representation Learning With Continuous Outputs

Efficient Contextual Representation Learning With Continuous Outputs Liunian Harold Li†, Patrick H. Chen∗, Cho-Jui Hsieh∗, Kai-Wei Chang∗ †Peking University ∗University of California, Los Angeles liliunian@pku.edu.cn, patrickchen@g.ucla.edu {chohsieh, kwchang}@cs.ucla.edu Abstract Contextualrepresentationmodels have achieved great success in improving various downstream natural language processing tasks. However, these language-model-based encoders are dif- ficult to train due to their large parameter size and high computational complexity. By carefully examining the training procedure, we observe that the softmax layer, which pre- dicts a distribution of the target word, often in- duces significant overhead, especially when the vocabulary size is large. Therefore, we re- visit the design of the output layer and consider directly predicting the pre-trained embedding of the target word for a given context. When appliedtoELMo, theproposedapproachachieves a 4-fold speedup and eliminates 80% trainable parameters while achieving competitive per- formance on downstream tasks. Further anal- ysis shows that the approach maintains the speed advantage under various settings, even when the sentence encoder is scaled up. 1 Introduction In recent years, text representation learning ap- proaches, such as ELMo (Peters et al., 2018a), GPT (Radford et al., 2018), BERT (Devlin et al., 2019), and GPT-2 (Radford et al., 2019), have been developed to represent generic contextual information in natural languages by training an encoder with a language model objective on a large unlabelled corpus. During the training process, the encoder is given part of the text and asked to predict the missing pieces. Prior studies show that the encoders trained in this way can capture generic contextual information of the input text and improve a variety of downstream tasks significantly. However, training contextual representations is known to be a resource-hungry process. For example, ELMo is reported to take about 2 weeks to train on a one-billion-token corpus with a vocabulary of 800,000 words using three GPUs.1 This slow training procedure hinders the development cycle, prevents fine-grained param- eter tuning, and makes training contextual repre- sentations inaccessible to the broader community. Recent work also raises concerns about the envi- ronmental implications of training such large models (Strubell et al., 2019). In addition, the suc- cess of these models stems from a large amount of data they used. It is challenging, if not impossible, to train a contextual representation model on a larger corpus with tens or hundreds of billions of tokens. In this work, we explore how to accelerate contextualrepresentation learning. We identify the softmax layer as the primary cause of inefficiency. This component takes up a considerable portion of all trainable parameters (80% for ELMo) and consumes a huge amount of training time. However, it is often not needed in the final model as the goal of contextual representation learning is to build a generic encoder. Therefore, it is rather a waste to allocate extensive computational resources to the softmax layer. InspiredbyKumarand Tsvetkov (2019), wecon- sider learning contextual representation models with continuous outputs. In the training process, the contextual encoder is learned by minimizing the distance between its output and a pre-trained target word embedding. The …

Content truncated to 500 words