Linking In-context Learning in Transformers to Human Episodic Memory

Thank you for your appreciation and insightful feedback. Below is our point-by-point response:

Weaknesses

1. You’re right that prior works have linked Transformers to neuroscience. In the introduction, we referenced Whittington et al. (2022) on the relationship between attention mechanisms and hippocampal place cells. While the attention mechanism is equivalent to the update rule in continuous-state Hopfield networks (Ramsauer et al., 2021), its direct connection to neuroscience or the brain is less clear. Generally it’s not surprising to find connections between Transformers and memory, as the attention mechanism directly generalizes the old idea of key-value memory networks. However, our focus is on emergent capabilities rather than explicitly designed model components, particularly the emergent behavior of interacting heads rather than individual ones. We will revise the introduction to specify this gap and avoid overstatement.
2. Your point about connecting human episodic memory to fitting CMR to heads is well-taken. CMR is designed to model episodic recall patterns, especially the asymmetric contiguity bias. It only captures behavior with temporal contiguity and forward asymmetry (or complete symmetry); no CMR parameterization would produce a CRP curve that’s temporally discontiguous or backwardly asymmetric. A good fit (low CMR distance) usually requires temporal contiguity and forward asymmetry, as observed in human subjects and attention heads. We provide a direct comparison with human data in global rebuttal point 2 and examples in Fig. R3, where high CMR distances reflect temporal discontiguity and/or backward asymmetry. We also did an ablation study suggesting a causal link between low CMR distances and a model’s ICL ability (see global rebuttal point 4).
3. This is a great observation, which we are exploring in other projects. Existing works suggest that context-dependent episodic mechanisms support adaptive control and flexible reinforcement learning, including generalization to novel decision and categorization tasks (e.g., Kumaran & McClelland, 2012; Gershman & Daw, 2017; Zhou et al., 2023; Giallanza et al., 2024). By linking context-driven episodic retrieval and ICL, we aim to reveal normative and mechanistic explanations of generalizable knowledge that are universal in humans and machines.

Questions

1. Thank you for the insightful question. We’ve noticed the scale difference between raw attention scores and those from human recall. We can speculate a few reasons. (i) It’s common to tune the temperature only at the last layer, so it’s unclear if Transformers benefit from variable temperatures in other layers. (ii) The range of attention scores might depend on the distributions of the model’s training and evaluation data (e.g., evaluation prompts). (iii) Unlike LLMs, humans may not fully optimize next-word predictions, as biological objectives are more complex and varied. Recall might constitute only one use case, with the cognitive infrastructure also engaged in tasks like spatial navigation and adaptive decision-making. The more moderate scale in humans may thus reflect tradeoffs between multiple objectives.
2. Thank you for the opportunity to clarify our claim. Section 5.2 states, “As the model's loss on the designed prompt decreases through training (Fig. 7a), the degree of temporal clustering increases, especially in layers where induction heads usually emerge.” You are right that we didn’t fine-tune for the target task and only evaluated the model with in-context prompting. We took different checkpoints of pre-trained models as-is, computed their losses on the designed prompt, and assessed their in-context capability and degree of temporal clustering at different training stages. This analysis helps identify when induction/CMR-like heads and ICL abilities emerge during training.
3. We appreciate your thoughtful question. Based on existing findings (e.g., Elhage et al., 2021), we doubt that attention patterns between K- and Q-composition would differ qualitatively in hard-wired 2-layer Transformers or larger models—they likely both exhibit CMR-like CRP patterns. Since ICL in large models may arise from more complex head-composition mechanisms (noted in Olsson et al., 2022), empirically identifying attention pattern differences will be challenging. The only reliable way to distinguish between mechanisms is to examine head interactions via attention patterns, outputs, and weights ($W_{QK}$, $W_{OV}$, etc.).
4. Thank you for noting this. For example, the head in Fig. 5d exhibits the highest attention on the previous occurrence of the current token and can be categorized as a duplicate token head (Wang et al., 2022). Transformer models trained to predict the next token also learn to encode several tokens in the future (Pal et al., 2023). We speculate that heads with low CMR distances and low induction-head matching scores might play a role here, as they encode distant token information more than an ideal induction head. Understanding these CMR-like heads' functional roles remains an open question.

References

• Ramsauer et al. (2021). Hopfield Networks is All You Need.
• Kumaran & McClelland. (2012). Generalization Through the Recurrent Interaction of Episodic Memories.
• Gershman & Daw. (2017). Reinforcement Learning and Episodic Memory in Humans and Animals: An Integrative Framework.
• Zhou et al. (2023). Episodic retrieval for model-based evaluation in sequential decision tasks.
• Giallanza et al. (2024) Toward the Emergence of Intelligent Control: Episodic Generalization and Optimization.
• Elhage et al. (2021). A mathematical framework for transformer circuits.
• Olsson et al. (2022) In-context Learning and Induction Heads.
• Pal et al. (2023). Future Lens: Anticipating Subsequent Tokens from a Single Hidden State.
• Wang et al. (2022). Interpretability in the wild: a circuit for indirect object identification in gpt-2 small.

Thank you for acknowledging the potential of our work and offering us the opportunity to clarify. Below is our point-by-point response.

Weaknesses

Thank you for raising the point about the target audience of this paper. Please see the global rebuttal, point 1. Given the interdisciplinary nature of this work, our aim is to offer new insights to readers knowledgeable in mechanistic interpretability, cognitive modeling, and/or neuroscience. Thus, we want to ensure that readers from these various fields (and hopefully beyond them!) will find our paper clear and accessible. To answer your specific points:

First, we intentionally included K-composition for readers from the mechanistic interpretability (MI) community. K-composition is widely studied (Elhage et al., 2021; Singh et al., 2024; Edelman et al., 2024, Ferrando et al. 2024), while Q-composition is rarely mentioned (see Appendix of Elhage et al. 2021). Since we compare CMR with Q-composition, these readers, who are mainly familiar with K-composition, may be particularly interested in the differences between K-composition and Q-composition/CMR. We agree with you that it can be confusing for readers who are less familiar with MI. We will include this motivation in Section 3.3 and inform the reader in the Introduction which section(s) may be the most relevant based on their background and interest.

Second, the attention scores we used are standard in the field, and commonly used in both the general ML literature (e.g., Dai et al., 2019) and MI (e.g., Elhage et al., 2021). However, we recognize that an explicit definition may be helpful for those from the psychology/neuroscience side: specifically, an attention score with respect to a specific token refers to the dot product of its query vector and the key vector of a (possibly different) token, scaled by the dimension of the key/query vector. In line 102, we mentioned that “We recorded the attention scores of each head (before softmax)...”, referring to the term inside the softmax function of attention as defined in Vaswani et al. (2017).

Questions

1. We mentioned CMR distance and CMR-fitted attention scores (q) in line 219 (Section 5.1) and referred to Appendix C.3. We visualized a subset of them in Fig. 5a-d. Essentially, each set of CMR-fitted scores corresponds to a CRP curve produced by setting the parameters of CMR to specific values, minimizing its CMR distance to the corresponding average attention scores.

2. The y-axis label is correct. To clarify, top induction heads should have high matching scores and low CMR distances. We plotted CMR distance to show the difference between CMR-like heads and induction heads, which we showed an example of in Fig. 5d and discussed in the second paragraph of Discussion. We will add this clarification in the main text to avoid confusion.

3. Thank you for this important question. We believe that the discovered connections have many important consequences for both fields. In a broad sense, both fields have benefited tremendously in the past from the discovery of similar connections, including the relevance of Hopfield networks, convolutional neural networks, recurrent neural networks, reinforcement learning models in neuroscience and machine learning. So far, a similar connection has been largely lacking for transformers, which many researchers in neuroscience, cognitive science, and deep learning view as being fundamentally different from the brain. In that context, the evidence we provide paves the way for similar translations between the current generation of AI algorithms and a century of research in human memory.

   Besides this broad connection, we can also provide specific examples of how our work might inform future research in both fields. From the perspective of mechanistic interpretability, one main contribution of our results is the mechanistic and behavioral characterization of a group of heads that may support emergent ICL abilities. One interesting implication of it is that the “lost in the middle” phenomenon seen in LLMs may be related to these heads, as humans also exhibit the same recall pattern. Understanding the connection could therefore inform ways to mitigate the problem by adopting known cognitive strategies – for example, adjusting the study schedule based on the serial position (Murphy et al., 2022). From the perspective of neuroscience, our results also suggest a common principle that underlie natural and artificial intelligence. As episodic mechanisms captured by CMR have been posited to support adaptive control and general decision-making (Lu et al., 2024; Giallanza et al., 2024; Zhou et al., 2023), this connection directly enables researchers to develop alternative mechanisms for episodic memory and its interactions with other cognitive functions based on more complex attention-head composition mechanisms (e.g., see the method analyzing N-th order virtual attention head in Elhage et al., 2021). The Discussion will spell out these two implications in more detail.

References

• Elhage et al. (2021). A mathematical framework for transformer circuits.
• Singh et al. (2024). What needs to go right for an induction head?
• Edelman et al. (2024). The evolution of statistical induction heads: In-context learning markov chains.
• Ferrando et al. (2024). A primer on the inner workings of transformer-based language models.
• Dai et al. (2019). Transformer-XL: Attentive Language Models beyond a Fixed-Length Context.
• Vaswani et al. (2017). Attention is All You Need.
• Murphy et al. (2021). Metacognitive control, serial position effects, and effective transfer to self-paced study.
• Lu et al. (2024). Episodic memory supports the acquisition of structured task representations.
• Giallanza et al. (2024) Toward the Emergence of Intelligent Control: Episodic Generalization and Optimization.
• Zhou et al. (2023). Episodic retrieval for model-based evaluation in sequential decision tasks.

Thank you for your interest and thoughtful feedback on our work. Below is our point-by-point response:

Weaknesses

1. We used repeated random tokens because: (1) it aligns with human free recall experiments, where random words are presented sequentially (Murdock, 1962); (2) it is a widely acknowledged definition of induction heads in mechanistic interpretability literature (Elhage et al., 2021; Olsson et al., 2022; Bansal et al., 2022; Nanda, 2022; Crosbie et al., 2024); (3) it uses off-distribution prompts to focus on abstract properties, avoiding potential confounds from normal token statistics (Elhage et al., 2021). However, understanding these CMR-like heads in naturalistic language tasks is important, so we include a new analysis in the global rebuttal (point 4).

2. Thank you for the suggestion. Please see the causal analysis in the global rebuttal (point 4).

3. Thank you for the suggestion on generalizability. Please see the global rebuttal (point 3), where we replicated the results on the three models suggested. We chose the Pythia series due to their shared architecture and training checkpoints, which are informative about the timeline when ICL abilities and CMR-like heads emerge.

   We focused on autoregressive models with causal attention because their induction heads are widely studied, and because CMR is autoregressive and causal. Whether the biological brain has a non-autoregressive objective (e.g., masked language modeling) is an open question. In-context learning is less studied in non-autoregressive models like BERTs (but see Samuel 2024), and little is known about their “induction heads.” How the behavior and mechanisms of induction heads in GPT-like models generalize to BERT-like models is an open research question.

4. While CMR is a behavioral model, neuroscience suggests it can also explain patterns of neural activity. In CMR, episodic retrieval occurs in two phases: (1) retrieving a word based on the current temporal context via matrix $\mathbf{M}^{\rm TF}$; (2) retrieving the temporal context associated with the word via matrix $\mathbf{M}^{\rm FT}$. These associative matrices represent the full set of episodic memories, instantiated in hippocampal synapses. The temporal context serves as both input and output of the retrieval process in CMR, aligning with the hippocampus's recurrent nature. Studies suggest the temporal context is represented in hippocampal subregions (e.g., CA1 and dentate gyrus, Sakon & Kahana, 2022; Dimsdale-Zucker et al., 2020, with $\mathbf{M}^{\rm TF}$ and $\mathbf{M}^{\rm FT}$ in CA3). Others propose the temporal context is represented in cortical regions providing input to the hippocampus (e.g., entorhinal cortex; Howard et al., 2005), with $\mathbf{M}^{\rm TF}$ and $\mathbf{M}^{\rm FT}$ in the broader hippocampal network. We will include this in the revised manuscript.

5. We agree that human data should be included to contextualize our findings. To this end, we added typical CRP curves from human studies (Fig. R1a) and a more extensive comparison of fitted parameters between top CMR-like heads and average human subjects from previous experiments (Fig. R1b). Please see the global rebuttal (point 2) and Fig. R1.

Questions

1. We see rich connections between our results and neural network models of episodic memory, which can benefit research in both directions. For example, Salvatore & Zhang (2024) found that when trained to maximize recall accuracy, a seq2seq model with attention shows the same recall pattern as the best-performing CMR, with intermediate states showing similar recall patterns to human subjects. Li et al. (2024) found that RNNs trained for free recall produced the same recall order as the optimal CMR and the method of loci (where people mentally “place” items in imagined locations and then “retrieve” them in the same order), even without explicitly optimizing the recall order. Giallanza et al. (2024) showed that a neural network implementation of CMR can explain flexible cognitive control in humans. These findings suggest that CMR-like behavior can emerge in neural networks with recurrence and/or attention mechanisms, making it efficient for prediction and advantageous for general decision-making.

   Our results are also consistent with work exploring connections between attention mechanisms in transformers and the hippocampal formation (Whittington et al., 2021). While that work focused on emergent place and grid cells in transformers with recurrent position encodings, the hippocampal subfields involved are also postulated to represent CMR components (see our reply to weakness, point 4). These results generally support our proposal linking the query-key-value retrieval mechanism to human episodic retrieval.

2. Please see our reply to weakness points 1 and 2.

3. Please see our reply to weakness point 2.

4. Please see our reply to weakness point 3.

5. Please see our reply to weakness point 4.

6. Please see our reply to weakness point 5.

Finally, thank you for the minor comments. We will modify the texts/labels accordingly.

References

• Murdock et al. (1962). The serial position effect of free recall.
• Elhage et al. (2021). A mathematical framework for transformer circuits.
• Olsson et al. (2022) In-context Learning and Induction Heads.
• Bansal et al. (2022). Rethinking the role of scale for in-context learning: An interpretability-based case study at 66 billion scale.
• Nanda. (2022). Induction mosaic.
• Crosbie et al. (2024). Induction Heads as an Essential Mechanism for Pattern Matching in In-context Learning.
• Samuel. (2024). BERTs are Generative In-Context Learners.
• Salvatore et al. (2024). Parallels between Neural Machine Translation and Human Memory Search: A Cognitive Modeling Approach.
• Li et al. (2024). A neural network model trained on free recall learns the method of loci.
• Giallanza et al. (2024) Toward the Emergence of Intelligent Control: Episodic Generalization and Optimization.

Thank you for recognizing the value of our work and offering constructive feedback. We provide a point-by-point response below.

Weaknesses

Please see below.

Limitations

1. We recognize the need for more clarity, particularly for those unfamiliar with CMR. In the revised manuscript, we will enhance Section 4.1 with a more intuitive explanation and add a pedagogical description in the appendix. Please refer to our global rebuttal (point 1) for details.

2. We agree that including human data is crucial for context. The revision will feature typical CRP curves from human studies (Fig. R1a) and a more comprehensive comparison of fitted parameters between top CMR-like heads and average human subjects (Fig. R1b). See our global rebuttal (point 2) and Fig. R1 for more information.

3. Thank you for the suggestion. As a baseline, we have included a descriptive model using Gaussian functions (with the same number of parameters as CMR). Its bell shape captures the basic aspects of temporal contiguity and forward/backward asymmetry. We found that, across 12 different models (GPT2, all Pythia models, Qwen-7B, Mistral-7B, Llama3-8B), CMR provides significantly better descriptions (lower distances) than the Gaussian function for the top induction heads (average CMR distance: 0.11 (top20) / 0.05 (top50) / 0.12 (top100) / 0.12 (top200), average Gaussian distance: 1.0 (top20) / 0.98 (top50) / 0.98 (top100) / 0.97 (top200), all p<0.0001). Importantly, Gaussian functions offer little, if any, insight into how these attention patterns might arise in the first place. In contrast, our CMR shows a direct link to the Q-composition induction head at the mechanistic level.

   Additionally, we are not aware of other substantially different models from the broad neuroscience literature that can give good – let alone meaningful – results as ours. Model-free reinforcement learning and other models without temporal components will perform poorly, regardless of the degree of freedom. Even with a temporal component, models like drift-diffusion models or random context models (Murdock, 1997) cannot capture the attention patterns of induction heads, as their CRP will be essentially flat (Howard & Kahana, 2002). In fact, CMR is well-poised to explain the properties of observed CRPs, because previous models failed to explain the asymmetric contiguity, and all current successful models share the same core dynamics as CMR. If you have a particular alternative model in mind that might provide insights into the mechanisms behind induction heads and better explanations than CMR, we are happy to perform a more specific comparison.

   Finally, to improve our results over correlational model-fitting, we performed an ablation study and found that removing heads with the lowest CMR distances causes significantly worse ICL performance, suggesting they are in fact necessary for ICL (see global rebuttal point 4 and Fig. R5).

4. Thank you for pointing out the unclear sentence. We have revised it to: “The training process leads to higher values of $\beta_\text{rec}$. Specifically, $\beta_\text{rec}$ values are higher than $\beta_\text{enc}$, highlighting the importance of temporal clustering during decoding for model performance.”

5. We appreciate the concerns and believe the revisions demonstrate that our results are not cherry-picked. In particular, we have changed the sentence in line 227 to:“... we found that the majority of heads in the intermediate and later layers of GPT2-small have lower CMR distances than earlier layers”. This is supported by the distribution of CMR distances of heads in each layer (Fig. R2a). Additionally, it is inaccurate to claim “a linear increase across layers” given the low proportion (Fig. 6a) and high CMR distance (Fig. R2a) for early-intermediate layers (~25% position). Furthermore, heads with lower CMR distances tend to occur in intermediate-late layers (rather than earlier or later layers) across twelve models of varying complexity (Fig. R2b). We discuss more evidence in the next response.

6. With a goal to identify heads with CMR-like attention scores, the threshold of 0.5 is informed by both the quantitative distribution of individual head CMR distances and exploratory qualitative analysis of the attention patterns. First, as the bottom panel of Fig. 5e shows, the distribution of individual head CMR distances in GPT2 can be divided roughly into two clusters: a cluster peaking around 0.1 (and extending to around 1), and a spread-out cluster around 2.4. We examined heads with ~1 CMR distance and found non-human-like CRP (e.g., Fig. R3a, b). Further, many heads with CMR distance around 0.6-0.8 again show CRP different from humans (e.g., Fig. R3c, d). Therefore, we operationally set the threshold of 0.5. We also tested a threshold of 0.1 and found no qualitative difference (Fig. R2c, d). Thus our threshold is chosen to balance the inclusion of CMR-like heads and the exclusion of non-CMR-like heads. We will include this clarification in the appendix.

7. Please refer to (Fig. R4) for an updated Fig. 7 with standard errors shown.

References

• Elhage et al. (2021). A mathematical framework for transformer circuits.
• Olsson, et al. (2022). In-context Learning and Induction Heads.
• Crosbie, J., & Shutova, E. (2024). Induction Heads as an Essential Mechanism for Pattern Matching in In-context Learning.
• Murdock, B. B. (1997). Context and mediators in a theory of distributed associative memory (TODAM2).
• Howard, M. W., & Kahana, M. J. (2002). A distributed representation of temporal context.