Towards Understanding Neural Machine Translation with Attention Heads' Importance

Abstract

Although neural machine translation has made great progress, and the Transformer has advanced the state-of-the-art in various language pairs, the decision-making process of the attention mechanism, a crucial component of the Transformer, remains unclear. In this paper, we propose to understand the model's decisions by the attention heads' importance. We explore the knowledge acquired by the attention heads, elucidating the decision-making process through the lens of linguistic understanding. Specifically, we quantify the importance of each attention head by assessing its contribution to neural machine translation performance, employing a Masking Attention Heads approach. We evaluate the method and investigate the distribution of attention heads' importance, as well as its correlation with part-of-speech contribution. To understand the diverse decisions made by attention heads, we concentrate on analyzing multi-granularity linguistic knowledge. Our findings indicate that specialized heads play a crucial role in learning linguistics. By retaining important attention heads and removing the unimportant ones, we can optimize the attention mechanism. This optimization leads to a reduction in the number of model parameters and an increase in the model's speed. Moreover, by leveraging the connection between attention heads and multi-granular linguistic knowledge, we can enhance the model's interpretability. Consequently, our research provides valuable insights for the design of improved NMT models. Keywords: neural machine translation; interpretability; linguistics.

1. Introduction

With the advent of attentional mechanisms [1] and sequence-to-sequence learning [2], neural machine translation (NMT) has rapidly advanced, narrowing the gap between machine and human language capabilities, particularly in understanding and generating translations. The Transformer model [3] excels across a spectrum of machine translation tasks, with its attention mechanism being a pivotal component that offers exceptional parallelism, facilitating efficient large-scale data processing. Despite these advancements, the intricacy of the model's structure and the opaque nature of decision making in NMT models render them akin to black boxes. This obscurity hinders a clear understanding of the specific decisions made and the precise attributes of attention learned. Moreover, the inherent complexity of machine translation tasks, coupled with the sophisticated model architectures designed to tackle them, exacerbates this challenge [4].

For humans, NMT models often appear as "black boxes", while interpretability serves as a "crystal ball" that elucidates the models' internal decision-making processes. This clarity not only aids researchers in debugging the models but also provides a theoretical foundation for trustworthy AI. The interpretability of NMT models has been tackled from two primary, complementary perspectives. One approach emphasizes visualizing the decision-making processes and simulating the models' internal mechanisms to cultivate genuine trust between humans and the models [5]. The other perspective aims to thoroughly understand the decisions made by the model from a human standpoint, identifying the model's learned content through language knowledge that is accessible to people [6].

We chose the state-of-the-art NMT model, the Transformer model, as our object of study. This model is a neural network based on the self-attention mechanism and includes multiple layers of encoder and decoder. We quantify each attention head's importance through its performance contribution, and analyze the importance distribution and its correlation with multi-granularity linguistic knowledge — parts of speech, dependencies, and syntax trees — revealing how heads make decisions and how unimportant ones can be pruned for efficiency and interpretability.

2. Background

Interpretability lacks a unified definition. Some researchers perceive it as the extent to which people can understand the reasons behind decisions, while others believe that interpretability refers to an individual's ability to consistently and accurately predict the outcomes of model predictions [10]. Due to the complexity of widely used neural network models, their internal logic and working principles are difficult for users to understand. Users are unclear about what comprehensible linguistic knowledge the model has acquired through training and how the model processes language information. Consequently, we propose that interpretability involves bridging the gap between models that are unintelligible to humans and the integration of human-comprehensible linguistic knowledge to elucidate model behavior.

In the realm of NMT interpretability, numerous studies have concentrated on explaining the importance of various model components. These studies enable an intuitive observation of the elements crucial to model predictions, thereby enhancing the interpretability of models. Recent inquiries into attention interpretability have honed in on the functionalities of attentional mechanisms and the implications of attention weights, with research on attention heads gaining increasing traction [11]. There has been a significant emphasis on dissecting the influence of diverse components within models for enhancing interpretability [12]. Attention as the driving force behind the Transformer, the state-of-the-art NMT model, is garnering escalating scholarly attention. Previous investigations underscored the intrinsic value of attention. For example, Serrano [13] employed intermediate representation erasure to ascertain the impact of attentional mechanisms, while Li [14] explored the alignment effect of attention weights. However, as recent studies indicate, relying solely on attention weights to interpret attentional processes is inadequate.

Linguistics delves into the structure, function, and essence of human language, encompassing areas such as language structure, phonology, word formation, syntax, and semantics. In the realm of interpretability, linguistic knowledge is particularly valued for its ability to render model decisions understandable to humans. Many researchers advocate for using linguistic knowledge as a basis for elucidating models [23]. Recent studies have increasingly interpreted NMT models using linguistic insights, including aspects like morphology [6] and POS [24]. However, these studies primarily focus on the impact of individual word information in NMT. Consequently, we have explored the utilization of inter-word interactions and multi-word information to shed light on model decisions [25]. Understanding the interplay between words and leveraging multiple-word information can make translations smoother, achieving the "faithfulness, expressiveness, and elegance" standard of translation. Additionally, this approach can improve the accuracy and fluency of translations involving complex sentence structures, such as idioms and colloquialisms.

This paper investigates the decisions made by the model's important heads at three different granularities of linguistic knowledge: POS, dependency relations, and syntactic trees. POS reflects the role of individual word information in a sentence in NMT and helps us understand how the model handles individual words. Dependencies reflect the dominant relationship between two words in a sentence and help us understand how the model handles inter-word relationships. The syntax tree reflects the relationships between words in a sentence and helps us understand how the model handles multiple words as well as syntactic structure. Finally, in Table 1, we compared two categories of interpretability: visualization and linguistic interpretation. We explained their strengths and limitations and provided explanations for some references and their main contributions in order to better understand the main trends.

3. Attention Heads' Importance

We quantify each attention head's importance by assessing its contribution to translation performance, employing a Masking Attention Heads approach: we mask important, unimportant, and random heads per the computed ranking, and infer ranking validity from translation-performance degradation. Masking shows that removing important heads degrades performance sharply while masking unimportant heads barely matters — validating the importance ranking. We then analyze the importance distribution across layers and the three mechanisms (encoder self-attention, decoder self-attention, encoder–decoder attention), and its correlation with heads' part-of-speech contributions, revealing how heads learn multi-granularity linguistic knowledge.

3.1. Experimental Setup

To perform a robust and reproducible evaluation of the disclosed approach, we chose a large-scale, publicly available dataset for our study. Specifically, for the Chinese–English translation task, we use the CWMT21 Chinese–English dataset, which contains 9.0M sentence pairs [27]. As detailed in Table 2, this dataset is segmented into training, testing, and validation sets in the ratio of 0.9:0.05:0.05. Following the standard NMT procedure, we adopt the standard byte pair encoding (BPE) [28] for all language pairs. We use the standard SacreBLEU score as the evaluation metric for translation performance. The BLEU score is a widely used metric for evaluating the quality of machine translation. It measures the similarity between the machine-generated translation and one or more reference translations based on n-gram precision and a brevity penalty [29].

We employ the state-of-the-art Transformer, which is composed of two parts: an encoder and a decoder. Each part is divided into six layers, each layer has eight attention heads. A particularity of this model is that it features three distinct attention mechanisms: encoder self-attention (Enc-Enc), decoder self-attention (Dec-Dec), and encoder–decoder attention (Enc-Dec), all of which use multi-headed attention (MHA). Our study analyzed the importance of the attention head to all of these three attention mechanisms. We analyzed the learning of attention heads in the Transformer model, which includes information on individual words (part of speech), inter-word relationship information (dependency relations), and overall sentence structure information (syntax tree). Based on the importance ranking of the attention heads, we established a relationship between these three types of human-understandable linguistic knowledge and head importance.

3.2. Multi-Headed Attention

NMT is a technique for translating one natural language x = {x1, . . . , xM} into another natural language y = {y1, . . . , yN}. In autoregressive NMT, each word of the target sentence is generated step by step. Therefore, the generation of the nth target word yn is influenced by the source sentence x and the part of the target sentence that has already been generated, denoted as y<n: P(y|x) = Π P(yn|y<n, x) (1), where x is the source sentence, y the target sentence, yn the nth target word generation, and y<n the previously generated words. The model generates an output word based on the source sentence x and the partial translation y<n, and the output values of numerous attention heads are combined to produce MHA per the attention formula: Attention(Q, K, V) = softmax(QK^T/√dk)V (2), where Q, K, V are three linearly transformed matrices and dk the vector dimension of Q, K: first compute query–key similarities, second derive weights with softmax, third average weights with corresponding values, and finally concatenate each layer's attention-head outputs: MultiHead(Q, K, V) = Concat(h1, . . . , hn)WO (3), where hi is the ith attention head output and WO a linear transformation of the concatenated head outputs.

In machine translation, masking is a crucial technique that helps the model learn and predict data more effectively. The application of this technique markedly boosts the model's efficiency and accuracy, rendering it an essential component of contemporary NMT systems. Masking enables the model to better understand and process sentences of varying lengths and to utilize contextual information more effectively, thereby enhancing the quality of translation. Notably, decoder masking prevents attending to future words during training, preserving the autoregressive property of generation.

3.3. Attention Head Analysis

We mask several important attention heads, unimportant attention heads, and random attention heads based on the attention heads' importance we acquired to show that it is accurate. By looking at the degradation in translation performance, we evaluate the validity of attention heads' importance. Figure 2 illustrates the effect of masking multiple attention heads on translation performance. There is a significant degradation in translation performance when masking several important attention heads, while masking multiple unimportant attention heads results in minimal impact. This outcome validates the accuracy of our attention head importance ranking. To ensure the reliability of our experimental results, we assessed the distribution of attention head importance and its correlation with the contribution of attention heads to POS.

Wang [12] posits that certain components of the Transformer model are more important than others, particularly within the attention layer. They suggest that the decoder self-attention layers are the least critical, whereas the upper encoder-attention layers in the decoder are more significant than the lower ones. We conducted two experiments to investigate the distribution of attention head importance and verify if it aligns with these suggested properties. First, we want to know which layers in the model are more important and which are less important. This can help us understand the importance of the attention layers. Therefore, the initial experiment aims to assess the importance of the attention heads in each layer. Specifically, the translation performance is evaluated as a function that aggregates the contributions of all attention heads within a layer. If this function demonstrates a significant influence, the layer is deemed crucial. Otherwise, it is considered non-essential. Then, we want to know whether the important attention heads are concentrated in a particular layer and whether retaining only a few attention heads in a certain layer can preserve performance — guiding model pruning and acceleration.

3.4. Correlation Analysis

Previous research has identified that important attention heads frequently possess specific and interpretable functions [18]. To verify whether the attention heads' importance is consistent with the interpretable function, we analyzed the correlation between attention head importance and attention head POS contribution. The attention head contributes to the accurate translation of each POS word, which is the attention head POS contribution. Specifically, we evaluate the impact on the translation accuracy of each POS word when masking the attention head. For more details see Section 4.1. We employed various statistical methods to evaluate these correlations, including Pearson correlation and KL divergence. The Pearson correlation measures the linear correlation between two variables and the KL divergence is an asymmetric measure of the difference between two probability distributions. The value of the Pearson correlation is easy to understand and directly reflects the strength of the linear relationship between variables. The KL divergence can sensitively capture subtle differences between two probability distributions and is widely used in machine learning.

Since the data contain negative values, we must normalize it to [0, 1] before computing the KL divergence. We use two normalization methods. The first is KLNorm = (xi − xmin)/(xmax − xmin). The second method takes the logarithm of the data in the interval (−∞, +∞) to obtain data in (0, +∞), then normalizes as KLSoftmax = xi/sum(x). The correlation between attention head importance and attention head POS contribution among three distinct relevance evaluation criteria is shown in Table 3. We have established a connection between human-understandable POS information and the importance of attention heads, significantly enhancing the model's interpretability and aiding in our understanding of why some attention heads are important.

4. Linguistic Knowledge and Part-of-Speech

In this section, we analyze how the attention head learns linguistic information at various granularities, which better reveals how NMT makes decisions. Through various linguistic knowledge granularities — POS, dependency relations, and syntax trees — we explore how NMT models process word information, inter-word relations, multi-word information, and syntactic structure information. We strive to elucidate model decisions that are beyond human comprehension by leveraging linguistic knowledge that is understandable to humans [32]. We masked several important and unimportant attention heads according to the attention heads' importance to identify the various linguistic knowledge that these attention heads had acquired. Three steps comprised our investigations. In the first step, the reference translation was annotated by StanfordCoreNLP 4.5.4. StanfordCoreNLP is an open-source natural language processing toolkit that integrates various NLP tools, enabling comprehensive linguistic analysis. Therefore, we chose it for our experiments. The accuracy of the reference translation and the translation generated from the complete model are compared in the second step. In the third step, we compare the accuracy of the translation generated by the pruning model and the reference translation.

We believe that when the important attention heads are masked, the performance of the model is not as good as the complete model, leading to a decrease in the accurate translation rates of various types of linguistic information and the occurrence of negative accuracy. This experiment helps us explore which key linguistic information the model has learned and which pieces of linguistic information are more important. On the other hand, when the less important attention heads are masked, it may have a positive effect on some linguistic information, providing new insights for model pruning. The following is a detailed explanation of specific linguistic knowledge: in the POS experiment, we masked several important and unimportant attention heads per head importance to identify different POS knowledge learned by these heads, evaluating the learned POS information via per-POS translation accuracy in masking-model translations. First, per the POS types in Penn Treebank [33] — a large-scale corpus with precise POS and syntactic annotations, extremely useful for language-model training and evaluation and widely referenced as an NLP benchmark — we used StanfordCoreNLP to POS-annotate the reference translation. POS information represented the function of individual word information in a sentence for NMT. For this study, we focused on the six most frequent POS types: noun (N), verb (V), adjective (ADJ), adverb (ADV), preposition or subordinating conjunction (IN), and determiner (DT). Since POS types split into subclasses (e.g. general, comparative, and superlative adjectives), we cluster nouns, verbs, adjectives, and adverbs individually [34]. We counted per-POS-type word counts in the reference translation as in Table 4, denoting total POS-type-p words in the reference as total(p,ref).

4.2. Dependency

In this section, we explored the knowledge of various dependencies learned by the attention heads by masking several important attention heads and unimportant attention heads to better understand the information about inter-word relationships learned by each attention head in the model, as in the experiments on POS. In the first step, we used StanfordCoreNLP to extract dependencies from reference translations, and this tool can extract all dependencies of a sentence. Dependency is linguistic knowledge that reveals the dominant relationship between two words. A sentence has a root that does not depend on other words, all other words depend on one word of the sentence, and no one word can depend on two or more words. For this study, we focused on the eight dependency types that are most frequently used: root, nsubj, obj, advmod, det, amod, nmod, and compound. We statistically counted the number of word combinations of each dependency type in the reference translation, as shown in Table 4, and the total number of word combinations of dependency type dep in the reference translation was indicated as total(dep,ref).

In the second step, we extracted the dependencies from the translations generated by the complete model. The number of correct translations for the word combination of each dependency in the translation is counted as N(dep,model), with the formula: N(dep,model) = Σ Σ φ(DEPref, DEPmodel) (9), where s denotes the number of sentences in the dataset and m denotes a word combination of dependency dep in a sentence. φ() is an indicator function, DEPref,j denotes all word combinations where dependency is dep in the jth sentence of the reference translation, and DEPmodel,j,i denotes the ith word combination where dependency is dep in the jth sentence of the translation generated by the complete model. If the word combination of DEPmodel,j,i exists in DEPref,j, the word combination is considered correctly translated and the indicator output is 1, otherwise 0. Third, the procedure was repeated on pruning-model translations, and the accuracy drop under important-head masking revealed which inter-word relations the crucial heads had learned.

4.3. Syntax Tree

In this section, we studied the change in the accuracy of syntactic tree information in the translation while masking several important attention heads and unimportant attention heads to evaluate the impact of attention heads on the model's processing of multiple-word information and overall syntactic structure [35]. A syntactic tree is an abstract syntactic structure that represents the structure of a sentence. In the first step, we constructed the syntactic analysis of the reference translation using StanfordCoreNLP, a tool that constructs syntactic trees of sentences. We summarize the patterns of the bottom three nodes in the syntactic tree paths of each word to better examine the syntactic structure knowledge acquired by the attention head because there are plenty of ways to combine syntactic tree paths. For this study, we focused on the five syntactic tree path patterns that are most frequently used: NP-NP-NN, NP-PP-IN, PP-NP-NN, NP-NP-DT, and VP-PP-IN. We statistically counted the number of words of each syntactic tree path pattern in the reference translation, as shown in Table 4, and the total number of words of each syntactic tree path pattern path in the reference translation was indicated as total(path,ref).

In the second step, we constructed syntactic trees for the translations generated by the complete model. We counted the total number of words in the translation generated by the complete model and the reference translation whose word paths have the same syntactic tree path pattern, denoted as N(path,model), with the formula: N(path,model) = Σ Σ φ(PATHref, PATHmodel) (13), where PATHref,j denotes the word where the syntactic tree path pattern is path in the jth sentence of the reference translation, and PATHmodel,j,i denotes the ith word where the syntactic tree path pattern is path in the jth sentence of the translation generated by the complete model. The other parameters were consistent with the experimental settings of dependency. Third, the procedure was applied to the pruning model, and results showed important heads learn multi-word syntactic structures to varying degrees — evidence that complex syntactic decisions concentrate in specialized heads worth keeping under pruning.

5. Conclusions

In this paper, we explore the importance of attention heads in understanding NMT. We confirmed the correctness of the importance of the attention head and investigated its distribution. The correlation between attention head importance and their contribution to POS was assessed using three evaluative criteria. Our investigation covered how attention heads acquire language knowledge at three different granularities: POS, dependency relationships, and syntactic trees, enhancing our understanding of their learning processes. We find that not all attention heads are important, and the important attention heads are more distributed in the encoder self-attention layer and the encoder–decoder attention layer. These findings have propelled us towards a deeper understanding of the internal structure of NMT models and have opened up new research directions for subsequent related studies.

Specifically, these discoveries can guide us in conducting in-depth analyses of critical components within the model to enhance its performance. Additionally, they also support pruning non-critical parts of the model to reduce parameters and improve operational efficiency. In addition, the important attention heads in Chinese–English machine translation contribute more to the generation of nouns and verbs in their translations. Finally, we discovered that the attention heads had varying degrees of POS, dependence, and syntax tree learning. These findings offer valuable guidance for future model design. For instance, when designing Chinese–English machine translation models, we should pay closer attention to the processing of linguistic elements such as nouns and adjectives. This ensures that the attention mechanism functions more effectively, thereby enhancing translation quality. Overall, we investigated the multigranularity of linguistic knowledge within the context of attention heads making decisions. For future work, we plan to explore the role of attention heads in learning other linguistic knowledge.