Research map / Large language models / Language, safety and interpretability
Cognitive Mirrors: Exploring the Diverse Functional Roles of Attention Heads in LLM Reasoning
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Introduces CogQA, a dataset of cognitively annotated subquestion decompositions, and uses multi-class probing with gradient×activation attribution to identify sparse "cognitive heads" in six LLMs, validated by head masking and activation steering.
Problem
Understanding which attention heads in LLMs take on specialized cognitive functions during complex multi-step reasoning, which prior work has studied mainly on simple or synthetic tasks.
Approach
Decompose reasoning questions into chain-of-thought subquestions labeled with eight cognitive functions, train a multi-class MLP probe on per-head activations, and attribute class importance to heads via gradient×activation; then validate heads causally by masking and by steering along functional directions.
Key results
- Masking the identified cognitive heads for Retrieval drops Llama3.1-8B retrieval accuracy from 84.71% (random) to 8.24%, while random masking causes only marginal degradation.
- Probe test accuracy on CogQA ranges from 75.18% (Yi-1.5-6B) to 84.71% (Qwen3-8B).
- Positive activation steering on Math Calculation heads raises Llama3.1-8B GSM8K_100 accuracy from 82 to 84 and Qwen3-4B from 91 to 92; on Extractive_QA, Qwen3-4B rises from 51.02 to 69.38 with retrieval-head steering.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Mask Retrieval heads instead of leaving the model unmodified (ablation) | Llama3.1-8B · Extractive_QA | accuracy (%): 57.14 -> 0 | Hurt |
| Enhance Retrieval heads along their functional direction | Llama3.1-8B · Extractive_QA | accuracy (%): 57.14 -> 63.26 | Helped |
| Mask Retrieval heads instead of leaving the model unmodified (ablation) | Qwen3-8B · Extractive_QA | accuracy (%): 57.14 -> 14.29 | Hurt |
| Enhance Retrieval heads along their functional direction | Qwen3-4B · Extractive_QA | accuracy (%): 51.02 -> 69.38 | Helped |
| Mask Math Calculation heads instead of leaving the model unmodified (ablation) | Llama3.1-8B · GSM8K_100 | accuracy (%): 82 -> 38 | Hurt |
| Enhance Math Calculation heads along their functional direction | Llama3.1-8B · GSM8K_100 | accuracy (%): 82 -> 84 | Helped |
| Enhance Semantic Understanding heads along their functional direction | Qwen3-4B · GSM8K_100 | accuracy (%): 91 -> 93 | Helped |
| Mask cognitive heads rather than an equal number of random heads (ablation) | Llama3.1-8B · CogQA retrieval intervention | acc: 84.71 -> 8.24 | Hurt |
| Mask cognitive heads rather than an equal number of random heads (ablation) | Llama3.1-8B · CogQA mathematical calculation intervention | acc: 83.08 -> 66.17 | Hurt |
| Use top-k important answer tokens rather than first generated token to identify heads (ablation) | Llama3.1-8B · Retrieval intervention, 30 heads masked | Retrieval(acc): 83.53 -> 46.47 | Helped |
| Use top-k important answer tokens rather than full-answer average to identify heads (ablation) | Llama3.1-8B · Retrieval intervention, 30 heads masked | Retrieval(comet): 73.93 -> 70.05 | Helped |
| Mask Retrieval or Knowledge Recall heads during downstream high-order tasks (ablation) | Llama3.1-8B-instruct · Math, Inference, Logical Reasoning, Decision-Making | accuracy (%): 0.00; 66.67; 88.24; 93.10; 57.14; 76.92 | Mixed |
Surprising findings
- Masking Math Calculation heads drops GSM8K_100 accuracy for Llama3.1-8B from 82 to 38, but the paper notes arithmetic and language functions otherwise stay largely unaffected.
- Masking retrieval heads causes masked-head accuracy to reach zero for some models, e.g. Llama3.2-3B semantic and syntactic accuracy drop to 0.00 in the intervention table.
Limitations
- Each subquestion is annotated with a single cognitive function and each head is assumed to serve one function, though real heads may serve multiple functions.
- Only eight predefined cognitive functions are considered, and the dataset is small (570 main QA, 3,402 subQAC triplets).
- No error bars or statistical significance reported.
Experimental setup
Datasets: AQuA CREAK ECQA e-SNLI GSM8K CogQA
Benchmarks: CogQA test set GSM8K_100 Extractive_QA (49 samples)
Metrics: COMET BLEU ROUGE semantic similarity accuracy probe test accuracy
Baselines: random head masking
Builds on: chain-of-thought prompting probing classifiers head pruning importance
Closest papers by meaning
- Longer Context, Deeper Thinking: Uncovering the Role of Long-Context Ability in Reasoning (NeurIPS 2025 · similarity 0.89)
- Large Language Models Think Too Fast To Explore Effectively (NeurIPS 2025 · similarity 0.88)
- Who Reasons in the Large Language Models? (NeurIPS 2025 · similarity 0.88)
- Attention Illuminates LLM Reasoning: The Uncovered Preplan-and-Anchor Rhythm Enables Fine-Grained Policy Optimization (ICML 2026 · similarity 0.88)
- Heads up! Large Language Models Can Perform Tasks Without Your Instruction via Selective Attention Head Masking (ICML 2025 · similarity 0.88)
- Eliciting Reasoning in Language Models with Cognitive Tools (NeurIPS 2025 · similarity 0.88)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.