How Far Does BERT Look At: Distance-based Clustering and Analysis of BERT ’s Attention (2020)

Abstract: Recent research on the multi-head attention mechanism, especially that in pre-trained models such as BERT, has shown us heuristics and clues in analyzing various aspects of the mechanism. As most of the research focus on probing tasks or hidden states, previous works have found some primitive patterns of attention head behavior by heuristic analytical methods, but a more systematic analysis specific on the attention patterns still remains primitive. In this work, we clearly cluster the attention heatmaps into significantly different patterns through unsupervised clustering on top of a set of proposed features, which corroborates with previous observations. We further study their corresponding functions through analytical study. In addition, our proposed features can be used to explain and calibrate different attention heads in Transformer models.

## introduction with the rapid development of neural network in nlp tasks these year, the transformer @xcite that uses multi-head attention (mha) mechanism is one recent huge leap @xcite . it has become a standard building block of recent nlp models. the transformer-based bert @xcite model further advances the model accuracy by introducing pre-training and has reached the state-of-the-art performance on many nlp tasks. beyond the general explanation which attributes the effectiveness of bert model to its capability of long-range dependency and contextual embeddings, more detailed analysis of the mha mechanism in bert still remains an active research topic @xcite @xcite . @xcite ; @xcite ; @xcite explore the linguistic information or importance of bert by inspecting hidden state embeddings. @xcite ; @xcite further investigate possible attention mechanism designs. several previous works on attention interpretability have analyzed the function and behavior of attention heads. @xcite finds that attention heads within a same layer tend to function similarly, and vast redundancy could be found in the attention heads. @xcite provides the study on relating the behavior of each head to linguistic merits such as dependency parser. @xcite manually annotates attention heatmaps and trains a cnn to classify attention heads. given these various analytical findings, from an algorithmic view, we are not aware of any research that provides a clear and reliable method to classify the attention heads without extra human intervention. in this work, we provide a simple set of features that can be used to reliably disentangle the attention heads into different categories through clustering method. different from heuristic analytics that looks at single inputs, our proposed attention head features are extracted corpus-wide and is thus more comprehensive in its discovered patterns. interestingly, our algorithmically discovers patterns that match well with previous study @xcite . on the other hand, our method could easily be scaled up to large datasets and is generalizable to other similar multi-head-attention-based models. ## background and motivation the bert model @xcite includes multiple layers of transformer @xcite encoders. its core component is multi-head attention mechanism (mha). given a sequence of input embeddings, the output contextual embedding is composed by the input sequence with different attention at each position. the attention weight is calculated as following, given the great success of bert-like models, researchers start to explore the working mechanisms of multi-head attention component. one common approach is to visualize the attention weight matrix @xcite @xcite , which often finds that attention heads in bert demonstrate several specific patterns. for example, some attention heads show a dominant stripe pattern on diagonal direction while some others show a dominant vertical stripe pattern. there are also attention heads with a relatively homogeneous distribution of attention weights. in this work, we refer this patterns as dense, vertical or diagonal with some samples in the top part of fig. 1 . however, these previous works rely on simple heuristic rules or human annotation, which makes them not scalable and error-prone. in this work, we propose a scalable and efficient unsupervised learning approach that automatically clusters different attention heads, which lets us systematically study their different roles in nlp tasks. ## distance feature based on the attention pattern findings in previous works, we design a distance feature that enables unsupervised clustering of attention heads. for each input token, we first accumulate its attention weights to the special tokens including '[cls]', '[sep]', ',' and '.', because the vertical stripe pattern shows strong focus on these tokens. we then accumulate its attention weights to nearby tokens within a window size, which lets us distinguish the diagonal stripe with strong locality and dense pattern with weak locality. for a input sample with length of l, we divide the sample to n windows, each of which has a size of l/n . eq. 2 calculates the window feature f (i) of i th window. the special token attention and window attention are exclusive and we normalize both them by the input sequence length l. as such, their summation for a give input token remains one. for an input sample, attention weight matrix from all tokens are averaged to get a low-dimensional representation of an individual head. with 16 windows, we can represent an attention head with a 21 (16 windows + self + 4 special tokens) dimensional feature for further analysis. ## clustering we now describe how we leverage the 21-dimensional distance feature to apply an unsupervised attention head clustering. we study attention behavior in the bert-large model foot_0 . we randomly pick 1,000 samples from the squad dataset @xcite and extract the averaged distance feature across all selected samples, which is used to represent an attention head. to verify if there exists clustering behavior in bert, we first use a non-linear dimension reduction algorithm t-sne @xcite to visualize all 384 heads (24 layers × 16 heads per layer). the result in fig. 2 clearly demonstrates four clusters. we then use k-means clustering algorithm (hartigan and wong, 1979) with 4 clusters and the result is shown as the colors. fig. 1 shows two examples (top) and the distance feature box-plot (bottom) for each cluster. we also show the distribution of each type of attention head across different layers in fig. 3 . our results show that an attention head shows a consistent weight pattern across different input samples and there exists a strong clustering behavior across different layers. we describe the key findings in details. dense: this type has dense attention patterns (fig. 1 (a) top). fig. 1(a ) bottom shows that the heads in this type have the highest attention at long range windows and relatively small attention weights to the special tokens. they account for 26.04% of total heads and mainly appear at the beginning layers. diagonal: this type has dominant diagonal patterns and a few vertical stripes (fig dense&vertical: this type mixes the dense and vertical patterns (fig. 1(c ) top). its distance features have large values at long range windows and special tokens. they account for 34.11% of total heads and mainly appear at middle layers. vertical: this type has mostly strong vertical patterns (fig. 1(d ) top). its distance features have large values at special tokens, which are much higher than other types. they account for 13.28% of total heads and mainly appear at higher layers. as a matter of fact, it is possible to exploit the proposed distance feature with other unsupervised clustering algorithms. our work chooses to use the k-means clustering algorithm as a representative one. with gaussian mixture model clustering @xcite , we also get similar results. as such, we choose to use four clusters in the following experiment based on the t-sne visualization result. moreover, we find the four clusters are human interpretable. we still get stable clustering results with 3 or 5 clusters, but the patterns can not be distinguished by human-being. we measure the stability of clustering and generality on a pre-trained language model and downstream tasks including question answering, sentence pair similarity regression and natural language inference. for the aforementioned tasks, we use datasets that include squad @xcite , mnli @xcite @xcite respectively. we also extract the features on finetuned bert-large models accordingly. all of these training settings are inline with @xcite and @xcite . we evaluate the stability of clustering following von luxburg (2010). the idea is to obtain multiple models with partial training data and evaluate the consistent samples among all fitted models. we randomly split 50% training data to fit clustering models and predict on the whole dataset. the stability of clustering algorithm is calculated as the proportion of input samples that is consistent on all models fitted with the random partial training set split. we report the clustering stability of 10 times clustering in tbl. 1. our result demonstrates that over 95% attention heads are clustered stably with a limited amount of data. this shows that the proposed method is generally applicable and scalable. the detailed results are shown in fig. 4 , which demonstrates similar type distribution among different models. in specific, all the four models share similar head type distribution described in fig. 3 . empirically, we find many heads have the same cluster after fine-tuning process. however, some heads in the highest layers also change. we suggest that higher layers are more relevant to task specific knowledge while lower layers handles more general linguistic features. particularly, pre-trained model and squad-fine-tuned model have more dense&vertical types at highest layer, which we will show to be responsible for global information in sec. 5. this is explained by the fact that question answering and language modeling are more difficult tasks as they use token embeddings from all positions for classification. on the contrast, the other two tasks only use the special token @xcite . ## experiments based on the attention clustering results, we now study which attention head type(s) is important and what range of information different head types process. in specific, we manipulate the different attention heads in bert-large model and observe the resulted accuracy impact on squad v1.1 development set. it should be mentioned that all these experiments are directly ablation without fine-tuning on the model parameters. ## conclusion in this work, we propose a feasible and scalable unsupervised clustering method to classify the attention heads. the distance based feature well captures the observations of previous works and is utilized to conduct stable unsupervised clustering. we believe our method is helpful for attention interpretability study. with such taxonomy, we further apply analytic experiments to explore the function of each behavior according to distance. we are also looking forward to further analysis the behavior and function of variant patterns with probing task datasets @xcite and analytic tools @xcite as our next plan. besides, there are several recent works focusing on the optimization of over-parameterized mha mechanism @xcite @xcite . our results reveal the important types and components of each type empirically. with our explorations, we hope our insights could cast light on novel design of interpretable attention mechanism.