Emergence of Episodic Memory in Transformers: Characterizing Changes in Temporal Structure of Attention Scores During Training (2025)

Abstract: We investigate in-context temporal biases in attention heads and transformer outputs. Using cognitive science methodologies, we analyze attention scores and outputs of the GPT-2 models of varying sizes. Across attention heads, we observe effects characteristic of human episodic memory, including temporal contiguity, primacy and recency. Transformer outputs demonstrate a tendency toward in-context serial recall. Importantly, this effect is eliminated after the ablation of the induction heads, which are the driving force behind the contiguity effect. Our findings offer insights into how transformers organize information temporally during in-context learning, shedding light on their similarities and differences with human memory and learning.

## introduction large language models (llms) have demonstrated a remarkable capacity for in-context learning. they are capable of adapting to new tasks using examples provided within the input prompt, without any parameter updates @xcite . the temporal position of tokens plays an important role in in-context learning. for instance, simply asking the model to repeat a sequence of words from the prompt requires the model to use the information about the temporal position of tokens. this resembles human learning, where the temporal organization of memory plays a critical role in recalling specific past episodes. previous work has demonstrated that some attention heads show a temporal induction property. these induction heads search the input prompt for the prior occurrence of the current token. if a match is found, they attend to the token that followed the previous presentation of the current token. this mechanism allows induction heads to effectively learn and reproduce sequences of tokens and it has been argued to contribute to the model's ability to perform tasks based on contextual information @xcite @xcite @xcite . the induction property is related to human episodic memory. numerous studies have demonstrated that episodic memory exhibits contiguity effect, where items or events that occur close together in time are more likely to be remembered together in memory recall @xcite @xcite @xcite . for example, if a person experiences a sequence of events within a short time span, they are more likely to recall these events together than if they were spread out over a longer period. this effect was systematically studied through free recall, a task where participants are presented with a list of items (e.g., words, pictures) and then asked to recall them in any order they choose. participants tend to recall items that were presented close together in the original list in clusters, indicating that the temporal proximity during encoding influences the retrieval process. this supports the idea that our memories are organized not just by the content of a stimulus but also by its temporal context @xcite . neuroscience studies have demonstrated that neural activity in the brain during recall reinstates the neural activity observed during encoding, indicating the retrieval of temporal context or a mental "jump back in time" @xcite . the contiguity effect has been recently studied in deep neural networks, including recurrent neural networks @xcite and pretrained transformer models @xcite where it was found that pretrained transformers contain attention heads that recover temporal context in a manner consistent with human episodic memory. furthermore, human-like episodic memory is found to improve performance in tasks that require processing over extended temporal context @xcite . in addition to temporal induction, transformers exhibit serial position effects, specifically recency and primacy @xcite @xcite @xcite . recency implies that tokens that are more recent in the prompt (closer to the present token) are going to be more important in generating the subsequent token. conversely, primacy implies the same property but for the tokens that are presented at the beginning of the prompt. the emergence of these effects in transformers can be shaped by positional encoding @xcite . positional encoding embeddings can be learned @xcite or fixed to vectors that convey positional information, such as rotformer, which uses rotary positional embeddings @xcite or early transformer models that used sinusoidal embeddings @xcite . these positional embeddings capture translation invariance, monotonicity and symmetry through distance metrics in their vector embeddings @xcite . recency and primacy are also characteristic of human memory, where we often exhibit better recall for items presented at the beginning (primacy effect) and end (recency effect) of a list @xcite @xcite . this parallel suggests that llms, like humans, may be sensitive to the temporal context of information, with recent and initial tokens playing a more significant role in shaping the model's internal representations and influencing its output. here we investigate the temporal aspects of attention heads and transfromer outputs during learning. we train two transformer models of different sizes (gpt-2 small and gpt-2 medium) on three datasets: wikitext-103 @xcite and two sampled datasets of fineweb (hugging-facefw, 2024) (with 1b tokens and 10b tokens). we adopt tools used in cognitive science to characterize temporal aspects of human memory and use them to examine attention heads and outputs of transformer models. specifically, we use lag-conditional recall probability (lag-crp) analysis which measures the probability of recalling an item a certain number of positions away (called lag) from the previously recalled item. we apply this analysis to transformer outputs and attention heads to understand how temporal relationships shape attention scores and token predictions. we characterize the impact of trainable positional encoding, model size, and number of training interactions on the emergence of contiguity effect and serial position effects. we also ablate induction heads to examine their role in shaping temporal properties of transformer outputs. our findings provide novel insights into similarities and differences between human memory and transformers, improving our understanding of in-context learning. ## results we trained gpt-2 small and gpt-2 medium transformer models. on the wikitext-103 dataset gpt-2 small converged after around 4000 iterations, and the larger gpt-2 medium model converged after around 2000 iterations. the convergence was determined by monitoring the validation loss. the models converged to perplexities similar to those of models with comparable size, including transformer-xl standard @xcite lamemo @xcite , hybrid h3 @xcite and trimelm long @xcite , all of which have perplexity above 20 (tab. 1). importantly, regardless of the magnitude of the positional encoding, the models converged to similar values of perplexity, consistent with results in @xcite . we also trained gpt-2 small on fineweb-1b and 10b. both models converged after around 10000 iterations. after training, most attention heads exhibited some form of structured temporal modulationrecency, primacy, contiguity, or a combination of these effects. for wikitext-103, heads in layers closer to the input were characterized by recency and primacy effects, while heads in layers closer to the output had strong recency and contiguity effects. attention scores of gpt-2 small as a function of lag across all heads before and after training are shown in fig. a3 and fig. 1 respectively. attention scores across all heads after training of gpt-2 medium model are shown in fig. a4 . fig. 2 shows attention scores as a function of lags for two representative attention heads from layers closer to the output. the top row (fig. 2a-b ) illustrates the raw values, while the middle and bottom rows (fig. 2c-f ) show results after the recency effect has been removed, highlighting the contiguity effect. for fineweb-1b and 10b attention scores across all heads after training are shown in fig. a8 and fig. a9 respectively. while similar to the results on wikitext-103, they often showed more complex, non-linear patterns across a wide range of lags. ## discussion we quantified temporal properties of attention patterns in transformer outputs using lag-crp analysis, commonly used for studying episodic memory and serial position effects in human memory experiments. by using multiple permutations of the input sequences, we were able to reduce the semantic effects of token similarity and isolate the temporal effects making it possible to observe primacy, recency and contiguity effects in the attention heads. unlike human memory experiments, where the contiguity effect is robust across a wide range of scales @xcite , supporting powerlaw decay of memory @xcite @xcite , we did not find evidence for retrieval of a broad temporal context in transformers. in fact, training typically had an impact of reducing the time constants of lag-crp to small values in the range of 2 to 4 lags. downstream analysis demonstrated a strong preference towards a serial recall that was eliminated after ablation of the induction heads. overall, we showed that tools from cognitive science can be used to better understand learning in transformers, providing valuable insights into the emergence of temporal structure during in-context learning. ## limitations our approach used relatively small models. training larger, instruction fine-tuned language models could provide additional insights into the temporal properties of in-context learning. while our approach was inspired by studies of episodic memory in humans, a number of methodological differences between the present analysis and human experiments (such as the fact that humans receive task instructions) prevent us from making direct parallels with human memory and learning.