Probing and Boosting Large Language Models Capabilities via Attention Heads
Dezhi Zhao, Xin Liu, Xiaocheng Feng, Hui Wang, Bing Qin
Abstract
Understanding the internal origins of capabilities in large language models (LLMs) is crucial for interpretability and efficient adaptation. However, the emergence of specific capabilities remains poorly understood, as most existing approaches rely on external signals (e.g., performance shifts or gradient similarities) with limited structural grounding. To address these issues, this paper proposes a lightweight and highly interpretable approach that links LLM capabilities to internal components by identifying correspondences at the level of attention heads. Specifically, we first define five fundamental capabilities, namely Mathematical Reasoning, Reading Comprehension, Commonsense Reasoning, Scientific Reasoning, and Professional Expertise, and employ probing techniques to detect the attention heads most predictive of each, thereby establishing capability-head mappings. For targeted instruction tuning, complex tasks are decomposed into these fundamental capabilities, and training data are selected accordingly. Experiments on LLaMA3.1-8B and Qwen2.5-7B show over 70% discrimination accuracy in identifying capabilities. On MMLU and BBH, our method improves accuracy by 1 to 1.5 points over the gradient-based method LESS and by 5 to 6 points over other intermediate-state baselines 1 .
Ask about this paper
Your agent reads all of it.
Lune indexed this paper to the last equation, along with the top-tier papers that cite it. Ask a question and the answer quotes them.
Your agent calls
Luneget_paper_fulltext
Free to start. No credit card required.
Terminal
Install the CLIlune papers fulltext aafcdf41-2c65-4f9d-b58c-bad8d5f7af0cCited by top-tier papers1
Ask how each one uses itBuilds on15
- Measuring Massive Multitask Language UnderstandingDan Hendrycks, Collin Burns, Steven Basart, Andy Zou et al.ICLR 2021 · 7,905 citations
- Finetuned Language Models are Zero-Shot LearnersJason Wei, Maarten Bosma, Vincent Y. Zhao, Kelvin Guu et al.ICLR 2022 · 4,966 citations
- TruthfulQA: Measuring How Models Mimic Human FalsehoodsStephanie Lin, Jacob Hilton, Owain EvansACL 2022 · 3,228 citations
- Learn to Explain: Multimodal Reasoning via Thought Chains for Science Question AnsweringPan Lu, Swaroop Mishra, Tanglin Xia, Liang Qiu et al.NeurIPS 2022 · 2,727 citations
- Inference-Time Intervention: Eliciting Truthful Answers from a Language ModelKenneth Li, Oam Patel, Fernanda B. Viégas, Hanspeter Pfister et al.NeurIPS 2023 · 1,549 citations
Related papers
- Compressed Sensing for Capability Localization in Large Language ModelsAnna Bair, Yixuan Xu, Mingjie Sun, Zico KolterICML 2026 · 1 citation
- BRIEF: Bi-level Coreset Selection for Efficient Instruction Tuning in LLMsChaoyuan Shen, Chi Zhang, Chengliang Chai, Jiacheng Wang et al.VLDB 2026 · 2 citations
- Cognitive Mirrors: Exploring the Diverse Functional Roles of Attention Heads in LLM ReasoningXueqi Ma, Jun Wang, Yanbei Jiang, Sarah M. Erfani et al.NeurIPS 2025 · 5 citations
- Interpreting and Improving Large Language Models in Arithmetic CalculationWei Zhang, Chaoqun Wan, Yonggang Zhang, Yiu-ming Cheung et al.ICML 2024 · 47 citations
- Neuron-Aware Data Selection in Instruction Tuning for Large Language ModelsXin Chen, Junchao Wu, Shu Yang, Runzhe Zhan et al.ICLR 2026 · 2 citations
