从记忆视角，你也能发明 Transformer
转自 You could have invented transformers from memory | Alex Wang
记忆召回动态演示：汉字“記”的检索
如今的大语言模型能够回忆起世界上各个角落生僻冷门的知识。就在最近，我赶着去参加朋友的桌游之夜，路上快迟到了，于是我让 Astra 提醒我《沙丘：帝国》（Dune: Imperium）的游戏规则。它在我开车时一路娓娓道来，帮我把规则记得一清二楚，让我能立刻上手再玩一局。我们早已习惯了向大模型询问任何问题，以至于把它们海量的知识储备视作理所当然。
模型究竟是如何记住如此庞大的信息的？而且，它们到底能存储多少记忆？
在这篇博文中，我将介绍一种非常直观的视角，将线性映射（Linear Map）、多层感知机（MLP）以及自注意力机制（Self-Attention）分别视为不同类型的“记忆存储器”（Memory Store）。基于这种直观理解，我们能够推导出各个网络层究竟能存储多少记忆的启发式容量公式。
衡量一个“良好”记忆存储器的核心指标，是它能否准确召回已存储的内容。毕竟，即使能存储无限多的事实，如果我们无法在需要时准确召回它们，也是徒劳无功的。我们存储的记忆越多，在尝试召回某一特定记忆时，各个记忆之间的干扰（Interference）就越严重。因此，在本文中我们将始终聚焦于一个关键量——干扰信号比（Interference-to-Signal Ratio），它衡量了在检索召回过程中其他干扰记忆对目标信号的阻碍程度。一个记忆存储器在干扰增长到压倒信号之前所能容纳的记忆数量上限，就是它的容量（Capacity）。
我们将从这一视角出发，展示如何从最基础的线性层逐步迈向更优的结构：先是迈向能够同时提升容量和召回质量的 MLP，再一路迈向拥有极佳召回能力且理论上能容纳无界数量记忆的自注意力机制。在 64 维键（Key）的设置下，这意味着容量从约 64 条记忆跃升至约 2,700 条，最终突破百万级条记忆。
以下是贯穿全文的核心术语定义：
- 记忆存储器（Memory Store）：一个能够保存 $M$ 条记忆（编号 $i = 1, \ldots, M$）的函数 $f$。
- 记忆（Memory）：一个键值对（Key-Value Pair），包含用于基于相似度进行检索匹配的键（Key） $k_i \in \mathbb{R} ^ {d_k}$，以及我们希望检索出的值（Value） $v_i \in \mathbb{R} ^ {d_v}$。
- 查询与召回：我们用输入 $x$ 对函数 $f$ 进行查询（Query），它会召回（Recall）输出 $f(x)$。
一个优秀的记忆存储器，在我们用 $k_i$（或与 $k_i$ 高度相似的向量）进行查询时，能够准确召回对应的 $v_i$。在文献中，这种按相似度进行召回的记忆机制被称为联想记忆（Associative Memory）。
外积记忆存储器
最基础的记忆存储器是线性的：它仅仅通过外积将记忆不断累加起来：
$$W = \sum_ { i=1} ^ {M} v_i k_i^\top = V^\top K$$
其中矩阵 $K \in \mathbb{R} ^ {M\times d_k}$ 和 $V \in \mathbb{R} ^ {M\times d_v}$ 的行分别由键的转置 $k_i^\top$ 与值的转置 $v_i^\top$ 构成。[1]
若要存入一条新记忆 $(k, v)$，我们只需在现有权重矩阵上直接加上外积项 $vk^\top$ 即可。[2]
旁白：一次梯度下降步即是在写入一条记忆
考虑神经网络内部一个输出为 $y = Wx$ 的线性层，设标量损失函数为 $L$，令 $\bar{y} = \nabla L(y)$ 为损失对该层输出的梯度。那么损失对权重矩阵 $W$ 的梯度为 $\bar{y} x^\top$。因此，学习率为 $\eta$ 的一步更新即为：
$$W \leftarrow W - \eta, \bar{y} x^\top.$$
这一递推更新公式本质上就是一个记忆写入操作：它用键 $x$ 和值 $-\eta\bar{y}$ 向权重矩阵中写入了一条新记忆。
关于此视角的深入探讨，可参见 Irie, Csordás & Schmidhuber，他们将通过梯度下降训练得到的线性层等价表述为在其所有训练输入上的注意力机制。
若要使用输入向量 $x$ 对该记忆存储器进行查询，只需将输入与该矩阵相乘：
$$f(x) = V^\top K x = \sum_ { i=1} ^ {M} v_i,(k_i^\top x).$$
由此可见，召回操作本质上是一次基于与存储键相似度的检索：我们将查询与每一个键进行比对计算内积，然后返回以这些相似度为权重的所有存储值的加权和。
召回过程中的干扰
既然记忆的存储如此简单——只需将外积持续累加到权重中即可，那么在召回时又会发生什么呢？
为简便起见，假设每个键都已归一化为单位长度，即 $k_j^\top k_j = 1$。此时我们用存储键 $k_j$ 进行查询：
$$f(k_j) = \underbrace{v_j}_ { \text{信号}} + \underbrace{\sum_ { i\neq j} v_i,(k_i^\top k_j)}_ { \text{干扰}},$$
这表明我们最终得到的输出，等于我们查询所针对的**目标记忆（Target Memory）**的真实值（即信号），加上来自所有其他与查询相似的键所对应的干扰项之和。
如果所有键彼此正交归一（Orthonormal），则干扰项完全为零，召回将百分之百精确。然而，在 $d_k$ 维实数空间 $\mathbb{R} ^ {d_k}$ 中最多只能存在 $d_k$ 个相互正交的方向。因此，一旦存储的记忆数量超过键的维度（即 $M > d_k$），我们便再也无法获得完美无扰的召回。
那么外积记忆存储器的容量究竟是多少？这里有一个直观的推导论证。式中共有 $M-1$ 个干扰项。若各个键为随机单位向量，则每个相似度内积 $k_i^\top k_j$ 的期望大小约为 $\pm 1/\sqrt{d_k}$。若各个值向量同样是随机且模长相近的，这 $M-1$ 个干扰项在累加时就像随机游走（Random Walk）一样展开，因此干扰项模长的平方大约为信号模长平方的 $(M-1)/d_k$ 倍，由此得出：
$$\frac{\left|\text{干扰}\right|}{\left|\text{信号}\right|} \sim \sqrt{\frac{M}{d_k}}.$$
将该干扰信号比设为 1，便能得出其临界容量为：
$$M \approx d_k,$$
也就是说，一旦存储的记忆数量达到键的维度大小，干扰便会与信号相当，从而将目标信号淹没。
下方的条形图展示了单次召回中，每个已存储记忆相对于目标记忆权重（红色）的贡献大小。每个小条都会向召回结果中掺入一部分它自身的值。如果将 $M$ 拖动超过 $d_k = 64$，可以看到零碎的小条不断累加堆积，直到召回彻底失效。切换按钮可体验我们在后文中构建的高级记忆存储器。
交互式记忆权重柱状图演示
如何减少干扰？
即便是在容量阈值（约 $d_k$ 条记忆）以下，外积存储器在召回时依然会产生不可忽视的显著干扰。[3] 为了提升容量上限，或是在容量范围内获得更纯净的召回结果，我们必须减少干扰。
最简单粗暴的方法是增大键维度 $d_k$，因为线性存储器的容量随着 $d_k$ 线性增长。但 $d_k$ 能做多大受制于硬件算力与显存。
另一种思路是进行算法层面的改良。我们可以在存储时更加精细：外积存储器写入每条记忆时彼此完全独立、互不关照，因此根本没有做任何消除彼此干扰的处理。或者，我们可以在召回时更加挑剔：线性存储器中的每一条记忆都会贡献干扰，即便它的键与当前查询的相似度极低；因此，我们可以通过抑制低相似度的键来显著提升召回的信噪比。
在接下来的内容中，我们将看到：更精细地存储外积记忆，将引领我们得到最优线性记忆存储器，其容量翻倍；而利用非线性函数进行选择性召回，则自然而然地推导出了 MLP 与自注意力机制。
最优线性记忆存储器
外积存储器在写入每条记忆时毫不顾及其他记忆的存在。与此相反，我们可以精心选取权重矩阵 $W$，使得召回过程同时最佳拟合所有已存储的记忆。这正是一个线性回归（Linear Regression）问题：[4]
$$\operatorname{minimize} _ W \sum _ { i=1}^M \left|W k_i - v_i\right| _ 2 ^ 2$$
当存储的记忆数量多于键维度（即 $M > d_k$）时，最小二乘解为：
$$W = \left(\sum_i v_i k_i^\top\right)\left(\sum_i k_i k_i^\top\right) ^ {-1} = V^\top K,(K^\top K) ^ {-1}.$$
这正是原始外积存储器 $V^\top K$ 加了一处精妙改变：在将查询与各个键进行比对之前，先乘上了逆矩阵 $(K^\top K) ^ {-1}$。这一变换消除了键之间的协方差 $K^\top K$，从而实现了对键的白化（Whitening），使存储器不再过度偏向于多个键所共享的共有方向。
当 $M \leq d_k$ 时，只要这些键线性无关，该回归问题就能精确无误地拟合每一条记忆。因此，最优存储器在召回时完全没有任何干扰，即使各个键并不正交也是如此。
而当 $M > d_k$ 时，由于整个记忆存储器仅有 $d_k$ 个可支配的独立方向，干扰便不可避免地开始出现。对于随机的键和值，干扰信号比变为：
$$\frac{\left|\text{干扰}\right|}{\left|\text{信号}\right|} \approx \sqrt{\frac{M}{d_k} - 1}.$$
该比值在下式处达到 1：
$$M = 2d_k,$$
这意味着最优线性存储器的容量足足是之前的两倍！
旁白：最优线性记忆存储器的数学推导
当 $M \leq d_k$ 时：若各个键线性无关，则键之间的内积矩阵（$M\times M$ 维）$KK^\top$ 可逆，此时：
$$W = V^\top (KK^\top) ^ {-1} K$$
能够完美拟合每一条记忆。当 $M < d_k$ 时，该回归问题是欠定（Underconstrained）的，存在无穷多个解，而上述解正是最小范数解（Min-norm solution）。矩阵 $(KK^\top) ^ {-1} K$ 的每一行 $\tilde{k}_i^\top$ 被称为 $K$ 各行的对偶基（Dual Basis）：根据构造，$(KK ^ \top) ^ {-1} KK ^ \top = I$，因此 $\tilde{k} _ i ^ \top k _ j = \delta _ { ij}$。由此得出 $f(k _ j) = \sum _ i v _ i, \tilde{k} _ i ^ \top k _ j = v _ j$，这意味着完全没有干扰，我们获得了完美无瑕的召回。
当 $M > d_k$ 时：令目标函数的梯度为零，可得 $W K^\top K = V^\top K$，故 $W = V^\top K,(K^\top K) ^ {-1}$。
两种情况下的解都可以统一写作 $W = V^\top (K^\top)^\dagger$，其中 $A^\dagger$ 表示矩阵 $A$ 的 Moore-Penrose 伪逆（Pseudoinverse）：当键线性无关时 $(K^\top)^\dagger = (KK^\top) ^ {-1} K$，而当键张成整个 $\mathbb{R} ^ {d_k}$ 时 $(K^\top)^\dagger = K(K^\top K) ^ {-1}$。现在我们可以分析 $f(k_j) = V^\top K,(K^\top K) ^ {-1} k_j$ 的干扰项。
对于已存入的键 $k_j$，分配给每个值 $v_i$ 的权重为 $k_i^\top (K^\top K) ^ {-1} k_j$，这恰好对应 $M\times M$ 投影矩阵 $P = K(K^\top K) ^ {-1} K^\top$ 的第 $(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$ 的列空间（即 $\mathbb{R} ^ {M}$ 的一个 $d_k$ 维子空间）的正交投影矩阵。投影矩阵的两条性质有助于我们解析干扰信号比：
信号项：利用矩阵迹（Trace）的循环性质：
$$\sum_ { j=1}^M P_ { jj} = \operatorname{tr}\left(K(K^\top K) ^ {-1} K^\top\right) = \operatorname{tr}\left((K^\top K) ^ {-1} K^\top K\right) = \operatorname{tr}\left(I_ { d_k}\right) = d_k,$$
这表明所有 $M$ 个信号权重的总和恰好等于 $d_k$。如果键是随机选取的，没有任何一个键处于特殊地位，因此对称地有 $P_ { jj} \approx d_k/M$。
干扰项：由于 $P$ 是对称投影矩阵，因此满足 $P^\top P = P^2 = P$。由此可得：
$$\sum_ { i\neq j} P_ { ij}^2 = \sum_ { i=1}^M P_ { ij}^2 - P_ { jj}^2 = (P^\top P)_ { jj} - P_ { jj}^2 = P_ { jj} - P_ { jj}^2 \approx \frac{d_k}{M} - \left(\frac{d_k}{M}\right)^2.$$
若各值向量随机独立且模长相近，则：
$$\frac{\left|\text{干扰}\right|}{\left|\text{信号}\right|} \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$ 维的单位向量，用其中一个已存储的键去查询存储器，并测算其干扰信号比：
图：在 $d_k = 64$ 下两种线性存储器的干扰信号比随 $M$ 变化曲线。当查询等于已存键时，外积记忆在 $M = 64$ 处干扰达到信号大小；而最优最小二乘记忆在 $M \leq 64$ 时完全无干扰，并在 $M \approx 128$ 处干扰达到信号大小。
外积存储器在 $M \approx d_k = 64$ 处达到其容量极限。而最优存储器在 $M \leq d_k$ 时召回完全精确，之后干扰从零开始按 $\sqrt{M/d_k - 1}$ 上升，并正如理论预测的那样在 $M = 2d_k = 128$ 附近达到容量极限。由于最优存储器抵消掉的干扰量是固定数量（相当于抵消了 $d_k$ 条记忆的干扰），随着 $M$ 的不断增大，两者的相对差距会逐步缩小。
然而，这种最优存储器对于相似度而言依然是线性的。它对无论高低的所有相似度都必须以相同的尺度进行缩放。若要取得突破，我们必须从根本上改变存储器的召回机制。
作为非线性记忆存储器的 MLP
线性记忆存储器的召回形式为：
$$f(x) = \sum_ { i=1}^M v_i (k_i^\top x).$$
如果我们有一种机制，能在内积相似度 $k_i^\top x$ 较小时强力抑制对应的权重，那么就能大幅削减低相似度键所带来的累积干扰。
为此，我们在计算值的加权和之前，让每个相似度先穿过一个非线性函数 $\sigma$：
$$f(x) = \sum_ { i=1} ^ {M} v_i,\sigma(k_i^\top x) = V^\top \sigma(Kx).$$
我们竟然直接从第一性原理“重新发明”了双层多层感知机（Two-layer MLP）！第一层权重矩阵的各行正是键（Keys），第二层权重矩阵的各列正是值（Values）。Sukhbaatar 等人首先指出了前馈层（Feedforward Layers）与记忆存储器之间的这种相似性，而 Geva 等人发现，训练好的 Transformer 中的前馈层往往就是以这种方式运作的，其键向量会针对输入中可识别的特定模式精确激活。
在理想状态下，门控激活函数 $\sigma$ 的设计应满足：当相似度较低时（$k_i^\top x \sim 0$ 甚至 $k_i^\top x < 0$），该项在求和中的贡献被极大抑制，以此降低干扰。ReLU 激活函数能够直接将负相似度的键归零，引入偏置项（Bias）还能平移这个截断阈值。[5] 而平方 ReLU（Squared 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 是以**叠加状态（Superposition）**来存储知识事实的：海量的事实共享相同的参数权重，输出结果是众多记忆的混合表现。Garcia 等人显式构建了这样一种网络：GLU MLP 的隐层单元计算输入的随机特征 $\phi(x)$，第二层权重矩阵是这些特征上的外积存储器 $\sum_i v_i \phi(k_i)^\top$，该 MLP 存储的事实数量远远超过其隐层单元数。即便如此，训练好的 Transformer 的 ReLU MLP 在面对特定输入时，往往只有不到 1% 的隐层单元被激活，因此每次召回实际上仅调用了极小一部分记忆（参见 Li 等人）。
旁白：Transformer 原作者们早已心知肚明
Transformer 开山之作其实早就指出了这一内在联系，但只出现于 arXiv 上的 v2 和 v3 版本中。在那里，Vaswani 等人明确将前馈层描述为“一种注意力机制”（a kind of attention），其两个权重矩阵的各行分别充当键和值，只是用 ReLU 取代了 Softmax。他们甚至尝试过将前馈层替换为对可学习键值参数的注意力机制，并取得了相当的翻译质量：一个版本略差，另一个版本略好，不过前者的训练速度慢了约 30%。这段论述在论文后续的所有版本中都被删去了。稍后我们将在下文深入讨论自注意力机制与这套理论的关联。
通过更高阶次减少干扰
对相似度取平方之所以奏效，是因为相比于高相似度权重，它对低相似度权重的压制剧烈得多。我们可以把这一思想推向极致：令 $\sigma(s) = \operatorname{ReLU}(s)^p$。此时目标信号保持不变（因为 $1^p = 1$），而相似度约为 $1/\sqrt{d_k}$ 的干扰项则被大幅压缩至 $d_k ^ {-p/2}$。这 $M$ 个干扰项以随机游走方式叠加，因此干扰信号比为：
$$\frac{\left|\text{干扰}\right|}{\left|\text{信号}\right|} \approx \sqrt{\frac{M}{d_k^p}}.$$
这意味着记忆容量约为：
$$M \approx d_k^p$$
条记忆，容量相比之前呈现出几何级数/多项式级的爆发式增长！
为了验证这一点，我们针对五种不同的 $\sigma$ 选择进行了相同的数值模拟实验：
图：在双对数坐标系下，干扰信号比随存储记忆数量 $M$ 的变化曲线。查询等于已存键时，线性记忆在 $M \approx 64$ 处达到容量极限，ReLU 记忆在约 128 处达到，平方 ReLU 和 SwiGLU 记忆在约 3,000 处达到，而三次方 ReLU 记忆在超过 16,384 处仍未失效。
由于干扰均以 $\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$。
作为非参数化记忆存储器的自注意力机制
将相似度权重提升到 $p$ 次幂，能够将容量提升至约 $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}$（这正是标准自注意力机制所采用的缩放值）[6] 时，完整指数函数与其截断为前 2、3、4 项泰勒展开后的表现对比：
图：锐度为 $\beta = \sqrt{d_k}$ 的指数函数及其截断为 2、3、4 项后的干扰信号比随 $M$ 变化曲线，附带线性记忆作为基准对比。
- 仅保留 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{\left|\text{干扰}\right|}{\left|\text{信号}\right|}
&\approx \sqrt{\sum_ { i\neq j} \left(\frac{w_i}{w_j}\right)^2} && \text{干扰项以随机游走方式累加} \
&= \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}, \mathbb{E}\left[e ^ {2\beta s}\right]} && \text{对其余 } M \text{ 条记忆求期望均值} \
&= \sqrt{M, e ^ {-2\beta}, e ^ {2\beta^2/d_k}} && s \sim \mathcal{N}(0, 1/d_k) \text{ 时，矩母函数 } \mathbb{E}\left[e ^ {ts}\right] = e ^ {t^2/(2d_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$ 呈现指数级增长，而不再是之前的多项式增长！当取标准设置 $\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^\top x)}{\sum_j \exp(\beta,k_j^\top x)} = V^\top \operatorname{softmax}(\beta K x).$$
这，正是自注意力机制（Self-Attention）。[7]
与 MLP（其键和值是固化在参数权重中的）不同，自注意力机制的键和值是从上下文中的各个 Token 动态计算而来的。每个 Token 存储了一条记忆，因此记忆条数 $M$ 就是上下文长度。这正是自注意力机制被称为“非参数化（Non-parametric）”的原因：它的记忆并不占据模型静态参数，存储器的大小随着上下文动态扩张。
再次回顾记忆条柱图，从 Softmax 开始观察。在 $\beta = 8$ 时，在滑块所能调节的范围内召回永远不会失败；但若将 $\beta$ 调低至 2 并增大 $M$，仅需几十条记忆召回就会失效。$M$ 每翻一番，都需要相应增大一点 $\beta$ 才能维持纯净召回，这正是下一小节的核心思想。
处理长上下文中的海量键值对
既然 $\beta$ 决定了对低相似度键的压制力度，那么随着键值对数量的膨胀，我们理应相应增大 $\beta$。那为什么不直接把 $\beta$ 设得极大呢？当 $\beta \to \infty$ 时，Softmax 会退化为不可微的硬性 $\operatorname{argmax}$，仅选取唯一的最高相似度键；而在训练过程中，过大的 $\beta$ 也会导致 Softmax 快速饱和，使得传递给其他所有键的梯度彻底消失，导致训练困难。因此，我们只需要随着 $M$ 的需求温和地调大 $\beta$。
在长上下文建模中，依据 $\log M$ 缩放注意力 Logits 正是实现这一目标的妙招。业界存在几种不同的缩放方案，它们通常相对于固定的参考长度 $L$（通常为预训练长度）来调整 $\beta$，确保在 $M = L$ 时 $\beta$ 恰好保持其训练设定值：
- Su (2021) / 追一科技苏剑林 基于在 $M$ 增长时维持注意力熵（Entropy）恒定的准则推导出了 $\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$ 的设置下，各规则在数值实验中的表现：
图：在 $d_k = 64$ 下，四种锐度规则对应的 Softmax 注意力干扰信号比随记忆数量 $M$（最高至 200 万对）的变化曲线。固定 $\beta = 8$ 时，干扰在约 120 万对处达到信号水平；而按对数缩放的规则下，干扰信号比随着 $M$ 的增长反而持续下降。
Qwen 的规则与苏剑林的规则在超出参考长度 $L$ 后完全重合，其干扰信号比随着存储容量的增加不升反降，在达到 200 万条记忆时甚至下降至约 0.004。而 YaRN 增速更快的缩放法则将该比值直接压低到了 $10 ^ {-4}$ 量级。
在这些方案中，让 $\beta$ 随 $\log M$ 动态扩展，使得自注意力机制理论上能够容纳近乎无界、无限多条拥有不同键的记忆！
总结与思考
以下是我们探讨过的所有记忆存储器汇总，以及在查询已存储键且键和值均为随机向量时的理论容量对比：
| 记忆存储器类型 | 数学表达 $f(x)$ | 理论容量 | $d_k = 64$ 时的容量 |
|---|---|---|---|
| 外积存储器 | $V^\top K x$ | $d_k$ | 64 |
| 最优线性存储器 | $V^\top K (K^\top K) ^ {-1} x$ | $2d_k$ | 128 |
| ReLU MLP | $V^\top \operatorname{ReLU}(Kx)$ | $2d_k$ | 128 |
| 平方 ReLU (GLU MLP) | $V^\top \operatorname{ReLU}(Kx)^2$ | $2d_k^2/3$ | 约 2,700 |
| SwiGLU MLP | $V^\top \left(\operatorname{Swish}(Kx) \odot Kx\right)$ | 约 $2d_k^2/3$ | 约 2,900 |
| 三次方 ReLU | $V^\top \operatorname{ReLU}(Kx)^3$ | $2d_k^3/15$ | 约 35,000 |
| 自注意力机制，$\beta = \sqrt{d_k}$ | $V^\top \operatorname{softmax}(\beta K x)$ | $e ^ {2\sqrt{d_k} - 2}$ | 约 120 万 |
| 自注意力机制，$\beta \propto \log M$ | $V^\top \operatorname{softmax}(\beta K x)$ | 实验范围内未触顶 | 超过 200 万 |
我们从最朴素的启发式外积记忆出发，一步一个脚印，逐步构建出了当今现代 Transformer 中所见到的所有标准网络层架构。
这种记忆存储器的视角，也为我们提供了一把极其清晰的量化标尺，用以评估不同架构在作为向量键值存储器时的本质容量差异。
值得特别说明的是，这种基于干扰信号比的容量估算，只是审视现代神经网络架构的诸多视角之一。我们的容量测算仅仅统计了单个网络层在干扰压倒信号之前所能准确召回的随机向量数量。在实际中，我们尚未考虑真实的训练数据分布，而训练数据深刻影响着键和值在真实潜空间中的表征方式！还有许多研究从训练后的网络层出发，以“比特（Bits）”为度量单位来测量模型存储的信息量：
- Morris 等人 测算得出，GPT 风格的模型每个参数大约能记住 3.6 比特的信息；并发现模型一旦填满其容量，便会开始泛化，对每个具体训练样本的死记硬背量随之减少。
- Allen-Zhu & Li 发现，语言模型每个参数大约能存储 2 比特的事实知识，即便模型被量化到 INT8 依然保持这一特性。
沿着这一脉络，还有大量激动人心的前沿方向值得深入探索：
- 单层网络的信息存储上限：线性存储方面的 Nichani, Lee & Bietti、Barnfield 等人；显式存储事实的 MLP 方面的 Dugan 等人、Garcia 等人；以及启发了本文干扰信号比估算方法的 Zhong 等人。
- Hopfield 网络与具有指数级容量的记忆存储器：Krotov & Hopfield 关于密集联想记忆（Dense Associative Memory）的开创性工作；将其与注意力机制紧密联系起来的 Demircigil 等人、Ramsauer 等人；以及 Krotov 等人 的综述。
- 推理过程中动态更新的短期记忆（Test-Time Memory）：快速权重（Fast Weights）、Mamba、Mamba-2、DeltaNet、Gated DeltaNet、TTT 层（Test-Time Training）、Titans、MesaNet 以及测试时回归（Test-Time Regression）。
- 用于补充或替代传统 MLP 的超大规模稀疏参数化记忆存储器：产品键记忆（Product-Key Memories）、规模化记忆层（Memory Layers at Scale）、UltraMem、百万专家混合架构（Mixture of A Million Experts）以及 Engram。
- 通过在权重外部追加记忆来扩展的非参数化记忆存储器：kNN-LM、RETRO、记忆化 Transformer（Memorizing Transformers）以及持续生长的非参数化网络（Growing Nonparametric Networks）。
- 定位与编辑训练后模型中的内部记忆：Geva 等人、Bietti 等人，以及知识编辑的里程碑算法 ROME 与 MEMIT。
- 压缩短期记忆，或将其固化为长期记忆：KVzip、要义词元（Gist Tokens）、Cartridges 以及注意力匹配（Attention Matching）。
- 持续学习方案：如稀疏记忆微调（Sparse Memory Finetuning）等。
我向来对“记忆”这一课题深深着迷，或许是因为我能切身体会到人类自身记忆的局限。也许在我内心深处，一直怀抱着这样的期许：通过理解机器究竟如何铭记万物，我也能更深地读懂我们自己。
脚注
- 有时，这种键和值的累加过程被称为“赫布学习”（Hebbian Learning），其核心思想在于：同步激活（相关联）的输入（键）与输出（值）会增强它们之间的突触连接强度（即权重）。 ↩
- 这种外积记忆存储器正是线性注意力机制（Linear Attention，亦称“快速权重 / Fast Weights”）所采用的状态表达形式：每个 Token 通过在矩阵中累加 $v k^\top$ 来存储记忆 $(k, v)$，以供后续 Token 查询检索（参见 Katharopoulos 等人；Schlag, Irie & Schmidhuber）。 ↩
- 此处的分析针对的是召回一个纯净无噪的值向量。如果每个值只需在一个固定的候选集合中指向正确答案（例如分类任务中挑选概率最高的 Token），那么召回结果只需比错误答案更接近正确答案即可；在此类设定下，一个 $d\times d$ 维的矩阵大约能容纳 $d^2$ 个事实（参见 Nichani, Lee & Bietti）。 ↩
- 我们同样可以利用这种回归目标函数，来设计序列模型在前向传播过程中动态写入的短期记忆（如线性注意力与自注意力）。具体可参见作者关于测试时回归的研究工作。 ↩
- 对于随机生成的键，ReLU 能够将干扰降低 $\sqrt{2}$ 倍。因为随机键的内积相似度为正或为负的概率均等，ReLU 会将大约一半的负权重归零，保留另一半正权重。这使得干扰的方差（权重平方和）减半，从而使干扰项自身的模长缩小了 $\sqrt{2}$ 倍。完整数学推导参见 Zhong 等人。 ↩
- 在标准的 Transformer 中，是对相似度除以 $\sqrt{d_k}$ 而非乘以该系数。但标准 Transformer 中的查询与键的向量长度通常约为 $\sqrt{d_k}$，使得完全匹配的一对键值具有 $q^\top k \approx d_k$。在本文的推导中，键和查询均被归一化为单位长度，完全匹配时的相似度为 1。因此相应的缩放系数对应为 $\beta = \sqrt{d_k}$。 ↩
- 在 Transformer 的标准符号表示中，相似度通常记作 $k_i^\top q$（其中 $q = W_Q x$），并除以缩放系数 $\sqrt{d_k}$。在本文中，我将投影矩阵 $W_Q$ 吸收融合到了键的定义中。Transformer 中的键与查询模长通常约为 $\sqrt{d_k}$，这一缩放与本文在单位长度键下取 $\beta = \sqrt{d_k}$ 是完全等价的。 ↩