Transformer 架构中的每一层本质上都是一个记忆存储单元。通过交互式演示与代码实现，可以直观地理解这一机制。
当前的大型语言模型（LLM）能够从全球范围内检索模糊的事实信息。例如，询问 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}$ 为需要检索的内容。通过输入 $x$ 查询 $f$，输出值为 $f(x)$。理想的存储器应在使用 $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$ 添加到现有权重矩阵中。
通过输入 $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{signal}} + \underbrace{\sum_{i\neq j} v_i\,(k_i\transpose k_j)}_{\text{interference}},$$这意味着检索结果包含目标记忆的值（信号），以及其他键与查询键相似所产生的干扰项。
如果所有键正交，干扰为零，回忆将是完美的。然而，$\reals^{d_k}$ 空间中仅有 $d_k$ 个正交方向，因此一旦存储的记忆数量超过关键维度（$M>d_k$），就无法实现完美回忆。
外部积存储的容量可通过直观推导得出。存在 $M-1$ 个干扰项。若键为随机单位向量，每个相似度 $k_i\transpose k_j$ 约为 $1/\sqrt{d_k}$。若值也是随机且长度相近，$M-1$ 个干扰项如同随机游走般叠加，其平方长度约为 $(M-1) /d_k$ 乘以信号强度。
$$\frac{\norm{\text{interference}}}{\norm{\text{signal}}} \sim \sqrt{\frac{M}{d_k}}.$$当该比率为 1 时：
$$M \approx d_k,$$此时干扰与信号同等大小，导致信号被淹没。
下图展示了单次回忆中各存储记忆的权重相对于目标记忆权重的比例。拖动滑块使 $M$ 超过 $d_k=64$，可见小条堆积导致回忆失败。
如何减少干扰？#
即使存储容量低于 $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}.$$该比率在以下情况达到临界点：
$$M = 2d_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$ 甚至 $< 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 MLPs 赋予门控独立权重，并使用 Swish 激活替代 ReLU。
实践中，MLP 以叠加方式存储事实：许多事实共享权重，输出是多种记忆的混合。Gracia 等人建立了明确模型：GLU MLP 的隐藏单元计算输入的随机特征 $\phi(x)$，第二个权重矩阵是针对这些特征的外部积存储 $\sum_i v_i \phi(k_i) \transpose$。尽管 MLP 存储的事实远多于其隐藏单元数，但训练好的 Transformer ReLU MLP 对于给定输入仅激活其隐藏单元的 1%，意味着每次回忆只利用了极小部分记忆（Li 等人）。
通过更高幂次减少干扰#
由于高次幂会放大较大权重并缩小较小权重，可使用 $\sigma(s) = \operatorname{ReLU}(s) ^p$ 进一步增强效果。信号保持不变（因 $1^p=1$），而相似性为 $1/\sqrt{d_k}$ 的干扰项缩小至 $d_k^{-p/2}$。因此，干扰与信号比率约为：
$$\frac{\norm{\text{interference}}}{\norm{\text{signal}}} \approx \sqrt{\frac{M}{d_k^p}}.$$这意味着容量约为：
$$M \approx d_k^p$$相比之前呈几何级数增长。
对 $\sigma$ 的五种选项进行相同实验：
由于干扰随 $\sqrt{M}$ 增长，在对数-对数图中线条几乎平行，仅高度不同，对应容量变化：
- 线性：容量 $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{干扰项像随机游走一样叠加} \\ &= \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{对其他 } M \text{ 个记忆取平均} \\ &= \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$ 呈指数增长，而非之前的多项式增长。当 $\beta = \sqrt{d_k} = 8$ 时，容量约为 $e^{14} \approx 1.2$ 百万条记忆，几乎是线性存储的 20,000 倍。长时间运行同一实验也显示容量约为 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).$$这就是自注意力机制。
与键和值为训练权重的 MLP 不同，自注意力从其上下文中的 token 计算键和值。每个 token 存储一个记忆，因此记忆数量 $M$ 等于上下文长度。这使得自注意力成为非参数化的：其记忆不是参数，存储库随上下文增长。
下图再次展示了条形图，起始于 softmax。在 $\beta = 8$ 时，滑块范围内回忆从未失败。降低 $\beta$ 至 2 并增加 $M$：几十条记忆后回忆开始失败。$M$ 每翻倍，需要稍高的 $\beta$ 来维持回忆功能，这引出了下一节的内容。
处理长上下文中的大量键值对#
由于 $\beta$ 决定了对低相似性键的抑制强度，随着键值对数量增加，可提高 $\beta$。为何不直接将 $\beta$ 设为极大值？当 $\beta \to \infty$ 时，softmax 变为硬 argmax，仅选择最相似的键，但这不可微分；且极大的 $\beta$ 会导致 softmax 饱和，其他键的梯度消失，从而损害训练效果。因此，仅根据 $M$ 的需求适度提高 $\beta$。
