Memory Retrieval in Transformers: Insights from The Encoding Specificity Principle

TL;DR

This paper explores how attention mechanisms in transformer-based language models function as memory retrieval systems, drawing parallels with human memory processes and providing insights into interpretability and unlearning.

Contribution

It introduces the keywords-as-cues hypothesis and identifies neurons in attention layers that encode context-defining keywords, linking psychological principles to model interpretability.

Findings

Attention weights reflect cue trace similarity.

Neurons selectively encode retrieval keywords.

Keywords can be extracted for downstream tasks.

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.