跳转至

Geometric Factual Recall in Transformers

会议: NeurIPS 2026
arXiv: 2605.12426
OpenReview: https://openreview.net/forum?id=34a614De3N
领域: 可解释性
关键词: 事实召回, 嵌入几何, 线性叠加, 多跳推理, 思维链

一句话总结

论文证明了在可学习嵌入下，Transformer能够将多关系事实直接编码为主体嵌入的线性叠加，并通过极小的ReLU MLP充当通用关系选择器实现对数维度事实召回，并在多跳查询下证明了CoT绕过容量下界瓶颈的几何机制。

研究背景与动机

大型语言模型具备存储海量世界知识并在生成时按需检索的能力。以往关于事实记忆的主流理论（如关联记忆假说）多将内部权重矩阵（特别是前馈全连接网络 MLP）视为键值映射存储器，假设输入嵌入为相互正交的随机向量，从而导致参数量必须随存储事实总数线性增长；而可解释性研究在观察到事实召回时，往往也将中间层 MLP 定位为主体-属性关联的主存储模块。然而，近期经验研究揭示了现实模型中不同实体的嵌入并非无序随机，而是自发展现出反映关系拓扑的低维几何结构，例如实体关系在隐藏状态中的线性可解码性。

这种理论假说与经验观察的核心矛盾在于：现有理论假设嵌入无结构，把全部记忆负担强行推给模型计算权重，无法解释为什么实际模型在高度压缩的参数空间下仍能高效操纵多关系事实。更严峻的是，当推理拓展到复合的多跳关系链（例如“某人妻子的母亲是谁”）时，单次前向传递中表示容量与模型深度的极限约束缺乏严格的信息论刻画，现代大模型究竟依靠何种机制突破多跳关系爆炸的理论壁垒亦缺乏完备阐释。

本文的切入角度是将词表嵌入本身视作可塑的几何表示空间，重新审视多关系共享属性场景下的记忆机制。核心 idea：Transformer利用可学习嵌入在实体向量中构建各关系属性的线性叠加，使极小规模的MLP直接退化为通用的关系条件选择器而非事实存储器，并在多跳查询中借助CoT的离散投影重新解耦几何表征，从而打破无CoT情形下的参数容量与维度瓶颈。

方法详解

整体框架

本文针对多关系事实召回问题建立形式化框架：设有 \(N\) 个主体与 \(R\) 个关系，每个关系对应一个双射函数 \(g_r: [N] \to [N]\)，映射到共享的 \(N\) 个属性池。对于单跳任务，模型输入序列为 \((s, r)\) 并预测 \(g_r(s)\)；对于 \(k\) 跳多跳查询，输入序列为 \((s_0, r_1, \dots, r_k)\)，目标输出为复合关系 \(y = (g_{r_k} \circ \cdots \circ g_{r_1})(s_0)\)。

在整体计算架构中，单层 Transformer 接收主体与关系 token，通过自注意力机制将关系条件与主体嵌入汇聚，随后送入前馈网络。区别于传统将事实硬编码进权重参数的思路，本方法揭示了双层分工：主体的初始输入嵌入将所有可能关系的属性特征进行超维度线性叠加，而后续的 MLP 仅充当关系选择网关，通过非线性激活函数对无关分支施加抑制并保留目标关系属性。多跳场景下，若采用直接预测，模型必须在隐空间静态维护指数膨胀的关系树；若引入思维链（CoT），模型则在每一步自回归生成中间实体，通过语言模型头将隐状态离散投影回词表，从而动态重置并复用单跳选择机制。

%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
    A["输入查询序列<br/>主体与关系条件"] --> B["自注意力聚合<br/>对齐主体叠加态与关系指示标记"]
    B --> C["单跳关系条件选择<br/>ReLU门控滤波与目标属性抽取"]
    C -->|直接单次前向| D["无CoT多跳计算<br/>遍历高维预计算答案树或宽MLP查表"]
    C -->|自回归离散发射| E["CoT离散词表重投影<br/>逐步检索1-hop局部邻域重置隐状态"]
    D --> F["输出最终属性预测"]
    E --> F

关键设计

1. 共享属性下的线性叠加嵌入与对数维度记忆：突破关联记忆的线性参数约束

在多关系共享属性设定下，若假设各属性编码相互正交，传统关联记忆理论要求参数量达到 \(\tilde{\Omega}(NR)\)。本文提出几何叠加构造：首先利用低失真随机编码引理，在维度为 \(4\log N\) 的超立方体顶点 \(\{\pm 1\}^{4\log N}\) 中为 \(N\) 个共享属性分配近似正交的向量 \(\mathbf{v}_1, \dots, \mathbf{v}_N\)。每个主体 \(i \in [N]\) 的嵌入向量并非随机独立采样，而是直接将其对应的全部 \(R\) 个目标属性向量沿坐标分块拼接，得到主体嵌入：

\[x_i = \begin{pmatrix} \mathbf{v}_{g_1(i)} \\ \vdots \\ \mathbf{v}_{g_R(i)} \\ 0 \end{pmatrix} \in \mathbb{R}^{4R\log N + 1}\]

该设计将所有事实关系全部吸收进词表嵌入的几何结构中，使得总嵌入维度收敛为 \(d = 4R\log N + 1\)。这种结构使得注意力层仅需使用均匀注意力将关系指示 \(r_j\) 与主体 \(x_i\) 相加，无需复杂的键值模式匹配。

2. 关系条件下的 ReLU 门控选择器：将 MLP 从记忆容器转变为布尔滤波器

以往研究通常认为 MLP 权重存储了知识三元组事实。本文构建了一个宽度仅为 \(O(R\log N)\) 的 3 层 ReLU MLP，其功能完全是无状态的关系选择器。定义关系距离过滤标量函数 \(f_j(z) = 2([-z + j]_+ + [z - j]_+)\)，当且仅当输入的查询关系编号 \(z = j\) 时 \(f_j(j) = 0\)，而在 \(z \neq j\) 时恒有 \(f_j(z) \ge 2\)。MLP 对叠加向量计算如下变换：

\[h(X) = \sum_{\ell=1}^R \left[ g_\ell(i) - f_\ell(j)\mathbf{1}_{4\log N} \right]_+ - \left[ -g_\ell(i) - f_\ell(j)\mathbf{1}_{4\log N} \right]_+\]

由于属性向量每个坐标均为 \(\pm 1\)，当 \(\ell \neq j\) 时，偏置项 \(-f_\ell(j) \le -2\) 使得 ReLU 内部输入严格为负并被直接归零；只有当 \(\ell = j\) 时，门控完全开启，恒等提取出目标属性向量 \(\mathbf{v}_{g_j(i)}\)。MLP 内部没有任何特定实体的常识参数，仅实现了按条件切除干扰分块的通用掩码机制。

3. 多跳无 CoT 场景的容量-深度权衡与信息论下界：显式刻画单向遍历瓶颈

针对 \(k\) 步复合事实召回 \(y = (g_{r_k} \circ \cdots \circ g_{r_1})(s_0)\)，论文证明了在无 CoT 辅助下单次前向传递的严格信息论下界：网络内部参数比特数 \(W\) 和嵌入比特数 \(D\) 必须满足 \(W \ge \max\{\mathcal{B}_{\text{global}}, \mathcal{B}_{\text{local}}\} - R \cdot D\)。通过在 \(R\)-正则有向边着色图上进行鸽巢原理与顶点不相交 \(R\)-叉子树计数，证明模型要么承受指数级的嵌入维度，要么承受宽度的线性膨胀： - 键值查找路径：保持紧凑维度 \(d = \tilde{O}(k)\)，但每层 MLP 必须退化为遍历 \(N \times R\) 事实表的暴力查表器，宽度达 \(O(NR)\)； - 嵌入预计算路径：允许嵌入维度膨胀至 \(d = \tilde{O}(R^k)\)，在初始主体嵌入中静态打包深度为 \(k\) 的整棵决策评估树，后续各层注意力与宽度为 \(\tilde{O}(R^k)\) 的布尔选择器逐层剪枝淘汰 \(R-1\) 个错误分支。

4. 思维链的离散词表重投影机制：动态解耦组合爆炸

与无 CoT 结构的组合爆炸形成对比，CoT 允许 Transformer 自回归生成中间实体 \(s_1, \dots, s_k\)。此时仅需单层 Transformer，嵌入维度压缩至 \(d = \tilde{O}(R + k)\)，MLP 宽度压缩至 \(\tilde{O}(R)\)。其核心机制在于：在自回归发射第 \(\ell\) 个中间 token \(s_\ell\) 时，语言模型非嵌入层（unembedding matrix）对词表进行内积取 \(\arg\max\)，将连续隐空间状态强制投射回离散的实体词表标识。当模型在下一步将 \(s_\ell\) 重新作为输入读入时，静态嵌入矩阵再次被动态检索，并直接注入该实体专属的局部 1-hop 叠加拓扑 \(\text{hop}(s_\ell) = [\text{bin}(g_1(s_\ell))^T, \dots, \text{bin}(g_R(s_\ell))^T]^T\)。这一离散跳转彻底打破了连续隐状态跨层传递信息时的信道容量下界，将指数级深层树展开降解为 \(k\) 次独立的单跳局部滤波。

实验关键数据

主实验

为了验证理论预测的几何记忆机制是否会在常规梯度下降优化中自然涌现，作者在合成共享属性数据集（\(N = 4096\)，遍历关系数 \(R \in [2, 16]\)，嵌入维度 \(d \in [32, 768]\)）上进行了系统评测，并与保持输入嵌入固定的随机正交基线（Nichani et al., 2025 所代表的纯关联记忆范式）进行对比。

模型配置 / 实体规模 关系数 \(R\) 最小饱和嵌入维度 \(d_{\min}\) (准确率 \(\ge 95\%\)) 最终全集准确率 (%) 属性线性探测准确率
可学习嵌入 (本文理论机制) (\(N=4096\)) \(R=2\) 32 \(100.0 \pm 0.0\) \(> 0.99\)
可学习嵌入 (本文理论机制) (\(N=4096\)) \(R=8\) 64 \(100.0 \pm 0.0\) \(> 0.99\)
可学习嵌入 (本文理论机制) (\(N=4096\)) \(R=16\) 128 \(100.0 \pm 0.0\) \(> 0.99\)
冻结随机嵌入 (关联记忆对照) (\(N=4096\)) \(R=2\) 128 \(100.0 \pm 0.0\) \(\approx 0.00\) (随机几率)
冻结随机嵌入 (关联记忆对照) (\(N=4096\)) \(R=8\) 256 \(100.0 \pm 0.0\) \(\approx 0.00\) (随机几率)
冻结随机嵌入 (关联记忆对照) (\(N=4096\)) \(R=16\) 512 \(100.0 \pm 0.0\) \(\approx 0.00\) (随机几率)

在真实预训练语言模型（Qwen2.5-0.5B、Qwen3-14B、Llama-3.1-8B、Llama-3.2-1B、Phi-4）针对 6 类实体域（人物、公司、建筑、电影、编程语言、物种）的 42 种真实事实关系探测中，主体末尾 token 隐状态与最终输出词表向量之间的仿射探针准确率同样展现出清晰的逐层几何富集趋势。

真实预训练模型 (平均跨多关系) 实体类别 输入层 (Layer-0) MRR 最佳隐藏层 MRR 最佳层 Hits@1 (%) 峰值相对网络深度
Qwen3-14B (40 layers) 人物 (people) 0.465 0.707 65.9 36/40 (90%)
Qwen3-14B (40 layers) 公司 (companies) 0.461 0.640 59.4 36/40 (90%)
Qwen3-14B (40 layers) 建筑 (buildings) 0.490 0.672 58.8 40/40 (100%)
Llama-3.1-8B (32 layers) 人物 (people) 0.436 0.701 65.1 28/32 (87.5%)
Llama-3.1-8B (32 layers) 建筑 (buildings) 0.308 0.547 46.2 32/32 (100%)
Phi-4 (40 layers) 人物 (people) 0.437 0.707 65.6 12/40 (30%)
Phi-4 (40 layers) 公司 (companies) 0.363 0.565 50.0 40/40 (100%)

消融实验

为了严格证明 MLP 学到的是无偏的关系选择机制而非记忆了特定实体的答案，作者设计了因果干预与零样本双射迁移实验：在固定训练收敛的 MLP 权重后，重新随机生成全新的双射映射 \(g_r'\)，并通过各关系线性读出矩阵的伪逆进行主体嵌入智能重初始化 \(s_x = \text{target} \cdot W_{\text{stack}}^+\)。

评估配置 干预/迁移设置 召回准确率 / 选择性得分 说明
可学习嵌入完整模型 原始双射测试 100.0% 正常收敛状态
因果属性微扰干预 最优秩截断伪逆扰动 Selectivity \(\approx 0.85 \sim 0.95\) 改变目标属性时，其余关系预测保持稳定
冻结 MLP 权重迁移 智能初始化 \(W_{\text{stack}}^+\) (全新双射) 高零样本准确率 (\(> 80\%\)) 未重新训练任何网络参数，直接提取新知识
冻结 MLP 权重迁移 随机初始化 (全新双射) 0.0% (完全随机) 证明必须具备线性叠加结构才能被正确路由
冻结嵌入对照模型 智能初始化 \(W_{\text{stack}}^+\) (全新双射) 0.0% (完全随机) 缺乏线性读出结构导致伪逆重构完全失效

针对多跳设置（\(N=2048, d=256\)），进一步消融了 CoT 与嵌入可学习性的必要性：在 \(k=4\) 跳与 \(R=16\) 极具挑战的设定下，可学习嵌入 + CoT 能够达到 \(99.8\%\) 的高准确率；若去除 CoT，模型在有限嵌入维度下无法承载指数增长的候选空间，准确率迅速崩溃至 \(10\%\) 以下；若冻结输入嵌入，即便引入 CoT，准确率也迅速退化至接近 \(0\%\)。

关键发现

  • 容量缩放法则的一致性：当嵌入可学习时，临界维度严格遵循 \(d_{\min} = \Theta(R \log N)\) 的理论渐近线；而冻结随机嵌入则呈现多项式或线性膨胀 \(d_{\min} = \Theta(NR)\)，证实了几何自组织对参数效率的巨大增益。
  • MLP 职能的彻底解耦：冻结网络后的零样本新知识迁移实验强有力地证明了，优化算法训练出的 MLP 并不是存储主体事实的查找表，而是纯粹的关系滤波算子。
  • CoT 本质是外部内存再路由：多跳消融证实，CoT 在事实推理中的巨大优势根源于将隐状态的连续容量瓶颈转换为离散符号再嵌入，每步自回归都重新激活静态词表的局部几何拓扑。

亮点与洞察

  • 颠覆了将 MLP 视为主体事实键值库的传统视角，首次从理论和经验双向证实了嵌入层内部的“线性超维度叠加 + 浅层门控滤波”即可胜任完备的事实记忆。
  • 严密建立了无 CoT 条件下多跳关系记忆的信息论下界，给出了指数级维度或线性参数爆炸的数学证明，为理解大模型为什么必须依赖思考链提供了全新的几何解释。
  • 提出的智能初始化逆推法（Smart Initialization via Pseudoinverse）为模型的事实编辑与知识无损注入提供了极具启发性的新思路——只需修改主体嵌入向量而无需微调网络权重。

局限与展望

  • 实验设置中将关系视为不相关的离散双射符号，忽略了自然语言知识中普遍存在的语义多义性、层次化包含关系以及长尾频次分布。
  • 理论证明主要针对单层与双层简化 Transformer 架构，真实深层模型中几何叠加与多层注意力头交织时的具体动力学演化仍有待进一步推广。
  • 现实自然语言中大量实体由多个子词（BPE tokens）构成，跨 token 聚合过程中几何叠加的稳健性尚未完全建立统一理论。

相关工作与启发

  • vs Nichani et al. (2025): 该工作基于正交随机嵌入假设，推导出需要 \(\tilde{\Omega}(NR)\) 参数的关联记忆下界且假设属性空间互不相交；本文放宽为更贴合实际的共享属性池，证明了通过嵌入自组织可将维度降至 \(O(R \log N)\)。
  • vs Meng et al. (ROME, 2022) / Geva et al. (2021): 经典知识定位研究强调中间层 MLP 充当知识查找键值对；本文揭示了另一种平行存在的优雅解——事实存储在几何嵌入中，MLP 仅作为通用条件选择器。
  • vs Dugan et al. (2025): 该工作仅关注单关系下的编码器-解码器 MLP 容量；本文将其泛化至多关系交织与多跳推理网络，并系统刻画了深度与容量的权衡。

评分

  • 新颖性: ⭐⭐⭐⭐⭐ 颠覆了事实记忆完全依赖 MLP 键值存储的传统教条，理论与因果实验均极具启发性。
  • 实验充分度: ⭐⭐⭐⭐⭐ 涵盖了受控合成实验、严格因果干预、零样本权重迁移以及 5 种前沿开源开源大模型的实证探测。
  • 写作质量: ⭐⭐⭐⭐⭐ 数学推导严谨自洽，论证逻辑环环相扣，图表与实验设计高度切合理论断言。
  • 价值: ⭐⭐⭐⭐⭐ 对大模型知识存储机制、模型编辑、模型压缩以及 CoT 推理边界具有重要的理论指导价值。

相关论文