Interpreting and Exploiting Functional Specialization in Multi-Head Attention under Multi-task Learning
Chong Li, Shaonan Wang, Yunhao Zhang, Jiajun Zhang, Chengqing Zong
Abstract
Transformer-based models, even though achieving super-human performance on several downstream tasks, are often regarded as a black box and used as a whole. It is still unclear what mechanisms they have learned, especially their core module: multi-head attention. Inspired by functional specialization in the human brain, which helps to efficiently handle multiple tasks, this work attempts to figure out whether the multi-head attention module will evolve similar function separation under multitasking training. If it is, can this mechanism further improve the model performance? To investigate these questions, we introduce an interpreting method to quantify the degree of functional specialization in multi-head attention. We further propose a simple multi-task training method to increase functional specialization and mitigate negative information transfer in multi-task learning. Experimental results on seven pre-trained transformer models have demonstrated that multi-head attention does evolve functional specialization phenomenon after multi-task training which is affected by the similarity of tasks. Moreover, the multi-task training strategy based on functional specialization boosts performance in both multi-task learning and transfer learning without adding any parameters. 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 c9e7c084-1de6-41da-aa92-30802419e02aCited by top-tier papers7
- Fine-Grained Activation Steering: Steering Less, Achieving MoreZijian Feng, Tianjiao Li, Zixiao Zhu, Hanzhang Zhou et al.ICLR 2026 · 6 citations
- SEEKR: Selective Attention-Guided Knowledge Retention for Continual Learning of Large Language ModelsJinghan He, Haiyun Guo, Kuan Zhu, Zihan Zhao et al.EMNLP 2024 · 4 citations
- Reallocating Attention Across Layers to Reduce Multimodal HallucinationHaolang Lu, Bolun Chu, WeiYe Fu, Guoshun Nan et al.CVPR 2026 · 3 citations
- HeadHunt-VAD: Hunting Robust Anomaly-Sensitive Heads in MLLM for Tuning-Free Video Anomaly DetectionZhaolin Cai, Fan Li, Ziwei Zheng, Haixia Bi et al.AAAI 2026 · 1 citation
- Semantics-Adaptive Activation Intervention for LLMs via Dynamic Steering VectorsWeixuan Wang, Jingyuan Yang, Wei PengICLR 2025
Builds on18
- Language Models are Few-Shot LearnersTom B. Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah et al.NeurIPS 2020 · 64,255 citations
- An Image is Worth 16x16 Words: Transformers for Image Recognition at ScaleAlexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn et al.ICLR 2021 · 21,477 citations
- DeBERTaV3: Improving DeBERTa using ELECTRA-Style Pre-Training with Gradient-Disentangled Embedding SharingPengcheng He, Jianfeng Gao, Weizhu ChenICLR 2023 · 394 citations
- Self-Attention Attribution: Interpreting Information Interactions Inside TransformerYaru Hao, Li Dong, Furu Wei, Ke XuAAAI 2021 · 282 citations
- Understanding and Improving Information Transfer in Multi-Task LearningSen Wu, Hongyang R. Zhang, Christopher RéICLR 2020 · 183 citations
Related papers
- The Stem Cell Hypothesis: Dilemma behind Multi-Task Learning with Transformer EncodersHan He, Jinho D. ChoiEMNLP 2021 · 111 citations
- Neuron Specialization: Leveraging Intrinsic Task Modularity for Multilingual Machine TranslationShaomu Tan, Di Wu, Christof MonzEMNLP 2024
- What's in Your Head? Emergent Behaviour in Multi-Task Transformer ModelsMor Geva, Uri Katz, Aviv Ben-Arie, Jonathan BerantEMNLP 2021
- Heads up! Large Language Models Can Perform Tasks Without Your Instruction via Selective Attention Head MaskingSenyu Han, Hongchuan Zeng, Kai Yu, Lu ChenICML 2025
- Resolving Token-Space Gradient Conflicts: Token Space Manipulation for Transformer-Based Multi-Task LearningWooseong Jeong, Kuk-Jin YoonICCV 2025 · 2 citations
