Interpreting and Exploiting Functional Specialization in Multi-Head Attention under Multi-task Learning (2023)

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 multi-tasking 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.

## introduction transformer, based on the multi-head attention module, has been the dominant model for downstream applications due to its impressive results @xcite @xcite . however, it is still being utilized as a whole black-box model, and little is known about the functions of each sub-module on the final prediction. simultaneously, although controversy still exists, there is overwhelming evidence that supports the idea of functional specialization in the human brain @xcite . such a functional specialization mechanism makes it easier for the human brain to handle multiple tasks and solve new problems. it can reuse existing resources and at the same time evolve specific regions to avoid the huge cost of redesigning. considering the benefits of functional specialization to human learning ability, it is interesting to explore whether a transformer model, especially its central module multi-head attention, would evolve a similar mechanism under multi-task training. if so, which factors will impact the degree of functional specialization in the multi-head attention module? and how to exploit this phenomenon to improve the generalization ability of transformerbased models? to investigate these questions, we first propose a method, called important attention-head pruning (iap), to quantify the degree of functional specialization in the multi-head attention of transformerbased models. iap first calculates the importance scores of each attention head on different tasks, then prunes the top important heads for each task to determine their impact on task performance. we apply our method to five different tasks with seven pre-trained transformers. results show that the multi-head attention module has evolved distinct functional specialization phenomena across different sizes of bert and pre-training methods. further quantitative analysis indicates that there is a negative correlation between task similarity and the functional specialization phenomenon. moreover, we propose a multi-task learning method, namely important attention-head training (iat), to promote the segregation of functions in the multi-head attention module by training only the most important part of attention heads for each task. experimental results on the glue dataset have demonstrated that our method alleviates the negative transfer among tasks and improves the performance of transformer-based models on both multi-task learning and transfer learning without additional parameters. to summarize, our main contributions are twofold: ## conclusions and future work in this paper, we conduct extensive dissociation experiments and observe that the brain-like functional specialization phenomenon does evolve in multi-head attention after dual-task or multi-task learning. furthermore, experimental results show that the performance and generalization ability of multi-task models can be improved by the multitask training method based on functional specialization. this work, inspired by neuroscience findings, studies the interpretation and improvement of neural networks, which we hope will promote more efforts on interdisciplinary work combining neuroscience and artificial intelligence. in the future, we plan to investigate more neural network modules that may arise the functional specialization phenomenon under multi-task learning. another direction is to design better methods exploiting this phenomenon to further improve multi-task learning models.