Linking In-context Learning in Transformers to Human Episodic Memory

Authors: ji-an  li
corey y zhou
marcus k benna
marcelo g mattar
Venue: NeurIPS_cc_2024_Conference
Published: 2024-05-15 05:27:13

Abstract: Understanding connections between artificial and biological intelligent systems can reveal fundamental principles of general intelligence. While many artificial intelligence models have a neuroscience counterpart, such connections are largely missing in Transformer models and the self-attention mechanism. Here, we examine the relationship between interacting attention heads and human episodic memory. We focus on induction heads, which contribute to in-context learning in Transformer-based large language models (LLMs). We demonstrate that induction heads are behaviorally, functionally, and mechanistically similar to the contextual maintenance and retrieval (CMR) model of human episodic memory. Our analyses of LLMs pre-trained on extensive text data show that CMR-like heads often emerge in the intermediate and late layers, qualitatively mirroring human memory biases. The ablation of CMR-like heads suggests their causal role in in-context learning. Our findings uncover a parallel between the computational mechanisms of LLMs and human memory, offering valuable insights into both research fields.

Main paper: 
 introduction: 
 Neural networks often bear striking similarities to biological intelligence. For instance, convolutional networks trained on computer vision tasks can predict neuronal activities in the visual cortex [1] [2] [3] [4] . Recurrent neural networks trained on spatial navigation develop neural representations similar to the entorhinal cortex and hippocampus [5, 6] , and those trained on reward-based tasks recapitulate biological decision-making behaviors [7] . Feedforward networks trained on category learning exhibit human-like attentional bias [8] . Identifying commonalities between artificial and biological intelligence offers unique insights into both model properties and the brain's cognitive functions.In contrast to this long tradition of drawing parallels between AI models and biology, exploration of the biological relevance of the Transformer architecture -originally proposed for natural language translation [9] -remains limited, with many researchers in neuroscience, cognitive science, and deep learning viewing it as being fundamentally different from the brain. So far, Transformers have been linked to generalized Hopfield networks [10] , neural activities in the language cortex [11] [12] [13] , and hippocampo-cortical circuit representations [14] . However, it remains unclear whether and how the emergent behavior and mechanisms of interacting attention heads relate to biological cognition.This study bridges this gap by examining the parallels between attention heads in Transformer models and episodic memory in biological cognition. We focus on "induction heads", a particular type of attention head in Transformer models and a crucial component of in-context learning (ICL) observed in LLMs [15] . ICL enables LLMs to perform new tasks on the fly during test time, relying solely on the context provided in the input prompt, without the need for additional fine-tuning or task-specific training [16, 17] . We show that induction heads share several parallel properties with the contextual maintenance and retrieval (CMR) model, an influential model of human episodic memory during free recall. Understanding the mechanisms of ICL is important for developing better models capable of performing unseen tasks, as well as for AI safety research, as the models could be instructed to perform malicious activities after being deployed in real-world scenarios.In sections below, we introduce the task in Section 2, Transformer and induction heads in Sections 3.1, 3.2, and the CMR model in Section 4.1. We show that induction heads and CMR are mechanistically similar in Sections 3.3 and 4.2 and behaviorally similar in Section 5.1. We further characterize the emergence of CMR-like behavior in Section 5.2 and its possible causal role in Section 5.3. Overall, our study provides evidence for a novel bridge between Transformer models and episodic memory. 
next-token prediction and memory recall: 
 Transformer models in language modeling are often trained to predict the next token [16] . ICL thus helps next-token prediction using information provided solely in the input prompt context. One way to evaluate a model's ICL is to run it on a sequence of repeated random tokens [15] (Fig. 1a ). For example, consider the prompt "[D]". Assuming that no structure between these tokens has been learned, the first occurrence of each token cannot be predicted -e.g., the first [C] cannot be predicted to follow the first [B] . At the second [B] , however, a model with ICL should predict [C] to follow by retrieving the temporal association in the first part of the context.Much like ICL in a Transformer model, human cognition is also known to perform associative retrieval when recalling episodic memories. Episodic retrieval is commonly studied using the free recall paradigm [18, 19] (Fig. 1b ). In free recall, participants first study a list of words sequentially, and are then asked to freely recall the studied words in any order [20] . Despite no requirements on recall order, humans often exhibit patterns of recall that reflect the temporal structure of the preceding study list. In particular, the retrieval of one word triggers the subsequent recall of other words studied close in time (temporal contiguity). Additionally, words studied after the previously recalled word are retrieved with higher probability than words studied before the previously recalled word, leading to a tendency of recalling words in the same temporal ordering of the study phase (forward asymmetry). These effects are typically quantified through the conditional response probability (CRP): given the most recently recalled stimulus with a serial position i during study, the CRP is the probability that the subsequently recalled stimulus comes from the serial position i+lag (see e.g., Fig. 4a ).3 Transformer models and induction heads 
residual stream and interacting heads: 
 The standard view of Transformers emphasizes the stacking of Transformer blocks. An alternative, mathematically equivalent view emphasizes the residual stream [21, 22] . Each token at position i in the input has its own residual stream z i (with a dimension of d model ) serving as a shared communication channel between model components at different layers (Fig. 1c ; the residual stream is shown as a blue path), such as self-attention and multi-layered perceptrons (MLP). The initial residual stream z (0) i contains token embeddings T E (vectors that represent tokens in the semantic space) and position embeddings P E (vectors that encode positions of each input token). Each model component reads from the residual stream, performs a computation, and additively writes into the residual stream. Specifically, attention heads at layer l read from all past z j (with j ≤ i) and write into the current z i as z(l) ′ i ← z (l-1) i + heads h H (h) (z (l-1) i ; {z (l-1) j} j≤i ), while MLP layers read from only the current z i and write into z i as z(l) i ← z (l) ′ i + MLP(z (l-1) ′ i). Readers unfamiliar with attention heads (e.g., attention scores and patterns) are referred to Appendix B. Other components like layer normalization are omitted for simplicity. At the final layer, the residual stream is passed through the unembedding layer to generate the logits (input to softmax) that predict the next token.Components in different layers can interact with each other through the residual stream [21] . As an important example, a first-layer head H 1 may write its output into the residual stream, which is later read by a second-layer head H 2 that writes its output to the residual stream for later layers to use.Previous mechanistic interpretability studies identified a type of attention heads critical for ICL, known as induction heads [21, 15, 23, 24] . Induction heads are defined by their match-then-copy behavior [15, 24] . They look back (prefix matching) over previous occurrences of the current input token (e.g., [B] ), determine the subsequent token (e.g., [C] if the past context included the pair [B][C]), and increase the probability of the latter -that is, after finding a "match", it makes a "copy" as the predicted next token (. .). To formalize this match-then-copy pattern, we use the induction-head matching score (between 0 and 1) to measure the prefix-matching behavior. We then use the copying score (between -1 and 1) to measure the copying behavior (see Appendix E). An induction head should have a large induction-head matching score and a positive copying score.We first examined the induction behaviors of attention heads in the pre-trained GPT2-small model [16] using the TransformerLens library [25] . To elicit induction behaviors, we constructed a prompt consisting of two repeats of a random-token sequence (see Section 2 and Appendix D). We recorded the attention scores of each head Ã (before softmax) and the attention patterns A (after softmax) for each pair of previous and current token positions (for definitions see Appendix B). Several heads in GPT2-small had a high induction-head matching score (Fig. 2a ), such as L5H1 (layer number 5, head number 1). In the first sequence repeat, this attention head attends mostly to the beginning-ofsequence token. In the second repeat, this head shows a clear "induction stripe" (Fig. 2b ) where it mostly attends to the token that follows the current token in the first repeat.We calculated attention scores as a function of relative position lags to further characterize the behavior of induction heads. This analysis is reminiscent of the CRP analysis on human recall data, as in Section 4.1. We found that induction heads' attention to earlier tokens follows a similar pattern as seen in human episodic recall (Fig. 2c , Fig. 5a-c ), including temporal contiguity (e.g., the average attention score for |lag| ≤ 2 is larger than for |lag| ≥ 4) and forward asymmetry (e.g., the average attention score for lag> 0 is larger than for lag< 0). The matching score and copying score describe the behavior of individual attention heads. However, they do not provide a mechanistic understanding of how the induction head works internally. To gain insights into the internal mechanisms of induction heads, we focus here on smaller transformer models, acknowledging that individual attention heads of larger LLMs likely exhibit more sophisticated behavior. Prior work has discovered two kinds of induction mechanisms in two-layer attention-only Transformer models: K-composition and Q-composition (Fig. 3a -b, Tab. S1) [21, 24] , characterizing how information from the first-layer head is composed to inform attention of the second-layer head.We provide an overview of both below. Our main focus in this paper is the comparison between Q-composition and CMR, but K-composition is provided as a point of comparison for readers familiar with mechanistic interpretability.In K-composition (Fig. 3a ), the first-layer "previous token" head uses the current token's position embedding, P E i , as the query, and a past token's position embedding P E j , as the key. When the match condition P E j = P E i-1 is satisfied (meaning j is the previous position of i), the head writes the previous token's token embedding, T E j , as the value into the residual stream z i . The second-layer induction head uses the current token's T E k as the query, and the previous token head's output T E j at residual stream z i as the key ("K-composition"). When the match condition T E j = T E k is satisfied, the head writes T E i (at residual stream z i ) as the value into the residual stream z k , effectively increasing the logit for the token that occurred at position i.In Q-composition (Fig. 3b ), the first-layer "duplicate token" head uses the current token's T E k as the query, and a past token's T E j as the key. When the match condition T E j = T E k is satisfied (meaning token k is a duplicate of token j), the head writes the token's P E j as the value into the residual stream z k . The second-layer induction head uses the duplicate token head's output P E j at residual stream z k as the query ("Q-composition") and a past token's P E i as the key. When the match condition P E j = P E i-1 is satisfied, the head writes T E i (at residual stream z i ) as the value into the residual stream z k , increasing the logit for the token that occurred at position i.In the following sections, we will reveal a novel connection between ICL and human episodic memory. We first introduce the CMR model of episodic memory, and then formally re-write it as a Q-composition induction head performing prefix matching, allowing us to link induction heads' attention biases to those known in human episodic memory. (c) CMR is similar to a Q-composition induction head, except that the context vector t j-1 is first updated by M FT into t j at position j, then directly used at position j + 1 (equal to i for the optimal match condition; shown by red lines).[28] effects. We provide a pedagogical introduction to CMR in Appendix C for readers without a background in cognitive science or episodic memory. Below, we briefly list the essentials of CMR.In CMR (Fig. 1d ), each word token is represented by an embedding vector f (e.g., one-hot; f i for the i-th word in a sequence). The core dynamic that drives both sequential encoding and retrieval isEQUATIONwhere t i is the temporal context at time step i, and t IN i is an input context associated with f i . β controls the degree of temporal drift between time steps (β enc for encoding/study phase and β rec for decoding/retrieval phase) and ρ is picked to ensure t i has unit norm. Specifically, during the encoding phase, t ∈ [0, 1] as t IN i = ((1 -γ FT )M FT pre + γ FT M FT exp )f i .The asymmetric contiguity bias arises from this slow evolution of temporal context: when 0 < β < 1, t i passes through multiple time steps, causing nearby tokens to be associated with temporally adjacent contexts that are similar to each other (temporal contiguity), i.e., ⟨t i , t j ⟩ is large if |i -j| is small. Additionally, t IN pre = M FT exp f i only enters the temporal context after time i. Thus t IN i,exp is associated with f j only for j > i (asymmetry). CMR also learns a second context-to-word (mapping t back to f ) memory matrix M TF (updated by each f i t T i-1 ). When an output is needed, CMR retrieves a mixed word embedding f IN i = M TF t i . If f j are one-hot encoded, we can simply treat f IN i as a (unnormalized) probability distribution over the input tokens. Or, CMR can compute the inner product ⟨f j , f IN i ⟩ for each cached word f j as input to softmax (with an inverse temperature τ -1 ) to recall a word.Intuitively, the temporal context resembles a moving spotlight with a fuzzy edge: it carries recencyweighted historical information that may be relevant to the present, where the degree of information degradation is controlled by ρ. Larger β's correspond to "sharper" CRPs with stronger forward asymmetry and stronger temporal clustering that are core features of human episodic memory. As a concrete example, consider n unique one-hot encoded words{f i }. If M FT pre = n i=1 f i f T i (i.e., the pre-experimental context associated with each word embedding is the word embedding itself) and γ FT = 0, Eq. 1 at decoding is reduced to t i = ρt i-1 + βf i = i j=1 βρ i-j f j , which is a linear combination of past word embeddings.  
cmr as an induction head: 
 We now proceed to map CMR to a Q-composition-like head (see Fig. 3c and Tab. S1 for details).To begin with, we note that the word f i seen at position i is the same as T E i , and the context vector t i-1 (before update) at position i is functionally similar to P E i . It follows that the set {f i , t i-1 } is functionally similar to the residual stream z i , updated by the head outputs.CMR experimental word-context retrieval as first-layer self-attention. The temporal context is updated byt IN i = ((1 -γ FT )M FT pre + γ FT M FT exp,i)f i at decoding. The memory matrix M FT exp,i acts as a first-layer duplicate token head, where the current word f i is the query, the past embeddings {f j } make up the keys, and the temporal contexts t j-1 associated with each f j are values. This head effectively outputs "What's the position (context vector) at which I encountered the same token f i ?" CMR pre-experimental word-context retrieval as MLP. The pre-experimental context t IN i = M FT pre f i (retrieved contextual information not present in the experiment) is the output of a linear fully-connected layer (functionally similar to MLP; not drawn).CMR evolution as residual stream updating. The context vector is updated by 3c ). CMR context-word retrieval as second-layer self-attention. The retrieved embedding is f IN k = M TF t k , where M TF acts as a second-layer induction head, where the temporal context t k is the query, the past contexts {t i-1 } make up the keys, and the embeddings f i associated with each t i-1 are values. This effectively implements Q-composition [21] , because t k , as the Query, is affected by the output of the first-layer M FT exp head. CMR word recall as unembedding. The final retrieved word probability is determined by the inner product between the retrieved memory f IN k and each studied word f j , similar to the unembedding layer generating the output logits from the residual stream.t k = ρt k-1 + βt IN k . Equivalently, the head M FT updates the information from {f k , t k-1 } to {f k , t k-1 , t k }. At the position k during recall, the updated context t k contains t IN k (≈ t j ) (Fig.CMR learning as a causal linear attention head. Both associative matrices of CMR are learned in a manner consistent with the formation of causal linear attention heads. Specifically, the wordto-context matrix is updated byM FT exp,i = M FT exp,i-1 + t i-1 f T i (with M FT exp,0 = 0), associating f i (key) and t i-1 (value). It is equivalent to a causal linear attention head, because M FT exp f k = ( i<k t i-1 f T i )f k = i<k t i-1 (f T i f k ) = i<k sim(f i , f k )t i-1 .Similarly, the context-to-word matrix, updated byM TF i = M TF i-1 + f i t T i-1 (with M TF 0 = 0), is equivalent to a causal linear attention head that associates t i-1 (key) with f i (value).To summarize, the CMR architecture resembles a two-layer transformer with a Q-composition linear induction head. It's worth noting that although we cast t i as the position embedding, unlike position embeddings that permit parallel processing in Transformer models, t i is recurrently updated in CMR (Eq. 1). It is possible that Transformer models might acquire induction heads with a similar circuit mechanism, where t i corresponds to autoregressively updated context information in the residual stream that serves as the input for downstream attention heads.Question: Does the paper fully disclose all the information needed to reproduce the main experimental results of the paper to the extent that it affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?Answer: [Yes] Justification: We fully disclose the models (available through the TransformerLens library [25] ; CMR implementation is included in the code submission), metrics (see Appendix E), model fitting (see Section 5.1), and head ablation (see Section 5.3). All figures included in the paper can be reproduced using the included code (see Appendix A).• The answer NA means that the paper does not include experiments.• If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not. • If the contribution is a dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable. • Depending on the contribution, reproducibility can be accomplished in various ways.For example, if the contribution is a novel architecture, describing the architecture fully might suffice, or if the contribution is a specific model and empirical evaluation, it may be necessary to either make it possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g., in the case of a large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed. • While NeurIPS does not require releasing code, the conference does require all submissions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. Guidelines:• The answer NA means that the paper does not include experiments.• The authors should answer "Yes" if the results are accompanied by error bars, confidence intervals, or statistical significance tests, at least for the experiments that support the main claims of the paper. • The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions). • The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.) • The assumptions made should be given (e.g., Normally distributed errors).• It should be clear whether the error bar is the standard deviation or the standard error of the mean. • We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurIPS Code of Ethics and the guidelines for their institution. • For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.Contextual maintenance and retrieval model (CMR)4.1 CMR in its original form CMR, an influential model of human episodic memory, provides a general framework to model memory recall as associative retrieval. It leverages a distributed representation called temporal context, which acts as a dynamic cue for recalling subsequent information based on recently seen words[26]. CMR explains the asymmetric contiguity bias in human free recall (see Fig.4and Fig.S1) and has been extended to more complex memory phenomena such as semantic[27] and emotional 
quantifying the similarity between an induction head and cmr: 
 We have shown that induction heads in pre-trained LLMs exhibit CMR-like attention biases (Fig. 2c , Fig. 5a-c ; also see Fig. S2 for heads unlike CMR) and established the mechanistic similarity between induction heads and CMR (Fig. 3 ). To further quantify their behavioral similarity, we propose the metric CMR distance, defined as the mean squared error (MSE) between the head's average attention scores and its CMR-fitted scores (see Appendix E, Fig. 5a-d ). In essence, we optimized the parameters (β enc , β rec , γ FT , τ -1 ) for each head to obtain a set of CMR-fitted scores that minimizes MSE.At the population level, heads with a large induction-head matching score and a positive copying score also have a smaller CMR distance (Fig. 5e ), suggesting that the CMR distance captures meaningful behavior of these heads. Notably, certain heads that are not typically considered induction heads (e.g., peaking at lag=0) can be well captured by CMR (Fig. 5d ).Consistent with prior findings that induction heads were primarily observed in the intermediate layers of LLMs [31] , we found that the majority of heads in the intermediate-to-late layers of GPT2-small have lower CMR distances (Fig. 6a , Fig. S3a ; for the choice of threshold see Appendix E.5). We also observed similar phenomena in a different set of LLMs called Pythia (Fig. 6b ), a family of models with shared architecture but different sizes, as well as three well-known models (Qwen-7B [32] , Mistral-7B [33] , Llama3-8B [34] , Fig. 6c ). We summarized these results in Fig. S3b .Additionally, to contextualize these CMR distances, we included the Gaussian distance using Gaussian functions as a baseline (with the same number of parameters as CMR), since its bell shape captures the basic aspects of temporal contiguity and forward/backward asymmetry. We found that, across 12 different models (GPT2, 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 (top 20), 0.05 (top 50), 0.12 (top 100), 0.12 (top 200); average Gaussian distance: 1.0 (top 20), 0.98 (top 50), 0.98 (top 100), 0.97 (top 200); all p < 0.0001), again suggesting that CMR is well-poised to explain the properties of observed CRPs.The Pythia models' different pretraining checkpoints allowed us to measure the emergence of CMRlike behavior. 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. For instance, intermediate layers of Pythia-70m (e.g., L3, L4) show the strongest temporal clustering that persists over training (Fig. S4a-b ). This, combined with an increasing inverse temperature (Fig. 7b ), suggests that attention patterns become more deterministic over training, while shaped to mirror human-like asymmetric contiguity biases. In fact, human subjects with better free recall performance tend to exhibit stronger temporal clustering and a higher inverse temperature [29] . For individual heads, those with higher induction-head matching scores (Fig. 7d ) (or similarly with smaller CMR distances, see Fig. 7f ) consistently exhibit greater temporal clustering (Fig. 7e , f respectively), as the fitted β's (both β enc and β rec ) were large. We also observed similar values of fitted βs in Qwen-7B, Mistral-7B, and Llama3-8B (Fig. S6b ). These fitted β's of these attention heads fall into a similar range as human recall data (Fig. 7c ). We interpret this in light of a normative view of the human memory system: in humans, the asymmetric contiguity bias with a β < 1 is not merely phenomenological; under the CMR framework, it gives rise to an optimal policy to maximize memory recall when encoding and retrieval are noisy [29] . In effect, a large β (but less than 1) in Eq. 1 provides meaningful associations beyond adjacent words to facilitate recall, such that even if the immediately following token is poorly encoded or the agent fails to decode it, information from close-by tokens encapsulated in the temporal context still allows the agent to continue decoding.In addition, we observed an increase in β enc and β rec (Fig. 7e , g) during the first 10 training checkpoints, when the model loss significantly drops. The training process leads to higher values of β rec . Specifically, β rec values are higher than β enc , highlighting the importance of temporal clustering during decoding for model performance. These results suggest that attention to temporally adjacent tokens with an asymmetric contiguity bias may support ICL in LLMs.While these CMR-like heads were identified using repeated random tokens, we ask whether they were also causally necessary for ICL in more naturalistic tasks. We thus performed an ablation study using natural language texts. Specifically, we ablated either the top 10% CMR-like heads (i.e., top 10% heads with the smallest CMR distances) or the same number of randomly selected heads in each model. We then computed the resultant ICL score on the sampled texts (2000 sequences, each with at least 512 tokens) from the processed version of Google's C4 dataset [35] . The ICL score is defined as the loss of the 500th token in the context minus the loss of the 50th token in the context, averaged over dataset examples [15] . Intuitively, a model with better in-context learning ability has a lower ICL score, as the 500th token is further into the context established from the beginning. We tested models with various model architectures and complexity, including GPT2, Pythia models, and Qwen-7B (Fig. 8 ). Most models showed a higher ICL score (worse ICL ability) if the top 10% CMR-like heads were ablated, compared to if the same number of randomly selected heads were ablated. This effect was particularly significant if the original model had a low ICL score (e.g., GPT2, Qwen-7B).Our finding therefore suggests that CMR-like heads are not merely an epiphenomenon, but essential underlying LLM's ICL ability, and that the CMR distance is a meaningful metric to characterize individual heads in LLMs. Nonetheless, our result needs to be interpreted cautiously: First, a Hydra effect has been noted where ablation of heads causes other heads to compensate [36] . Second, our analysis cannot confirm a direct causal role of the CMR-like behavior of these CMR-like heads in ICL, since it is possible that characteristics other than the episodic-memory features in these CMR-like heads might causally contribute to ICL. Finally, the ICL scores for the original Pythia models are closer to 0 (worse ICL performance) than other models', suggesting either weaker ICL ability of the Pythia series, or larger distributional differences between their training data and our evaluation texts. 
discussion: 
 This study bridges LLMs and human episodic memory by comparing Transformer models' induction heads and CMR. We revealed mechanistic similarities between CMR and Q-composition induction heads and identified CMR-like attention biases (i.e., asymmetric contiguity) in pre-trained LLMs. Notably, CMR-like heads emerge in LLMs' intermediate-to-late layers, evolve towards a state akin to human memory biases, and may play a causal role in ICL. These findings offer significant connections between the current generation of AI algorithms and a century of human memory research.From a neuroscience and cognitive science perspective, CMR's link to induction heads might reveal normative principles of memory and hippocampal processing, echoing the role of the hippocampus in pattern prediction and completion [37] [38] [39] . While CMR is a behavioral model, it also explains neural activity patterns: associative matrices that represent episodic memories may be instantiated in hippocampal synapses. The temporal context aligns with the hippocampus' recurrent nature, likely involving subregions including CA1, dentate gyrus [40] , CA3 [41] ), and cortical areas projecting to the hippocampus such as the entorhinal cortex [42] .Our results show strong connections to neural network models of episodic memory. For example, neural networks with attention [43] or recurrence [44] trained for free recall exhibit the same recall pattern as the optimal CMR. [45] showed that a neural network implementation of CMR can explain humans' flexible cognitive control. Our results also align with research connecting attention mechanisms in Transformers to the hippocampal formation [14] . While prior work focused on emergent place and grid cells in Transformers, the hippocampal subfields involved are also postulated to represent CMR components [40] . These results support our proposal that query-key-value attention mechanisms link to biological episodic retrieval, suggesting that CMR-like behavior emerges naturally in neural networks with proper objectives.The link to induction heads could enable researchers to develop alternative mechanisms for episodic memory. For instance, K-composition and Q-composition induction circuits might serve as alternative models to CMR. K-composition and Q-composition require positional encoding, which we speculate could be implemented by grid-like cells with periodic activations tracking space and time [46] . Further, the interactions between episodic memory and other advanced cognitive functions [47] [48] [49] [50] [51] might be understood based on more complex attention-head composition mechanisms (e.g., N-th order virtual attention head [21] ).From the perspective of LLM mechanistic interpretability, we offer a more detailed behavioral description and reinterpret induction heads through the lens of CMR and the asymmetric contiguity bias. First, CMR-like heads are not limited to induction heads -some attention heads are well captured by CMR despite not meeting the traditional induction head criterion (e.g., Fig. 5d ). Second, we speculate that heads with low CMR distances and low induction-head matching scores may encode multiple future tokens observed in LLMs [52] , capturing distant token information better than ideal induction heads. Third, we observed a scale difference between raw attention scores and those from human recall, as discussed in Appendix G.1. Lastly, though CMR relies on recurrently updated context vectors that are different from the K-composition and Q-composition mechanisms, we posit that deeper Transformer models may develop similar mechanisms via autoregressively updated information in the residual stream, a possibility yet to be explored. Overall, our results suggest a fuller view of ICL mechanisms in LLMs, where heads learn attention biases akin to human episodic memory, empowering next-token and future-token predictions.These connections with CMR may shed light on intriguing features and functions in LLMs. For instance, the "lost in the middle" phenomenon may be related to these heads [53] , as humans exhibit similar recall patterns [19, 26] . Understanding the connection could suggest strategies to mitigate the problem, e.g., adjusting study schedules based on the serial position [54] . Second, CMR not only applies to individual words but also to clusters of items [55, 56] , suggesting these heads may process hierarchically organized text chunks [57] . Third, episodic mechanisms captured by CMR are posited to support adaptive control [47] [48] [49] and flexible reinforcement learning [50, 51] , suggesting similar roles for these heads in more complex cognitive functions of LLMs. Finally, it was proposed that Transformers can be implemented in a biologically plausible way [58] , and we outline an alternative proposal mapping Transformer models to the brain (Appendix G.2).Our study has several limitations. First, while we use CMR to characterize these heads' behavior, it is unclear if CMR can serve as a mechanistic model in larger Transformer models. Second, it is unclear if our conclusions hold for untested Transformer models. Third, our causal analysis lacks the power to narrow down the causal role of specific CMR-like characteristics. In addition, while ICL score is a primary metric for measuring ICL ability [15, 59] , a systematic evaluation using specific datasets and tasks would allow for a stronger causal claim. Addressing these limitations is a key future direction.All code is available at https://github.com/corxyz/icl-cmr.For each attention head h, the standard self-attention mechanism calculates its output asH (h) (z (l-1) i ; {z (l-1) j }) = W (h) O j W (h) V z (l-1) j A (h) ij , where W (h) O is the d model × d head output projection matrix, W (h) Vis the d head × d model value projection matrix, and A (h) ij is the attention pattern attending from the query position i to the key position j. The attention pattern matrix is A (h) = Softmax( Ã(h) ) = Softmax(z (l-1) T j W T Q W K z (l-1) j / √ d head ),In this section, we offer a more detailed description of the Context Maintenance and Retrieval (CMR) model of human episodic memory, focusing on its core components and underlying dynamics. CMR operates in two distinct experimental phases: an encoding/learning phase, where the model memorizes a list of words, and a retrieval/recall phase, where it freely recalls the learned words. The CMR model consists of four primary components: the evolving temporal context (t), the word representation or embedding (f ), the word-to-context association matrix (M FT ), and the context-to-word association matrix (M TF ).The temporal context vector (t) continuously evolves over time, incorporating recency-weighted historical information during both encoding and retrieval. The core dynamics of the temporal context t are governed by the update rule Eq. 1.For illustrative purposes, in the rest of this section, we consider an example sequence consisting of five distinct one-hot word embeddings f 1 , f 2 , ..., f 5 . Each embedding is represented by a 6-dimensional vector, such that f i contains a 1 at the i-th position, and 0 everywhere else. The 6-th dimension represents a dummy stimulus that only serves to maintain the unit norm of the context vector but otherwise irrelevant to the experiment. Additionally, we index encoding/learning time with j and retrieval/recall time with k to explicitly mark the two phases.For simplicity, we set the pre-experimental word-to-context memory matrix M FT pre (6 × 6) to the identity matrix. As the result of the distinctiveness assumption, each word embedding at encoding directly induces the same one-hot vector as the input context t IN j = M FT exp f j = f j . Let t 0 = (0, 0, 0, 0, 0, 1) T . Thus given a sequential presentation of f 1 , f 2 , ..., f 5 , the corresponding context vectors t j are t 1 = (β, 0, 0, 0, 0, ρ) T t 2 = (ρβ, β, 0, 0, 0, ρ 2 ) T . . .t 5 = (ρ 4 β, ρ 3 β, ρ 2 β, ρβ, β, ρ 5 ) T ,where β is the temporal drift parameter, and ρ = 1 -β 2 ensures that t j has unit norm. Distinct parameters β enc and β rec are used by the encoding and retrieval phases respectively to capture the phase-specific temporal dynamics. Conceptually, the temporal context reflects a recency-weighted history of past information, with the parameter ρ determining the rate of degradation over time.More generally, the input context is determined by both pre-experimental and experimental contexts:t IN j = M FT f j = ((1 -γ FT )M FT pre + γ FT M FT exp )f j .The memory matrix M FT is therefore a weighted combination of pre-existing associations learned prior to encoding (the pre-experimental memory matrix M FT pre ) and associations learned during the encoding phase (the experimental memory matrix M FT exp ). M FT pre is commonly initialized as the identity matrix at the beginning of each experiment and held constant throughout the experiment. On the other hand, M FT exp is usually initialized as the zero matrix. It is updated during the encoding phase at the presentation of each new word by associating the context vector t j-1 (value) with the embedding of the just encountered word f j (key): 3c , first layer).M FT j,exp ← M FT j-1,exp + t j-1 f T j (Fig.Suppose at the beginning of the experiment, the experimental word-to-context matrix M FT 0,exp is initialized as the zero matrix. At encoding time j = 1, the experimental matrix M FT j,exp is updated by the pairwise association of the previous context and the current word:M FT 1,exp = M FT 0,exp + t 0 f T 1 = M FT 0,exp +        0 0 0 0 0 1        [1 0 0 0 0 0] =       0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 1 0 0 0 0 0      .Similarly, at time j = 2, the experimental matrix incorporates the new context-word pair asM FT 2,exp = M FT 1,exp + t 1 f T 2 = M FT 1,exp +        β 0 0 0 0 ρ        [0 1 0 0 0 0] =       0 β 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 1 ρ 0 0 0 0       .The same update procedure continues until the end of the word sequence. Critically, note that the product of M FT j,exp and any f j+l with l > 0 is the zero vector. During retrieval, the most recently recalled word f k (query) retrieves its corresponding input contextt IN k = M FT f k = (1 -γ FT )M FT pre f k + γ FT M FT exp f k = (1 -γ FT )M FT pre f k + γ FT l (t l-1 f T l )f k = (1 -γ FT )M FT pre f k + γ FT l t l-1 (f T l f k ).Let f j denote the word encountered at encoding time j with j < k. Assuming f k = f j (i.e., the most recent recall f k matches a previously encoded word f j ), we havet IN k = (1 -γ FT )M FT pre f j + γ FT t j-1 .Furthermore, because we assumed M FT pre to be the identity matrix,t IN k = (1 -γ FT )f j + γ FT t j-1 . (S1)This input context is subsequently reinstated along with the current context t k-1 at (retrieval) time k, resulting in an updated context t k via Eq. 1. Thus t k also incorporates information from f j (the previous instance of the most recent recall) and t j-1 (the context immediately before presentation of the previous instance).Importantly, the context t j at encoding is also a linear combination of f j and t j-1 (by Eq. 1),t j = βt IN j + ρt j-1 ,wheret IN j = ((1 -γ FT )M FT pre + γ FT M FT exp )f j .However, because all word embeddings are one-hot and distinct, M FT exp has not learned any associations involving the word f j prior to time j. Therefore, the second term is zero, andt j = βf j + ρt j-1 . (S2)Comparing Eq. S2 and Eq. S1 thus reveals that t IN k , thus t k , become most similar to t j . Finally, the context-to-word association matrix M TF is responsible for retrieving the context associated with a given word. M TF is usually initialized as the zero matrix prior to experiments. During encoding, it is updated to associate each word embedding f j (value) with the corresponding context vector t j-1 (key):M TF j ← M TF j-1 + f j t T j-1(Fig. 3 , second layer). In practice, we chose to maintain M TF as the transpose of M FT exp . Because the updated context t k is most similar to t j (measured by their inner product), the retrieved vectorf IN k+1 = M TF t k = l (f l t T l-1 )t k = l f l (t T l-1 t k )will be the most similar to the word embedding f j+1 . Thus, CMR will most likely recall f j+1 following a recall of f k = f j . In the main text, we specified i = j + 1.Consequently, the CMR model explains two key phenomena observed in human free recall experiments: the temporal contiguity effect and the forward asymmetry effect. The temporal contiguity effect arises because (i) t k integrates the information of t j-1 , and (ii) the context vector t drifts over time, making t j-1 more similar to context vectors associated with nearby words in the sequence, favoring the retrieval of words close to f j . In addition, the forward asymmetry effect occurs because (i) t k integrates the information of t IN j , and (ii) t IN j contributes to t j = t i-1 and its following context vectors, but does not affect earlier context vectors. As a result, t k preferentially retrieves words that follow f j , rather than those that precede it.We used the prompt of repeated random tokens due to several reasons (1) it aligns with human free recall experiments, where random words are presented sequentially [18] ; (2) it is a widely acknowledged definition of induction heads in mechanistic interpretability literature [21, 15, 31, 60] ;(3) it uses off-distribution prompts to focus on abstract properties, avoiding potential confounds from normal token statistics [21] .The prompt is constructed by taking the top N=100 most common English tokens with a leading space (to avoid unwanted tokenization behavior), which are tokens with the largest biases in the unembedding layer of GPT2-small (or the Pythia models). The prompt concatenated two copies of the permuted word sequence and had a total length of 2N + 1 (one end-of-sequence token at the beginning). The two copies correspond to the study (encoding) and recall phases, respectively.Formally, we define the induction-head target pattern (i.e., attention probability distribution of an ideal induction head) Ā = (ā d,s ) over a sequence of tokens {x i } asEQUATIONHere, d is the destination position and s is the source position. We then assess the extent to which each attention head performs this kind of prefix matching [15, 25] . Specifically, the induction-head matching score for a head with attention pattern A = (a d,s ) is defined as( d s a d,s ād,s )/( d s a d,s ) ∈ [0, 1](Fig. 2a-b ). A head that always performs ideal prefix matching will have an induction-head matching score of 1.Additionally, an induction head should write to the current residual stream to increase the corresponding logit of the attended token (token copying). We adopt the copying score [21] to measure each head's tendency of token copying. In particular, consider the circuitW = W U W O W V W E for one head,where W E defines token embeddings, W V computes the value of each token from the residual stream (i.e., aggregated outputs from all earlier layers), W O computes the head's output using a linear combination of the value vectors and the head's attention pattern, and W U unembeds the output to predict the next token. The copying score of the head is equal to( k λ k )/( k |λ k |) ∈ [-1, 1],where λ k 's are the eigenvalues of the matrix W . Since copying requires positive eigenvalues (corresponding to increased logits), an induction head should have a positive copying score. An ideal induction head will have a copying score of 1.For a head with attention scores Ã, the average attention score αlag is defined asαlag = 1 N -|lag| * 2 |lag|<s≤N -|lag| ãs+N,s+lag ,where N is the length of the first repeat in the prompt. Thus if lag = 0, αlag quantifies how much the first instance of a token is attended to on average; for lag = 1, αlag quantifies the average amount of attention distributed to the immediately following token in the first repeat. We used lag ∈ [-5, 5] throughout the paper.The metric CMR distance is defined asd CMR = min q    lag (q lag -αlag ) 2 /Var (α lag ) N lag    .Here, N lag is the number of distinct lags, and q lag is obtained by calculating the CRP using CMR with specific parameter values. Note we did not consider the full set of possible q but only the subset given by the combinations of parameters β enc = 0.05, 0.1, . . . , 1, β rec = 0, 0.05, . . . , 1, and γ FT = 0, 0.1, . . . , 1 for model fitting.The metric Gaussian distance, serving as a baseline, is defined asd Gaussian = min g    lag (g lag -αlag ) 2 /Var (α lag ) N lag    ,where g lag is the Gaussian function, defined asg lag = c 1 exp - (lag -c 2 ) 2 2c 2 3 + c 4 ,and c 1 , c 2 , c 3 , c 4 are coefficients.To identify heads with CMR-like attention scores and exclude non-CMR-like heads, we operationally chose a threshold of 0.5 in the main text. This threshold 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 in 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. Second, we examined heads with 1 CMR distance and found non-human-like CRP (e.g., Fig. S2a, b ). Further, many heads with CMR distance around 0.6-0.8 again show CRP different from humans (e.g., Fig. S2c, d ). Finally, we also tested a threshold of 0.1 and found no qualitative difference (Fig. Fig. S3c, d ).K-composition Q-composition CMR Representation (residual stream) before H 1 P E i & T E i [At i] P E j & T E j [At j] P E k & T E k [At k] P E i & T E i [At i] P E j & T E j [At j] P E k & T E k [At k] t i-1 & f i [At i] t j-1 & f j [At j] t k-1 & f k [At k] First-layer head H 1previous token head duplicate token head word-to-context matrixM FT Query P E i [At i] T E k [At k] f k [At k] Key P E j [At j] T E j [At j] f j [At j] Optimal Q-K match condition P E j+1 = P E i T E j = T E k f j = f k Activation Softmax Softmax Linear Value T E j [At j] P E j [At j] t j-1 [At j] Output T E j [At i] * P E j [At k] * t j-1 [At k] Representation (residual stream) after H 1 P E i & T E i & T E j [At i] P E j & T E j [At j] P E k & T E k [At k] P E i & T E i [At i] P E j & T E j [At j] P E k & T E k & P E j [At k] t i-1 & t i & f i [At i] t j-1 & t j & f j [At j] t k-1 & t k & f k [At k] † second-layer head H 2 Type induction head induction head context-to-word matrix M TF Query T E k [At k] P E j [At k] t j [At k] † Key T E j [At i] P E i [At i] t i-1 [At i] Optimal Q-K match condition T E j = T E k P E j+1 = P E i † † t j = t i-1 † † Activation Softmax Softmax Linear Value T E i [At i] T E i [At i] f i [At i] Output T E i [At k] * T E i [At k] * f IN k+1 [At k] Representation (residual stream) after H 2 P E i & T E i & T E j [At i] P E j & T E j [At j] P E k & T E k & T E i [At k] P E i & T E i [At i] P E j & T E j [At j] P E k & T E k & P E j & T E i [At k] t i-1 & t i & f i [At i] t j-1 & t j & f j [At j] t k-1 & t k & f k & f IN k+1 [At k]Table S1 : Comparison of mechanisms of induction heads and CMR. P E i : position embedding for the token at index/position i. T E i : token embedding for the token at index/position i. t i-1 : context vector associated with the word at index/position i. f i : word embedding for the word at index/position i.[At i]: information available at index/position i. * : The output is approximate (due to the weighted-average over all previous tokens). † : t k is updated from t k-1 by t j-1 (due to M FT exp ) and t IN j (due to M FT pre ), thus containing the information of t j (updated from t j-1 by t IN j ). † † : The optimal Q-K match condition in Q-composition requires transformation from P E j to P E j+1 , implemented by the W Q W K matrix of H 2 . The optimal Q-K match condition in CMR requires transformation from t j-1 to t j , implemented by t IN j .G Supplementary discussionWe have observed 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 readout layer, and it's unclear if Transformers benefit from variable temperatures in other attention 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.Here, we propose that the MLP layers might be mapped to the cortex and the self-attention layers might be mapped to the hippocampus. The model parameters could be encoded by slowly updated synapses in the cortex and hippocampus, with the key-value associations stored in fast Hebbian-like hippocampal synapses. The residual stream updated by MLP and attention layers may be akin to the activation-based working memory quickly updated by the cortico-hippocampal circuits.The induction-head matching scores and copying scores of each head in GPT2-small and all Pythia models are computed using Google Colab Notebook. All models were pretrained and accessible through the TransformerLens library [25] with MIT License and used as is. See Table S2 for details. Table S2 : Details of compute resources used to compute induction head metrics. All models were pretrained and accessible through the TransformerLens library [25] with MIT License. The numbers in the "Computing time" column indicate the total number of minutes it took to compute all scores for all heads across all checkpoints where available.We used an internal cluster to compute the subset of q (CRP) for model fitting. CRP enc = 0.9, rec = 0.0 enc = 0.9, rec = 0.1 enc = 0.9, rec = 0.2 enc = 0.9, rec = 0.3 enc = 0.9, rec = 0.4 enc = 0.9, rec = 0.5 enc = 0.9, rec = 0.6 enc = 0.9, rec = 0.7 enc = 0.9, rec = 0.8 enc = 0.9, rec = 0.9 enc = 0.9, rec = 1.0 Question: For each theoretical result, does the paper provide the full set of assumptions and a complete (and correct) proof?Answer: [NA] Justification: This paper does not introduce any new theorems, formulas, or lemmas to be proved.• The answer NA means that the paper does not include theoretical results.• All the theorems, formulas, and proofs in the paper should be numbered and crossreferenced. • All assumptions should be clearly stated or referenced in the statement of any theorems.• The proofs can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged to provide a short proof sketch to provide intuition. • Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in or supplemental material. • Theorems and Lemmas that the proof relies upon should be properly 
 
 Appendix:  


Decision: This submission presents mechanistic and behavioral correspondences between transformer language models and a classical model of human episodic memory, the contextual maintenance and retrieval (CMR) model. Connections of this kind are of great potential interest to the NeurIPS community because of a growing gap between neuro-inspired and engineering-based approaches to designing machine learning systems, and this submission bridges these two subdomains nicely. The submission and author responses themselves are thorough and precise, and the authors provided open-source code. Reviewers had no qualms about the technical details of the paper but raised concerns about the readability of the paper and the causal status of the investigations. The authors responded to the latter by providing preliminary causal evidence in the author response, and promised to substantially improve the paper's readability to address the former. 

Since no major extra content beyond the submission and the rebuttal will be required to address the reviewers' main concerns, I am inclined to accept this paper but request that authors indeed follow through with the promised changes, listed below for accountability (though perhaps I have missed some):

1. Improve clarity and accessibility of Sections 3 and 4.

2. Include comparisons to human behavioral data and evaluations with more models presented in the global response.

3. Address the causal analysis more carefully.

4. Expand the discussion on practical implications of their findings for both mechanistic interpretability and neuroscience fields.

 I also request the following improvements in formatting for the camera-ready:

- Figures at the top of the page, rather than breaking up text.

- Better global placement of figures near their relevant text, especially Figure 3.

- Some other formatting solution for Section 4.2 than run-in numbering, which makes this poorly readable as a recipe.

Reviews: review: paper_summary: paper_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.
strength_weakness: strength_weakness: 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: questions: I have mentioned them in the limitations section below.
limitations: 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?
paper_summary: paper_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.
strength_weakness: strength_weakness: 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: 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.
limitations: 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.
paper_summary: 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.
strength_weakness: strength_weakness: 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: 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?
limitations: limitations: They are discussed appropriately.
paper_summary: paper_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.
strength_weakness: strength_weakness: 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: 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?
limitations: limitations: Stated in the paper.

Comments: comment: Thank you for your feedback and for raising the score. We appreciate your recognition of the project's potential. We understand the limitations of the rebuttal process and ensure that the revised paper will address concerns about Sections 3 and 4, integrating the various components more cohesively. We are committed to ensuring our work provides clear, valuable insights to the community.
Thanks for acknowledging our effort. We will discuss contextualization of episodic memory mechanisms in language modeling and ICL, e.g., on relating episodic memory to flexible language processing in humans (Duff & Brown-Schmidt, 2012). 

We will include the discussion of implications and limitations of our causal analysis using the ICL score. We agree that our results demonstrate the causal importance of CMR-like heads, while it would be more challenging to narrow down the causal role of CMR-like characteristics (e.g., ablating specific characteristics but keeping all other intact). More generally, the ablation of selected heads remains the primary method in the MI community to test a causal claim about the characteristics of attention heads on model functions (Nanda et al., 2023; Olsson et al., 2022; Wang et al., 2022; Meng et al., 2023; Chan et al., 2022). Additionally, we note that the ICL score is one of the main metrics that the MI community uses to measure the ICL ability (Lee et al., 2024; Olsson et al., 2022), but we agree that a systematic evaluation of ICL ability using multiple datasets and tasks will provide stronger evidence. These limitations and considerations will be explicitly addressed in our discussion section.
Thank you all for your valuable feedback. We appreciate the recognition of our work's potential impact and interdisciplinary approach. In addition to the analyses in the rebuttal, we acknowledge the concerns about Sections 3 and 4 and commit to thorough implementations of feedback to enhance accessibility. Our goal is to present a final version that addresses all concerns, improves clarity, and strengthens the paper's impact across the diverse NeurIPS community.
Thank you for your thorough review and constructive feedback throughout this process. We're grateful for your recognition of our additional analyses and the strengthening of our paper. We appreciate your guidance and we commit to implementing these revisions effectively.
We appreciate your acknowledgment of our response and the increase in your score. We understand the limitations of the review process. Our final revision will address all the concerns raised, integrating the clarifications provided in our rebuttal. We are committed to providing clear value to the community.
Thank you for your comprehensive and thoughtful rebuttal. I appreciate the additional analyses you've conducted, particularly on more recent models like Llama3-8B, Mistral-7B, and Qwen-7B, as well as the ablation study demonstrating the causal link between CMR-like heads and in-context learning performance. Your inclusion of human behavioral data comparisons and expanded discussion on biological plausibility have also strengthened the paper significantly.

I believe the other reviewers have raised valid concerns, particularly regarding the clarity of Sections 3 and 4. I strongly advise revising these sections based on the collective feedback we've provided. Given your thorough responses and plans for improvement, I trust you will address these issues effectively.

Considering these planned improvements, I am increasing my score to 7 (Accept). In case the paper gets accepted, I believe all these changes (if implemented in the final version) will make your work more impactful and valuable to the NeurIPS community. Good luck!
Thank you for the revisions. While I appreciate the efforts made, it's challenging to assess the changes in Sections 3 and 4 within the constraints of the NeurIPS rebuttal which I believe is key for acceptance. I do believe this project has a lot of scope but there is a lot of moving parts and it needs to be better merged together for the community to benefit from these insights. However, I believe the paper is now in better shape, and I am raising my score to a 4.
Thanks a lot for your reponse. I have increased my score to 4 but also acknowledge my high uncertainty. To make a more certain judgement, I would need to see the fully revised paper (which sadly is not possible in this review process).
Thank you for these detailed responses. I appreciate the effort going into the additional experiments on different models and datasets. I do think that the causal claim should be handled carefully. The ablation shows that low-CMR-distance heads are more important for ICL vs. random heads, but it doesn't directly transfer to a conceptual causal link about the removal of episodic-memory-like characteristics in the model to be responsible for the drop in ICL abilities, as it could be something else about these heads and their interaction with other heads (especially given that the ICL score used is only a heuristic metric; though I am not familiar to what extent this particular score is used to approximate ICL abilities in the interpretability field. I imagine ideally we would want a set of datasets/tasks to derive a performance-based ICL metric). In general, as the authors mentioned in the response and are exploring in other projects, I think if the authors can contextualize the results in more discussions on what and why episodic-memory-like mechanisms may be key to language modeling and ICL, it would enhance the paper further.