Cognitive Mirrors: Exploring the Diverse Functional Roles of Attention Heads in LLM Reasoning
Authors: Xueqi Ma
Jun Wang
Yanbei Jiang
Sarah Monazam Erfani
Tongliang Liu
James Bailey
Venue: NeurIPS 2025 poster ()
Keywords: Large language models
interpretability
attention heads
Primary area: deep_learning

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-label 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, and vary in number and distribution across different cognitive functions, and they 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.

Cognitive Mirrors:
Exploring the Diverse Functional Roles of Attention Heads in LLM Reasoning
Xueqi Ma Jun Wang
Yanbei Jiang'
Sarah Monazam Erfani!
Tongliang Liu
James Bailey'
'The University of Melbourne
"The University of Sydney
Ammzoon
(xueqim, jun2, yanbeij]@studen unimelb. edu.
(sarah.erfani, erfami, (aileyj)@unimen edu. au tongliang. liu@sydney. edu.
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 mecha- nisms. Understanding these mechanisms is crucial improve their reasoning abilities. Drawing inspiration from the interplay between neural processes and human cognition, we propose novel interpretability framework systematically analyze the roles and behaviors of attention heads, which are key components of LLMs. We introduce CogQA, 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 apply- ing 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, and vary in number and distribution across different cognitive functions, and they display interactive and hierarchical structures. We further show that cognitive heads play vital role in reasoning tasks-removin them leads to performance degradation, while augmenting them enhances reasoning accuracy. These insights offer deeper understanding of LLM reasoning and suggest important implications for model design, training and fine-tuning strategies. The code is available at httpa/ub..ommssuuuddin///eeeeeee
Introduction
built neural networks that mimic the structure of the human brain, have demonstrated exceptional performance across various natural language processing (NLP) tasks, often exceeding human capabilities. This has sparked growing interest in exploring the potential similarities between the cognitive processes of LLMs and the human brain. Prior studies have demonstrated that LLMs can predict brain responses to natural language [8] 28). indicating a functional alignment between artificial models and biological systems. However, to the best of our knowledge, systematic efforts align reasoning processes between LLMMs and human cognitive agents remain scarce. When solving complex reasoning tasks (e.g.. mathematical multiple- choice question; Figure TT, the human brain engages a network of specialized regions: the frontal lobe recalls relevant knowledge (35] language areas (e.g.. Wernicke's 's and Broca's) support semantic processing
Large language models (LLMs) 14 30 38
20], and the parietal and prefrontal cortices carry out higher-order reasoning [5]
'This work is not related to Amazon.
39th Conference on Neural Information Processing Systems (NcurIPS 2025).

---

Question:
What percent is 120 of 90?
Input
Options:
(A) 133 1/3% (B) 134 1/3 % (c) 135 1/3 % (D) 140 1/3 % (E) 143 1/3 %
Knowledge Recall
Semantic Understanding
Math
Decision Making
is the formula calculate percentage of number relative to another?
What the part and
in 'Whot percent 120 90?
Calculate 120 90
Which option corresponds to 133 1/3%?
1100%
FRENITA
FRINTAL
USE
LODE
LIBE
IEFROCAMPU.
CEPPOCAMPUS
HPPPOCANP.
The formula is (part/ whole) x100%
The part is 120 and the whole is 90.
133 1/3
Option (A)
Figure I: To solve complex question, the human brain engages multiple regions to perform distinct cognitive functions necessary for generating response. We explore whether there are specific attention heads in LLM play functional roles in producing answers.
Analogously, recent research suggests that components within LLMs may also take specialized roles. For example, multi-hess head attention mechanisms in transformers [31] have been found to handle distinct functions, such information retrieval [36] or maintaining answer consistency [17]. pointing toward a form of architectural division of labor. However, most these findings are based on relatively simple tasks [40], leaving open how such specialization operates under complex, multi-step reasoning scenarios. In parallel, prompting techniques like chain-of-then of-thouuhhttt (CoT) 34 have been shown to improve LLM performance by decomposing complex problems into intermediate steps, a strategy reminiscent of human groblem-solving like the example in Figure [I] We hypothesize that such prompting may activate and coordinate specialized components within the model. Thus, analyzing the behavior attention heads under CoT reasoning could contribute insights for deeper understanding of the internal workings of L.MSs and how they process complex tasks.
In this work, we present a novel interpretability framework to systematically analyze the cognitive roles of attention heads during complex reasoning. To facilitate this, introduce Cognitive Ques- tion& Answering (AogQA), a benchmark dataset that decomposes natural language questions into structured subquestions annotated with fine-gre ieeegrained cognitive functions, such retrieval, logical inference, and knowledge recall. Leveraging CogQA, we develop a multi-class probing method to identify and characterize attention heads responsible for distinct cognitive operations within the transformer architecture.
We conduct extensive experiments on three major LLM families, including LLaMA 30 Qwen (38. and [39]. Our results reveal the existence of cognitive heads that consistently exhibit universality, sparsity, and layered functional organization actoss architectures. Further analysis of the correlations among these cognitive heads reveals clear functional clustering, with heads grouping based on cognitive roles, and uncovers hierarchical structure in which lower-level heads modulate higher-level ones oneoooneeeeee mirroring the modular and distributed processing observed in the human cortex 23
Furthermore, we validate functional importance of these heads by showing that their removal degrades performance on complex tasks and leads to specific error patterns, while their enhancement improves reasoning capahilities. Our findings shed light on the structured cognitive architecture embedded in LLMs and open avenues for function- aware model design and analysis.

---

2 CogQA
In this section, we present a detailed account of our benchmark dataset CogQA's construction and key characteristics. Although extensive existing benchmark collections span wide array of NLP tasks, to our knowledge no resource explicitly evaluates LLM reasoning across diverse cognitive functions. To address this, we introduce CogQA, dataset containing 570 main questions and 3,402 subquestions. Each example comprises a question, its answer, and an annotation specifying the cognitive function required for resolution.
2.1 Cognitive Function
To systematically capture the cognitive processes involved in complex reasoning tasks, we categorize cognitive functions into two groups: low-level functions and high-order functions, inspired by estab- lished frameworks in cognitive science 4 [2] Low-level functions primarily involve information retrieval and linguistic analysis, while high-order functions engage more abstract reasoning, problem- solving, and meaiii..-making.. Detailed descriptions of these cognitive functions provided in Appendix(AA4 A.4
The low-level cognitive functions include:
Retrieval: locating relevant information from an external source or prior context. Knowledge Recall: accessing stored factual or procedural knowledge from memory.
Semantic Understanding: interpreting the meaning of words, phrases, or concepts.
Syntactic Understanding: analyzing the grammatical structure of a sentence.
The high-order cognitive functions include:
Mathematical Calculation: performing arithmetic or numerical operations.
Logical Reasoning: drawing conclusions based on formal logical relationships.
Inference: deriving implicit information that is not directly stated.
Decision-Making: selecting the best outcome among alternatives based on reasoning.
This categorization reflects natural progression from basic information processing complex cognitive integration. Both the human brain and LLMs encompass a wide range of functional modules. Our focus in this work is specifically on reasoning- related cognitive functions. By identifying and organizing these eight core reasoning functions, we can more clearly examine how LLMs handle different types of thinking steps, in a way that is both systematic and easy interpret.
2.2 Data Collections
Based on categorization of cognitive functions, sampled 750 diverse questions from NLP reasoning benchmarks, selecting 150 examples from each of AQuA [18] CREAK (24), ECQA (2) e-SNLI [7), and GSM8K [1I]. These datasets cover range of reasoning types, including logical, mathematical, and commonsense reasoning. Using the CoT paradigm, we prompted GPT-4o [16] to decompose each question into subquestions, each targeting a single cognitive function. The prompt encourages structured, step-by-step reasoning, with each subquestion being clear, answerable, and sequentially dependent. This yields a of aablguestioo-answer-oogiiiiv function (subQAC) triples for each QA pair: subQACs = {(a, a; G)) where each contains subquestion qi its concise answer dg and the corresponding cognitive function label q. The prompt for generating subquestions and examples are list Appendix A. and Appendix A.6. respectively.
2.3 Data Filtering and Annotation
Recent advances have made increasingly feasible to use LLMs for dataset construction, owing to their strong reasoning abilities and capacity to generate high-quality annotations at scale [3 Al- though our dataset constructed automatically using an LLM to reduce manual effort, we implement a strict two-stage human verification pipeline to ensure data quality and mitigate hallucinations. In the first stage, three expert annotators independently assess whether the subquestions are logically structured and align with natural human reasoning. QA pairs with inconsistent or incoherent decom- positions are filtered out. In the second stage, annotators verify and, necessary, relabel the cognitive function associated with each subquestion ensure alignment with the intended mental process.
3

---

Finally, we validate the subanswers by cross-checking them using the GPT-o4-min model [25]. followed by human adjudication where discrepancies arise. Details of the annotation process and rubric can be found in Appendix A.5 This multi-step filtering ensures that each retained subQAC
triple reflects coherent, interpretable reasoning step grounded in core cognitive functions. After this refinement, our final dataset contains 570 main QA and 3,402 validated subQAC triplets.
3 Cognitive Function Detections
Given the CogQA dataset, we aim to identify which attention heads in LLMs are associated with specific cognitive functions. We adopt a probing-bart based framework, widely used interpretability technique in which an auxiliary classifier is trained to predict properties from intermediate model representations [ 6 [9). We frame this multi-class classification task: for each cognitively extract head activations (see Section 3.1. train classifier and compute importance scores to identify contributing heads (see Section 3.2). Unlike prior work focusing on a single-class, our method captures many-to-many relationships between heads and functions, enabling a more detailed analysis of functional specialization and overlap compared to prior single-class approaches.
annotated subquestion,
3.1 Head Feature Extraction
Given large language model M, we generate an answer aM for each subquestion 9, derived from a main question Qt To support coherent multi-step reasoning, include preceding subquestions and their answers as contextual input, emulating the incremental reasoning process observed human cognition.
During inference, input tokens are embedded and processed through successive transformer layers. At each layer, attention and feedforward operations update the residual stream, which is ultimately decoded into token predictions. For each generated token i, we extract attention head outputs X, {x]" (i) |-= 1,..,.., L. 1,.... .M) across all layers, where x;"" denotes the value vector from the m-th head in layer projected into the residual stream, with M the number of heads per layer and L the total number of layers.
Let N₂ denote the number of tokens in the generated answer az To isolate semantically informative content relevant reasoning, we select the top-k most important tokens. determined by prompting GPT-o4-mini [25] (skilled in reasoning), yielding an index set Ia with I&l = k (Top-k (k 5) token examples
in Appendix I0). For each index j Ig. extract the corresponding attention head activations Xj. and compute the averaged activation feature for the m-th head in layer as E = Ee jEI, z!" (i). This results in a full set of head-level features X 1.... L, m = 1,...,, M].
{
incorporate layer-wise information by computing the average activation T; = A म. for each layer. We then augment each head-level vector with its corresponding layer summary, resulting in enriched features zmm' [r- Tt.. For each subQA triplet (4:, ai, ci), the final input to the probing classifier is given by { 1, ... L, =1...... 1.......0. M).
Given prior findings suggesting that cognitive functions may vary by layer depth 40.
3.2
Heads Importance
For the CogQA dataset with N subQA pairs, we collect all activations to construct the probing dataset:
{ D Prooee = (aa)" e) 1=1 ,L LI, m E (L.
(I)
We split the dataset into training and validation sets with a 4:1 ratio. Each attention head feature is first passed through a trainable linear projection for dimensionality reduction, followed by two-layer MLP that performs multi-class classification over cognitive functions (training details are
We include ablation study in Appendix A.9t analyze the impact of using alternatiee token positions.

---

Knowledge Recall
Retrieval
Logical Reasoning
Decislon-mmaking
0.0030+
0.025
0.0020
Hoads Impo PAATEE
Layer
Somantic Understanding
12
24
20
Syntactic Understanding
12
24
30
12
244
20
Math Calculation
12
24
30
0.0015
0.0010
0.0006
Head
The existence of cognitive heads in Llama3.1 1-8B-instruch responsible for eight distinct functions in complex reasoning tasks. The x-axis represents the head index, while the y-axis indicates the layer index.
Figure 2:
provided in Appendix A.3). To interpret the contribution of individual heads to each function, we use gradient-based attribution method. Specifically, for each function class c, we compute the contribution of each head feature via the gradientx activatioo technique:
Dge Or,
1/00 ((a.c)-Doote Poote
Zz
where žj the j-th head input feature, and De the classifier's predicted logit for class e. This yields an importance score for each attention head with respect to each cognitive function. We aggregate the scores into matrix I E (L-MM)) where each row corresponds to function class and each column specific head in specific layer.
We hypothesize that attention heads with higher importance scores contribute more significantly to each cognitive function. By ranking heads according to their importance, can identify which heads and layers are specialized for specific functions. Subsequent targeted interventions on these heads validate the effectiveness of this approach.
4 Experiments
We conduct series of experiments three LLM families across various model scales, including LLaMA (30 -llama3.--B--nssnnn and Llamma3.-3B-instruct) Qwen (38) (Qwen3-8B and Qwen3- 4B), and Yi 39 (Yi1.5-9B and Yi1.5-6B). Our goal is to identify cognitive attention heads associated with specific reasoning functions and evaluate their roles via targeted interventions. By selectively masking these heads, we assess their functional significance in supporting downstream performance. We evaluate our method in terms of functional alignment, consistency across models, and causal impact on reasoning tasks. Results confirm the existence of sparse, snuction-specifi heads and highlight their critical contribution to structured cognitive processing within LLMs.
4.1 Properties of Cognitive Heads
Our analysis reveals that cognitive head importance large language models exhibits three key properties: sparsity and universality, and layered functional organization. To illustrate these characteristics, we present the heatmap of attention head importance scores across eight cognitive functions in Llama3.1 1-8.--8--iim 8B-instruct (Figure 2).
5

---

Table 1: Intervention results (%) of cognitive heads vs. random heads across cognitive functions: Retrieval, Knowledge Recall, Semantic Understanding, Syntax Understanding, Math Calculation, Inference, Logic Reasoning, and Decision Making. Lower values indicate more effective intervention outcomes, suggesting that the corresponding heads play a greater role in the cognitive function.
Information Extraction and Analycis Functions
Higher-Order Processing Functions
Model
Inter_Heeud
Retrieval
Recall
Semantic:
Syntectic
Math
Inference
Logic
Decision
comeet
acc
comet
acc
comett
acc
comnet
acc
comeet
nec
cornet
acc
comett
acc
comeet
aCc
L1ama3.1-8B
random
90.83
84.71
87.85
83.84
91,44
97.50
87.81
66.17
94.25
83.08
91.90
70.18
91.39
54.69
97.64
90.91
cognitive
44.96
8.24
56.93
38.38
81.98
75.00
69.20
40.00
87.81
66.17
76.65
52.63
52.07
4.69
56.02
4.55
Llama332-3B
random
87.89
86.47
76.35
68.69
90.54
90.00
75.82
40.00
94.98
69.65
95.66
85.96
92.75
76.56
93.30
81.82
cognitive
49.47
17.06
13.13
10.00
43.62
92.01
53.60
7.02
0.00
0.00
3-B
cognitive
random
92.81
59.19
38.24
75.29
89.90
30.30
53.54
92.73
42.50
88.60
80,00
89
92.66
60.20
94.45
72.77
24.56
94.15
20.31
96.52
31,82
Qwen3-4B
random
94.17
84.71
84.61
77.78
86,91
77.50
98.15
80.00
87. 17.15
44.78
96.89
87.72
92.00
75.00
94.79
72.73
cognitive
80.13
64.71
63.10
35.35
65.95
60.00
46.25
0.00
82.40
46.27
84.88
64,91
82.79
39.06
45.49
13:64
Yi-1,5-9B
random
86.83
79.41
82.02
54.55
77.40
35.00
81.53
60.00
76.04
36.32
89
83
36.84
87.53
42. 19
86.27
63.64
cognitive
52.76
21.76
45.99
9.09
47.25
2.50
48.10
40.00
54.22
16.92
52.41
15.79
82.75
26.56
62.85
18.18
Yi-1.5-6B
random
80.64
69.41
68.82
38.38
77,83
55,00
69.61
60.00
73.33
43.78
77.71
22.81
81.65
29.69
88.54
72.73
cognitive
49,90
15.29
68.23
41.41
49,54
2.50
42.92
0,00
76.64
43.78
68.53
14.04
44.94
0.00
86,28
50,00
TogK Anmuracy
Random Accuracy
ToyK Comet
Random Comet
Retrieval
Knowledge Recall
Math Calculation
Inference
1.0
Score
0.6
0.4
0.2
0.0
32 Masked Heads
Masked Heads
Masked Heads
# Masked Heads
Figure 3: The performance of 1-aama..--B-insse by masking out top K cognitive heads vs K random heads on retrieval, knowledge recall, math calculation, and inference.
Sparsity and Universality: As shown in Figure 2 each cognitive function activates only a small number of high-importance attention heads, revealing strikingly sparse pattern. In Llama3.1 1-8B- instruct, fewer than 7% of all heads have importance scores above 0.001 across the eight functions, suggesting that only compact subset of heads meaningfully contribute to task performance. This sparsity is not uniform: Retrieval contains the highest proportion of salient heads (6.45% exceeding 0.01), while Inference has the fewest (3.42%). These results highlight that LLMs rely on highly specialized, localized components for different cognitive abilities. Importantly, we observe that this sparse functional organization is consistent across different model architectures and sizes. Additional heatmaps for five other models are provided in Appendix A. supporting the universality of this phenomenon.
Layered Functional Organization: In addition to sparsity, attention heads show a structured distribution across model layers. Retrieval-related heads cluster primarily in the middle layers, while math-related heads appear more frequently in higher layers. This structured, task-dependent localization points an emergent modular organization, where different layers support distinct cognitive operations. Further, we identify cognitive heads by selecting those before the elbow point of each function's 's descending importance curve (Appendix A.2). and find notable variation in head counts across functions (Appendix A.8). For example, in the LLaMA family, mathematical calculation requires fewer heads (59 in Llama3.1 -Sn-stsruttt 35 in draws on substantially more (139 and 98, respectively), reflecting differences in representational and computational complexity.
while inference
4.2 Functional Contributions of Cognitive Heads
After identifying the cognitive heads associated with each function, we examine their functional roles by evaluating the model's behavior on the CogQA test set under targeted interventions. We perform

---

Table Intervention results (%) of different cognitive heads and random heads across Retrieval and
Knowledge Recall functions.
Model
Inter Head
Retrieval (comet)
Retricval (acc)
Recall (comet)
Recall (acc)
Llama3. 1-BB
random
90.83
84.71
87.85
83.84
Llama3. 1omma3.1-8B 1-8B
retrieval
44.96
8.24
72.05
33,33
Llama3. 1-8B
recall
86.79
75.29
56.93
38.38
Owen3-8B
random
92.81
75.29
89.90
53.54
Qwen3-8B
retrieval
59.19
38.24
79.26
57.58
Qwen3-8B
recall
83,31
71.18
64.81
30.30
head ablation by scaling the output of a specific attention head with small factor (e.g., 0.001), effectively suppressing its contribution:
W Vdz/n
mask
Softmax
Specifically, compare model performance when masking identified cognitive heads versus masking an equal number of randomly selected heads. To quantify the impact of masking, we use several standard evaluation metrics including COMET [27]. BLEU 26), ROUGE [9. and semantic similarity to compare the model's 's outputs before and after intervention. We define an output as unaffected if the BLEU score exceeds 0.8, or either the ROUGE or semantic similarity scores surpass 0.6, and compute accuracy accordingly.
As shown in Table masking cognitive heads leads significant decline in performance, whereas masking an equal number of random heads results in only marginal degradation across all LLMs. In some cases, masking the identified cognitive heads causes the accuracy drop zero, indicating that the model cannot execute the corresponding function without them. This sharp contrast highlights the essential role cognitive heads play in enabling specific reasoning capabilities. To further validate the functional specialization, we conduct experiments where we mask the retrieval heads during the evaluation of knowledge recall (Recall), and conversely, mask knowledge recall heads during the evaluation of retrieval performance. The results in Table 2 show that masking the corresponding cognitive heads causes a significantly larger performance drop than masking others.
We further investigate the performance of model under different numbers of masked attention heads. As shown in Figure 3 increasing the number of randomly masked heads has minimal impact on In contrast, masking cognitive heads results in
overall performance of significant drop in performance across various functions. Notably, masking heads associated with Retrieval and Knowledge Recall causes pronounced degradation in their respective functions, whereas functions such as Math Calculation and Inference exhibit more resilience. This suggests that certain cognitive functions depend more heavily on specific, distinguishable attention heads, while others are distributed more broadly across the model.
4.3 Relationship Among Cognitive Heads
While cognitive heads are specialized for distinet functions, understanding their relationships is crucial for revealing how complex reasoning emerges from their cooperation.
Functional Clustering: Inspired by neuroscience findings that related cognitive functions localize in overlapping brain regions (c.g., prefrontal cortex for reasoning and inference (5)), we investigate whether LLM attention heads show similar patterns. We rank each head's 's importance across eight cognitive functions, form ranking vectors, and apply principal component analysis (PCA) to visualize their organization (Figure 4T. The results reveal clear clustering: heads linked to reasoning, inference, and decision -aaiin group closely, while those related to mathematical computation form distinct cluster in Llama and Qwen, and lie adjacent to reasoning heads in Yi. Lower-level functions also show moderate clustering. These patterns suggest a modular functional architecture in LLMs akin to that in the human brain.
Hierarchical Structure: Human problem solving often involves hierarchical reasoning, where lower-level functions such retrieval and comprehension support higher-level inference and decision-
7

---

Mathe
Math
Decision,
elggical
Logial
Info
Math
PC
Semantic
Retreval
PC
Decision
Infereng
Logicl
PC
Retrieval
Syntactic
Sema ntic
Retrieval
Recal
Rocall
Semantic
Sybeetic
PC 1
PC 1
PC 1
(a) Alama3.1-88
(b) Qwen3-4B
(c) Yi-1.5-6B
Figure 4: PCA visualization of the 8 function heads' clustering in three models.
Table 3: Study the influence of low-level cognitive heads for high-order function Llama3.1-80
instruct. Accuracy is measured based on BLEU, ROUGE, and semantic similarity scores.
Retrieval Knowledge Semantic Syntactic
Math
Inference
Logical
Decision
0.00
0.00
0.00
0.00
100
100
100
100
/
0.00
0.00
0.00
0.00
100
100
100
100
V
X
66.67
88.24
93.10
57.14
33.33
42.8666
76.92
100
100
23.08
0.00
0.00
making. The CogQA dataset captures this structure through subquestions progressing from simple information extraction complex reasoning. We test if LLMs reflect this hierarchy by masking attention heads tied to carly-stage functions and measuring the effect on later tasks. For instance, to assess how Retrieval affects Math Calculation, we suppress Retrieval-related heads throughout the subquestions. Answers from earlier Retrieval are used as priors for later math reasoning, allowing us observe how disrupting low-level functions can propagate and impair higher-level reasoning along the chain. As Table ssoos,, masking retrieval or knowledge recall heads causes significant performance drops in subsequent dacision-making steps, whereas masking syntactic understanding heads has minimal impact. This provides evidence for an emergent hierarchical organization in LLMs, where foundational cognitive functions underpin advanced reasoning.
4.4 Influence of Cognitive Heads Downstream Tasks
In this section, we investigate how cognitive heads influence downstream tasks through both negative interventions (masking out cognitive funetion heads) and positive interventions (shifting heads toward specific functions). We conduct experiments on two tasks: math task using 100 GSM8K samples (GSMSK_ 100) and a retrieval task with 49 samples from an extractive QA dataset. The Extractive_ QA pairs are generated by GPT-4o, with answers extracted directly from the source paragraph.
Negative Intervention: We perform negative intervention by masking corresponding cognitive heads (Math Calculation heads for GSMSK 100 and Retrieval heads for Extractive _A)). effectively suppressing their activations. As shown in Table this causes significant performance drops across models, confirming these heads' functional roles. Notably, after masking, performance converges to similarly low level across different I.IMss, regardless of model size or original accuracy. This expected, as the crucial cognitive heads responsible for specific functions are disabled, making difficult for the model to arrive at correct answers.
For math, the remaining 30% accuracy likely stems from factors: (1) memorized answers in the base model, and (2) simple questions not requiring actual computation. For retrieval, masking Retrieval heads almost completely abolishes the model's retrieval ability across scales. This indicates that cognitive functions are indeed localized in subset of heads, and masking them leads to a systematic degradation, irrespective of model capacity. The negative intervention example further shows that, masking the Math Calculation heads leads to errors in arithmetic tasks, while retrieval and language functions remain largely unaffected. This confirms that these cognitive heads are crucial for specific functions and highlights the robustness and generalizability of our method.

---

Positive Intervention: We calculate the activation directions of different cognitive functions using the CogQA dataset. For each function, the activation direction of head at layer and index h is computed as:
dir D. [i] Eig [x]
(4)
where 2f (i) denotes the activation of head at layer I and index ht and Donnect and D Dincorret represent the sets of samples answered correctly and incorrectly, respectively. Then we estimate the standard deviation of activations [7 along the cognitive function direction to be of, and shift original head
activation as xf (i) + x (i) + aof dirř', where a is a parameter.
The experimental results in Table show that enhancing the activation of retrieval heads along their corresponding functional directions leads to improved performance on the retrieval task. For the math task, questions often engage multiple cognitive functions. We observe cases where computations correct but semantic understanding fails. By positively intervening on either Math Calculation or Semantic Understanding heads, performance consistently improves, highlighting the collaborative nature of cognitive functions in complex reasoning and the need for joint interventions a direction leave for future work.
The example below illustrates a positive intervention on Math Calculation heads in Qwen3-4B: the model initially produces incorrect answer (red) but generates the correct one (blue) after intervention. This demonstrates that adjusting cognitive head activations along their functional directions reliably enhances performance on complex reasoning tasks.
Negative Intervention Case
Query: Question: Solve -42%r 27 1167 and 130 + 4*c = 372 for r.
Original Output: Okay, so need to solve this system of equations: -42r + 27c = 1167 and 130r + 4c 372 First equation: -42*4 + 27*(-37) -168 999 -1167 Correct. Second equation: 130#4 4*(-37) 520 148 372. Correct. Yes, that's right. So everything checks out. So think that's the correct answer. Therefore, r is 4. Negative Intervention (Math) Output: ... To climinate equations to be the same (or negatives). The coefficients are 27 and 4. The least common multiple of 27 and 4 is 100? Wait, 27 is 3 and 4 is 2 LCM is 2 3 = 4 27 = 100. So if can make the coefficients of с to be 100 and 100, or 27 and 4. Wait, but 27 and 4. So I multiply the first equation by and the second equation by 27, then the coefficients of would be 27*4 = 100 and 4*27 = 100. Then, if subtract the equations, the c terms would be 100 and 100, but since they are both positive, subtracting the equations would eliminate
need the coefficients of in both
Positive Intervention Case
Query: On school trip to the seashore, Alan and his friends collected shells. Alan collected four times as many shells Ben did. Ben got late start and only collected a third of what Laurie did. If Laurie collected 36 shells how many did Alan collect?
Original Output: Alan collected four times as many shells as Ben did, and Ben collected third of what Laurie did. Since Laurie collected 36 shells, Ben collected 36 / 3 12 shells. Therefore, Alan collected * 12 48 shells. However, the question asks for the number of shells Alan collected, and it is given that Alan collected four times as many shells as Ben did. Therefore, the correct answer is 4 12 = 48 4/3 72.
Positive Intervention (Math) Output S.nce Ben collected third of what Lauric did, he collected 36 / 3 12 shells. Alan collected four times as many shells as Ben did, so he
collected 4 12 48 shells....Th correct answer is 48
5 Related Works
Neural Networks and the Brain Neural networks have long been studied as computational models of the brain, with early work linking artificial neurons to the biological mechanisms perception and learning [9] Convolutional neural networks (CNNs), in particular, have been shown to capture representations similar those in the visual cortex [37] and more recent studies suggest that the functional modularity observed in deep networks gives rise to brain-like specialization [3] in vision

---

Table 4: The LLM performance on GSM8k 100 and Extractive _Q by Negative and Positive Intervention (Inter.). Math accuracy (%) is measured by exact answer match; Extractive _QA accuracy checks if the original paragraph answer appears in the response.
Dataset
Method
Head
LJama3.1-8B
LJama3.2-3B
Qwen3-8B
Qwen3-4B
Buse
$7.14
36.73
57.14
51.02
Extractive QA
Negative Inter.
Retrieval Retrieval
0
14.29
12.24
Positive Inter.
63.26
44.90
61.22
69.38
Base
82
64
94
91
Negative Inter.
Math
34
GSMBK_ 100
Positive Inter.
Math
84
66
94
Positive Inter.
Semantic
84
65
94
93
task. More recently, LLMs have exhibited striking parallels with buman brain activity during language processing. In particular, transformer-based models, such GPT-2. produce internal representations that align with neural responses in language-selectnn seleeteveeeetvee brain regions [8] 28) However, prior work mostly focuses on perception and language representations, with limited study on higher-level cognitive functions like reasoning. We instead analyze LLMs' behavior in complex reasoning tasks to explore their alignment with human cognitive functions and functional specialization.
Functional Specialization of Attention Heads Recent years have witnessed growing interest in understanding the functional roles of attention heads in Transformer-based models, forming core component of mechanistic interpretabiliity research. Early work by [10] demonstrated that individual heads BERT capture specific linguistic phenomena such as syntactic dependencies and coreference, indicating a degree of functional specialization. Building on this, [32] proposed a pruning- based approach to identify important heads by measuring their contribution to downstream performance, showing that many heads are redundant. Subsequent studies extended this analysis to decoder-only large language models (LL.s). 21 explored functional decomposition in such models, leading to the identification of distinct attention heads responsible for tasks such as pattern induction 22 truthfulness [17,, information retrieval (36), and safety alignment [4]]. For broader survey, see [40]. Despite these advances, most prior work focuses on isolated heads and evaluates them in relatively simple or synthetic tasks. In contrast, we investigate functionally specialized heads under more complex reasoning settings by aligning attention head behavior with human cognitive functions.
6 Limitations and Future works
While our study provides an initial framework for analyzing the cognitive functions of attention heads, several limitations remain. We examine only eight predefined functions, which, though representative, may not cover the full range of LLM capabilities; future work could expand this taxonomy with finer- grained functions. Each CogQA subquestion is labeled with a single cognitive function, though real reasoning may involve multiple functions. Likewise, we assume one head corresponds to one function, even though head may support multiple context-dependdnt roles. These complexities fully addressed in our current framework. Excluding subquestions with incorrect subanswers could improve multi-class probing, and further investigation is needed to understand heads serving multiple functions. While our study focuses on analysis rather than application, identifying cognitively relevant heads could inform model design, including dynamic head activation, improved saain-of-thoughh prompting, targeted fine-tuning, or modular architectures directions we leave for future exploration.
not
7 Conclusions
We propose an interpretability framework that connects attention heads in large language models (LLMs) to human cognitive functions involved in reasoning. To support this, we introduce CogQA, cognitively grounded dataset, along with a multi-class classification approach identify specialized heads associated with specific reasoning tasks. Our analysis across multiple LLM families and scales demonstrates that attention heads exhibit universality, sparsity, intrinsic roles, and dynamic, hierarchical organization. These findings indicate that L.Ms internally organize reasoning processes manner akin to human cognition, laying the groundwork for more interpretable and cognitively informed language models.
10

---

Acknowledgments
This work is partially supported by the following Australian Research Council (ARC) projects: FT220100318, DP220102121, LP220100527, 122200017,,,,,,,2 LP22020099, DP230101534.
References
[1] Josh Achiam, Steven Adler, Sandhini Agarwal, Lama Ahmad, Ilge Akkaya, Florencia Leoni Aleman, Diogo Almeida, Janko Altenschmidt, Sam Altman, Shyamal Anadkat, et al. Gpt-4 technical report. arXiv preprint arXiv:2303.08774, 2023.
[2] Shourya Aggarwal, Divyanshu Mandowara, Vishwajeet Agrawal, Dinesh Khandelwal, Parag Singla, and Dinesh Garg. Explanations for CommonsenseeQ:: New Dataset and Models. In Proceedings of the 59th Anmual Meeting of the Association for Computational Linguistics and the 11th International Joint Conference Natural Language Processing (Volume 1: Long Papers), Online, 2021. Association for Computational Linguistics.
[3] Guillaume Alain and Yoshua Bengio. Understanding intermediate layers using linear classifier probes. arXiv preprint arXiv:1610.0166 2016.
[4] John R Anderson. Rules of the mind. Psychology Press, 2014.
[5] Lawrence W Barsalou. Cognitive psychology: An overview for cognitive scientists. Psychology Press, 2014.
[6] Yonatan Belinkov. Probing classifiers: Promises, shortcomings, and advances. Computational Linguistics, 48(1):207-219 2022.
[7] Oana-Maria Camburu, Tim Rocktäschel, Thomas Lukasiewicz, and Phil Blunsom. e-snli: Natural language inference with natural language explanations. In Samy Bengio, Hanna M. Wallach, Hugo Larochelle, Kristen Grauman, Nicolò Cesa-Bianchi, and Roman Garnett, edi- tors, Advances in Neural Information Processing Systems 31: Annual Conference on Neural Information Processing Systems 2018, NeurIPS 2018, December 3-8, 2018, Montréal, Canada, pages
2018.
[8] Charlotte Caucheteux, Alexandre Gramfort, and Jean-Rémi King. Deep language algorithms predict semantic comprehension from brain activity. Scientific reports, 12(1):16327, 2022.
[9] Lin Chin- Rouge: A package for automatic evaluation of summaries. In Proceedings of the Workshop on Text Summarization Branches Out, 2004, 2004.
[10] Kevin Clark, Urvashi Khandelwal, Omer Levy, and Christopher D Manning. What does bert look at? an analysis of bert's attention. arXiv preprini 1rXiv:1906.04:: 2019.
[11] Karl Cobbe, Vineet Kosaraju, Mohammad Bavarian, Mark Chen, Heewoo Jun, Lukasz Kaiser, Matthias Plappert, Jerry Tworek, Jacob Hilton, Reiichiro Nakano, Christopher Hesse, and John Schulman. Training verifiers to solve math word problems. CoRR, abs/21 14168, 2021.
[12] Adele Diamond. Executive functions. Annual review of psychology, 64(1))135-168, 2013.
[13] Katharina Dobs, Julio Martinez, Alexander JE Kell, and Nancy Kanwisher. Brain-like func- tional specialization emerges spontaneously in deep neural networks. Science advances, 8(11):eab18913 2022.
[14] Aaron Grattafiori, Abhimanyu Dubey, Abhinav Jauhri, Abhinav Pandey, Abhishek Kadian, Ahmad Al-Dahle, Aiesha Letman, Akhil Mathur, Alan Schelten, Alex Vaughan, et al. The llama 3 herd of models. arXiv preprint arxiv:2407.2178. 2024.
[15] Edward M Hubbard, Manuela Piazza, Philippe Pinel, and Stanislas Dehacne. Interactions between number and space in parietal cortex. Nature reviews neuroscience, 6(6):435-448, 2005.
[16] Aaron Hurst, Adam Lerer, Adam P Goucher, Adam Perelman, Aditya Ramesh, Aidan Clark, AJ Ostrow, Akila Welihinda, Alan Hayes, Alec Radford, et al. Gpt-4o system card. arXiv preprint arXiv:2410.21276 2024.
11

---

[17] Kenneth Li, Oam Patel, Fernanda Viégas, Hanspeter Pfister, and Martin Wattenberg. Inference- time intervention: Eliciting truthful answers from a language model. Advances in Neural Information Processing Systems, (5:41451-41530 2023.
[18] Wang Ling, Dani Yogatama, Chris Dyer, and Phil Blunsom. Program induction by rationalc generation: Learning solve and explain algebraic word problems. In ACL, 2017.
[19] Warren S. McCulloch and Walter H. Pitts. A logical calculus of the ideas immanent in nervous activity. In Margaret A. Boden, editor, The Philosophy of Artificial Intelligence, Oxford readings philosophy, pages 22-39. Oxford University Press, 1990.
[20] Patric Meyer, Axel Mecklinger, Thomas Grunwald, Juergen Fell, Christian E Elger, and An- gela D Friederici. Language processing within the human medial temporal lobe. Hippocampus, 15(4):451-455) 2005.
[21] Paul Michel, Omer Levy, and Graham Neubig. Are sixteen heads really better than one? Advances in neural information processing systems, 32, 2019.
[22] Catherine Olsson, Nelson Elhage, Neel Nanda, Nicholas Joseph, Nova DasSarma, Tom Henighan, Ben Mann, Amanda Askell, Yuntao Bai, Anna Chen, et al. In-context learning and induction heads. arXiv preprint arXiv:2209. 1899 2022.
[23] Yumie Ono, Xian Zhang. J Adam Noah, Swethasri Dravida, and Joy Hirsch. Bidirectional connectivity between broca's 's area and wernicke's 's during interactive verbal communication. Brain connectivity, 12(3):210-222, 2022.
[24] Yasumasa Onoe, Michael J. Q. Zhang, Hunsol Choi, and Greg Durrett. CREAK: A dataset for commonsense reasoning over entity knowledge. In Proceedings of the Neural Information Processing Systems Track on Datasets and Benchmarks I, NeurlPS Datasets and Benchmarks 2021, December 2021, virtual, 2021.
[25] OpenAL. Introducing o3 and 04 mini. https://opeaai.com/iii introducing-o3-anee -4---a-n--4--ii10 o4-mini/ 2024. Accessed: 2025-05-05.
[26] Kishore Papineni, Salim Roukos, Todd Ward, and Wei-Jing Zhu. Bleu: a method for automatic evaluation of machine translation. In Proceedings of the 40th annual meeting of the Association for Computational Linguistics, pages 31 111-318, 2002.
[27] Ricardo Rei, Craig Stewart, Ana Farinha, and Alon Lavie. Comet: A neural framework for mt evaluation. arXiv preprint arXiv:2009.00023 2020.
[28] Martin Schrimpf, Idan Asher Blank, Greta Tuckute, Carina Kauf, Eghbal A Hosseini, Nancy Kanwisher, Joshua B Tenenbaum, and Evelina Fedorenko. The neural architecture of language: Integrative modeling converges predictive processing. Proceedings of the National Academy of Sciences, 118(45)ee105566111 18, 2021.
[29] Ian Tenney, Dipanjan Das, and Ellie Pavlick. Bert rediscovers the classical nlp pipeline. arXiv preprint 1905.09905.055000 2019.
[30] Hugo Touvron, Thibaut Lavril, Gautier Izacard, Xavier Martinet, Marie-Anne Lachaux, Timo- thée Lacroix, Baptiste Rozière, Naman Goyal, Eric Hambro, Faisal Azhar, et al. Llama: Open and efficient foundation language models. arXiv preprint arXiv:2302.13971 2023.
[31] Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention all you need. Advances in neural information processing systems, 30, 2017.
[32] Elena Voita, David Talbot, Fedor Moiseev, Rico Sennrich, and Ivan Titov. Analyzing multi-head self-attention: Specialized heads do the heavy lifting, the rest can be pruned. arXiv preprint arXiv:190..09418. 2019.
12

---

[33] Xinru Wang, Hannah Kim, Sajjadur Rahman, Kushan Mitra, and Zhengjie Miao. Human- collaborative annotation through effective verification of LLM labels. In Florian 'Floyd' Mueleer, Penny Kyburz, Julie R. Williamson, Corina Sas, Max L. Wilson, Phoebe O. Toups Dugas, and Irina Shklovski, editors, Proceedings of the CHI Conference on Human Factors in Computing Systems, CHI 2024, Honolulu, HI, USA, May 1-16, 2024, pages 303:1-303:21. ACM, 2024.
[34] Jason Wei, Xuezhi Wang, Dale Schuurmans, Maarten Bosma, Brian Ichter, Fei Xia, Ed H. Chi, Quoc Le, and Denny Zhou. Chain-of ofainnof-thought prompting elicits reasoning in large language models. In Sanmi Koyejo, S. Mohamed, A. Agarwal, Danielle Belgrave, K. Cho, and A. Oh, editors, Advances in Neural Information Processing Systems 35: Annual Conference on Neural Information Processing Systems 2022, NeurIPS 2022, New Orleans, LA, USA, November 28 December 9, 2022, 2022.
[35] Mark Wheeler, Donald T Stuss, and Endel Tulving. Toward a theory of episodic memory: the frontal lobes and autonoetic consciousness. Psychological bulletin, 121(3):331, 1997.
[36] Wenhao Wu, Yizhong Wang, Guangxuan Xiao, Hao Peng, and Yao Fu. Retrieval head mecha- nistically explains long-context factuality, 2024, URL https://arxiv. org/abs/2404, 15574.
[37] Daniel LK Yamins, Ha Hong, Charles F Cadieu, Ethan A Solomon, Darren Seibert, and James J DiCarlo. Performance-opiimizee hierarchical models predict neural responses in higher visual cortex. Proceedings of the national academy of sciences, 11 111(23):619-8622 2014.
[38] An Yang. Baosong Yang. Beichen Zhang, Binyuan Hui, Bo Zheng, Bowen Yu, Chengyuan Li, Dayiheng Liu, Fei Huang, Haoran Wei, al. Qwen2. 5 technical report. arXiv preprint arXiv:2412.5511 15115, 2024.
[39] Alex Young, Bei Chen, Chao Li, Chengen Huang, Ge Zhang, Guanwei Zhang, Guoyin Wang, Heng Li, Jianccheng Zhu, Jianqun Chen, et al. Yi: Open foundation models by 01. ai. arXiv preprint arXiv:2403.00 04652,, 2024.
[40] Z. Zheng, Y Wang, Y Huang, S Song, M Yang, B Tang, F Xiong, and Z Li. Attention heads of large language models: survey. arxiv 2024. arXiv preprint arXiv:2409.03755
[41] Zhenhong Zhou, Haiyang Yu, Xinghua Zhang, Rongwu Xu, Fei Huang, Kun Wang, Yang Liu, Junfeng Fang, and Yongbin Li. On the role of attention heads in large language model safety, arXiv preprint arXiv:2410. 13708, 2024.
NeurIPS Paper Checklist
Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification: Yes, the main claims in the abstract and introduction accurately reflect the paper's s contributions and scope.
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made in the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or
NA answer to this question will not be perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.
It is fine to include aspirational goals as motivation as long as it is clear that these goals are not attained by the paper.
2. Limitations
Question: Does the paper discuss the limitations of the work performed by the authors?
13

---

Answer: [Yes]
Justification: We discussed our limitation in discussion section.
Guidelines:
The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors encouraged to create separate "Limitations" section in their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be.
The authors should reflect on the scope of the claims made, e.g. if the approach was only tested on a few datasets or with a few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution low or images are taken in low lighting. Or speech-to-text system might be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.
The authors should discuss the computational efficiency the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness.
While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, worse outcome might be that reviewers discover limitations that aren't 't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an impor- tant role in developing norms that preserve the integrity of the community, Reviewers will be specifically instructed to not penalize honesty concerning limitations.
3. Theory assumptions and proof'
Question: For each theoretical result, does the paper provide the full set of assumptions and a complete (and correct) proof?
Answer: [NA]
Justification:
We don't 't have theoretical result
Guidelines:
The answer NA means that the paper does not include theoretical results.
All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated or referenced the statement of any theorems.
The proofs can either appear in the main paper or the supplemental material, but they appear in the supplemental material, the authors are encouraged to provide short proof sketch to provide intuition.
Inversely, any informal proof provided in the core of the paper should complemented by formal proofs provided in appendix or supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental result reproducibility
Question: Does the paper fully disclose all the information needed to reproduce the main perimental results the paper to the extent that affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided not)?
Answer: [Yes]
Justification: Yes. The paper provides sufficient information to reproduce the main exper- imental results. We release the dataset and describe the experimental setup, intervention methods, model training procedures, and evaluation metrics in detail in the main paper, ensuring transparency and reproducibility of the core findings.
14

---

Guidelines:
The answer NA means that the paper does not include experiments.
If the paper includes experiments, No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
If the contribution is dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is a novel urchitecture, describing the architecture fully might suffice, or if the contribution a specific model and empirical evaluation, may be necessary to either make it possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g., in the case of large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.
While NeurIPS does not require releasing code, the conference does require all submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) If the contribution is primarily new algorithm, the paper should make it clear how to reproduce that algorithm. (b) the contribution is primarily a new model architecture, the paper should describe the architecture clearly and fully.
(c) If the contribution is new model (e.g., a large language model), then there should either be way to access this model for reproducing the results or a way to reproduce the model (e.g., with an open-source dataset instructions for how to construct the dataset).
(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome describe the particular way they provide for reproducibility. the case of closed-source models, may be that access to the model limited in some way (e.g., registered users), but it should be possible for other researchers to have some path to reproducing or verifying the results.
5. Open access to data and code
Question: Does the paper provide open access to the data and code, with sufficient instruc- tions faithfully reproduce the main experimental results, described in supplemental material?
Answer: [Yes]
Justification:
We have released the complete GitHub repository with dataset and code.
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurIPS code and data submission guidelines (https://nips... cc/
for more details.
While we encourage the release of code and data, we understand that this might not be possible, so "No" an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for new open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurIPS code and data submission guidelines (https
hhipss.c//ublic/uuddds wtwpee///oooeSeeeeeoeoeoooo
www.obiiiiiiioooooioiiooo for more details.
The authors should provide instructions data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. only subset of experiments are reproducible, they should state which ones are omitted from the script and why.
15

---

At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
Providing as much information as possible in supplemental material (appended to the paper) recommended, but including URLs to data and code is permitted.
6. Experimental setting/details
Question: Does the paper specify all the training and test details (c.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Answer: [Yes]
Justification: Yes, we give details about how we select data, and how we constructed our dataset. Also the training configuration.
Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented in the core of the paper to a level of detail that is necessary to appreciate the results and make sense of them.
The full details can be provided either with the code, in appendix, or as supplemental material.
7. Experiment statistical significance
Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Answer: [No]
Justification: We believe that our experiment does not require this.
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results are accompanied by enror bars, confi- dence intervals, or statistical significance tests, at least for the experiments that support the main claims of the paper.
The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions).
The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)
The assumptions made should be given (e.g., Normally distributed errors).
should be clear whether the error bar is the standard deviation the standard error of the mean.
It is OK to report 1-sigma error bars, but should state it. The authors should preferably report a 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful to show in tables or figures symmetric error bars that would yield results that are out of range (e.g. negative error rates).
If error bars are reported in tables or plots, The authors should explain the text how they were calculated and reference the corresponding figures tables in the text.
8. Experiments compute resources
For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time of execution) needed reproduce the experiments?
Question:
Answer: [No]
Justification: We are working with inference only, compute resources is not the factor of any of our experiments.
Guidelines:
16

---

The answer NA means that the paper does not include experiments.
The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage.
The paper should provide the amount of compute required for each of the individual experimental runs as well estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't make into the paper).
9. Code of ethics
Question: Does the research conducted in the paper conform, in every respect, with the NeurlPS Code of Ethics
Answer:
[Yes]
Justification:
We don't have any code of ethics issues this paper
Guidelines:
The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.
If the authors answer No, they should explain the special circumstances that require deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g.. if there is special consid- cration due to laws or regulations their jurisdiction).
10. Broader impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [NA]
Justification:
We believe this is not related our work.
Guidelines:
The answer NA means that there is no societal impact of the work performed.
If the authors answer NA No, they should explain why their work has no societal impact why the paper does not address societal impact.
Examples of negative societal impacts include potential malicious unintended uses (c.g., disinformation, generating fake profiles, surveillance), fairness considerations deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
The conference expects that many papers will be foundational research and not tied particular applications, let alone deployments. However, if there is direct path to any negative applications, the authors should point it out. For example, it is legitimate point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.
If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition attacks, mechanisms for monitoring misuse, mechanisms to monitor how system learns from feedback over time, improving the efficiency and accessibility of ML..
11. Safeguards
Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, scraped datasets)?
Question:
Answer: [NA]
17

---

Justification:
All data we used are published public dataset.
Guidelines:
The answer NA means that the paper poses no such risks.
Released models that have high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
recognize that providing effective safeguards challenging, and many papers do not require this, but we encourage authors to take this into account and make best faith effort.
Licenses for existing assets
Question: Are the creators or original owners of assets (e.g.. code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?
Answer: [Yes]
Justification: We cited all models, and dataset used in this paper.
Guidelines:
The answer NA means that the paper does not existing assets.
The authors should cite the original paper that produced the code package dataset.
The authors should state which version the asset used and, possible, include URL
The name of the license (e.g., CC-BY 4.0) should be included for each asset.
For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided.
If assets are released, the license, copyright information, and terms of in the package should be provided. For popular datasets, wwpppssiptcoomeo thcode. com datasets has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
For existing datasets that re-packaged, both the original license and the license of the derived asset (if has changed) should provided.
this information is not available online, the authors are encouraged to reach out to the asset's creators.
13. New assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [Yes]
Justification: We introduce new annotated dataset support cognitive function analysis, which is submitted in the supplementary materials. We also provide detailed documentation describing its construction process, structure, and usage guidelines in the paper.
Guidelines:
The answer NA means that the paper does release new assets.
Researchers should communicate the details the dataset/code/moodl as part their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent was obtained from people whose asset used.
At submission time, remember to anonymize your assets (if applicable). You can either create anonymized URL or include an anonymized zip file.
14. Crowdsourcing and research with human subjects
18

---

Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well as details about compensation (if any)?
Answer: [NA]
Justification: Not realted
Guidelines:
The answer NA means that the paper does not involve crowdsourcing research with human subjects.
Including this information in the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then as much detail as possible should be included in the main paper.
According to the NeurIPS Code of Ethies, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional review board (IRB) approvals or equivalent for research with human subjects
Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or equivalent approval/review based the requirements of your country or institution) were obtained?
Question:
Answer: [NA]
Justification: Not related
Guidelines:
The answer NA means that the paper does not involve crowdsourcing research with human subjects.
Depending on the country in which research is conducted, IRB approval (or equivalent)
may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurIPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.
16. Declaration of LLM usage
Question: Does the paper describe the usage of LLMs if it is an important, original, or non-standard component of the core methods in this research? Note that if the LLM is used only for writing, editing, formatting purposes and does not impact the core methodology, scientific rigorousness, or originality of the research, declaration is not required.
Answer: [Yes]
Justification: We used LLMs as an integral part of our dataset construction process. Specif- ically, LLMs were used to generate intermediate reasoning steps and candidate answers under controlled prompting. All prompts, generation procedures, and filtering steps are clearly documented in the main paper and supplementary materials ensure transparency and reproducibility.
Guidelines:
The answer NA means that the core method development in this research does not involve LLMs as any important, original, or non non-standard components.
Please refer to our LLM policy (https: /tta://///e/ppp// cc//oonferencee/22255LL.. for what should or should not be described.
19

---

A
Appendix
A.1 The cognitive function distribution of other models
We present the heatmaps for the remaining five models in this section. The results reveal a notable universality in the sparsity patterns of attention heads across different architectures. Moreover, models within the same family tend to exhibit similar sparsity distributions. For instance, Llama3.2-3B (Figure ()) and Llama3.1-88 1-8B (Figure 2 share comparable patterns, do Qwen3-4B (Figure 7 and Qwen3-88 8B (Figure 6). as well as Yi-1.5-66 1.5-6B (Figure 9 and Yi-1.5-9B (Figure (8. This consistency is likely due to the shared architectural design and similar pretraining data within each model family,
Knowledge Recall
Retrieval
Logical Reasoning
Decbion-mmking
0.0090+
+0.000
Impo 2000
+0.0000
Layer
Undorstanding
Syntactic Lndorstanding
-0.0000
2
0.0000
Head
Figure 5: The existence of cognitive heads in Llama3.2 2-----insttuue responsible for eight distinct functions in complex reasoning tasks. The x-axis represents the head index, while the y-axis indicates the layer index.
Knowledge Recall
Retrieval
Logical Reasoning
Decisonnmmaknng
2
24
2,0005
+0.000
imporancoo
Layer
Semantic Underslanding
12
10
24
Syntectic Understanding
12
Inference
1/2
00
Math Calculation
-0,0000
0.c0n5
0.0005
Head
Figure 6: The existence of cognitive heads in Qwen3-8B responsible for eight distinct functions in complex reasoning tasks. The x-axis represents the head index, while the y-axis indicates the layer index.
20

---

Knowledge Recall
Retrieval
Logical Reasoning
Deciiinnmaking
22
0,0095
+0,0090
HoddS Importono
Layer
Semantic Understanding
12
10
Syntactic Understanding
12
Inference
12
00
Math Calculation
-0,000
-0.c0n5
0.0005
2
Head
Figure 7: The existence of cognitive heads in Qwen3-4B responsible for eight distinct functions in complex reasoning tasks. The x-axis represents the head index, while the y-axis indicates the layer index.
Knowledge Recall
Retrieval
Logical Reasoning
Decisonnmmkinn
22
0,0018
Importane
Layer
Semantic Understanding
12
10
24
Syntectic Understanding
12
Inference
Math Calculation
-0,0010
-0.000
-0.0005
-.com
30 20
Head
Figure 8: The existence of cognitive heads in Yi-1.5-9B responsible for eight distinct functions in complex reasoning tasks. The x-axis represents the head index, while the y-axis indicates the layer index.
A.2 Importance curve
We ranked the importance scores and identified the elbow point, as illustrated in Figure 10
21

---

Knowledge Recall
Retrieval
Logical Reasoning
Deciiinnmaking
24
50
-0,000
Importa
Layer
Semantic Understanding
10
24
Syntactic Understanding
12
Inference
00
Math Calculation
-0,0015
-0.0000
0.0005
1
Head
Figure 9: The existence of cognitive heads in Yi-1.5-6B responsible for eight distinct functions in complex reasoning tasks. The x-axis represents the head index, while the y-axis indicates the layer index.
Table 5: The test accuracy (%) of probing method on different LLMs.
Dataset 1.ama3.1-8B-insteer
Llama3.23B-iustruet
Qwen3-8B
Qwen3-4B
Vi-1.5-9B
Yi-1.5-6B
CogQA
84.71
80.79
77.56
75.18
Decision makiig
inference
Knowledge e Recall
Logical Reaconing
6.00e-03
4.50e 033
3,00e-03
Importance
0.00e+00
Math Calculation
Retrieval
Semantic Understanding
Syntacti Understanding
6.00e-03
4.50e-03
3.00e-03
1.50a-03
256 512 7681024
256 512 7681024
256 512 7681024
256 512 7681024
Heads
Figure 10: Importance curve for eight functions, Llama3.1 1-8B-instruct.
A.3 MLP
We train two-layer multi-class MLP for cognitive function classification. The first layer applics shared linear projection to each multi-head representation vector, reducing each to a 64-dimensional embedding. These embeddings are then flattened and concatenated into single vector of size 64 x Sumberofheadds This vector is fed into hidden layer with 512 units, followed by a ReLU activation and a dropout with rate of 0.3. The final output layer maps the 512-dimensional hidden representation to the set of cognitive function labels.
The model is trained using the Adam optimizer with learning rate of 10-1 and cross-entropy loss. Training proceeds for 100 epochs. The test accuracy of our classification method across all L.MM models is summarized in the T
22

---

A.4 Prompt for Generating CogQA
Prompt
Prompt: You expert in analytical logical reasoning. You will be given question along with its chain-of-thought process. Your task to break the question down into subquestions based the chain-of soain-of-thought process, ensuring that all necessary steps for solving the problem and constructing the logical chain are included to simulate critical thinking. Decompose the Question: Identify and formulate the key subquestions required to solve the main question logically. Fill in Missing Steps: Ensure that all essential reasoning steps are explicitly stated.
NOTE: The information of chain-of-thought cannot be used directly if it doesn't exist in main query. Each subquestion should be derived solely from the main query and the preceding subquestion. Answer the Subquestions: Provide clear, step-by-step solutions for each subquestion. Annotate Cognitive Skills: Identify and label the specific cognitive abilities required to answer each subquestion. If you believe other cognitive skills are relevant, you may also consider incorporating them. You will be given predefined labels along with their descriptions. Your goal is to enhance the logical reasoning process by making it explicit and structured.
signittivesskilsse **Retrieva Refers to the process of fetching relevant information from input text, typically involving the extraction of specific words, phrases, or sentences directly from the original text. **Knowledge Recall**; Involves the storage and recall of domain-specite specifcc knowledge, such as concepts from math, physics, biology, etc. This is typically the internal knowledge base of language model. (Corresponding to the memory head) **Semantic Understanding Refers to the ability to comprehend and extract meaning from text symbols by recognizing relationships between words, phrases, and concepts. goes beyond syntactic understanding by grasping context, intent, and underlying knowledge. **Syntactic Understanding": Involves the ability analyze and interpret the grammatical structure of sentences, including the roles and relationships of words, phrases, and clauses within the language. **Math Calculation**; Refers to the process of performing arithmetic mathematical operations to ohtain a result. It involves applying mathematical concepts, such addition, subtraction, multiplication, division, and calculus), to solve problems or derive values from given inputs. **Inference**; Involves drawing conclusions based on existing evidence or information. follows logical rules to deduce new statements or decisions from given information. **Logical Reasoning**: The process of drawing conclusions based on a of premises, following established rules of logic, used to ensure that decisions of people are coherent, consistent, and based on sound principles. ** on previous information or analysis. <cognitive_skiln Here the question: <question> question <question> Here is the chain-of-thought: <chain-of <f cot <chain-of-thoughe
complex operations (e.g.. algebra,
The process of making choice selection task based
Note
Your task is to break the question down into detailed subquestions, ensuring each subquestion can answered using only one specific cognitive skill. You need to create structured and explicit reasoning process that simulates critical thinking while maintaining clarity and precision. The subquestion needs to be easy to answer and the answer needs to be concise The information of chain-of chain-of-thought cannot be used directly if it doesn't 't exist in main query. Each subquestion should be derived solely from the main query and the preceding subquestion. You CAN NOT retrieval information from chain-of-thougee retrieval from question. Your output should be formatted as a list of JSON objects, where each object represents subquestion, its answer, and the required cognitive skill. You should use the most efficient logic to analyze the problem and minimize the number of subquestions. Output format "subquestion": "<Subquestion text>". "answer": "<Concise answer>". Scognitive_skille "<Assigned cognitive skill>" "subquestion": "<Subquestion text>", "answer": "<Concise answer>". "cognitive_skine skill": "<Assigned cognitive skill>" Your answer:
but you
23

---

A.5 Annotations
To ensure the quality and reliability of the decomposed subQAC triplets in the CogQA dataset, we design rigorous stati- annotation pipeline, combining expert review and verification. The goal to verify the logical validity of subquestions, the correctness of their associated cognitive function labels, and the accuracy of the answers.
Stage 1: Validating Subquestion Decomposition In the first stage, we evaluate whether the generated subquestions are logically sound and align with natural human reasoning. For each QA pair, three expert annotators (with backgrounds in linguisties or cognitive science) independently assess the validity each subquestion. subquestion is marked true if it meaningfully contributes to answering the main question and follows a logical reasoning trajectory. Otherwise, it is marked false.
If a subquestion depends on prior information as the question text or the inswe--from another subquestion, the subquestion order must reflect this dependency. While some subquestions can be answered in parallel and are order-independent, others have prerequisite relationships that require a specific sequence. As the overall reasoning structure often forms a graph, where both sequential and parallel dependencies coexist. During LLM inference, we include the previous subquestions and their corresponding subanswers in the prompt as prior information. Thus, the critical factor is not the ordering alone, but whether the prompt provides the necessary context answer the current subquestion accurately.
We apply the following filtering criteria:
Al-Human Agreement: If any annotator considers fewer than 60% of the subquestions valid, the entire QA decomposition is discarded. Inter-Annotator Agreement: A subquestion deemed invalid if at least two annotators mark it as false. If over 40% of the subquestions in a QA pair are invalid under this rule, the whole QA pair is removed.
This filtering ensures that the retained QA decompositions follow coherent, cognitively plausible reasoning chains.
Stage Verifying Cognitive Function Labels In the second stage, annotators evaluate the correctness of the cognitive function label assigned to each SubQAC triplet (4:, Q; G.). Three annotators independently mark each label as true or false. When discrepancies occur, annotators collaboratively reassign the correct cognitive label to ensure alignment with the underlying mental operation.
This step ensures that the categorization of subquestions accurately reflects established distinc- tions between information retrieval, semantic understanding, logical reasoning, and other cognitive processes.
Stage 3: Answer Verification via Model and Human Review In the final stage, we verify the correctness of each answer a, using both automated and manual procedures. We employ the 04-mini model (25), known for its logical reasoning capabilities, re-evaluate GPT-4o-generated answers. If 04-mini disagrees with GPT-4o, it provides an alternative answer. A human annotator then compares both answers and resolves discrepancies by supplying the correct one when necessary. Given the generally objective nature of answers, only one annotator is required for this task.
Annotation Outcome Following this multi-stage process, we retain 570 validated QA pairs, yielding a total of 3,402 high-quality SubQAC triplets. Notably, we augment certain cognitive functions to ensure balance across categories. As a result, the original 570 QA pairs were expanded to 720 (including some duplicates), with each duplicated pair potentially associated with distinet subquestions and cognitive functions.
A.6 CogQA Example
Table p illustrative examples from the CogQA dataset. The main question and its correspond- ing answer are taken from the original dataset. Based on an analysis of the main question, a sequence of sub-questions, their answers, and associated cognitive function labels are generated in order.
24

---

Table 6: Two examples from the CogQA dataset showing a main question, its final answer, and breakdown into subquestions with answers and their corresponding cognitive function labels.
Example 1:
Main Question
A onc-year subscription newspaper is offered with 45% discount. How much does the discounted subscription cost if subscription normally costs $80?
Answer
We calculate first the discount: 80 45/ 100 = S36. So, the discounted subscription amounts to 80 36 S44.
Subquestion
Answer
Cognitive Label
What is the normal cost of a one-year subscription to the news-
$80
Retrieval
paper?
What the discount percentage offered on subscription?
45%
Retrieval
3. How much is the discount amount in dollars for the subscription?
$36
Math Calculation
What the cost of subscription after applying discount?
$44
Math Calculation
Example 2:
Main Question
What does every person talk out of? Options: name hide mother and father
mouth heart
Answer
By mouth, talking is done. Every person talk out mouth.
Subquestion
Answer
Cognitive Label
What the primary function of talking?
communi- cate verbally.
Knowledge Recall
Which part of the human body is primarily used for verbal communication?
Mouth
Knowledge Recall
3. Based on the options provided, which option corresponds to the part used for verbal communication?
Mouth
Decision-making
A.7 Prompt for Question Asking
Prompt
Prompt: You are an expert in analytical and logical reasoning. You will be given main question and prior knowledge in chain-of-thought (CoT) format. Your task answer follow-up subquestion using the information provided. Here the main question:
Here the prior knowledge in question>
amain__uestion question </mmain_questions
(CoT) format:
<prior_knen cillo__kwwwedge>>>> cot
Here the subquestion:
<subquestion> subquestion </subquestion>
Instructions:
Answer the subquestion carefully.
You can use the information in the -noolldeeeeeeeeeee to help you answer the subquestion.
Your response should be clear and concise.
Stick to factual reasoning based on provided COT.
Do not include any explanation, commentary, code.
-Do not output anything after the closing square bracket 'J'
Only output your final answer using this format:
["answer": "<Your answer heres")
Your answer:
25

---

Table 7: Count (C) and percentage (%) of attention heads exceeding elbow thresholds for each cognitive function across six models.
Model
Recall Retrieval
Logical
Decision Semantic Syntactic Inference
Math
C
%
C
%
%
C
C
%
Llama3. 1-B.a...---B-istrree
105
10.3
118
11.5
142
13.9
124
12.1
60
5.9
81
7.9
139
13.6
59
5.8
Llama3.2-BB-insiue
95
14.1
62
9.2
14.1
12.9
13.4
63
9.4
98
14.6
35
5.2
Qwen3-8B
119
10.3
115
10.0
114
9.9
87
7.6
68
5.9
108
9.4
178
15.5
61
5.3
Qwen3-4B
115
10.0
94
8.2
120
10.4
170
14.8
143
12.4
106
9.2
109
9.5
99
8.6
Yi-1.5-9B
200
13.0
134
8.7
134
8.7
174
11.3
218
14.2
140
9.1
173
11.3
167
10.9
Yi-1.5-6B
118
11.5
90
8.8
200
19.5
9.1
9.7
142
13.9
146
14.3
67
6.5
A.8
The number of cognitive heads for different LLMs
The number of cognitive heads for each model is shown Table
A.9 Ablation study Different position of head activation
In the main experiments, we use the top-k generated tokens and average their multi-head attention vectors. We also explore alternative strategies for extracting representations, including using the first generated token, the last generated token, the first meaningful token, and the average of all generated tokens. The corresponding results shown in Table
Here, first is the first token, last the last token, meaning_fire first the first meaning token (excluding formatting), top-k is the top-k most semantically important tokens, full is all tokens in the answer. We observe that top-k token masking leads to most significant performance drop when masking the top- 30 identified heads, indicating higher precision in identifying retrieval-relevane heads. Interestingly, last, meaning first, full, and top-k show similar performance trends. This is because different tokens in the output contribute to answering the question, and as the number of masked cognitive heads increases, the influence of token using decreases. Additionally, for Retrieval, the full answer usually meaningful, whereas others like Math Calculation require semantically meaningful tokens. Based on these results, we choose top-k as our final setting.
Table 8: Attention heads associated with cognitive functions are selected based on different token positions. Accuracy and COMET scores evaluated after intervention; lower values indicate better outcomes.
Model
Head num
Token use
Retrieval(comet)
Retrieval(acc)
Math(comet)
Math(acc)
Llama3.1-8 1-8B
30
first
90.51
83,53
91.13
73.13
Llama3. 1-8B
30
last
86.86
81.76
90.04
68.66
Llama3.1-88 1-8B
30
meaning
88.13
79.41
89.72
68.66
Llama3.1 1-8B
full
73.93
47.06
89.92
69.15
Llama3.1 1.mma3.1-81 1-8B
30
lop-k
70.05
46.47
89.32
67.16
Llama3. 10mma3.1-88 1-BB
50
first
93.28
89.41
94.46
89.57
Llama3. 1-8B
50
last
64.39
41.18
92.05
70.15
Llama3. 1-8B
50
meaning first
62.90
34.12
84.60
60.69
Llama3. 1-8B
full
89.01
78.11
Llama3. 1-BBB
50
top-k
65.64
47.76
89.65
70.15
A.10
Examples of top-k tokens
The selected tokens are intended semantically represent the generated answer. Below are examples for different cognitive functions for Llama3.1 1-iniiit-tns
26

---

Table 9: Examples of question decomposition with cognitive functions and token selection. Main Question
Subquestion
Cognitive Func- Answer Selected To- kens
tion
Given the sentence "A surfboarder catches the waves." can we conclude that surfboarder in the water."? (Options: yes / it is not possible to tell / no)
What is typically required for surf- boarder catch waves?
Knowledge call
Re- The surf- ['surfboarder',
boarder needs 'needs', 'be', 'in', 'water']
to be in the water.
Is the following fac- tual statement? "Due to its high density, countries around the world use Pal- ladium mint coins." (Options: yes / no)
What is the state- ment in question?
Retrieval
The statement in question is: Due its high density, countries around the world use Palladium to mint coins.
['high',
'density', 'Palladium', 'mint', coins']
A one-year subscription a newspaper is of- fered with 45% dis- count. How much does the discounted subscrip- tion cost if subscription normally costs $80?
How much is the discount amount in dollars for the subscription?
Math Calculation
36
['36]]
We can see that the selected tokens semantically represent the answer. Note that we use all tokens when the number of tokens is fewer than 5.
27


Reviews:
[{"review_id": "ZreIL9CJn3", "reviewer": "Reviewer_rBWy", "summary": "The paper introduces CogQA, a dataset for evaluating the use of specialized attention heads for cognitive functions similar to the brain in chain-of-thought reasoning tasks. The functions are created by decomposing problems from existing frameworks into sub-questions and classified into low-level and high-order functions. The authors use probes to predict sub-question class and identify which linear heads contribute to the classification.", "questions": "1. Have you investigated whether the question decomposition follows the implicit sequential reasoning assumption? Can it affect the quality of the generated dataset?\n\n2. As the probe used is a 2-layers MLP, classification may come from non-linear interactions between several heads (across layers). Have you investigated this hypothesis?\n\n3. Can you detail more precisely how you select the top-k most important tokens in the LLM's answer?", "limitations": "yes", "rating": 5, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 4, "clarity": 3, "significance": 2, "originality": 3, "strengths_and_weaknesses": "The paper discusses an interesting problem and goes beyond what is done in most of the functional analysis literature by tackling CoT reasoning tasks instead of single-token prediction. The writing is clear and easy to follow and the experiments back up the claims made in the paper. While the findings are not fundamentally novel, I believe they are significant and of interest for the research community.\n\nHowever, I have a few concerns:\n\n1. The CoT decomposition is not generated by the evaluated model but by another model (GPT-4o when creating the dataset), one possible consequence is that these heads are not significantly used by the model in practice as it may not generate the same decomposition or may even not be able to plan the right reasoning steps.\n\n2. Probing only tells that a head contains the necessary information for classification, and not that this head is actually used in the computation or that it executes the associated function. In particular, if my understanding is correct, the probe takes information from every layer at once (and not from a particular layer), so the classification may result from non-linear interactions between heads and even across layers (the quantity of noise in the head importance in Figure 2 supports this hypothesis).\n\n3. The decomposition method divides a problem into a list of sub-questions, implicitly assuming that the reasoning structure is sequential while it could be graph-based (e.g. diamond shape: solve two independent subproblems and combine the results to solve the full task).\n\n4. minor typo: final line of section 2.2 \"for generating subqeustions\"", "was_revised": true, "final_justification": "The discussion with the authors resolved my concerns and answered my questions and misconceptions about the paper.\n - The issues I had with the subquestion decomposition have been adressed\n - The concerns about the non-linear interactions that could be captured from probing have similarly been resolved\n\nOverall, this is a very interesting work and I keep my positive score.", "rebuttal": "Dear Reviewer,\n\nThank you for your constructive comments and suggestions. They have been invaluable in helping us enhance the quality and clarity of our paper. Please find our point-by-point responses to your concerns below.\n\n**Q1. The CoT decomposition is not generated by the evaluated model but by another model (GPT-4o when creating the dataset), one possible consequence is that these heads are not significantly used by the model in practice as it may not generate the same decomposition or may even not be able to plan the right reasoning steps.**\n\n**A.**  Our intention is to use chain-of-thought (CoT) not as an output of the model under analysis, but as a lens to study the model’s internal structure. The CoT decompositions are carefully designed and verified by human annotators to reflect the human-like cognitive steps required to solve the problem.\n\nInstead of prompting the model with the full question, we query it with individual subquestions, each aligned with a specific cognitive function. During prompting, previous subquestions and their ground-truth subanswers are provided as prior context. This setup allows us to observe which attention heads are activated as the model processes each cognitive step. In doing so, we can probe the internal mechanisms underlying distinct reasoning capabilities, offering insights into how different heads support specific cognitive functions—even if the model does not decompose problems this way under standard inference.\n\n**Q2. Probing only tells that a head contains the necessary information for classification, and not that this head is actually used in the computation or that it executes the associated function. In particular, if my understanding is correct, the probe takes information from every layer at once (and not from a particular layer), so the classification may result from non-linear interactions between heads and even across layers (the quantity of noise in the head importance in Figure 2 supports this hypothesis). As the probe used is a 2-layers MLP, classification may come from non-linear interactions between several heads (across layers). Have you investigated this hypothesis?**\n\n**A.** Thank you for raising this insightful point. Inspired by findings in neuroscience [1-2] where specific regions of the human brain are activated when performing certain tasks, we hypothesize that particular attention heads are similarly activated when a task mainly relies on a specific cognitive function. These heads exhibit distinct activation patterns that set them apart from other unrelated heads. shows more important for classification.\n\nAdditionally, the interventions (positive and negative) complement probing by demonstrating the causal impact of those heads on task performance, suggesting that at least some of the probed signals are functionally utilized. \n\nIn our initial experiments, we trained separate logistic regression classifiers for each head in each layer. However, results on Llama3-8B show that cognitive heads identified using only activation features and multiple independent classifiers did not outperform those selected by our current single-classifier approach in subsequent intervention experiments.\n\n[1] John R Anderson. Rules of the mind. Psychology Press, 2014.\n\n[2] Poldrack, Russell A. \"Can cognitive processes be inferred from neuroimaging data?.\" Trends in cognitive sciences 10.2 (2006): 59-63.\n\n**Q3.The decomposition method divides a problem into a list of sub-questions, implicitly assuming that the reasoning structure is sequential while it could be graph-based (e.g. diamond shape: solve two independent subproblems and combine the results to solve the full task).**\n\n**A.** In constructing CogQA, annotators are instructed to ensure that each subquestion aligns with natural human reasoning. Specifically, if a subquestion depends on prior information—such as the question text or the answer—from another subquestion, the subquestion order must reflect this dependency. While some subquestions can be answered in parallel and are order-independent, others have prerequisite relationships that require a specific sequence. \n\nWe acknowledge that the overall reasoning structure often forms a graph, where both sequential and parallel dependencies coexist.\nDuring LLM inference, we include the previous subquestions and their corresponding subanswers in the prompt as prior information. Thus, the critical factor is not the ordering alone, but whether the prompt provides the necessary context to answer the current subquestion accurately.\n\n**Q4. Can you detail more precisely how you select the top-k most important tokens in the LLM's answer?**\n\n**A.** We select the top-k most important tokens by prompting LLM.\nGiven a reference answer and a predicted answer to a sub-question. We first predict whether the prediction answer is consistent to the reference answer, \"True\" and \"False\".\n\nIf \"True\" is returned, extract and return the 5 most semantically important tokens from the prediction answer.\n   - These tokens should reflect the core meaning of the answer.\n   - Avoid common stopwords unless they are crucial to the semantics.\n \n**Q5. minor typo: final line of section 2.2 \"for generating subqeustions\"**\n\n**A.** Thank you for pointing out this typo. We will carefully proofread the manuscript to ensure clarity and correctness throughout."}, {"review_id": "d6fV17VMqr", "reviewer": "Reviewer_wHF7", "summary": "The paper investigates the cognitive role of attention heads in Transformers. The central assumption is that transformers exhibit cognitive functions analogous to those in humans, and that these functions are embedded within attention heads. To explore this hypothesis, the authors design a dataset where each question is annotated with a specific cognitive function. They then use a linear probe classifier to identify the cognitive role of each attention head. In the experiments, the authors analyze the distribution of cognitive functions across attention heads and perform interventions to evaluate the impact of individual heads on model performance.", "questions": "Please see the weaknesses part.", "limitations": "Yes, the authors include a discussion of the limitations of their work.", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 3, "originality": 3, "strengths_and_weaknesses": "**Strengths**\n- The paper connects insights from cognitive science with interpretability research, which is an interesting direction.\n- The idea of linking attention heads to cognitive functions is intriguing.\n- The paper introduces a new dataset, which may be valuable for future research.\n\n**Weaknesses**\n- The definitions of each cognitive function are vague and ambiguous. The proposed categories may not be atomic or mutually exclusive. For example, decision-making and inference often co-occur with other classes, and it is unclear why mathematical computation would be considered distinct from or incompatible with logical reasoning. Applying a classifier to force these functions into exclusive categories may not be a valid approach.\n\n- The method of using GPT-4o-mini to \"extract key words\" is not a sound approach:\n  - It is unclear whether key words alone can fully capture or identify a cognitive function.\n  - Even if key words do provide a comprehensive representation in some cases, it is difficult to justify that GPT-4o-mini and the target model share the same mechanism or metric for determining word importance.\n\n- It is unclear why adding the layer-averaged attention vector as a feature would resolve the issue of cognitive functions varying by depth, given that this information still originates from the same layer. Moreover, since the goal is to identify the distinct role of each attention head, incorporating the average activation may fail to disentangle the influence of individual heads from the collective effect of the layer.\n\n- The underlying assumption that each attention head corresponds to a single cognitive function is quite strong. A counterexample is the first layer of an induction head, whose role is primarily to provide contextual information rather than perform a distinct cognitive function.\n\n- It is unclear how the reported accuracy is calculated. If accuracy is based solely on whether the model answers correctly, it becomes difficult to attribute failure cases to the cognitive function roles of specific attention heads. Suppressing an essential head across all examples could lead to uniformly degraded performance, as shown in Table 1, without offering insight into functional specialization.\n\n- The paper derives its results from a single dataset. To demonstrate that the identified functional roles are generalizable, the authors should evaluate the attention heads—identified using the proposed dataset—on out-of-distribution tasks. For example, they could suppress the attention head associated with mathematical calculation and assess its impact on a benchmark like GSM8K to test whether the identified function holds in a different setting.\n\n- Regarding the steering experiment in Section 4.4, quantitative results evaluated on multiple examples are needed, rather than relying on a single-case study.\n\n\n- Reproducibility is limited, as the authors do not provide code at the time of review.\n\n**Minor:**\n\n- In line 138, the set annotation should be under the summation symbol.", "was_revised": true, "final_justification": "The authors provide sufficient evidence and details. I have no other concerns.", "rebuttal": "Dear Reviewer,\n\nWe sincerely appreciate the time and effort you’ve invested in reviewing our paper. Below, we provide detailed responses to the questions you raised.\n\n**Q1. About definitions of Cognitive Functions and the Validity of the proposed method.**\n\n**A.** 1. Inspired by cognitive science [1-2] and recent LLM attention research [3], we focus on eight key cognitive functions, split into low- and high-level categories (see Subsection 2.1). To better align with human cognitive processes, we incorporate human annotations to guide the function labeling for subquestions. A brief overview of these functions is provided in subsection 2.1, with details in Appendix A.4.\n\n2. About decision-making: we define it as \"the process of making a choice in a selection task based on prior information or analysis.\" This prior information often results from other cognitive processes such as inference or mathematical computation. Therefore, in decision-making tasks (e.g., multiple-choice questions), the model typically does not need to perform new reasoning, but rather selects the correct answer based on the output of earlier processes. This aligns with our observation that decision-making is frequently associated with retrieval and inference/logical reasoning, consistent with its cognitive definition (see Subsection 4.3). Importantly, cognitive functions in LLMs—like those in the human brain—are inherently complex and interdependent.\n\nFurthermore, we distinguish between mathematical computation and logical reasoning. For instance, the question “What is the total amount Lee earned from mowing 16 lawns?” requires mathematical computation, whereas “Can a man who is swimming also go hiking?” involves logical reasoning. These tasks have distinct cognitive demands and should not be conflated.\n\n3. In our current setting, we treat each instance as primarily reflecting a dominant cognitive function to facilitate analysis and probing, rather than to assert hard functional boundaries. Our experimental results in subsection 4.2 and 4.4 in the paper demonstrate the effectiveness of this approach. An additional experiment investigating the influence of masking one class of cognitive heads on other cognitive tasks (Response to Q5) further demonstrates that the cognitive heads are indeed relatively disentangled.\n\n**Q2. About using GPT-4o-mini to extract key words.**\n\n**A.** 1. In CogQA, each subquestion is annotated with a predefined cognitive function, which is determined primarily based on the subquestion itself. Furthermore, certain attention heads are more likely to be activated when generating key tokens associated with specific cognitive functions, thus, we use GPT-4o-mini to extract key tokens from the generated subanswers, to help identify which attention heads are activated during the generation of these function-relevant tokens. \n\n2. We do not assume that GPT-4o-mini and the target model share identical importance metrics. Rather, we use these tokens to isolate content-relevant attention patterns and to reduce noise in the probing process.\n\n**Q3. About adding the layer-averaged attention vector.**\n\n**A.** It has been suggested that cognitive functions may vary by layer depth [3], meaning that even if two attention heads across different layers exhibit similar activation patterns aligned with a cognitive function, they might still play distinct roles due to their positional context in the model. Thus, layer information is crucial for interpreting functional specialization. We include the layer vector as a coarse signal to differentiate these roles.\n\nWe compare using the layer-averaged vector (Layer_True) versus not (Layer_False) in the following table. Masking heads with Layer_True causes a larger accuracy drop, showing the value of including layer information.\n\n**Table. The results of Llama3.1-8B-instruct when masking Retrieval heads.**\n|Model|Inter_Head|comet|acc|\n|-|-|-|-|\n|Llama3.1-8B|layer_True|70.05|46.47|\n|Llama3.1-8B|layer_False|69.88|53.53|\n\n**Q4. About the underlying assumption of attention head and the example of induction head**\n\n**A.** For the example of induction head, while certain heads in the first layer are primarily responsible for providing contextual information that is also related to induction, there exist some heads that are directly involved in performing induction. While they are both related to induction, the difference in their roles reflects their varying importance for the induction task. In other words, masking an induction head with high functional importance tends to result in a greater performance drop than masking a lower-layer contextual head.\n\nIn our approach, we quantify and rank the importance of attention heads for a given cognitive function, identifying the top ones as cognitive heads. This does not imply that all other heads are entirely non-functional or irrelevant; rather, they contribute less significantly compared to the top-ranked ones.\n\nImportantly, while prior work often focuses on a single function or mechanism, our study extends to multiple cognitive functions and adopts a more fine-grained perspective. We aim to reveal dominant functional tendencies in heads, not strict one-to-one mappings.  We will clarify and illustrate this in the revised version.\n\n**Q5.  About accuracy and evidence for functional specialization.**\n\n**A.** As described in Subsection 4.2 (lines 202–206), we use several standard evaluation metrics, including COMET, BLEU, ROUGE, and semantic similarity (measured by the cosine similarity of embeddings). The overall accuracy is computed as a combination of BLEU, ROUGE, and semantic similarity scores. \n\nBeyond Table 1, examples in Subsection 4.4 demonstrate that a positive intervention on Math Calculation heads in Qwen3-4B can correct originally incorrect mathematical computations, whereas a negative intervention can cause originally correct computations to become incorrect. Notably, the examples used in Subsection 4.4 are not from the CogQA dataset.\n\nTo further validate the functional roles of cognitive heads, we conduct experiments where we mask the retrieval heads during the evaluation of knowledge recall (Recall), and conversely, mask knowledge recall heads during the evaluation of retrieval performance. As shown in the following table, masking the corresponding cognitive heads causes a significantly larger performance drop than masking others, confirming their functional specialization.\n\n**Table. Intervention results (%) of different cognitive heads and random heads across Retrieval and Recall functions.**\n|Model|Inter_Head|Retrieval (comet)|Retrieval (acc)|Recall (comet)|Recall (acc)|\n|-|-|-|-|-|-|\n|LLama3.1-8B|random|90.83|84.71|87.85|83.84|\n|LLama3.1-8B|retrieval|44.96|8.24|72.05|33.33|\n|LLama3.1-8B|recall|86.79|75.29|56.93|38.38|\n|Qwen3-8B|random|92.81|75.29|89.90|53.54|\n|Qwen3-8B|retrieval|59.19|38.24|79.26|57.58|\n|Qwen3-8B|recall|83.31|71.18|64.81|30.30|\n\n**Q6. Experiments to demonstrate that the identified functional roles are generalizable.**\n\n**A.** Thank you for your suggestions. To evaluate the generalization of the identified functional heads, we present results on data beyond the CogQA dataset in Subsection 4.4. To further quantify the results, we conduct experiments on two tasks: a math task using 100 GSM8K samples (GSM8K_100) and a retrieval task with 49 samples from an extractive QA dataset. The Extractive_QA pairs are generated by GPT-4o, with answers extracted directly from the source paragraph.\n\nWe suppress Math Calculation heads for GSM8K_100 and Retrieval heads for Extractive_QA. As shown in the table, this causes significant performance drops across models, confirming these heads’ consistent functional roles. Math accuracy is measured by exact answer match; Extractive_QA accuracy checks if the original paragraph answer appears in the response.\n\n**Table. LLM performance on GSM8k_100 and Extractive_QA by negative intervention. Inter. is the negative intervention.**\n|Dataset|Head|Llama3.1-8B|Llama3.2-3B|Qwen3-8B|Qwen3-4B|\n|-|-|-|-|-|-|\n|GSM8k\\_100|Math (Base)|82|64|94|91|\n|GSM8k\\_100|Math (Inter.)|38|34|34|37|\n|Extractive\\_QA|Retrieval (Base)|57.14|36.73|57.14|51.02|\n|Extractive\\_QA|Retrieval (Inter.)|0|0|14.29|12.24|\n\n**Q7. More quantitative results about steering experiment.**\n\n**A.** We intervene on the cognitive heads in LLMs using activation adjustments on the GSM8K_100 and Extractive_QA datasets. The table below shows that boosting retrieval head activations along their functional directions improves retrieval task performance.\n\nMath tasks often involve multiple cognitive functions. We observe that in some cases, the numerical computations are correct, but the model’s understanding of the question is semantically flawed. By positively intervening on either the Math Calculation or Semantic Understanding heads, we observe consistent improvements in performance. This highlights the need for intervention methods that adjust multiple cognitive heads simultaneously. We leave this for future work.\n\n**Table. LLM performance on GSM8k_{100} and Extractive_QA by positive intervention. Inter. is the positive intervention.**\n|Dataset|Head|Llama3.1-8B|Llama3.2-3B|Qwen3-8B|Qwen3-4B|\n|-|-|-|-|-|-|\n|GSM8K|Math(Base)|82|64|91|94|\n|GSM8K|Math(Inter.)|84|66|92|94|\n|GSM8K|Semantic(Base)|82|64|91|94|\n|GSM8K|Semantic(Inter.)|84|65|93|94|\n|Extractive_QA|Retrieval(Base)|57.14|36.73|57.14|51.02|\n|Extractive_QA|Retrieval(Inter.)|63.26|44.90|61.22|69.38|\n\n**Q8. About code and the set annotation in line 138.**\n\nWe will open-source our code upon acceptance to ensure reproducibility. Thanks for noting the mistake in line 138; we will correct it and review the manuscript thoroughly.\n\n[1] John R Anderson. Rules of the mind. Psychology Press, 2014.\n\n[2] Adele Diamond. Executive functions. Annual review of psychology, 64(1):135–168, 2013\n\n[3] Zheng, Z., et al. \"Attention heads of large language models: A survey. arXiv 2024.\" arXiv preprint arXiv:2409.03752."}, {"review_id": "T6SHY1MDYb", "reviewer": "Reviewer_Nip4", "summary": "This paper introduces a new dataset for evaluating the cognitive roles of different attention heads in large language models (LLMs). In particular, the dataset is constructed by prompting ChatGPT to decompose original questions into sub-questions focusing on different roles (e.g., retrieval and logical reasoning). Its analysis centers around training linear classifiers to predict the cognitive roles based on features from different heads and layers, and estimating the importance of attention heads with gradient-based saliency methods. Experiments highlight several interesting characteristics of the attention heads, e.g., sparsity of important heads and their structure across layers.", "questions": "(1) Please clarify the usage of subquestion and linear classifier during the analysis.\n\n(2) What are most of the important attention heads located in a few middle layers regardless of cognitive roles?\n\n(3) Please consider performing the masking experiments with a different baseline (as suggested in the comments above).\n\n(4) It would be reasonable to include more quantitative analysis for the positive/negative interventions.", "limitations": "yes", "rating": 4, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 3, "originality": 3, "strengths_and_weaknesses": "As far as I am concerned, this paper has the following strengths:\n\n(1) It is an intuitive idea to understand LLMs based on the distinct roles of their attention heads. The paper performs a collection of analyses, and facilitates understanding from different perspectives.\n\n(2) The paper presents a new dataset, which could be useful for subsequent studies of related fields.\n\n(3)  It is interesting that manipulating attention heads with lower-level cognitive roles leads to considerable impacts on higher-level cognitive reasoning.\n\nHowever, it also has several limitations that can be improved:\n\n(1) It is not very clear how the hierarchical structure experiment is performed (Table 2). Are the authors consistently manipulating attention heads of lower cognitive functions across all subquestions, or only permuting the initial subquestions related to low-level cognitive roles and evaluating decision-making questions after? If it is the former, how would attention heads affect the chain-of-thought reasoning?\n\n(2) This may be related to the previous comment. What are the key advantages of question decomposition? It is also feasible to determine the cognitive roles of the original questions without decomposition. \n\n(3) The authors argue that the attention heads have layered functional organization. Nevertheless, this is mostly obvious only for math vs other cognitive functions. Looking at Figure 2 and Figure 5-9, it appears that the important attention heads are commonly located within a few layers (e.g., layer 13-15 in Figure 2), regardless of cognitive functions and models. Any idea why these layers are so special?\n\n(4) Table 1 shows that masking important attention heads leads to significant performance drop, which random masking does not. However, this could be related to the observations that important attention heads are concentrated in a few layers, without considering the cognitive roles (e.g., randomly masking attention heads in layer 1-10 may not affect the performance, while masking layers 13-15 would). I would suggest considering a different baseline, which performs the permutation based on attention heads identified by a different cognitive role (e.g., masking attention heads for retrieval while evaluating the performance for recall).\n\n(5) Analysis in Section 4.4 looks interesting, but only two qualitative examples are provided. In addition, both of them are about math calculation.\n\n(6) The proposed method simultaneously considers features from all attention heads across layers during linear probing, which leads to inputs of very high dimension. Considering the size of the dataset (~3000 samples), do the linear classifiers have overfitting issues? What is the accuracy of these classifiers on predicting the cognitive roles?", "was_revised": true, "final_justification": "The rebuttal addresses most of my concerns raised in the initial review, and I have decided to raise my score.", "rebuttal": "Dear Reviewer,\n\nWe deeply appreciate the time and effort you have dedicated to reviewing our work. Below, we address the questions you raised.\n\n**Q1.  It is not very clear how the hierarchical structure experiment is performed (Table 2). Are the authors consistently manipulating attention heads of lower cognitive functions across all subquestions, or only permuting the initial subquestions related to low-level cognitive roles and evaluating decision-making questions after? If it is the former, how would attention heads affect the chain-of-thought reasoning?**\n\n**A.** For the hierarchical structure experiment—such as analyzing the influence of Retrieval heads on Mathematical Calculation—we intervene by masking Retrieval-related heads across all subquestions in the LLM. The answers generated from earlier subquestions then serve as prior information for subsequent subquestions involving math calculation. This setup enables us to investigate how suppressing low-level cognitive heads (e.g., Retrieval) can indirectly impact higher-level functions (e.g., Math Calculation) by propagating through the chain of reasoning.\n\nAdditionally, as shown in Table 1, randomly masking certain heads can also have a direct effect on cognitive functions.\n\nThus, in this experiment, low-level cognitive heads influence high-level reasoning via two mechanisms: (1) indirectly, by affecting intermediate subquestion answers within the chain-of-thought process; and (2) directly, by masking a part of heads.\n\n**Q2. This may be related to the previous comment. What are the key advantages of question decomposition? It is also feasible to determine the cognitive roles of the original questions without decomposition.  Please clarify the usage of subquestion and linear classifier during the analysis.**\n\n**A.** Thank you for the valuable question. We would like to clarify the key advantages of using question decomposition in our analysis:\n\n1. Reasoning tasks typically involve the presence—but rarely the explicit absence—of certain cognitive functions. Treating the full question as a single unit results in a lack of well-defined negative samples (i.e., tasks that do not require a given function), which makes probing-based analysis less effective. Moreover, incorrect answers do not necessarily serve as reliable negative examples, since a wrong answer can arise from various causes unrelated to the absence of a cognitive function. By decomposing questions into subquestions, we obtain fine-grained annotations that clearly delineate the cognitive functions involved in each part, enabling precise classification and more effective probing.\n\n2. Reasoning tasks are diverse, with different tasks involving different and often multiple cognitive functions. Subquestions allow us to disentangle these functions and examine them systematically and independently, thereby reflecting a more human-like interpretable model. \n\n3. Importantly, question decomposition facilitates building a comprehensive interpretability framework that accounts for both the independent and collaborative roles of attention heads across cognitive functions. This framework offers practical benefits for model design, modular training, and fine-tuning strategies.\n\n**Q3. The authors argue that the attention heads have layered functional organization. Nevertheless, this is mostly obvious only for math vs other cognitive functions. Looking at Figure 2 and Figure 5-9, it appears that the important attention heads are commonly located within a few layers (e.g., layer 13-15 in Figure 2), regardless of cognitive functions and models. Any idea why these layers are so special? What are most of the important attention heads located in a few middle layers regardless of cognitive roles?**\n\nFrom Figures 2 and 5–9, we observe that many cognitively important attention heads are located in the middle layers. This finding is consistent with [1], which suggests that intermediate layers can encode even richer representations. A recent study on VLMs [2] further supports this, demonstrating that middle layers play a critical role. Notably, while math-related heads tend to emerge at deeper layers, other high-level cognitive functions—such as decision-making—also become more prominent in the higher layers for some LLMs. In general, shallow layers primarily capture surface-level patterns, whereas mid-to-late layers are responsible for constructing task-specific and abstract representations closely tied to cognitive processing.\n\n[1] Skean, Oscar, et al. \"Layer by layer: Uncovering hidden representations in language models.\" ICML 2025.\n\n[2] Jiang, Zhangqi, et al. \"Devils in middle layers of large vision-language models: Interpreting, detecting and mitigating object hallucinations via attention lens.\" CVPR 2025.\n\n**Q4.  Please consider performing the masking experiments with a different baseline.**\n\n**A.** To further validate the functional roles of cognitive heads, we conduct experiments where we mask the retrieval heads during the evaluation of knowledge recall (Recall), and conversely, mask knowledge recall heads during the evaluation of retrieval performance. As shown in the following table, masking the corresponding cognitive heads leads to a significantly larger performance drop compared to masking other heads, supporting their functional specialization.\n\n\n**Table. Intervention results (%) of different cognitive heads and random heads across Retrieval and Recall functions.**\n|Model|Inter_Head|Retrieval (comet)|Retrieval (acc)|Recall (comet)|Recall (acc)|\n|-|-|-|-|-|-|\n|LLama3.1-8B|random|90.83|84.71|87.85|83.84|\n|LLama3.1-8B|retrieval|44.96|8.24|72.05|33.33|\n|LLama3.1-8B|recall|86.79|75.29|56.93|38.38|\n|Qwen3-8B|random|92.81|75.29|89.90|53.54|\n|Qwen3-8B|retrieval|59.19|38.24|79.26|57.58|\n|Qwen3-8B|recall|83.31|71.18|64.81|30.30|\n\n**Q5. Analysis in Section 4.4 looks interesting, but only two qualitative examples are provided. In addition, both of them are about math calculation.**\n\n**A.** To further quantify the results, we conduct experiments on two tasks: a math task and a retrieval task. For the math task, we use 100 samples from GSM8K (denoted as GSM8K\\_100), and for the retrieval task, we use an extractive QA dataset consisting of 49 samples, as described in [1]. The Extractive\\_QA dataset is constructed by prompting GPT-4o to generate question–answer pairs from a given paragraph, where the answer is an original sentence extracted from that paragraph.\n\nFor negative intervention, we suppress the Math Calculation heads when evaluating on GSM8K\\_100, and the Retrieval heads when evaluating on Extractive\\_QA. As shown in the following table, this suppression leads to a significant drop in performance on both tasks across different models, demonstrating that these heads exhibit consistent and transferable functional roles.\n\n**Table. LLM performance on GSM8k_100 and Extractive_QA by negative intervention. Inter. is the negative intervention.**\n|Dataset|Head|Llama3.1-8B|Llama3.2-3B|Qwen3-8B|Qwen3-4B|\n|-|-|-|-|-|-|\n|GSM8k\\_100|Math (Base)|82|64|94|91|\n|GSM8k\\_100|Math (Inter.)|38|34|34|37|\n|Extractive\\_QA|Retrieval (Base)|57.14|36.73|57.14|51.02|\n|Extractive\\_QA|Retrieval (Inter.)|0|0|14.29|12.24|\n\nFor positive intervention, we intervene on the cognitive heads in LLMs using positive activation adjustments. The experimental results in the following table show that enhancing the activation of retrieval heads along their corresponding functional directions leads to improved performance on the retrieval task.\n\nFor the math task, each question often involves multiple cognitive functions. We observe that in some cases, the numerical computations are correct, but the model’s understanding of the question is semantically flawed. By positively intervening on either the Math Calculation or Semantic Understanding heads, we observe consistent improvements in performance. This indicates the importance of developing intervention methods that can simultaneously adjust multiple cognitive functional heads which collaborate during complex tasks. We leave this as an avenue for future work.\n\n**Table. LLM performance on GSM8k_100 and Extractive_QA by positive intervention. Inter. is the positive intervention.**\n|Dataset|Head|Llama3.1-8B|Llama3.2-3B|Qwen3-8B|Qwen3-4B|\n|-|-|-|-|-|-|\n|GSM8K|Math(Base)|82|64|91|94|\n|GSM8K|Math(Inter.)|84|66|92|94|\n|GSM8K|Semantic(Base)|82|64|91|94|\n|GSM8K|Semantic(Inter.)|84|65|93|94|\n|Extractive_QA|Retrieval(Base)|57.14|36.73|57.14|51.02|\n|Extractive_QA|Retrieval(Inter.)|63.26|44.90|61.22|69.38|\n\n[1] Wu, Wenhao, et al. \"Retrieval head mechanistically explains long-context factuality.\" arXiv preprint arXiv:2404.15574 (2024).\n\n**Q6. The proposed method simultaneously considers features from all attention heads across layers during linear probing, which leads to inputs of very high dimension. Considering the size of the dataset (~3000 samples), do the linear classifiers have overfitting issues? What is the accuracy of these classifiers on predicting the cognitive roles?**\n\nOur experimental results show that the average accuracy for predicting cognitive roles on test dataset of CogQA is around 80\\% which demonstrates that the classifiers can effectively capture the relevant features.\n\n**Table. Test accuracy for LLMs.**\n| Dataset | Llama3.2-3B | Qwen3-8B | Yi-1.5-9B |\n|---------|-------------|-------------|----------|\n| CogQA   | 79.80       | 84.71    | 80.79    | 77.56     |"}, {"review_id": "VFIDjJ19kM", "reviewer": "Reviewer_8YmT", "summary": "The paper presents an analysis on the roles of attention heads in LLMs in reasoning tasks, aided by their introduced dataset called CogQA. It contains complex questions that reduce into subquestions, similar to a chain-of-thought design. They conduct experiments across several model architecture and model sizes, highlighting how LLMs internally organize reasoning processes.", "questions": "Questions\n\n- One of the cognitive functions, “Knowledge Recall” is defined as accessing stored factual or procedural knowledge from memory. However, datasets like CREAK from which the proposed dataset is sourced from contains multiple entities which the model may not be aware of, leading to faults in the analysis of this cognitive function. An example of Knowledge Recall from the training data, “What is the density characteristic of Palladium?”. This knowledge may not be known by every model, making it an unfair comparison.\n\n- The authors make the claim of layered functional organization wherein retrieval related heads cluster primarily in the middle layers while math-related heads appear more frequently in higher layers. Although this is apparent in Llama3.1-8B, the same trend is not too apparent in other models like Qwen3-8B and Yi-1.5-9B. Please provide a discussion and potential reasoning for the same.\n\n- Table 1 provides the accuracy change in a cognitive function after suppressing the attention heads for that particular task. However, performing cross task evaluation (i.e., suppressing cognitive heads for task like Retrieval and evaluating on Knowledge Recall) would provide more insights. Although a similar analysis in provided between lower and higher level cognitive functions, intra-level function relations would substantiate the Functional Clustering results presented in Section 4.3.\n\n- Section 4.4 discusses positive and negative interventions. However, no empirical results (e.g. change in model accuracy on Math tasks after positive intervention on Math related attention heads) are provided. Change in accuracy due to positive/negative interventions in attention heads corresponding to specific cognitive tasks would further solidify their analysis.", "limitations": "yes", "rating": 4, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 3, "originality": 4, "strengths_and_weaknesses": "Strengths:\n\n- The paper is well written and provides a clear flow of thought. In multiple instances, they draw parallels with human cognition.\n- The experiment design is well thought of and evaluates their approach across multiple model families and model sizes.\n\nWeaknesses:\n\n- The paper provides no discussion on how the cognitive functions were decided. Moreover, adding examples for each of the cognitive functions would bring more clarity.\n- Value of hyperparameter k used in top-k most important tokens selection is missing\n- Table 2 doesn’t specify which model is evaluated", "was_revised": true, "final_justification": "The authors resolved my concerns regarding the theory behind cognitive functions, change in results due to positive and negative interventions and the effect of interventions w.r.t model size. I would like to maintain my score. This is a technically solid paper, with good evaluation. But similar works have been proposed in the past and the impact of this work is limited to sub-field of LLM reasoning.", "rebuttal": "Dear Reviewer,\n\nThank you for your thoughtful comments and suggestions. They have been instrumental in improving the clarity and quality of our paper. Below, we provide detailed, point-by-point responses to your concerns.\n\n**Q1. The paper provides no discussion on how the cognitive functions were decided. Moreover, adding examples for each of the cognitive functions would bring more clarity.**\n\n**A.** Inspired by research in cognitive science [1-2] and recent findings on the roles and mechanisms of attention heads in large language models (LLMs) [3], we focus on eight representative cognitive functions grouped into two categories: low-level and high-level functions, each relevant to complex reasoning tasks. To better align with human cognitive processes, we incorporate human annotations to guide the function labeling for subquestions. \n\nWe appreciate your suggestion regarding the addition of examples. To improve clarity, we will include representative examples for each cognitive function in the Appendix. Furthermore, the complete CogQA dataset—with all annotated subquestions—is provided in the supplementary materials for reference.\n\n[1] John R Anderson. Rules of the mind. Psychology Press, 2014.\n\n[2] Adele Diamond. Executive functions. Annual review of psychology, 64(1):135–168, 2013\n\n[3] Zheng, Z., et al. \"Attention heads of large language models: A survey. arXiv 2024.\" arXiv preprint arXiv:2409.03752.\n\n**Q2. Value of hyperparameter k used in top-k most important tokens selection is missing.**\n\n**A.** In our experiments, we set k = 5. This value was chosen to balance capturing enough key information while minimizing the inclusion of less relevant or noisy tokens. We will clarify this setting in the revised manuscript.\n\n**Q3. Table 2 doesn’t specify which model is evaluated.**\n\n**A.** Thank you for pointing this out. Table 2 evaluates on llama3.1\\_8B\\_instruct model. We will carefully review and clarify this throughout the final version.\n\n**Q4. One of the cognitive functions, “Knowledge Recall” is defined as accessing stored factual or procedural knowledge from memory. However, datasets like CREAK from which the proposed dataset is sourced from contains multiple entities which the model may not be aware of, leading to faults in the analysis of this cognitive function. An example of Knowledge Recall from the training data, “What is the density characteristic of Palladium?”. This knowledge may not be known by every model, making it an unfair comparison.**\n\n**A.** Our goal is to explore the functional roles of attention heads rather than to evaluate whether a model possesses specific factual knowledge. For example, when asked, “What is the density characteristic of Palladium?”, even if the model does not know the answer, it still engages in the process of attempting to recall factual knowledge—much like a human trying to retrieve uncertain information. This process inherently activates attention heads related to the Knowledge Recall function. Of course, incorrect answers may correspond to insignificant cognitive head activation, so the straightforward approach is to exclude them from the analysis. However, according to our observations, such cases are relatively rare. Overall, our analysis focuses on identifying the functional roles involved in the cognitive process and does not involve comparisons between different LLMs.\n\n**Q5. The authors make the claim of layered functional organization wherein retrieval related heads cluster primarily in the middle layers while math-related heads appear more frequently in higher layers. Although this is apparent in Llama3.1-8B, the same trend is not too apparent in other models like Qwen3-8B and Yi-1.5-9B. Please provide a discussion and potential reasoning for the same.**\n\n**A.** As shown in Fig. 6 for Qwen3-8B and Fig. 8 for Yi-1.5-9B, retrieval-related heads are predominantly located in the middle layers, with some also appearing in the deeper layers. This variation may stem from differences in pretraining strategies and data across LLMs. Importantly, our main claim is not that the exact layer locations of functional heads are fixed across models, but rather that cognitive functions tend to exhibit layered preferences, which can be consistently observed across architectures.\n\n**Q6.  Table 1 provides the accuracy change in a cognitive function after suppressing the attention heads for that particular task. However, performing cross task evaluation (i.e., suppressing cognitive heads for task like Retrieval and evaluating on Knowledge Recall) would provide more insights. Although a similar analysis in provided between lower and higher level cognitive functions, intra-level function relations would substantiate the Functional Clustering results presented in Section 4.3.**\n\n**A.** Following your suggestions, we conduct experiments where we mask the retrieval heads during the evaluation of knowledge recall (Recall), and conversely, mask knowledge recall heads during the evaluation of retrieval performance. As shown in the following table, intra-level masking of cognitive heads leads to a certain degree of performance degradation. However, as shown in Table 2, masking either Retrieval or Knowledge Recall heads results in a complete failure on high-level reasoning tasks, with performance dropping to zero.\n\nThis stark contrast suggests that low-level cognitive heads not only play a crucial role within their own level but also significantly influence high-level functions through the chain-of-thought (CoT) reasoning process. \n\n\n**Table. Intervention results (%) of different cognitive heads and random heads across Retrieval and Recall functions.**\n|Model|Inter_Head|Retrieval (comet)|Retrieval (acc)|Recall (comet)|Recall (acc)|\n|-|-|-|-|-|-|\n|LLama3.1-8B|random|90.83|84.71|87.85|83.84|\n|LLama3.1-8B|retrieval|44.96|8.24|72.05|33.33|\n|LLama3.1-8B|recall|86.79|75.29|56.93|38.38|\n|Qwen3-8B|random|92.81|75.29|89.90|53.54|\n|Qwen3-8B|retrieval|59.19|38.24|79.26|57.58|\n|Qwen3-8B|recall|83.31|71.18|64.81|30.30|\n\n\n**Q7. Change in accuracy due to positive/negative interventions in attention heads corresponding to specific cognitive tasks would further solidify their analysis.**\n\n**A.** To further quantify the results, we conduct experiments on two tasks: a math task and a retrieval task. For the math task, we use 100 samples from GSM8K (denoted as GSM8K\\_100), and for the retrieval task, we use an extractive QA dataset consisting of 49 samples, as described in [1]. The Extractive\\_QA dataset is constructed by prompting GPT-4o to generate question–answer pairs from a given paragraph, where the answer is an original sentence extracted from that paragraph.\n\nFor negative intervention, we suppress the Math Calculation heads when evaluating on GSM8K\\_100, and the Retrieval heads when evaluating on Extractive\\_QA. As shown in the following table, this suppression leads to a significant drop in performance on both tasks across different models, demonstrating that these heads exhibit consistent and transferable functional roles.\n\n**Table. LLM performance on GSM8k_100 and Extractive_QA by negative intervention. Inter. is the negative intervention.**\n|Dataset|Head|Llama3.1-8B|Llama3.2-3B|Qwen3-8B|Qwen3-4B|\n|-|-|-|-|-|-|\n|GSM8k\\_100|Math (Base)|82|64|94|91|\n|GSM8k\\_100|Math (Inter.)|38|34|34|37|\n|Extractive\\_QA|Retrieval (Base)|57.14|36.73|57.14|51.02|\n|Extractive\\_QA|Retrieval (Inter.)|0|0|14.29|12.24|\n\nFor positive intervention, we intervene on the cognitive heads in LLMs using positive activation adjustments. The experimental results in the following table show that enhancing the activation of retrieval heads along their corresponding functional directions leads to improved performance on the retrieval task.\n\nFor the math task, each question often involves multiple cognitive functions. We observe that in some cases, the numerical computations are correct, but the model’s understanding of the question is semantically flawed. By positively intervening on either the Math Calculation or Semantic Understanding heads, we observe consistent improvements in performance. This indicates the importance of developing intervention methods that can simultaneously adjust multiple cognitive functional heads which collaborate during complex tasks. We leave this as an avenue for future work.\n\n**Table. LLM performance on GSM8k_{100} and Extractive_QA by positive intervention. Inter. is the positive intervention.**\n|Dataset|Head|Llama3.1-8B|Llama3.2-3B|Qwen3-8B|Qwen3-4B|\n|-|-|-|-|-|-|\n|GSM8K|Math(Base)|82|64|91|94|\n|GSM8K|Math(Inter.)|84|66|92|94|\n|GSM8K|Semantic(Base)|82|64|91|94|\n|GSM8K|Semantic(Inter.)|84|65|93|94|\n|Extractive_QA|Retrieval(Base)|57.14|36.73|57.14|51.02|\n|Extractive_QA|Retrieval(Inter.)|63.26|44.90|61.22|69.38|\n\n[1] Wu, Wenhao, et al. \"Retrieval head mechanistically explains long-context factuality.\" arXiv preprint arXiv:2404.15574 (2024)."}]

Author rebuttal:

Meta-review: This work mechanistically analyzes roles and behaviors of attention heads inside LLMs and introduces a dataset to divide complex questions into multiple steps with retrieval or logical reasoning cognitive functions. Furthermore, they locate universal sparse, interactive, and hierarchical cognitive attention heads across different LLMs. In addition, removing or augmenting those heads can affect reasoning tasks.

The main strength of this work is to conduct a comprehensive and mechanistic analysis of cognitive attention heads on their proposed dataset across different LLMs. In addition, the following located attention head intervention experiment demonstrates how those heads affect reasoning tasks.

For weakness, several reviewers raised concerns about the cognitive question decomposition setting, the vague definition of distinct cognitive functions, and the analysis of how attention heads behave (layered functional organization) by masking cognitive function relevant or random masking leading to different performance drops.

In general, different reviewers’ questions were well solved regarding those concerns above during the discussion.

Based on reviewer-author discussions, I recommend that authors add those extra explanations into the camera-ready version, such as how LLMs process the CoT decomposition evaluated by another model, what if reasoning structure is not sequential, new experimental results for probing, how to evaluate and select semantically relevant tokens, etc.

To summarize, I would recommend this work as poster for NeurIPS.
Decision: Accept (poster)