Memory in Language Models

Where Is Factual Memory Located in Language Models?

Problem framing

Here we revise the mechanistic question:

A necessary distinction:

Memory type	Source	Typical mechanism
Contextual memory	Tokens in the current prompt	Retrieved and routed through Attention
Parametric memory	Information learned during training	Encoded across model weights
Factual recall	Runtime transformation of a prompt into an answer	Joint computation involving attention, MLPs, and the residual stream

Attention as Contextual Memory

Attention can be interpreted as a soft dictionary lookup:

$$\mathrm{Attention}(Q,K,V)=\mathrm{softmax}\left(\frac{Q{K}^{\top }}{\sqrt{d_k}}\right)V.$$

• Keys: specify what information matches a query.
• Values: specify what information is returned.
• Attention weights: determine the mixture of retrieved values.
• Keys and values are generally constructed from the current context.

The factual-memory literature asks whether an analogous lookup occurs over information stored in the model’s parameters.

MLPs as Parametric Associative Memories

Feed-Forward Layers as Key–Value Memories

Mor Geva et al., Transformer Feed-Forward Layers Are Key-Value Memories, interpret Transformer MLPs as collections of learned key–value associations.

For an MLP,

$$\mathrm{MLP}(x)=W_{\mathrm{out}}\sigma (W_{\mathrm{in}}x),$$

a neuron-wise decomposition gives

$$\mathrm{MLP}(x)=\sum _i\sigma (k_i^{\top }x)v_i.$$

• $k_i$: a row of $W_{\mathrm{in}}$ acting as a learned key.
• $\sigma (k_i^{\top }x)$: the degree to which the input matches the key.
• $v_i$: a column of $W_{\mathrm{out}}$ contributing an output direction.
• The MLP output is a mixture of activated parameter-stored values.

Main empirical findings:

• Lower-layer keys tend to match shallow or lexical patterns.
• Higher-layer keys tend to match more semantic patterns.
• MLP outputs shift the residual stream toward tokens compatible with the activated patterns.
• Individual MLP memories are refined across subsequent residual layers.

Conclusion: MLPs behave like associative memories, but the paper does not establish that each fact occupies one unique neuron or address.

Knowledge Neurons

Damai Dai et al., Knowledge Neurons in Pretrained Transformers, identify MLP neurons associated with particular relational facts in BERT.

Method:

1. Prompt the model using a factual cloze task:

   The capital of Ireland is [MASK].

2. Use attribution methods to identify neurons contributing strongly to the correct answer.
3. Suppress or amplify those neuron activations.
4. Measure the change in the probability of the factual answer.

Findings:

• Suppressing selected neurons reduces expression of their associated facts.
• Amplifying them can increase factual-answer probability.
• Selected neurons can sometimes be used to update or erase factual behavior.

Conclusion: Small neuron sets can be causally associated with factual expression, but factual representations are not necessarily atomically localized.

Causal Localization of Factual Recall

Causal Tracing

Kevin Meng et al., Locating and Editing Factual Associations in GPT, introduce causal tracing, also called representation denoising. Given:

The Space Needle is located in the city of Seattle.

the procedure uses three runs.

1. Clean run
   • Run the model normally.
   • Record hidden states $h_{\ell ,t}^{\mathrm{clean}}$.
2. Corrupted run
   • Add noise to the subject representation.
   • The probability of the correct answer falls.
3. Restoration run
   • Restore one clean activation inside the corrupted computation:

$$h_{\ell ,t}^{\mathrm{corrupt}}\leftarrow h_{\ell ,t}^{\mathrm{clean}}.$$

If this restoration recovers the answer, the activation lies on a causally important path.

Intervention	Question answered
Ablation	Does performance deteriorate when this component is removed?
Causal tracing	Does restoring this component rescue a corrupted computation?
Activation patching	What information does a component transmit between conditions?
Weight editing	Can changing these parameters modify the model’s factual behavior?

ROME Localization Result

The causal-tracing experiments found two broad sites:

• Middle-layer MLP activity at the final subject token

  • Strongly associated with recovering factual information.

  • Interpreted as an early factual-association recall stage.

• Later attention activity near the output position

  • Routes or transfers recovered information into the position used for prediction.

A simplified pipeline is:

$$\begin{array}{rlrl}\text{subject representation}& \stackrel{\text{middle-layer MLPs}}{\to }\text{attribute-compatible representation}& \stackrel{\text{later attention}}{\to }\text{output position}& \stackrel{\text{unembedding}}{\to }\text{answer token}.\end{array}$$

Example:

$$\text{France}\stackrel{\text{parametric association}}{\to }\text{Paris direction}\stackrel{\text{routing}}{\to }P(\text{Paris})\uparrow .$$

Rank-One Model Editing

ROME treats a middle-layer MLP as an associative memory and applies a rank-one parameter update to change a factual mapping:

$$W^{\prime }=W+u{v}^{\top }.$$

Desired effect:

Eiffel Tower → Rome

instead of:

Eiffel Tower → Paris

A successful edit should satisfy:

• Reliability: the edited prompt produces the new fact.

• Generalization: paraphrased prompts also produce the new fact.

• Specificity: unrelated facts remain approximately unchanged.

Initial ROME conclusion: factual associations are sufficiently localized in middle-layer MLP computations to permit direct, targeted editing.

Why Localization Is Not Equivalent to Storage

Localization–Editing Mismatch

Peter Hase et al., Does Localization Inform Editing?, test whether causal localization identifies the best parameters to edit.

Natural hypothesis:

$$\text{high causal-tracing score}\Rightarrow \text{best layer for editing}.$$

Their results substantially weaken this hypothesis:

• Layers identified by causal tracing were not reliably the layers where editing worked best.

• Facts could be changed by editing parameters outside the supposedly localized site.

• Layer identity itself could predict editing effectiveness better than localization scores.

A component may be causally necessary because it:

• routes information;

• changes representation format;

• performs error correction;

• exposes information to the output;

• participates in—but does not uniquely store—the association.

Conclusion: causal tracing locates important computational dependencies, not necessarily unique physical storage sites.

Factual Recall as a Multi-Stage Pipeline

Argument Passing and Function Application

Ang Lv et al., Interpreting Key Mechanisms of Factual Recall in Transformer-Based Language Models, propose a staged interpretation of factual recall.

For:

The capital of France is ...

the pipeline is approximately:

1. Topic extraction

   • Task-relevant attention heads identify or copy the topic token: France.

2. Activation control

   • A following MLP amplifies or suppresses information contributed by different heads.

3. Implicit function application

   • A deeper MLP transforms the topic representation toward the answer:

$$\mathrm{getCapital}(\text{France})\approx \text{Paris}.$$

4. Output regulation

   • Late components may strengthen or suppress the correct answer.

Conclusion: attention and MLPs have complementary roles.

Component	Typical role in factual recall
Attention heads	Identify entities, relations, or relevant positions; route information
MLPs	Transform and amplify representations; apply association-like mappings
Residual stream	Shared communication channel carrying intermediate representations
Unembedding	Converts residual directions into vocabulary logits

Knowledge as a Distributed Circuit

Knowledge Circuits

Yunzhi Yao et al., Knowledge Circuits in Pretrained Transformers, move from neuron- or layer-level localization to causal subgraphs.

A knowledge circuit may include:

• information heads;
• relation heads;
• mover heads;
• MLP blocks;
• token positions;
• residual-stream edges.

A typical abstract circuit is:

$$\text{entity identification}\to \text{relation computation}\to \text{attribute construction}\to \text{answer transport}.$$

Possible functional roles:

• Information heads: extract subject or entity information.
• Relation heads: identify what attribute is requested.
• MLPs: construct or strengthen relation-specific answer directions.
• Mover heads: transfer relevant information to the output position.

Conclusion: factual knowledge is better modelled as a sparse computational circuit than as a single knowledge neuron.

Redundant, Distributed, Non-Contiguous Retrieval

Hail Hochman et al., Factual Retrieval in LLMs Is a Redundant, Distributed and Non-Contiguous Process, use iterative activation patching to identify minimal layer subsets required for attribute computation.

Findings:

• Minimal factual-retrieval paths usually involve multiple layers.
• Participating layers may be non-adjacent.
• Computation can skip intermediate layers.
• Several functionally equivalent paths may retrieve the same fact.
• The same layer may participate in many different factual computations.

This helps explain:

• why individual-neuron ablations often have incomplete effects;
• why causal localization and optimal editing locations can disagree;
• why successful recall depends on the prompt and intervention;
• why factual behavior can survive partial model pruning.

Whole-Layer Ablation

Cornerstone Layers

Yang Zhang et al., Investigating Layer Importance in Large Language Models, evaluate individual layers using Shapley-style importance estimation and layer ablation. Reported pattern:

• Certain early cornerstone layers have disproportionately large global importance.
• Removing one can cause catastrophic performance loss.
• Removing many non-cornerstone layers causes comparatively modest degradation.
• Layer-importance profiles vary across architectures.

However:

$$\text{large ablation effect}\nRightarrow \text{memory storage site}.$$

An early layer may be indispensable because later layers expect representations in a particular format.

Analogy:

Removing a database parser may prevent access to every stored record, but the parser is not itself the database.

Conclusion: whole-layer ablation measures global functional dependency, not specifically factual-memory localization.

Comparison of Major Claims

Evolution of the Literature

Work	Unit studied	Method	Main conclusion
Geva et al.	MLP neurons and layers	Pattern and output analysis	MLPs behave like learned key–value memories
Dai et al.	Individual neurons	Attribution and activation intervention	Small neuron sets can causally affect factual expression
Meng et al.	Layer-position activations and MLP weights	Causal tracing and ROME	Middle MLPs mediate editable factual associations
Hase et al.	MLP layers	Localization versus editing comparison	Causal localization does not reliably identify optimal edit sites
Lv et al.	Heads, MLPs, residual states	Component decomposition and intervention	Recall is an attention–MLP pipeline
Yao et al.	Computational subgraphs	Circuit discovery	Facts are expressed by collaborative knowledge circuits
Zhang et al.	Whole layers	Layer ablation and Shapley importance	Some layers are globally indispensable, but not necessarily storage sites
Hochman et al.	Layer paths	Iterative activation patching	Retrieval paths are distributed, redundant, and non-contiguous

Design Principles for Interpreting Ablation Results

Causal Interpretation Rules

• Ablation establishes dependency, not exclusive storage.
• Restoration establishes causal mediation, conditional on the corruption used.
• Decodability does not establish use: information may be linearly recoverable but ignored by downstream computation.
• Editability does not establish original storage: an alternative parameter site may override existing behavior.
• Prompt dependence matters: different prompts may activate different computational paths.
• Redundancy masks localization: removing one route may expose a substitute route.
• Granularity changes conclusions: neuron-, head-, layer-, and circuit-level analyses answer different questions.
• Residual connections complicate removal: skipping a block is not equivalent to deleting all information associated with it.

Major Drawbacks and Warnings

• Polysemanticity: one neuron or direction may participate in unrelated concepts.
• Superposition: many features may share the same representational dimensions.
• Corruption sensitivity: causal-tracing results depend on how the subject representation is damaged.
• Metric sensitivity: answer probability, exact accuracy, calibration, and generation quality can rank layers differently.
• Model-family dependence: GPT-2, BERT, Llama, OPT, and mixture-of-experts models need not share identical localization patterns.
• Fact-expression ambiguity: changing an answer may alter retrieval, suppression, routing, or output decoding rather than stored knowledge itself.
• Behavioral editing risk: a successful edit may add a competing pathway instead of replacing the original representation.

Final Synthesis

Where Does the Memory Come From?

Parametric factual knowledge is encoded across learned weights, with MLP parameters playing a prominent association-like role. Its expression at inference time is produced by a distributed causal circuit involving attention, MLP transformations, residual-stream communication, and output decoding.

The thing to remember is:

$$\boxed{\text{stored parameters}\ne \text{single storage location}\ne \text{runtime recall path}}$$

A useful operational model is:

$$\begin{array}{rlrlrl}\text{prompt parsing}& \to \text{entity/relation identification}& \to \text{contextual routing by attention}& \to \text{association-like transformation by MLPs}& \to \text{distributed refinement}& \to \text{answer transport and decoding}.\end{array}$$

Historical progression:

$$\boxed{\text{MLPs as memories}\to \text{knowledge neurons}\to \text{middle-layer localization}\to \text{localization–editing mismatch}\to \text{distributed redundant circuits}}$$