Transformer中的几何事实召回

摘要

Transformer语言模型如何记忆事实关联？常见观点将内部权重矩阵视为嵌入对上的联想记忆，要求参数数量随事实数量线性增长。我们提出一种替代性的几何记忆形式的理论与经验解释：在这种形式中，学习到的嵌入直接编码关系结构，而MLP扮演着性质不同的角色。在单层transformer必须记忆从主体到共享属性集的随机双射这一受控设置中，我们证明对数嵌入维度就足够：主体嵌入编码其关联属性向量的线性叠加，而小型MLP并非作为联想键值映射，而是充当关系条件选择器，通过ReLU门控提取相关属性。我们进一步将这些结果扩展到多跳设置——即关系查询链，例如“x的妻子母亲是谁？”——提供了有无思维链两种构造方法，展示了可证明的容量-深度权衡，并辅以匹配的信息论下界。实验发现，梯度下降能够发现具有精确预测结构的解。当主体嵌入被适当重新初始化时，训练后的MLP能够零样本地迁移到全新的双射上，这表明它学习到的是一种通用选择机制，而非记忆任何特定的事实集合。

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, 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 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.

Transformer中的几何事实召回

Geometric Factual Recall in Transformers

摘要

Support