Short-term memory in neural language models

1 INTRODUCTION . Language models ( LMs ) are computational systems trained to predict upcoming tokens based on past context . To perform this task well , they must construct a coherent representation of the text , which requires establishing relationships between words that occur at non-adjacent time points . Despite their simple learning objective , LMs based on contemporary artificial neural network architectures perform well on language processing tasks that require maintenance and retrieval of context spanning multiple words in a sentence . For example , LMs learn to correctly match the grammatical number of the subject and a corresponding verb across multiple intervening words ; for example , they prefer the correct The girls standing at the desk are tall , to the incorrect The girls standing at the desk is tall ( Linzen et al. , 2016 ; Marvin & Linzen , 2018 ; Gulordava et al. , 2018 ; Futrell et al. , 2018 ) . Moreover , language representations learned by large-scale LMs can be fine-tuned or directly used Memory for language in neural LMs for solving downstream natural language processing ( NLP ) tasks ( Devlin et al. , 2019 ; Brown et al. , 2020 ) . Together , these results suggest that LMs extract linguistically meaningful signals and that , over the course of learning , they develop a short-term memory capacity : the ability to store and access recent past context for processing . But what is the nature of the memory processes that LMs learn ? Are these memory processes able to access individual tokens from the recent past verbatim , or is the memory system more implicit , so that only an aggregate gist of the prior context is available to subsequent processing ? The ability to access arbitrary individual tokens from the past could be computationally powerful , akin to the “ working memory ” systems thought to enable flexible human cognitive capacities ( Baddeley , 2003 ) . In this work , we introduce a paradigm for characterizing the short-term memory abilities of a language model . We apply it to two particular architectures of neural language models that possess the architectural ingredients to hold past items in memory : attention-based transformers ( Vaswani et al. , 2017 ) and long short-term memory networks ( Hochreiter & Schmidhuber , 1997 , LSTM ) . Whereas an LSTM incorporates the past by reusing the results of processing from previous time steps and through dedicated memory cells , the transformer receives the internal representations of each of the previous tokens as input . These architectural ingredients alone , however , are not sufficient for a model to have memory . Whether or not the model puts this memory capacity to use depends on whether the training task ( next word prediction ) requires it — the parameters controlling the activation of context representations and subsequent retrieval computations are in both cases learned . Our goal was to determine whether and when the LMs we study maintain and retrieve verbatim representations of individual prior tokens . First , we measured the detail of the context representation : does the LM maintain a verbatim representation of all prior tokens and their order , or does it instead combine multiple prior tokens into a summary representation , like a semantic gist ? Second , we considered the resilience of the memory to interference : after how many intervening tokens do the representation of prior context become inaccessible ? Third , we considered the content-invariance of the context representations : does the resilience of prior context depend on semantic coherence of the prior information , or can arbitrary and unrelated information sequences be retrieved ? Our paradigm , which is inspired by benchmark tasks for models of human short-term memory ( Oberauer et al. , 2018 ) , reveals whether or not and under what conditions the LM can encode and retrieve lists of nouns embedded in vignettes ( see Fig . 1 ) . By comparing LM surprisal on repeated lists of nouns , we found that : • The transformers we tested can perform verbatim recall of past nouns and are robust to interference . The transformers retrieved precise representation of nouns ( word identity and word order ) from the context and did so over long-range intervening text ( > 400 tokens ) . • The LSTM we tested stores a semantic gist of past nouns over short distances . The LSTM ’ s noun list retrieval capacity was limited to a semantic summary of the earlier nouns in the list , did not preserve the order or identity of the original nouns , and only succeeded across a small number of intervening tokens . • Verbatim recall in transformers improves with training corpus size . The transformer trained on a larger dataset showed much larger effects compared to transformer trained on smaller corpus . • Verbatim recall in transformers requires depth . Memory retrieval capacity emerged with 6- and 12- layer models , but not in smaller 1- and 3-layer models . • Order-sensitivity of transformers relies on learned attention patterns . After randomly permuting attention weights , the transformer no longer showed ordering retrieval capacity . 2 RELATED WORK . Previous studies examined how properties of linguistic context influenced average LM performance . Khandelwal et al . ( 2018 ) showed that LSTM LMs use a window of approximately 200 tokens of past context and word order information of the past 50 words , in the service of predicting the next token in natural language sequences . Subramanian et al . ( 2020 ) applied a similar analysis to a transformer LM and showed that LM loss on test-set sequences was not sensitive to context perturbations beyond 50 tokens . These studies focused on how specific features of past context affect LM performance on novel input in the test set . In contrast , our paradigm tests for the ability of LMs to exactly retrieve nouns that are repeated from prior context . 3 METHODS . 3.1 PARADIGM : LISTS OF NOUNS IN CONTEXT . Noun lists were embedded in brief vignettes ( Figure 1 , A and B ) . Each vignette opened with a preface string ( e.g . “ Before the meeting , Mary wrote down the following list of words : ” ) . This string was followed by a list of nouns ( the first list ) , which were separated by commas ; the list-final noun was followed by a full stop ( e.g . “ county , muscle , vapor. ” ) . The first list was followed by an intervening text , which continued the narrative established by the preface string ( “ After the meeting , she took a break and had a cup of coffee. ” ) . The intervening text was followed by a short prompt string ( e.g . “ After she got back , she read the list again : ” ) after which another list of nouns , either identical to the first list or different from it , was presented ( we refer to this list as the second list ) . 3.2 SEMANTIC COHERENCE OF NOUN LISTS . We used two types of word lists : arbitrary and semantically coherent ( all word lists are shown in Tables 1 and 2 in the Appendix ) . Arbitrary word lists ( e.g . “ device , singer , picture ” ) were composed of randomly sampled nouns from the Toronto word pool.1 Semantically coherent word lists were sampled from the categorized noun word pool,2 which contains 32 lists , each of which contains 32 semantically related nouns ( e.g . “ robin , sparrow , heron , ... ” ) . We ensured that the nouns in each list were part of the LSTM ’ s vocabulary . After ensuring there were at least 10 valid , in-vocabulary nouns per semantic set ( as this was the maximal list length we considered ) , we were able to construct 23 nouns lists . Finally , to reduce the variance attributable to tokens occurring in specific positions , we generated 10 “ folds ” of each list by circularly shifting the 1http : //memory.psych.upenn.edu/files/wordpools/nouns.txt 2http : //memory.psych.upenn.edu/files/wordpools/catwpool.txt tokens in the first list 10 times . In this way , each noun in each list was tested in all possible ordinal positions . This procedure resulted in a total of 23× 10 = 230 noun lists . 3.3 LANGUAGE MODELS . LSTM We used an LSTM released by van Schijndel et al . ( 2019 ) ,3 which had 2 hidden layers with 400-dimensional hidden states each and had a vocabulary of 28,439 entries . The model contained 25.3 million trainable parameters . It was trained on a 40-million word subset of the Wikitext-103 corpus ( Merity et al. , 2016 ) and had a test-set perplexity of 79.6 . The model ’ s vocabulary contained 28,439 entries . We selected this checkpoint as our starting point because van Schijndel et al . ( 2019 ) showed that increasing the number of model parameters and size of training set beyond this point did not significantly improve the model ’ s performance on the subject-verb agreement task . To test for the potential role of model/training set size on our task , we also evaluated a larger , 2-layer LSTM with 1600-dimensional hidden states ( 132 million trainable parameters ) , trained by van Schijndel et al . ( 2019 ) on 80 million tokens which had a test-set perplexity of 108.5 . Transformer We trained a transformer LM on ∼ 40 million subword tokens of the Wikitext-103 benchmark ( i.e . comparable to the corpus size used for training the LSTM checkpoint ) . We ensured that the vocabulary size was matched to the size of the LSTM vocabulary ( 28,439 entries ) by retraining the tokenizer ( see Appendix section A.4.2 ) . We trained both the 12-layer GPT-2 architecture ( known as “ GPT-2 small ” , 107.7 million trainable parameters ) and , as a point of comparison , smaller , 1- , 3- , and 6-layer transformers ( 29.7 , 43.9 , and 65.2 million trainable parameters , respectively ) . The context window was set to 1024 tokens and embedding dimension was kept at 768 across the architectures ( see Section A.4 for training details ) . The perplexities for the 12- , 6- , 3- and 1-layer models on the Wikitext-103 test set were 40.3 , 46.7 , 60.1 , and 93.2 , respectively . Unless indicated otherwise , the label “ Wikitext-103 transformer ” refers to the 12-layer checkpoint in the figures . We also evaluated the pretrained transformer LM released by Radford et al . ( 2019 ) , which we refer to simply as GPT-2 . This model was trained on the WebText corpus , which consists of approximately 8 million online documents . We used the GPT-2-small checkpoint which has 12 attention layers and 768-dimensional embedding layer . The model contains 124 million parameters and has a vocabulary of 50,257 entries . We used the maximum context size of 1024 tokens . 3.4 SURPRISAL . For each token wt in our sequence , we computed the negative log likelihood ( surprisal ) : surprisal ( wt ) = − log2 P ( wt|w1 , . . . , wt−1 ) . In cases when the transformer byte-pair encoding tokenizer split a noun into multiple tokens—e.g . “ sparrow ” might be split into “ sp ” and “ arrow ” —we summed the surprisals of the resulting tokens . Quantifying retrieval : repeat surprisal To quantify how the memory trace of the first list affected surprisal on the second list , we measured the ratio between the surprisal on the second list and the surprisal on the first list : repeat surprisal = s̄ ( L2 ) s̄ ( L1 ) ×100 where s̄ ( L1 ) refers to mean surprisal across non-initial nouns in the first list and s̄ ( L2 ) to median surprisal across all non-initial nouns in the second list . We take a reduction in surprisal on second lists to indicate the extent to which an LM has retrieved tokens from the first list .

Summary: This paper studies how LSTMs and Transformers represent prior context. In particular, this work adapts benchmark tasks for human working memory to neural language models. In this paradigm, the model is presented with <preface text> <a list of words> <intervening text> <the same list of words>, and perplexity / surprisal is measured for each of the words when the list of words appears a second time at the end of the input. Their results lead to a variety of insights about how LSTMs and Transformers trained on varying amounts of data with varying amounts of depth handle context.