Alleviating the Inequality of Attention Heads for Neural Machine Translation (2022)

Abstract: Recent studies show that the attention heads in Transformer are not equal. We relate this phenomenon to the imbalance training of multi-head attention and the model dependence on specific heads. To tackle this problem, we propose a simple masking method: HeadMask, in two specific ways. Experiments show that translation improvements are achieved on multiple language pairs. Subsequent empirical analyses also support our assumption and confirm the effectiveness of the method.

## introduction recently, more and more novel network structures of neural machine translation(nmt) have been proposed @xcite @xcite , among which transformer @xcite achieves the best results. one important difference between transformer and other translation models is its multi-head attention mechanism. some interesting phenomena of the attention heads are discovered recently. @xcite find that only a small subset of heads appear to be important for the translation task and vast majority of heads can be removed without seriously affecting performance. @xcite also find that several heads can be removed from trained transformer models without statistically significant degradation in test performance. it turns out that not all heads are equally important. we speculate that this can be attributed to the imbalanced training of multi-head attention, as some heads are not trained adequately and contribute little to the model. however, this can be turned into the bottleneck for the whole model. for an analogy, if a soccer player gets used to using the right foot and spares more training opportunities for it, it will * * work was done while at nju be stronger and stronger. as a result, the right foot is further relied on, while the left foot receives less training and gradually turns into the limitation. in this paper, we firstly empirically confirm the inequality in multi-head attention. then a new training method with two variants is proposed to avoid the bottleneck and improve the translation performance. further analyses are also made to verify the assumption. ## head inequality this is equivalent to the taylor expansion method from @xcite . in transformer base @xcite , there are 3 types of attention (encoder self attention, decoder self attention, encoder-decoder attention) with 6 layers per type and 8 heads per layer. therefore, it amounts to 144 heads. we divide them into 8 groups with 18 heads (12.5%) each group according to their importance i h , among which, 1-18 are the most important and so on. we then mask different groups of the heads. as is shown in figure 1 , masking a group of unimportant heads has little effect on the translation quality while masking important heads leads to a significant drop of performance. surprisingly, almost half of the heads are not important, as it makes almost no difference whether they are masked or not. mask different groups of heads (sorted by importance) we also gradually masking more heads group by group in the ascending order and descending order, respectively. as is shown in figure 2 , the line starting with unimportant heads drops much slower than the one starting with important ones. it fully illustrates the inequality of different heads. figure 1 and figure 2 further demonstrates the inequality of the importance of attention heads. a simple assumption for explanation is that some heads coincidentally get more updating opportunities in the early stage, which makes the model learning to depend on them gradually. as a result, the model increasingly draws a strong connection with these specific heads while this local dependence prevents the rest attention heads from adequate training and restricts the overall capacity. ## headmask since the problem refers to the unfair training of attention heads, it is natural for us to explicitly balance the training chances. we propose a simple method: headmask, which masks certain heads during training in two specific ways. ## related works recently, many analytical works about multi-head attention come out @xcite @xcite @xcite . and for the inequality of the networks, some studies focus on the model level @xcite , layer level @xcite , and neuron level @xcite . for the mask algorithm, there are also works on the layer level @xcite , word level @xcite , and neuron level @xcite . different from them, we mainly study the attention level and conduct a statistical analysis. ## conclusion in this paper, we empirically validate the inequality of attention heads in transformer and come up with an assumption of imbalanced training. correspondingly, we propose a specific method in two ways to resolve the issue. experiments show the improvements on multiple language pairs. and detailed analysis shows the alleviation of the problem and the effectiveness of our techniques.