Cognitive Mirrors: Exploring the Diverse Functional Roles of Attention Heads in LLM Reasoning

2025-12-03

AI Summary

Understanding the intrinsic cognitive functions of attention heads in large language models (LLMs) and their causal impact on reasoning remains an open challenge. Method: We construct CogQA, a multi-step cognitive decomposition dataset, and design interpretable probes alongside chain-of-thought subproblem decomposition to systematically identify “cognitive heads” responsible for retrieval, logical inference, and other reasoning subtasks. Contribution/Results: We discover, for the first time, cross-model functional specialization of attention heads across diverse LLM families—characterized by sparsity, heterogeneous distribution, and hierarchical interaction patterns. Causal intervention experiments (pruning and enhancement) confirm that cognitive heads exert direct influence on reasoning performance: pruning critical heads significantly degrades accuracy, while targeted enhancement yields up to a 12.7% improvement. These findings establish a novel paradigm for interpretable AI and reasoning optimization, bridging mechanistic interpretability with practical model refinement.

Technology Category

• Cognitive Modeling & Cognitive Systems: Conceptual Inference and Reasoning
• Machine Learning: Large Multimodal Models (LMMs)
• Knowledge Representation and Reasoning: Action, Change, and Causality

Application Category

• Semantics and Knowledge: Data modeling to support human-machine intelligence, including LLMs agents, intelligent system behavior, explanations, and user-friendly interactions
• Search and Retrieval-Augmented AI: Large language models for search
• Graph Algorithms and Modeling for the Web: Foundation models and LLMs for Web-related graphs

Abstract

Large language models (LLMs) have achieved state-of-the-art performance in a variety of tasks, but remain largely opaque in terms of their internal mechanisms. Understanding these mechanisms is crucial to improve their reasoning abilities. Drawing inspiration from the interplay between neural processes and human cognition, we propose a novel interpretability framework to systematically analyze the roles and behaviors of attention heads, which are key components of LLMs. We introduce CogQA, a dataset that decomposes complex questions into step-by-step subquestions with a chain-of-thought design, each associated with specific cognitive functions such as retrieval or logical reasoning. By applying a multi-class probing method, we identify the attention heads responsible for these functions. Our analysis across multiple LLM families reveals that attention heads exhibit functional specialization, characterized as cognitive heads. These cognitive heads exhibit several key properties: they are universally sparse, vary in number and distribution across different cognitive functions, and display interactive and hierarchical structures. We further show that cognitive heads play a vital role in reasoning tasks - removing them leads to performance degradation, while augmenting them enhances reasoning accuracy. These insights offer a deeper understanding of LLM reasoning and suggest important implications for model design, training, and fine-tuning strategies.

Problem

• Analyzes attention heads' roles in LLM reasoning
• Identifies functionally specialized cognitive heads in LLMs
• Explores cognitive heads' impact on reasoning performance

Innovation

• Proposed interpretability framework analyzing attention heads
• Introduced CogQA dataset with step-by-step subquestions
• Identified cognitive heads via multi-class probing method