Transformer通过几何方式记忆事实,用向量叠加和选择机制替代传统键值存储。
Geometric Factual Recall in Transformers

- 用线性叠加的嵌入编码关系结构,小型MLP作为条件选择器提取信息
- 嵌入维度仅需对数级即可完整记忆随机双射关系
- 模型可零样本迁移至新事实,说明学到的是通用选择机制
Transformer语言模型如何记忆事实关联?传统观点认为内部权重矩阵作为嵌入对之间的关联记忆,参数量随事实数量线性增长。本文提出一种几何记忆机制:学习到的嵌入直接编码关系结构,MLP起质变作用。在单层Transformer需记忆从主体到共享属性集的随机双射的控制实验中,我们证明对数级嵌入维度即足够:主体嵌入以相关属性向量的线性叠加形式编码,小型MLP通过ReLU门控实现关系条件选择,而非键值映射。我们将结果扩展至多跳场景(如“x的妻子的母亲是谁?”),构造出有无链式思考的解法,并给出可证明的容量-深度权衡,辅以匹配的信息论下界。实验表明梯度下降能发现符合预测结构的解。训练后,当主体嵌入重新初始化时,MLP可零样本迁移到全新双射关系,揭示其学习的是通用选择机制而非特定事实集合。
原文摘要 · Abstract (English)
How do transformer language models memorize factual associations? A common view casts internal weight matrices as associative memories over pairs of embeddings, requiring parameter counts that scale linearly with the number of facts. We develop a theoretical and empirical account of an alternative, \emph{geometric} form of memorization in which learned embeddings encode relational structure directly, and the MLP plays a qualitatively different role. In a controlled setting where a single-layer transformer must memorize random bijections from subjects to a shared attribute set, we prove that a logarithmic embedding dimension suffices: subject embeddings encode \emph{linear superpositions} of their associated attribute vectors, and a small MLP acts as a relation-conditioned selector that extracts the relevant attribute via ReLU gating, and not as an associative key-value mapping. We extend these results to the multi-hop setting -- chains of relational queries such as ``Who is the mother of the wife of $x$?'' -- providing constructions with and without chain-of-thought that exhibit a provable capacity-depth tradeoff, complemented by a matching information-theoretic lower bound. Empirically, gradient descent discovers solutions with precisely the predicted structure. Once trained, the MLP transfers zero-shot to entirely new bijections when subject embeddings are appropriately re-initialized, revealing that it has learned a generic selection mechanism rather than memorized any particular set of facts.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。