A Dynamic Head Importance Computation Mechanism for Neural Machine Translation

Akshay Goindani,Manish Shrivastava

DOI: https://doi.org/10.48550/arXiv.2108.01377

2021-08-03

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.

Computation and Language

What problem does this paper attempt to address?

The problem that this paper attempts to solve is the redundancy problem existing in the Multi - Head Attention in Neural Machine Translation (NMT). Specifically, the Multi - Head Attention improves the performance of the Transformer model in various applications such as neural machine translation and text classification by using multiple attention heads. However, one limitation of this mechanism is that multiple heads may focus on the same part of the input, causing these heads to become redundant and thus the model resources are not fully utilized.

To solve this problem, the author proposes a Dynamic Head Importance Computation Mechanism (DHICM), which aims to dynamically calculate the importance of each head during the training process. By introducing an additional attention layer, using the output and input of the Multi - Head Attention to calculate the importance of each head, and adding an additional loss function to prevent the model from assigning the same score to all heads, it is able to identify more important heads and thus improve the model performance. Experimental results show that DHICM performs better than the traditional Transformer model in translation tasks of different language pairs, especially in the case of less training data.