Title: Linking In-context Learning in Transformers to Human Episodic Memory
Venue: NeurIPS 2024
Accepted: yes
Keywords: in-context learning

Transformer

induction head

episodic memory

mechanistic interpretability
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

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.

1 Introduction

Neural networks often bear striking similarities to biological intelligence. For instance, convolutional networks trained on computer vision tasks can predict neuronal activities in the visual cortex [1-4]. Recurrent neural networks trained on spatial navigation develop neural representations similar to the entorhinal cortex and hippocampus [5, 6], and those trained on reward-based tasks recapitulate biological decision-making behaviors [7]. Feedforward networks trained on category learning exhibit human-like attentional bias [8]. Identifying commonalities between artificial and biological intelligence offers unique insights into both model properties and the brain's cognitive functions.

In contrast to this long tradition of drawing parallels between AI models and biology, exploration of the biological relevance of the Transformer architecture - originally proposed for natural language translation [9] - remains limited, with many researchers in neuroscience, cognitive science, and deep learning viewing it as being fundamentally different from the brain. So far, Transformers have been linked to generalized Hopfield networks [10], neural activities in the language cortex [11-13],

Corey Y Zhou ∗

a

b

c

d

token

embed

𝐻

1

…

MLP

𝐻

2

…

unembed

logits

…

position

token

embed

unembed

logits

…

position

…

Light

Cat

Table

Water

Time

Study phase

…

…

Cat

(Studied word?)

Recall phase

Time

Time

…

Light

Cat

Table

Water

…

Cat

(What's next?)

f (word)

t (context)

𝑀

FT

𝑡

IN

𝑀

TF

and hippocampo-cortical circuit representations [14]. However, it remains unclear whether and how the emergent behavior and mechanisms of interacting attention heads relate to biological cognition.

This study bridges this gap by examining the parallels between attention heads in Transformer models and episodic memory in biological cognition. We focus on 'induction heads', a particular type of attention head in Transformer models and a crucial component of in-context learning (ICL) observed in LLMs [15]. ICL enables LLMs to perform new tasks on the fly during test time, relying solely on the context provided in the input prompt, without the need for additional fine-tuning or task-specific training [16, 17]. We show that induction heads share several parallel properties with the contextual maintenance and retrieval (CMR) model, an influential model of human episodic memory during free recall. Understanding the mechanisms of ICL is important for developing better models capable of performing unseen tasks, as well as for AI safety research, as the models could be instructed to perform malicious activities after being deployed in real-world scenarios.

In sections below, we introduce the task in Section 2, Transformer and induction heads in Sections 3.1, 3.2, and the CMR model in Section 4.1. We show that induction heads and CMR are mechanistically similar in Sections 3.3 and 4.2 and behaviorally similar in Section 5.1. We further characterize the emergence of CMR-like behavior in Section 5.2 and its possible causal role in Section 5.3. Overall, our study provides evidence for a novel bridge between Transformer models and episodic memory.

2 Next-token prediction and memory recall

Transformer models in language modeling are often trained to predict the next token [16]. ICL thus helps next-token prediction using information provided solely in the input prompt context. One way to evaluate a model's ICL is to run it on a sequence of repeated random tokens [15] (Fig. 1a). For example, consider the prompt '[A][B][C][D][A][B][C][D]'. Assuming that no structure between these tokens has been learned, the first occurrence of each token cannot be predicted - e.g., the first [C] cannot be predicted to follow the first [B]. At the second [B], however, a model with ICL should predict [C] to follow by retrieving the temporal association in the first part of the context.

Much like ICL in a Transformer model, human cognition is also known to perform associative retrieval when recalling episodic memories. Episodic retrieval is commonly studied using the free recall paradigm [18, 19] (Fig. 1b). In free recall, participants first study a list of words sequentially, and are then asked to freely recall the studied words in any order [20]. Despite no requirements on recall order, humans often exhibit patterns of recall that reflect the temporal structure of the preceding study list. In particular, the retrieval of one word triggers the subsequent recall of other words studied close in time (temporal contiguity). Additionally, words studied after the previously recalled word are

retrieved with higher probability than words studied before the previously recalled word, leading to a tendency of recalling words in the same temporal ordering of the study phase (forward asymmetry). These effects are typically quantified through the conditional response probability (CRP): given the most recently recalled stimulus with a serial position i during study, the CRP is the probability that the subsequently recalled stimulus comes from the serial position i + lag (see e.g., Fig. 4a).

3 Transformer models and induction heads

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 [21, 22]. Each token at position i in the input has its own residual stream z i (with a dimension of d model ) serving as a shared communication channel between model components at different layers (Fig. 1c; the residual stream is shown as a blue path), such as self-attention and multi-layered perceptrons (MLP). The initial residual stream z (0) i contains token embeddings 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 a computation, and additively writes into the residual stream. Specifically, attention heads at layer l read from all past z j (with j ≤ i ) and write into the current z i as z ( l ) ′ i ← z ( l -1) i + ∑ heads h H ( h ) ( z ( l -1) i ; { z ( l -1) j } j ≤ i ) , while MLP layers read from only the current z i and write into z i as z ( l ) i ← z ( l ) ′ i + MLP ( z ( l -1) ′ i ). Readers unfamiliar with attention heads (e.g., attention scores and patterns) are referred to Appendix B. Other components like layer normalization are omitted for simplicity. At the final layer, the residual stream is passed through the unembedding layer to generate the logits (input to softmax) that predict the next token.

Components in different layers can interact with each other through the residual stream [21]. As an important example, a first-layer head H 1 may write its output into the residual stream, which is later read by a second-layer head H 2 that writes its output to the residual stream for later layers to use.

3.2 Induction heads and their attention patterns

Previous mechanistic interpretability studies identified a type of attention heads critical for ICL, known as induction heads [21, 15, 23, 24]. Induction heads are defined by their match-then-copy behavior [15, 24]. They look back ( prefix matching ) over previous occurrences of the current input token (e.g., [B]), determine the subsequent token (e.g., [C] if the past context included the pair [B][C]), and increase the probability of the latter - that is, after finding a 'match', it makes a 'copy' as the predicted next token (. . . [B][C] . . . [B] → [C]). To formalize this match-then-copy pattern, we use the induction-head matching score (between 0 and 1) to measure the prefix-matching behavior. We then use the copying score (between -1 and 1) to measure the copying behavior (see Appendix E). An induction head should have a large induction-head matching score and a positive copying score.

We first examined the induction behaviors of attention heads in the pre-trained GPT2-small model [16] using the TransformerLens library [25]. To elicit induction behaviors, we constructed a prompt consisting of two repeats of a random-token sequence (see Section 2 and Appendix D). We recorded the attention scores of each head ˜ A (before softmax) and the attention patterns A (after softmax) for each pair of previous and current token positions (for definitions see Appendix B). Several heads in GPT2-small had a high induction-head matching score (Fig. 2a), such as L5H1 (layer number 5, head number 1). In the first sequence repeat, this attention head attends mostly to the beginning-ofsequence token. In the second repeat, this head shows a clear 'induction stripe' (Fig. 2b) where it mostly attends to the token that follows the current token in the first repeat.

We calculated attention scores as a function of relative position lags to further characterize the behavior of induction heads. This analysis is reminiscent of the CRP analysis on human recall data, as in Section 4.1. We found that induction heads' attention to earlier tokens follows a similar pattern as seen in human episodic recall (Fig. 2c, Fig. 5a-c), including temporal contiguity (e.g., the average attention score for | lag | ≤ 2 is larger than for | lag | ≥ 4 ) and forward asymmetry (e.g., the average attention score for lag > 0 is larger than for lag < 0 ).

a

b

c

Source token

First repeat

Second repeat

Head

Destination token

Layer

Second repeat

First repeat

Induction head matching score

Attention pattern (L5H1)

Conditional response probability

(L5H1)

0.5

10

10

score

5

Attention

0

一

-5

0

Lag

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 a mechanistic understanding of how the induction head works internally. To gain insights into the internal mechanisms of induction heads, we focus here on smaller transformer models, acknowledging that individual attention heads of larger LLMs likely exhibit more sophisticated behavior. Prior work has discovered two kinds of induction mechanisms in two-layer attention-only Transformer models: K-composition and Q-composition (Fig. 3a-b, Tab. S1) [21, 24], characterizing how information from the first-layer head is composed to inform attention of the second-layer head. We provide an overview of both below. Our main focus in this paper is the comparison between Q-composition and CMR, but K-composition is provided as a point of comparison for readers familiar with mechanistic interpretability.

In K-composition (Fig. 3a), the first-layer 'previous token' head uses the current token's position embedding, PE i , as the query, and a past token's position embedding PE j , as the key. When the match condition PE j = PE i -1 is satisfied (meaning j is the previous position of i ), the head writes the previous token's token embedding, TE j , as the value into the residual stream z i . The second-layer induction head uses the current token's TE k as the query, and the previous token head's output TE j at residual stream z i as the key ('K-composition'). When the match condition TE j = TE k is satisfied, the head writes TE i (at residual stream z i ) as the value into the residual stream z k , effectively increasing the logit for the token that occurred at position i .

In Q-composition (Fig. 3b), the first-layer 'duplicate token' head uses the current token's TE k as the query, and a past token's TE j as the key. When the match condition TE j = TE k is satisfied (meaning token k is a duplicate of token j ), the head writes the token's PE j as the value into the residual stream z k . The second-layer induction head uses the duplicate token head's output PE j at residual stream z k as the query ('Q-composition') and a past token's PE i as the key. When the match condition PE j = PE i -1 is satisfied, the head writes TE i (at residual stream z i ) as the value into the residual stream z k , increasing the logit for the token that occurred at position i .

In the following sections, we will reveal a novel connection between ICL and human episodic memory. We first introduce the CMR model of episodic memory, and then formally re-write it as a Q-composition induction head performing prefix matching, allowing us to link induction heads' attention biases to those known in human episodic memory.

4 Contextual maintenance and retrieval model (CMR)

4.1 CMRin its original form

CMR, an influential model of human episodic memory, provides a general framework to model memory recall as associative retrieval. It leverages a distributed representation called temporal context , which acts as a dynamic cue for recalling subsequent information based on recently seen words [26]. CMR explains the asymmetric contiguity bias in human free recall (see Fig. 4 and Fig. S1) and has been extended to more complex memory phenomena such as semantic [27] and emotional

𝑃𝐸𝑖

𝑇𝐸𝑖

𝑃𝐸 𝑗

𝑇𝐸 𝑗

𝑃𝐸𝑘

𝑇𝐸𝑘

𝐻

1

𝑃𝐸𝑖

𝑇𝐸𝑖

𝑇𝐸 𝑗

𝐻

2

𝑃𝐸𝑘

𝑇𝐸𝑘

𝑇𝐸𝑖

…

Query

Value

Key

𝑃𝐸𝑘

𝑇𝐸𝑘

Output

𝑃𝐸𝑖

𝑇𝐸𝑖

𝑃𝐸𝑗

𝑇𝐸𝑗

𝑃𝐸𝑘

𝑇𝐸𝑘

𝐻

1

𝑃𝐸𝑖

𝑇𝐸𝑖

𝐻

2

𝑃𝐸𝑘

𝑇𝐸𝑘

𝑃𝐸 𝑗

𝑇𝐸𝑖

…

𝑃𝐸𝑘

𝑇𝐸𝑘

𝑃𝐸𝑗

K-composition

Q-composition

𝑡

𝑖-

1

𝑓 𝑖

𝑡 𝑗 -

1

𝑓 𝑗

𝑡𝑘-

1

𝑓 𝑘

𝑀

FT

𝑡

𝑖-

1

𝑓 𝑖

𝑡

𝑖

𝑀

TF

…

𝑡𝑘-

1

𝑓 𝑘

𝑡𝑘

𝑡 𝑗 -

1

𝑓 𝑗

𝑡 𝑗

Study phase

Recall phase

𝑡𝑘-

1

𝑓 𝑘

𝑡𝑘

𝑓 𝑘

+1

CMR

a

b

c

Condition

𝑃𝐸𝑗

=

𝑃𝐸𝑖-

1

𝑇𝐸𝑗

=

𝑇𝐸𝑘

𝑇𝐸𝑗

=

𝑇𝐸𝑘

𝑃𝐸𝑖-

1 =

𝑃𝐸𝑗

𝑓 𝑗

=

𝑓 𝑘

𝑡

𝑖-

1 =

𝑡

𝑗

[28] effects. We provide a pedagogical introduction to CMR in Appendix C for readers without a background in cognitive science or episodic memory. Below, we briefly list the essentials of CMR.

In CMR (Fig. 1d), each word token is represented by an embedding vector f (e.g., one-hot; f i for the i -th word in a sequence). The core dynamic that drives both sequential encoding and retrieval is

where t i is the temporal context at time step i , and t IN i is an input context associated with f i . β controls the degree of temporal drift between time steps ( β enc for encoding/study phase and β rec for decoding/retrieval phase) and ρ is picked to ensure t i has unit norm. Specifically, during the encoding phase, t IN i represents a pre -experimental context associated with the i -th word as t IN i = M FT pre f i , where M FT pre is a pre-fixed matrix. At each time step, a word-to-context (mapping f to t ) memory matrix M FT exp learns the association between f i and t i -1 (i.e., M FT exp is updated by t i -1 f T i ). During the decoding (retrieval) phase, t IN i is a mixture of pre-experimental ( t IN pre = M FT pre f i ) and experimental contexts ( t IN exp = M FT exp f i ). The proportion of these two contexts is controlled by an additional parameter γ FT ∈ [0 , 1] as t IN i = ((1 -γ FT ) M FT pre + γ FT M FT exp ) f i . The asymmetric contiguity bias arises from this slow evolution of temporal context: when 0 < β < 1 , t i passes through multiple time steps, causing nearby tokens to be associated with temporally adjacent contexts that are similar to each other (temporal contiguity), i.e., ⟨ t i , t j ⟩ is large if | i -j | is small. Additionally, t IN pre = M FT exp f i only enters the temporal context after time i . Thus t IN i, exp is associated with f j only for j > i (asymmetry).

CMR also learns a second context-to-word (mapping t back to f ) memory matrix M TF (updated by each f i t T i -1 ). When an output is needed, CMR retrieves a mixed word embedding f IN i = M TF t i . If f j are one-hot encoded, we can simply treat f IN i as a (unnormalized) probability distribution over the input tokens. Or, CMR can compute the inner product ⟨ f j , f IN i ⟩ for each cached word f j as input to softmax (with an inverse temperature τ -1 ) to recall a word.

Intuitively, the temporal context resembles a moving spotlight with a fuzzy edge: it carries recencyweighted historical information that may be relevant to the present, where the degree of information degradation is controlled by ρ . Larger β 's correspond to 'sharper' CRPs with stronger forward asymmetry and stronger temporal clustering that are core features of human episodic memory. As a concrete example, consider n unique one-hot encoded words { f i } . If M FT pre = ∑ n i =1 f i f T i (i.e., the pre-experimental context associated with each word embedding is the word embedding itself) and γ FT = 0 , Eq. 1 at decoding is reduced to t i = ρ t i -1 + β f i = ∑ i j =1 βρ i -j f j , which is a linear combination of past word embeddings.

a

b

c

Human

CMR

4.2 CMRas an induction head

We now proceed to map CMR to a Q-composition-like head (see Fig. 3c and Tab. S1 for details).

To begin with, we note that the word f i seen at position i is the same as TE i , and the context vector t i -1 (before update) at position i is functionally similar to PE i . It follows that the set { f i , t i -1 } is functionally similar to the residual stream z i , updated by the head outputs.

CMRexperimental word-context retrieval as first-layer self-attention . The temporal context is updated by t IN i = ((1 -γ FT ) M FT pre + γ FT M FT exp ,i ) f i at decoding. The memory matrix M FT exp ,i acts as a first-layer duplicate token head, where the current word f i is the query, the past embeddings { f j } make up the keys, and the temporal contexts t j -1 associated with each f j are values. This head effectively outputs 'What's the position (context vector) at which I encountered the same token f i ?'

CMR pre-experimental word-context retrieval as MLP . The pre-experimental context t IN i = M FT pre f i (retrieved contextual information not present in the experiment) is the output of a linear fully-connected layer (functionally similar to MLP; not drawn).

CMRevolution as residual stream updating . The context vector is updated by t k = ρ t k -1 + β t IN k . Equivalently, the head M FT updates the information from { f k , t k -1 } to { f k , t k -1 , t k }. At the position k during recall, the updated context t k contains t IN k ( ≈ t j ) (Fig. 3c).

CMRcontext-word retrieval as second-layer self-attention . The retrieved embedding is f IN k = M TF t k , where M TF acts as a second-layer induction head, where the temporal context t k is the query, the past contexts { t i -1 } make up the keys, and the embeddings f i associated with each t i -1 are values. This effectively implements Q-composition [21], because t k , as the Query , is affected by the output of the first-layer M FT exp head.

CMRword recall as unembedding . The final retrieved word probability is determined by the inner product between the retrieved memory f IN k and each studied word f j , similar to the unembedding layer generating the output logits from the residual stream.

CMRlearning as a causal linear attention head . Both associative matrices of CMR are learned in a manner consistent with the formation of causal linear attention heads. Specifically, the wordto-context matrix is updated by M FT exp ,i = M FT exp ,i -1 + t i -1 f T i (with M FT exp , 0 = 0 ), associating f i (key) and t i -1 (value). It is equivalent to a causal linear attention head, because M FT exp f k = ( ∑ i<k t i -1 f T i ) f k = ∑ i<k t i -1 ( f T i f k ) = ∑ i<k sim ( f i , f k ) t i -1 . Similarly, the context-to-word

a

e

b

c

d

matrix, updated by M TF i = M TF i -1 + f i t T i -1 (with M TF 0 = 0 ), is equivalent to a causal linear attention head that associates t i -1 (key) with f i (value).

To summarize, the CMR architecture resembles a two-layer transformer with a Q-composition linear induction head. It's worth noting that although we cast t i as the position embedding, unlike position embeddings that permit parallel processing in Transformer models, t i is recurrently updated in CMR (Eq. 1). It is possible that Transformer models might acquire induction heads with a similar circuit mechanism, where t i corresponds to autoregressively updated context information in the residual stream that serves as the input for downstream attention heads.

5 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. 2c, Fig. 5a-c; also see Fig. S2 for heads unlike CMR) and established the mechanistic similarity between induction heads and CMR (Fig. 3). To further quantify their behavioral similarity, we propose the metric CMR distance , defined as the mean squared error (MSE) between the head's average attention scores and its CMR-fitted scores (see Appendix E, Fig. 5a-d). In essence, we optimized the parameters ( β enc , β rec , γ FT , τ -1 ) for each head to obtain a set of CMR-fitted scores that minimizes MSE.

At the population level, heads with a large induction-head matching score and a positive copying score also have a smaller CMR distance (Fig. 5e), suggesting that the CMR distance captures meaningful behavior of these heads. Notably, certain heads that are not typically considered induction heads (e.g., peaking at lag=0) can be well captured by CMR (Fig. 5d).

Consistent with prior findings that induction heads were primarily observed in the intermediate layers of LLMs [31], we found that the majority of heads in the intermediate-to-late layers of GPT2-small have lower CMR distances (Fig. 6a, Fig. S3a; for the choice of threshold see Appendix E.5). We also observed similar phenomena in a different set of LLMs called Pythia (Fig. 6b), a family of models with shared architecture but different sizes, as well as three well-known models (Qwen-7B [32], Mistral-7B [33], Llama3-8B [34], Fig. 6c). We summarized these results in Fig. S3b.

Additionally, to contextualize these CMR distances, we included the Gaussian distance using Gaussian functions as a baseline (with the same number of parameters as CMR), since its bell shape captures the basic aspects of temporal contiguity and forward/backward asymmetry. We found that, across 12 different models (GPT2, Pythia models, Qwen-7B, Mistral-7B, Llama3-8B), CMR provides significantly better descriptions (lower distances) than the Gaussian function for the top induction heads (average CMR distance: 0.11 (top 20), 0.05 (top 50), 0.12 (top 100), 0.12 (top 200); average

a

b

c

a

b

d

c

e

f

g

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 to measure the emergence of CMRlike behavior. As the model's loss on the designed prompt decreases through training (Fig. 7a), the degree of temporal clustering increases, especially in layers where induction heads usually emerge. For instance, intermediate layers of Pythia-70m (e.g., L3, L4) show the strongest temporal clustering that persists over training (Fig. S4a-b). This, combined with an increasing inverse temperature (Fig. 7b), suggests that attention patterns become more deterministic over training, while shaped to mirror human-like asymmetric contiguity biases. In fact, human subjects with better free recall performance tend to exhibit stronger temporal clustering and a higher inverse temperature [29].

*

**

*

n.s.

***

n.s.

n.s.

***

***

n.s.

For individual heads, those with higher induction-head matching scores (Fig. 7d) (or similarly with smaller CMR distances, see Fig. 7f) consistently exhibit greater temporal clustering (Fig. 7e, f respectively), as the fitted β 's (both β enc and β rec ) were large. We also observed similar values of fitted β s in Qwen-7B, Mistral-7B, and Llama3-8B (Fig. S6b). These fitted β 's of these attention heads fall into a similar range as human recall data (Fig. 7c). We interpret this in light of a normative view of the human memory system: in humans, the asymmetric contiguity bias with a β < 1 is not merely phenomenological; under the CMR framework, it gives rise to an optimal policy to maximize memory recall when encoding and retrieval are noisy [29]. In effect, a large β (but less than 1) in Eq. 1 provides meaningful associations beyond adjacent words to facilitate recall, such that even if the immediately following token is poorly encoded or the agent fails to decode it, information from close-by tokens encapsulated in the temporal context still allows the agent to continue decoding.

In addition, we observed an increase in β enc and β rec (Fig. 7e, g) during the first 10 training checkpoints, when the model loss significantly drops. The training process leads to higher values of β rec . Specifically, β rec values are higher than β enc , highlighting the importance of temporal clustering during decoding for model performance. These results suggest that attention to temporally adjacent tokens with an asymmetric contiguity bias may support ICL in LLMs.

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) or the same number of randomly selected heads in each model. We then computed the resultant ICL score on the sampled texts (2000 sequences, each with at least 512 tokens) from the processed version of Google's C4 dataset [35]. The ICL score is defined as the loss of the 500th token in the context minus the loss of the 50th token in the context, averaged over dataset examples [15]. Intuitively, a model with better in-context learning ability has a lower ICL score, as the 500th token is further into the context established from the beginning. We tested models with various model architectures and complexity, including GPT2, Pythia models, and Qwen-7B (Fig. 8). Most models showed a higher ICL score (worse ICL ability) if the top 10% CMR-like heads were ablated, compared to if the same number of randomly selected heads were ablated. This effect was particularly significant if the original model had a low ICL score (e.g., GPT2, Qwen-7B).

Our finding therefore suggests that CMR-like heads are not merely an epiphenomenon, but essential underlying LLM's ICL ability, and that the CMR distance is a meaningful metric to characterize individual heads in LLMs. Nonetheless, our result needs to be interpreted cautiously: First, a Hydra effect has been noted where ablation of heads causes other heads to compensate [36]. Second, our analysis cannot confirm a direct causal role of the CMR-like behavior of these CMR-like heads in ICL, since it is possible that characteristics other than the episodic-memory features in these CMR-like heads might causally contribute to ICL. Finally, the ICL scores for the original Pythia models are closer to 0 (worse ICL performance) than other models', suggesting either weaker ICL ability of the Pythia series, or larger distributional differences between their training data and our evaluation texts.

6 Discussion

This study bridges LLMs and human episodic memory by comparing Transformer models' induction heads and CMR. We revealed mechanistic similarities between CMR and Q-composition induction

heads and identified CMR-like attention biases (i.e., asymmetric contiguity) in pre-trained LLMs. Notably, CMR-like heads emerge in LLMs' intermediate-to-late layers, evolve towards a state akin to human memory biases, and may play a causal role in ICL. These findings offer significant connections between the current generation of AI algorithms and a century of human memory research.

From a neuroscience and cognitive science perspective, CMR's link to induction heads might reveal normative principles of memory and hippocampal processing, echoing the role of the hippocampus in pattern prediction and completion [37-39]. While CMR is a behavioral model, it also explains neural activity patterns: associative matrices that represent episodic memories may be instantiated in hippocampal synapses. The temporal context aligns with the hippocampus' recurrent nature, likely involving subregions including CA1, dentate gyrus [40], CA3 [41]), and cortical areas projecting to the hippocampus such as the entorhinal cortex [42].

Our results show strong connections to neural network models of episodic memory. For example, neural networks with attention [43] or recurrence [44] trained for free recall exhibit the same recall pattern as the optimal CMR. [45] showed that a neural network implementation of CMR can explain humans' flexible cognitive control. Our results also align with research connecting attention mechanisms in Transformers to the hippocampal formation [14]. While prior work focused on emergent place and grid cells in Transformers, the hippocampal subfields involved are also postulated to represent CMR components [40]. These results support our proposal that query-key-value attention mechanisms link to biological episodic retrieval, suggesting that CMR-like behavior emerges naturally in neural networks with proper objectives.

The link to induction heads could enable researchers to develop alternative mechanisms for episodic memory. For instance, K-composition and Q-composition induction circuits might serve as alternative models to CMR. K-composition and Q-composition require positional encoding, which we speculate could be implemented by grid-like cells with periodic activations tracking space and time [46]. Further, the interactions between episodic memory and other advanced cognitive functions [47-51] might be understood based on more complex attention-head composition mechanisms (e.g., N-th order virtual attention head [21]).

From the perspective of LLM mechanistic interpretability, we offer a more detailed behavioral description and reinterpret induction heads through the lens of CMR and the asymmetric contiguity bias. First, CMR-like heads are not limited to induction heads - some attention heads are well captured by CMR despite not meeting the traditional induction head criterion (e.g., Fig. 5d). Second, we speculate that heads with low CMR distances and low induction-head matching scores may encode multiple future tokens observed in LLMs [52], capturing distant token information better than ideal induction heads. Third, we observed a scale difference between raw attention scores and those from human recall, as discussed in Appendix G.1. Lastly, though CMR relies on recurrently updated context vectors that are different from the K-composition and Q-composition mechanisms, we posit that deeper Transformer models may develop similar mechanisms via autoregressively updated information in the residual stream, a possibility yet to be explored. Overall, our results suggest a fuller view of ICL mechanisms in LLMs, where heads learn attention biases akin to human episodic memory, empowering next-token and future-token predictions.

These connections with CMR may shed light on intriguing features and functions in LLMs. For instance, the 'lost in the middle' phenomenon may be related to these heads [53], as humans exhibit similar recall patterns [19, 26]. Understanding the connection could suggest strategies to mitigate the problem, e.g., adjusting study schedules based on the serial position [54]. Second, CMR not only applies to individual words but also to clusters of items [55, 56], suggesting these heads may process hierarchically organized text chunks [57]. Third, episodic mechanisms captured by CMR are posited to support adaptive control [47-49] and flexible reinforcement learning [50, 51], suggesting similar roles for these heads in more complex cognitive functions of LLMs. Finally, it was proposed that Transformers can be implemented in a biologically plausible way [58], and we outline an alternative proposal mapping Transformer models to the brain (Appendix G.2).

Our study has several limitations. First, while we use CMR to characterize these heads' behavior, it is unclear if CMR can serve as a mechanistic model in larger Transformer models. Second, it is unclear if our conclusions hold for untested Transformer models. Third, our causal analysis lacks the power to narrow down the causal role of specific CMR-like characteristics. In addition, while ICL score is a primary metric for measuring ICL ability [15, 59], a systematic evaluation using specific datasets and tasks would allow for a stronger causal claim. Addressing these limitations is a key future direction.

References

Daniel LK Yamins, Ha Hong, Charles F Cadieu, Ethan A Solomon, Darren Seibert, and James J DiCarlo. Performance-optimized hierarchical models predict neural responses in higher visual cortex. Proceedings of the national academy of sciences , 111(23):8619-8624, 2014.

Daniel LK Yamins and James J DiCarlo. Using goal-driven deep learning models to understand sensory cortex. Nature neuroscience , 19(3):356-365, 2016.

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 is most brain-like? BioRxiv , page 407007, 2018.

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.

Christopher J Cueva and Xue-Xin Wei. Emergence of grid-like representations by training recurrent neural networks to perform spatial localization. arXiv preprint arXiv:1803.07770 , 2018.

Andrea Banino, Caswell Barry, Benigno Uria, Charles Blundell, Timothy Lillicrap, Piotr Mirowski, Alexander Pritzel, Martin J Chadwick, Thomas Degris, Joseph Modayil, et al. Vectorbased navigation using grid-like representations in artificial agents. Nature , 557(7705):429-433, 2018.

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-reinforcement learning system. Nature neuroscience , 21(6):860-868, 2018.

Catherine Hanson, Leyla Roskan Ça˘ glar, and Stephen José Hanson. Attentional bias in human category learning: The case of deep learning. Frontiers in Psychology , 9, 2018.

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, S. 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.

Hubert Ramsauer, Bernhard Schäfl, Johannes Lehner, Philipp Seidl, Michael Widrich, Thomas Adler, Lukas Gruber, Markus Holzleitner, Milena Pavlovi´ c, Geir Kjetil Sandve, et al. Hopfield networks is all you need. arXiv preprint arXiv:2008.02217 , 2020.

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.

Martin Schrimpf, Idan Blank, Greta Tuckute, Carina Kauf, Eghbal A Hosseini, Nancy Kanwisher, Joshua Tenenbaum, and Evelina Fedorenko. The neural architecture of language: Integrative reverse-engineering converges on a model for predictive processing. BioRxiv , 2020.

Richard Antonello, Aditya Vaidya, and Alexander Huth. Scaling laws for language encoding models in fmri. Advances in Neural Information Processing Systems , 36, 2024.

James CR Whittington, Joseph Warren, and Timothy EJ Behrens. Relating transformers to models and neural representations of the hippocampal formation. arXiv preprint arXiv:2112.04035 , 2021.

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.11895 , 2022.

Alec Radford, Jeffrey Wu, Rewon Child, David Luan, Dario Amodei, Ilya Sutskever, et al. Language models are unsupervised multitask learners. OpenAI blog , 1(8):9, 2019.

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:1877-1901, 2020.

Murdock and B Bennet. The serial position effect of free recall. Journal of Experimental Psychology , 64:482-488, 1962.

Marc W Howard and Michael J Kahana. A distributed representation of temporal context. Journal of mathematical psychology , 46(3):269-299, 2002.

Roger Ratcliff and Gail McKoon. Does activation really spread? Psychological Review , 88(5): 454-462, 1981. doi: 10.1037/0033-295x.88.5.454.

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 Thread , 1, 2021.

Rupesh Kumar Srivastava, Klaus Greff, and Jürgen Schmidhuber. Highway networks. arXiv preprint arXiv:1505.00387 , 2015.

Gautam Reddy. The mechanistic basis of data dependence and abrupt learning in an in-context classification task. arXiv preprint arXiv:2312.03002 , 2023.

Aaditya K 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 arXiv:2404.07129 , 2024.

Neel Nanda. Transformerlens, 2022. URL https://github.com/neelnanda-io/ TransformerLens .

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.

Lynn Lohnas, Sean Polyn, and Michael Kahana. Expanding the scope of memory search: Modeling intralist and interlist effects in free recall. Psychological review , 122:337-363, 04 2015. doi: 10.1037/a0039036.

Rivka T Cohen and Michael J. Kahana. Retrieved-context theory of memory in emotional disorders. bioRxiv , page 817486, 2019.

Qiong Zhang, Thomas L. Griffiths, and Kenneth A. Norman. Optimal policies for free recall. Psychological Review , 130(4):1104-1124, 2023. doi: 10.1037/rev0000375.

Kevin Wang, Alexandre Variengien, Arthur Conmy, Buck Shlegeris, and Jacob Steinhardt. Interpretability in the wild: a circuit for indirect object identification in gpt-2 small. arXiv preprint arXiv:2211.00593 , 2022.

Hritik Bansal, Karthik Gopalakrishnan, Saket Dingliwal, S. Bodapati, Katrin Kirchhoff, and Dan Roth. Rethinking the role of scale for in-context learning: An interpretability-based case study at 66 billion scale. In Annual Meeting of the Association for Computational Linguistics , 2022.

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.16609 , 2023.

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.06825 , 2023.

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.

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.

Callum McDougall, Arthur Conmy, Cody Rushing, Thomas McGrath, and Neel Nanda. Copy suppression: Comprehensively understanding an attention head, 2023. URL https://arxiv. org/abs/2310.04625 .

Kimberly L Stachenfeld, Matthew M Botvinick, and Samuel J Gershman. The hippocampus as a predictive map. Nature neuroscience , 20(11):1643-1653, 2017.

Anna C Schapiro, Nicholas B Turk-Browne, Matthew M Botvinick, and Kenneth A Norman. Complementary learning systems within the hippocampus: a neural network modelling approach to reconciling episodic memory with statistical learning. Philosophical Transactions of the Royal Society B: Biological Sciences , 372(1711):20160049, 2017.

Li Ji-An, Fabio Stefanini, Marcus K Benna, and Stefano Fusi. Face familiarity detection with complex synapses. Iscience , 26(1), 2023.

John J Sakon and Michael J Kahana. Hippocampal ripples signal contextually mediated episodic recall. Proceedings of the National Academy of Sciences , 119(40):e2201657119, 2022.

Halle R. Dimsdale-Zucker, Maureen Ritchey, Arne D. Ekstrom, Andrew P. Yonelinas, and Charan Ranganath. Ca1 and ca3 differentially support spontaneous retrieval of episodic contexts within human hippocampal subfields. Nature Communications , 2020.

Marc 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 , 112(1):75, 2005.

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.

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.

Tyler Giallanza, Declan Campbell, and Jonathan D Cohen. Toward the emergence of intelligent control: Episodic generalization and optimization. Open Mind , 8:688-722, 2024.

Benjamin J Kraus, Mark P Brandon, Robert J Robinson, Michael A Connerney, Michael E Hasselmo, and Howard Eichenbaum. During running in place, grid cells integrate elapsed time and distance run. Neuron , 88(3):578-589, 2015.

Qihong Lu, Ali Hummos, and Kenneth Norman. Episodic memory supports the acquisition of structured task representations, 2024.

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 24702986. doi: 10.1162/opmi_a_00143. URL https://doi.org/10.1162/opmi_a_00143 .

Dharshan Kumaran and James L. McClelland. Generalization through the recurrent interaction of episodic memories. Psychological Review , 119:573 - 616, 2012.

Corey Y. Zhou, Deborah Talmi, Nathaniel D. Daw, and Marcelo G. Mattar. Episodic retrieval for model-based evaluation in sequential decision tasks, 2023.

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.

Koyena Pal, Jiuding Sun, Andrew Yuan, Byron Wallace, and David Bau. Future lens: Anticipating subsequent tokens from a single hidden state. In Jing Jiang, David Reitter, and Shumin Deng, editors, Proceedings of the 27th Conference on Computational Natural Language Learning (CoNLL) , pages 548-560. Association for Computational Linguistics, 2023. doi: 10.18653/v1/2023.conll-1.37. URL https://aclanthology.org/2023.conll-1.37 .

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.

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/s13421-021-01204-y.

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-140, 2016. doi: 10.1016/j. jml.2015.10.002.

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: 10.1038/s41598-019-46908-z.

Christopher A. 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.e5, 2016.

Leo Kozachkov, Ksenia V Kastanenka, and Dmitry Krotov. Building transformers from neurons and astrocytes. Proceedings of the National Academy of Sciences , 120(34):e2219150120, 2023.

Ivan Lee, Nan Jiang, and Taylor Berg-Kirkpatrick. Is attention required for icl? exploring the relationship between model architecture and in-context learning ability. In The Twelfth International Conference on Learning Representations .

Joy Crosbie and Ekaterina Shutova. Induction heads as an essential mechanism for pattern matching in in-context learning. arXiv preprint arXiv:2407.07011 , 2024.

A Code

All code is available at https://github.com/corxyz/icl-cmr .

B Attention mechanisms in Transformer models

For each attention head h , the standard self-attention mechanism calculates its output as H ( h ) ( z ( l -1) i ; { z ( l -1) j } ) = W ( h ) O ∑ j W ( h ) V z ( l -1) j A ( h ) ij , where W ( h ) O is the d model × d head output projection matrix, W ( h ) V is the d head × d model value projection matrix, and A ( h ) ij is the attention pattern attending from the query position i to the key position j . The attention pattern matrix is A ( h ) = Softmax ( ˜ A ( h ) ) = Softmax ( z ( l -1) T j W T Q W K z ( l -1) j / √ d head ) , where ˜ A ( h ) represents the attention score matrix; W ( h ) K and W ( h ) Q are the d head × d model key and query projection matrices respectively.

C Detailed introduction of CMR

In this section, we offer a more detailed description of the Context Maintenance and Retrieval (CMR) model of human episodic memory, focusing on its core components and underlying dynamics. CMR operates in two distinct experimental phases: an encoding/learning phase, where the model memorizes a list of words, and a retrieval/recall phase, where it freely recalls the learned words. The CMR model consists of four primary components: the evolving temporal context ( t ), the word representation or embedding ( f ), the word-to-context association matrix ( M FT ), and the context-to-word association matrix ( M TF ).

The temporal context vector ( t ) continuously evolves over time, incorporating recency-weighted historical information during both encoding and retrieval. The core dynamics of the temporal context t are governed by the update rule Eq. 1.

For illustrative purposes, in the rest of this section, we consider an example sequence consisting of five distinct one-hot word embeddings f 1 , f 2 , ..., f 5 . Each embedding is represented by a 6-dimensional vector, such that f i contains a 1 at the i -th position, and 0 everywhere else. The 6-th dimension represents a dummy stimulus that only serves to maintain the unit norm of the context vector but otherwise irrelevant to the experiment. Additionally, we index encoding/learning time with j and retrieval/recall time with k to explicitly mark the two phases.

For simplicity, we set the pre-experimental word-to-context memory matrix M FT pre ( 6 × 6 ) to the identity matrix. As the result of the distinctiveness assumption, each word embedding at encoding directly induces the same one-hot vector as the input context t IN j = M FT exp f j = f j . Let t 0 = (0 , 0 , 0 , 0 , 0 , 1) T . Thus given a sequential presentation of f 1 , f 2 , ..., f 5 , the corresponding context vectors t j are

where β is the temporal drift parameter, and ρ = √ 1 -β 2 ensures that t j has unit norm. Distinct parameters β enc and β rec are used by the encoding and retrieval phases respectively to capture the phase-specific temporal dynamics. Conceptually, the temporal context reflects a recency-weighted history of past information, with the parameter ρ determining the rate of degradation over time.

More generally, the input context is determined by both pre-experimental and experimental contexts:

The memory matrix M FT is therefore a weighted combination of pre-existing associations learned prior to encoding (the pre-experimental memory matrix M FT pre ) and associations learned during the encoding phase (the experimental memory matrix M FT exp ). M FT pre is commonly initialized as the

identity matrix at the beginning of each experiment and held constant throughout the experiment. On the other hand, M FT exp is usually initialized as the zero matrix. It is updated during the encoding phase at the presentation of each new word by associating the context vector t j -1 (value) with the embedding of the just encountered word f j (key):

(Fig. 3c, first layer).

Suppose at the beginning of the experiment, the experimental word-to-context matrix M FT 0 , exp is initialized as the zero matrix. At encoding time j = 1 , the experimental matrix M FT j, exp is updated by the pairwise association of the previous context and the current word:

Similarly, at time j = 2 , the experimental matrix incorporates the new context-word pair as

The same update procedure continues until the end of the word sequence. Critically, note that the product of M FT j, exp and any f j + l with l > 0 is the zero vector.

During retrieval, the most recently recalled word f k (query) retrieves its corresponding input context

Let f j denote the word encountered at encoding time j with j < k . Assuming f k = f j (i.e., the most recent recall f k matches a previously encoded word f j ), we have

Furthermore, because we assumed M FT pre to be the identity matrix,

This input context is subsequently reinstated along with the current context t k -1 at (retrieval) time k , resulting in an updated context t k via Eq. 1. Thus t k also incorporates information from f j (the previous instance of the most recent recall) and t j -1 (the context immediately before presentation of the previous instance).

Importantly, the context t j at encoding is also a linear combination of f j and t j -1 (by Eq. 1),

where t IN j = ((1 -γ FT ) M FT pre + γ FT M FT exp ) f j . However, because all word embeddings are one-hot and distinct, M FT exp has not learned any associations involving the word f j prior to time j . Therefore, the second term is zero, and

Comparing Eq. S2 and Eq. S1 thus reveals that t IN k , thus t k , become most similar to t j .

Finally, the context-to-word association matrix M TF is responsible for retrieving the context associated with a given word. M TF is usually initialized as the zero matrix prior to experiments. During encoding, it is updated to associate each word embedding f j (value) with the corresponding context vector t j -1 (key):

(Fig. 3, second layer). In practice, we chose to maintain M TF as the transpose of M FT exp . Because the updated context t k is most similar to t j (measured by their inner product), the retrieved vector

will be the most similar to the word embedding f j +1 . Thus, CMR will most likely recall f j +1 following a recall of f k = f j . In the main text, we specified i = j +1 .

Consequently, the CMR model explains two key phenomena observed in human free recall experiments: the temporal contiguity effect and the forward asymmetry effect. The temporal contiguity effect arises because (i) t k integrates the information of t j -1 , and (ii) the context vector t drifts over time, making t j -1 more similar to context vectors associated with nearby words in the sequence, favoring the retrieval of words close to f j . In addition, the forward asymmetry effect occurs because (i) t k integrates the information of t IN j , and (ii) t IN j contributes to t j = t i -1 and its following context vectors, but does not affect earlier context vectors. As a result, t k preferentially retrieves words that follow f j , rather than those that precede it.

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 are presented sequentially [18]; (2) it is a widely acknowledged definition of induction heads in mechanistic interpretability literature [21, 15, 31, 60]; (3) it uses off-distribution prompts to focus on abstract properties, avoiding potential confounds from normal token statistics [21].

The prompt is constructed by taking the top N=100 most common English tokens with a leading space (to avoid unwanted tokenization behavior), which are tokens with the largest biases in the unembedding layer of GPT2-small (or the Pythia models). The prompt concatenated two copies of the permuted word sequence and had a total length of 2 N +1 (one end-of-sequence token at the beginning). The two copies correspond to the study (encoding) and recall phases, respectively.

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 = (¯ a d,s ) over a sequence of tokens { x i } as

Here, d is the destination position and s is the source position. We then assess the extent to which each attention head performs this kind of prefix matching [15, 25]. Specifically, the induction-head matching score for a head with attention pattern A = ( a d,s ) is defined as

(Fig. 2a-b). A head that always performs ideal prefix matching will have an induction-head matching score of 1.

Additionally, an induction head should write to the current residual stream to increase the corresponding logit of the attended token (token copying). We adopt the copying score [21] to measure each head's tendency of token copying. In particular, consider the circuit W = W U W O W V W E for one head, where W E defines token embeddings, W V computes the value of each token from the residual stream (i.e., aggregated outputs from all earlier layers), W O computes the head's output using a linear combination of the value vectors and the head's attention pattern, and W U unembeds the output to predict the next token. The copying score of the head is equal to

where λ k 's are the eigenvalues of the matrix W . Since copying requires positive eigenvalues (corresponding to increased logits), an induction head should have a positive copying score. An ideal induction head will have a copying score of 1.

E.2 CRP analysis of attention heads

For a head with attention scores ˜ A , the average attention score ˜ α lag is defined as

where N is the length of the first repeat in the prompt. Thus if lag = 0 , ˜ α lag quantifies how much the first instance of a token is attended to on average; for lag = 1 , ˜ α lag quantifies the average amount of attention distributed to the immediately following token in the first repeat. We used lag ∈ [ -5 , 5] throughout the paper.

E.3 CMRdistance

The metric CMR distance is defined as

Here, N lag is the number of distinct lags, and q lag is obtained by calculating the CRP using CMR with specific parameter values. Note we did not consider the full set of possible q but only the subset given by the combinations of parameters β enc = 0 . 05 , 0 . 1 , . . . , 1 , β rec = 0 , 0 . 05 , . . . , 1 , and γ FT = 0 , 0 . 1 , . . . , 1 for model fitting.

E.4 Gaussian distance

The metric Gaussian distance , serving as a baseline, is defined as

where g lag is the Gaussian function, defined as

and c 1 , c 2 , c 3 , c 4 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 a threshold of 0.5 in the main text. This threshold is informed by both the quantitative distribution of individual head CMR distances and exploratory qualitative analysis of the attention patterns. First, as the bottom panel in Fig. 5e shows, the distribution of individual head CMR distances in GPT2 can be divided roughly into two clusters: a cluster peaking around 0.1 (and extending to around 1), and a spread-out cluster around 2.4. Second, we examined heads with 1 CMR distance and found non-human-like CRP (e.g., Fig. S2a,b). Further, many heads with CMR distance around 0.6-0.8 again show CRP different from humans (e.g., Fig. S2c, d). Finally, we also tested a threshold of 0.1 and found no qualitative difference (Fig. Fig. S3c, d).

F Comparison of composition mechanisms of induction heads and CMR

PE i : position embedding for the token at index/position i .

TE i : token embedding for the token at index/position i .

t i -1 : context vector associated with the word at index/position i .

f i : word embedding for the word at index/position i .

[ At i ] : information available at index/position i .

† : t k is updated from t k -1 by t j -1 (due to M FT exp ) and t IN j (due to M FT pre ), thus containing the information of t j (updated from t j -1 by t IN j ).

†† : The optimal Q-K match condition in Q-composition requires transformation from PE j to PE j +1 , implemented by the W Q W K matrix of H 2 . The optimal Q-K match condition in CMR requires transformation from t j -1 to t j , implemented by t IN j .

G Supplementary discussion

G.1 Scale difference

We have observed the scale difference between raw attention scores and those from human recall. We can speculate a few reasons. (i) It's common to tune the temperature only at the last readout layer, and it's unclear if Transformers benefit from variable temperatures in other attention layers. (ii) The range of attention scores might depend on the distributions of the model's training and evaluation data (e.g., evaluation prompts). (iii) Unlike LLMs, humans may not fully optimize next-word predictions, as biological objectives are more complex and varied. Recall might constitute only one use case, with the cognitive infrastructure also engaged in tasks like spatial navigation and adaptive decision-making. The more moderate scale in humans may thus reflect tradeoffs between multiple objectives.

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 MLP and attention layers may be akin to the activation-based working memory quickly updated by the cortico-hippocampal circuits.

H Experiment compute resources

The induction-head matching scores and copying scores of each head in GPT2-small and all Pythia models are computed using Google Colab Notebook. All models were pretrained and accessible through the TransformerLens library [25] with MIT License and used as is. See Table S2 for details.

Weused 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-2699v3, which has 72 threads and 512GB RAM per thread. This computation took a total of 90 hours.

Finally, model fitting was done on a 2023 MacBook Pro with 16GB RAM. All experiments were completed within 1 hour regardless of model size. This section contains all experiments we conducted that required non-trivial computational resources.

I Additional figures

0.0

0.5

1.0

CRP

enc

=0.1,

rec

=0.0

enc

=0.1,

rec

=0.1

enc

=0.1,

rec

=0.2

enc

=0.1,

rec

=0.3

enc

=0.1,

rec

=0.4

enc

=0.1,

rec

=0.5

enc

=0.1,

rec

=0.6

enc

=0.1,

rec

=0.7

enc

=0.1,

rec

=0.8

enc

=0.1,

rec

=0.9

enc

=0.1,

rec

=1.0

0.0

0.5

1.0

CRP

enc

=0.2,

rec

=0.0

enc

=0.2,

rec

=0.1

enc

=0.2,

rec

=0.2

enc

=0.2,

rec

=0.3

enc

=0.2,

rec

=0.4

enc

=0.2,

rec

=0.5

enc

=0.2,

rec

=0.6

enc

=0.2,

rec

=0.7

enc

=0.2,

rec

=0.8

enc

=0.2,

rec

=0.9

enc

=0.2,

rec

=1.0

0.0

0.5

1.0

CRP

enc

=0.3,

rec

=0.0

enc

=0.3,

rec

=0.1

enc

=0.3,

rec

=0.2

enc

=0.3,

rec

=0.3

enc

=0.3,

rec

=0.4

enc

=0.3,

rec

=0.5

enc

=0.3,

rec

=0.6

enc

=0.3,

rec

=0.7

enc

=0.3,

rec

=0.8

enc

=0.3,

rec

=0.9

enc

=0.3,

rec

=1.0

0.0

0.5

1.0

CRP

enc

=0.4,

rec

=0.0

enc

=0.4,

rec

=0.1

enc

=0.4,

rec

=0.2

enc

=0.4,

rec

=0.3

enc

=0.4,

rec

=0.4

enc

=0.4,

rec

=0.5

enc

=0.4,

rec

=0.6

enc

=0.4,

rec

=0.7

enc

=0.4,

rec

=0.8

enc

=0.4,

rec

=0.9

enc

=0.4,

rec

=1.0

0.0

0.5

1.0

CRP

enc

=0.5,

rec

=0.0

enc

=0.5,

rec

=0.1

enc

=0.5,

rec

=0.2

enc

=0.5,

rec

=0.3

enc

=0.5,

rec

=0.4

enc

=0.5,

rec

=0.5

enc

=0.5,

rec

=0.6

enc

=0.5,

rec

=0.7

enc

=0.5,

rec

=0.8

enc

=0.5,

rec

=0.9

enc

=0.5,

rec

=1.0

0.0

0.5

1.0

CRP

enc

=0.6,

rec

=0.0

enc

=0.6,

rec

=0.1

enc

=0.6,

rec

=0.2

enc

=0.6,

rec

=0.3

enc

=0.6,

rec

=0.4

enc

=0.6,

rec

=0.5

enc

=0.6,

rec

=0.6

enc

=0.6,

rec

=0.7

enc

=0.6,

rec

=0.8

enc

=0.6,

rec

=0.9

enc

=0.6,

rec

=1.0

0.0

0.5

1.0

CRP

enc

=0.7,

rec

=0.0

enc

=0.7,

rec

=0.1

enc

=0.7,

rec

=0.2

enc

=0.7,

rec

=0.3

enc

=0.7,

rec

=0.4

enc

=0.7,

rec

=0.5

enc

=0.7,

rec

=0.6

enc

=0.7,

rec

=0.7

enc

=0.7,

rec

=0.8

enc

=0.7,

rec

=0.9

enc

=0.7,

rec

=1.0

0.0

0.5

1.0

CRP

enc

=0.8,

rec

=0.0

enc

=0.8,

rec

=0.1

enc

=0.8,

rec

=0.2

enc

=0.8,

rec

=0.3

enc

=0.8,

rec

=0.4

enc

=0.8,

rec

=0.5

enc

=0.8,

rec

=0.6

enc

=0.8,

rec

=0.7

enc

=0.8,

rec

=0.8

enc

=0.8,

rec

=0.9

enc

=0.8,

rec

=1.0

0.0

0.5

1.0

CRP

enc

=0.9,

rec

=0.0

enc

=0.9,

rec

=0.1

enc

=0.9,

rec

=0.2

enc

=0.9,

rec

=0.3

enc

=0.9,

rec

=0.4

enc

=0.9,

rec

=0.5

enc

=0.9,

rec

=0.6

enc

=0.9,

rec

=0.7

enc

=0.9,

rec

=0.8

enc

=0.9,

rec

=0.9

enc

=0.9,

rec

=1.0

5

1 0 1

5

Lag

0.0

0.5

1.0

CRP

enc

=1.0,

rec

=0.0

5

1 0 1

5

Lag

enc

=1.0,

rec

=0.1

5

1 0 1

5

Lag

enc

=1.0,

rec

=0.2

5

1 0 1

5

Lag

enc

=1.0,

rec

=0.3

5

1 0 1

5

Lag

enc

=1.0,

rec

=0.4

5

1 0 1

5

Lag

enc

=1.0,

rec

=0.5

5

1 0 1

5

Lag

enc

=1.0,

rec

=0.6

5

1 0 1

5

Lag

enc

=1.0,

rec

=0.7

5

1 0 1

5

Lag

enc

=1.0,

rec

=0.8

5

1 0 1

5

Lag

enc

=1.0,

rec

=0.9

5

1 0 1

5

Lag

enc

=1.0,

rec

=1.0

a

b

c

d

a

b

c

d

a

c

b

a

c

f

g

h

i

j

k

l

m

n

o

p

q

r

s

t

u

v

w

x

b

e

d

a

b

d

c

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 5.2). We also suggest a 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 can be expected to generalize to other settings.

It is fine to include aspirational goals as motivation as long as it 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 no 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 a 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 on the factors that influence the performance of the approach. For example, a facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or a 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, a 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 important role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.

3. Theory Assumptions and Proofs

Question: For each theoretical result, does the paper provide the full set of assumptions and a complete (and correct) proof?

Answer: [NA]

Justification: This paper does not introduce any new theorems, formulas, or lemmas to be proved.

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 crossreferenced.

All assumptions should be clearly stated or referenced in the statement of any theorems.

The proofs can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged to provide a short proof sketch to provide intuition.

Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material.

Theorems and Lemmas that the proof relies upon should be properly referenced.

4. Experimental Result Reproducibility

Question: Does the paper fully disclose all the information needed to reproduce the main experimental results of the paper to the extent that it affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?

Answer: [Yes]

Justification: We fully disclose the models (available through the TransformerLens library [25]; CMR implementation is included in the code submission), metrics (see Appendix E), model fitting (see Section 5.1), and head ablation (see Section 5.3). All figures included in the paper can be reproduced using the included code (see Appendix A).

Guidelines:

The answer NA means that the paper does not include experiments.

If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.

If the contribution is a dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.

Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is a novel architecture, describing the architecture fully might suffice, or if the contribution is a specific model and empirical evaluation, it may be necessary to either make it possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g., in the case of a large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.

While NeurIPS does not require releasing code, the conference does require all submissions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example

(a) If the contribution is primarily a new algorithm, the paper should make it clear how to reproduce that algorithm.

(b) If the contribution is primarily a new model architecture, the paper should describe the architecture clearly and fully.

(c) If the contribution is a new model (e.g., a large language model), then there should either be a way to access this model for reproducing the results or a way to reproduce the model (e.g., with an open-source dataset or instructions for how to construct the dataset).

(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome to 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 instructions 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/ public/guides/CodeSubmissionPolicy ) 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 a 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: //nips.cc/public/guides/CodeSubmissionPolicy ) 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. If only a 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) is recommended, but including URLs to data and code is permitted.

6. Experimental Setting/Details

Question: Does the paper specify all the training and test details (e.g., data splits, hyperparameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?

Answer: [Yes]

Justification: 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 TransformerLens [25], we omitted their training details.

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

Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?

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, confidence intervals, or statistical significance tests, at least for the experiments that support the main claims of the paper.

The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions).

The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)

The assumptions made should be given (e.g., Normally distributed errors).

It should be clear whether the error bar is the standard deviation or the standard error of the mean.

It 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 a 96% CI, 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 are 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

Question: For each experiment, does the paper provide sufficient information on the computer resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?

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 in the paper (e.g., preliminary or failed experiments that didn't make it into the paper).

9. Code Of Ethics

Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics https://neurips.cc/public/EthicsGuidelines ?

Answer: [Yes]

Justification: We have reviewed the Neurips Code of Ethics and, to the best of our knowledge, we note no violation to the Code of Ethics by our current work.

Guidelines:

The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.

If the authors answer No, they should explain the special circumstances that require a deviation from the Code of Ethics.

The authors should make sure to preserve anonymity (e.g., if there is a special consideration due to laws or regulations in their jurisdiction).

10. Broader Impacts

Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?

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 or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment 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 particular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point it out. For example, it is legitimate to 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 to attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML).

11. Safeguards

Question: Does the paper describe 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)?

Answer: [NA]

Justification: We did not release any new datasets or new models. The paper only analyzes existing pretrained language models using a novel approach. We did not perceive any risks.

Guidelines:

The answer NA means that the paper poses no such risks.

Released models that have a 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 a best faith effort.

12. Licenses for existing assets

Question: 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?

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

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 not 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 is used and, if possible, include a 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.

If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswithcode.com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of a dataset.

For existing datasets that are re-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 creators.

13. New Assets

Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?

Answer: [Yes]

Justification: 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).

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 to 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 crowdsourcing nor research with human subjects.

Including this information in the supplemental material is fine, but if the main contribution 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 Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.

Institutional Review Board (IRB) Approvals or Equivalent for Research with Human Subjects

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 an 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 on 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 we expect authors to adhere to the NeurIPS Code of Ethics and the guidelines for their institution.

For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS:\nThe 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.\n\nWEAKNESSES:\nThe 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.\n\nQUESTIONS:\nSince 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?\n\nLIMITATIONS:\nStated in the paper.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-06"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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.\n\nSTRENGTHS:\nFrom 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.\n\nWEAKNESSES:\nWhile 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.\n\nQUESTIONS:\nThe 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?\n\nLIMITATIONS:\nThey are discussed appropriately.", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 2, "SOUNDNESS_CORRECTNESS": 2, "DATE": "2024-07-10"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS:\n1. 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.\n\nWEAKNESSES:\n1. 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.\n\nQUESTIONS:\nI 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.\n\nLIMITATIONS:\nOverall, 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.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-11"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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.\n\nSTRENGTHS:\n1. 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.\n\nWEAKNESSES:\n1. The paper lacks clarity, making it difficult to understand.\n2. The analysis has several issues, which will be discussed in the limitations section.\n\nQUESTIONS:\nI have mentioned them in the limitations section below.\n\nLIMITATIONS:\nWhile 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?", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 2, "SOUNDNESS_CORRECTNESS": 2, "DATE": "2024-07-12"}