Understanding Factual Recall in Transformers via Associative Memories

📄 arXiv: 2412.06538v1 📥 PDF

作者: Eshaan Nichani, Jason D. Lee, Alberto Bietti

分类: cs.LG, cs.CL, cs.IT, stat.ML

发布日期: 2024-12-09

💡 一句话要点

通过联想记忆理解Transformer中的事实性知识回忆能力

🎯 匹配领域: 支柱二：RL算法与架构 (RL & Architecture) 支柱九：具身大模型 (Embodied Foundation Models)

关键词: Transformer 联想记忆 事实性知识回忆 自注意力机制 存储容量

📋 核心要点

  1. 大型语言模型在事实性知识回忆方面表现出色，但对其内部机制的理解仍有挑战。
  2. 该论文提出浅层Transformer可以通过联想记忆的组合实现接近最优的存储容量，并进行了理论证明。
  3. 通过合成的事实性知识回忆任务，验证了Transformer的自注意力机制和MLP层均可作为联想记忆存储事实。

📝 摘要（中文）

大型语言模型在事实性知识回忆方面表现出令人印象深刻的能力。先前的工作表明，在事实性知识回忆任务上训练的Transformer可以以与其参数数量成比例的速度存储信息。本文表明，浅层Transformer可以使用联想记忆的组合来获得接近最优的存储容量。首先证明了线性联想记忆和MLP联想记忆的存储容量都与参数数量呈线性关系。然后，引入了一个合成的事实性知识回忆任务，并证明了具有单层自注意力后接MLP的Transformer，只要自注意力参数或MLP参数的总数（在对数因子范围内）与事实的数量呈线性关系，就可以在该任务上获得100%的准确率。特别地，Transformer可以在使用值矩阵或MLP作为联想记忆来存储事实数据集之间进行权衡。最后，通过分析在事实性知识回忆任务上训练的简化线性注意力模型的梯度流轨迹，补充了这些表达性结果，表明该模型表现出顺序学习行为。

🔬 方法详解

问题定义：论文旨在理解Transformer模型如何存储和回忆事实性知识。现有方法缺乏对Transformer内部机制的深入理解，特别是其存储容量和信息存储方式。Transformer模型通常被视为黑盒，难以解释其事实性知识回忆能力的来源。

核心思路：论文的核心思路是将Transformer模型中的自注意力机制和MLP层视为联想记忆。通过理论分析和实验验证，证明这些模块可以有效地存储和回忆事实性知识，并且其存储容量与参数数量呈线性关系。Transformer可以通过权衡使用自注意力机制或MLP层来存储事实。

技术框架：论文主要研究了具有单层自注意力和MLP的浅层Transformer模型。首先，对线性联想记忆和MLP联想记忆的存储容量进行了理论分析。然后，设计了一个合成的事实性知识回忆任务，用于评估Transformer模型的性能。最后，分析了简化线性注意力模型的梯度流轨迹，以理解模型的学习行为。

关键创新：论文的关键创新在于将Transformer模型中的自注意力机制和MLP层与联想记忆联系起来，并证明了它们可以有效地存储和回忆事实性知识。此外，论文还揭示了Transformer模型在存储事实时可以在自注意力机制和MLP层之间进行权衡。

关键设计：论文设计了一个合成的事实性知识回忆任务，其中包含多个事实，每个事实由一个查询和一个答案组成。Transformer模型的输入是查询，输出是答案。论文使用交叉熵损失函数来训练模型。此外，论文还分析了简化线性注意力模型的梯度流轨迹，以理解模型的学习行为。具体来说，论文研究了权重矩阵的更新过程，并观察到模型表现出顺序学习行为，即先学习一些事实，然后再学习其他事实。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

论文证明了浅层Transformer可以通过联想记忆的组合获得接近最优的存储容量。在合成的事实性知识回忆任务中，只要自注意力参数或MLP参数的总数与事实的数量呈线性关系，Transformer就可以获得100%的准确率。此外，论文还揭示了模型表现出顺序学习行为。

🎯 应用场景

该研究成果有助于更好地理解大型语言模型的事实性知识存储和回忆机制，为改进模型设计、提高知识表示能力和增强模型的可解释性提供理论基础。潜在应用包括知识图谱构建、问答系统、信息检索等领域，并有助于开发更可靠、更可控的AI系统。

📄 摘要（原文）

Large language models have demonstrated an impressive ability to perform factual recall. Prior work has found that transformers trained on factual recall tasks can store information at a rate proportional to their parameter count. In our work, we show that shallow transformers can use a combination of associative memories to obtain such near optimal storage capacity. We begin by proving that the storage capacities of both linear and MLP associative memories scale linearly with parameter count. We next introduce a synthetic factual recall task, and prove that a transformer with a single layer of self-attention followed by an MLP can obtain 100% accuracy on the task whenever either the total number of self-attention parameters or MLP parameters scales (up to log factors) linearly with the number of facts. In particular, the transformer can trade off between using the value matrices or the MLP as an associative memory to store the dataset of facts. We complement these expressivity results with an analysis of the gradient flow trajectory of a simplified linear attention model trained on our factual recall task, where we show that the model exhibits sequential learning behavior.