Taking Notes on the Fly Helps Language Pre-Training

1 INTRODUCTION . Unsupervised language pre-training , e.g. , BERT ( Devlin et al. , 2018 ) , is shown to be a successful way to improve the performance of various NLP downstream tasks . However , as the pre-training task requires no human labeling effort , a massive scale of training corpus from the Web can be used to train models with billions of parameters ( Raffel et al. , 2019 ) , making the pre-training computationally expensive . As an illustration , training a BERT-base model on Wikipedia corpus requires more than five days on 16 NVIDIA Tesla V100 GPUs . Therefore , how to make language pre-training more efficient and less resource-intensive , has become an important research direction in the field ( Strubell et al. , 2019 ) . Our work aims at improving the efficiency of language pre-training methods . In particular , we study how to speed up pre-training through better data utilization . It is well-known that in a natural language data corpus , words follow a heavy-tail distribution ( Larson , 2010 ) . A large proportion of words appear only very few times and the embeddings of those ( rare ) words are usually poorly optimized and noisy ( Bahdanau et al. , 2017 ; Gong et al. , 2018 ; Khassanov et al. , 2019 ; Schick & Schütze , 2020 ) . ∗Equal Contribution . Work done during internships at Microsoft Research Asia . †Correspondence to : dihe @ microsoft.com Unlike previous works that sought to merely improve the embedding quality of rare words , we argue that the existence of rare words could also slow down the training process of other model parameters . Taking BERT as an example , if we imagine the model encounters the following masked sentence during pre-training : COVID-19 has cost thousands of lives . Note that “ COVID-19 ” is a rare word , while also the only key information for the model to rely on to fill in the blank with the correct answer “ lives ” . As the embedding of the rare word “ COVID-19 ” is poorly trained , the Transformer lacks concrete input signal to predict “ lives ” . Furthermore , with noisy inputs , the model needs to take longer time to converge and sometimes even can not generalize well ( Zhang et al. , 2016 ) . Empirically , we observe that around 20 % of the sentences in the corpus contain at least one rare word . Moreover , since most pre-training methods concatenate adjacent multiple sentences to form one input sample , empirically we find that more than 90 % of input samples contain at least one rare word . The large proportion of such sentences could cause severe data utilization problem for language pre-training due to the lack of concrete semantics for sentence understanding . Therefore , learning from the masked language modeling tasks using these noisy embeddings may make the pre-training inefficient . Moreover , completely removing those sentences with rare words is not an applicable choice either since it will significantly reduce the size of the training data and hurt the final model performance . Our method to solve this problem is inspired by how humans manage information . Note-taking is a useful skill which can help people recall information that would otherwise be lost , especially for new concepts during learning ( Makany et al. , 2009 ) . If people take notes when facing a rare word that they don ’ t know , then next time when the rare word appears , they can refer to the notes to better understand the sentence . For example , we may meet the following sentence somewhere beforehand : The COVID-19 pandemic is an ongoing global crisis . From the sentence , we can realize that “ COVID-19 ” is related to “ pandemic ” and “ global crisis ” and record the connection in the notes . When facing “ COVID-19 ” again in the masked-language-modeling task above , we can refer to the note of “ COVID-19 ” . It is easy to see that once “ pandemic ” and “ global crisis ” are connected to “ COVID-19 ” , we can understand the sentence and predict “ lives ” more easily , as illustrated in Figure 1 . Mapped back to language pre-training , we believe for rare words , explicitly leveraging cross-sentence information is helpful to enhance semantics of the rare words in the current sentence to predict the masked tokens . Through this more efficient data utilization , the Transformer can receive better input signals which leads to more efficient training of its model parameters . Motivated by the discussion above , we propose a new learning approach called “ Taking Notes on the Fly ” ( TNF ) to improve data utilization for language pre-training . Specifically , we maintain a note dictionary , where the keys are rare words and the values are historical contextual representations of them . In the forward pass , when a rare word w appears in a sentence , we query the value of w in the note dictionary and use it as a part of the input . In this way , the semantic information of w saved in the note can be encoded together with other words through the model . Besides updating the model parameters , we also update the note dictionary . In particular , we define the note of w in the current sentence as the mean pooling over the contextual representations of the words nearby w. Then we update w ’ s value in the note dictionary by a weighted linear combination of w ’ s previous value and w ’ s note in the current sentence . TNF introduces little computational overhead at pre-training since the note dictionary is updated on the fly during the forward pass . Furthermore , different from the memory-augmented neural networks ( Santoro et al. , 2016 ; Guu et al. , 2020 ) , the note dictionary is only used to improve the training efficiency of the model parameters , while not served as a part of the model . When the pre-training is finished , we discard the note dictionary and use the trained Transformer encoder during the fine-tuning of downstream tasks , same as all previous works . We conduct experiments using BERT and ELECTRA ( Clark et al. , 2019 ) as TNF ’ s backbone methods . Results show that TNF significantly expedites BERT and ELECTRA , and improves their performances on downstream tasks . BERT-TNF and ELECTRA-TNF ’ s training times are both 60 % less than their corresponding backbone models when reaching the same performance . When trained with the same number of iterations , BERT-TNF and ELECTRA-TNF outperform the backbone methods on both the average GLUE score and the majority of individual tasks . We also observe that even in the downstream tasks where rare words only take a neglectable proportion of the data ( i.e . 0.47 % ) , TNF also outperforms baseline methods with a large margin . It indicates that TNF improves the pre-training of the entire model . 2 RELATED WORK . Efficient BERT pre-training . The massive energy cost of language pre-training ( Strubell et al. , 2019 ) has become an obstacle to its further developments . There are several works aiming at reducing the energy cost of pre-training . Gong et al . ( 2019 ) observes that parameters in different layers have similar attention distribution , and propose a parameter distillation method from shallow layers to deep layers . Another notable work is ELECTRA ( Clark et al. , 2019 ) , which develops a new task using one discriminator and one generator . The generator corrupts the sentence , and the discriminator is trained to predict whether each word in the corrupted sentence is replaced or not . Orthogonal to them , we focus on improving pre-training efficiency by finding ways to utilize the data corpus better . Therefore , it can be applied to all of the methods above to further boost their performances . Representation of rare words . It is widely acknowledged that the quality of rare words ’ embeddings is significantly worse than that of popular words . Gao et al . ( 2019 ) provides a theoretical understanding of this problem , which illustrates that the problem lies in the sparse ( and inaccurate ) stochastic optimization of neural networks . Several works attempt to improve the representation of rare words using linguistic priors ( Luong et al. , 2013 ; El-Kishky et al. , 2019 ; Kim et al. , 2016 ; Santos & Zadrozny , 2014 ) . But the improved embedding quality is still far behind that of popular words ( Gong et al. , 2018 ) . Sennrich et al . ( 2015 ) develops a novel way to split each word into sub-word units . However , the embeddings of low-frequency sub-word units are still difficult to train ( Ott et al. , 2018 ) . Due to the poor quality of rare word representations , the pre-training model built on top of it suffers from noisy input semantic signals which lead to inefficient training . We try to bypass the problem of poor rare word representations by leveraging cross-sentence information to enhance input semantic signals of the current sentence for better model training . Memory-augmented BERT . Another line of work close to ours uses memory-augmented neural networks in language-related tasks . Févry et al . ( 2020 ) and Guu et al . ( 2020 ) define the memory buffer as an external knowledge base of entities for better open domain question answering tasks . Khandelwal et al . ( 2019 ) constructs the memory for every test context at inference , to hold extra token candidates for better language modeling . Similar to other memory-augmented neural networks , the memory buffer in these works is a model component that will be used during inference . Although sharing general methodological concepts with these works , the goal and details of our method are different from them . Especially , our note dictionary is only maintained in pre-training for efficient data utilization . At fine-tuning , we ditch the note dictionary , hence adding no extra time or space complexity to the backbone models . 3 TAKING NOTES ON THE FLY . 3.1 PRELIMINARY . In this section , we use the BERT model as an example to introduce the basics of the model architecture and training objective of language pre-training . BERT ( Bidirectional Encoder Representation from Transformers ) is developed on a multi-layer bidirectional Transformer encoder , which takes a sequence of word semantic information ( token embeddings ) and order information ( positional embeddings ) as input , and outputs the contextual representations of words . Each Transformer layer is formed by a self-attention sub-layer and a position-wise feed-forward sub-layer , with a residual connection ( He et al. , 2016 ) and layer normalization ( Ba et al. , 2016 ) applied after every sub-layer . The self-attention sub-layer is referred to as ” Scaled Dot-Product Attention ” in Vaswani et al . ( 2017 ) , which produces its output by calculating the scaled dot products of queries and keys as the coefficients of the values , i.e. , Attention ( Q , K , V ) = Softmax ( QKT√ d ) V. ( 1 ) Q ( Query ) , K ( Key ) , V ( Value ) are the hidden representations outputted from the previous layer and d is the dimension of the hidden representations . Transformer also extends the aforementioned self-attention layer to a multi-head version in order to jointly attend to information from different representation subspaces . The multi-head self-attention sub-layer works as follows , Multi-head ( Q , K , V ) =Concat ( head1 , · · · , headH ) WO ( 2 ) headk =Attention ( QW Q k , KW K k , V W V k ) , ( 3 ) where WQk ∈ Rd×dK , WKk ∈ Rd×dK , WVk ∈ Rd×dV are projection matrices . H is the number of heads . dK and dV are the dimensions of the key and value separately . Following the self-attention sub-layer , there is a position-wise feed-forward ( FFN ) sub-layer , which is a fully connected network applied to every position identically and separately . The FFN sub-layer is usually a two-layer feed-forward network with a ReLU activation function in between . Given vectors { h1 , . . . , hn } , a position-wise FFN sub-layer transforms each hi as FFN ( hi ) = σ ( hiW1+b1 ) W2+b2 , where W1 , W2 , b1 and b2 are parameters . BERT uses the Transformer model as its backbone neural network architecture and trains the model parameters with the masked language model task on large text corpora . In the masked language model task , given a sampled sentence from the corpora , 15 % of the positions in the sentence are randomly selected . The selected positions will be either replaced by special token [ MASK ] , replaced by randomly picked tokens or remain the same . The objective of BERT pre-training is to predict words at the masked positions correctly given the masked sentences . As this task requires no human labeling effort , large scale data corpus is usually used to train the model . Empirically , the trained model , served as a good initialization , significantly improves the performance of downstream tasks .

Summary: The paper proposes an external memory architecture. When encountering the rare words (with a frequency between 100-500), the method will store the average contextualized word embedding of nearby words into a dictionary. Next time it encounters the same rare word, it will retrieve the average embedding and input it into BERT encoder. The experiment results show that given the same number of training steps, adding the external memory improves the MLM loss and significantly improves the results on RTE (Recognizing Textual Entailment) dataset, which leads to a slightly better GLUE score. The experiment also shows that keeping the external memory during the fine-tuning stage slightly degrades the performance.