[Paper Review] Attention Heads of Large Language Models: A Survey

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

Topic Modeling4 citations

TL;DR

This survey proposes a four-stage cognitive framework—Knowledge Recalling, In-Context Identification, Latent Reasoning, and Expression Preparation—to systematically categorize the functions of attention heads in large language models (LLMs). It classifies attention heads by their roles in these stages, reviews modeling-free and modeling-required discovery methods, and outlines evaluation benchmarks, offering a structured foundation for mechanistic interpretability in LLMs with implications for improving model robustness and alignment with human cognition.

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.

Motivation & Objective

• To address the black-box nature of large language models (LLMs) by investigating the functional roles of attention heads.
• To develop a four-stage cognitive framework inspired by human thought processes to classify attention head functions in LLMs.
• To systematize and categorize existing experimental methodologies for discovering specialized attention heads into modeling-free and modeling-required approaches.
• To summarize evaluation benchmarks and highlight limitations in current interpretability research, including lack of generalizability and theoretical grounding.
• To propose future research directions, including robustness analysis, comprehensive interpretability frameworks, and integration with machine psychology.

Proposed method

• The authors propose a four-stage cognitive framework—Knowledge Recalling, In-Context Identification, Latent Reasoning, and Expression Preparation—based on cognitive neuroscience and psychology to map LLM reasoning to human-like thought processes.
• They classify attention heads according to their functional roles within each stage, identifying commonalities and collaborative dynamics across heads.
• The survey categorizes discovery methods into two types: modeling-free (e.g., activation patching, ablation) and modeling-required (e.g., circuit training, causal scrubbing).
• The authors review existing evaluation tasks and benchmarks used to validate head functions, such as the IOI task and color object task.
• They integrate insights from mechanical interpretability and machine psychology to enrich the understanding of LLM internals from anthropomorphic and behavioral perspectives.
• The framework is applied to recent LLMs like LLaMA and GPT, focusing on state-of-the-art research rather than outdated models like BERT.

Experimental results

Research questions

• RQ1 How can attention heads in LLMs be systematically categorized based on their functional roles in a human-like reasoning process?
• RQ2 What are the key differences and trade-offs between modeling-free and modeling-required methods for discovering specialized attention heads?
• RQ3 To what extent do discovered attention head mechanisms generalize across diverse downstream tasks?
• RQ4 How can the collaborative dynamics of multiple attention heads across reasoning stages be modeled and understood?
• RQ5 What role can machine psychology and behavioral analysis play in interpreting and improving LLM mechanisms?

Key findings

• The four-stage cognitive framework successfully maps attention head functions to distinct phases of reasoning, enabling systematic classification of heads across knowledge recall, context identification, latent reasoning, and expression preparation.
• Modeling-free methods such as activation patching and ablation are widely used for head discovery, while modeling-required methods like causal scrubbing offer deeper mechanistic insights but require more complex setup.
• Many discovered circuits, such as those in the IOI and color object tasks, lack validation across broader task distributions, indicating limited generalizability.
• Current research predominantly focuses on isolated head functions, with minimal exploration of collaborative mechanisms among multiple heads, highlighting a key research gap.
• The absence of mathematical proofs for most proposed mechanisms raises concerns about coincidental or non-robust findings, underscoring the need for stronger theoretical validation.
• Future research should prioritize robustness analysis, comprehensive interpretability frameworks, and integration with machine psychology to close the gap between LLMs and human cognition.