Linking In-context Learning in Transformers to Human Episodic Memory

Summary

This work establishes a correspondence between 1) induction heads, known to functionally contribute to LLM's in-context learning, and 2) the CMR model of human episodic memory. They showed mechanistic equivalences between the model components, relating the Q-composition realization of induction heads to CMR. They also replicated an analysis of human episodic recall over different temporal token position distances on the attention scores of the induction heads in pre-trained LLMs and found human-like patterns.

Strengths

The paper establishes a nice link between transformers and human memory mechanisms, which helps the field think about the role of induction heads as how episodic memory mechanisms may contribute to language generation, especially from the angle of drifting temporal context. The normative interpretation is very interesting. The overall presentation is clear and easy to follow.

Weaknesses

The stated missing connection between transformer and neuroscience seems overstated given prior work relating the attention mechanism to the Hopfield network and transformers with hippocampus representations. I'm uncertain how informative the CMR distance metric is since CMR seems to flexibly fit to many curves. A strong link to human episodic memory would require the specific characteristics of temporal contiguity and forward asymmetry to be present in a given head. The authors make a nice connection, but it would be great if the paper can dive deeper into why this may be a meaningful connection. For example, considering that episodic memory is about context-dependent retrieval of prior-related information, but ICL can lead to some degree of generalization to novel tasks.

Questions

Since the model CRP curve uses pre-softmax attention scores, the actual transition probabilities would be of a very different scale with human recall CPR. I'm curious how the authors interpret such a scale difference. The beginning of section 5.2 seems to suggest some fine-tuning for the target task, which was a surprise as I thought the model evaluation was purely in-context prompting. Could the authors clarify? Would K-composition still show the same CRP-like patterns? Is there a way to differentiate the two, e.g. by empirical measures or by creating hard-wired 2-layer models? In Fig5, many GPT2 heads have really small CMR distance but don’t have high induction head matching scores. What could the functional role of these other low CMR-distance heads be?

Rating

6

Confidence

4

Soundness

3

Presentation

4

Contribution

3

Limitations

Stated in the paper.

Summary

The authors examine the relationship between attention heads in transformers and human episodic memory. They demonstrate that induction heads are behaviorally, functionally, and mechanistically similar to the contextual maintenance and retrieval model (CMR) of human episodic memory. In particular, they find that CMR-like heads often emerge in the intermediate model layers and that their behavior qualitatively mirrors the memory biases seen in humans.

Strengths

From a conceptual perspective, I found this a very interesting idea, trying to link recent results from mechanistic interpretability and older results from cognitive science/modeling. In general, I think there is solid scope for such insights and information transfer. Due to that, the general framing was also well-motivated.

Weaknesses

While I think there is something interesting at the basis of this work, I don’t believe that it is ready yet in its current shape. The writing, in particular, was very hard to follow, thus making it hard for me to judge the actual contents of the paper. A particularly important question to ask in this context is: who is the target audience? There are only very few people (if any) with the right background knowledge to follow the article in its present form. You pretty much find no one well-versed in mechanistic interpretability, cognitive modeling, and neuroscience. I would consider myself to have a strong background in cognitive modeling and reasonably solid knowledge of mechanistic interpretability but I was lost. To give two concrete examples of how the writing could be improved: (1) K-composition is introduced in Section 3.3. However, it is essentially completely irrelevant to the paper, and hence just confusing the reader. (2) Attention scores are a crucial concept but they are not defined anywhere.

Questions

The paper mentions CMR-fitted scores. What are these exactly? Is the y-axis label in Figure 7C correct? Shouldn’t it be matching scores instead of CMR distance? Looking beyond an exchange of ideas between two fields, are there any direct practical consequences of the discovered connection?

Rating

4

Confidence

2

Soundness

2

Presentation

1

Contribution

3

Limitations

They are discussed appropriately.

Summary

This paper explores connections between in-context learning (ICL) capabilities of large language models (LLMs) and human episodic memory. Specifically, the authors draw parallels between induction heads in Transformers and the Contextual Maintenance and Retrieval (CMR) model of human episodic memory. They demonstrate behavioural and mechanistic similarities and show that CMR-like attention patterns emerge in intermediate layers of LLMs during training. The work provides a novel and a very interesting (in my opinion at least) perspective on understanding ICL mechanisms through the lens of cognitive science models. I recommend weak acceptance, contingent on addressing some of the weaknesses identified, particularly regarding generalizability.

Strengths

1. Novelty: The paper presents an original connection between two previously separate areas of research - mechanistic interpretability of LLMs and cognitive models of human memory. I think that this approach offers insights in both fields. 2. Thorough analysis: The authors provide a detailed comparison between induction heads and CMR, including both mechanistic and behavioural analysis. The step-by-step mapping between CMR and Q-composition induction heads is particularly illuminating. 3. Empirical support: The paper includes interesting experiments on a few pre-trained LLMs (GPT-2 and Pythia models of various sizes) to support their claims. The analysis of how CMR-like behaviours emerge during training adds valuable insights. 4. Clear exposition: The paper is well-written and logically structured, making it accessible to readers from both machine learning and cognitive science backgrounds. 5. Broader impact: As in my first point, I think the work opens up new avenues for research in both AI interpretability and cognitive science, potentially leading to improved understanding in both fields.

Weaknesses

1. Limited scope of experiments: The experiments focus on a specific prompt design (repeating random tokens) which may not fully capture the complexities of natural language processing. While this is acknowledged as a limitation, it would strengthen the paper to include some analysis on more naturalistic language tasks. 2. Lack of causal analysis: While the paper shows correlations between CMR-like behaviours and model performance, it doesn't establish a causal link. It remains unclear whether these behaviours are necessary for ICL or merely a byproduct of training. 3. Generalizability: The findings are primarily based on GPT-2 and Pythia models. It's not clear how well these results generalize to other Transformer architectures or non-autoregressive models. Testing on a broader range of models would strengthen the claims. Why for example use computational resources for a Pythia-12b model and not for much more advanced and well-known models, such as Llama3-8B, Mistral-7B or Qwen-7B? All of those models are also available in TransformerLens, and I believe this would make this work more relevant to the community. 4. Limited exploration of biological plausibility: While the paper draws interesting parallels to human memory, it doesn't deeply explore the biological plausibility of the proposed mechanisms. A more thorough discussion of how these findings relate to neural implementations of episodic memory would enhance the paper's impact. 5. Lack of quantitative comparison to human data: The paper shows qualitative similarities to human memory biases, but doesn't provide quantitative comparisons to human behavioural data. Such comparisons could further validate the proposed connections, although I understand that this might be a bit too much for a single NeurIPS submission.

Questions

I have a few questions/suggestions for the authors: 1. Have you considered the connection between your results and neural-network models of episodic memory? A discussion on this (as well as the next comments) would be great. 2. How might the findings change if applied to more complex, natural language tasks rather than repeated random tokens? 3. Can you provide any insights into whether CMR-like behaviours are causally necessary for ICL, or if they might be an epiphenomenon? 4. How do you think these findings might generalize to non-autoregressive models or other Transformer variants? 5. Could you elaborate on potential neural implementations of the proposed mechanisms and their biological plausibility? 6. Have you considered comparing your model's behaviour quantitatively to human behavioural data on episodic memory tasks? Also, I have a couple of minor comments: - When describing the residual stream, I wouldd add one of the Schmidhuber citations mentioned in the Antropic paper cited here that originally described this idea (e.g. Highway Networks, Srivastava et al., 2015). - Leave a space before the citation in line 73. - Line 121 Tab. tab:comparison. I think you forgot to use /ref{.}. - Letter labels b, c and d in the caption of Figure 1 as well as all letters in Figure 2 and 3 are not bold. All other letters are so please be consistent.

Rating

7

Confidence

4

Soundness

3

Presentation

4

Contribution

3

Limitations

Overall, the authors' transparency about limitations is commendable and aligns well with the NeurIPS guidelines, although there is still room for improvement. The authors have partially addressed limitations in their Discussion section, acknowledging some key issues such as the use of random tokens as input instead of natural language, limited mechanistic interpretability offered by CMR for larger models, and potential lack of generalization to untested Transformer models. However, several important limitations identified in the paper's weaknesses are not fully addressed (see weaknesses). Potential negative societal impacts are deemed minimal for this type of foundational research.

Summary

The paper compares LLMs to a neuroscience model of human episodic memory, particularly highlighting the similarities between induction heads (responsible for in-context learning in transformers) and CMR.

Strengths

1. The work is original and offers a deeper understanding of induction heads by linking them to a well-studied model of human memory. 2. The potential significance is high, but improvements in analysis and clarity are needed for the community to build on these ideas.

Weaknesses

1. The paper lacks clarity, making it difficult to understand. 2. The analysis has several issues, which will be discussed in the limitations section.

Questions

I have mentioned them in the limitations section below.

Rating

4

Confidence

2

Soundness

2

Presentation

2

Contribution

3

Limitations

While the idea of the paper is promising, there are several points in the analysis that need to be addressed. I am open to changing my score if these points are addressed with new analysis. 1. **Clarity**: Section 4.1, explaining the CMR model, needs to be rewritten for better understanding, especially since the ML community might not be familiar with this model. 2. **Comparison to Human Data**: The paper claims similarities to human memory but doesn't compare its results with human data. Including such comparisons would illustrate how the temporal contiguity and forward asymmetry biases appear in humans. 3. **Hyperparameter fitting**: Fitting that many hyperparameters to each layer could give good fits to so many models and therefore does not prove anything. It lacks hypothesis testing or baselines. For instance, you argue that it is similar to Q-composition-like heads, and therefore fitting the model to Q-compositions vs K-compositions and showing the difference in fits would give more value to this CMR distance and interpretation. Because I can just choose any model from any neuroscience litterature and just fit it and I believe it will give me some good fits with enough degree of freedom. Therefore, some baselines would be interesting. The reader needs a reference. 4. **Hyperparameter interpretation**: The meaning of hyperparameters needs more exploration. The only mention of this is the discussion of temporal clustering with the β values which is quite vague. For instance, you say that in Figure 7 "The increase in β_rec was particularly prominent, indicating the importance of temporal clustering during decoding for model performance." which from my interpration seem to say that it is particularly prominent relative to β_end, which does not seem visually to be the case. But maybe I missed something. 5. **Qualitative Trends** As for the above point, I felt like sometimes you just cherry pick some qualitative trends that are supposed to match your claims and they are not even clear. One example is line 227: "we found that the majority of heads in the intermediate layers of GPT2-small have lower CMR distances (Fig. 6a). " Again, maybe I got this wrong but this is really not what the figure shows, it shows a linear increase across layers. 6. The CMR distance < 0.5 grouping seems arbitrary. Do you get similar results with < 0.1 and < 1 for instance? 7. Could you add error bars in Figure 7?