从记忆存储机制推导 Transformer 各核心层

如今的大型语言模型能够回忆起世界各个角落的冷门知识。最近，我因为参加朋友的桌游之夜迟到了，所以我让 Astra 提醒我《Dune: Imperium》的规则。它在我开车前往的路上欣然为我复述，帮助我记起足够清晰的规则以便继续游玩。我们已经习惯向模型询问任何事情，以至于把它们广博的知识视为理所当然。

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

在本文中，我将介绍一种直观的方式，把线性映射、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$。¹

要存储一条新的记忆 $(k, v)$，我们只需将 $vk\transpose$ 加到现有权重上即可。²

旁注：一次梯度下降步骤即写入一条记忆

考虑一个网络中具有标量损失 $L$ 的线性层 $y = Wx$，并设 $\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,$$

也就是说，一旦存储的记忆数量与键的维度相当，干扰就会与信号一样大，从而将信号淹没。

下图中的每根柱表示在一次回忆中某个存储记忆的权重相对于目标记忆权重（红色）的大小。每个柱都会将其值的一部分累加到回忆中。把 $M$ 拖到 $d_k = 64$ 之后，就能看到这些小柱逐渐堆积，直到回忆失败。其他按钮可以切换到本文后面将介绍的其他记忆存储方式。

如何减少干扰？#

即使在其约 $d_k$ 个记忆的容量以下，外积存储在回忆时仍会产生显著的干扰。³ 为了提高容量，或者在容量以下更干净地回忆，我们需要减少干扰。

一种粗暴的方法是把键的维度 $d_k$ 取得更大，因为线性存储的容量随 $d_k$ 增长。而 $d_k$ 能取多大，则受限于硬件条件。

另外，我们可以在算法上做改进。我们可以更仔细地存储记忆：外积存储在添加每条记忆时不考虑其他记忆，因此它对抵消它们之间的干扰没有任何作用。我们也可以更具选择性地进行读取：线性存储中的每条记忆都会产生干扰，即使其键与查询的相似度较低也是如此，因此我们可以通过抑制相似度较低的键来改进读取。

在本文余下的部分，我们将看到，更仔细地存储外积记忆可以给我们最优的线性记忆存储，其容量大约是原来的两倍。通过非线性进行选择性读取则可以得到 MLP 和自注意力。

最优的线性记忆存储#

外积存储在写入每条记忆时不考虑其他记忆。相反，我们可以选取权重 $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$ 个记忆的干扰。

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

外积存储在 $M \approx d_k = 64$ 时达到其容量。最优存储在 $M \leq 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. 指出了前馈层与记忆存储之间的这种相似性，Geva et al. 则发现训练好的 transformer 中的前馈层经常以这种方式工作：键会在输入中可识别的模式上激活。

理想情况下，门控函数 $\sigma$ 的选择应使得当相似度较低时（$k_i\transpose x \sim 0$，甚至当 $k_i\transpose x < 0$ 时），它在求和中对应的项被抑制，从而减少干扰。ReLU 激活会把相似度为负的键归零，添加偏置则可以移动该阈值⁵。平方 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{干扰}}}{\norm{\text{信号}}} \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!},$$

其中尖锐度 $\beta$ 设置了较高幂次的权重。下面是用 $\beta = \sqrt{d_k}$（自注意力所使用的值）⁶对指数以及它在 2、3、4 项处截断的级数进行的相同实验：

使用 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 120$ 万个记忆，几乎是线性存储的 20,000 倍。同一实验的较长时间运行也将容量定为约 120 万。

在实践中，我们可能不希望在实现中天真地对数取指数，因为小数可能会下溢，大数可能会溢出。幸运的是，如上面的推导所示，干扰仅取决于权重之间的比值，因此我们可以将每个权重除以其总和，而不改变干扰信号比，同时保证数值稳定性：

$$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).$$

这就是自注意力。⁷

与 MLP 不同（MLP 的键和值都是训练得到的权重），自注意力从其上下文中的 token 计算出自己的键和值。每个 token 存储一份记忆，因此记忆的数量 $M$ 等于上下文长度。这正是自注意力非参数化的原因：它的记忆不是参数，而且存储会随上下文增长。

再看一下柱状图，从 softmax 开始。在 $\beta = 8$ 时，在滑动条范围内召回从不失败，因此把 $\beta$ 降到 2 并提高 $M$：召回在几十份记忆之后失效。每把 $M$ 翻倍都需要再增大一点 $\beta$ 才能让召回继续工作，这就是下一节要讨论的核心思想。

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

由于 $\beta$ 决定了我们对不太相似的键的抑制力度，我们可以在键值对数量增加时同步增大 $\beta$。为什么不直接把 $\beta$ 设得很大？当 $\beta \to \infty$ 时，softmax 就变成一个硬性的 argmax，只挑出最相似的键，而这是不可微的；并且非常大的 $\beta$ 已经会让训练效果变差，因为 softmax 出现饱和，传递到其它每个键的梯度都会消失。所以我们只按 $M$ 的需求来提高 $\beta$。用 $\log M$ 对注意力 logits 做缩放，就是一种面向长上下文的、能够做到这一点的小技巧。它有几种变体，每种都把 $\beta$ 相对于一个固定的参考长度 $L$（通常是训练长度）来设定，从而在 $M = L$ 时保持训练时用到的 $\beta$ 值：

• Su (2021) 通过让注意力熵在 $M$ 增长时保持不变，推导出 $\log M$ 的缩放，并做了归一化，使得在 $M = L$ 时与标准注意力一致。在任何长度下，logits 都乘以 $\log M / \log L$。Scalable-Softmax 形式相同，把 logits 乘以 $s \log M$，其中 $s$ 由学习得到。
• 最初的 Qwen 模型使用的是 LogN-Scaling，它只在上下文长度超过 $L$ 时才施加相同的缩放因子。
• Qwen2.5-1M 使用的是 YaRN 的注意力缩放，它把 logits 乘以 $(0.1 \log(M/L) + 1)^2$，因此其增长与 $(\log M)^2$ 同阶。

下面给出在我们的实验中每条规则的表现，其中 $L = 1024$ 份记忆时 $\beta = \sqrt{d_k} = 8$：

在超过参考长度 $L$ 之后，Qwen 的规则与 Su 的规则重合，此时干扰-信号比会随着存储增大而下降，到 2 百万份记忆时降到约 0.004。YaRN 的缩放增长更快，把该比值压到约 $10^{-4}$。 在每种情况下，把 $\beta$ 按 $\log M$ 缩放，原则上就允许存储数量本质上不受限制、键互不相同的记忆。

结语#

下表汇总了我们讨论过的各种记忆存储，列出了当用存储的键进行查询、且键与值都是随机时它们的容量：

我们从一个启发式的记忆存储——外积记忆——出发，逐步构建出如今现代 transformer 中的所有标准层。 这一视角也让我们能够估计它们在容量上的差异，也就是它们作为向量值键值存储的能力。

我应该说明的是，这种容量估计只是观察现代神经网络架构的一种方式。我们的容量估计只计算在干扰主导信号之前，单个层能够回忆多少个随机向量。事实上我们根本没有考虑训练数据，而训练数据会影响键和值中实际被表示的内容！还有一些工作考虑了训练后的层，以比特为单位衡量存储的信息量。Morris 等人估计 GPT 类模型每个参数大约记忆 3.6 比特，并发现一旦模型填满了它的容量，它就会开始泛化，对每个训练样本的记忆会更少。Allen-Zhu 和 Li发现语言模型每个参数存储大约 2 比特的事实性知识，即使量化到 int8 也是如此。

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

• 单个层能存储多少信息：Nichani、Lee 和 Bietti，Barnfield 等人针对线性存储，以及Dugan 等人，Garcia 等人针对显式事实存储的 MLP，还有Zhong 等人，他们的工作启发了本文中的干扰-信号比估计。
• Hopfield 网络和具有指数容量的记忆存储：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
• 持续学习方法，例如sparse memory finetuning

我一直被记忆所吸引，也许是因为我深切地感受到自己的记忆有多么有限。也许我的一部分希望，通过理解机器如何记忆，我也能对自己有所了解。

1. 有时这种键与值的累积被称为"赫布学习"，因为突触（权重）的强度是由同时激活（具有相关性）的输入（键）和输出（值）来增强的。 ↩︎

2. 这种外积记忆存储用于线性注意力的状态，也称为快速权重：每个标记通过将 $v k\transpose$ 添加到一个矩阵来存储记忆 $(k, v)$，供后续标记进行查询（Katharopoulos 等人；Schlag、Irie 与 Schmidhuber）。 ↩︎

3. 这是为了回忆出一个干净的数值向量。如果每个值只需在固定集合中指向正确答案，例如选出最可能的标记，那么回忆只需比指向错误答案更接近正确答案即可，而一个 $d\times d$ 矩阵可容纳约 $d^2$ 条事实（Nichani、Lee 与 Bietti）。 ↩︎

4. 我们同样可以利用这个回归目标，来设计序列模型在前向传播过程中写入的短期记忆，例如线性注意力和自注意力。参见我关于测试时回归的工作。 ↩︎

5. 对于随机键，ReLU 将干扰降低 $\sqrt 2$ 倍。随机键的相似度为正或为负的概率相同，因此 ReLU 将约一半的权重置零，另一半保持不变。这会把干扰的方差——即权重平方之和——减半，因此干扰本身缩小 $\sqrt{2}$ 倍。完整推导参见 Zhong 等人。 ↩︎

6. Transformer 将相似度除以 $\sqrt{d_k}$ 而不是相乘。但它们的查询和键长度约为 $\sqrt{d_k}$，所以匹配对有 $q\transpose k \approx d_k$。我们的键和查询长度为 1，因此匹配对的相似度为 1。所以我们按 $\beta = \sqrt{d_k}$ 进行缩放。 ↩︎

7. 在 Transformer 符号中，相似度为 $k_i\transpose q$，其中 $q = W_Q x$，按 $1/\sqrt{d_k}$ 缩放。我将 $W_Q$ 折入到键中。Transformer 中键和查询的长度约为 $\sqrt{d_k}$，因此该缩放对于像我们这样的单位长度键而言，对应 $\beta = \sqrt{d_k}$。↩︎