Transformer verbatim in-context retrieval across time and scale (2024)

Abstract: To predict upcoming text, language models must in some cases retrieve in-context information verbatim. In this report, we investigated how the ability of language models to retrieve arbitrary in-context nouns developed during training (across time) and as language models trained on the same dataset increase in size (across scale). We then asked whether learning of in-context retrieval correlates with learning of more challenging zero-shot benchmarks. Furthermore, inspired by semantic effects in human short-term memory, we evaluated the retrieval with respect to a major semantic component of target nouns, namely whether they denote a concrete or abstract entity, as rated by humans. We show that verbatim in-context retrieval developed in a sudden transition early in the training process, after about 1% of the training tokens. This was observed across model sizes (from 14M and up to 12B parameters), and the transition occurred slightly later for the two smallest models. We further found that the development of verbatim in-context retrieval is positively correlated with the learning of zero-shot benchmarks. Around the transition point, all models showed the advantage of retrieving concrete nouns as opposed to abstract nouns. In all but two smallest models, the advantage dissipated away toward the end of training.

## introduction in language models (lms), successful prediction of upcoming words depends on in-context information. for example, when given the context prompt "the novel's plot and symbolism are centered around three objects: a centipede, a parachute, and a waterfall. the first and most important object in the list is the ___", an lm must retrieve the noun (centipede) out of all in-context tokens to correctly predict the continuation. in human cognitive science, this ability to flexibly retrieve items from recent context is known as short-term memory and first, studying learning trajectories of lm capacities offers complementary insights to studying only performance of fully-trained models (e.g. @xcite . previous work on lm learning trajectories showed that transformers learn nexttoken prediction by undergoing a sudden transition ("phase change") early during training, which coincides with the development of attention heads that attend to repeated tokens @xcite . does verbatim retrieval follow a similar learning trajectory? second, the ability to retrieve and predict incontext tokens verbatim (i.e. identity-based matching) can be viewed as a rudimentary form of the more flexible zero-shot learning, where the relevant in-context information is not necessarily given verbatim and must possibly be retrieved based on fuzzy, similarity-based matching @xcite . how does successful learning of verbatim retrieval relate to lm's zero-shot performance on more challenging benchmark tasks? third, while the successful retrieval of arbitrary nouns underscores the flexibility of transformer short-term memory, this approach neglects that the lexicon of natural language is not a set of unorganized, arbitrary words -instead, it has semantic structure. two prominent semantic categories are concrete and abstract nouns. concrete nouns (e.g. "hammer") have sensory referents, whereas abstract nouns (e.g. "justice") do not have a straightforward sensory component. word concreteness affects human cognitive processing. children typically acquire concrete words, especially nouns, earlier than abstract words @xcite . in certain short-term memory paradigms, humans are better at recalling concrete than abstract words @xcite . importantly, the two word categories differ also in their distributional properties: concrete words occur in a semantically narrower range of contexts compared to abstract words (schulte im @xcite . is the transformer retrieval affected by whether nouns refer to concrete vs. abstract entities? to address these questions, we evaluated verbatim in-context retrieval on the pythia suite of language models @xcite . leveraging the fact that the suite includes pretrained lms ranging from 14m to 12b parameters in scale and their intermediate training checkpoints across the entire learning epoch, we evaluated how retrieval develops over the course of training and across model sizes (fig. 1 , b and c ). additionally, the pythia suite contains zero-shot evaluations on various benchmarks for each lm checkpoint. to test how in-context retrieval relates to lm's zero-shot performance, we correlated the learning trajectory of the retrieval against the learning trajectories on zero-shot benchmarks (fig. 1 , d ). finally, to test the role of noun semantics for in-context retrieval, we evaluated how noun concreteness, as rated by human participants @xcite , affected retrieval over the course of training (fig. 1 , e ). the main contributions of the current work are: a) in all models, verbatim retrieval developed in a sudden transition early during training, after about 1% training tokens elapsed, and remained constant during the rest of training, b) learning of verbatim retrieval was positively correlated with learning of zero-shot task performance, and c) around the transition point, lms showed an advantage to retrieve concrete rather than abstract nouns. this advantage almost entirely diminished towards the end of training. ## related work several recent studies investigated the behavior of lms in domain of either verbatim or in-context retrieval more generally. @xcite developed a paradigm to test the short-term memory ability (in-context retrieval) of lms. they showed that gpt-2 can retrieve the identity and ordering of repeated arbitrary nouns, but have only tested a single fully-trained lm and did not investigate learning trajectories. @xcite compared lm (gpt-2) and human word prediction performance on spans of repeated text. they reported that lms' next word prediction performance diverges from human performance on subsequent repetitions. they showed that gpt-2 performance aligned better with humans if its attention heads had a bias towards recent context. @xcite investigated in-context retrieval of facts (e.g. retrieval of the capital city given a country name) and how such retrieval was affected by the pretraining statistics of retrieved facts. they showed that lms (pythia) could override retrieval of (counterfactual) in-context information and instead retrieved the fact that has a higher frequency of occurrence in training data (e.g., even when given the in-context counterfactual "the capital city of poland is london" they tend to predict the statistically more likely "warszaw"). the current report is also related to the recent work on lm interpretability and the role of attention heads in specific forms of retrieval. several studies @xcite @xcite have identified circuits of attention heads that detect repeated incontext tokens and their previous continuations; the computations governing the behavior investigated presently. these studies focused either on how such attention mechanisms are learned and how they affect generic next-word prediction @xcite or how these attention mechanisms govern the retrieval of proper nouns as direct objects in sentences @xcite or factual knowledge (direct objects, @xcite . here, we complement these lines of work and investigate retrieval as the ability of lms to retrieve lists of arbitrary combinations of common nouns (unlikely seen co-occuring during training) and their semantic properties. ## discussion we showed that transformer lms learned verbatim retrieval in a sudden transition, early in training, with the performance remaining stable over the course of training. the sharp onset of retrieval capacity around 1-2b tokens in training (approximately 1% total training data) is in line with the results reported by @xcite who showed that the lm loss over in-context tokens started dropping suddenly 1-2% tokens in training (between 2.5b and 5b tokens). once the learning change had occurred, the lms became better at predicting repeated text -which is what was tested in the current work. the learning trajectory of verbatim retrieval also coincides with the lms' learning trajectories on zero-shot benchmarks. this was reflected in the generally high and robust correlations across training for select tasks in our results. specifically, an abrupt change around 1b tokens in training was observed in the task of predicting the last token of a narrative passage (lambada, @xcite , multiple choice exams @xcite , and in the winogrande benchmark @xcite which requires pronoun resolution based on common-sense reasoning. retrieving in-context information (e.g. lists of nouns) verbatim is a basic computation needed for solving a zero-shot multiple-choice task: given a prompt with only in-context instructions (that is, the question and the list of possible answers), an lm system must index and retrieve (i.e. increase the probability of) the token representing the correct answer. in this sense, retrieving the correct in-context tokens is a necessary step. it is evident, however, that it is not sufficient and that verbatim retrieval must be learned along with other computations. consider the lambada and winogrande benchmarks, where the task is to predict the passage-or sentence-final word which itself is not predictable on the basis of immediately preceding words. to take an example from the winogrande benchmark: "robert woke up at 9am while samuel woke up at 6am, so he had less time to get ready for school" foot_3 . the task is to answer who the pronoun "he" refers to (robert or samuel). to this end, the lm must first establish that 9am is later than 6am -a distinct computational step indicating that "he" refers to "robert" -and only then retrieve the name to be predicted as the response. in the final experiment, we show that around the transition point (after ≃1b training tokens), when the capacity for verbatim retrieval occurs, noun semantics affect the retrieval -models showed an advantage to retrieve concrete, as opposed to abstract nouns. why would lm in-context retrieval be sensitive to noun semantics? in humans, concrete words, especially nouns, tend to be acquired earlier in development compared to abstract words @xcite . this advantage is presumably conferred by hearing words for concrete objects and concurrently observing or interacting with the objects the heard words refer to in the world. lms as text-based statistical learners by construction have no direct access to word semantics via experience or text-external data @xcite . nevertheless, text statistics, governed by human language use, can serve as a cue to the semantic structure of language -in this case, the lexicon. it is an empirical question whether and what aspects of the linguistic system are in fact recovered by lms in the service of the next-word prediction objective and subsequently reflected in the lm behavior or internal mechanisms @xcite . we speculate that earlier in training, lms are leveraging the fact that concrete nouns tend to be used in more predictable, less diverse contexts (schulte im @xcite where presumably token repetition would be more likely to occur. however, once the lms and the training compute scale in size, this distributional difference no longer confers an important advantage for retrieval. the phenomenon of concreteness advantage early, but not later in training underscores the general notion that with the increasing amounts of training data, lms as machine learning systems become incommensurate with human learners (see also @xcite , who operate on the order(s) of magnitude smaller amount of learning data, at least in terms of number of words -recent estimates point to around 100m words by adolescence @xcite . future work. in this study we investigated retrieval across a diverse set of nouns, and broken down by a core semantic dimension. however, lms are statistical learners. a dimension of future work will be to disentangle the learning sources that lms leverage to perform retrieval. in a recent study, @xcite showed that pretraining frequency can override the retrieval of counterfactual in-context information. an lm is more likely to predict a proper noun that was frequently occurring in pretraining, e.g. "warszaw", even when the counterfactual in-context prompt suggests it should retrieve a different name ( "the capital of poland is london. what is the capital of poland?"). our present results do not speak directly to this issue as our paradigm does not involve counterfactuals. it is based on lists of arbitrary nouns that unlikely frequently co-occurred in pretraining data. however, it would be important to establish whether and to what extent the in-context retrieval in general is governed by the pretraining frequencies of individual common nouns and to what degree this capacity is robust to pretraining statistics. finally, our measure of verbatim retrieval is a behavioral measure insofar that it only takes into account the output of the lm. the field of model interpretability has seen an increased interest in recent years and aims to reverse engineer the computations of lms (e.g. @xcite @xcite . future work could focus on investigating the internal mechanisms and their causal role in transformer in-context retrieval. there is consistent evidence suggesting that lms develop dedicated attention heads @xcite @xcite governing the retrieval capacity. whereas this line work frequently focuses on interpretability for practical purposes (e.g. better control of lm output in downstream applications), it would be valuable to simultaneously develop a more fine-grained computational characterization of lm mechanisms interpretable with respect to cognitive science constructs like the short-term memory @xcite . in cognitive neuroscience, language features derived from transformer lms (contextualized word embeddings) are currently among the best performing when it comes to predicting brain data recorded in human language processing tasks (e.g. @xcite @xcite . however, these high-dimensional features and the resulting statistical fits are frequently hard to interpret. coupled with loose theoretical motivations such high predicting models can be right for the wrong scientific reasons (see e.g. antonello and huth, 2023). a better characterization of lm mechanisms in terms of cognitive capacities (e.g. @xcite would be instrumental in understanding how and why lms succeed in modeling human brain and cognitive data. ## conclusion retrieving information from context is an important capacity of transformer language models. in this work, we investigated how the ability to retrieve repeated nouns from context develops across lm training and scale and its dependence on whether the retrieved nouns denote concrete or abstract entities. retrieval was learned early in training across scale and once learned, it remained stable. retrieval learning was robustly correlated with learning of zero-shot task performance. around the point when the in-context retrieval was learned, the models showed advantage to retrieving concrete as opposed to abstract nouns and the advantage dissipated as the models saw more training data.