Linking In-context Learning in Transformers to Human Episodic Memory
Authors: Li Ji-An
Corey Yishan Zhou
Marcus K. Benna
Marcelo G Mattar
Venue: NeurIPS 2024 poster (2024)
Keywords: in-context learning
Transformer
induction head
episodic memory
mechanistic interpretability
Primary area: neuroscience_and_cognitive_science
TL;DR: This paper explores a striking similarity between the in-context learning mechanisms of Transformer models and human episodic memory, revealing parallel computational processes in artificial and biological intelligence systems.

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.

Linking In-context Learning in Transformers to Human Episodic Memory
Li Ji-An
Corey Zhou
Neurosciences Graduate Program University Californi, San Diego ji1095Gucad.ens edu
Department of Cognitive Science University of California, San Diego yiz329Qucsd. edu
Marcus K. Benna Department of Neurobiology University of California, San Diego mbenna@ucsd. edu
Marcelo G. Mattar
Department Psychology
New York University marcelo. wwttarrnyu.....mmmmm edu
Abstract
Understanding connections between artificial and biological intelligent systems can reveal fundamental principles of general intelligence. While many artificial intelligence models have neuroscience counterpart, such connections are largely missing Transformer models and the self-attention mechanism. Here, we ex- amine the relationship between interacting attention heads and human episodic memory. We focus on induction heads, which contribute to in-context learning Transformer-bast 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 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.
Introduction
Neural networks often bear striking similarities to biological intelligence. For instance, convolutional networks trained computer vision tasks can predict neuronal activities in the visual cortex [+ 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- 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 '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 [ remains limited, with many researchers in neuroscience, cognitive science, and deep learning viewing as being fundamentally different from the brain. So far, Transformers have been linked to generalized Hopfield networks (10), neural activities in the language cortex [1HE3.
Equua contribution. Co-senior author.
38th Conference on Neural Information Processing Systems (NcurIPS 2024).

---

d
logits
ligits
Linembed
rembed
Light
Table
Water
Cat
(word)
Time
H2
MLP
1N
(conlext)
Col
Table
embed
embeed
Watar
Study phose
Recill
position
position
Figure 1: Tasks and model architectures. (a) Next-token prediction task. The ICL of pre-trained LLMs is evaluated on a sequence of repeated random tokens (okens("...(AT]" FF..AA][B][CID]]..........]..... [AIBB]]III[ ". c.g., [A]=light, [B] ccat, [C]=table, [D]=water) by predicting the next token (e.g., "... [A][B][C][]] ... "..[B]+++ ?). (b) Human memory recall task. During the study phase, the subject sequentially presented with list of words memorize. During the recall phase, the subject is required to recall the studied words in any order. (c) Transformer architecture, centering on the residual stream. The blue path is the residual stream of the current token, and the grey path represents the residual stream of a past token. Hi and H2 are attention heads. MLP is the multilayer perceptron. (d) Contextual maintenance and retrieval model. The word vector f is retrieved from the context vector t via M and the context vector is updated by the word vector via MFT (see main text for details).
and hippocampo-cortical circuit representations 14] However, 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 mesk-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 ICI. is important for developing better models capable of performing unseen tasks, as well as for Al 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 [.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 novel bridge between Transformer models and episodic memory.
2 Next-token prediction and memory recall
Transformer models language modeling are often trained to predict the next token [6]. 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 on a sequence of repeated random tokens [15] (Fig. [b). For example, consider the prompt "[A][B [[CID][A][][[[I]] these tokens has been learned, the first occurrence of each token cannot be predicted - c.g., the first [C] cannot be predicted to follow the first [B]. At the second [B], however, model with ICL should predict [C] to follow by retrieving the temporal association in the first part of the context.
Assuming that no structure between
Much like ICL in a Transformer model, human cognition is also known to perform associative retrieval when recalling episodic memories. Episodic retrieval commonly studied using the free recall paradigm [18] [9] (Fig. [)) 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 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 l-l (see e.g., Fig. Fa).
3 Transformer models and induction heads
3.1 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 [2] Each token at position d in the input has its own residual stream 24 (with a dimension of dmodel) serving as a shared communication channel between model components at different layers (Fig. II: the residual stream is shown as a blue path), such as self-attention and multi-layered perceptrons (MLP). The initial
22).
residual stream zi (0) contains token embeddings TE (vectors that represent tokens in the semantic space) and position embeddings PE (vectors that encode positions of each input token). Each model component reads from the residual stream, performs computation, and additively writes into the residual stream. Specifically, attention heads at layer read from all past 2j (with j S i) and write
into the current 24 as 2/ < 2, ( Enonds h (h) (2(1-1), (2 (())))))
(""))))))) Bys)). while MLP layers read from only the current 2; and write into 2; as 2. + 2,- + MLP(22 (1
((11)) ). Readers unfamiliar with attention heads (c.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, first-layer head H1 may write its output into the residual stream, which is later read by a second-layer head H2 that writes its output to the residual stream for later layers to use.
3.2 Induction heads and their attention patterns
Previous mechanistic interpretability studies identified 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, 224. 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), and increase the probability of the latter that is, after finding "match", it makes "copy" as the predicted next token (... [B][[ ... [B] - (CD. To formalize this match-then-copy pattern, we use the induction-heart matching score (between 0 and 1) to measure the mreti--mathhing 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-heatt matching score and positive copying score.
We first examined the induction behaviors of attention heads in the pre-trained GPT2-small model 116] using the Transformerl.cn library 25 To elicit induction behaviors, we constructed prompt consisting of two repeats of handom-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 BB. Several heads in GPT2-small had a high induction-head matching score (Fig. 24), such as LSHI (layer number 5, head number 1). In the first sequence repeat, this attention head attends mostly to the beginning-of sequence token. In the second repeat, this head shows a clear "induction stripe" (Fig. 2b) where 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 similar pattern as seen in human episodic recall (Fig. 2 Fig. 5p-c). including temporal contiguity (e.g., the average attention score for |lag] < 2 is larger than for [lag] 2 4) and forward asymmetry (e.g., the average
attention score for lag 0 is larger than for lag< 0).
3

---

a
Induction head matching score
b
Attention pattern (LSH1)
Conditinal resporse probebiltt (LSH1)
Seavstinpps FaNE
Destmelion token
Altantion SCOO
Layer
Head
Source
token
Lag
Figure 2: Induction heads the GPT2-small model. (a) Several heads in GPT2 have a relatively
large induction-head matching score. (b) The attention pattern of the L5HI head, which has the largest induction-hest matching score. The diagonal line ("induction stripe") shows the attention from the destination token in the second repeat to the source token in the first repeat. scores the L5HI head averaged over all tokens in the designed prompt as a function of the relative position lag (similar to CRP). Error bars show the SEM across tokens.
(c) The attention
3.3 K-composition and Q-composition induction heads
The matching score and copying score describe the behavior of individual attention heads. However, they do not provide mechanistic understanding of how the induction head works internalyy. 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. Bb. Tab. [T 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. 3R), the first-layer "previous token" head uses the current token's 's position embedding, PE, the query, and a past token's 's position embedding PE,, as the key. When the match condition Р = - is satisfied (meaning j is the previous position of i), the head writes the previous token's token embedding, FEj, as the value into the residual stream 2,. The second-layer induction head uses the current token's T& as the query. and the previous token head's output T F, at residual stream 2; the key ("K-composition") When the match condition TE, = TÉ is satisfied, the head writes E'E, (at residual stream z;) as the value into the residual stream =k. effectively increasing the logit for the token that occurred position
In Q-composition (Fig, (D), the first-layer "duplicate token" head uses the current token's T& as the query, and past token's 's ТE as the key. When the match condition TE, TEg is satisfied (meaning token is a duplicate of token j), the head writes the token's 's PEş as the value into the induction head uses the duplicate token head's output PE
residual stream 2g. The second- at residual stream 24 as the query ""O-coompoittion") and a past token's 's PF; as the key. When the match condition PE,- PE;-1 is satisfied, the head writes T TF; (at residual stream 2,) the value into the residual stream žy., 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 Q-composition induction head performing prefix matching, allowing us to link induction heads' attention biases to those known in human episodic memory.
4 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 (66). CMR explains the asymmetric contiguity bias in human free recall (see Fig. and Fig. SI and has been extended to more complex memory phenomena such semantic [27] and emotional

---

Query
Earditr
K-compostion
Q-composition
CMR
PC (7
l/z
H2
M
H
74 - 76
PS
PS 7kg
PS: FNG
Study
Recall
Figure 3: Comparison of composition mechanisms of induction heads and CMR. All panels correspond the optimal Q-K match condition (j 1). See main text and Tab. [ST for details. (a) K-composition induction head. The first-layer head's output serves the Key of the second -hooddllayrr head. (b) Q-composition induction head. The first-layer head's output serves as the Query of the second-layer head. (c) CMR is similar to Q-composition induction head, except that the context vector 1 is first updated by MFT into i at position j. then directly used at position j + (equal to 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, briefly list the essentials of CMR.
In CMR (Fig. [],, each word token is represented by an embedding vector f (c.g.. one-hot; f, for the i-th word in sequence). The core dynamic that drives both sequential encoding and retrieval is
t; pt St;
where t; is the temporal context at time step i, and !NN is an input context associated with f. B controls the degree of temporal drift between time steps (Bene for encoding/study phase and Srec for decoding/retrieval phase) and is picked to ensure t; has unit norm. Specifically, during the encoding phase, ! represents a pre-experimental context associated with the i-th word as t. Mprof, where MFT pro a pre-fixed matrix. At each time step, a word-to-context (mapping f to t) memory matrix exp learns the association between f, and -; (i.e., MFT is updated by t During the decoding (retrieval) phase, t is mixture of pre-experimental (t f) and experimental contexts IN MFT f.). The proportion of these two contexts is controlled by an additional pre + 7FT
parameter TFT E [0, exp
arises from this slow evolution of temporal context: when exp ff, The asymmetric contiguity bias < 1, t, 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,, t2)) is large if li is small. Additionally, + MFT f only enters the temporal context after time Thus eexx is associated with f, only for j > (asymmetry).
as E = ((1 - 7FT) M
CMR also learns second context-to-word (mapping t back to f) memory matrix TF (updated by each f,t?_,). When an output is needed, CMR retrieves mixed word embedding f.N TF, If are one-hot encoded, we can simply treat T!N (unnormalized) probability distribution over the input tokens. Or, CMR can compute the inner product (f,, f, F!N) for each cached word fy as input to softmax (with inverse temperature 7 11 recall word.
Intuitively, the temporal context resembles a moving spotlight with fuzzy edge: it carries recency- weighted historical information that may be relevant to the present, where the degree of information degradation is controlled by p. Larger fi'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.). If M pre Σi f.f (i.e.. the pre-experimental context associated with each word embedding is the word embedding itself)
and 7FT 0. Eq. I at decoding is reduced to t; pt,-1 3f, ΣE= 3p--2f ff, which is linear combination of past word embeddings.
5

---

Human
CMR
CRP
0.5
0.
0.0
0.3
9.4
8
18g
0.1
9.1
1agg
Lag
Figure 4: The conditional response probability (CRP) as a function of position lags in a human experiment and different parametrization of CMR. (a) CRP of participants (N=171) in the PEERS dataset, reproduced from [29]. "Top 10%" refers to participants whose performance was in the top 10th percentile of the population when recall started from the beginning of the list. They have sharper CRP with a larger forward asymmetry than other subjects. (b) Left, CMR with "sequential chaining" behavior (Beuc = 3rec 1, WFT 0). The recall has exactly the same order the study phase without skipping over any word. Right, CMR with moderate updating at both encoding and retrieval, resulting in human-like free recall behavior (Bene Bree 7
7FT = 0). Recall is more
likely than not to have the same order as during study and sometimes skips words. (c) Same as (b Right) except with 7 0.5 (Left) and FFT 1 (Right). For more examples, see Fig. ST
4.2 CMR as an induction head
We now proceed to map CMR to Q-composition-lik head (see Fig. 3c and Tab. ST for details),
To begin with, note that the word seen at position is the same TEg and the context vector t--1 (before update) at position i is functionally similar to PE, It follows that the set (f,, f;- 1 is functionally similar to the residual stream 2, updated by the head outputs.
CMR experimental word-context retrieval as first-layer self-attention. The temporal context is updated by t!NN ((1 TFT M pre + TFT. ]f; at decoding. The memory matrix M.T. acts as a first-layer duplicate token head, where the current word f; is the query, the past embeddings } make up the keys, and the temporal contexts f associated with each f, are values. This head effectively outputs "What's 's the position (context vector) at which I encountered the same token f,?"
CMR pre-experimental word-context retrieval MLP. The pre-experimental context t.NN MFTf, (retrieved contextual information not present in the experiment) is the output of linear fully-connected layer (functionally similar to MLP; not drawn).
CMR evolution as residual stream updating. The context vector is updated by t; ptx-1 + S Equivalently, the head F updates the information from (fR, ---1 to (fg, - t& At the position during recall, the updated context t contains !N ( t2)) (Fig. 3c).
CMR word retrieval second-layer self-attention. MT where TFF acts second- induction head, where the temporal context tę is the query, the past contexts (t,-1) } make up the keys, and the embeddings f, associated with each t;-1 values. This effectively implements Q-composition (2I), because tg. as the Query, is affected by the output of the first-layer MFT head. exp
The retrieved embedding is fN
CMR word recall as unembedding. The final retrieved word probability is determined by the inner product between the retrieved memory ff and each studied word fy. similar to the unembedding layer generating the output logits from the residual stream.
CMR learning as a causal linear attention head. Both associative matrices of CMR are learned in manner consistent with the formation of causal linear attention heads. Specifically, the word- to-context matrix is updated by MET, M + t--1f," (with MEJ.0 0), associating f, (key) and t;- (value). It is equivalent to a causal linear attention head, because MFT = (Eick 1f ffx Eien (f.f(f fx) Ever sim(f, t. Similarly, the context-to-word

---

1310
ONS GOMG
5.00
OKR
un
3.44
CWn) saparro t.12
54-
CMT obilnce (955)
Figure 5: CMR distance provides meaningful descriptions for attention heads in GPT2. (a-c) Average attention scores and the CMR-fitter attention scores of example induction heads (with a non-zero induction-head matching score and positive copying score). (d) Average attention scores and the CMM--itted attention scores of duplicate token head [30] that is traditionally not considered an induction head but can be well-captured by the CMR. (e) (Top) CMR distance (measured by MSE) and the induction-heal matching score for each head. (Bottom) Histogram of the CMR distance.
matrix, updated by M.P M + ftt (with MT 0), is equivalent causal linear attention head that associates 1 (key) with f. (value).
To summarize, the CMR architecture resembles two-layer transformer with Q-composition linear induction head. It's worth noting that although we cast t, as the position embedding, unlike position embeddings that permit parallel processing in Transformer models, is recurrently updated in CMR (Eq. []]. It possible that Transformer models might acquire induction heads with a similar circuit mechanism, where t, corresponds to autoregressively updated context information in the residual stream that serves as the input for downstream attention heads.
Experiments
5.1 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. 2 Fig.5h-c 5p--; also see Fig. S22 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 the mean squared error (MSE) between the head's is average attention scores and its CMR-fitted scores (see Appendix E Fig. -dd).. In essence, we optimized the parameters CMR-fitted scores that minimizes MSE
(Bene, Bree, 7FT, T 1) for each head obtain
At the population level, heads with a large induction-head matching score and a positive copying score also have a smaller CMR distance (Fig. (5), 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. 5t).
Consistent with prior findings that induction heads were primarily observed in the intermediate layers of L.MSs (3I, we found that the majority of heads in the intermediate-to-lee layers of GPT2-small have lower CMR distances (Fig. h.. Fig. S3a; for the choice of threshold see Appendix E.5). We also observed similar phenomena in different set of LLMs called Pythia (Fig. (6)) family of models with shared architecture but different sizes, as well as three well- well-known models (Qwen-7B [32]
Mistral-7B [33. Llama3-8B
34
Fig. (6). We summarized these results in Fig. S3p.
Additionally, to contextualize these CMR distances, we included the Gaussian distance using Gaussian functions as 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
7

---

200m
420m
60
60
3.4b
8
20
100
100
100
20
Figure 6: CMR distances vary with relative layer positions in LLMs. (a-b) Percentage of heads with CMR distance less than 0.5 in different layers. Also see Fig. S3 for the threshold of 0.1. (a) GPT2-small. (b) Pythia models across different model sizes (label indicates the number of model parameters). CMR distances are computed based on the last model checkpoint. (c) Qwen-7B, Mistral- 7B. and Llama-8B models. Heads with lower CMR distances often emerge in the interaediate-to-late layers.
0.1
250
0.5
0.4
1:0
58
1.0
Chockpoiit
Chedpciit
Amm
0.6
1.0
0.6
20 tpp 10
0.2
tp 100
top
0.0
0.0
100
50
100
50
100
50
100
Checkpoint
Crecksein
Chockpoint
Checkpoint
g
DKR calarcs
0.7
1.0
0.4
top 20
0.2
Iap 58
top
Tap
0.0
Chockporit
Crectsoie
Chedipsin
Checkponnt
Figure 7: Strong asymmetric contiguity bias arises as model performance improves. (a) Model loss on the designed prompt as function of training time. Loss is recorded every 10 training checkpoints. (b) Average fitted inverse temperature increases in the intermediate layers of Pythia-70m as training progresses. Values are averaged across heads with CMR distance lower than 0.5 in each layer. (c) Comparison of fitted Senc and Srec in Pythia's 's top CMR-like heads and in existing human studies. (d) CMR distance of top induction heads in Pythia models as a function of training time. Heads are selected based on the highest induction-
matching scores across all Pythia models (e.g., "top 20" corresponds to twenty heads with the highest induction-head matching scores). (e) Fitted CMR temporal drift parameters Seac(left)) Sree (middle), 7FT (right) as a function of training time in attention heads with the highest Induction-heat matching scores. (f-g) Same as c-d but for top CMR-like heads (e.g., "top 20" corresponds to those with the lowest CMR distances), demonstrating differences between top induction heads and top CMR-like heads. Shaded regions indicate standard error, except (b) which indicates the range (the scale factor 7 is non-negative).
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.
5.2 CMR-like heads develop human-like temporal clustering over training
The Pythia models' different pretraining checkpoints allowed us measure the emergence of CMR- like behavior. As the model's loss on the designed prompt decreases through training (Fig. Z 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. (4p-b). This, combined with an increasing inverse temperature (Fig. 7)) suggests that attention patterns become more deterministic over training, while shaped to mirror human-like asymmetric contiguity biases. fact, human subjects with better free recall performance tend to cxhibit stronger temporal clustering and higher inverse temperature [29]
8

---

Original
ICL score
0
Top 10% ablated
Randem a ablated
3072--ma4... quetiTT pythtt 1 2ythar-2 pythan 2yt 2ytha-100m pythe A2 Byshion3a
Figure 8: CMR-like heads are causally relevant for ICL. ICL scores are evaluated for intact models (Original), models with the top 10% CMR-like heads ablated (Top 10% ablated), and models with randomly selected heads ablated (Random ablated). Lower scores indicate better ICL abilities, with error bars showing SEM across sequences. ***: p < 0.001, **: p < 0.01, *: p < 0.05, n.s.: p 2 0.1.
For individual heads, those with higher induction-head matching scores (Fig. 7 (or similarly with smaller CMR distances, see Fig. 7 consistently exhibit greater temporal clustering (Fig. Ze. f respectively), as the fitted 3's (both Benc and Bree) were large. We also observed similar values of fitted 3s in Qwen-7B, Mistral-7B and Llama3-88 (Fig. Sop). These fitted I's of these attention heads fall into a similar range as human recall data (Fig. 7p.. We interpret this in light of a normative view of the human memory system: in humans, the asymmetric contiguity bias with a B < is not merely phenomenological; under the CMR framework, it gives rise to an optimal policy to maximize memory recall when encoding and retrieval noisy [29] In effect, a large B (but less than 1) in Eq. [ provides meaningful associations beyond adjacent words to facilitate recall, such that even if the immediately following token is poorly encoded the agent fails decode it, information from close-by tokens encapsulated in the temporal context still allows the agent to continue decoding.
In addition, observed an increase in Senc and Sree (Fig. 7e. g) during the first 10 training checkpoints, when the model loss significantly drops. The training process leads to higher values of Bree: Specifically, Srec values are higher than Bene highlighting the importance of temporal clustering during decoding for model performance. These results suggest that attention to temporally adjacent tokens with asymmetric contiguity bias may support ICI. in LLMs.
5.3 CMR-like heads are causally relevant for ICL capability
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) 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 B5 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 [5)- Intuitively, a model with better in-context learning ability has 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 like heads 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, Hydra effect has been noted where ablation of heads causes other heads to compensate [36). Second, our analysis cannot confirm 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-memon memory features in these CMR-like heads might causally contribute 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, larger distributional differences between their training data and our evaluation texts.
Discussion
This study bridges L.MMs 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 causal role in ICI. These findings offer significant connections between the current generation of AI algorithms and 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 (7-39). While CMR is 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 CAI, dentate gyrus [40]. CA3 [4]), and cortical areas projecting to the hippocampus such as the entorhinal cortex 42
Our results show strong connections neural network models of episodic memory. For example, neural networks with attention recall pattern as the optimal CMR. [45] showed that neural network implementation of CMR can explain humans' flexible cognitive control. Our results also align with research connecting attention mechanisms in Transformers 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-valle attention mechanisms link to biological episodic retrieval, suggesting that CMR-like behavior emerges naturally in neural networks with proper objectives.
4 or recurrence 44 trained for free recall exhibit the same
the hippocampal formation 14] While prior work focused on
The link to induction heads could enable researchers to develop alternative mechanisms for episodic memory. For instance, K-composition and O-composition induction eircuits 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 [4745] might be understood based on more complex attention-head composition mechanisms (e.g.. N-th order virtual attention head (1)).
From the perspective of LLM mechanistic interpretability, we offer more detailed behavioral description and reinterpret induction heads through the lens of CMR and the asymmetric contiguity bias. First, CMR-like heads not limited induction heads some attention heads are well captured by CMR despite not meeting the traditional induction head criterion (e.g., Fig. 5tD) Second, speculate that heads with low CMR distances and low induction-heant 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.I 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 be explored. Overall, our results suggest fuller view of ICI. mechanisms in L.Mss, 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 56). suggesting these heads may process hierarchically organized text chunks [57]. Third, episodic mechanisms captured by CMR are posited to support adaptive control [7 49 and flexible reinforcement learning [50] suggesting similar
roles for these heads more complex cognitive functions of L.MMs. 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 -ike characteristics. In addition, while ICL score is a primary metric for measuring ICL ability [15] 59]. systematic evaluation using specific datasets and tasks would allow for a stronger causal claim. Addressing these limitations is key future direction.
10

---

Acknowledgements
M.K.B was supported by NIH ROINS125298. M.K.B and L.J.A. were supported by the Kavli Institute for Brain and Mind. The authors thank the Apart Lab for Interpretability Hackathon 3.0, and thank the support from Swarma Club and Gost-ChaIGPT Reading Group supported by the Save 2050 Programme jointly sponsored by Swarma Club and X-Order. In addition, the authors thank X. Li, H. Xiong, and the anonymous reviewers for their feedback.
References
[1] Daniel LK Yamins, Ha Hong, Charles F Cadieu, Ethan A Solomon, Darren Seibert, and James DiCarlo. Performance-optmmized hierarchical models predict neural responses in higher visual cortex. Proceedings of the national academy of sciences, 111 111(23):861-8622 2014.
[2] Daniel LK Yamins and James DiCarlo. Using goal-driven deep learning models understand sensory cortex. Nature neuroscience, 19(3):356-365. 2016.
[3] Martin Schrimpf, Jonas Kubilius, Ha Hong. Najib J Majaj, Rishi Rajalingham, Elias B Issa, Kohitij Kar, Pouya Bashivan, Jonathan Prescott-Roy Franziska Geiger, et al. Brain-score: Which artificial neural network for object recognition most brain-like? BioRxiv, page 407007, 2018.
[4] Sun Minni, Li Ji-An, Theodore Moskovitz, Grace Lindsay, Kenneth Miller, Mario Dipoppa. and Guangyu Robert Yang. Understanding the functional and structural differences across excitatory and inhibitory neurons. bioRxiv, page 680439, 2019.
[5] Christopher Cueva and Xue-Xin Wei. Emergence grid-like representations by training recurrent neural networks to perform spatial localization. arXiv preprint arXiv:1803.07701 2018.
[6] Andrea Banino, Caswell Barry, Benigno Uria, Charles Blundell, Timothy Lillicrap. Piotr Mirowski, Alexander Pritzel, Martin Chadwick, Thomas Degris, Joseph Modayil, et al. Vector- based navigation using grid-like representations in artificial agents. Nature, 557(7705):499-433 2018.
[7] Jane X Wang. Zeb Kurth-Nelson Dharshan Kumaran, Dhruva Tirumala, Hubert Soyer, Joel Z Leibo, Demis Hassabis, and Matthew Botvinick. Prefrontal cortex as a meta-reinforcemenn learning system. Nature neuroscience, 21(6))860-88) 868, 2018.
[8] Catherine Hanson, Leyla Roskan Çaglar, and Stephen José Hanson. Attentional bias in human category learning: The case of deep learning. Frontiers in Psychology, 9, 2018.
[9] Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N. Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. In Isabelle Guyon, Ulrike von Luxburg, Samy Bengio, Hanna M. Wallach, Rob Fergus, V. N. Vishwanathan, and Roman Garnett, editors, Advances in Neural Information Processing Systems 30: Annual Conference on Neural Information Processing Systems 2017, December 4-9, 2017, Long Beach, CA, USA, pages 5998-6008. 2017.
[10] Hubert Ramsauer, Bernhard Schäfl, Johannes Lehner, Philipp Seidl, Michael Widrich, Thomas Adler, Lukas Gruber, Markus Holzleitner, Milena Pavlović, Geir Kjetil Sandve, et al. Hopfield networks is all you need. arXiv preprint arXiv:2008.022 02211, 2020.
[11] Ariel Goldstein, Avigail Grinstein-Dabush Mariano Schain, Haocheng Wang. Zhuoqiao Hong, Bobbi Aubrey, Mariano Schain, Samuel A Nastase, Zaid Zada, Eric Ham, et al. Alignment of brain embeddings and artificial contextual embeddings in natural language points to common geometric patterns. Nature communications, 15(1):2768, 2024.
[12] Martin Schrimpf, Idan Blank, Greta Tuckute, Carina Kauf, Eghbal A Hosseini, Nancy Kan- wisher, Joshua Tenenbaum, and Evelina Fedorenko. The neural architecture of language: Integrative reverse-engineering converges on model for predictive processing. BioRxiv, 2020.
11

---

[13] Richard Antonello, Aditya Vaidya, and Alexander Huth. Scaling laws for language encoding models in fmri. Advances in Neural Information Processing Systems, 36, 2024.
[14] James CR Whittington, Joseph Warren, and Timothy EJ Behrens. Relating transformers to mod- els and neural representations of the hippocampal formation. arXiv preprint arXiv:2112.04033, 2021.
[15] Catherine Olsson, Nelson Elhage, Neel Nanda, Nicholas Joseph, Nova DasSarma, Tom Henighan, Ben Mann, Amanda Askell, Yuntao Bai, Anna Chen, et al. In-context learning and induction heads. arXiv preprint arXiv:2209. 1895, 2022.
[16] Alec Radford, Jeffrey Wu, Rewon Child, David Luan, Dario Amodei, Ilya Sutskever, et al. Language models are unsupervised multitask learners. OpenAl blog, 18)):, 2019.
[17] Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared D Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, et al. Language models are few-shot learners. Advances in neural information processing systems, 33: 33:1877-1901, 2020.
[18] Murdock and B Bennet. The serial position effect of free recall. Journal of Experimental Psychology, 64:482-488, 1962.
[19] Mare W Howard and Michael Kahana. A distributed representation of temporal context. Journal of mathematical psychology, 46(3):269-29) 299, 2002.
[20] Roger Ratcliff and Gail McKoon. Does activation really spread? Psychological Review, 88(5): 454 462, 1981, doi: 10. 103700033-299935 455770033--95x.8..5555..4
[21] Nelson Elhage, Neel Nanda, Catherine Olsson, Tom Henighan, Nicholas Joseph, Ben Mann, Amanda Askell, Yuntao Bai, Anna Chen, Tom Conerly, et al. A mathematical framework for transformer circuits. Transformer Circuits Threud, 1, 2021.
[22] Rupesh Kumar Srivastava, Klaus Greff, and Jürgen Schmidhuber. Highway networks. arXiv preprint
2015.
[23] Gautam Reddy. The mechanistic basis of data dependence and abrupt learning in an in-context classification task. arXiv preprint 0rXiv:2312.00300 2023.
[24] Aaditya Singh, Ted Moskovitz, Felix Hill, Stephanie CY Chan, and Andrew M Saxe. What needs to go right for an induction head? a mechanistic study of in-context learning circuits and their formation. arXiv preprint (00::24040712999 2024.
[25] Neel Nanda. Transformerlens, 2022.
URL https://githab..cmm/nelllnaaaaa
ormeeLess
[26] Sean M Polyn, Kenneth A Norman, and Michael J Kahana. A context maintenance and retrieval model of organizational processes in free recall. Psychological review, 116(1):129, 2009.
[27] Lynn Lohnas, Sean Polyn, and Michael Kahana. Expanding the scope of memory search: Modeling intralist and interlist effects in free recall. Psychological review, 1 122:337-363, (44 2015. doi: 10. 1037/200339036.
[28] Rivka T Cohen and Michael Kahana. Retrieved-conter theory of memory in emotional disorders. bioRxiv, page 817486, 2019.
[29] Qiong Zhang, Thomas L. Griffiths, and Kenneth A. Norman. Optimal policies for free recall. Psychological Review, 130(4): 104-1124, 2023. doi: 10. 00.10377/e000003355
[30] Kevin Wang, Alexandre Variengien, Arthur Conmy, Buck Shlegeris, and Jacob Steinhardt. Interpretability in the wild: circuit for indirect object identification in gpt-2 small. arXiv preprint arXiv:221 1.00:22 00593, 2022.
[31] Hritik Bansal, Karthik Gopalakrishnan, Saket Dingliwal, S. Bodapati, Katrin Kirchhoff, and Dan Roth. Rethinking the role of scale for in-context learning: An interpretabilltt-baase case study at 66 billion scale. Annual Meeting of the Association for Computational Linguistics, 2022.
12

---

[32]
Jinze Bai, Shuai Bai, Yunfei Chu, Zeyu Cui, Kai Dang, Xiaodong Deng, Yang Fan, Wenbin Ge, Yu Han, Fei Huang, et al. Qwen technical report. arXiv preprint arXiv:2309.1666 2023.
[33] Albert Q Jiang, Alexandre Sablayrolles, Arthur Mensch, Chris Bamford, Devendra Singh Chaplot, Diego de las Casas, Florian Bressand, Gianna Lengyel, Guillaume Lample, Lucile Saulnier, et al. Mistral 7b. arXiv preprint arXiv:2310.06622, 2023.
[34] Abhimanyu Dubey, Abhinav Jauhri, Abhinav Pandey, Abhishek Kadian, Ahmad Al-Dahle, Aiesha Letman, Akhil Mathur, Alan Schelten, Amy Yang, Angela Fan, et al. The llama 3 herd of models. arXiv preprint arXiv:2407.21783 2024.
[35] Colin Raffel, Noam Shazeer, Adam Roberts, Katherine Lee, Sharan Narang, Michael Matena, Yanqi Zhou, Wei Li, and Peter J Liu. Exploring the limits of transfer learning with a unified text-to-text transformer. Journal of machine learning research, 21(140):1-67, 2020.
[36] Callum McDougall, Arthur Conmy, Cody Rushing, Thomas McGrath, and Neel Nanda. Copy suppression: Comprehensively understanding an attention head, 2023, URL htthtt:sp///xii 2ttbbs//211000455
[37] Kimberly L Stachenfeld, Matthew M Botvinick, and Samuel J Gershman. The hippocampus as a predictive map. Nature neuroscience, 20(1 1):1643-1653, 2017.
[38] Anna C Schapiro, Nicholas B
Matthew M Botvinick, and Kenneth A Norman. Complementary learning systems within the hippocampus: neural network modelling approach to reconciling episodic memory with statistical learning. Philosophical Transactions of the Royal Society B: Biological Sciences,
2017.
[39] Li Ji-An, Fabio Stefanini, Mareus K. Benna, and Stefano Fusi. Face familiarity detection with complex synapses. Iscience, 26(1), 2023.
[40] John Sakon and Michael Kahana. Hippocampal ripples signal contextually mediated episodic recall. Proceedings of the National Academy of Sciences, 1 119(40)=e2011657119 2022.
[41] Halle R. Dinsdale-Zucker, Maureen Ritchey, Ame D) Ekstrom, Andrew P. Yonelinas, and Charan Ranganath. Cal and ca3 differentially support spontaneous retrieval of episodic contexts within human hippocampal subfields. Nature Communications, 2020.
[42] Mare W Howard, Mrigankka S Fotedar, Aditya V Datey, and Michael E Hasselmo. The temporal context model in spatial navigation and relational learning: toward a common explanation of medial temporal lobe function across domains. Psychological review, 1 112(1):75, 2005.
[43] Nikolaus Salvatore and Qiong Zhang, Parallels between neural machine translation and human memory search: A cognitive modeling approach. In Proceedings of the Annual Meeting of the Cognitive Science Society, 2024.
[44] Moufan Li, Kristopher T Jensen, and Marcelo G Mattar. A neural network model trained on free recall learns the method of loci. In Proceedings of the Annual Meeting of the Cognitive Science Society, volume 46, 2024.
[45] Tyler Giallanza, Declan Campbell, and Jonathan D Cohen. Toward the emergence of intelligent control: Episodic generalization and optimization. Open Mind, 8:688-722, 2024.
[46] Benjamin J Kraus, Mark P Brandon, Robert J Robinson, Michael A Connemny,, Michael E Hasselmo, and Howard Fichenbaum. During running in place, grid cells integrate elapsed time and distance run. Neuron, 88(3):578-85 2015.
[47] Qihong Lu, Ali Hummos, and Kenneth Norman. Episodic memory supports the acquisition of structured task representations, 2024.
[48] Tyler Giallanza, Declan Campbell, and Jonathan D. Cohen. Toward the Emergence of Intelligent Control: Episodic Generalization and Optimization. Open Mind, 8:688 722, 2024. ISSN 2470- 2986. doi: 10. 10.1162/opmi a_ 00143. URL https://doiiiii http::::///ooi.r//1.... 1162/opmi. oopii. Lammii___000 a_0443
[49] Dharshan Kumaran and James I.. McClelland. Generalization through the recurrent interaction of episodic memories. Psychological Review, 119:573 - 616, 2012.
13

---

[50] Corey Y. Zhou, Deborah Talmi, Nathaniel D. Daw. and Marcelo G. Mattar. Episodic retrieval for model-based evaluation sequential decision tasks, 2023.
[51] Samuel J. Gershman and Nathaniel D. Daw. Reinforcement learning and episodic memory in humans and animals: An integrative framework. Annual Review of Psychology, 68:101-128, 2017.
[52] Koyena Pal, Jiuding Sun, Andrew Yuan, Byron Wallace, and David Bau. Future lens: An- ticipating subsequent tokens from single hidden state. In Jing Jiang, David Reitter, and Shumin Deng, editors, Proceedings of the 27th Conference on Computational Natural Lan- guage Learning (CoNLL), pages 548-560. Association for Computational Linguistics, 2023. doi: 10. 21118633/122223.com 1.37. URL UBRLLtttp:://a/lmmttoooyy wrg.llooy..o.o//// 2023g/2223..2 wrtg//22..oonn111.... con11-1.37
[53] Nelson F Liu, Kevin Lin, John Hewitt, Ashwin Paranjape, Michele Bevilacqua, Fabio Petroni, and Percy Liang. Lost in the middle: How language models use long contexts. Transactions of the Association for Computational Linguistics, 12:157-173, 2024.
[54] Dillon Murphy, Michael Friedman, and Alan Castel. Metacognitive control, serial position effects, and effective transfer to self-paced study. Memory & Cognition, 50:144-159, 2021. doi: 10.3758//13441--22--00002
[55] Neal W Morton and Sean M Polyn. A predictive framework for evaluating models of semantic organization in free recall. Journal of memory and language, 86:119-1400 2016. doi: 10.1016/j. jml.2015.10.002.
[56] Michelangelo Naim, Mikhail Katkov, Stefano Recanatesi, and Michail Tsodyks. Emergence of hierarchical organization in memory for random material. Scientific Reports, 9:10448, 07 2019. doi: 00038//41998-01-444400
[57] Christopher Baldassano, Janice Chen, Asieh Zadbood, Jonathan W. Pillow, Uri Hasson, and Kenneth A. Norman. Discovering event structure in continuous narrative perception and memory. Neuron, 95:709-721.cs 2016.
[58] Leo Kozachkov, Ksenia V Kastanenka, and Dmitry Krotov. Building transformers from neurons and astrocytes. Proceedings of the National Academy of Sciences, 120(34):c2211151120, 2023.
[59] Ivan Lee, Nan Jiang, and Taylor Bong-Kirkpattrck. Is attention required for icl? exploring the relationship between model architecture and in-context learning ability. In The Twelfth International Conference on Learning Representations.
[60] Joy Crosbie and Ekaterina Shutova. Induction heads as an essential mechanism for pattern matching in in-context learning. arXiv preprint arXiv: 2024.
14

---

A
Code
All code is available allcppe//ggbatb..ooo//rrrr//////
B
Attention mechanisms in Transformer models
For each attention head h, the standard self-attention mechanism calculates its output as (2 (2f (2 whth) Ej w/th 2 A where W-O is the dmodel dhead output projection matrix, with)) is the dhead x dmoddl value projection matrix, and A is the attention pattern attending from the query position i to the key position j. The attention pattern matrix is A Softmax(A(A)
Softmax (z) (t-1) WJWK22)-1) /Vduead), where A represents the attention score matrix; W and W are the dhead x diuodel key and query projection matrices respectively.
C Detailed introduction of CMR
In this section, we offer a more detailed description the Context Maintenance and Retrieval (CMR) model of human episodic memory, focusing its components and underlying dynamics. CMR operates in two distinct experimental phases: an encoding/learning phase, where the model memorizes list of words, and 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 (MFT), T, and the context-to- word association matrix TF)..
The temporal context vector (t) continuously evolves over time, incorporating recenty-weighted historical information during both encoding and retrieval. The core dynamics of the temporal context t governed by the update rule Eq.
For illustrative purposes, the rest this section, we consider an example sequence consisting of five distinct one-hot word embeddings fJ, f2, .... f5. Each embedding is represented by 6-dimensional vector, such that f, contains represents 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 & to explicitly mark the two phases.
at the i-th position, and everywhere else. The 6-th dimension
the pre-experimental word-to-context memory matrix M (6 x 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 IN M fy. Let to (0, O, O, O, 0, 1)". Thus given a sequential presentation of fj fz fs. the corresponding context vectors t, are
For simplicity,
t (F
t2 (08,8,0,0.0.p33
t;
where 3 is the temporal drift parameter, and /1 32 ensures that t, has unit norm. Distinet parameters Bene and froe are used by the encoding and retrieval phases respectively to capture the phase-specific temporal dynamics. Conceptually, the temporal context reflects a secency-weighter history of past information, with the parameter / determining the rate of degradation over time.
More generally, the input context determined by both pre-experimental and experimental contexts:
IN MF ((1 7FT M + M
PT
The memory matrix FT is therefore weighted combination of pre-existing associations learned prior to encoding (the pre-experimental memory matrix MET) and associations learned during the encoding phase (the experimental memory matrix MFT), MFT is commonly initialized as the
15

---

identity matrix at the beginning of each experiment and held constant throughout the experiment. On the other hand, Mxx usually initialized as the zero matrix. updated during the encoding phase at the presentation of each new word by associating the context vector f (value) with the embedding of the just encountered word f (key):
M exx f M 1-xp + 1f
(Fig. B first layer).
Suppose at the beginning of the experiment, the experimental word-to-contett matrix M is initialized the zero matrix. At encoding time j = 1, the experimental matrix M is updated by the pairwise association of the previous context and the current word:
M
M
0.epp + t
exxp
= M
0,
[1
0
0
0
0
0
0
0
0
0
0
0
0
0
0
0
0
0
Similarly, at time j 2, the experimental matrix incorporates the new context-word pair
M.T
2,exp
M.T
1,oxp
t
1
0
0
8
0
0
0
0
0
0
0
0
0
0
1
The update procedure continues until the end of the word sequence. Critically, note that the product of M.T j.exp and any f/+1 with > 0 is the zero vector.
During retrieval, the most recently recalled word f (query) retrieves its corresponding input context
%
M FT fg
(1 TPT
M FT pre
la + FFT
FF Cr. exp
(1 YPT M FT + MF
E(t
>FT M
FFT E4-1(
Let denote the word encountered at encoding time with j < k. Assuming fú f, (i.c., the most recent recall fç matches a previously encoded word fz). we have
INN (1 7FT M FT,, + 7FTtJ-1-
tk
16

---

Furthermore, because we assumed M to be the identity matrix,
IN (1 - (F-FFF 7Frts-1-
(S1)
This input context is subsequently reinstated along with the current context f 1 at (retrieval) time k, resulting in updated context t. via Eq. [ previous instance the most recent recall) and 1 (the context immediately before presentation of the previous instance).
Thus t also incorporates information from f (the
Importantly, the context tú at encoding is also a linear combination of fj and 1 (by Eq. E.
f
St + ptz-1,
where t.N = ((1 TPT M MFFTM FT f. However, because all word embeddings are one-hot exp and distinct, M has learned any associations involving the word fj prior time j. Therefore, the second term is zero, and
t
Bf, + ptz-1.
(S2)
Comparing Eq. S2 and Eq. S1 thus reveals that &N thus Lg. become similar tj-
Finally, the context-to- association matrix M is responsible for retrieving the context associ- ated with a given word. M is usually initialized as the zero matrix prior experiments. During encoding, it is updated to associate each word embedding (value) with the corresponding context vector 1 (key):
MPP + M + f,tt-1
Fig.. second layer). In practice, we chose to maintain T as the transpose of exp' Because the updated context ty is most similar to ty (measured by their inner product), the retrieved vector
P!N = M th
will be the most similar to the word embedding fs+1- Thus, CMR will most likely recall fy+1 following a recall of fş = fy. In the main text, we specified i +
Consequently, the CMR model explains two key phenomena observed in human free recall experi- ments: the temporal contiguity effect and the forward asymmetry effect. The temporal contiguity effect arises because (i) t& integrates the information of 1--1, and (ii) the context vector t drifts over time, making 1 more similar to context vectors associated with nearby words in the sequence, favoring the retrieval of words close to fz. In addition, the forward asymmetry effect occurs because (i) ta integrates the information of !NN and (ii) t. contributes to t, -4-1 and its following context vectors, but does not affect earlier context vectors. a result, t. preferentially retrieves words that follow fj. rather than those that precede it.
D Prompt design
We used the prompt of repeated random tokens due to several reasons (1) it aligns with human free recall experiments, where random words presented sequentially [8]] (2) it is widely acknowledged definition of induction heads in mechanistic interpretability literature [2]] 60]; (3) it uses off-distribution prompts to focus abstract properties, avoiding potential confounds from normal token statistics (1
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 total length of 2N + (one end-of-sequeed token at the beginning). The two copies correspond to the study (encoding) and recall phases, respectively.
17

---

E Metrics & definitions
E.1 Metrics for induction heads
Formally, we define the induction-head target pattern (i.e., attention probability distribution of an ideal induction head) A (ād,a) over sequence of tokens {x;}
as
f1 if 1- va and < <d 0dd (0
otherwise
Here, d is the destination position and is the source position. We then assess the extent to which each attention head performs this kind of prefix matching [15 Specifically, the induction-head matching score for head with attention pattern A (ddds) defined as
ΣΣ od e
// ad.s) E [0,1]
(Fig- 2b. A head that always performs ideal prefix matching will have an induction-head matching score of 1.
Additionally, an induction head should write the current residual stream to increase the correspond- ing 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 circuit W WoWWoWW WW-Wr WE for one head, where We defines token embeddings, Wv computes the value of each token from the residual stream (i.e., aggregated outputs from all earlier layers), Wo computes the head's output using a linear combination of the value vectors and the head's 's attention pattern, and Wy unembeds the output to predict the next token. The copying score of the head is equal to
xx) Σ (el) E [-1,1],
where Ag's are the eigenvalues of the matrix W. Since copying requires positive eigenvalues (corresponding to increased logits), induction head should have a positive copying score. An ideal induction head will have copying score of
E.2
CRP analysis of attention heads
For head with attention scores A, the average attention score alag defined as
Σ
Ölag
No+Naa
N
lagl 2
[lag] <eS [lag]
where N is the length of the first repeat in the prompt. Thus if lag uag quantifies how much the first instance of a token is attended to on average; for lag = 1, clng quantifies the average amount of attention distributed to the immediately following token in the first repeat. We used lag [-5,5] throughout the paper.
E.3
CMR distance
The metric CMR distance is defined as
(glag
2gg )2/22222 /Va (árag)
dCMR min
Nua
Here, Ning is the number distinct lags, and (lag is obtained by calculating the CRP using CMR with specific parameter values. Note we did not consider the full set of possible but only the subset given by the combinations of parameters Banc 0.05, 0.1,..,, Pree O, 0.05...... 1, and 7FT = 0, 0.1,..... ..... 1 for model fitting.
18

---

E.4
Gaussian distance
The metric Gaussian distance, serving as a baseline, is defined as
(gag
Ögag / /Va (Crag)
dGaussian min
g
Nug
where Olg is the Gaussian function, defined as
(lag 02)2
Plag 01 exp
C4.
203
and C1. C2. Ca. are coefficients.
E.5 Choosing thresholds for CMR distance
To identify heads with CMR-like attention scores and exclude non-CMR-like heads, we operationally chose 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, the bottom panel in Fig. 5 shows, the distribution of individual head CMR distances in GPT2 can divided roughly into two clusters: a cluster peaking around 0. (and extending to around 1), and spread-out cluster around 2.4. Second, examined heads with 1 CMR distance and found non-human-lic CRP (e.g.. Fig. b2.. Further, many heads with CMR distance around 0.6-0.8 again show CRP different from humans (e.g., Fig. [S2p, d). Finally, also tested threshold of 0.1 and found no qualitative difference (Fig. Fig. S3c. d).
19

---

F
Comparison of composition mechanisms of induction heads and CMR
K- K-composition PE, & TE; [ РE & TE, [A PF & TE2 [
Q D-composition PE; & TE, [At PE, TE [AL PE T [At
CMR
Representation
& Ji [Ata]
(residual stream)
4
f4 [At
before H1
fg fu []
word-to-context
Type
previous token head
duplicate token head
matrix M
Query
PE, [AL PE, At 5l
TEA [At
Ju [AU
First-layer
Key
TE, A
J, At
head H1
Optimal
Q-K match
PE PE,
TEJ TE
J3 SK
condition
Activation
Softmax
Softmax
Linear
Value
TE, [AT
PE, [AL
A J
Output
TE, [AUi]
PE, [AL
1, [At k]
Representation (residual stream) after H,
PE, & TH, F. & TE, [Ari]
PF, & TE, [At
4:-1 & 1, & fi [At
PE, & TE; [AL]]
PE, & TE, [At]]
13 & f.y & fy [At
PE & TEg [
PE & TE PE, [ALR]
4--1 lR fx [Ack]
context-to-wond
Тype
induction head
induction head
matrix M
Query
TF4 [
PE, [t
F, At k
second-layer
Key
TE, [AL
PE; [At
( [ s
head H2
Optimal
Q-K match
TE₃- TE
PE PE, 11
4 4-1 ft
condition
Activation
Sofimax
Soltmax
Linear
Value
TE; [At
TE: [At
ft [At
Output
TE, [ALK]"
TE, [AL
fel [A k]
PE & TE, [At
4-11 t, & fa [At
Representation (residual stream) after H/
PE & TE, & TE₁ [AUi] PE₁ & TE, [AL]] PE T & TE [ALR]
PE, & TE, [AL]]
1--1 & t2 f3 [At
PE & TE
fk-1 & fh.
& PE, & TE, [ALA]
& fx & SEN [Atk]
Table S1: Comparison of mechanisms of induction heads and CMR.
PE, position embedding for the token at index/position i.
TE; token embedding for the token at index/position i.
4-11 context vector associated with word at index/position
J: word embedding for the word at index/position i.
At i]: information available at index/position
* The output is approximate (due to the weighted-average over all previous tokens).
:: te is updated from fx-1 by ty-1 (due to MET)) and eJN (due to MFTr), thus containing the
information of t, (updated from 1--1 by r!j")
tt. The optimal Q-K match condition in Q-composition requires transformation from PE PE implemented by the WoWk matrix of H2. The optimal Q-K match condition in CMR requires transformation from 15-1 to ty, implemented by ejN.
G Supplementary discussion
G.1 Scale difference
We have observed the scale difference between raw attention scores and those from human recall. We can speculate 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 '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.
20

---

G.2 Mapping components in Transformer models to the brain
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 MIP and attention layers may be akin to the activation-based working memory quickly updated by the cortico crrtic--ippoccmmal cireuits.
H Experiment compute resources
The induction-hert head matching scores and copying scores of each head in GPTl-small and all Pythia models are computed using Google Colab Notebook. All models were pretrained and accessible through the TransformerLeess library [25] with MIT License and used is. See Table S2 for details.
Transformer Model
Type of compute worker
RAM (GB)
Storage (GB)
Computing time (minutes)
GPT2-small
CPU
12.7
225.8
7ychia-70m-dedupeed-aa
CPU
12.7
225.8
Pychia-160mmdeddupeeea
CPU
12.7
225.8
5
4ytiia--10mm-ddduuedd-00
CPU
12.7
225,8
15
Pythia- Pythia-Ib-dctee Ib-deduped-v()
CPU
12.7
225.8
45
High-RAM CPU
51.0
225.8
56
Pythia-2, 2.8--dddudddde---
High-RAM CPU
51.0
225.8
Pythia-6 6.9b 9b-deduped-v0
TPU v2
334.6
225.3
111
Pythia- 2b-deduped-vv
TPU v2
334.6
225.3
205
Qwen-7b
TPU v2
334.6
225.3
7
Mistral-7b Llama3-8b
TPU v2
334.6
225.3
5
TPU v2
334.6
225.3
Table $2: Details of compute resources used to compute induction head metries. All models were pretrained and accessible through the TransformerL.en 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 heads across checkpoints where available.
used an internal cluster to compute the subset of q (CRP) for model fitting. The internal cluster has 6 nodes with Dual Xeon E5-2699v3, which has 72 threads and 256GB RAM per thread, plus 4 nodes with Dual Xeon E5- 269969,,,, which has 72 threads and 512GB RAM per thread. This computation took total of 90 hours.
Finally, model fitting was done on a 2023 MacBook Pro with 16GB RAM. All experiments were completed within hour regardless of model size. This section contains all experiments we conducted that required non-trivial computational resources.
21

---

I
Additional figures
CH
Av-07 bu- Clab Ad
Clipe A2
Aw Ae te
de- As durts
Au
Ac a1 to Au- +s fosttA-tt -tt Ac -o fo Au- to tt
A
A
A
-21
-05
-21
-00
-21
-11
-51
18g
1ag
Lng
Lag
L9g
Lng
Lgg
Figure $1: Example conditional response probability distributions produced by different param- eterizations of CMR. Black lines correspond 7FT 0 (i.c., only the pre-experimental contexts are used during retrieval). Grey lines correspond to 7FT = (i.c., only the experimental contexts are used during retrieval).
22

---

a
CMR distance 1.2
L4H11
b
CMR distance
LOH3
Penc 0.4, Prec 0.0, vr 0.0
Penc
0.45, 0.0, Yr 0.0
-60
Attentian scores
-65
Attention scores
-20
-70
-21
-75
-80
-23
Lag
Lag
d
CMR distance 0.72
CMR distance 0.54
Pun: 0.45, Ar 0.15, Vrr 0.0
Buc
0.05, frue 0.0 05, Yor 0.1
-3.0
Attention scores
Attention scores
-3.2
-3.4
-3.6
Lag
Lag
Figure S2: Example GPT2 attention heads with CMR distance roughly between 0.5 and 1.0.
a
2.5
% CMR distance 0.1
60
CMR distance
1.5
Count
40
1.0
0.5
50
100
20
40
60
80
100
25
50
75
100
Relative layer position (%)
Relative layer pesition (%)
Relative layer position
60
% CMR distaicce < 0.1
50
60
CMR distannce
40
30
160m
2.8b
410m
10
6.9b
Relative layer position (%)
Relative layer position (96)
Figure S3: CMR distances vary with relative layer positions in LLMs. (a) Average CMR distance across heads in each layer of GPT2-small Error bars indicate standard errors. (b) Histogram of the relative positions of the two layers with the lowest average CMR distances in each of the following models: GPT2-small, all Pythia models, Qwen-7B, Mistral-7B, and Llama3-8B. (c) Percentage of heads with CMR distance less than 0.1 in GPT2-small. (d) Same in (c) except for Pythia models.
23

---

b
1.0
Avg Barc
Avg 9ax
0.5
L2
L3
Avg Bass
Aug Broe
Avg VIT
0.5
0.0
LS
1.0
Avg Bar
Avg Aec
Aug NT
0.0
Chockpoint
Chackpaint
Chockpoint
Chackpoint
Chackpoint
Chockpoint
Figure S4: Fitted CMR parameters of Pythia-70m model over training. Solid lines indicate values averaged across heads with a CMR distance lower than 0.5 in the corresponding layer. Shaded areas indicate the range fitted values (the lower edge indicates the minimum value; the upper edge indicates the maximum value). (a-b) Fitted CMR temporal drift parameters Bene (a), and free (b) in different layers as a function of training time. (c) Fitted experimental context mix parameter YFT in different layers as a function of training time.
24

---

1.2
140
CAR
0.7
221
3:2
100
30
124
16
100
Chacspoit
Choqgent
Oedppet
Chacgpes
20
D40
2.2
050
h
1.30
1.43
1.00
0.4
235
2.2
100
50
100
56
100
100
Chectpoit
Chocopoit
Choctpoit
Ohtppire
Chocopeit
0.50
4 0.50
2.00
6.00
1325
0.93
10
1.00
1.53
1.00
535
123
0.00
1.43
6.00
100
50
339
104
50
100
300 Cholpeet
Checlecet
Coocpeott
Chocooot
Occhpett
0.0
4 0.0
8 850
0.22
625
1.00
1.30
3.00
0.7
032
0.00
0.30
0:00
100
357
50
106
53
330
393
Chilkee
Crippe
Choope
Choopent
W
X
030
3.50
155
025
125
123
Figure S5: Behavior of top attention heads in different Pythia models. (a) CMR distance of the top CMR-like heads in Pythia-70m as a function of training time. Heads are selected by ranking all attention heads by their CMR distances (e.g.. "top 3" heads correspond to the three lowest CMR distances). (b) Fitted CMR temporal drift parameters Sene Srec (middle), 7FT (right) of the top CMR-like heads in Pythia-70m of the top CMR-like heads in Pythia-70n 70m as a function of training time. (d-x) Same as a-c except for Pythia- 160m (d-f), Pythia-4100 40m (g-i), Pythia- lb (j-I), Pythia-1.44 1.4b (m-o), Pythia-2.8h 2.8b (p-r), Pythia- 6.9b (s-u), and Pythia-120 12b (v-x).
function of training time. (c) Average fitted inverse temperature
25

---

2.00
0A8172
0.75
0:00
0.25
PE000C 0
1:00
Figure S6: Generalizing results to additional models. (a) Percentage of heads with a CMR distance less than 0.5 in different layers. (b) Fitted CMR temporal drift parameters (left,, 8ree (middle), TFT (right) in attention heads with the highest induction-head matching scores (i.e., top induction heads). (e) CMR distance of top induction heads in different models. (d) Same (b) except for top CMR-lid heads (i.e., attention heads with the lowest CMR distances).
26

---

NeurIPS Paper Checklist
1. Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification: Our current paper shows a novel connection between LLMs and human episodic memory, both mechanistically (Section 3.3 and 4.2 and behaviorally (Section 5.1
and (.2). 5.2). We also suggest potential causal link between attention heads in LLMs that exhibit human-like memory patterns and the model's in-context learning ability. Both the abstract and introduction highlight these contributions and discuss the potential implications for research on mechanistic interpretability and human cognition.
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made in the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or
NA answer to this question will not be perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results
be expected to generalize other settings.
is fine to include aspirational goals as motivation as long is clear that these goals are not attained by the paper.
2. Limitations
Question: Does the paper discuss the limitations of the work performed by the authors?
Answer: [Yes]
Justification:
We have mentioned a few limitations of the current work in Discussion
Guidelines:
The answer NA means that the paper has limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors are encouraged to create separate "Limitations" section in their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g.. independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be.
The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on a few datasets or with a few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution low or images are taken in low lighting. Or speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.
The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness.
While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, worse outcome might be that reviewers discover limitations that aren't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an impor- tant role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed not penalize honesty concerning limitations.
3. Theory Assumptions and Proofs
27

---

Question: For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?
Answer: [NA]
Justification: This paper does not introduce any new theorems, formulas, or lemmas to be proved.
Guidelines:
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 cross- referenced.
All assumptions should be clearly stated referenced in the statement of any theorems.
The proofs can either appear the main paper or the supplemental material, but they appear in the supplemental material, the authors are encouraged to provide 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 appendix or supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental Result Reproducibility
Does the paper fully disclose all the information needed to reproduce the main ex- perimental 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)?
Question:
Answer: [Yes]
Justification: We fully disclose the models (available through the TransformerLeens library [25:: CMR implementation is included in the code submission), metrics (see Appendix E). model fitting (see Section (.1]) and head ablation (see Section 5.3). All figures included in the paper can reproduced using the included code (see Appendix
Guidelines:
The answer NA means that the paper does not include experiments.
If the paper includes experiments, No answer 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 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 novel architecture, describing the architecture fully might suffice, or if the contribution a specific model and empirical evaluation, 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 good way to accomplish this, but reproducibility instructions for how to replicate the results, access to a hosted model (e.g.. in the case of large language model), releasing of model checkpoint, other means that appropriate to the research performed.
also be provided via detailed
While NeurIPS does not require releasing code, the conference does require all submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) the contribution is primarily a new algorithm, the paper should make clear how to reproduce that algorithm.
(b) If the contribution is primarily new model architecture, the paper should describe the architecture clearly and fully.
(c) the contribution a new model (e.g., a large language model), then there should either be way to access this model for reproducing the results or way to reproduce the model (e.g., with an open-source dataset instructions for how to construct the dataset).
28

---

(d)
We recognize that reproducibility may be tricky in some cases, in which case authors are welcome describe the particular way they provide for reproducibility. In the case of closed-source models, it may be that access to the model is limited in some way (e.g., to registered users), but it should be possible for other researchers to have some path to reproducing or verifying the results.
5. Open access to data and code
Question: Does the paper provide open access to the data and code, with sufficient instruc- tions to faithfully reproduce the main experimental results, as described in supplemental material?
Answer: [Yes]
Justification: All scripts are included in the supplementary materials (see Appendix A where documentation is provided to reproduce all results and figures in the current paper.
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurIPS code and data submission guidelines (https://nips.. cc/
for more details.
While we encourage the release of code and data, we understand that this might not be possible, so "No" is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for new open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurIPS code and data submission guidelines (https
/ wttpps.cc/public////
www.hoooimisiooooooooocc.o... for more details.
The authors should provide instructions on data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. only subset of experiments are reproducible, they should state which ones are omitted from the script and why.
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
Providing as much information as possible in supplemental material (appended to the paper) recommended, but including URLs to data and code permitted.
6. Experimental Setting/Details
Question: Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Answer: [Yes]
We provide written descriptions of the metrics, model parameters, model fitting, and model ablation procedure in the main text and the Appendix. Since all LLM models are pre-trained and accessible through TransformerLeens 25 we omitted their training details.
Justification:
Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented in the core of the paper to a level of detail that is necessary to appreciate the results and make sense of them.
The full details can be provided either with the code, in appendix, or as supplemental material.
7. Experiment Statistical Significance
Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Question:
Answer: [Yes]
Justification:
We included error bars and statistical significance whenever appropriate.

---

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, confi- dence intervals, or statistical significance tests, 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.
is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have 96% CL, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that out of range (e.g. negative error rates).
If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables in the text.
8. Experiments Compute Resources
For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time of execution) needed reproduce the experiments?
Question:
Answer: [Yes]
Justification: We provided full details about the compute resources required to reproduce all experiments discussed in the paper (see Appendix H
Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage.
The paper should provide the amount of compute required for each of the individual experimental runs as well as estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported didn't 'tttt make it into the paper).
paper (e.g., preliminary or failed experiments that
9. Code Of Ethics
Question: Does the research conducted the paper conform, in every respect, with the NeurlPS Code of Ethics httpge//hhorrpp...c//lll////heha
Answer: [Yes]
Justification: We have reviewed the Neurips Code of Ethics and, to the best our knowledge, note violation the Code of Ethics by current work.
Guidelines:
The answer NA means that the authors have not reviewed the NeurlPS Code of Ethics.
If the authors answer No, they should explain the special circumstances that require deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g., there special consid- eration due to laws or regulations in their jurisdiction).
10. Broader Impacts
Question: Does the paper discuss both potential positive societal impacts and negative socictal impacts of the work performed?
30

---

Answer: [NA]
Justification: We briefly discussed in the Introduction that the current work may help to improve interpretability of ICL and prevent malicious use of LLM. Since the paper is mainly foundational research and does not propose new technologies that can be applied or deployed, we did not perceive any direct societal impact.
Guidelines:
The answer NA means that there is no societal impact of the work performed.
If the authors answer NA or No. they should explain why their work has no societal impact or why the paper does not address societal impact.
Examples of negative societal impacts include potential malicious unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g.. deploymentt of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
The conference expects that many papers will be foundational research and not tied to partieular applications, let alone deployments. However, if there is direct path to any negative applications, the authors should point it out. For example, it is legitimate point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.
If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition attacks, mechanisms for monitoring misuse, mechanisms to monitor how system learns from feedback over time, improving the efficiency and accessibility of MI.).
11. Safeguards
Does the paper deseribe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Question:
Answer: [NA]
Justification: We did not release any new datasets or new models. The paper only analyzes existing pretrained language models using novel approach. We did not perceive any risks.
Guidelines:
The answer NA means that the paper poses such risks.
Released models that have high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make best faith effort.
12. Licenses for existing assets
Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?
Question:
Answer: [Yes]
Justification: The only existing assets we used (without modification) were the pretrained language models accessible through the open-source library TransformerLens, The library
31

---

is credited wherever appropriate throughout the paper and we mention the license details in the Appendix.
Guidelines:
The answer NA means that the paper does use existing assets.
The authors should cite the original paper that produced the code package or dataset.
The authors should state which version of the asset used and, possible, include URL.
The name of the license (e.g., CC-BY 4.0) should be included for each asset.
For scraped data from a particular source (e.g.. website). the copyright and terms of service of that source should be provided.
assets are released, the license, copyright information, and terms use in package should be provided. For popular datasets, paperssii thoode com/ com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
For existing datasets that are re- e-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.
If this information is not available online, the authors are encouraged to reach out to the asset's 's creators.
13. New Assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [Yes]
The only new assets introduced in the paper consist of the code for model fitting and analysis. We provided detailed documentation alongside the released code (see Appendix A).
Justification:
Guidelines:
The answer NA means that the paper does not release new assets.
Researchers should communicate the details of the dataset/code/model as part of their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent was obtained from people whose asset is used.
At submission time, remember anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.
14. Crowdsourcing and Research with Human Subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well as details about compensation (if any)?
Answer: [NA]
Justification:
This paper does not involve crowdsourcing or studies on human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsoureing nor research with human subjects.
Including this information in the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then as much detail as possible should be included in the main paper.
According to the NeurIPS Code of Ethies, workers involved data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional Review Board (IRB) Approvals or Equivalent for Research with Human Subjects
32

---

Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or equivalent approval/review based on the requirements of your country or institution) were obtained?
Answer: [NA]
Justification: This paper does not involve crowdsourcing or studies human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Depending on the country in which research is conducted, IRB approval (or equivalent)
may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
We recognize that the procedures for this may vary significantly between institutions and locations, and 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 the institution conducting the review.
33


Reviews:
[{"review_id": "UyZ3SgJNvT", "reviewer": "Reviewer_rPH1", "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.", "questions": "I have mentioned them in the limitations section below.", "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.\n1. **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.\n2. **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.\n3. **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.\n4. **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.\n5. **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.\n6. The CMR distance < 0.5 grouping seems arbitrary. Do you get similar results with < 0.1 and < 1 for instance?\n7. Could you add error bars in Figure 7?", "rating": 4, "confidence": 2, "soundness": 2, "presentation": 2, "contribution": 3, "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.\n2. 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.\n2. The analysis has several issues, which will be discussed in the limitations section.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for recognizing the value of our work and offering constructive feedback. We provide a point-by-point response below.\n\n### Weaknesses \nPlease see below.\n\n### Limitations\n1. 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.\n\n2. 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.\n\n3. 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. \n\n    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.\n\n    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). \n\n4. 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.”\n\n5. 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.\n\n6. 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.\n\n7. Please refer to (Fig. R4) for an updated Fig. 7 with standard errors shown.\n\n**References**\n- Elhage et al. (2021). A mathematical framework for transformer circuits.\n- Olsson, et al. (2022). In-context Learning and Induction Heads.\n- Crosbie, J., & Shutova, E. (2024). Induction Heads as an Essential Mechanism for Pattern Matching in In-context Learning.\n- Murdock, B. B. (1997). Context and mediators in a theory of distributed associative memory (TODAM2).\n- Howard, M. W., & Kahana, M. J. (2002). A distributed representation of temporal context.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "wmMGWg0S0O", "reviewer": "Reviewer_xjvL", "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.\n\nI recommend weak acceptance, contingent on addressing some of the weaknesses identified, particularly regarding generalizability.", "questions": "I have a few questions/suggestions for the authors:\n\n1. 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.\n\n2. How might the findings change if applied to more complex, natural language tasks rather than repeated random tokens? \n\n3. Can you provide any insights into whether CMR-like behaviours are causally necessary for ICL, or if they might be an epiphenomenon?\n\n4. How do you think these findings might generalize to non-autoregressive models or other Transformer variants?\n\n5. Could you elaborate on potential neural implementations of the proposed mechanisms and their biological plausibility?\n\n6. Have you considered comparing your model's behaviour quantitatively to human behavioural data on episodic memory tasks?\n\nAlso, I have a couple of minor comments:\n- 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). \n- Leave a space before the citation in line 73.\n- Line 121 Tab. tab:comparison. I think you forgot to use /ref{.}.\n- 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": "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.", "rating": 7, "confidence": 4, "soundness": 3, "presentation": 4, "contribution": 3, "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.\n\n2. 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.\n\n3. 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.\n\n4. Clear exposition: The paper is well-written and logically structured, making it accessible to readers from both machine learning and cognitive science backgrounds.\n\n5. 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.\n\n2. 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.\n\n3. 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.\n\n4. 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.\n\n5. 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your interest and thoughtful feedback on our work. Below is our point-by-point response:\n### Weaknesses\n1. 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).\n2. Thank you for the suggestion. Please see the causal analysis in the global rebuttal (point 4).\n3. 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.\n\n    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.\n\n4. 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.\n5. 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.\n### Questions\n1. 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.\n\n    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.\n2. Please see our reply to weakness points 1 and 2.\n3. Please see our reply to weakness point 2.\n4. Please see our reply to weakness point 3.\n5. Please see our reply to weakness point  4.\n6. Please see our reply to weakness point 5.\n\nFinally, thank you for the minor comments. We will modify the texts/labels accordingly.\n\n**References**\n- Murdock et al. (1962). The serial position effect of free recall. \n- Elhage et al. (2021). A mathematical framework for transformer circuits. \n- Olsson et al. (2022) In-context Learning and Induction Heads. \n- Bansal et al. (2022). Rethinking the role of scale for in-context learning: An interpretability-based case study at 66 billion scale.\n- Nanda. (2022). Induction mosaic.\n- Crosbie et al. (2024). Induction Heads as an Essential Mechanism for Pattern Matching in In-context Learning.\n- Samuel. (2024). BERTs are Generative In-Context Learners. \n- Salvatore et al. (2024). Parallels between Neural Machine Translation and Human Memory Search: A Cognitive Modeling Approach. \n- Li et al. (2024). A neural network model trained on free recall learns the method of loci. \n- Giallanza et al. (2024) Toward the Emergence of Intelligent Control: Episodic Generalization and Optimization.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "a8m9uwc1eN", "reviewer": "Reviewer_uMwU", "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.", "questions": "The paper mentions CMR-fitted scores. What are these exactly?\n\nIs the y-axis label in Figure 7C correct? Shouldn’t it be matching scores instead of CMR distance?\n\nLooking beyond an exchange of ideas between two fields, are there any direct practical consequences of the discovered connection?", "limitations": "They are discussed appropriately.", "rating": 4, "confidence": 2, "soundness": 2, "presentation": 1, "contribution": 3, "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.\n\nTo give two concrete examples of how the writing could be improved:\n\n(1) K-composition is introduced in Section 3.3. However, it is essentially completely irrelevant to the paper, and hence just confusing the reader.\n\n(2) Attention scores are a crucial concept but they are not defined anywhere.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for acknowledging the potential of our work and offering us the opportunity to clarify. Below is our point-by-point response.\n\n### Weaknesses\nThank 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:\n\nFirst, 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.\n\nSecond, 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).\n\n### Questions\n1. 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.\n2. 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.\n3. 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. \n\n    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. \n\n**References**\n- Elhage et al. (2021). A mathematical framework for transformer circuits.\n- Singh et al. (2024). What needs to go right for an induction head? \n- Edelman et al. (2024). The evolution of statistical induction heads: In-context learning markov chains.\n- Ferrando et al. (2024). A primer on the inner workings of transformer-based language models.\n- Dai et al. (2019). Transformer-XL: Attentive Language Models beyond a Fixed-Length Context.\n- Vaswani et al. (2017). Attention is All You Need.\n- Murphy et al. (2021). Metacognitive control, serial position effects, and effective transfer to self-paced study.\n- Lu et al. (2024). Episodic memory supports the acquisition of structured task representations.\n- Giallanza et al. (2024) Toward the Emergence of Intelligent Control: Episodic Generalization and Optimization. \n- Zhou et al. (2023). Episodic retrieval for model-based evaluation in sequential decision tasks.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "o0XCTj7ZSu", "reviewer": "Reviewer_18Dm", "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.", "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.\n\nThe 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?\n\nWould 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?\n\nIn 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": "Stated in the paper.", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 4, "contribution": 3, "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.\n\nThe 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.\n\nI'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.\n\nThe 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your appreciation and insightful feedback. Below is our point-by-point response:\n\n### Weaknesses\n1. 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.\n2. 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).\n3. 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.\n\n### Questions\n1. 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.\n2. 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.\n3. 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.). \n4. 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.\n\n**References**\n- Ramsauer et al. (2021). Hopfield Networks is All You Need.\n- Kumaran & McClelland. (2012). Generalization Through the Recurrent Interaction of Episodic Memories.\n- Gershman & Daw. (2017). Reinforcement Learning and Episodic Memory in Humans and Animals: An Integrative Framework.\n- Zhou et al. (2023). Episodic retrieval for model-based evaluation in sequential decision tasks.\n- Giallanza et al. (2024) Toward the Emergence of Intelligent Control: Episodic Generalization and Optimization. \n- Elhage et al. (2021). A mathematical framework for transformer circuits. \n- Olsson et al. (2022) In-context Learning and Induction Heads. \n- Pal et al. (2023). Future Lens: Anticipating Subsequent Tokens from a Single Hidden State.\n- Wang et al. (2022). Interpretability in the wild: a circuit for indirect object identification in gpt-2 small.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:
We thank the reviewers for their constructive comments. Here we respond to questions asked by multiple reviewers:
## 1 Clarity and accessibility of our paper
While some reviewers praised the clarity and accessibility of our paper, others felt that there was room for improvement. Given the interdisciplinary nature of our work, we hope to reach a broad audience in ML/mechanistic interpretability, cognitive modeling, and neuroscience. Accordingly, we will revise Sections 3 and 4 to provide more intuition about induction heads and CMR before diving into the technical details. Given the space limit, we also plan to include a more detailed, pedagogical description of induction heads and CMR in the appendix for the revised paper.

## 2 Comparison to human data
We thank the reviewers for pointing out the missing human data. We analyzed the CRPs (Fig. R1a) of the human free recall data from Zhang et al. (2023). The curves exhibit a forward asymmetry and temporal contiguity as described in Section 4.1. The top-performing subjects have a sharper CRP with a larger forward asymmetry, compared to other subjects. Correspondingly, the fitted $\beta_\text{rec}$ in CMR is higher for top-performing subjects (typically around 0.7 or higher). Similarly, our results suggest that top induction/CMR-like heads often exhibit larger $\beta_\text{rec}$.

We also performed a more extensive literature review to determine the distribution of $\beta$ observed in existing human studies (Fig. R1b). The average fitted values of $\beta_\text{enc}$ and $\beta_\text{rec}$ for human subjects vary from study to study (0.2-0.9). Our results suggest that the top CMR-like heads exhibit a degree of temporal clustering consistent with that in humans.
 
## 3 Generalizability of our results
We agree that verifying the generalizability of our results on more LLM models with different architectures is crucial. We performed the same analysis on three additional models suggested by Reviewer xjvL (Llama3-8B, Mistral-7B, Qwen-7B) and the findings were similar (Fig. R6). Specifically, CMR-like heads appear most frequently in intermediate-late layers (Fig. R6a), and the top heads generally have high temporal clustering (high $\beta_\text{enc}$, even higher $\beta_\text{rec}$; Fig. R6b, d).

## 4 Causal analysis of CMR-like heads on a more naturalistic language task
We appreciate the reviewers’ questions about why CMR and CMR distance should be of interest to our understanding of large language models. Our results as shown in the paper mainly suggest a correlational relationship between small CMR distances (“CMR-like heads”) and in-context learning (ICL) on repeated random tokens. 

To test whether CMR-like heads are causally necessary for ICL in more naturalistic sentences, we conducted an ablation study. We ablated the top 10% CMR-like heads (i.e., top 10% heads with the smallest CMR distances) in each model and computed the resultant ICL score (evaluated on the first 1000 texts of the processed version of Google’s C4 dataset (allenai/c4 on huggingface) with at least 512 tokens, except pythia-12b, which was evaluated using the first 500 due to time constraint). Specifically, 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” (Olsson et al., 2022). 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. R5, ranked by original ICL score from lowest to highest). 

Most models (except Pythia-1B) showed a higher ICL score (worse ICL ability) if the top 10% CMR-like heads were ablated, compared to the case where 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). We note, however, there are a few complications when explaining the results: we lack time to average over enough samples to further reduce the error bars; generally there might be a Hydra effect where ablation of heads causes other heads to compensate (McGrath et al 2023); the ICL scores for original models are closer to 0 (weaker ICL performance) in Pythia series than other architectures, suggesting either Pythia series’ ICL abilities are weaker, or the distributional differences between their training data and evaluation data are larger. 

In short, our finding suggests that CMR-like behavior is not merely an epiphenomenon; it is essential underlying LLM’s ICL ability. CMR distance thus provides a meaningful metric to characterize individual heads in LLMs.

Meta-review: 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.
Decision: Accept (poster)