MLP是Hebbian的：为Transformer构建高效的事实存储MLP

# MLP 是赫布型记忆体：为 Transformer 构建高效的事实存储 MLP  
来源：https://arxiv.org/html/2607.10034  

Roberto Garcia¹†  
Jerry Liu¹\*†  
Ronny Junkins²\*  
Sabri Eyuboglu²  
Atri Rudra³  
Chris Ré²  

¹斯坦福大学计算与数学工程研究所  
²斯坦福大学计算机科学系  
³布法罗大学计算机科学与工程系  

###### 摘要  

大型语言模型（LLMs）将事实知识存储在其参数中。尽管近期研究表明这些知识位于 MLP 层，但现有的事实存储构造性与机制性可解释性模型未能解释一个令人惊讶的实证现象：LLM 以信息论最优速率存储事实。本文中，我们为此现象提供了一个理论解释。我们首次开发出与 Transformer 兼容的事实存储 MLP 闭式构造，该构造满足 LLM 中经验观察到的三个性质：它 (i) 达到最优事实存储缩放，(ii) 处理任意输入/输出几何结构，以及 (iii) 能在 Transformer 内部工作。我们的关键在于分析 MLP 的*解码间隔*，而此前研究仅关注 MLP 事实存储。在等向性嵌入下，我们的构造实现了信息论最优的存储容量缩放，在匹配事实数量下所需的参数比先前构造少 \(10^{10}\) 到 \(10^{4}\) 倍。对于任意的键和值嵌入，我们证明该构造能达到相同的存储容量缩放，仅受嵌入几何结构决定的惩罚因子影响。此外，我们展示了所构造的 MLP 可在 Transformer 模块内以最优容量缩放用于事实回忆任务，在匹配事实数量下参数比先前构造少 15 到 63 倍。最后，作为概念验证，我们展示了事实存储 MLP 通过用新 MLP 替换 Transformer 的 MLP 实现*模块化事实编辑*。  

---
²通信作者：[email protected] (https://arxiv.org/html/2607.10034v1/mailto:[email protected])，[email protected] (https://arxiv.org/html/2607.10034v1/mailto:[email protected])。  
³代码发布于 https://github.com/HazyResearch/hebbian-mlps。  

## 1 引言  

大型语言模型（LLMs）在数学、科学和法律等领域的出色表现 (Google DeepMind, 2024 (https://arxiv.org/html/2607.10034#bib.bib53); Guha et al., 2023 (https://arxiv.org/html/2607.10034#bib.bib18); Saab et al., 2024 (https://arxiv.org/html/2607.10034#bib.bib15))，部分归因于它们能在参数中存储海量知识 (Petroni et al., 2019 (https://arxiv.org/html/2607.10034#bib.bib16); Meng et al., 2023a (https://arxiv.org/html/2607.10034#bib.bib38))。先前研究表明，Transformer 中的知识存储于多层感知机（MLP）中，形式为键值映射，即*事实* (Geva et al., 2021 (https://arxiv.org/html/2607.10034#bib.bib33); Dai et al., 2022 (https://arxiv.org/html/2607.10034#bib.bib39))。然而，尽管有这些发现，事实存储 MLP 仍然未得到充分理解。虽然先前工作在理解和建模 MLP 事实存储方面取得了重要进展，但现有模型未能捕捉到 LLM 事实存储的三个经验观察性质：MLP 必须 (i) 达到最优事实存储缩放，(ii) 处理任意输入/输出几何结构，(iii) 能在 Transformer 内部工作。  

机制可解释性工作 (Geva et al., 2021 (https://arxiv.org/html/2607.10034#bib.bib33); Dai et al., 2022 (https://arxiv.org/html/2607.10034#bib.bib39)) 假设 MLP 将事实存储在单个神经元中，但这些模型导致次优的存储容量缩放。最近，Nichani et al. (2024 (https://arxiv.org/html/2607.10034#bib.bib40)) 通过引入具有理论容量保证的 MLP*权重建构* (NTK MLP) 来研究 LLM 事实存储。然而，现有构建 (i) 在理论和实证上均未达到 LLM 中观测到的信息论最优容量缩放；(ii) 局限于等向性（例如均匀球形）嵌入分布，而 LLM 嵌入是各向异性的 (Ethayarajh, 2019 (https://arxiv.org/html/2607.10034#bib.bib81); Razzhigaev et al., 2024 (https://arxiv.org/html/2607.10034#bib.bib82))；(iii) 无法被 Transformer 模块用于事实回忆任务，例如回答“法国的首都是什么？”。  

我们的核心洞察是研究 MLP 解码间隔的缩放（第 2.1 节）。先前工作关注无噪键查询下的事实存储，而我们的解码间隔研究使我们能够开发出首个闭式 MLP 构造，该构造可被 Transformer 用于事实回忆。因此，我们从理论上证明*Transformer 模块能够实现信息论最优的事实存储*——首次提供了与预训练 LLM 中实证容量缩放一致的理论解释。我们的构造是首个捕捉 LLM 事实存储所有三个经验观测性质的构造：  

参照图注  

图 1: (A) 解码间隔示意图：上图为大解码间隔的 MLP，下图为低解码间隔的 MLP。大解码间隔的 MLP 具有更大的“误差容限”，在查询受到扰动的版本 \(k_1\) 时更为鲁棒。(B) MLP 是信息论最优的：我们展示了 MLP 是核空间中的赫布记忆，且 MLP 的存储容量以信息论最优速率缩放。(C) Transformer 模块是信息论最优的：我们展示了在注意力噪声有界的条件下，Transformer 模块的存储容量以信息论最优速率缩放。  

- • **最优事实存储容量与间隔（第 3、4 节）**。我们迈向匹配 LLM 实证最优事实存储缩放的第一步是证明：1) MLP 是赫布核记忆，2) 赫布核记忆能达到渐近最优的事实存储缩放。我们通过开发一个闭式 MLP 构造来证明这一点，该构造等价于具有草图化二次核的赫布记忆，并证明它对于 \(F\) 个事实使用 \(W\) 个参数能达到事实存储容量 \(F = \Theta(W / \log W)\)（第 4.3 节）。理论上，在等向性嵌入下，我们的构造将先前景构造的最优性差距缩小了 \(\log^{11} F\) 倍。此外，在匹配事实数量下，NTK 基线所需的参数比我们最好的数据依赖核构造多 \(10^{10}\) 到 \(10^{4}\) 倍（图 2c）。  
- • **处理任意嵌入几何结构（第 4 节）**。我们接下来研究任意嵌入几何结构如何影响 MLP 间隔和存储容量的缩放。我们展示了将 MLP 间隔和容量缩放推广到任意嵌入几何结构，会在等向性嵌入下推导的信息论最优缩放中引入四个嵌入几何统计量作为乘法因子（定理 4.3）。直观上，这些统计量根据键和值嵌入的聚集程度来惩罚间隔和容量缩放。实证上，我们展示了广义边界精确刻画了经验解码间隔的缩放（\(R^2 \geq 0.95\)；图 5c）。此外，我们发现在各向异性嵌入下，我们的构造与训练过的 MLP 之间的容量差距仍然存在。  
- • **可在 Transformer 内用于事实回忆的 MLP（第 5 节）**。为了理解 MLP 在 LLM 中的使用，我们发现 MLP 需要非平凡的解码间隔才能在 Transformer 模块内用于事实回忆。注意力层会产生不完美的、有噪的查询，因此被查询的事实存储 MLP 应对噪声具有鲁棒性。基于此洞察，我们首次从理论和实证上证明，Transformer 模块可以从具有最优事实存储容量的 MLP 中检索事实，以解决事实回忆任务（定理 5.2，图 3c）。此外，在这些 Transformer 实验中，在匹配事实数量下，NTK 基线所需的参数比使用我们数据依赖构造的 Transformer 模块多出大约 15 到 63 倍。最后，作为概念验证，我们展示了事实存储 MLP 能够通过替换 Transformer 中的 MLP 来实现*模块化事实编辑*（第 5.2 节）。我们的方法 **MLP 交换** 实现了近乎完美的*事实编辑分数*——正确地编辑目标事实同时避免脱靶效应——而先前的最先进方法在编辑 10% 的事实集时，分数会下降到仅约 30%。  

总之，我们的工作向着理解 Transformer 中的 MLP 迈出了构建性的一步。我们提出了一种事实存储 MLP 构造，它实现了最优的间隔和事实存储容量，提供了任意嵌入下的可证明解码间隔保证，并可在 Transformer 模块内以最优容量进行事实回忆。我们还展示了在模块化事实编辑中的应用，为 LLM 中鲁棒且模块化的知识操作提供了一条路径。  

## 2 预备知识  

### 2.1 形式化事实知识  

##### 事实集与存储。  
给定键嵌入 \(\mathbf{K}\in\mathbb{R}^{|\mathbf{K}|\times d}\) 和值嵌入 \(\mathbf{V}\in\mathbb{R}^{|\mathbf{V}|\times d}\)，一个**事实集**是一个映射 \(f: [|\mathbf{K}|] \to [|\mathbf{V}|]\)。我们将第 \(i\) 个键和值嵌入分别记为 \(\mathbf{k}_i\) 和 \(\mathbf{v}_i\)。  

###### 定义 2.1 (事实存储)  
给定嵌入 \(\mathbf{K}\) 和 \(\mathbf{V}\)，模型 \(g_\theta: \mathbb{R}^d \to \mathbb{R}^d\) **存储**事实集 \(f: [|\mathbf{K}|] \to [|\mathbf{V}|]\)，若对所有 \(i \in [|\mathbf{K}|]\) 和所有 \(j \neq f(i) \in [|\mathbf{V}|]\)，有  
\[
\langle g_\theta(\mathbf{k}_i), \mathbf{v}_{f(i)} \rangle > \langle g_\theta(\mathbf{k}_i), \mathbf{v}_j \rangle.
\]  
(1)  

值得注意的是，该定义等价于语言建模中正确的 softmax 解码。  

###### 定义 2.2 (间隔)  
\(g_\theta\) 在事实 \(i\) 上对抗竞争事实 \(j \neq f(i)\) 的**间隔**为  
\[
\gamma_{i,j} := \langle g_\theta(\mathbf{k}_i), \mathbf{v}_{f(i)} \rangle - \langle g_\theta(\mathbf{k}_i), \mathbf{v}_j \rangle,
\]  
(2)  

最小间隔为 \(\gamma_{\min} := \min_{i, j \neq f(i)} \gamma_{i,j}\)。注意，存储一个事实集（按照第 2.1 节的意义）等价于 \(\gamma_{\min} > 0\)。  

##### 事实存储成本与容量。  
为衡量参数效率，我们定义一个模型类存储固定嵌入上*所有*事实集所需的最小参数预算。  

###### 定义 2.3 (事实存储成本与容量)  
模型类 \(\mathbf{g}\) 在嵌入 \(\mathbf{K}\) 和 \(\mathbf{V}\) 上的**事实存储成本**是表示*所有*可能事实集所需的最小参数数量：  
\[
W(\mathbf{g}; \mathbf{K}, \mathbf{V}) = \min\left\{ \#(\theta) \;\Big|\; \forall f: [|\mathbf{K}|] \to [|\mathbf{V}|], \; \exists\,\theta \text{ s.t. } \mathbf{g}_\theta \text{ stores } f \right\}.
\]  
(3)  

相应的**事实存储容量**是在固定参数预算下可存储的最大事实数量。  

###### 定理 2.4 (信息论下界)  
假设每个参数有恒定的比特数，嵌入 \(\mathbf{K}\) 和 \(\mathbf{V}\) 对于*任何*模型类 \(\mathbf{g}\) 的事实存储成本满足  
\[
W(\mathbf{g}; \mathbf{K}, \mathbf{V}) = \Omega(|\mathbf{K}| \log [|\mathbf{V}|]).
\]  
证明见附录 B.1。  

### 2.2 模型类别  

本文中，我们研究两个模型类别：门控单隐层 MLP 和赫布记忆 (Kohonen, 1972 (https://arxiv.org/html/2607.10034#bib.bib74); Hopfield, 1982 (https://arxiv.org/html/2607.10034#bib.bib75); Bubeck et al., 2020 (https://arxiv.org/html/2607.10034#bib.bib43); Cabannes et al., 2024 (https://arxiv.org/html/2607.10034#bib.bib76); Nichani et al., 2024 (https://arxiv.org/html/2607.10034#bib.bib40))。  

**MLP**。我们考虑形式为 \(g_\theta: \mathbb{R}^d \to \mathbb{R}^{d_v}\) 的模型  
\[
\text{MLP}(\mathbf{x}) = g_\theta(\mathbf{x}) = \mathbf{B}\!\left( (\mathbf{A}\mathbf{x}) \odot \sigma(\mathbf{G}\mathbf{x}) \right),
\]  
(4)  

其中 \(\mathbf{x} \in \mathbb{R}^d\)，\(\mathbf{A}, \mathbf{G} \in \mathbb{R}^{m \times d}\)，\(\mathbf{B} \in \mathbb{R}^{d_v \times m}\)。该族包括现代语言模型中使用的 SwiGLU 风格 MLP (Shazeer, 2020 (https://arxiv.org/html/2607.10034#bib.bib77); Yang et al., 2025 (https://arxiv.org/html/2607.10034#bib.bib51); DeepSeek-AI et al., 2025 (https://arxiv.org/html/2607.10034#bib.bib78); Dubey et al., 2024 (https://arxiv.org/html/2607.10034#bib.bib1))。我们的显式构造 (第 4.1 节) 使用 \(\sigma = \mathrm{id}\)。  

**赫布记忆**。这些线性模型是形式为 \(g_\mathbf{W}: \mathbb{R}^d \to \mathbb{R}^{d_v}\) 的映射  
\[
g_\mathbf{W}(\mathbf{x}) = \mathbf{W} \mathbf{x},
\]  
其中 \(\mathbf{x} \in \mathbb{R}^d\)，\(\mathbf{W} \in \mathbb{R}^{d_v \times d}\)。一个事实集通过以下方式存储：  
\[
\mathbf{W} = \sum_{j=1}^{|\mathbf{K}|} \mathbf{v}_{f(j)} \mathbf{k}_j^\top.
\]  
(5)  

第 3 节展示了 MLP 可以重新解释为核特征空间中的赫布记忆。  

### 2.3 相关工作  

与本文最接近的两项工作是 Nichani et al. (2024 (https://arxiv.org/html/2607.10034#bib.bib40)) 和 Zhong et al. (2025 (https://arxiv.org/html/2607.10034#bib.bib79))。Nichani et al. (2024 (https://arxiv.org/html/2607.10034#bib.bib40)) 首次给出了事实存储 MLP 的显式构造，并展示了近乎最优的事实存储容量（差一个多对数因子），但其分析局限于等向性嵌入，且仅研究可分离性条件 \(\gamma_{\min} > 0\)。Zhong et al. (2025 (https://arxiv.org/html/2607.10034#bib.bib79)) 开发了注意力与 MLP 的统一联想记忆视角，但侧重于平均检索保真度而非最坏情况间隔。相比之下，我们推导了超越等向性的显式间隔界，识别出间隔远离零是 Transformer 可用性的条件，并展示了构造的 MLP 可以集成到 Transformer 模块中进行事实回忆。此外，我们还在模块化事实编辑方面展示了初步应用。