Interpreting Context Look-ups in Transformers: Investigating Attention-MLP Interactions

📄 arXiv: 2402.15055v2 📥 PDF

作者: Clement Neo, Shay B. Cohen, Fazl Barez

分类: cs.CL, cs.AI, cs.LG

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

备注: Accepted to EMNLP 2024 Main Conference

💡 一句话要点

提出一种方法研究Transformer中注意力与MLP的交互

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

关键词: 大型语言模型 注意力机制 多层感知器 模型解释 文本生成 深度学习

📋 核心要点

  1. 现有研究对注意力机制和多层感知器的独立性进行了探讨，但两者的交互作用仍未得到充分理解。
  2. 本研究提出了一种新方法，识别下一个令牌神经元并分析其与注意力头的相互作用，以揭示其在预测中的角色。
  3. 实验结果表明，某些注意力头能够识别特定上下文并激活相应的下游神经元，从而提高了下一个令牌的预测准确性。

📝 摘要（中文）

理解大型语言模型（LLMs）的内部工作机制对于推动其理论基础和实际应用至关重要。尽管注意力机制和多层感知器（MLPs）已被独立研究，但它们之间的交互仍然未被充分探索。本研究调查了注意力头和下一个令牌神经元在LLMs中如何相互作用以预测新词。我们提出了一种方法来识别下一个令牌神经元，找到高度激活它们的提示，并确定负责的上游注意力头。我们的发现揭示了一些注意力头识别与预测令牌相关的特定上下文，并相应地激活下游的令牌预测神经元。这一机制为理解注意力头如何与MLP神经元协同工作以进行下一个令牌预测提供了更深刻的见解。我们的研究为进一步探索LLMs的复杂工作机制及其对文本生成和理解的影响奠定了基础。

🔬 方法详解

问题定义：本论文旨在解决大型语言模型中注意力机制与多层感知器之间交互作用的理解问题。现有方法未能有效揭示这两者如何协同工作以进行下一个令牌预测。

核心思路：论文提出了一种系统的方法，通过识别下一个令牌神经元及其激活的提示，分析注意力头的作用，从而深入理解其交互机制。

技术框架：研究流程包括三个主要模块：首先，识别下一个令牌神经元；其次，找到能够高度激活这些神经元的提示；最后，确定与这些神经元相关的上游注意力头，并生成相应的解释。

关键创新：本研究的创新点在于系统性地揭示了注意力头与下一个令牌神经元之间的具体交互机制，填补了现有研究的空白，提供了新的视角。

关键设计：在方法设计中，采用了特定的激活阈值来识别下一个令牌神经元，并通过自动化生成解释来分析注意力头的活动，确保了结果的可靠性和可重复性。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

实验结果显示，某些注意力头能够有效识别与预测令牌相关的上下文，并相应激活下游神经元，从而显著提高了下一个令牌的预测准确性。具体性能提升幅度尚未提供，但研究为后续优化提供了重要依据。

🎯 应用场景

该研究的潜在应用领域包括自然语言处理、文本生成和对话系统等。通过深入理解注意力机制与MLP的交互，研究者可以优化模型设计，提高文本生成的质量和准确性，推动智能助手等应用的发展。

📄 摘要（原文）

Understanding the inner workings of large language models (LLMs) is crucial for advancing their theoretical foundations and real-world applications. While the attention mechanism and multi-layer perceptrons (MLPs) have been studied independently, their interactions remain largely unexplored. This study investigates how attention heads and next-token neurons interact in LLMs to predict new words. We propose a methodology to identify next-token neurons, find prompts that highly activate them, and determine the upstream attention heads responsible. We then generate and evaluate explanations for the activity of these attention heads in an automated manner. Our findings reveal that some attention heads recognize specific contexts relevant to predicting a token and activate a downstream token-predicting neuron accordingly. This mechanism provides a deeper understanding of how attention heads work with MLP neurons to perform next-token prediction. Our approach offers a foundation for further research into the intricate workings of LLMs and their impact on text generation and understanding.