Roles and Utilization of Attention Heads in Transformer-based Neural Language Models (2020)

Abstract: Sentence encoders based on the transformer architecture have shown promising results on various natural language tasks. The main impetus lies in the pre-trained neural language models that capture long-range dependencies among words, owing to multi-head attention that is unique in the architecture. However, little is known for how linguistic properties are processed, represented, and utilized for downstream tasks among hundreds of attention heads inside the pre-trained transformer-based model. For the initial goal of examining the roles of attention heads in handling a set of linguistic features, we conducted a set of experiments with ten probing tasks and three downstream tasks on four pre-trained transformer families (GPT, GPT2, BERT, and ELECTRA). Meaningful insights are shown through the lens of heat map visualization and utilized to propose a relatively simple sentence representation method that takes advantage of most influential attention heads, resulting in additional performance improvements on the downstream tasks.

## introduction sentence encoders in transformer architectures as in gpt, bert @xcite @xcite and electra @xcite have shown promising results on various natural language understanding (nlu) tasks, such as question answering, text entailment and natural language inference (nli) @xcite , owing to their pre-training capabilities in modeling languages. the pre-training effects of the transformer-based approaches are known to be crucial for obtaining superior performance in various downstream nlu tasks. the main impetus lies in capturing longrange dependencies among words obtainable with bidirectional learning and self-attention @xcite and sufficiently varied corpora of a large quantity @xcite . despite all the recent successes of the transformer-based models, little is known for how linguistic properties are processed and represented internally when the architectures are used. given that self-attention heads are unique in the family of transformer architectures, we attempt to answer the question of how basic linguistic properties are captured with the attention heads across the models and used for downstream tasks. once we figure out the roles of attention heads in "storing" various linguistic properties, we should be able to modulate them to maximize the performance of the downstream tasks. given the motivation, we analyze several publicly available pre-trained transformer encoders (bert, gpt, gpt2, and electra) trained with different model capacities ranging from 144 to 384 attention heads and 12 to 24 layers. considering the output vector from each attention head of an encoder as a mini-sentence embedding, we examine whether certain linguistic properties are "stored" in embeddings among ten sentence probing tasks @xcite that cover surface, syntactic, and semantic information and require different linguistic properties (e.g. the depth of a parsed sentence). each of the probing tasks is treated as if it were a downstream task for the examination; a classifier is attached for each of the primitive linguistic properties. in order to predict the depth of the parse tree, for example, an n-ary classifier is connected, where n is the number of possible depths. in order to aggregate and summarize the performance results out of all the attention heads, we construct an accuracy heat map for each probing task, where the patterns across layers and attention heads can be recognized easily. by examining the heat map, we can observe the patterns of how the attention heads contribute to the accuracy of each probing task, including whether an individual attention head is contributing to multiple linguistic features together or just specialized for a particular feature. aiming at producing improved sentence representation, we use the analysis result that allows for selecting and concatenating the outputs of superior attention heads. the sentence representations from the hidden layers and the top-n attention heads are compared to check whether using only influential attention heads selectively could help certain downstream tasks. this attempt is in contrast with the common approach of using the output of the last layers of a transformer-based encoder as the representation that is fed into a downstream task. our hypothesis is those final representations from the top of the transformer-based encoders might not be the best not only in carrying primitive linguistic properties of the language but also for downstream tasks. all the source code is publicly available foot_0 . the major contribution of our research is twofold: 1) we suggest an analysis method which helps understand where linguistic properties are learned and represented along attention heads in transformer architectures and 2) we show that using analysis results, attention heads can be maximally utilized for performance gains during the fine-tuning process on the downstream tasks and for capturing linguistic properties. ## related work several studies looked into the representations learned by a neural network for various language properties @xcite . a similar line of work focused on learned linguistic features inside the word and sentence embeddings. they used downstream tasks in order to probe surface information, syntactic and semantic information @xcite . some recent work looked inside the sentence encoders with various depths, by analyzing the hidden states at a layer-level @xcite and even at a neuron-level @xcite . @xcite attempted to understand linguistic characteristics learned in a series of pre-trained encoder models by jointly analyzing their behaviors across different nlp tasks. for studying attention mechanisms, there have been two streams of work: 1) visual analysis of attention weights to associate various functionalities and 2) analysis of the characteristics of the output representations from individual attention heads. for the first category, vig and jesse (2019) developed a visualization tool for attention weights of bert and gpt2 and identified notable heads but without any quantitative analysis. @xcite showed the extent to which attention agrees with traditional alignments in neural machine translation (mt). @xcite and @xcite on the other hand argued that attention rarely provides an explanation of model predictions. they showed through attention map analysis that attention weights frequently are not correlated with other measures of feature importance. for the second category that attempts to discover various roles attention heads play, @xcite studied the characteristics of individual attention heads from the transformer, pretrained with an mt task and evaluated on a limited suite of linguistic tasks, pos tagging, ner tagging, and chunking. similarly, @xcite showed that some attention heads are specialized for dependency parsing and coreference resolution. @xcite showed through an ablation study that some dedicated heads have a significant role in mt and revealed the dynamics of attention heads during the training process. @xcite provided a method to identify the major role of each attention head in a transformer model trained for mt. the two studies are limited to mt and a particular transformer model, bert. unlike the recent studies mentioned above, our analysis is more comprehensive in its scope for generalizability. the analysis probes a variety of surface, syntactic, and semantic information at sentence levels with different transformer encoders pre-trained on language modeling tasks. more importantly, our work goes beyond an analysis and suggests a method of utilizing the analysis results for performance gains on several downstream tasks. it not only proposes a simple yet new method for the downstream tasks but also validates the analysis of the attention mechanisms. to the best of our knowledge, this is the first attempt to do an in-depth analysis of the seven recent pre-trained encoders for their internal workings in handling linguistic features, not to mention the newly proposed way for improvements on the downstream tasks. l and h denote the number of stacked encoding layers and the number of attention heads packed within each encoding layer, respectively. ## methodology consider a transformer-based encoder m , typically with a stack of l identical layers, each of which makes use of multi-head self-attention, and a two sub-layer feed-forward network coupled with layer normalization and residual connection (see figure 1a ). for a given input sequence @xmath0, x 2 , . . . , x n ), each word embedding xis concatenated with a positional encoding and fed into the encoder layer to produce an attention head output h i,@xmath1, respectively. then a series of sub-layers produce hidden states of the i-th encoding layer z @xmath2. for all pre-trained encoders, d @xmath3. since the transformer-based encoders encode the input sequence word by word, z i and h i,j are produced individually for given word x k along the input sequence x. in order to produce a sequencelevel representation, we need to select one of the input representations of the sequence. since the selection method depends on the chosen pre-trained model, we defer a detailed discussion to section 4.1. for now, we assume z i and h i,j have been already determined with the specific word chosen from the input sequence and consider it as the sentence-level representation. ## conclusion while recent research demonstrated the capability of the transformer-based encoders for generating rich sentence representations, the roles of individual self-attention heads were hardly unknown. furthermore, little is known for whether and how we can utilize them for better capturing linguistic properties and eventually improving the performance of downstream tasks for which the embeddings are constructed. one of the major contributions of this paper is to fill the void by inspecting where and how the attention heads are "trained" internally for classi-fication tasks corresponding to different linguistic properties and for the downstream tasks. the analysis results clearly show a tendency through the comprehensive heat maps that syntactic and semantic information is mainly handled from the lower layers to the upper layers. we also showed that understanding the roles of attention heads in handling task-specific information can help to develop adaptive sentence representations, by selecting influential attention heads and testing them for the three downstream tasks. the additional performance gains obtained by the simple method show that this approach of using the anatomy of the transformer models and the attention heads is promising in utilizing expensive pre-trained transformer models to their maximal extent. furthermore, we explored how the hundreds of attention heads underwent performance variation during the fine-tuning process on the downstream tasks, revealing the internal behaviors with the proposed analysis method. the analysis of syntacticsemantic score distributions revealed that individual attention heads capture both syntactic and semantic information. it also showed that the amount of both syntactic and semantic information handled by the heads vary from layer to layer, sometimes showing that the last layer contributes much less especially with large models. while the empirical results are strong, additional work remains to further our understanding of the internal workings of the transformer architecture and its role in building such strong language models for a variety of tasks. immediate attention should be paid to the investigation of how heat maps would vary during the extensive pre-training so that we have a better understanding of the dynamics of the learning processes.