Characterizing Verbatim Short-Term Memory in Neural Language Models (2022)

Abstract: When a language model is trained to predict natural language sequences, its prediction at each moment depends on a representation of prior context. What kind of information about the prior context can language models retrieve? We tested whether language models could retrieve the exact words that occurred previously in a text. In our paradigm, language models (transformers and an LSTM) processed English text in which a list of nouns occurred twice. We operationalized retrieval as the reduction in surprisal from the first to the second list. We found that the transformers retrieved both the identity and ordering of nouns from the first list. Further, the transformers’ retrieval was markedly enhanced when they were trained on a larger corpus and with greater model depth. Lastly, their ability to index prior tokens was dependent on learned attention patterns. In contrast, the LSTM exhibited less precise retrieval, which was limited to list-initial tokens and to short intervening texts. The LSTM’s retrieval was not sensitive to the order of nouns and it improved when the list was semantically coherent. We conclude that transformers implemented something akin to a working memory system that could flexibly retrieve individual token representations across arbitrary delays; conversely, the LSTM maintained a coarser and more rapidly-decaying semantic gist of prior tokens, weighted toward the earliest items.

## 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 in contexts that require maintenance and retrieval of dependencies span- before the meeting, mary wrote down a list of words: county, muscle, vapor. after the meeting, mary took a break and... after she got back, she read the list again: lm input sequence preface string first list intervening text prompt string second list context in memory retrieval change in surprisal paradigm 1) how detailed is lm memory of nouns (identity and ordering)? figure 1: characterizing verbatim memory retrieval in neural language models. in our paradigm, language models processed english text in which a list of nouns occurred twice. we operationalized retrieval as the reduction in surprisal from the first to the second list presentation. we measured retrieval while varying: a) set size, b) the structure of the second list, c) the length of the intervening text, and d) the content and structure of the intervening text. here, we introduce a paradigm (fig. 1 ), inspired by benchmark tasks for models of human shortterm memory @xcite , for characterizing short-term memory abilities of lms. we apply it to two particular neural lm architectures that possess the architectural ingredients to hold past items in memory: attention-based transformers @xcite and long short-term memory networks @xcite . whereas lstms incorporate the past by reusing the results of processing from previous time steps through dedicated memory cells, transformers use 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. we hypothesize that 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 is to determine whether and when the lms we study maintain and retrieve verbatim representations of individual prior tokens. first, we measure 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 consider the resilience of the memory to interference: after how many intervening tokens do the representation of prior context become inaccessible? third, we consider 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? ## related work previous studies examined how properties of linguistic context influenced next-word prediction accuracy in transformer and lstm lms trained on text in english. @xcite showed that lstm lms use a window of approximately 200 tokens of past context and word order informa-tion of the past 50 words, in the service of predicting the next token in natural language sequences. @xcite 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. o'connor and andreas (2021) investigated whether fine-grained lexical and sentential features of context are used for next-word prediction in transformer lms. they showed that transformers rely predominantly on local word co-occurrence statistics (e.g. trigram ordering) and the presence of open class parts of speech (e.g. nouns), and less on the global structure of context (e.g. sentence ordering) and the presence of closed class parts of speech (e.g. function words). in contrast with these studies, which focused on how specific features of past context affect lm performance on novel input at test time, our paradigm tests for the ability of lms to retrieve nouns that are exactly repeated from prior context. in a separate line of work bearing on memory maintenance in lstms, @xcite studied an lstm's capacity to track subject-verb agreement dependencies. they showed that lstm lms relied on a small number of hidden units and the gating mechanisms that control memory contents. here, we are similarly concerned with memory characteristics that support lm performance, but -akin to behavioral tests in cognitive science -we infer the functional properties of lm memory by manipulating properties of repeated noun lists and observing the effects these manipulations have on the behavior (surprisal) of the lm rather than on its internal representation. a third related area of research proposes architectural innovations that augment rnns and lstms with dedicated memory components (e.g. @xcite or improve the handling of context and memory in transformers (see @xcite . here, we are not concerned with improving architectures, but with developing a paradigm that allows us to study how lms put to use their memory systems, whether those are implicit or explicit. ## transformer results we first describe the results of our experiments with the two largest transformer models, the offthe-shelf gpt-2 and the 12-layer transformer we trained; lstm results are discussed in section 5, and results with smaller transformers are discussed towards the end of this section. the transformers retrieved prior nouns and their order; this capacity improved when the model was trained on a larger corpus. we tested whether the transformers could retrieve the identity and order of 10-token noun lists (arbitrary or semantically coherent). to this end, we constructed vignettes in which the second list was either (a) identical to the first list, (b) a permutation of the first list, or (c) a list of novel nouns not present in the first list. 4 we then measured retrieval as reduction in surprisal from first to second list. when the two transformers were presented with second lists that were repeated version of the first ones (blue in fig. 2 , b and c ), token-by-token surprisal decreased compared to novel tokens, suggesting that the transformers were able to access 4 novel nouns in the string were introduced by randomly selecting a list of nouns from one the 22 remaining lists in the noun pool. in semantically coherent lists, novel nouns were drawn from a different semantic category than the nouns in the first list. in order to contextualize the magnitude of these retrieval effects, we computed the relative surprisal across all tokens in lists except the first one (fig. 3 ). when the first and second lists were identical (e.g. with @xmath0, the wikitext-103 transformer's median relative surprisal was at 88% of the first list, compared to 92% for the permuted lists, and 99% for the novel lists. in gpt-2, repeat surprisal was only 2% of the first list, much lower than the 58% for the permuted lists, and 96% of the novel list. retrieval in gpt-2 was robust to the exact phrasing of the text that introduced the lists. replacing the subject 'mary' with 'john' in the vignette, replacing the colon with a comma or randomly permuting the preface or the prompt strings did not affect the results (fig. 7 , right, appendix a). by contrast, the same perturbations reduced retrieval effects for wikitext-103 (fig. 7 , left, appendix a), supporting the conclusion that larger training corpus size contributes to robustness of transformer retrieval. transformer retrieval was robust to the number of items being retrieved. in studies of human short-term memory, performance degrades as the number of items that need to be retained increases ("set-size effects", @xcite . is our lms' short-term memory similarly taxed by increasing the set size? we varied the number of tokens to be held in memory with n @xmath1, 5, 7, 10}. for this comparison, the length of the intervening text was kept at 26 tokens. results reported in fig. 3 show that for gpt-2, verbatim recall was, for the most part, consistent across the different set sizes. repeat surprisal increased monotonically with set size only when the order of nouns in second list, either semantically coherent or arbitrary, was permuted. 5 for the smaller wiktiext-103 transformer, repeat surprisal showed a slight increase with set size further indicating that retrieval robustness increases with training corpus size. transformer retrieval was robust to the length and content of intervening text, but scrambling the intervening text reduced retrieval of order information. for how long are individual items retained in the memory of the lm? we tested this by varying the length of the intervening text for n @xmath2, 99, 194, 435} (see fig. 1 , panel b). to generate longer intervening text samples, we continued the narrative established by the initial preface string ("before the meeting, mary wrote down the following list of words:"). all intervening text strings ended with the same prompt string ("when she got back, she read the list again:") which introduced the second list. memory retrieval in gpt-2 was largely invariant to the size of the intervening text between the first and second lists (fig. 3 , b and c , respectively ). the wikitext-103 transformer exhibited small repeat surprisal increase over intervening text length, suggesting it's memory retrieval was less robust to maintenance over long distances compared to gpt-2. all in all, the results suggest that the two transformers were retrieving prior nouns using a form of direct indexing of the relevant words from the input buffer, rather than implementing a generic memory heuristic, such as predicting that the nouns that have occurred in the most recent 20 tokens will recur. increasing the length of well-formed, semantically coherent intervening text does not, then, interfere with memory retrieval in the transformer. in models of human memory, current context, such as immediately preceding text, can indeed be used as a cue for recalling the encoded items (kahana, 2020). does the transformers' capacity to retrieve copies of past nouns rely on the content and structure of the intervening text? we tested this by creating incongruent and scrambled versions of the longest intervening text (435 tokens). an incongruent condition was created by using intervening text that was syntactically well-formed but semantically incongruent with respect to the preface. the scrambled version was created by randomly permuting the tokens of the intervening text. the transformers' retrieval of past tokens was largely unaffected by the specific content of the intervening text, as long as the intervening text was coherent/well-formed (fig. 4 ). however, in gpt-2, median surprisal across permuted arbitrary lists of nouns increased by 8% when the intervening text was scrambled (fig. 4 , bottom) compared to wellformed text. this suggests that gpt-2 relied on narrative coherence of the intervening text, rather than its aggregate semantic content alone, as a cue for retrieving the ordering information of arbitrary word lists. transformer verbatim recall is learned, guided by attention, and requires increase in size. having shown that the transformer lms could flexibly and robustly retrieve words and their ordering verbatim from short-term memory (figs. 3 and 4 ), we next asked: is this ability learned, or does it derive directly from the architecture? to address this question, we re-ran the experiment with varying number of tokens in lists with a randomly initialized transformer model (architecture as in section 3.3). this random-weights model was unable to retrieve words or their order: for example, repeat surprisal remained at 100% relative to first lists regardless of whether or not the nouns in the second list have appeared before (fig. 8 , top, appendix verbatim retrieval as a function of set size and intervening text w₁ w₂ w� w₁ w₂ w� dragon pillow danger dragon pillow danger set size first list intervening text second list preface 26 99 194 435 intervening text len. (n. tokens) 80 90 100 110 repeat surprisal (%) arbitrary list 26 99 194 435 intervening text len. (n. tokens) semantically coherent list second list repeated permuted novel lstm (merity et al, 2017) 26 99 194 435 intervening text len. (n. tokens) 0 50 100 repeat surprisal (%) arbitrary list 26 99 194 435 intervening text len. (n. tokens) semantically coherent list second list repeated permuted novel transformer (radford et al, 2019) 3 5 7 10 set size (n. tokens) 0 50 100 repeat surprisal (%) arbitrary list 3 5 7 10 set size (n. tokens) semantically coherent list second list repeated permuted novel transformer (radford et al, 2019) 3 5 7 10 set size (n. tokens) 80 90 100 110 repeat surprisal (%) arbitrary list 3 5 7 10 set size (n. tokens) semantically coherent list second list repeated permuted novel lstm (merity et al, 2017) a) lstm (merity et al, 2017) b) transformer (12-layer, wt-103) c) transformer (radford et al, 2019) 26 99 194 435 intervening text len. (n. tokens) 80 90 100 110 repeat surprisal (%) arbitrary list 26 99 194 435 intervening text len. (n. tokens) semantically coherent list second list repeated permuted novel lstm (merity et al, 2017) 3 5 7 10 set size (n. tokens) 60 80 100 repeat surprisal (%) arbitrary list 3 5 7 10 set size (n. tokens) semantically coherent list second list repeated permuted novel transformer (wt103, 12-layer) 26 99 194 435 intervening text len. (n. tokens) 60 80 100 repeat surprisal (%) arbitrary list 26 99 194 435 intervening text len. (n. tokens) semantically coherent list second list repeated permuted novel transformer (wt103, 12-layer) in ta c t s c ra m b le d in c o n g ru e n t intervening text 0 50 100 repeat surprisal (%) arbitrary list in ta c t s c ra m b le d in c o n g ru e n t intervening text semantically coherent list second list repeated permuted novel transformer (radford et al, 2019) verbatim retrieval as a function of intervening text type a) b) c) in ta c t s c ra m b le d in c o n g ru e n t intervening text 60 80 100 repeat surprisal (%) arbitrary list in ta c t s c ra m b le d in c o n g ru e n t intervening text semantically coherent list second list repeated permuted novel lstm (merity et al, 2017) in ta c t s c ra m b le d in c o n g ru e n t intervening text 60 80 100 repeat surprisal (%) arbitrary list in ta c t s c ra m b le d in c o n g ru e n t intervening text semantically coherent list second list repeated permuted novel transformer (wt103, 12-layer) figure 4: lm memory retrieval for different intervening texts. we plot relative list-averaged surprisal over all non-initial tokens in lists. points show group median (over n @xmath3. error bars denote 95% confidence interval around the median (bootstrap estimate). note that in the top-row plots y-axis starts at 60%. next we tested whether the transformers' ability to recall past tokens depended on the attention mechanism @xcite which allows it, in principle, to use all past words, weighted according to their relevance, for next word prediction. to test for the role of attention in verbatim retrieval, we randomly permuted the rows of key and query matrices in each of the 12 attention layers of gpt-2 and reran the experiment with varying number of tokens in lists. the shuffled-attention model retained some capacity to retrieve past nouns (fig. 8 , bottom, appendix a), but the effect was greatly reduced. for example, repeat surprisal for lists of @xmath4, compared with 3% for the intact model. intriguingly, this shuffled-attention model showed the same surprisal for repeated and permuted lists, indicating that it was no longer accessing word order information from the original list. thus, the attention mechanism is necessary for transformers to index past nouns and their order from memory. finally, a deep layered architecture is a key characteristic of transformers and performance typically scales with model size @xcite . does the capacity to perform verbatim recall depend on model size? to address this question, we trained transformers with 1, 3, 6 and 12 layers on the wikitext-103 dataset. consistent with the hypothesis that size -in addition to architecture -is crucial, the smaller 1-and 3-layer models showed a modest verbatim recall capacity, but were not sensitive to order (e.g. the 3-layer model shows 85% repeat surprisal for repeated and permuted lists of @xmath5, fig. 5 ). sensitivity to order progressively emerged in 6-and 12-layer models, where in the 12-layer model repeat surprisal levels were 5% lower for repeated relative to permuted 10-token lists (fig. 5 ). while this result confirms that even transformers trained on smaller amounts of text can exhibit short-term memory with sufficient increase in complexity, it remains unclear whether it is the increased depth or the parameter count alone that contribute to this increase in performance. ## lstm results the lstm retrieves gist-like memories over short intervening distances, facilitated by se-mantic coherence. the lstm language model expected nouns in the second list to belong to the same semantic category as the first list, and especially to the category of the earliest nouns in the first list. if the intervening text was no longer than 26 tokens, lstm repeat surprisal across noninitial token positions (fig. 3 , a ) showed a modest decrease (5%) relative to first list, but only when the nouns in the first and second lists came from the same semantic category. examining surprisal values broken down by token position in the list (fig. 2 , top) shows that in semantically coherent lists of nouns, surprisal was higher for novel lists than for repeated or permuted lists, but this memory effect was only present for tokens near the beginning of the list. in light of this limited evidence for retrieval in the lstm across 26 intervening tokens, we examined whether the lstm retrieves more successfully over shorter intervals. we reduced the intervening text to 4 tokens of coherent text ("before the meeting, mary wrote down the following lists of words. one was: <first @xmath6: <second @xmath7. in this short-range retrieval setting, we now observed a small reduction of relative repeat surprisal of 5% and 4% for arbitrary lists of 3 or 5 nouns, respectively, as well as a stronger reductions ranging from 12% (3-token list) to 5% (10-token list) for semantically coherent lists (fig. 6 ). overall, the reduction in surprisal was comparable for repeated and permuted lists, indicating that the lstm did not predict that words would occur in their original order. taken together, the experiments described in the section suggest that the lstm retrieves a semantic gist of the prior list, rather than individual tokens in order. consistent with this notion of an aggregate semantic memory, we found that retrieval was stronger for semantically coherent lists, for which an aggregated semantic representation would be closer to each of the individual words in the list. ## discussion short-term memory-the capacity to temporarily store and access recent context for current processing-is a crucial component of language prediction. in this paper, we introduced a paradigm for characterizing a language model's short-term memory capabilities, based on retrieval of verbatim content (sequences of nouns) from prior context, and used this paradigm to analyze lms with trans-verbatim retrieval of words with increasing transformer size the transformers we tested were able to access verbatim information -individual tokens and their order -from past context. furthermore, this verbatim retrieval was learned and largely resilient to interference from intervening context. this indicates that the models (especially those trained on the largest corpora) implemented, via learning, a highresolution memory system. the ability to access individual tokens may in turn support functions that rely on token indexing, akin to the functionality of the general-purpose working memory (wm) buffer proposed in cognitive science @xcite . such flexible wm could subserve the reported ability of transformers to rapidly generalize to new tasks at runtime @xcite , also known as "in-context learning". indeed, in concurrent work to ours, @xcite observed that small (2 or 3-layer) attention-only transformers developed attention heads that functioned as so-called "induction heads". these effectively performed pattern matching by looking over the past context for any occurrences of the current token and predicting the same (or similar) sequence completions. attention heads that learned this basic inductive computation were also shown to perform more general in-context learning for complex tasks such as language translation. similarly, it has been sug-gested that in standard rnns such meta-learning requires a short-term memory mechanism known as fast weights @xcite which can be thought of as analogous to self-attention in transformers @xcite . however, a highly resilient verbatim memory system could also be disadvantageous if it causes the lm to place too much confidence on verbatim features of prior context for next-word prediction. indeed, text generated from a transformer lm's predictions can be highly repetitive @xcite -it is possible that an over-reliance on accessing short-term memory may underlie this tendency. in contrast to the transformers, the lstm model only retrieved a coarse semantic category of previous lists, without fine-grained information about word order, and was only able to do so when the intervening text was short. this is in spite of the fact that the lstm had a larger parameter count than the transformer models and obtained comparable perplexity on wikitext103 ( is tracking a shallow representation of context always a limitation? not necessarily. humans frequently maintain a "good-enough" (i.e. gistlike) representation of context @xcite . when the potential for memory capacity is limited (e.g. when context must be compressed to a single hidden state as in an rnn) maintaining a broad, gist-like -as opposed to token-specificmemory of context may be more efficient overall. the memory paradigm and the measure of repeat surprisal introduced here allowed us to pinpoint computational differences in how neural lms put their architectural capacities to use for storing and accessing context in short-term memory when processing english text. while our decision to use autoregressive (left-to-right) lms was ultimately based on our initial cognitive psycholinguistic motivation, it may be fruitful to apply our paradigm to other classes of transformer models, for example, bidirectional encoder-only transformers such as bert @xcite and encoderdecoder models such as t5 @xcite . these architectures have gained traction in applied nlp settings and it would be informative to test whether this paradigm can provide diagnostic value for lm performance on other benchmarks. similarly, if the compressed context representation in lstms serves as a short-term memory bottleneck, it would be instructive to test lstm lm architectures when explicitly augmented with attention @xcite or a copy-mechanism @xcite . finally, our attention-ablation experiment in the transformer was performed uniformly across layers; future studies could focus on targeted ablations of specific attention heads to pinpoint the mechanistic locus of short-term memory @xcite . ## conclusions pretrained language models, and self-supervised predictive learning broadly, have received increased attention in terms of their (in)sufficiency as a framework for achieving feats of human-like language processing @xcite . here, akin to the line of work evaluating cognitive linguistic capacities of neural lms @xcite , we tested the ability of language models to perform an important aspect of human intelligence for natural language -flexibly accessing items from short-term memory -and showed that the transformer model, even though not trained with a short-term memory objective, retrieved remarkably detailed representations of past context. this capacity emerged from training: a transformer trained on a small amount of data showed more modest retrieval abilities. the retrieval abilities of lstm lms, by contrast, were different; the lstm maintained a summary representation of the list, which was not sensitive to word order. we conclude that our paradigm can illuminate the memory systems that arise in neural language models.