Analyzing Linear Maps, MLPs, and Self-Attention as Memory Stores in LLMs

Memory in Large Language Models

Large language models (LLMs) demonstrate extensive knowledge recall, a capability often taken for granted. This analysis investigates the mechanisms behind this memory, specifically how much information models can store and retrieve effectively. The core idea is to view different neural network components as distinct types of memory stores.

Evaluating Memory Stores

A key characteristic of an effective memory store is its ability to recall stored information accurately. As the number of stored memories increases, interference between them can hinder recall. The analysis introduces the "interference-to-signal ratio" to measure this phenomenon. The capacity of a memory store is defined as the number of memories it can hold before this interference becomes as significant as the signal itself.

Architectural Comparison of Memory

The analysis progresses from simpler linear layers to Multi-Layer Perceptrons (MLPs) and then to self-attention mechanisms, evaluating each as a memory store. Linear layers have limited capacity, while MLPs offer improvements in both capacity and recall. Self-attention is presented as having very good recall and a potentially unbounded memory capacity. For instance, with 64-dimensional keys, the memory capacity can jump from approximately 64 memories in a linear layer, to about 2,700 in an MLP, and over a million in a self-attention mechanism.

Defining Memory Stores

In this context, a memory store is a function that holds 'M' memories, each consisting of a key-value pair. A key is a vector used for lookup based on similarity, and a value is the information to be retrieved. The goal of a good memory store is to recall the correct value when queried with its corresponding key or a similar vector. This type of memory store, which recalls by similarity, is known as an associative memory.