A Dynamic Head Importance Computation Mechanism for Neural Machine Translation (2021)

Abstract: Multiple parallel attention mechanisms that use multiple attention heads facilitate greater performance of the Transformer model for various applications e.g., Neural Machine Translation (NMT), text classification. In multi-head attention mechanism, different heads attend to different parts of the input. However, the limitation is that multiple heads might attend to the same part of the input, resulting in multiple heads being redundant. Thus, the model resources are under-utilized. One approach to avoid this is to prune least important heads based on certain importance score. In this work, we focus on designing a Dynamic Head Importance Computation Mechanism (DHICM) to dynamically calculate the importance of a head with respect to the input. Our insight is to design an additional attention layer together with multi-head attention, and utilize the outputs of the multi-head attention along with the input, to compute the importance for each head. Additionally, we add an extra loss function to prevent the model from assigning same score to all heads, to identify more important heads and improvise performance. We analyzed performance of DHICM for NMT with different languages. Experiments on different datasets show that DHICM outperforms traditional Transformer-based approach by large margin, especially, when less training data is available.

## introduction transformer based nmt systems perform well on multiple translation tasks @xcite . multi-head attention is a very important component of the transformer model @xcite . multiple heads improve performance compared to a single head, as they allow the model to jointly look at different subspaces, and hence capture enhanced features from sentences. for example, a head can capture positional information by attending to adjacent tokens, or it can capture syntactic information by attending to tokens in a particular syntactic dependency relation @xcite . however, the performance of the transformer-base model with 8 heads at each layer is only 1 bleu point higher than that of a similar model with just a single head at each layer @xcite . this is due to the fact that majority of the heads learn similar weights, and therefore, multiple heads attend to the same parts of the input. hence, most of the heads are redundant, leading to an increased computational complexity without improving performance. to avoid this redundancy, one approach is to prune the redundant heads based on certain importance score. in this work, we focus on designing an importance computation method to compute the importance score for each head. some recent work has analyzed the importance of heads by considering average attention weights of each head at some specific position @xcite . however, average of attention weights is a static measure of the head importance as it does not consider the varying importance of each head with respect to the input. the importance of a head is dynamic, as a head can be very important for a particular word, but can be less important for other words. thus, in this work, we propose a dynamic head importance computation mechanism (dhicm) to calculate the importance score for each head, and this can be later utilized to design a pruning strategy. our key idea is to apply a second level attention on the outputs of all heads, to dynamically calculate the importance score for each head, that varies with the input, while training. we also propose to add a new loss term to prevent our approach from assigning equal importance to all heads. note that we apply dhicm for both self attention heads and encoder-decoder attention heads present in the encoder and decoder of the transformer architecture. to evaluate the performance of our method, we considered multiple translation tasks with different language pairs such as hindi-english, belarusian-english, german-english. results show that dhicm achieves a much higher performance compared to the standard transformer model, particularly, in low-resource conditions where much less training data is available. moreover, dhicm requires only ∼d 2 additional parameters (d is the word embedding dimension), that is much less than the total number of parameters in the transformer base model. the transformer model has a large number of hyperparameters, due to which, it is computationally challenging to search for their optimal values. thus, much of the previous work used default values of the hyperparameters @xcite . however, these are not guaranteed to yield optimal performance on different datasets. grid search over all hyperparameters is computationally intensive due to the exponential number of combinations across all possible values. therefore, in this work, we perform grid search over a subset of hyperparameters, i.e., architecture hyperparameters and regularisation hyperparameters, and experiments show that the hyperparameter values obtained from our method yield significantly better performance compared to the default values. to summarize, our work makes the following major contributions: ## results t-optimal achieved 5.49 bleu score. thus, t-optimal outperformed t-base by 1.4 bleu points. moreover, dhicm achieved 6.29 bleu score, thus outperforming t-optimal by 0.8 bleu points. we also compared the performance of our method with the multi-hop multi-head attention model @xcite on wmt17-cs de-en dataset. we observed that dhicm outperforms @xcite by 1.77 bleu points. since the hyperaparameters for both t-optimal and dhicm are same, we can see that the gain in the performance of our method is due to the proposed second layer attention over the multihead attention. in addition, our proposed loss function (section 3.1) prevents the model from assigning the same importance to all heads. thus, we are able to filter more important heads. ## analysis our proposed approach dhicm outperforms tbase and t-optimal by a large margin in the low resource conditions. we further analyzed the performance of the baseline model and dhicm, and observed that dhicm learns better word alignment especially, in low resource conditions. one of the reasons for learning better alignment can be that for each word, all heads are not equally important. the second level attention that we designed in our model allows the tokens to pay more attention to the heads that capture more relevant information for translation. since the heads that are more relevant receive more attention, the parts of the input to which these heads attend, in turn receive more attention, and thus, the alignment becomes better. for example, providing more attention to the heads that capture the syntactic or semantic information, and relatively less attention to the heads that capture positional information. this justifies our hypothesis mentioned in section 3.1. we also verified this using the encoder-decoder attention distribution of the models shown in figure 1 (low resource conditions) and figure 2 (high resource conditions). the decoder of the transformer model uses the outputs of the encoder to generate the tokens in the target language. each generated token pays some attention to each token in the source language. the attention distribution matrix shows the attention paid by the generated tokens in the target sentence (rows) to the tokens in the source sentence (columns). in figure 1a and figure 2a , we can see that most of the tokens on the source side get similar attention for the baseline approach. moreover, the highest attention a source token receives is approximately 0.12 and 0.5 in figure 1a and figure 2a , respectively. this implies that the most important source token for translation does not receive enough attention, resulting in a poor word alignment. on the contrary, for dhicm (figure 1b and figure 2b ), we observe a large variance in the distribution of the attention paid by a target token to the source tokens. thus, more appropriate source tokens receive higher attention scores (∼ 0.8) in dhicm, leading to a better word alignment, as shown in both figure 1b and figure 2b . also when 160k training sentences are used for iwslt14, although the performance of the baseline and dhicm is similar, dhicm learns better word alignments compared to the baseline (shown in figure 2 ), as dhicm helps the model to pay more attention to more relevant source tokens. moreover, dhicm allows the model to pay higher attention (∼ 0.8) to the appropriate source words compared to the baseline model where highest attention received by a source token is ∼ 0.5. this shows that for both low resource and high resource conditions, dhicm helps the model to pay higher (a) traditional transformer-base model (b) dhicm figure 2: encoder-decoder attention distribution (from model trained using 160k sentence pairs of iwslt14 training set) we also analysed the additional attention layer introduced in dhicm. we compute the attention paid by each token to each head. using the second level attention, we compute the attention paid by a particular token to all the heads and plot the attention values to create an attention distribution matrix. figure 3 shows the attention distribution for the second level attention added on top of the multi-head self attention in the last layer of the encoder. the attention distribution matrix shows the attention paid by each source token (rows) to all the 4 heads (columns). the distribution shows that each token pays different amount of attention to each head, and this justifies our hypothesis that all heads are not equally important. also, different tokens pay different amount of attention to a particular head, which also supports our hypothesis that the importance of a head is dynamic in nature, i.e., it varies as the input token changes. the attention distribution matrix also shows that the additional loss term indeed allows the model to compute different importance scores for different heads. in figure 3 , we can see that the second head gets the least attention from all the tokens. this shows that our proposed method identifies the least important heads, and thus, by incorporating dhicm, an appropriate pruning strategy can be developed to prune the least important heads. ## related work some recent work has shown that most of the heads in a multi-head attention model become redundant during test time @xcite . @xcite analyzed the heads in a multi-head attention model, based on some importance score that is calculated after the model is fully trained. in contrast, in this work, we propose to calculate the importance scores dynamically while training. a recent work @xcite proposed to apply attention on top of the output of multihead attention. however, they apply an additional attention layer only on the encoder, whereas, in our proposed method, we apply the second level attention on both encoder and decoder, that helps the generated target words to pay significant attention to appropriate source words, which in turn enhances the encoder-decoder attention distribution as shown in figure 1b . moreover, their proposed approach might learn equal attention weights for the additional attention layer, which would make all the heads equally important. in such a case, their approach would perform similar to transformer base model, even after adding more number of parameters compared to the standard transformer. to address this, we add an extra loss term in our method, to penalize for learning similar weights for the second level attention. this helps our method to compute different importance scores for different heads. furthermore, during the calculation of the final attention, they transform the output of each head using a different transformation matrix for each head, while our proposed approach dhicm uses a single transformation matrix for the outputs of all heads. thus, dhicm learns much fewer number of parameters in addition to achieving greater performance (the number of additional parameters learned in their approach is 550k, whereas dhicm learns 500k additional parameters). ## conclusion and future work in this work, we proposed an effective dynamic head importance computation mechanism (dhicm) to dynamically calculate the importance of different heads during training. our idea is to calculate the importance with an additional attention layer along with the standard multi-head attention. we also proposed a loss function to prevent our method from computing equal importance for all heads, which together with the second-level attention facilitates to dynamically identify heads that are most important to the input word. thus, the target words generated pay significantly high attention to the more appropriate/relevant source words. we also performed extensive hyperparameter tuning on a subset of hyperparameters, and observed that the optimal hyper-parameters obtained from our search yield a much higher bleu score compared to the default hyper-parameters. experiments on multiple translation tasks show that dhicm outperforms the standard transformer model by a large margin, especially in low resource settings. in the future, we will use the importance scores of the heads computed using dhicm and implement a strategy for pruning the less important heads. we would also like to explore further in the direction of reducing redundancy in multi-head attention.