How much can a matrix remember?

1The problem attention solves

A sequence model reads vectors $x_1, x_2, \dots, x_T \in \R^d$, one per token. At the input these are token embeddings; deeper in the network they are the hidden states of the previous layer. At position $t$ the model has to produce an output that may depend on anything seen so far. Which earlier tokens matter is not fixed by position. In "the cat that the dog chased ran away", the verb "ran" needs "cat", five tokens back, and in another sentence the relevant word will be somewhere else. The model has to find earlier information by its content.

The simplest content-based lookup uses the dot product. For two vectors, $a^\top b = \sum_m a_m b_m = \|a\|\,\|b\|\cos\angle(a,b)$ is large when they point in similar directions, zero when they are orthogonal, and negative when they point apart. So position $t$ could score every earlier token by $x_t^\top x_i$ and return an average of the $x_i$ weighted by these scores.

This asks one vector to do three different jobs. As the thing doing the searching, $x_t$ describes what position $t$ is looking for. As the thing being searched, $x_i$ describes how position $i$ can be found. As the result, $x_i$ is also what gets returned. Using a single vector for all three causes concrete failures. A token scores highest against tokens that resemble itself, so it mostly retrieves near-copies of itself. The score is symmetric, so if "ran" finds "cat", "cat" finds "ran" equally well, although what each needs from the other is different. And what is returned is the raw token, with no choice about which of its features to pass on.

The repair is to give each job its own learned linear map:

The query $\Q{t}$ is what position $t$ is looking for. The key $\K{i}$ is how position $i$ advertises itself to queries. The value $\V{i}$ is what position $i$ hands over when it is found. Queries and keys are compared by a dot product, so they must live in the same space, $\Q{t}, \K{i} \in \R^{d_k}$. Values are never compared with anything, so their dimension $d_v$ is a separate choice. The three matrices are learned by gradient descent along with the rest of the network; this post takes them as given and asks what happens once the vectors exist.

2Softmax attention

The scores $\Q{t}^\top\K{i}$ can be any real numbers. To average values we want nonnegative weights that sum to one. Exponentiating makes every score positive, and dividing by the total makes the weights sum to one:

The exponential does more than make weights positive. A score higher by $1$ gets $e \approx 2.7$ times the weight, so if one key matches the query clearly better than the others, its value dominates the average and the lookup behaves almost like picking a single entry. The sum runs over $i \le t$ so that position $t$ never uses tokens that come after it, which is what makes the model usable for generation. This operation is attention (Vaswani et al., 2017). We leave out the $1/\sqrt{d_k}$ scaling and the use of several heads in parallel; neither changes anything below.

Now look at where $\Q{t}$ appears. It sits inside every exponential and inside the denominator, which couples all the keys together. Before the query arrives there is nothing useful to compute in advance, so every pair $(\K{i}, \V{i})$ has to be kept until it is needed. This store is the key-value cache, and it grows with $t$. Each new token must be compared with all earlier ones, so producing $T$ tokens costs time proportional to $T^2$.

Suppose instead we are only allowed a fixed number of numbers, a matrix $S \in \R^{d_v\times d_k}$, and we must write each pair into it as it arrives and discard the pair afterwards. Two questions follow. How should a pair be written? And how many pairs can such a matrix hold before reads start to go wrong?

3Removing the softmax

Replace $\exp(\Q{t}^\top\K{i})$ by $\Q{t}^\top\K{i}$ and drop the denominator (Katharopoulos et al., 2020). Each term $\V{i}\,(\K{i}^\top\Q{t})$ is a vector times a number, and that number is linear in $\Q{t}$. Writing $\K{i}^\top\Q{t}$ and then multiplying by $\V{i}$ is the same as first forming the matrix $\V{i}\K{i}^\top$ and then applying it to $\Q{t}$. Since every term has this form, the query can be pulled out of the sum:

The matrix $\Vb\Kb^\top \in \R^{d_v\times d_k}$ is the outer product of $\Vb$ and $\Kb$, with entries $(\Vb\Kb^\top)_{ab} = v_a k_b$. Pulling $\Q{t}$ out of the sum is the entire derivation. Everything the past contributes to any future read is collected into $S_t$ before the query is known. The pairs can be thrown away, and $S_t$ is updated by adding one term per step:

A model that carries a fixed-size state from one step to the next and updates it with each input is a recurrent network, and $S_t$ is its state. Storing a pair means adding its outer product to the state. Retrieving means multiplying the state by a query.

The same matrix appeared long before transformers, as the correlation matrix memory (Kohonen, 1972; Anderson, 1972), and Schlag et al. (2021) show that linear attention is a fast weight programmer, a network whose weights are rewritten at every step. To understand what this memory can and cannot do, we need to understand the object being added.

4One stored pair

Apply $\Vb\Kb^\top$ to an arbitrary input $x \in \R^{d_k}$:

Read the right-hand side from the inside out. The scalar $\Kb^\top x$ depends only on the component of $x$ along $\Kb$. That number is then used as a coefficient on $\Vb$. So the map ignores every direction of the input perpendicular to $\Kb$, and it can only ever output multiples of $\Vb$. The set of inputs sent to zero, the null space, is the set of vectors perpendicular to $\Kb$. The set of possible outputs, the image, is the line through $\Vb$. A matrix whose image is a single line has rank 1, and every rank-1 matrix is an outer product of this kind.

Querying with the key itself gives $(\Vb\Kb^\top)\Kb = \Vb\,\|\Kb\|^2$. The stored value comes back scaled by the squared length of the key, so for exact recall we want keys of length one. We assume $\|\K{i}\|=1$ from here on.

In the figures $d_k = d_v = 2$, so keys, queries and values are arrows in a plane and everything can be drawn. Nothing in the argument depends on the dimension being 2. The left panel of each figure is the space of keys and queries, the right panel is the space of values, and the matrix in between is the map from one to the other.

5Two pairs, and the Gram matrix

Store a second pair. The memory is $S = \V{1}\K{1}^\top + \V{2}\K{2}^\top$, and a read is the sum of what each outer product does on its own:

The query is dropped onto each key separately, and each drop weights its own value. Querying with $\K{1}$ gives $\V{1}\,\|\K{1}\|^2 + \V{2}\,(\K{2}^\top\K{1}) = \V{1} + \cos\theta\,\V{2}$, where $\theta$ is the angle between the keys. The second term is interference: part of the other value comes back as well, in proportion to how much the keys overlap.

The same computation works for any number of pairs. Stack the $n$ keys as the columns of $\KK = [\K{1}\;\cdots\;\K{n}] \in \R^{d_k\times n}$ and the values as the columns of $\VV = [\V{1}\;\cdots\;\V{n}] \in \R^{d_v\times n}$. The sum of outer products is $S = \VV\KK^\top$, and querying with every stored key at once gives

The $n \times n$ matrix $G$ of all dot products between keys is the Gram matrix. Column $j$ of $\VV G$ is the read for $\K{j}$, namely $\sum_i G_{ij}\,\V{i}$, so column $j$ of $G$ lists how much of each value comes back. The diagonal entries are $\|\K{j}\|^2 = 1$, the wanted value at full strength. The off-diagonal entries are the cosines between keys, and each one is the weight of a wrong value leaking into a read.

Every stored value is read back exactly if and only if $G = I$, which means the keys are orthonormal: unit length and mutually perpendicular. The values can point anywhere. In $\R^{d_k}$ at most $d_k$ vectors are mutually perpendicular, so at most $d_k$ pairs can be stored without interference, and in the plane a third key always overlaps the first two. The capacity depends on the key dimension alone. The values are only written down; the keys are what has to be kept apart. Schlag et al. (2021) discuss this limit for linear transformers.

6The best a fixed matrix can do

Adding outer products is one way to fill $S$. Is it the best? Suppose all $n$ pairs are available and we choose the matrix whose reads are closest to the stored values, measured by the total squared error:

The subscript $F$ denotes the Frobenius norm, the square root of the sum of squared entries; the two expressions are the same sum. This is a least squares problem. When the keys are linearly independent, $G$ is invertible and the solution is

Every value is recalled exactly, even for keys that overlap. To see what $G^{-1}$ does, write $S^\star = \sum_i \V{i}\,\tilde{k}_i^\top$, where $\tilde{k}_i$ is column $i$ of $\KK G^{-1}$. These vectors satisfy

So $\tilde{k}_i$ has dot product one with its own key and zero with every other key. The vectors $\tilde{k}_i$ are called the dual basis of the keys. The least squares memory stores each value under the dual of its key, and since the dual of $\K{1}$ is perpendicular to $\K{2}$, querying with $\K{2}$ picks up nothing of $\V{1}$.

There is a second way to read the same formula. A query in the span of the keys can be written in exactly one way as $\Qb = \sum_j c_j \K{j}$, and taking the dot product with $\tilde{k}_i$ shows $c_i = \tilde{k}_i^\top\Qb$. Hence $S^\star\Qb = \sum_i c_i \V{i}$: the least squares memory is the linear map that sends each key to its value, applied to the coordinates of the query in the basis of keys. The additive memory weights $\V{i}$ by $\K{i}^\top\Qb$, a perpendicular drop onto each key separately; the least squares memory weights it by $c_i$, which is found by walking to the query along the key directions. The two agree exactly when the keys are orthonormal.

Exact recall of overlapping keys has a cost, and Figure 3 shows it. For two unit keys at angle $\theta$, the dual vectors have length $\|\tilde{k}_i\| = 1/\sin\theta$. As the keys approach each other the dual vectors grow without bound, and $S^\star$ stretches the direction in which the two keys differ by a large factor. The largest stretch of a matrix, its largest singular value $\sigma_{\max}$, is the long semi-axis of the ellipse. A query that is slightly off $\K{1}$ then gives a read that is far from $\V{1}$. Least squares separates nearby keys by amplifying the small difference between them, and it amplifies small errors in the query by the same factor.

When the keys are linearly dependent, which is unavoidable once $n > d_k$, $G$ is not invertible. The minimizer is then $S^\star = \VV\KK^{+}$, where $\KK^{+}$ is the pseudoinverse, and $S^\star$ is the best compromise in squared error rather than an exact memory. This optimal linear associative memory was described by Kohonen and Ruohonen (1973). Linear attention cannot use it. Forming $G$ requires the dot product between every pair of keys, and a recurrent memory discards each key after writing it.

7How much can a matrix remember?

Section 5 answered this for keys we choose. A trained network does not choose its keys to be perpendicular; they are $W_K x_i$ for whatever tokens arrive, and they overlap. To see what happens then, take the simplest model of overlapping keys: $n$ keys drawn independently and uniformly from the unit sphere in $\R^{d_k}$, and values drawn independently with mean zero and $\mathbb{E}\|\V{i}\|^2 = 1$. We measure the average squared read error $e(n) = \mathbb{E}\,\|S\K{j} - \V{j}\|^2$, which is on the same scale as a stored value.

Additive memory. The error in reading $\K{j}$ is the interference term, $S\K{j} - \V{j} = \sum_{i\neq j}\V{i}\,(\K{i}^\top\K{j})$. Different values are uncorrelated, so the cross terms average to zero and

The expectation can be found without integrating. The distribution of keys does not change under rotations, so we may rotate until $\K{j}$ is the first coordinate axis. Then $\K{i}^\top\K{j}$ is just the first coordinate of $\K{i}$. The squared coordinates of a unit vector add up to 1, and by symmetry none of the $d_k$ coordinates is special, so each squared coordinate has mean $1/d_k$. With $n - 1$ interfering pairs,

The error grows in a straight line from the first extra pair onward. There is no threshold: the interference reaches the size of the signal at $n = d_k + 1$, and keeping the error below a fraction $\varepsilon$ allows only about $\varepsilon\, d_k$ pairs.

Least squares memory. For $n \le d_k$, random keys are linearly independent with probability one, so $e_{\text{ls}}(n) = 0$. For $n > d_k$ the reads are $S^\star\KK = \VV\,\KK^{+}\KK$, and $\KK^{+}\KK$ is the orthogonal projection of $\R^n$ onto a $d_k$-dimensional subspace (the row space of $\KK$). Each row of $\VV$ is a vector in $\R^n$ holding one coordinate of all $n$ values. If the values are Gaussian, this row points in a uniformly random direction independent of the keys, and projecting it onto a $d_k$-dimensional subspace of $\R^n$ keeps, on average, a fraction $d_k/n$ of its squared length. The lost fraction is the error:

This memory is exact up to exactly $d_k$ pairs and then degrades, reaching half the signal at $n = 2d_k$.

This answers the question in the title for the two memories we have. A matrix in $\R^{d_v\times d_k}$ can hold $d_k$ pairs exactly if it is fitted with all the keys in hand. Filled by adding outer products, as linear attention does, it recalls a single pair exactly and holds about $\varepsilon\, d_k$ pairs at relative error $\varepsilon$. The value dimension appears in neither answer. Everything a linear associative memory can store is limited by how many directions its keys can occupy.

References

• Anderson, J. A. (1972). A simple neural network generating an interactive memory. Mathematical Biosciences, 14(3–4), 197–220.
• Choromanski, K., Likhosherstov, V., Dohan, D., Song, X., Gane, A., Sarlos, T., et al. (2021). Rethinking attention with Performers. International Conference on Learning Representations.
• Katharopoulos, A., Vyas, A., Pappas, N., & Fleuret, F. (2020). Transformers are RNNs: Fast autoregressive transformers with linear attention. International Conference on Machine Learning.
• Kohonen, T. (1972). Correlation matrix memories. IEEE Transactions on Computers, C-21(4), 353–359.
• Kohonen, T., & Ruohonen, M. (1973). Representation of associated data by matrix operators. IEEE Transactions on Computers, C-22(7), 701–702.
• Schlag, I., Irie, K., & Schmidhuber, J. (2021). Linear transformers are secretly fast weight programmers. International Conference on Machine Learning.
• Tsai, Y.-H. H., Bai, S., Yamada, M., Morency, L.-P., & Salakhutdinov, R. (2019). Transformer dissection: A unified understanding of transformer's attention via the lens of kernel. Empirical Methods in Natural Language Processing.
• Vaswani, A., Shazeer, N., Parmar, N., Uszkoreit, J., Jones, L., Gomez, A. N., Kaiser, Ł., & Polosukhin, I. (2017). Attention is all you need. Advances in Neural Information Processing Systems.