你本可以从记忆中推导出 Transformer

# 你本可以从记忆中"发明"Transformer 来源：https://alexwang.ai/posts/associative-memory/ 从记忆中检索一个字符。随着记忆数量增多，召回效果会变差。 如今的大语言模型能够从世界各个角落回忆起冷僻的事实。最近，我去参加朋友的桌游之夜迟到了，于是我让 Astra 提醒我《沙丘：帝国》的规则。在我开车的路上，它愉快地为我讲述规则，帮我记住了足够再次上手游戏的要点。我们已经习惯了向模型询问任何事情，以至于对它们庞大的知识习以为常。模型是如何记住如此多信息的？*它们到底能存储多少？* 在本文中，我将通过一种直觉性的方式来理解线性映射、MLP 和自注意力作为不同类型记忆存储的角色。从这种直觉出发，我们可以推导出每层能存储多少记忆的一些经验法则。 --- 区分"好"记忆存储的关键特性是它召回所存储内容的能力。毕竟，即使我们能存储无限多条事实，如果我们无法在需要时召回它们，那也没有多大用处。存储的记忆越多，当我们试图召回其中一条时，它们彼此之间*相互干扰*就越严重。因此，在整篇文章中，我们将聚焦于一个量——**干扰信号比**，它衡量在召回过程中其他记忆造成的干扰程度。在干扰增长到与信号一样大之前，一个存储器能容纳的记忆数量就是它的**容量**。我们将建立这一视角，并展示如何从线性层走向更好的东西——MLP，它同时提升了容量和召回质量，然后一路走到自注意力，它具有非常好的召回能力，并且可以存储潜在无限数量的记忆。在 64 维键的情况下，这意味着从大约 64 条记忆跃升到大约 2,700 条，再到超过一百万条。 以下是我们将在全文中使用的术语。一个**记忆存储**是一个函数 $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$。在文献中，通过相似度进行召回的记忆存储被称为*联想记忆*。 ## 外积记忆存储 最基本的记忆存储是线性的：它简单地通过外积来累积记忆， $$W = \sum\_{i=1}^{M} v\_i k\_i\transpose = V\transpose K$$ 其中 $K \in \reals^{M \times d\_k}$ 和 $V \in \reals^{M \times d\_v}$ 的行分别是键 $k\_i\transpose$ 和值 $v\_i\transpose$。[^1] 要存储一条新记忆 $(k, v)$，我们只需将 $vk\transpose$ 加到现有的权重上。[^2] [^1]: 为简洁起见省略了偏差项。 [^2]: 这就是赫布式学习规则（Hebbian learning rule）："一起激活的神经元连接在一起。" **旁注：梯度下降的一步就是写入一条记忆** 考虑网络中的一个线性层 $y = Wx$，标量损失为 $L$，令 $\bar y = \grad L\(y\)$ 为损失对该层输出的梯度。关于 $W$ 的梯度为 $\bar y x\transpose$，所以学习率为 $\eta$ 的一步更新为： $$W \leftarrow W - \eta\, \bar y x\transpose.$$ 这个递推关系就是一次记忆更新：它用键为 $x$、值为 $-\eta\bar y$ 的新记忆来更新权重。关于这种视角的更多信息，参见 [Irie, Csordás & Schmidhuber](https://arxiv.org/abs/2202.05798)，他们将通过梯度下降训练的线性层表达为对其所有训练输入的注意力。 要用输入 $x$ 查询记忆存储，我们将输入乘以矩阵： $$f\(x\) = V\transpose K x = \sum\_{i=1}^{M} v\_i\,(k\_i\transpose x).$$ 所以召回就是基于与存储键的相似度进行查找：我们将查询与每个键比较，然后返回按这些相似度加权的值之和。 ### 召回会引入干扰 存储记忆很容易：只需不断向外积权重添加即可。但当我们试图召回一条记忆时会发生什么？为了简化，我们假设每个键都归一化为长度为 1，即 $k\_j\transpose k\_j = 1$。那么 $$f\(k\_j\) = \underbrace{v\_j}\_{\text{信号}} + \underbrace{\sum\_{i \neq j} v\_i\,(k\_i\transpose k\_j)}\_{\text{干扰}},$$ 这意味着我们得到了**目标记忆**（我们用其键查询的那条）的值，加上来自每条键与查询相似的其他记忆的干扰。如果所有键都是正交归一的，干扰为零，召回是精确的。但 $\reals^{d\_k}$ 只有 $d\_k$ 个正交方向，因此一旦我们存储的记忆数超过键的维度：$M > d\_k$，就无法获得完美的召回。 外积存储的容量是什么？这里有一个直观的论证。存在 $M-1$ 个干扰项。如果键是随机单位向量，每个相似度 $k\_i\transpose k\_j$ 大约为 $\pm 1/\sqrt{d\_k}$。如果值也是随机的且长度相似，这 $M-1$ 个干扰项会像随机游走一样叠加，因此干扰的平方长度大约是信号的 $(M-1)/d\_k$ 倍，于是 $$\frac{\norm{\text{干扰}}}{\norm{\text{信号}}} \sim \sqrt{\frac{M}{d\_k}}.$$ 将这个比值设为 1，得到容量为 $$M \approx d\_k,$$ 一旦我们存储的记忆数与键的维度相同，干扰就和信号一样大，将信号淹没。 下方每根柱子代表一次召回中某条已存储记忆的权重（相对于目标记忆的权重，以红色显示）。每根柱子都会为召回贡献一点它的值。将 $M$ 拖过 $d\_k = 64$，观察小柱子不断累积直到召回失败。其他按钮切换到我们在本文后面将构建的记忆存储。 已存储的记忆数 *M* = 32　锐度 *β* = 8　目标记忆　每条其他已存储记忆，相对于目标记忆的权重　干扰 / 信号 = 召回失败 ## 如何减少干扰？ 即使在其约 $d\_k$ 条记忆的容量之下，外积存储在召回时也会产生显著的干扰。[^3] 要提高容量，或在容量之下获得更干净的召回，我们需要*减少干扰*。最直接的方式是增大键的维度 $d\_k$，因为线性存储的容量随 $d\_k$ 增长。$d\_k$ 能有多大受硬件约束限制。或者，我们可以进行算法上的改进。我们可以更仔细地存储记忆：外积存储在添加每条记忆时不考虑其他记忆，因此它无法抵消它们之间的干扰。我们也可以更有选择性地召回：线性存储中的每条记忆都贡献干扰，即使它的键与查询的相似度很低，所以我们可以抑制相似度较低的键来改进召回。 [^3]: 尽管如此，在实践中它效果还不错。 在本文的后续部分，我们将看到更仔细地存储外积记忆可以得到最优的线性记忆存储，其容量约为两倍。通过非线性进行选择性召回则给了我们 MLP 和自注意力。 ## 最优线性记忆存储 外积存储在写入每条记忆时不考虑其他记忆。相反，我们可以选择权重 $W$，使得召回能够同时拟合所有已存储的记忆。这是一个线性回归问题：[^4] $$\operator\*{minimize}\_W \sum\_{i=1}^M \norm{W k\_i - v\_i}\_2^2.$$ [^4]: 一个有趣的历史注记：这正是 Anderson 和 Kohonen 在 1984 年提出的形式，他们独立于神经网络之外发展了联想记忆的研究。 当存储的记忆数超过键的维度，$M > d\_k$，最小二乘解为 $$W = \parens\*{\sum\_i v\_i k\_i\transpose}\parens\*{\sum\_i k\_i k\_i\transpose}\inverse = V\transpose K\,(K\transpose K)\inverse.$$ 这就是外积存储 $V\transpose K$ 加上一个改动：在我们将查询与键比较之前，先乘以 $(K\transpose K)\inverse$，它抵消了键的协方差 $K\transpose K$。这*白化*了键，使存储不再过度加权许多键共享的方向。 当 $M \leq d\_k$，只要键是线性无关的，回归可以精确拟合每条记忆。因此最优存储完全没有干扰地召回，即使键不是正交的。当 $M > d\_k$，我们开始产生干扰，因为记忆存储只有 $d\_k$ 个方向可以使用。对于随机的键和值，我们的干扰信号比变为 $$\frac{\norm{\text{干扰}}}{\norm{\text{信号}}} \approx \sqrt{\frac{M}{d\_k} - 1}.$$ 这个比值在 $$M = 2d\_k$$ 处等于 1，所以最优存储的容量是之前的*两倍*！ **旁注：推导最优线性记忆存储** **当 $M \leq d\_k$ 时。** 如果键是线性无关的，键之间的内积矩阵 $KK\transpose$（$M\times M$）是可逆的，且 $$W = V\transpose (KK\transpose)\inverse K$$ 精确拟合每条记忆。当 $M < d\_k$ 时，问题是欠约束的，有无穷多解，这个解是最小范数解。$(KK\transpose)\inverse K$ 的行 $\tilde k\_i\transpose$ 被称为 $K$ 的行的对偶基：由构造可知 $(KK\transpose)\inverse KK\transpose = I$，因此 $\tilde k\_i\transpose k\_j = \delta\_{ij}$。于是 $f\(k\_j\) = \sum\_i v\_i\, \tilde k\_i\transpose k\_j = v\_j$，*没有干扰*，我们有完美的召回。 **当 $M > d\_k$ 时。** 将目标函数的梯度设为零得到 $W K\transpose K = V\transpose K$，因此 $W = V\transpose K\,(K\transpose K)\inverse$。两个解都可以写成 $W = V\transpose (K\transpose)\pinv$，其中 $A\pinv$ 表示 $A$ 的伪逆：当键线性无关时 $(K\transpose)\pinv = (KK\transpose)\inverse K$，当键张成 $\reals^{d\_k}$ 时 $(K\transpose)\pinv = K(K\transpose K)\inverse$。 现在我们可以分析 $f\(k\_j\) = V\transpose K\,(K\transpose K)\inverse k\_j$ 的干扰。对于一个已存储的键 $k\_j$，每个值 $v\_i$ 上的权重是 $k\_i\transpose (K\transpose K)\inverse k\_j$，这恰好是 $M\times M$ 矩阵 $P = K(K\transpose K)\inverse K\transpose$ 的第 $(i, j)$ 个元素： $$f\(k\_j\) = \sum\_{i=1}^M v\_i\, P\_{ij} = \underbrace{P\_{jj}\, v\_j}\_{\text{信号}} + \underbrace{\sum\_{i \neq j} P\_{ij}\, v\_i}\_{\text{干扰}}.$$ $P$ 是到 $K$ 列空间的投影，它是 $\reals^M$ 中的一个 $d\_k$ 维子空间。关于投影的两条事实帮助我们分析干扰信号比： - **信号项：** 由迹的循环性质， $$\sum\_{j=1}^M P\_{jj} = \trace{K(K\transpose K)\inverse K\transpose} = \trace{(K\transpose K)\inverse K\transpose K} = \trace{I\_{d\_k}} = d\_k,$$ 所以 $M$ 个信号权重加起来等于 $d\_k$。如果键是随机的，没有哪个键特殊，因此 $P\_{jj} \approx d\_k/M$。 - **干扰项：** $P$ 是对称投影矩阵，因此 $P\transpose P = P^2 = P$。于是 $$\sum\_{i \neq j} P\_{ij}^2 = \sum\_{i=1}^M P\_{ij}^2 - P\_{jj}^2 = (P\transpose P)\_{jj} - P\_{jj}^2 = P\_{jj} - P\_{jj}^2 \approx \frac{d\_k}{M} - \parens\*{\frac{d\_k}{M}}^2.$$ 如果值是随机的且长度相似， $$\frac{\norm{\text{干扰}}}{\norm{\text{信号}}} \approx \frac{\sqrt{P\_{jj} - P\_{jj}^2}}{P\_{jj}} = \sqrt{\frac{M}{d\_k} - 1}.$$ 与外积存储的 $\sqrt{M/d\_k}$ 相比，$-1$ 就是最优存储抵消掉的 $d\_k$ 条记忆的干扰量。 我们可以数值地验证预测。在 $d\_k = 64$ 维空间中随机采样 $M$ 条单位长度键的记忆，用其中一条已存储的键查询记忆存储器，测量干扰信号比：召回失败：干扰大于信号 0.11 10 16 64 256 1024 4096 16384 已存储记忆数 *M*（键维度 $d\_k = 64$）干扰 / 信号 最优：$M \leq 64$ 时无干扰 外积（赫布式）记忆，$M = 16$：干扰/信号 = 0.46 外积（赫布式）记忆，$M = 32$：干扰/信号 = 0.68 外积（赫布式）记忆，$M = 64$：干扰/信号 = 0.97 外积（赫布式）记忆，$M = 128$：干扰/信号 = 1.39 外积（赫布式）记忆，$M = 256$：干扰/信号 = 1.97 外积（赫布式）记忆，$M = 512$：干扰/信号 = 2.81 外积（赫布式）记忆，$M = 1024$：干扰/信号 = 3.90 外积（赫布式）记忆，$M = 2048$：干扰/信号 = 5.63 外积（赫布式）记忆，$M = 4096$：干扰/信号 = 8.01 外积（赫布式）记忆，$M = 8192$：干扰/信号 = 11.38 外积（赫布式）记忆，$M = 16384$：干扰/信号 = 15.90 外积 最优（最小二乘）记忆，$M = 16$：干扰/信号 = 0.00 最优（最小二乘）记忆，$M = 32$：干扰/信号 = 0.00 最优（最小二乘）记忆，$M = 64$：干扰/信号 = 0.00 最优（最小二乘）记忆，$M = 65$：干扰/信号 = 0.08 最优（最小二乘）记忆，$M = 66$：干扰/信号 = 0.14 最优（最小二乘）记忆，$M = 68$：干扰/信号 = 0.23 最优（最小二乘）记忆，$M = 72$：干扰/信号 = 0.35 最优（最小二乘）记忆，$M = 80$：干扰/信号 = 0.49 最优（最小二乘）记忆，$M = 96$：干扰/信号 = 0.71 最优（最小二乘）记忆，$M = 112$：干扰/信号 = 0.86 最优（最小二乘）记忆，$M = 128$：干扰/信号 = 1.00 最优（最小二乘）记忆，$M = 256$：干扰/信号 = 1.72 最优（最小二乘）记忆，$M = 512$：干扰/信号 = 2.63 最优（最小二乘）记忆，$M = 1024$：干扰/信号 = 3.81 最优（最小二乘）记忆，$M = 2048$：干扰/信号 = 5.56 最优（最小二乘）记忆，$M = 4096$：干扰/信号 = 7.95 最优（最小二乘）记忆，$M = 8192$：干扰/信号 = 11.34 最优（最小二乘）记忆，$M = 16384$：干扰/信号 = 15.91 最优（最小二乘） 外积存储在 $M \approx d\_k = 64$ 处达到其容量。最优存储在 $M = d\_k$ 之前精确召回。超过之后，干扰从零开始以 $\sqrt{M/d\_k - 1}$ 的速度增长，存储在 $M = 2d\_k = 128$ 附近达到容量，与预测一致。由于它抵消了固定数量记忆的干扰，两种存储之间的差距随 $M$ 增大而缩小。 这个最优记忆存储仍然是对相似度的线性函数。它必须以相同的因子缩放大和小的相似度。要做得更好，我们需要改变存储的召回方式。 ## MLP 作为非线性记忆存储 线性记忆存储的召回方式为 $$f\(x\) = \sum\_{i=1}^M v\_i \(k\_i\transpose x\).$$ 如果我们有办法在 $k\_i\transpose x$ 较小时抑制权重，那么我们就可以减少低相似度键带来的干扰。让我们在计算值的加权和之前，先将每个相似度通过一个非线性函数 $\sigma$： $$f\(x\) = \sum\_{i=1}^{M} v\_i\,\sigma\(k\_i\transpose x\) = V\transpose \sigma\(Kx\).$$ 我们刚刚从第一性原理重新发现了一个两层 MLP！第一个权重矩阵的行是键，第二个权重矩阵的列是值。[Sukhbaatar et al.](https://arxiv.org/abs/1907.01470) 指出了