You could have invented transformers from memory
本文为原文前 6,000 字符的节选翻译,完整内容请查看原文。
You could have invented transformers from memory
你本可以凭记忆发明 Transformer。
Every layer is a memory store
每一层都是一个存储器
[code][interactive app] [代码][交互式应用]
Contents
目录
Retrieving a character from memory. The recall worsens as the number of memories grows. 从记忆中检索字符。随着记忆数量的增加,回想效果会变差。
Large language models today can recall obscure facts from every corner of the world. Most recently, I was running late for a friend’s board game night, so I asked Astra to remind me the rules to Dune: Imperium. It gladly narrated them to me as I drove over and helped me remember the rules well enough for me to play again. We’re so used to asking models about anything that we take their vast knowledge for granted. 如今的大型语言模型可以回想起世界各个角落的冷门事实。最近,我因为去参加朋友的桌游之夜迟到了,于是我让 Astra 提醒我《沙丘:帝国》的规则。它在我开车过去时愉快地为我讲述了规则,并帮我记住了足够多以便再次游玩的内容。我们已经习惯于向模型询问任何事情,以至于认为它们拥有的广博知识是理所当然的。
How do models remember so much information? And how much can they store? 模型是如何记住这么多信息的?它们又能存储多少信息?
In this post, I’ll go through an intuitive way to look at linear maps, MLPs, and self-attention as different kinds of memory stores. From this intuition we can derive some heuristics for how much memory each layer can store. 在这篇文章中,我将通过一种直观的方式,将线性映射、多层感知机(MLP)和自注意力机制视为不同类型的存储器。基于这种直觉,我们可以推导出关于每一层能存储多少记忆的一些启发式方法。
The key property that characterizes a “good” memory store is its ability to recall what it stored. After all, even if we could store infinitely many facts, it’s not very helpful if we can’t recall them when we need to. The more memories we store, the more they interfere with each other when we try to recall one. So throughout the post we’ll focus on one quantity, the interference-to-signal ratio, which measures how much the other memories get in the way during recall. The number of memories a store can hold before the interference grows as large as the signal is its capacity. 表征一个“好”存储器的关键属性是它回想所存储内容的能力。毕竟,即使我们能存储无限多的事实,如果无法在需要时回想起来,那也没什么用。我们存储的记忆越多,在尝试回想其中一个时,它们之间的干扰就越大。因此,在整篇文章中,我们将关注一个量——干扰信号比,它衡量了在回想过程中其他记忆造成了多大的阻碍。当干扰增长到与信号一样大时,存储器所能容纳的记忆数量即为其容量。
We’ll develop this perspective and show how you can go from a linear layer to something better, an MLP, which improves both capacity and recall, and then all the way to self-attention, which has very good recall and can store a potentially unbounded number of memories. With 64-dimensional keys, that’s a jump from about 64 memories, to about 2,700, to over a million. 我们将发展这一视角,并展示如何从线性层进阶到更好的 MLP(它同时改善了容量和回想效果),进而一直到自注意力机制(它具有非常好的回想效果,并且可以存储潜在无限数量的记忆)。对于 64 维的键,这意味着从大约 64 个记忆跃升至约 2,700 个,再到超过一百万个。
Here are the terms we’ll use throughout. A memory store is a function $f$ that holds $M$ memories, for $i = 1, \ldots, M$. A memory is one key-value pair: a key $k_i \in \reals^{d_k}$, used for lookup based on similarity, and a value $v_i \in \reals^{d_v}$, which is what we want to retrieve. To use $f$, we query it with an input $x$, and it recalls an output $f(x)$. A good memory store recalls $v_i$ when we query it with $k_i$, or with another vector similar to $k_i$. In the literature, a memory store that recalls by similarity is called an associative memory. 以下是我们通篇将使用的术语。存储器是一个函数 $f$,它持有 $M$ 个记忆,对于 $i = 1, \ldots, M$。一个记忆是一个键值对:一个用于基于相似度查找的键 $k_i \in \reals^{d_k}$,以及一个我们想要检索的值 $v_i \in \reals^{d_v}$。要使用 $f$,我们用输入 $x$ 对其进行查询,它会回想出一个输出 $f(x)$。一个好的存储器在用 $k_i$ 或与 $k_i$ 相似的另一个向量查询时,能够回想起 $v_i$。在文献中,这种通过相似度进行回想的存储器被称为联想记忆。
An outer product memory store
外积存储器
The most basic memory store is linear: it simply accumulates memories via outer products, 最基本的存储器是线性的:它仅仅通过外积来累积记忆,
$$W = \sum_{i=1}^{M} v_i k_i\transpose = V\transpose K$$ $$W = \sum_{i=1}^{M} v_i k_i\transpose = V\transpose K$$
where the rows of $K \in \reals^{M\times d_k}$ and $V \in \reals^{M\times d_v}$ are the keys $k_i\transpose$ and the values $v_i\transpose$. To store a new memory $(k, v)$, we simply add $vk\transpose$ to the existing weights. 其中 $K \in \reals^{M\times d_k}$ 和 $V \in \reals^{M\times d_v}$ 的行分别是键 $k_i\transpose$ 和值 $v_i\transpose$。要存储一个新的记忆 $(k, v)$,我们只需将 $vk\transpose$ 加到现有的权重中。
Aside: a gradient descent step writes a memory
旁注:梯度下降步骤写入一个记忆
Take a linear layer $y = Wx$ inside a network with a scalar loss $L$, and let $\bar y = \grad L(y)$ be the gradient of the loss with respect to the layer’s output. The gradient with respect to $W$ is $\bar y x\transpose$, so a step with learning rate $\eta$ is 考虑网络中一个带有标量损失 $L$ 的线性层 $y = Wx$,令 $\bar y = \grad L(y)$ 为损失相对于该层输出的梯度。相对于 $W$ 的梯度是 $\bar y x\transpose$,因此学习率为 $\eta$ 的一步更新为
$$W \leftarrow W - \eta, \bar y x\transpose.$$ $$W \leftarrow W - \eta, \bar y x\transpose.$$
This recurrence is a memory update: it updates the weights with a new memory with key $x$ and value $-\eta\bar y$. For more on this view, see Irie, Csordás & Schmidhuber, who express a linear layer trained by gradient descent as attention over all of its training inputs. 这种递归是一种记忆更新:它用一个键为 $x$、值为 $-\eta\bar y$ 的新记忆来更新权重。关于这一观点的更多信息,请参阅 Irie, Csordás & Schmidhuber 的研究,他们将通过梯度下降训练的线性层表达为对所有训练输入的注意力。
To query the memory store with an input $x$, we multiply the input by the matrix: 要用输入 $x$ 查询存储器,我们将输入乘以该矩阵:
$$f(x) = V\transpose K x = \sum_{i=1}^{M} v_i,(k_i\transpose x).$$ $$f(x) = V\transpose K x = \sum_{i=1}^{M} v_i,(k_i\transpose x).$$
So recall is a lookup by similarity to the stored keys: we compare the query with every key, and return the sum of the values weighted by those similarities. 因此,回想是通过与存储键的相似度进行查找:我们将查询与每个键进行比较,并返回由这些相似度加权后的值之和。
Recall picks up interference
回想会产生干扰
So storing memories is easy: just keep adding outer products to the weights. But what happens when we try to recall one? 所以存储记忆很容易:只需不断将外积加到权重中即可。但当我们尝试回想其中一个时会发生什么呢?
Let’s assume for simplicity that every key is normalized to have length 1, i.e. $k_j\transpose k_j = 1$. Then 为了简单起见,假设每个键都归一化为长度 1,即 $k_j\transpose k_j = 1$。那么
$$f(k_j) = \underbrace{v_j}{\text{signal}} + \underbrace{\sum{i\neq j} v_i,(k_i\transpose k_j)}{\text{interference}},$$ $$f(k_j) = \underbrace{v_j}{\text{signal}} + \underbrace{\sum_{i\neq j} v_i,(k_i\transpose k_j)}_{\text{interference}},$$
meaning we get back the value of the target memory, the one whose key we queried with, plus interference from every other memory whose key is similar to the query. 这意味着我们得到了目标记忆的值(即我们用其键进行查询的那个),加上来自每个键与查询相似的其他记忆的干扰。
If the keys were all orthonormal, the interference would be zero and recall would be exact. But $\reals^{d_k}$ only has $d_k$ orthogonal directions, so we cannot get perfect recall once we store more memories than the key dimension: $M > d_k$. 如果键都是正交归一的,干扰将为零,回想也将是精确的。但 $\reals^{d_k}$ 只有 $d_k$ 个正交方向,因此一旦我们存储的记忆数量超过键的维度($M > d_k$),就无法获得完美的回想。
What is the outer-product store’s capacity? Here’s an intuitive argument. There are $M-1$ interference terms. If the keys are random unit vectors, each similarity $k_i\transpose k_j$ is about $\pm 1/\sqrt{d_k}$. If the values are random too, with similar lengths, the $M-1$ interference terms add up like a random walk, so the interference’s squared length is about $(M-1)/d_k$ times the signal’s, and 外积存储器的容量是多少?这里有一个直观的论证。有 $M-1$ 个干扰项。如果键是随机单位向量,每个相似度 $k_i\transpose k_j$ 大约为 $\pm 1/\sqrt{d_k}$。如果值也是随机的,且长度相似,那么 $M-1$ 个干扰项会像随机游走一样累加,因此干扰的平方长度大约是信号的 $(M-1)/d_k$ 倍,并且
$$\frac{\norm{\text{interference}}}{\norm{\text{signal}}} \sim \sqrt{\frac{M}{d_k}}.$$ $$\frac{\norm{\text{interference}}}{\norm{\text{signal}}} \sim \sqrt{\frac{M}{d_k}}.$$
Setting this ratio to 1 gives a capacity of 将此比率设为 1 可得出容量为
$$M \approx d_k,$$ $$M \approx d_k,$$
once we store as many memories as the key dimension, the interference is as large as the signal, drowning it out. 一旦我们存储的记忆数量达到键的维度,干扰就会变得和信号一样大,从而将其淹没。
Each bar below is one stored memory’s weight in a single recall, relative to the target memory’s weight (in red). Every bar adds a bit of its value to the recall. Drag $M$ past $d_k = 64$ to watch the small bars pile up until recall fails. The other buttons switch to memory stores we’ll build later in the post. 下方每个条形图代表单次回想中一个已存储记忆的权重,相对于目标记忆的权重(红色)。每个条形都会将其值的一小部分添加到回想中。将 $M$ 拖动超过 $d_k = 64$,观察小条形堆积直到回想失败。其他按钮可切换到我们稍后将在文中构建的存储器。
How do we reduce interference?
我们如何减少干扰?
Even below its capacity of about $d_k$ memories, the outer-product store produces significant interference during recall. To raise the capacity, or recall more cleanly below it, we need to reduce interference. 即使在低于其约 $d_k$ 个记忆的容量时,外积存储器在回想过程中也会产生显著的干扰。为了提高容量,或在容量以下更清晰地回想,我们需要减少干扰。
The brute-force way is to make the key dimension $d_k$ large. 最简单粗暴的方法是增大键的维度 $d_k$。