Alleviating the Inequality of Attention Heads for Neural Machine Translation

Proceedings of the 29th International Conference on Computational Linguistics, pages 5246–5250 October 12–17, 2022. 5246 Alleviating the Inequality of Attention Heads for Neural Machine Translation Zewei Sun1,∗, Shujian Huang2,3, Xin-Yu Dai2, Jiajun Chen2 1 ByteDance AI Lab 2 State Key Laboratory for Novel Software Technology, Nanjing University 3 Peng Cheng Laboratory, China sunzewei.v@bytedance.com, {huangsj,daixinyu,chenjj}@nju.edu.cn Abstract Recent studies show that the attention heads in Transformer are not equal (Voita et al., 2019; Michel et al., 2019). We relate this phe- nomenon to the imbalance training of multi- head attention and the model dependence on specific heads. To tackle this problem, we pro- pose a simple masking method: HeadMask, in two specific ways. Experiments show that translation improvements are achieved on mul- tiple language pairs. Subsequent empirical analyses also support our assumption and con- firm the effectiveness of the method. 1 Introduction Recently, more and more novel network structures of neural machine translation(NMT) have been proposed (Bahdanau et al., 2015; Barone et al., 2017; Gehring et al., 2017; Vaswani et al., 2017), among which Transformer (Vaswani et al., 2017) 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. Voita et al. (2019) find that only a small subset of heads appear to be important for the translation task and vast ma- jority of heads can be removed without seriously affecting performance. Michel et al. (2019) 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 im- balanced training of multi-head attention, as some heads are not trained adequately and contribute lit- tle 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. 2 Head Inequality Following Michel et al. (2019), we define the im- portance of an attention head h as Ih = Ex∼X ���� ∂L(x) ∂ξh ���� (1) where L(x) is the loss on sample x and ξ is the head mask variable with values in {0, 1}. Intu- itively, if headh is important, switching ξh will have a significant effect on the loss. Applying the chain rule yields the final expression for Ih: Ih = Ex∼X ����Atth(x)T ∂L(x) ∂Atth(x) ���� (2) This is equivalent to the Taylor expansion method from Molchanov …

Content truncated to 500 words