PaperlayerView exact revision on arXiv

Source-linked AI summary

Linking In-context Learning in Transformers to Human Episodic Memory

Li Ji-An, Corey Y. Zhou, Marcus K. Benna, Marcelo G. Mattar

arXiv:2405.14992v2cs.CLcs.LG

TL;DR

The paper addresses limited understanding of how interacting Transformer heads relate to biological cognition by comparing induction heads with human episodic memory and CMR. It reports behavioral and mechanistic parallels, CMR-like biases in intermediate-to-late layers, and a possible role in ICL.

  • Problem

    How interacting Transformer attention heads relate to biological cognition remains unclear, despite links between biology and other AI architectures.

    Source: §1 · pp. 1–2

  • Method

    The study compares induction heads, which contribute to LLM in-context learning, with CMR, a model of human episodic memory.

    Source: §1 · p. 2

  • Results

    Induction heads show behavioral and mechanistic similarities to CMR, while CMR-like heads emerge in intermediate-to-late layers and may play a causal role in ICL.

    Source: §1 · p. 2, §6 · p. 10

  • Takeaways & Limitations

    The findings provide evidence for a computational bridge between Transformer models and human episodic memory.

    Source: §1 · p. 2, §6 · p. 10

    The conclusions may not hold for untested Transformer models, and the causal analysis cannot isolate the contribution of specific CMR-like characteristics.

    Source: §6 · p. 10

Abstract

from arXiv · show

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

Prior work has linked several AI architectures to biological systems, but how interacting Transformer attention heads relate to biological cognition remains unclear. This study examines induction heads and reports parallels with human episodic memory and CMR.

1 Introduction

  • Earlier work connected convolutional, recurrent, and feedforward networks with biological activity or behavior, while Transformer biology remained less explored.Existing Transformer links included Hopfield networks, language-cortex activity, and hippocampo-cortical representations.Source: §1 · pp. 1–2
  • The study examines whether interacting Transformer attention heads relate to biological cognition by comparing induction heads with episodic memory and CMR.Induction heads are a crucial component of in-context learning (ICL) in LLMs.Source: §1 · p. 2
  • The paper reports that induction heads and CMR are mechanistically and behaviorally similar, and characterizes CMR-like behavior and its possible role in ICL.The authors present the findings as evidence for a novel bridge between Transformer models and episodic memory.Source: §1 · p. 2

2 Next-token prediction and memory recall

The section describes parallel tests of associative retrieval: Transformers predict from repeated-token prompts, while people freely recall studied words. Both tasks reveal how earlier sequence structure informs retrieval.

2 Next-token prediction and memory recall

  • In repeated-token prompts, ICL is assessed by whether a model predicts a token that followed the current token in an earlier sequence repeat.The first occurrence cannot be predicted when the tokens have no learned structure, but the repeated context provides a temporal association.Source: §2 · p. 2
  • In free recall, participants study words sequentially and later recall them in any order, yet their recalls reflect the study-list's temporal structure.Retrieving one word can cue words studied close in time.Source: §2 · p. 2
  • Conditional response probability (CRP) measures the chance that the next recalled word came from a given lag relative to the most recently recalled word.Human recall shows temporal contiguity and a forward asymmetry, with words studied later recalled more often than words studied earlier.Source: §2 · pp. 2–3

3 Transformer models and induction heads

The section explains how Transformer components communicate through a shared residual stream, then describes induction-head behavior and two mechanisms for composing attention across layers. It positions Q-composition as a link to CMR.

3.1 Residual stream and interacting heads

  • Each token has a residual stream that components read from and write to, allowing heads in different layers to pass information between computations.The final residual stream is used to generate logits for next-token prediction.Source: §3.1 · p. 3

3.2 Induction heads and their attention patterns

  • Induction heads show temporal contiguity and forward asymmetry, attending more to nearby than distant tokens and to positive than negative lags.Figure 2 illustrates the L5H1 induction stripe and plots its attention by relative position lag.Source: §3.2 · p. 3, Figure 2 · p. 4

3.3 K-composition and Q-composition induction heads

  • K-composition uses a previous-token head's written token embedding as a key, whereas Q-composition uses a duplicate-token head's written position embedding as a query.Both are induction mechanisms in which a first-layer head's information informs second-layer attention.Source: §3.3 · p. 4
  • The authors rewrite Q-composition as a prefix-matching induction head to relate induction-head attention biases to those found in human episodic memory.Their main mechanistic comparison is between Q-composition and CMR; K-composition serves as a point of comparison.Source: §3.3 · p. 4

4 Contextual maintenance and retrieval model (CMR)

Section 4 describes CMR’s evolving temporal context and maps its memory operations onto a two-layer Q-composition Transformer circuit. It connects context updating and retrieval to attention-head operations.

4.1 CMR in its original form

  • CMR models episodic recall as associative retrieval, using temporal context that evolves from prior state and input context during encoding and retrieval.β controls temporal drift, while ρ is chosen to keep the context vector unit norm.Source: §4.1 · pp. 4–5
  • When 0 < β < 1, slowly evolving context makes nearby items temporally contiguous, while context entering only after an item yields forward asymmetry.Larger β values correspond to sharper CRPs, stronger forward asymmetry, and stronger temporal clustering.Source: §4.1 · p. 5
  • Figure 4 contrasts exact sequential chaining at βenc = βrec = 1 with human-like free recall at βenc = βrec = 0.7, which sometimes skips words.The human comparison includes 171 PEERS participants; panel c varies γFT.Source: Figure 4 · p. 6

4.2 CMR as an induction head

  • In the Transformer mapping, a first-layer duplicate-token head retrieves the context associated with a matching word, while an MLP-like layer supplies pre-experimental context.The first-layer operation asks where the same token was encountered; the pre-experimental context represents information absent from the experiment.Source: §4.2 · p. 6
  • At retrieval, MTF acts as a second-layer Q-composition head, and CMR’s associative memory updates correspond to causal linear attention.The retrieved memory is compared with studied words to produce recall probabilities, analogous to unembedding logits.Source: §4.2 · pp. 6–7
  • CMR resembles a two-layer Q-composition head: a recurrently updated context serves downstream attention, unlike fixed position embeddings.The authors suggest Transformer heads might acquire a similar circuit using autoregressively updated residual-stream context.Source: Figure 3 · p. 5, §4.2 · pp. 6–7

5 Experiments

The experiments assess how closely induction heads resemble CMR behavior, how CMR-like patterns develop during training, and whether CMR-like heads matter for in-context learning. They find CMR-like behavior across model layers and training, while head ablations are associated with worse ICL scores.

5.1 Quantifying the similarity between an induction head and CMR

  • Induction-head matching and positive copying scores were associated with smaller CMR distances, and CMR also captured some heads not typically considered induction heads.One such example is a head whose activity peaks at lag=0.Source: §5.1 · p. 7
  • Most heads with lower CMR distances appeared in intermediate-to-late layers of GPT2-small, Pythia, Qwen-7B, Mistral-7B, and Llama3-8B.The layer-wise figure reports the percentage of heads with CMR distance below 0.5.Source: §5.1 · p. 7, Figure 6 · p. 8
  • Across 12 models, CMR significantly outperformed an equal-parameter Gaussian baseline for top induction heads: average distances were 0.11 (top 20), 0.05 (top 50), 0.12 (top 100), and 0.12 (top 200), all p < 0.0001.Source: §5.1 · pp. 7–8

5.2 CMR-like heads develop human-like temporal clustering over training

  • As Pythia’s loss on the designed prompt fell, temporal clustering increased, especially in intermediate layers, alongside rising inverse temperature.The increase in βenc and βrec occurred during the first 10 training checkpoints, when loss significantly dropped.Source: §5.2 · pp. 8–9
  • Heads with stronger induction-head matching or smaller CMR distances showed greater temporal clustering, with fitted β values in a range similar to human recall data.Under CMR, asymmetric contiguity bias with β < 1 provides an optimal recall policy when encoding and retrieval are noisy.Source: §5.2 · p. 9

5.3 CMR-like heads are causally relevant for ICL capability

  • Most tested models had worse ICL scores after ablating the top 10% of CMR-like heads than after ablating the same number of random heads.The authors caution that ablation may trigger compensation by other heads and cannot isolate the causal role of episodic-memory features from other head properties.Source: §5.3 · p. 9

6 Discussion

The discussion presents CMR-like induction-head behavior as a connection between Transformer in-context learning and human episodic memory. It outlines implications for memory research and mechanistic interpretability while noting limits on the causal and model-wide conclusions.

6 Discussion

  • CMR-like attention biases emerge in intermediate-to-late LLM layers, evolve toward human memory biases, and may play a causal role in in-context learning.The discussion identifies asymmetric contiguity as a CMR-like bias.Source: §6 · pp. 9–10
  • CMR-like behavior may connect Transformer attention to hippocampal episodic retrieval, alongside prior neural-network models of memory and recall.The proposed link is consistent with work on neural networks trained for free recall and hippocampal representations.Source: §6 · p. 10
  • The link could motivate alternative episodic-memory mechanisms, although K- and Q-composition circuits require positional encoding.The proposed mapping of positional encoding to grid-like cells tracking space and time is speculative.Source: §6 · p. 10
  • CMR offers a behavioral lens on induction heads, but some CMR-fit heads do not meet traditional induction-head criteria, and deeper-model mechanisms remain speculative.The discussion also proposes that some heads may encode multiple future tokens.Source: §6 · p. 10

B Attention mechanisms in Transformer models

C Detailed introduction of CMR

  • CMR combines an evolving temporal context with word representations and word-to-context and context-to-word association matrices.Its context reflects recency-weighted history, with separate temporal dynamics for encoding and retrieval.Source: C Detailed introduction of CMR · pp. 15, 17
  • During recall, a recalled word reinstates its associated input context alongside the current context, favoring recall of the following studied word.The account links this process to temporal contiguity and forward asymmetry in human free recall.Source: C Detailed introduction of CMR · p. 17

D Prompt design

  • The prompt design uses two copies of a permuted sequence of common English tokens to represent study and recall phases.The tokens are selected from the 100 most common English tokens with a leading space.Source: D Prompt design · p. 17

E.1 Metrics for induction heads

The induction-head metrics assess prefix matching and token copying, two components of the match-then-copy behavior.

E.1 Metrics for induction heads

  • An ideal prefix-matching head scores 1 on the induction-head matching metric.The target pattern attends to earlier positions whose preceding token matches the current token.Source: §E.1 · p. 18
  • An ideal token-copying head scores 1, and induction heads should have positive copying scores because copying increases logits.The copying score measures a head's tendency to increase the attended token's logit.Source: §E.1 · p. 18

E.2 CRP analysis of attention heads

The CRP analysis summarizes average attention scores at different token-position lags.

E.2 CRP analysis of attention heads

  • The lag profile averages attention scores by offset; lag 0 measures attention to first-occurrence tokens, and lag 1 measures attention to the immediately following token.The analysis uses lags from −5 to 5.Source: §E.2 · p. 18

E.3 CMR distance

The section defines distances for comparing attention-head lag behavior with CMR predictions and a Gaussian baseline.

E.3 CMR distance

  • CMR distance compares lag-wise CRPs with CMR predictions drawn from a specified parameter grid.The grid varies encoding and retrieval parameters and the word-to-context association parameter.Source: §E.3 · p. 18

E.4 Gaussian distance

  • Gaussian distance serves as a baseline using a Gaussian function with four coefficients.The passage identifies the function and coefficients but does not provide their values.Source: §E.4 · p. 19

E.5 Choosing thresholds for CMR distance

The section sets an operational cutoff for identifying CMR-like heads, compares induction-head and CMR composition mechanisms, and discusses scale differences and a proposed mapping between Transformer components and brain systems.

E.5 Choosing thresholds for CMR distance

  • The study uses a CMR-distance threshold of 0.5, informed by head-distance clusters and qualitative attention-pattern analysis; testing 0.1 produced no qualitative difference.Heads near 1 and many around 0.6–0.8 showed non-human-like conditional response probabilities.Source: §E.5 · p. 19

F Comparison of composition mechanisms of induction heads and CMR

G.1 Scale difference

  • The authors report a scale difference between raw attention scores and human recall, speculating about temperature tuning, data distributions, and humans’ multiple objectives.They suggest recall is only one use of cognitive infrastructure that also supports spatial navigation and adaptive decision-making.Source: §G.1 · p. 20

G.2 Mapping components in Transformer models to the brain

  • The authors propose mapping MLP layers to cortex and self-attention layers to hippocampus, with key-value associations stored in fast Hebbian-like hippocampal synapses.They also compare residual-stream updates to activation-based working memory updated by cortico-hippocampal circuits.Source: §G.2 · p. 21

H Experiment compute resources

The resource section describes how model metrics and CMR fitting were computed. The additional figures show example CMR patterns and track CMR distances and fitted parameters across layers, training, and models.

H Experiment compute resources

  • Induction-head matching and copying scores were computed for pretrained models using Google Colab; Table S2 describes computation times for scoring heads across available checkpoints.The models were accessed through TransformerLens and used as is.Source: H Experiment compute resources · p. 21, Table S2 · p. 21
  • Computing the CRP subset for model fitting took 90 hours on an internal cluster, while fitting experiments ran on a 2023 MacBook Pro and each took under one hour.The cluster comprised six nodes with 256GB RAM per thread and four nodes with 512GB RAM per thread.Source: H Experiment compute resources · p. 21

I Additional figures

  • Figure S1 contrasts conditional response distributions for γFT = 0, using only pre-experimental contexts, and γFT = 1, using only experimental contexts; Figure S2 gives example GPT2 heads with CMR distances roughly 0.5–1.0.Source: Figure S1 · p. 22, Figure S2 · p. 23
  • Figure S3 describes layerwise CMR distances and the relative positions of the two lowest-distance layers across GPT2-small, Pythia models, Qwen-7B, Mistral-7B, and Llama3-8B.It also shows percentages of heads below 0.1 in GPT2-small and Pythia models.Source: Figure S3 · p. 23
  • Figure S4 tracks fitted βenc, βrec, and γFT parameters across Pythia-70m training for heads with CMR distance below 0.5.Solid lines average across selected heads; shading shows the fitted-value range.Source: Figure S4 · p. 24
  • Figure S5 tracks CMR distances and fitted parameters for the top CMR-like heads across Pythia models from 70m to 12b.Top heads are ranked by lowest CMR distance; the figure also shows their average fitted inverse temperature.Source: Figure S5 · p. 25
  • Figure S6 presents CMR-distance proportions across layers and compares fitted parameters and distances for top induction heads, top CMR-like heads, and additional models.Top induction heads are selected by induction-head matching score, while top CMR-like heads have the lowest CMR distances.Source: Figure S6 · p. 26
Loading 2405.14992v2…