[论文解读] Attention Heads of Large Language Models: A Survey

Zifan Zheng, Yezhaohui Wang+6 more|arXiv (Cornell University)|Sep 5, 2024

Topic Modeling被引用 4

一句话总结

本综述提出一个四阶段认知框架——知识回忆、上下文识别、潜在推理和表达准备——以系统化地分类大型语言模型（LLMs）中注意力头的功能。该框架根据注意力头在各阶段的角色进行分类，回顾了无需建模和需要建模的发现方法，并概述了评估基准，为LLMs的机械可解释性提供了结构化基础，对提升模型鲁棒性及与人类认知的一致性具有启示意义。

ABSTRACT

Since the advent of ChatGPT, Large Language Models (LLMs) have excelled in various tasks but remain as black-box systems. Understanding the reasoning bottlenecks of LLMs has become a critical challenge, as these limitations are deeply tied to their internal architecture. Among these, attention heads have emerged as a focal point for investigating the underlying mechanics of LLMs. In this survey, we aim to demystify the internal reasoning processes of LLMs by systematically exploring the roles and mechanisms of attention heads. We first introduce a novel four-stage framework inspired by the human thought process: Knowledge Recalling, In-Context Identification, Latent Reasoning, and Expression Preparation. Using this framework, we comprehensively review existing research to identify and categorize the functions of specific attention heads. Additionally, we analyze the experimental methodologies used to discover these special heads, dividing them into two categories: Modeling-Free and Modeling-Required methods. We further summarize relevant evaluation methods and benchmarks. Finally, we discuss the limitations of current research and propose several potential future directions.

研究动机与目标

• 为解决大型语言模型（LLMs）的黑箱特性，通过研究注意力头的功能角色。
• 基于人类思维过程的启发，提出一个四阶段认知框架，以分类LLMs中注意力头的功能。
• 系统化并分类现有实验方法，用于发现专用注意力头，分为无需建模和需要建模两类方法。
• 总结评估基准，并指出当前可解释性研究中的局限性，包括缺乏泛化能力和理论基础。
• 提出未来研究方向，包括鲁棒性分析、全面的可解释性框架，以及与机器心理学的整合。

提出的方法

• 作者基于认知神经科学和心理学，提出一个四阶段认知框架——知识回忆、上下文识别、潜在推理和表达准备——以将LLM的推理过程映射到类人思维过程。
• 根据注意力头在各阶段的功能角色进行分类，识别头之间的共性及协作动态。
• 该综述将发现方法分为两类：无需建模（如激活修补、消融）和需要建模（如电路训练、因果擦除）。
• 作者回顾了用于验证头功能的现有评估任务和基准，如IOI任务和颜色物体任务。
• 整合机械可解释性和机器心理学的洞见，从拟人化和行为视角丰富对LLM内部机制的理解。
• 该框架应用于近期LLMs（如LLaMA和GPT），聚焦于最先进研究，而非过时模型（如BERT）。

实验结果

研究问题

• RQ1 如何基于类人推理过程中的功能角色，系统化地对LLMs中的注意力头进行分类？
• RQ2 无需建模与需要建模的发现专用注意力头方法之间，关键差异和权衡是什么？
• RQ3 所发现的注意力头机制在多大程度上能泛化到多样化的下游任务？
• RQ4 如何建模和理解跨推理阶段的多个注意力头之间的协作动态？
• RQ5 机器心理学和行为分析在解释和改进LLM机制方面发挥什么作用？

主要发现

• 四阶段认知框架成功地将注意力头功能映射到推理的不同阶段，实现了对知识回忆、上下文识别、潜在推理和表达准备阶段头的系统分类。
• 无需建模方法（如激活修补、消融）广泛用于头的发现，而需要建模方法（如因果擦除）虽需更复杂的设置，但能提供更深层次的机制洞察。
• 许多已发现的电路（如IOI和颜色物体任务中的电路）缺乏在更广泛任务分布上的验证，表明其泛化能力有限。
• 当前研究主要聚焦于单个头的功能，对多个头之间协作机制的探索极少，凸显了关键的研究空白。
• 大多数所提机制缺乏数学证明，引发对偶然发现或非鲁棒结果的担忧，强调了更强理论验证的必要性。
• 未来研究应优先关注鲁棒性分析、全面的可解释性框架，以及与机器心理学的整合，以弥合LLMs与人类认知之间的差距。