Memory Retrieval in Transformers: Insights from The Encoding Specificity Principle

TL;DR

This paper investigates memory retrieval in transformer-based LLMs through the lens of the Encoding Specificity Principle, positing that attention implements cue-based, content-addressable memory via $Q$ as retrieval cues, $K$ as memory indices, and $V$ as content storage. It provides empirical support across two experiments: (1) attention-swapping experiments showing that manipulating $V$ and/or $K$ drives memory retrieval errors and hallucinations consistent with a memory-structured attention system, and (2) K-matrix perturbations demonstrating that keywords act as retrieval cues and that disabling them impairs memory recall; a small subset of neurons appears specialized for keyword-driven memory. The work identifies model-specific memory neurons and demonstrates a method to extract keywords, enabling downstream use in unlearning and privacy-preserving applications, while acknowledging limitations in perturbation scope and multi-word cue handling. Overall, it argues that attention in transformers behaves in memory-like, ESP-consistent ways, offering a pathway to targeted memory manipulation and improved interpretability. $Attention = softmax\left( \frac{QK^T}{\sqrt{d_k}} \right) V$ captures the core computation linking cues, traces, and content within this memory framework.

Abstract

While explainable artificial intelligence (XAI) for large language models (LLMs) remains an evolving field with many unresolved questions, increasing regulatory pressures have spurred interest in its role in ensuring transparency, accountability, and privacy-preserving machine unlearning. Despite recent advances in XAI have provided some insights, the specific role of attention layers in transformer based LLMs remains underexplored. This study investigates the memory mechanisms instantiated by attention layers, drawing on prior research in psychology and computational psycholinguistics that links Transformer attention to cue based retrieval in human memory. In this view, queries encode the retrieval context, keys index candidate memory traces, attention weights quantify cue trace similarity, and values carry the encoded content, jointly enabling the construction of a context representation that precedes and facilitates memory retrieval. Guided by the Encoding Specificity Principle, we hypothesize that the cues used in the initial stage of retrieval are instantiated as keywords. We provide converging evidence for this keywords-as-cues hypothesis. In addition, we isolate neurons within attention layers whose activations selectively encode and facilitate the retrieval of context-defining keywords. Consequently, these keywords can be extracted from identified neurons and further contribute to downstream applications such as unlearning.

Figures (18)

• Figure 1: Illustration of the hypothesized memory retrieval process in Transformer models, grounded in cue-based retrieval theories and the Encoding Specificity Principle. The example depicts a successful retrieval event contingent on a strong cue–trace match and the availability of the relevant content.
• Figure 2: $QKV$ swapping experiment results.
• Figure 3: Mean reciprocal rank of layer-head-dimension for each model across its respective book datasets.
• Figure 4: Radar graphs showingh average overall performance for all evaluated models when perturbed with different methods (individual book results can be found in Appendix \ref{['sec:eval_individual_book_all_model']}). All results are normalized to show higher is better.
• Figure 5: Top Memory Coefficient for Llama 2-7b. (Left): Top attention heads (Right): Top layer-head-dimension triplets