Key-value memory in the brain

Abstract: Classical models of memory in psychology and neuroscience rely on similarity-based retrieval of stored patterns, where similarity is a function of retrieval cues and the stored patterns. While parsimonious, these models do not allow distinct representations for storage and retrieval, despite their distinct computational demands. Key-value memory systems, in contrast, distinguish representations used for storage (values) and those used for retrieval (keys). This allows key-value memory systems to optimize simultaneously for fidelity in storage and discriminability in retrieval. We review the computational foundations of key-value memory, its role in modern machine learning systems, related ideas from psychology and neuroscience, applications to a number of empirical puzzles, and possible biological implementations.

• The paper demonstrates that separate key and value representations enhance both memory discriminability and fidelity.
• It uses computational frameworks and machine learning parallels, such as the transformer’s self-attention mechanism, to model memory storage and retrieval.
• Implications suggest bridging neural mechanisms with artificial systems, prompting further experimental validations in neuroscience and AI.

Key-Value Memory in the Brain

The paper "Key-value memory in the brain" explores the application of key-value memory systems to understanding memory processes in both artificial intelligence and neuroscience. Key-value memory systems store distinct representations for retrieval (keys) and storage (values). This structure allows optimization for both discriminability and fidelity of memory, potentially offering insights into how the brain's memory processes may parallel high-performance computational models like transformers.

Introduction to Key-Value Memory Systems

Key-value memory systems differentiate between retrieval keys and storage values. Fundamentally, this allows distinct optimization for storage fidelity and retrieval discriminability, addressing the gap where traditional models rely solely on similarity-based retrieval without distinct representations. Key-value systems conceptualize inputs as transformed into both keys and values, stored in memory and accessed by matching queries to keys. This paper posits that such a framework may align with processes observed within the human brain, specifically proposing a division within the hippocampus and cortical structures.

Computational Foundations

The paper reviews the computational underpinnings of key-value memories, emphasizing their role in storing and recollecting memories based on optimal key-value pairings. It revisits correlations and kernels as a basis for these systems, where the interaction between keys and queries can utilize varying similarity kernels. This aligns with frameworks established in machine learning, such as the self-attention mechanism in transformers. Essentially, keys link to memories and queries, assessing similarity to retrieve associated values.

Neurobiological Substrates and Key Learning

The biological implementation of key-value memories involves hypothesizing the existence of learning rules akin to Hebbian learning—wherein synaptic strength increases with coincidental neuron firing. However, drawing exact parallels to biological systems involves incorporating non-Hebbian rules, emphasizing sparse learning patterns and potentially implicating the hippocampus in key storage. This section lays the groundwork for seeing parallels between artificial memories and biological memory capacities.

Figure 1: Two architectures for key-value memory. Black symbols denote vectors and blue symbols denote matrices. (Left) Input $\mathbf{x}$.

Figure 2: Optimization of key and value representations. Each point represents an event in the memory and belongs to one of (A) two or (B) three classes, represented by different colors. In each case, the evolution of key (Top row) and value (Bottom row) representations during the optimization process is shown; each row shows (Left) Random initialization, (Middle) trajectory of representations during the optimization process, with the final positions marked by gray points, (Right) final configuration.

Cognitive and Neural Evidence

The paper explores cognitive evidence, supporting a division between key and value memories, implying distinct roles for the hippocampus (keys) and neocortex (values). Behavioral observations such as retrieval interference indicate that forgetting often results from retrieval failures, not storage limits. Memory retrieval similarities between cognitive processes and key-value frameworks reinforce this proposed relationship between artificial and biological systems.

Figure 3: Forgetting and reactivation of memory events. A one-layer feedforward neural network is trained on two tasks sequentially, Task 1 and 2, constructed using the MNIST and FashionMNIST datasets, respectively. (A) The evolution of the test classification accuracy for the two tasks as a function of training epochs. After epoch 5, the training dataset changes from Task 1 to Task 2; resulting in forgetting of Task 1 as the model learns Task 2. (B) The accuracy of the trained model on Task 1 as a function of the value of the artificial scaler beta used to amplify the keys in all key-value memory pairs corresponding to Task 1 learning.

Implications and Future Research Directions

Key-value memory systems could potentially transform approaches to memory in both AI systems and our understanding of human cognition. It offers frameworks wherein both discriminability and fidelity are optimized, creating parallels between advanced AI systems like transformers and biological memory processes. Testing these ideas further with direct experimental study could illuminate the validity of these analogies, exploring the potential reversibility of certain neurological memory phenomena and examining how these models can resolve cognitive issues.

Conclusion

The paper provides a comprehensive exploration of key-value memory systems as a useful lens through which to view both artificial and biological memories. By distilling complex processes down to a separable key and value format, the parallels between breakthrough AI and neuroscience become apparent, hinting at integrated insights that could advance both fields. The exploration of this framework suggests future research could bridge gaps in understanding memory storage and retrieval, both in engineered systems and within the brain.