You could have invented transformers from memory

你可以从记忆中发明出Transformer | Alex Wang

你可以从记忆中发明出Transformer

每一层都是一个记忆存储器

[code] [交互应用]

目录

从记忆中检索一个字符。随着记忆数量的增加，回忆效果会变差。

当今的大型语言模型可以从世界各个角落回忆起冷门的事实。最近，我因为要参加朋友的桌游之夜而迟到了，所以我让Astra提醒我《沙丘：帝国》的规则。它愉快地向我叙述了这些规则，帮助我记住了足够的内容，让我能够再次游玩。我们已经习惯了向模型询问任何问题，以至于我们理所当然地认为它们拥有广博的知识。

模型是如何记住这么多信息的？它们能存储多少信息？

在这篇文章中，我将介绍一种直观的方式来理解线性映射、多层感知机（MLP）和自注意力机制作为不同类型的记忆存储器。从这个直觉出发，我们可以推导出每层能存储多少记忆的一些启发式方法。

一个“优秀”记忆存储器的关键特性是其回忆存储内容的能力。毕竟，即使我们能够存储无限多的事实，如果无法在需要时回忆它们，这种能力也没有太大用处。我们存储的记忆越多，在尝试回忆某一个记忆时，它们之间的干扰就越严重。因此，本文将重点关注一个指标——干扰与信号比，它衡量在回忆过程中其他记忆的干扰程度。当干扰增长到与信号相当时，存储器能够容纳的记忆数量即为其容量。

我们将发展这种观点，并展示如何从线性层过渡到更优的结构——多层感知机（MLP），它在容量和回忆能力方面都有提升，最终达到自注意力机制，它具有极佳的回忆能力，理论上可以存储无限数量的记忆。使用64维的键，这种提升从大约64个记忆跃升到约2700个，再到超过一百万个。

以下是全文将使用的术语。记忆存储器是一个函数 $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

附注：梯度下降步骤会写入一个记忆

在网络中取一个线性层 $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 的研究，他们将通过梯度下降训练的线性层表达为对其所有训练输入的注意力。

要使用输入 $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,$$

一旦存储的记忆数量达到键维度 $d_k$，干扰的大小就与信号相当，从而淹没信号。

下面的每个条形图表示单次回忆中一个存储记忆的权重，相对于目标记忆的权重（以红色显示）。每个条形图都会为其值贡献一部分到回忆中。拖动 $M$ 超过 $d_k = 64$，观察小条形图堆积直到回忆失败。其他按钮将切换到我们稍后在本文中构建的记忆存储。

score

线性

ReLU

ReLU²

softmax

存储的记忆

M

=

32

锐度

β

8

目标记忆

其他存储记忆，权重相对于目标记忆

干扰 / 信号 =

回忆失败

我们如何减少干扰？ #

即使在容量约为 $d_k$ 记忆的范围内，外积存储在回忆过程中仍会产生显著的干扰。3 要提高容量，或在容量以下更清晰地回忆，我们需要减少干扰。

一种粗暴的方法是增加键维度 $d_k$，因为线性存储的容量随 $d_k$ 增大而增大。$d_k$ 能达到多大受硬件限制。

另一种方法是进行算法改进。我们可以更谨慎地存储记忆：外积存储在添加每个记忆时不会考虑其他记忆，因此无法抵消它们的干扰。或者我们可以更选择性地回忆：线性存储中的每个记忆都会产生干扰，即使其键与查询的相似度较低，因此我们可以通过抑制相似度较低的键来改进回忆。

在本文的其余部分，我们将看到，通过更谨慎地存储外积记忆，我们可以获得最优的线性内存存储，其容量大约是原来的两倍。通过非线性选择性回忆，我们可以得到多层感知机（MLPs）和自注意力机制。

最优线性内存存储 #

外积存储在写入每个记忆时不会考虑其他记忆。相反，我们可以选择权重 $W$，使得回忆能够同时适应所有存储的记忆。这是一个线性回归问题：

$$\operatorname*{minimize}_W \sum_{i=1}^M \norm{W k_i - v_i}_2^2.$$

当存储的记忆数量超过键维度，即 $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{interference}}}{\norm{\text{signal}}} \approx \sqrt{\frac{M}{d_k} - 1}.$$

当这个比值等于 1 时：

$$M = 2d_k,$$

最优存储的容量是之前的两倍！

附注：推导最优线性内存存储

当 $M \leq d_k$。如果键是线性无关的，键之间的内积构成的 $M\times M$ 矩阵 $KK\transpose$ 是可逆的，且：

$$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{signal}} + \underbrace{\sum_{i\neq j} P_{ij}\, v_i}_{\text{interference}}.$$

$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{interference}}}{\norm{\text{signal}}} \approx \frac{\sqrt{P_{jj} - P_{jj}^2}}{P_{jj}} = \sqrt{\frac{M}{d_k} - 1}.$$

与外积存储的 $\sqrt{M/d_k}$ 相比，这个 $-1$ 表示最优存储抵消了 $d_k$ 个记忆的干扰量。

我们可以通过数值验证这些预测。我们随机采样 $M$ 个在 $d_k = 64$ 维空间中单位长度的键，用其中一个存储的键查询记忆存储，并测量干扰与信号的比率：

外积存储在 $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 等人指出前馈层与记忆存储之间的这种相似性，Geva 等人发现训练后的 Transformer 中的前馈层通常以这种方式工作，键会在输入中可识别的模式上激活。

理想情况下，门控函数 $\sigma$ 的选择应使得当相似性较低（$k_i\transpose x \sim 0$ 或甚至 $k_i\transpose x < 0$）时，求和中的对应项被抑制以减少干扰。ReLU 激活函数会将负相似性的键置零，添加偏置可以调整这个阈值 5 。平方 ReLU 激活函数进一步抑制了相似性接近零的键。相似性为 0.1 的值变为 0.01，而相似性为 1 的值保持为 1。事实上，平方 ReLU 是 MLP 激活函数中较少被提及但已在大规模应用中验证的有效改进之一（So 等人）。

平方 ReLU MLP 是一种 GLU MLP，其门控和输入投影共享权重，因为 $\operatorname{ReLU}(Kx)^2 = \operatorname{ReLU}(Kx) \odot Kx$。当今语言模型中的 SwiGLU MLP 给门控分配了独立的权重，并使用 Swish 激活函数代替 ReLU。

在实际应用中，多层感知机（MLP）通过叠加方式存储事实：多个事实共享相同的权重，输出是多个记忆的混合。Garcia 等人构建了一个显式的模型：GLU MLP 的隐藏单元计算输入的随机特征 $\phi(x)$，第二个权重矩阵是这些特征的外积存储 $\sum_i v_i \phi(k_i)\transpose$，且 MLP 存储的事实数量远多于其隐藏单元的数量。然而，训练后的 Transformer 的 ReLU MLP 对于特定输入可能仅激活 1% 的隐藏单元，因此每次回忆仅依赖于记忆的一小部分（Li 等人）。

Transformer 的作者早已知晓这一点

原始 Transformer 论文也提到了这一点，但仅在 arXiv 的 v2 和 v3 版本中。在这些版本中，Vaswani 等人将前馈层描述为“一种注意力机制”，其两个权重矩阵的行分别作为键和值，且使用 ReLU 替代 softmax。他们甚至尝试用学习得到的键值参数上的注意力替换前馈层，得到的翻译质量几乎相同：一个版本稍差，另一个稍好，但第一个版本训练速度慢约 30%。该段内容在后续所有版本中均被删除。我们将在下文探讨自注意力如何与这些内容相关联。

通过更高次幂减少干扰 #

平方相似度有效的原因在于，它对低相似度权重的压缩效果远大于对高相似度权重的压缩。我们可以进一步使用 $\sigma(s) = \operatorname{ReLU}(s)^p$。信号保持不变，因为 $1^p = 1$，而相似度为 $1/\sqrt{d_k}$ 的干扰项会缩小到 $d_k^{-p/2}$。$M$ 个干扰项的总和类似于随机游走，因此干扰与信号的比值为

$$\frac{\norm{\text{interference}}}{\norm{\text{signal}}} \approx \sqrt{\frac{M}{d_k^p}}.$$

这意味着容量约为

$$M \approx d_k^p$$

的记忆，比之前几何级数增长！

为验证这一点，我们对五种 $\sigma$ 的选择重复了之前的实验：

由于干扰随 $\sqrt{M}$ 增长，因此在对数-对数图上这些线几乎平行。只有它们的高度发生变化，从而影响容量，即每条线与 1 的交点：

• 线性：容量为 $M \approx d_k = 64$。

• ReLU：将负相似度置零使干扰减少 $\sqrt{2}$，因此容量增长至 $M \approx 2d_k = 128$。

• 平方 ReLU：每个干扰项从约 $1/\sqrt{d_k}$ 缩小到约 $1/d_k$，因此容量增长至 $M \approx 2d_k^2/3 \approx 2700$，约为线性存储的 40 倍。

• SwiGLU MLP，其门控共享输入投影的权重（如平方 ReLU 所示），$\sigma(s) = \operatorname{Swish}(s)\cdot s = s^2/(1 + e^{-s})$：其容量几乎与平方 ReLU 相匹配。

• 立方 ReLU：每个干扰项缩小到约 $d_k^{-3/2}$，因此容量增长至 $M \approx 2d_k^3/15 \approx 35000$。

自注意力作为非参数记忆存储 #

将权重提升到某个幂次可使容量达到约 $d_k^p$。但如果能将这一方法推向极限会怎样？与其选择一个幂次，不如使用无限多个幂次。实际上，指数函数是所有幂次的加权和，

$$e^{\beta s} = 1 + \beta s + \frac{(\beta s)^2}{2} + \frac{(\beta s)^3}{6} + \cdots = \sum_{p=0}^{\infty} \frac{(\beta s)^p}{p!},$$

with the sharpness $\beta$ setting the weight on the higher powers. Here is the same experiment with $\beta = \sqrt{d_k}$, the value self-attention uses, 6 for the exponential and its series truncated after 2, 3 and 4 terms:

使用 2 项时，容量约为 40，低于线性存储器，因为常数项使每个键的权重约为 1。使用 3 项时约为 450，使用 4 项时约为 3,100，接近平方 ReLU 的容量。完整指数曲线远低于虚线，因此其容量已超出图表范围。

要确定指数的容量，我们用存储的键 $k_j$ 进行查询，因此目标内存的相似度为 1，权重 $w_j = e^{\beta}$。其他所有内存的权重 $w_i = e^{\beta s_i}$，其中对于随机单位键，相似度 $s_i$ 大致呈高斯分布，均值为 0，方差为 $1/d_k$。对于长度相近的随机值，

$$ \begin{aligned} \frac{\norm{\text{interference}}}{\norm{\text{signal}}} &\approx \sqrt{\sum_{i\neq j} \parens*{\frac{w_i}{w_j}}^2} && \text{interference terms add up like a random walk} \\ &= \sqrt{\sum_{i\neq j} e^{2\beta (s_i - 1)}} && w_i / w_j = e^{\beta s_i} / e^{\beta} \\ &\approx \sqrt{M\, e^{-2\beta}\, \ev{e^{2\beta s}}} && \text{average over the } M \text{ other memories} \\ &= \sqrt{M\, e^{-2\beta}\, e^{2\beta^2/d_k}} && \ev{e^{ts}} = e^{t^2/(2d_k)} \text{ for } s \sim \mathcal{N}(0, 1/d_k) \\ &= \sqrt{M}\, e^{\beta^2/d_k - \beta}. \end{aligned} $$

将比率设为 1 时，容量为

$$M \approx e^{2\beta(1 - \beta/d_k)},$$

只要 $\beta$ 远低于 $d_k$，容量就随 $\beta$ 指数增长，而不是像之前那样随 $d_k$ 多项式增长！当 $\beta = \sqrt{d_k} = 8$ 时，容量约为 $e^{14} \approx 1.2$ 百万条记忆，几乎是线性存储器的 2 万倍。对同一实验进行更长时间的运行，容量也约为 1.2 百万条。

在实践中，我们可能不希望在实现中直接对数字进行指数运算，因为小数字可能会下溢，大数字可能会上溢。幸运的是，如上链所示，干扰仅取决于权重之间的比率，因此我们可以将每个权重除以它们的总和，而不改变干扰与信号的比率，同时确保数值稳定性：

$$f(x) = \sum_{i=1}^{M} v_i\,\frac{\exp(\beta\,k_i\transpose x)}{\sum_j \exp(\beta\,k_j\transpose x)} = V\transpose \operatorname{softmax}(\beta K x).$$

这就是自注意力。7

与 MLP 不同，MLP 的键和值是训练的权重，而自注意力从其上下文中的标记计算键和值。每个标记存储一个记忆，因此记忆数量 $M$ 是上下文长度。这就是自注意力非参数化的原因：它的记忆不是参数，存储会随着上下文增长。

这是再次显示的条形图，从 softmax 开始。在 $\beta = 8$ 时，滑块范围内的回忆从未失败，因此将 $\beta$ 降低到 2 并增加 $M$：回忆在几十个记忆后失败。每次 $M$ 翻倍都需要稍微增加 $\beta$ 来保持回忆功能，这就是下一节的思路。

处理长上下文中的许多键值对 #

由于 $\beta$ 决定了我们抑制相似度较低键的强度，因此随着键值对数量的增加，我们可以相应地提高 $\beta$。为什么不直接让 $\beta$ 非常大呢？当 $\beta \to \infty$ 时，softmax 会退化为一个硬 argmax，只选择最相似的键，这不具备可微性，而且非常大的 $\beta$ 本身就会导致训练效果变差，因为 softmax 会饱和，其他所有键的梯度都会消失。因此我们只将 $\beta$ 提高到 $M$ 所需的程度。用 $\log M$ 对注意力对数进行缩放是处理长上下文的一种技巧，它能实现这一目标。这种方法有多种版本，每种版本都相对于一个固定的参考长度 $L$（通常是训练长度）设置 $\beta$，以确保当 $M = L$ 时 $\beta$ 保持其训练后的值：

• Su (2021) 从保持注意力熵恒定的角度推导出 $\log M$ 的缩放因子，通过标准化使 $M = L$ 时与标准注意力对齐。在每个长度上，对数都会乘以 $\log M / \log L$。Scalable-Softmax 的形式相同，但使用 $s \log M$ 进行缩放，其中 $s$ 是可学习的参数。

• 原始 Qwen 模型使用 LogN-Scaling，仅在上下文长度超过 $L$ 时应用相同的缩放因子。

• Qwen2.5-1M 使用 YaRN 的注意力缩放，将对数乘以 $(0.1 \log(M/L) + 1)^2$，因此其增长速度与 $(\log M)^2$ 相当。

在我们的实验中，每种规则的表现如下，当 $L = 1024$ 个记忆时 $\beta = \sqrt{d_k} = 8$：

Qwen 的规则和 Su 的规则在参考长度 $L$ 之后会重合，此时随着存储量的增加，干扰与信号的比率会下降，到 200 万个记忆时约为 0.004。YaRN 更快的增长速度将比率进一步降低到约 $10^{-4}$。在每种情况下，通过 $\log M$ 缩放 $\beta$ 都允许存储几乎无限数量的具有不同键的记忆。

结论 #

我们访问过的记忆存储及其容量（当使用存储的键进行查询且键和值为随机时）如下：

| 记忆存储 | $f(x)$ | 容量 | $d_k = 64$ 时的容量 | |------------------|-------------------------------------------|-----------------|----------------------| | 外积 | $V\transpose K x$ | $d_k$ | 64 | | 最优线性 | $V\transpose K (K\transpose K)\inverse x$ | $2d_k$ | 128 | | ReLU MLP | $V\transpose \operatorname{ReLU}(Kx)$ | 平方 ReLU (GLU MLP) | $V\transpose \operatorname{ReLU}(Kx)^2$ | $2d_k^2/3$ | 约 2,700 | | SwiGLU MLP | $V\transpose \parens*{\operatorname{Swish}(Kx) \odot Kx}$ | 约 $2d_k^2/3$ | 约 2,900 | | 立方 ReLU | $V\transpose \operatorname{ReLU}(Kx)^3$ | $2d_k^3/15$ | 约 35,000 | | 自注意力，$\beta = \sqrt{d_k}$ | $V\transpose \operatorname{softmax}(\beta K x)$ | $e^{2\sqrt{d_k} - 2}$ | 约 120 万 | | 自注意力，$\beta \propto \log M$ | 未在实验中达到 | 超过 200 万 |

我们从启发式记忆存储——外积记忆开始，逐步构建到现代变压器中所有标准层。这种视角也让我们能够估计它们在容量和作为向量值键值存储能力方面的差异。

我需要指出的是，这种容量估算方法只是观察现代神经网络架构的一种方式。我们的容量估算仅计算单层在干扰主导信号之前能回忆的随机向量数量。实际上，我们根本没有考虑训练数据，而训练数据会影响键和值中实际存储的内容！还有其他研究考虑了训练后的层，通过比特数衡量存储的信息。Morris 等人估计，GPT 风格的模型每参数大约记忆 3.6 比特，并发现当模型填满其容量后，开始泛化并记忆每个训练示例的更少内容。Allen-Zhu & Li 发现，即使量化到 int8，语言模型每参数存储约 2 比特的事实知识。

从这里出发，还有许多相关方向同样引人入胜：

• 一层能存储多少信息：Nichani, Lee & Bietti、Barnfield 等人针对线性存储，Dugan 等人、Garcia 等人针对显式事实存储的 MLP，以及 Zhong 等人启发了此处的干扰与信号比估算。

• 具有指数级容量的记忆存储：Krotov & Hopfield 的密集联想记忆，Demircigil 等人、Ramsauer 等人将其与注意力联系起来，以及 Krotov 等人的综述。

• 推理过程中更新的短期记忆：fast weights、Mamba、Mamba-2、DeltaNet、Gated DeltaNet、TTT layers、Titans、MesaNet 和 test-time regression。

• 用于补充或替代我们 MLP 的巨大但稀疏的参数化记忆存储：product-key memories、Memory Layers at Scale、UltraMem、Mixture of A Million Experts 和 Engram。

• 通过在权重之外追加记忆来扩展的非参数化记忆存储：kNN-LM、RETRO、Memorizing Transformers 和 growing nonparametric networks。

• 在训练模型中查找和编辑记忆：Geva 等人、Bietti 等人、ROME 和 MEMIT。

• 压缩短期记忆或将其转化为长期记忆：KVzip、gist tokens、Cartridges 和 attention matching。

• 类似稀疏记忆微调的持续学习方法。

我一直对记忆着迷，也许是因为我本能地感受到自己记忆的局限性。也许我内心的一部分希望，通过理解机器如何记忆，我也会对自身有所了解。

• 有时这种键和值的积累被称为“赫布学习”，因为突触（权重）的强度会随着同时激活（相关）的输入（键）和输出（值）而增强。↩︎

• 这种外积记忆存储用于线性注意力的状态，也称为 fast weights：每个标记通过将 $v k\transpose$ 添加到一个矩阵中来存储记忆 $(k, v)$，后续标记会查询这个矩阵（Katharopoulos 等人；Schlag, Irie & Schmidhuber）。↩︎

• 这是用于回忆干净的值向量。如果每个值只需要在固定集合中指向正确答案，例如选择最可能的标记，那么回忆只需更接近正确答案而不是错误答案，一个 $d\times d$ 矩阵可以存储约 $d^2$ 个事实（Nichani, Lee & Bietti）。↩︎

• 我们也可以使用这种回归目标来设计序列模型在前向传递过程中书写的短期记忆，例如线性注意力和自注意力。参见我关于 test-time regression 的工作。↩︎

• 对于随机键，ReLU 会将干扰降低 $\sqrt{2}$ 倍。随机键的相似度等概率为正或负，因此 ReLU 会使约一半权重归零，另一半保持不变。这使干扰的方差（即权重平方和）减半，从而让干扰本身缩小 $\sqrt{2}$ 倍。完整推导请参见 Zhong 等人的研究。 ↩︎

• Transformer 模型将相似度除以 $\sqrt{d_k}$ 而非乘以。但其查询和键的长度约为 $\sqrt{d_k}$，因此匹配对的 $q\transpose k \approx d_k$。我们的键和查询长度为 1，匹配对的相似度为 1。因此我们通过 $\beta = \sqrt{d_k}$ 进行缩放。 ↩︎

• 在 Transformer 表示法中，相似度为 $k_i\transpose q$（其中 $q = W_Q x$），并乘以 $1/\sqrt{d_k}$。我将 $W_Q$ 合并到键中。Transformer 的键和查询长度约为 $\sqrt{d_k}$，因此这种缩放方式与单位长度键（如我们的设计）的 $\beta = \sqrt{d_k}$ 相匹配。 ↩︎

引用本文

复制

@article{wang2026associativememory,
  author = {Alex Wang},
  title = {You could have invented transformers from memory},
  journal = {alexwang.ai},
  year = {2026},
  month = {Oct},
  url = {https://alexwang.ai/posts/associative-memory/}
}

你也可能喜欢

• 2026年10月9日 矩阵梯度，无需痛苦

• 2026年4月28日 为突破性进展定价

• 2026年2月21日 模型即公司