如果 2017 年那篇《Attention Is All You Need》从未发表，给你一个问题：「如何让机器记住尽可能多的东西，并且想取哪条就取哪条」，你能从零开始，把 Transformer 一层一层地推导出来吗？
很多人可能觉得一头雾水。10 月 15 日，斯坦福大学计算机科学博士、前 Cartesia 研究科学家 Alex Wang 发表了一篇博客《You could have invented transformers from memory》（你本可以凭「记忆」发明 Transformer），迅速在社交媒体引发关注。
正如标题所展示的，这篇文章做了一件很「野心勃勃」的事：它提出一个统一的理解视角 —— 我们熟悉的神经网络层，可以被看作不同形式的「联想记忆存储器」。
线性层、MLP、自注意力，看似差异很大的三种结构，其实都可以从「存记忆、取记忆」的角度理解。衡量它们记忆能力的关键指标，是信息之间的干扰：存储的记忆越多，彼此干扰越严重。谁能在干扰淹没目标信号之前装下更多记忆，谁的「容量」就越大。
沿着这个思路，作者像搭积木一样完成了一次有趣的推导：从最朴素的外积记忆出发，引入非线性激活，就自然「长」出了 MLP；进一步提高非线性程度，又与 GLU 等门控结构建立了联系。
而当你想要「用上所有幂次」来压制干扰时，指数函数登场了。经过归一化处理，便得到了 Softmax，自注意力的数学形式也随之浮现。
效果差异是可以量化的。在作者设定的随机记忆实验中，同样使用 64 维的 Key，最简单的线性存储器容量约为 64 条，平方 ReLU 构成的 GLU 式 MLP 能存储约 2700 条，而采用固定缩放系数的 Self-Attention，理论容量达到约 120 万条。
从 64 到 120 万，容量差了近两万倍。当然，这里的容量对应的是理想化条件下的随机向量记忆任务，并不意味着真实的 Transformer 能记住比 MLP 多两万倍的知识。
更有意思的是，Transformer 中著名的 1/d_k 缩放因子，以及 Qwen 的 LogN-Scaling、YaRN 等长上下文扩展方法采用的注意力缩放策略，也能从记忆干扰的角度得到解释。
而这篇文章与 2017 年 Transformer 原始论文之间，还有一段鲜为人知的联系。Alex Wang 发现，《Attention Is All You Need》的早期 arXiv 版本中，研究者就曾将前馈网络描述为一种特殊的 Attention，甚至尝试用基于可学习 Key-Value 参数的 Attention 替换前馈层，取得了相近的机器翻译效果。但这段讨论后来从论文中消失了。
九年之后，Alex Wang 重新拾起了这条线索，将线性层、MLP 和 Self-Attention 串联到同一个数学框架中。如果说 2017 年的《Attention Is All You Need》让 Transformer 成为深度学习的主流架构，那么这篇博客则邀请读者换一个角度，重新理解这些核心组件为何有效。
从记忆这个古老的问题出发，那些看似各自独立的神经网络设计，原来可以沿着同一条思路被重新推导出来。
社区的反馈很能说明问题。一位 HN 读者感叹：「真希望 13 年前上那些枯燥的 MLP 课程时，就有人这样讲。」
另一位读者指出，softmax 通过压制弱激活来提升记忆容量的解释「极其清晰且有趣」。还有一位通信工程师发现了奇妙的呼应：干扰信噪比、信道容量，这些通信领域的老概念，原来可以直接用来理解 AI。
当然，也有严谨的讨论：有评论提醒，「注意力即记忆检索」只是一个有用的类比，并不能完全解释训练动态，作者本人也坦然承认这一点。但正如他所言，理解神经网络的「记忆」如何工作，可能是通向可解释性的一把钥匙（Alex 在此处转发了 OpenAI 可解释团队创建人 Nick Cammarata 前两天的一篇长文，参见《神经网络可解释性被 AI 解决了？90% 神经元只有一个功能》）。
那么，这篇文章具体是如何一步步从「记忆」推导出 Transformer 的？下面我们进入正文。
博客链接：https://alexwang.ai/posts/associative-memory/
从记忆出发，
你本可以发明Transformer
从记忆中提取某个角色。记忆数量越多，回忆效果越差。
如今，大语言模型能够回忆起世界各个角落的冷门知识。最近有一次，我赶着参加朋友的桌游之夜，快要迟到了，于是请 Astra 帮我回顾《沙丘：帝国》的规则。它欣然在我开车前往的途中讲解规则，帮我回忆到足以再次上手的程度。我们已经习惯了什么都向模型询问，以至于把它们渊博的知识视为理所当然。
模型是如何记住这么多信息的？它们又能存储多少信息？
本文将介绍一种直观的视角，把线性映射、MLP 和自注意力看作不同类型的记忆存储器。从这种直觉出发，我们可以推导出一些启发式估计，判断每一层能存储多少记忆。
衡量一个记忆存储器是否「好」的关键，在于它能否回忆起所存储的内容。毕竟，即使我们能够存储无限多的事实，如果需要时无法回忆起来，也没有多大帮助。存储的记忆越多，在尝试回忆其中一条时，它们相互之间的干扰就越强。
因此，本文将始终关注一个量：干扰信号比，它衡量回忆过程中其他记忆造成的干扰有多大。在干扰增大到与信号同样强之前，存储器能够容纳的记忆数量，就是它的容量。
我们将逐步展开这一视角，展示如何从线性层出发，得到更好的结构 —— 同时改善容量和回忆效果的 MLP，再一路走向自注意力。自注意力的回忆效果非常好，能够存储的记忆数量还可能没有上限。对于 64 维的键，这意味着记忆数量从大约 64 条，跃升到大约 2700 条，再到超过一百万条。
下面是全文使用的术语。记忆存储器是一个函数 𝑓，其中保存了 𝑀 条记忆，编号为 𝑖 =1,…,𝑀。一条记忆就是一个键值对：键 用于根据相似度进行查找，值
则是我们希望检索的内容。使用 𝑓 时，我们以输入 𝑥 对其进行查询，它会回忆出输出 𝑓(𝑥)。好的记忆存储器在我们用 𝑘_𝑖 或与 𝑘_𝑖 相似的另一个向量查询时，能够回忆出 𝑣_𝑖。在本文中，这种基于相似度进行回忆的记忆存储器称为联想记忆。
外积记忆存储器
最基本的记忆存储器是线性的：它只是通过外积累积记忆，
其中， 和
的各行分别是键
和值
。
要存储一条新记忆 (𝑘,𝑣)，只需将 加到已有权重上。
补充：一步梯度下降就是写入一条记忆
考虑一个具有标量损失 𝐿 的网络中的线性层 𝑦 =𝑊𝑥，令 为损失相对于该层输出的梯度。相对于 𝑊 的梯度是
，因此，学习率为 𝜂 的一步更新为
这个递推式就是一次记忆更新：它用一条键为 𝑥、值为 的新记忆更新权重。关于这一视角的更多内容，参见 Irie、Csordás 与 Schmidhuber 的研究（https://arxiv.org/abs/2202.05798）；他们将经梯度下降训练的线性层表示为对其全部训练输入的注意力。
要用输入 𝑥 查询记忆存储器，我们将输入乘以该矩阵：
因此，回忆就是按照与已存储键的相似度进行查找：将查询与每一个键比较，再返回以这些相似度加权的值之和。
回忆会带来干扰
所以，存储记忆很容易：不断向权重中添加外积即可。但当我们试图回忆其中一条时，会发生什么？
为简单起见，假设每个键都已归一化为长度 1，即 。那么
这意味着，我们得到的是目标记忆的值 —— 即我们用其键进行查询的那条记忆 —— 再加上其他所有键与查询相似的记忆所产生的干扰。
如果所有键都相互正交，干扰就为零，回忆也会完全准确。但 只有 𝑑_𝑘 个正交方向，因此，一旦存储的记忆数量超过键的维度，即 𝑀 >𝑑_𝑘，我们就无法做到完美回忆。
外积存储器的容量是多少？下面给出一个直观论证。干扰项共有 𝑀 −1 个。如果键是随机单位向量，每个相似度 大约为
。如果值也是随机的，而且长度相近，那么这 𝑀 −1 个干扰项就像随机游走一样累加，因此干扰的长度平方大约是信号长度平方的 (𝑀 −1)/𝑑_𝑘 倍，于是
令这个比值等于 1，就得到容量 𝑀≈𝑑_𝑘, 一旦存储的记忆数量与键的维度相同，干扰就会与信号一样强，将信号淹没。
下方每根柱子表示一次回忆中某条已存储记忆的权重，相对于目标记忆的权重而言（目标记忆以红色标出）。每根柱子都会向回忆结果中加入一点对应的值。将 𝑀 拖动到超过 𝑑_𝑘 =64，就能看到小柱子不断累积，直到回忆失败。其他按钮用于切换到本文后面将构建的记忆存储器（交互图参见原文）。
如何减少干扰？
即使尚未达到大约 𝑑_𝑘 条记忆的容量，外积存储器在回忆时也会产生显著干扰。要提高容量，或者在未达到容量时获得更纯净的回忆结果，我们就需要减少干扰。
最直接的办法是增大键的维度 𝑑_𝑘，因为线性存储器的容量会随 𝑑_𝑘 增长。𝑑_𝑘 能有多大，受到硬件条件的限制。
另一条路是改进算法。我们可以更仔细地存储记忆：外积存储器添加每条记忆时都不考虑其他记忆，因此完全没有抵消相互干扰。或者，我们可以更有选择性地回忆：在线性存储器中，每条记忆都会造成干扰，即使其键与查询的相似度很低；所以，抑制相似度较低的键，可以改善回忆效果。
在本文余下部分，我们将看到，更仔细地存储外积记忆，会得到容量大约翻倍的最优线性记忆存储器。通过非线性进行有选择的回忆，则会得到 MLP 和自注意力。
最优线性记忆存储器
外积存储器写入每条记忆时都不考虑其他记忆。我们也可以选择权重 𝑊，使回忆结果同时拟合所有已存储的记忆。这是一个线性回归问题：
当存储的记忆数量超过键的维度，即 𝑀 >𝑑_𝑘 时，最小二乘解为
这就是外积存储器 ，但有一个变化：在将查询与键比较之前，先将查询乘以
，以消除键的协方差
的影响。这会使键「白化」，因此存储器不再对许多键共有的方向赋予过高的权重。
当 𝑀 ≤𝑑_𝑘 时，只要键线性无关，回归就能精确拟合每一条记忆。因此，即使键并不正交，最优存储器也能在完全没有干扰的情况下进行回忆。
当 𝑀 >𝑑_𝑘 时，干扰开始出现，因为记忆存储器只有 𝑑_𝑘 个方向可用。对于随机的键和值，干扰信号比变为
这个比值在下列条件下等于 1：
𝑀=2𝑑_𝑘,
因此，最优存储器的容量是之前的两倍！
补充：推导最优线性记忆存储器
当 𝑀 ≤𝑑_𝑘 时。如果键线性无关，那么由键之间的内积构成的 𝑀 ×𝑀 矩阵 可逆，并且
能够精确拟合每条记忆。当 𝑀 <𝑑_𝑘 时，该问题的约束不足，有无穷多个解，而这里给出的是最小范数解。 的各行
称为 𝐾 各行的对偶基：根据构造，
，因此
。于是
，不存在干扰，我们获得了完美回忆。
当 𝑀 >𝑑_𝑘 时。令目标函数的梯度为零，得到，所以
。两种解都可以写成
，其中 𝐴† 表示 𝐴 的伪逆：当键线性无关时，
；当键张成
时，
。
现在可以分析 的干扰。对于一个已存储的键 𝑘_𝑗，每个值 𝑣_𝑖 的权重为
，它恰好是 𝑀 ×𝑀 矩阵
的第 (𝑖,𝑗) 个元素：
𝑃 是到 𝐾 的列空间上的投影，这个列空间是 的一个 𝑑_𝑘 维子空间。投影的两个性质可以帮助我们分析干扰信号比：
信号项：根据迹的循环性质，
，因此，𝑀 个信号权重之和为 𝑑_𝑘。如果键是随机的，没有哪个键是特殊的，所以
。
干扰项：𝑃 是对称投影矩阵，因此 𝑃⊤𝑃 =𝑃2 =𝑃。于是
. 如果值是随机的，而且长度相近，
与外积存储器的 相比，这里的 −1 对应于最优存储器消除的、相当于 𝑑_𝑘 条记忆所产生的干扰。
我们可以用数值实验检验预测。随机采样 𝑀 条记忆，键为 𝑑_𝑘 =64 维的单位长度向量；用其中一个已存储的键查询记忆存储器，并测量干扰信号比：
外积存储器在 𝑀 ≈𝑑_𝑘 =64 时达到容量。最优存储器在 𝑀 =𝑑_𝑘 之前都能精确回忆。超过这一点，干扰按 从零开始上升，存储器在 𝑀 =2𝑑_𝑘 =128 附近达到容量，与预测一致。由于它消除的是相当于固定数量记忆产生的干扰，随着 𝑀 增大，两种存储器之间的差距逐渐缩小。
不过，这种最优记忆存储器对于相似度仍然是线性的。它必须用同一个系数缩放较小和较大的相似度。要做得更好，我们需要改变存储器的回忆方式。
作为非线性记忆存储器的 MLP
线性记忆存储器的回忆方式是
如果有办法在 较小时抑制对应权重，就能减少相似度较低的键造成的干扰。
在计算值的加权和之前，让每个相似度先经过一个非线性函数 𝜎：
我们刚刚从第一性原理重新发现了两层 MLP！第一个权重矩阵的各行是键，第二个权重矩阵的各列是值。Sukhbaatar 等人指出了前馈层与记忆存储器之间的这种相似性；Geva 等人发现，训练后的 Transformer 的前馈层往往就是这样工作的，其中的键会对输入中可辨认的模式产生激活。
理想情况下，我们选择的门控函数 𝜎 应使相似度较低，甚至
）时，求和中的对应项受到抑制，从而减少干扰。ReLU 激活会将相似度为负的键的权重置零，加入偏置则可以移动这一阈值。平方 ReLU 激活会进一步抑制相似度接近零的键。相似度 0.1 会变成 0.01，而相似度 1 仍然是 1。实际上，平方 ReLU 是一种不太为人所知、但已在大规模应用中得到验证的 MLP 激活函数改进。
平方 ReLU MLP 就是门控与输入投影共享权重的 GLU MLP，因为 。如今语言模型中的 SwiGLU MLP 为门控配置了独立的权重，并用 Swish 激活替代 ReLU。
在实践中，MLP 以叠加的方式存储事实：许多事实共享同一组权重，输出则是多条记忆的混合。
Garcia 等人显式构建了这样的模型：GLU MLP 的隐藏单元计算输入的随机特征 𝜙(𝑥)，第二个权重矩阵是建立在这些特征之上的外积存储器，而 MLP 存储的事实数量远超隐藏单元数量。尽管如此，对于给定输入，训练后的 Transformer 中的 ReLU MLP 可能只激活低至 1% 的隐藏单元，因此每次回忆只调用一小部分记忆。
Transformer 的作者早就知道这一点
最初的 Transformer 论文也指出了这种联系，但只出现在 arXiv 的 v2 和 v3 版本中。在这些版本里，Vaswani 等人将前馈层描述为「一种注意力」，其中两个权重矩阵的行充当键和值，ReLU 则取代 softmax。他们甚至尝试用对学习到的键值参数施加注意力来替换前馈层，得到的翻译质量大致相同：一个版本略差，另一个略好，不过第一个版本的训练速度慢了约 30%。此后的所有论文版本都删除了这段内容。下面我们会介绍自注意力与这些内容的关系。
用更高次幂减少干扰
对相似度取平方之所以有效，是因为它对低相似度权重的压缩远大于对高相似度权重的压缩。我们可以用 进一步推进这一思路。信号保持不变，因为
，而相似度为
的干扰项会缩小到
。𝑀 个干扰项像随机游走一样累加，因此干扰信号比为
这意味着容量大约为
条记忆，比之前按几何级数增加！
为检验这一点，我们对五种 𝜎 重复前面的实验：
由于干扰按 √𝑀 增长，双对数图上的线条几乎平行。变化的只有线条高度，以及随之变化的容量 —— 即每条线与 1 相交的位置：
线性：容量为 𝑀 ≈𝑑_𝑘 =64。
ReLU：将负相似度置零，使干扰缩小为原来的 1 / √2，因此容量增加到 𝑀 ≈2𝑑_𝑘 =128。
平方 ReLU：每个干扰项从约
缩小到约 1/𝑑_𝑘，因此容量增加到
，大约是线性存储器的 40 倍。
SwiGLU MLP：与平方 ReLU 一样，门控与输入投影共享权重，
；其容量与平方 ReLU 几乎相同。
立方 ReLU：每个干扰项缩小到约
，因此容量增加到
。
作为非参数记忆存储器的自注意力
对权重取幂，可以将容量提高到大约 。但如果把这一思路推向极限呢？与其选择一个幂次，不如使用无穷多个幂次。事实上，指数函数就是所有幂次的加权和，
其中，锐度 𝛽 决定高次幂的权重。下面采用自注意力所使用的值，对指数函数以及分别保留前 2、3、4 项的截断级数进行同样的实验：
保留 2 项时，容量约为 40，低于线性存储器，因为常数项为每个键赋予了约为 1 的权重。保留 3 项时，容量约为 450；保留 4 项时约为 3,100，接近平方 ReLU。完整指数函数的曲线始终远低于虚线，因此它的容量超出了图表范围。
为了求出指数函数的容量，我们用一个已存储的键 𝑘_𝑗 进行查询，因此目标记忆的相似度为 1，权重为 。其他每条记忆的权重为
；对于随机单位键，相似度 𝑠_𝑖 近似服从均值为 0、方差为 1/𝑑_𝑘 的高斯分布。当值随机且长度相近时，
令这个比值等于 1，得到容量
只要 𝛽 远小于 𝑑_𝑘，容量就随 𝛽 指数增长，而不再像之前那样随 𝑑_𝑘 多项式增长！当 时，容量约为
百万条记忆，接近线性存储器的 20000 倍。将同一实验运行到更大的规模，得到的容量也约为 120 万。
实际实现时，我们可能不希望直接计算指数，因为较小的数可能下溢，较大的数可能上溢。幸运的是，如上面的等式链所示，干扰只取决于权重之间的比值，因此我们可以将每个权重除以所有权重之和，在不改变干扰信号比的同时保证数值稳定性：
这就是自注意力。
MLP 的键和值是通过训练得到的权重，而自注意力的键和值由上下文中的 token 计算得到。每个 token 存储一条记忆，因此记忆数量 𝑀 就是上下文长度。正因如此，自注意力是非参数的：它的记忆不是参数，存储器会随着上下文增长。
下面再次给出柱状图，这次从 softmax 开始。当 𝛽 =8 时，滑块范围内的回忆始终不会失败，因此可以将 𝛽 降到 2，再增大 𝑀：存储几十条记忆后，回忆就会失败。𝑀 每翻一倍，都需要稍微增大 𝛽 才能维持回忆效果，这就是下一节的基本思路。
原文为交互图
处理长上下文中的大量键值对
由于 𝛽 决定了我们抑制低相似度键的力度，可以随着键值对数量增加而增大 𝛽。为什么不干脆把 𝛽 设得极大？当 𝛽 →∞ 时，softmax 会变成仅选择最相似键的硬 argmax，而它不可微；即使 𝛽 只是非常大，也会让训练效果变差，因为 softmax 会饱和，传向其他所有键的梯度都会消失。因此，我们只按 𝑀 的需要增大 𝛽。用 log𝑀 缩放注意力 logits，是实现这一点的一种长上下文技巧。它有几种不同版本，各自都相对于固定的参考长度 𝐿（通常是训练长度）设置 𝛽，以确保 𝑀 =𝐿 时，𝛽 保持训练时的值：
Su（2021）通过要求注意力的熵在 𝑀 增长时保持不变，推导出 log𝑀 的缩放形式，并进行归一化，使其在 𝑀 =𝐿 时与标准注意力一致。对于所有长度，logits 都乘以 log𝑀/log𝐿。Scalable-Softmax 具有相同形式：将 logits 乘以 𝑠log𝑀，其中 𝑠 通过学习得到。
最初的 Qwen 模型使用 LogN-Scaling，仅在上下文长度超过 𝐿 后应用相同的系数。
Qwen2.5-1M 使用 YaRN 的注意力缩放，将 logits 乘以
，因此其增长形式为
。
下面是在我们的实验中各条规则的表现，其中在 𝐿 =1024 条记忆时，：
原文为交互图
超过参考长度 𝐿 后，Qwen 的规则与 Su 的规则重合，干扰信号比随存储器增大而下降，在 200 万条记忆时降至约 0.004。YaRN 增长更快的缩放系数将这一比值压低至约。在这些情况下，让 𝛽 随 log𝑀 缩放，都能使存储器容纳具有不同键、数量基本上没有上限的记忆。
结语
下面汇总了我们讨论过的记忆存储器，以及在键和值随机、使用已存储的键查询时，它们各自的容量：
我们从一种启发式的记忆存储器 —— 外积记忆 —— 出发，逐步构建出了如今现代 Transformer 中的所有标准层。这一视角也让我们得以估计它们在容量上的差异，也就是它们作为向量值键值存储器的能力差异。
我应该说明，这种容量估计只是观察现代神经网络架构的一种方式。我们的容量估计，仅计算单层在干扰主导信号之前能够回忆多少个随机向量。实际上，我们完全没有考虑训练数据，而训练数据会影响键和值究竟表示什么！还有一些工作考虑了训练后的层，以比特衡量存储的信息。Morris 等人估计，GPT 类模型每个参数能记忆约 3.6 比特，并发现模型一旦填满容量，就会开始泛化，对每条训练样本的记忆反而减少。Allen-Zhu 与 Li 发现，语言模型每个参数存储约 2 比特的事实知识，即使量化到 int8 也是如此。
由此出发，还有许多同样引人入胜的相关方向（参考资料见原文）：
一层能够存储多少信息；
Hopfield 网络与具有指数容量的记忆存储器；
在推理期间更新的短期记忆；
用于补充或替代 MLP 的巨大而稀疏的参数化记忆存储器；
通过在权重之外追加记忆来扩展的非参数记忆存储器；
寻找并编辑训练后模型内部的记忆；
压缩短期记忆，或将其转化为长期记忆；
持续学习方法，例如稀疏记忆微调。
我一直对记忆着迷，也许是因为我切身体会到自己的记忆多么有限。或许，我内心的一部分也希望，通过理解机器如何记忆，我也能对自己有所了解。
© THE END
转载请联系本公众号获得授权
投稿或寻求报道：[email protected]