Self-Attention Limits Working Memory Capacity of Transformer-Based Models
作者: Dongyu Gong, Hantao Zhang
分类: cs.CL, cs.AI, q-bio.NC
发布日期: 2024-09-16 (更新: 2024-11-16)
备注: 10 pages, 12 figures
💡 一句话要点
Transformer自注意力机制限制了其工作记忆容量，影响N-back任务表现
🎯 匹配领域: 支柱九：具身大模型 (Embodied Foundation Models)
关键词: Transformer模型 自注意力机制 工作记忆 N-back任务 注意力熵 认知计算 深度学习 语言模型
📋 核心要点
- 大型语言模型在N-back任务中表现出工作记忆容量限制，但缺乏对其内在机制的解释。
- 该研究假设Transformer的自注意力机制是造成工作记忆容量限制的原因，并进行实验验证。
- 实验发现注意力分散程度与N-back任务难度相关，为改进模型架构提供了新思路。
📝 摘要（中文）
最近关于基于Transformer的大型语言模型(LLMs)的研究揭示了其工作记忆容量的显著限制，类似于在人类行为研究中发现的现象。具体而言，这些模型在N-back任务上的性能随着N的增加而显著下降。然而，对于这种现象产生的原因，仍然缺乏机制上的可解释性。受到行为科学中执行注意理论的启发，我们假设Transformer模型中的自注意力机制可能是其工作记忆容量限制的原因。为了验证这一假设，我们训练了vanilla decoder-only Transformer来执行N-back任务，并发现注意力分数在训练过程中逐渐聚集到N-back的位置，表明该模型通过学习一种策略来掌握任务，即关注当前位置和N-back位置之间的关系。至关重要的是，我们发现注意力分数矩阵的总熵随着N的增加而增加，表明注意力分数的离散可能是N-back任务中观察到的容量限制的原因。我们的发现因此提供了关于注意在人类和人工智能中共同作用的见解。此外，本研究揭示的自注意力机制的局限性可以为未来设计具有增强的工作记忆容量和认知能力的更强大的模型架构提供信息。
🔬 方法详解
问题定义：论文旨在探究Transformer模型在N-back任务中表现出的工作记忆容量限制的原因。现有研究缺乏对这种现象的内在机制的解释，无法指导模型架构的改进。N-back任务要求模型记住N步之前的信息，模拟了人类的工作记忆过程，是评估模型记忆能力的重要手段。现有方法难以解释为何Transformer模型在N增大时性能会显著下降。
核心思路：论文的核心思路是，Transformer模型的自注意力机制是造成其工作记忆容量限制的关键因素。作者认为，自注意力机制在处理N-back任务时，需要对不同位置的信息进行加权平均，当N增大时，注意力需要分散到更多的位置，导致信息丢失和性能下降。通过分析注意力权重的分布，可以揭示模型在处理N-back任务时的策略和局限性。
技术框架：论文采用decoder-only Transformer模型，并使用标准的N-back任务进行训练和测试。模型的输入是序列化的数据，输出是预测的N步之前的信息。训练过程中，使用交叉熵损失函数来优化模型参数。为了分析注意力机制，作者提取了注意力权重矩阵，并计算了其熵值。熵值越高，表示注意力分布越分散。
关键创新：论文的关键创新在于将Transformer模型的工作记忆容量限制与自注意力机制的特性联系起来。通过分析注意力权重的分布，揭示了注意力分散是导致性能下降的原因。这一发现为改进Transformer模型架构，提升其工作记忆容量提供了新的思路。
关键设计：论文的关键设计包括：1) 使用vanilla decoder-only Transformer模型，避免其他因素的干扰；2) 使用标准的N-back任务，便于与其他模型进行比较；3) 计算注意力权重矩阵的熵值，量化注意力分散程度；4) 分析注意力权重在不同位置的分布，揭示模型学习到的策略。
🖼️ 关键图片
📊 实验亮点
实验结果表明，随着N的增加，Transformer模型在N-back任务上的性能显著下降。同时，注意力权重矩阵的熵值也随着N的增加而增加，表明注意力分散程度与任务难度呈正相关。这些结果支持了自注意力机制是工作记忆容量限制的关键因素的假设。
🎯 应用场景
该研究成果可应用于提升大型语言模型在需要长期记忆和推理的任务中的表现，例如对话系统、机器翻译和文本摘要。通过改进自注意力机制或设计新的注意力机制，可以增强模型的工作记忆容量，使其能够处理更复杂的任务。此外，该研究也为理解人类认知机制提供了借鉴。
📄 摘要（原文）
Recent work on Transformer-based large language models (LLMs) has revealed striking limits in their working memory capacity, similar to what has been found in human behavioral studies. Specifically, these models' performance drops significantly on N-back tasks as N increases. However, there is still a lack of mechanistic interpretability as to why this phenomenon would arise. Inspired by the executive attention theory from behavioral sciences, we hypothesize that the self-attention mechanism within Transformer-based models might be responsible for their working memory capacity limits. To test this hypothesis, we train vanilla decoder-only transformers to perform N-back tasks and find that attention scores gradually aggregate to the N-back positions over training, suggesting that the model masters the task by learning a strategy to pay attention to the relationship between the current position and the N-back position. Critically, we find that the total entropy of the attention score matrix increases as N increases, suggesting that the dispersion of attention scores might be the cause of the capacity limit observed in N-back tasks. Our findings thus offer insights into the shared role of attention in both human and artificial intelligence. Moreover, the limitations of the self-attention mechanism revealed in the current study could inform future efforts to design more powerful model architectures with enhanced working memory capacity and cognitive capabilities.