Linking In-context Learning in Transformers to Human Episodic Memory

📄 arXiv: 2405.14992v2 📥 PDF

作者: Li Ji-An, Corey Y. Zhou, Marcus K. Benna, Marcelo G. Mattar

分类: cs.CL, cs.LG

发布日期: 2024-05-23 (更新: 2024-10-31)

备注: NeurIPS 2024

💡 一句话要点

揭示Transformer上下文学习与人类情景记忆的关联，发现类CMR机制

🎯 匹配领域: 支柱九：具身大模型 (Embodied Foundation Models)

关键词: Transformer模型 上下文学习 情景记忆 注意力机制 归纳头 CMR模型 大型语言模型

📋 核心要点

  1. Transformer模型缺乏与神经科学的明确联系，特别是自注意力机制与人类认知功能的对应关系有待探索。
  2. 论文的核心思想是Transformer中的归纳头与人类情景记忆的CMR模型存在相似性，它们在上下文学习中扮演重要角色。
  3. 实验结果表明，LLM中存在类CMR头，且其位置和功能与人类记忆偏差相符，消融实验验证了其在上下文学习中的作用。

📝 摘要（中文）

本研究旨在探索人工智能系统与生物智能系统之间的联系，从而揭示通用智能的基本原理。尽管许多人工智能模型都有神经科学对应物，但Transformer模型和自注意力机制在这方面的联系很大程度上缺失。本文研究了交互注意力头与人类情景记忆之间的关系，重点关注有助于Transformer大型语言模型（LLM）进行上下文学习的归纳头。研究表明，归纳头在行为、功能和机制上与人类情景记忆的上下文维护和检索（CMR）模型相似。对预训练LLM的分析表明，类CMR头通常出现在中间层和后期层，在质量上反映了人类的记忆偏差。对类CMR头的消融研究表明，它们在上下文学习中起着因果作用。这些发现揭示了LLM的计算机制与人类记忆之间的相似之处，为这两个研究领域提供了宝贵的见解。

🔬 方法详解

问题定义：现有的大型语言模型（LLM）在上下文学习方面表现出色，但其内部机制与人类认知功能的联系尚不明确。特别是，Transformer模型中的自注意力机制，其工作原理与人类记忆系统是否存在关联，以及这种关联如何影响LLM的性能，是当前研究的痛点。

核心思路：论文的核心思路是Transformer模型中的特定注意力头（归纳头）与人类情景记忆的上下文维护和检索（CMR）模型具有相似性。通过分析这些注意力头的行为、功能和机制，可以揭示LLM的上下文学习能力与人类记忆之间的潜在联系。这种联系有助于理解LLM的内部工作原理，并为改进LLM的设计提供新的思路。

技术框架：该研究主要包括以下几个阶段：1) 分析预训练LLM中注意力头的行为，识别出具有CMR模型特征的注意力头（类CMR头）；2) 比较类CMR头与人类情景记忆在行为和功能上的相似性；3) 通过消融实验，验证类CMR头在上下文学习中的因果作用；4) 分析类CMR头在LLM中的位置和分布，以及它们与人类记忆偏差的关联。

关键创新：该研究最重要的技术创新点在于发现了Transformer模型中的归纳头与人类情景记忆的CMR模型之间的对应关系。这种对应关系不仅揭示了LLM的上下文学习能力与人类记忆之间的潜在联系，也为理解LLM的内部工作原理提供了一个新的视角。此外，通过消融实验验证了类CMR头在上下文学习中的因果作用，进一步证实了这种联系的有效性。

关键设计：研究中使用了预训练的LLM，并分析了其中注意力头的权重和激活。通过比较注意力头的行为与CMR模型的预测，识别出类CMR头。消融实验通过移除或修改类CMR头的权重，观察LLM在上下文学习任务上的性能变化。此外，研究还分析了类CMR头在LLM中的位置和分布，以及它们与人类记忆偏差的关联。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

研究发现，LLM中存在与人类情景记忆的CMR模型相似的注意力头（类CMR头），这些头通常出现在中间层和后期层，且其位置和功能与人类记忆偏差相符。消融实验表明，移除类CMR头会导致LLM在上下文学习任务上的性能下降，证实了其在上下文学习中的因果作用。

🎯 应用场景

该研究成果可应用于提升大型语言模型的性能和可解释性。通过理解LLM与人类记忆的关联，可以设计更高效的上下文学习机制，并改进LLM的训练方法。此外，该研究还有助于开发更具通用性和适应性的AI系统，以及更好地理解人类认知功能。

📄 摘要（原文）

Understanding connections between artificial and biological intelligent systems can reveal fundamental principles of general intelligence. While many artificial intelligence models have a neuroscience counterpart, such connections are largely missing in Transformer models and the self-attention mechanism. Here, we examine the relationship between interacting attention heads and human episodic memory. We focus on induction heads, which contribute to in-context learning in Transformer-based large language models (LLMs). We demonstrate that induction heads are behaviorally, functionally, and mechanistically similar to the contextual maintenance and retrieval (CMR) model of human episodic memory. Our analyses of LLMs pre-trained on extensive text data show that CMR-like heads often emerge in the intermediate and late layers, qualitatively mirroring human memory biases. The ablation of CMR-like heads suggests their causal role in in-context learning. Our findings uncover a parallel between the computational mechanisms of LLMs and human memory, offering valuable insights into both research fields.