You could have invented transformers from memory

Transformers can retain a vast array of facts from around the globe. Recently, the author needed to remember the rules of the board game Dune: Imperium, and used a language model to help. This demonstrates the models' impressive memory capabilities. But how do these models remember so much, and what limits their storage? The author explores this question by examining linear maps, Multi-Layer Perceptrons (MLPs), and self-attention as different types of memory stores.

A "good" memory store should be able to retrieve what it stored, but as we add more memories, they interfere with each other and make recall more difficult. The interference-to-signal ratio measures this interference. The capacity of a memory store is the number of memories it can hold before interference outweighs the original signal.

The author explains how linear layers, MLPs, and self-attention stores work, and how they improve both capacity and recall. By using 64-dimensional keys, the capacity increases from about 64 memories to about 2,700, and potentially over a million. The key terms used include memory store, which is a function f that holds M memories, each with a key-value pair.

The query function f(x) retrieves the stored value v_i when given the key k_i, or a similar vector. A memory store that recalls by similarity is called an associative memory. The basic linear store simply accumulates memories through outer products, where the rows of K ∈ ℝ^(M×d_k) and V ∈ ℝ^(M×d_v) are the keys and values. To store a new memory (k, v), it adds vk^T to the existing weights.

This update is equivalent to a gradient step with learning rate η. The query function multiplies the input by the matrix W, and recall is a lookup by similarity to the stored keys. When keys are normalized to have length 1, the interference is zero and recall is exact. However, in practice, with random unit vectors, interference is significant once we store as many memories as the key dimension.

The author discusses different approaches to reducing interference and increasing capacity, such as increasing the key dimension d_k or improving the algorithm by storing memories more carefully or recalling selectively.