Venue: NAACL 2019

# Information Aggregation for Multi-Head Attention with

Routing-by-Agreement

Jian Li\({}^{1}\) Baosong Yang\({}^{2}\) Zi-Yi Dou\({}^{3}\) Xing Wang\({}^{4}\) Michael R. Lyu\({}^{1}\) Zhaopeng Tu\({}^{4}\)

\({}^{1}\)Department of Computer Science and Engineering, The Chinese University of Hong Kong {jianli,lyu}@cse.cuhk.edu.hk

\({}^{2}\)University of Macau

nlp2ct.baosong@gmail.com   dzdu@andrew.cmu.edu

\({}^{4}\)Tencent AI Lab

{brightxwang,zptu}@tencent.com

###### Abstract

Multi-head attention is appealing for its ability to jointly extract different types of information from multiple representation subspaces. Concerning the information aggregation, a common practice is to use a concatenation followed by a linear transformation, which may not fully exploit the expressiveness of multi-head attention. In this work, we propose to improve the information aggregation for multi-head attention with a more powerful _routing-by-agreement_ algorithm. Specifically, the routing algorithm iteratively updates the proportion of how much a part (i.e. the distinct information learned from a specific subspace) should be assigned to a whole (i.e. the final output representation), based on the agreement between parts and wholes. Experimental results on linguistic probing tasks and machine translation tasks prove the superiority of the advanced information aggregation over the standard linear transformation.

## 1 Introduction

Attention model becomes a standard component of the deep learning networks, contributing to impressive results in machine translation Bahdanau et al. (2015); Luong et al. (2015), image captioning Xu et al. (2015), speech recognition Chorowski et al. (2015), among many other applications. Its superiority lies in the ability of modeling the dependencies between representations without regard to their distance. Recently, the performance of attention is further improved by multi-head mechanism Vaswani et al. (2017), which parallelly performs attention functions on different representation subspaces of the input sequence. Consequently, different attention heads are able to capture distinct linguistic properties of the input, which are embedded in different subspaces Raganato and Tiedemann (2018). Subsequently, a linear transformation is generally employed to aggregate the partial representations extracted by different attention heads Vaswani et al. (2017); Ahmed et al. (2018).

Most existing work focus on extracting informative or distinct partial-representations from different subspaces (e.g. Lin et al., 2017; Li et al., 2018), while few studies have paid attention to the aggregation of the extracted partial-representations. Arguably, information extraction and aggregation are both crucial for multi-head attention to generate an informative representation. Recent studies in multimodal learning show that a straightforward linear transformation for fusing features in different sets of representations usually limits the extent of abstraction Fukui et al. (2016); Ben-Younes et al. (2017). A natural question arises: _whether the straightforward linear transformation is expressive enough to fully capture the rich information distributed in the extracted partial-representations?_

In this work, we provide the first answer to this question. We propose to empirically validate the importance of information aggregation in multi-head attention, by comparing the performance of the standard linear function and advanced aggregation functions on various tasks. Specifically, we cast information aggregation as the _assigning-parts-to-wholes_ problem Hinton et al. (2011), and investigate the effectiveness of the routing-by-agreement algorithm - an appealing alternative to solving this problem Sabour et al. (2017); Hinton et al. (2018). The routing algorithm iteratively updates the proportion of how much a part should be assigned to a whole, based on the agreement between parts and wholes. We leverage the routing algorithm to aggregate the information distributed in the extracted partial-representations.



We evaluate the performance of the aggregated representations on both linguistic probing tasks as well as machine translation tasks. The probing tasks Conneau et al. (2018) consists of 10 classification problems to study what linguistic properties are captured by input representations. Probing analysis show that our approach indeed produces more informative representation, which embeds more syntactic and semantic information. For translation tasks, we validate our approach on top of the advanced Transformer model Vaswani et al. (2017) on both WMT14 English\(\Rightarrow\)German and WMT17 Chinese\(\Rightarrow\)English data. Experimental results show that our approach consistently improves translation performance across languages while keeps the computational efficiency.

The main contributions of this paper can be summarized as follows:

* To our best knowledge, this is the first work to demonstrate the necessity and effectiveness of advanced information aggregation for multi-head attention.
* Our work is among the few studies (_cf._Gong et al. (2018); Zhao et al. (2018); Dou et al. (2019)) which prove that the idea of capsule networks can have promising applications on natural language processing tasks.

## 2 Background

Attention mechanism aims at modeling the relevance between representation pairs, thus a representation is allowed to build a direct relation with another representation. Instead of performing a single attention function, Vaswani et al. (2017) found it is beneficial to capture different context features with multiple individual attention functions, namely multi-head attention. Formally, attention function maps a sequence of query \({\bf Q}=\{{\bf q}_{1},\ldots,{\bf q}_{J}\}\) and a set of key-value pairs \(\{{\bf K},{\bf V}\}=\{({\bf k}_{1},{\bf v}_{1}),\ldots,({\bf k}_{M},{\bf v}_{ M})\}\) to outputs, where \({\bf Q}\in\mathbb{R}^{J\times d}\), \(\{{\bf K},{\bf V}\}\in\mathbb{R}^{M\times d}\). More specifically, multi-head attention model first transforms **Q**, **K**, and **V** into \(H\) subspaces with different, learnable linear projections:

\[{\bf Q}_{h},{\bf K}_{h},{\bf V}_{h}={\bf QW}_{h}^{Q},{\bf K}{\bf W}_{h}^{K},{ \bf V}{\bf W}_{h}^{V}, \tag{1}\]

where \(\{{\bf Q}_{h},{\bf K}_{h},{\bf V}_{h}\}\) are respectively the query, key, and value representations of the \(h\)-th head. \(\{{\bf W}_{h}^{Q},{\bf W}_{h}^{K},{\bf W}_{h}^{V}\}\in\mathbb{R}^{d\times\frac{ d}{H}}\) denote parameter matrices associated with the \(h\)-th head, where \(d\) represents the dimensionality of the model hidden states. Furthermore, \(H\) attention functions are applied in parallel to produce the output states \(\{{\bf O}_{1},\ldots,{\bf O}_{H}\}\), among them:

\[{\bf O}_{h}=\textsc{Att}({\bf Q}_{h},{\bf K}_{h}){\bf V}_{h}, \tag{2}\]

where \({\bf O}_{h}\in\mathbb{R}^{J\times\frac{d}{H}}\), Att\((\cdot)\) is an attention model. In this work, we use scaled dot-product attention Luong et al. (2015), which achieves similar performance with its additive counterpart Bahdanau et al. (2015) while is much faster and more space-efficient in practice Vaswani et al. (2017).

Finally, the \(H\) output states are concatenated and linearly transformed to produce the final state:

Concat: \[\widehat{\mathbf{O}}= [{\bf O}_{1},\ldots,{\bf O}_{H}],\] (3) Linear: \[{\bf O}= \widehat{\mathbf{O}}{\bf W}^{O},\] (4)

where \({\bf O}\in\mathbb{R}^{J\times d}\) denotes the final output states, \({\bf W}^{O}\in\mathbb{R}^{d\times d}\) is a trainable matrix.

As shown in Equations 3 and 4, the conventional multi-head attention uses a straightforward concatenation and linear mapping to aggregate the output representations of multiple attention heads. We argue that this straightforward strategy may not fully exploit the expressiveness of multi-head attention, which can benefit from advanced information aggregation by exploiting the intrinsic relationship among the learned representations.

## 3 Related Work

Our work synthesizes two strands of research work, namely _multi-head attention_ and _information aggregation_.

### Multi-Head Attention

Multi-head attention has shown promising empirical results in many NLP tasks, such as machine translation Vaswani et al. (2017); Domhan (2018), semantic role labeling Strubell et al. (2018), and subject-verb agreement task Tang et al. (2018). The strength of multi-head attention lies in the rich expressiveness by using multiple attention functions in different representation subspaces.

Previous work show that multi-head attention can be further enhanced by encouraging individual attention heads to extract distinct information. For example, Lin et al. (2017) introduce a penalization term to reduce the redundancy of attention weights among different attention heads. Li et al. (2018) propose disagreement regularizations 

to encourage different attention heads to capture distinct features, and Yang et al. (2019) model the interactions among attention heads. Shen et al. (2018) explicitly use multiple attention heads to model different dependencies of the same word pair, and Strubell et al. (2018) employ different attention heads to capture different linguistic features. Our approach is complementary to theirs, since they focus on extracting distinct information while ours aims at effectively aggregating the extracted information. Our study shows that information aggregation is as important as information extraction for multi-head attention.

### Information Aggregation

Information aggregation in multi-head attention (e.g. Equations 3 and 4) aims at composing the partial representations of the input captured by different attention heads to a final representation. Recent work shows that representation composition benefits greatly from advanced functions beyond simple concatenation or mean/max pooling. For example, Fukui et al. (2016) and Ben-Younes et al. (2017) succeed on fusing multi-modal features (e.g., visual features and textual features) more effectively via employing the higher-order bilinear pooling instead of vector concatenation or element-wise operations. In NLP tasks, Peters et al. (2018) aggregate layer representations with linear combination, and Dou et al. (2018) compose deep representations with layer aggregation and multi-layer attention mechanisms.

Recently, the routing-by-agreement algorithm, which origins from the capsule networks Hinton et al. (2011), becomes an appealing alternative to representation composition. The majority of existing work on capsule networks has focused on computer vision tasks, such as MNIST tasks Sabour et al. (2017); Hinton et al. (2018), CIFAR tasks Xi et al. (2017), and object segmentation task LaLonde and Bagci (2018). The applications of capsule networks in NLP tasks, however, have not been widely investigated to date. Zhao et al. (2018) testify capsule networks on text classification tasks and Gong et al. (2018) propose to aggregate a sequence of vectors via dynamic routing for sequence encoding. Dou et al. (2019) use routing-by-agreement strategies to aggregate layer representations dynamically. Inspired by these successes, we apply the routing algorithms to multi-head attention on both linguistic probing and machine translation tasks, which demonstrates the necessity and effectiveness of advanced information aggregation for multi-head attention.

## 4 Approach

In this work, we cast information aggregation in multi-head attention as the problem of _assigning-parts-to-wholes_. Specifically, each attention head extracts different linguistic properties of the same input Raganato and Tiedemann (2018), and the goal of information aggregation is to compose the partial representations extracted by different heads to a whole representation. An appealing solution to this problem is the _routing-by-agreement_ algorithm, as shown in Figure 1.

The routing algorithm consists of two layers: _input capsules_ and _output capsules_. The input capsules are constructed from the transformation of the partial representations extracted by different attention heads. For each output capsule, each input capsule proposes a distinct "voting vector", which represents the proportion of how much the information is transformed from this input capsule (i.e parts) to the corresponding output capsule (i.e. wholes). The proportion is iteratively updated based on the agreement between the voting vectors and the output capsule. Finally, all output capsules are concatenated to form the final representation.

### Routing-by-Agreement

Mathematically, the input capsules \({\mathbf{\Omega}^{in}=\{{\mathbf{\Omega}}_{1}^{in},\dots,{\mathbf{\Omega}}_{H}^{in}\}\) with \({\mathbf{\Omega}^{in}\in\mathbb{R}^{n\times d}}\) are constructed from the outputs of multi-head attention:

\[{\mathbf{\Omega}}_{h}^{in}=f_{h}(\widehat{\mathbf{O}}), \tag{5}\]

where \(f_{h}(\cdot)\) is a distinct non-linear transformation function associated with the input capsule

Figure 1: Illustration of routing-by-agreement.

 

\({\mathbf{\Omega}}_{h}^{in}\). Given \(N\) output capsules, each input capsule \({\mathbf{\Omega}}_{h}^{in}\) propose \(N\) "vote vectors" \({\bf V}_{h\to*}=\{{\bf V}_{h\to 1},\ldots,{\bf V}_{h\to N}\}\), which is calculated by

\[{\bf V}_{h\to n}={\mathbf{\Omega}}_{h}^{in}{\bf W}_{h\to n}, \tag{6}\]

Each output capsule \({\mathbf{\Omega}}_{n}^{out}\) is calculated as the normalization of its total input, which is a weighted sum over all "vote vectors" \({\bf V}_{*\to n}\):

\[{\mathbf{\Omega}}_{n}^{out}=\frac{\sum_{h=1}^{H}C_{h\to n}{\bf V}_{ h\to n}}{\sum_{h=1}^{H}C_{h\to n}}, \tag{7}\]

The weight \(C_{h\to n}\) with \(\sum_{n}C_{h\to n}=1\) measures the agreement between vote vector \({\bf V}_{h\to n}\) and output capsule \({\mathbf{\Omega}}_{n}^{out}\), which is determined by the iterative routing as described in the next section. Note that \(\sum_{h=1}^{H}C_{h\to n}\) is not necessarily equal to \(1\). After the routing process, following Gong et al. (2018), we concatenate the \(N\) output capsules to form the final representation: \({\bf O}=[{\mathbf{\Omega}}_{1}^{out},\ldots,{\mathbf{\Omega}}_{N}^{out}]\). To make the dimensionality of the final output be consistent with that of hidden layer (i.e. \(d\)), we set the dimensionality of each output capsule be \(\frac{d}{N}\).

### Routing Mechanisms

In this work, we explore two representative routing mechanisms, namely _simple routing_Sabour et al. (2017) and _EM routing_Hinton et al. (2018), which differ at how the agreement weights \(C_{h\to n}\) are calculated.

#### 4.2.1 Simple Routing

Algorithm 1 lists a straightforward implementation of routing mechanism. \(B_{h\to n}\) measures the degree that the input capsule \({\mathbf{\Omega}}_{h}^{in}\) should be coupled to the output capsule \({\mathbf{\Omega}}_{n}^{out}\), which is initialized as all \(0\) (Line 2). The agreement weights \(C_{h\to n}\) are then iteratively refined by measuring the agreement between vote vector \({\bf V}_{h\to n}\) and output capsule \({\mathbf{\Omega}}_{n}^{out}\) (Lines 4-6), which is implemented as a simple scalar product \({\mathbf{\Omega}}_{n}^{out}\cdot{\bf V}_{h\to n}\) (Line 6).

To represent the probability that the output capsule \({\mathbf{\Omega}}_{n}^{out}\) is activated, Sabour et al. (2017) use a non-linear "squashing" function:

\[{\mathbf{\Omega}}_{n}^{out} = \frac{||{\mathbf{\Omega}}_{n}^{out}||^{2}}{1+||{\mathbf{\Omega}}_{n}^{ out}||^{2}}\frac{{\mathbf{\Omega}}_{n}^{out}}{||{\mathbf{\Omega}}_{n}^{out}||}, \tag{8}\]

```
1:procedureEM Routing(V, \(T\)):
2:\(\forall{\bf V}_{h\rightarrow*}\): \(C_{l\to n}=1/N\)
3:for\(T\) iterations do
4:\(\forall{\bf\Omega}_{n}^{out}\): M-step(V, \(C\))\(\triangleright\) hold \(C\) constant, adjust (\({\mathbf{\mu}}_{n},{\mathbf{\sigma}}_{n},A_{n}\))
5:\(\forall{\bf V}_{h\rightarrow*}\): E-step(V, \({\mathbf{\mu}},{\mathbf{\sigma}},A\)) \(\triangleright\) hold (\({\mathbf{\mu}},{\mathbf{\sigma}},A\)) constant, adjust \(C_{h\rightarrow*}\)
6:\(\forall{\bf\Omega}_{n}^{out}\): \({\mathbf{\Omega}}_{n}^{out}=A_{n}*{\mathbf{\mu}}_{n}\)return\({\mathbf{\Omega}}\)
```

**Algorithm 2** Iterative EM Routing.

The EM algorithm alternates between an E-step and an M-step. The E-step determines, for 

each datapoint (i.e. input capsule), the probability of agreement (i.e. \(C\)) between it and each of the Gaussians (i.e. output capsules). The M-step holds the agreement weights constant, and for each Gaussian (i.e. output capsule) consists of finding the mean of these weighted datapoints (i.e. input capsules) and the variance about that mean.

_M-Step_ for each Gaussian (i.e. \({\bf\Omega}^{out}_{n}\)) consists of finding the mean \({\mathbf{\mu}}_{n}\) of the votes from input capsules and the variance \({\mathbf{\sigma}}_{n}\) about that mean:

\[{\mathbf{\mu}}_{n} = \frac{\sum_{h=1}^{H}C_{h\to n}{\mathbf{V}}_{h\to n}}{\sum_{ h=1}^{H}C_{h\to n}}, \tag{10}\] \[({\mathbf{\sigma}}_{n})^{2} = \frac{\sum_{h=1}^{H}C_{h\to n}({\mathbf{V}}_{h\to n}-{\mathbf{\mu}}_{n})^{2} }{\sum_{h=1}^{H}C_{h\to n}}. \tag{11}\]

The incremental cost of using an active capsule \({\bf\Omega}^{out}_{n}\) is

\[\chi_{n}=\sum_{i}\big{(}\log({\mathbf{\sigma}}_{n}^{i})+\frac{1+\log(2\pi)}{2} \big{)}\sum_{h=1}^{H}C_{h\to n},\]

where \({\mathbf{\sigma}}_{n}^{i}\) denotes the \(i\)-th dimension of the variance vector \({\mathbf{\sigma}}_{n}\). The activation probability of capsule \({\bf\Omega}^{out}_{n}\) is calculated by

\[A_{n}=logistic\big{(}\lambda(\beta_{A}-\beta_{\mu}\sum_{h=1}^{H}C_{h\to n}- \chi_{n})\big{)},\]

where \(\beta_{A}\) is a fixed cost for coding the mean and variance of \({\bf\Omega}^{out}_{n}\) when activating it, \(\beta_{\mu}\) is another fixed cost per input capsule when not activating it, and \(\lambda\) is an inverse temperature parameter set with a fixed schedule. We refer the readers to Hinton et al. (2018) for more details.

_E-Step_ adjusts the assignment probabilities \(C_{h\to*}\) for each input \({\bf\Omega}^{in}_{h}\). First, we compute the negative log probability density of the vote \({\mathbf{V}}_{h\to n}\) from \({\bf\Omega}^{in}_{h}\) under the Gaussian distribution fitted by the output capsule \({\bf\Omega}^{out}_{n}\) it gets assigned to:

\[P_{h\to n}=\sum_{i}\frac{1}{\sqrt{2\pi({\mathbf{\sigma}}_{n}^{i})^{2}}}\exp(- \frac{({\mathbf{V}}_{h\to n}^{i}-{\mathbf{\mu}}_{n}^{i})^{2}}{2({\mathbf{\sigma}}_{n}^ {i})^{2}}).\]

Again, \(i\) denotes the \(i\)-th dimension of the vectors \(\{{\mathbf{V}}_{h\to n},{\mathbf{\mu}}_{n},{\mathbf{\sigma}}_{n}\}\). Accordingly, the agreement weight is re-normalized by

\[C_{h\to n}=\frac{A_{n}P_{h\to n}}{\sum_{n^{\prime}=1}^{N}A_{n^{\prime}}P_{h \to n^{\prime}}}. \tag{12}\]

## 5 Experiments

In this section, we evaluate the performance of our proposed models on both linguistic probing tasks and machine translation tasks.

### Linguistic Probing Tasks

#### 5.1.1 Setup

TasksRecently, Conneau et al. (2018) designed 10 probing tasks to study what linguistic properties are captured by input representations. A probing task is a classification problem that focuses on simple linguistic properties of sentences. 'SeLen' is to predict the length of sentences in terms of number of words. 'WC' tests whether it is possible to recover information about the original words given its sentence embedding. 'TrDep' checks whether an encoder infers the hierarchical structure of sentences. In 'ToCo' task, sentences should be classified in terms of the sequence of top constituents immediately below the sentence node. 'Bshift' tests whether two consecutive tokens within the sentence have been inverted. 'Tense' asks for the tense of the main-clause verb. 'SubNam' focuses on the number of the subject of the main clause. 'ObjNm' tests for the number of the direct object of the main clause. In 'SOMO', some sentences are modified by replacing a random noun or verb with another noun or verb and the classifier should tell whether a sentence has been modified. 'CoIn' benchmark contains sentences made of two coordinate clauses. Half of the sentences are inverted the order of the clauses and the task is to tell whether a sentence is intact or modified. We conduct probing tasks to study whether the routing-based aggregation benefits multi-head attention to produce more informative representation.

Data and ModelsThe models on each classification task are trained and examined using the open-source dataset provided by Conneau et al. (2018), where each task is assigned 100k sentences for training and 10k sentences for validating and testing. Each of our probing model consists of 3 encoding layers followed by a MLP classifier. For each encoding layer, we employ a multi-head self-attention block and a feed-forward block as in TransformerBase, which have achieved promising results on several NLP tasks Dehghani et al. (2018); Devlin et al. (2018). The mean of the top encoding layer is served as the sentence 

representation passed to the classifier. The difference between the compared models merely lies in the aggregation mechanism of multiple attention heads: "Base" uses a standard concatenation and linear transformation, "Simple" and "EM" are assigned simple routing and EM routing algorithms, respectively. For routing algorithms, the number of output capsules and routing iterations are empirically set to 512 and 3.

#### 5.1.2 Results

Table 1 lists the classification accuracies of the three models on the 10 probing tasks. We highlight the best accuracies in bold. Several observations can be made here.

First, _routing-based models produce more informative representation._ The representation produced by encoders with routing-based aggregation outperforms that by the baseline in most tasks, proving that routing mechanisms indeed aggregate attention heads more effectively. The only exception is the sentence length classification task ('Seene'), which is consistent with the conclusion in Conneau et al. (2018): as a model captures deeper linguistic properties, it will tend to forget about this superficial feature.

Second, _EM routing outperforms simple routing by embedding more syntactic and semantic information._ As shown in the last row, EM routing for multi-head aggregation consistently achieves best performances on most syntactic and semantic tasks. Especially on task 'TrDep', 'Tense' and 'ObjNm', EM routing-based model surpasses the baseline more than 3 points, demonstrating that EM routing benefits multi-head attention to capture more syntax structure and sentence meaning. Simple routing, however, underperforms the baseline model in some cases such as 'StrDep' and 'SubNm'. We attribute the superiority of EM routing to generating more accurate agreement weights with the Gaussian estimation.

### Machine Translation Tasks

#### 5.2.1 Setup

DataWe conduct experiments on the widely-used WMT2014 English\(\Rightarrow\)German (En\(\Rightarrow\)De) and WMT2017 Chinese\(\Rightarrow\)English (Zh\(\Rightarrow\)En) machine translation tasks. For the En\(\Rightarrow\)De task, the dataset consists of 4.6M sentence pairs. We use newstest2013 as the development set and newstest2014 as the test set. For the Zh\(\Rightarrow\)En task, we use all of the available parallel data, consisting of about 20.6M sentence pairs. We use newsdev2017 as the development set and newstest2017 as the test set. We employ byte-pair encoding (BPE) Sennrich et al. (2016) with 32K merge operations for both language pairs. We use the case-sensitive 4-gram NIST BLEU score Papineni et al. (2002) as evaluation metric, and bootstrap resampling Koehn (2004) for statistical significance test.

ModelsWe implement the proposed approaches on top of the advanced Transformer model Vaswani et al. (2017). We follow Vaswani et al. (2017) to set the configurations and have reproduced their reported results on the En\(\Rightarrow\)De task. The _Base_ and _Big_ models differ at hidden size (512 vs. 1024) and number of attention heads (8 vs. 16). All the models are trained on eight NVIDIA P40 GPUs where each is allocated with a batch size of 4096 tokens.

Transformer consists of three attention components: encoder-side self-attention, decoder-side self-attention and encoder-decoder attention, all of which are implemented as multi-head attention. For the information aggregation in multi-head attention, we replace the standard linear transformation with the proposed routing mechanisms. We experimentally set the number of iterations to 3 and the number of output capsules as model hidden size, which outperform other configurations during investigation.

\begin{table}
\begin{tabular}{c|c c|c c c|c c c c} \multirow{2}{*}{**Model**} & \multicolumn{2}{c|}{**Surface**} & \multicolumn{3}{c|}{**Syntactic**} & \multicolumn{3}{c}{**Semantic**} \\ \cline{2-10}  & SeLen & WC & TrDep & ToCo & BShif & Tense & SubNm & ObjNm & SOMO & CoIn \\ \hline Base & **97.22** & 97.92 & 44.48 & 84.44 & 49.30 & 84.20 & 87.66 & 82.94 & 50.24 & 68.77 \\ \hline Simple & 97.10 & **98.85** & 43.37 & 86.15 & 49.87 & **88.22** & 87.25 & 85.07 & 48.77 & 69.12 \\ Em & 96.26 & 98.75 & **47.72** & **87.00** & **51.82** & 88.17 & **89.97** & **86.40** & **51.55** & **69.86** \\ \end{tabular}
\end{table}
Table 1: Classification accuracies on 10 probing tasks of evaluating the linguistic properties (“Surface”, “Synctetic”, and “Semantic”) learned by sentence encoder. “Base” denotes the standard linear transformation, “Simple” is the simple routing algorithm, and “Em” is the EM routing algorithm.

 



#### 5.2.2 Component Analysis

Table 2 lists the results on the En\(\Rightarrow\)De translation task with Transformer-Base. As seen, the proposed routing mechanism outperforms the standard aggregation in all cases, demonstrating the necessity of advanced aggregation functions for multi-head attention.

Routing Mechanisms(Rows 3-4) We first apply simple routing and EM routing to encoder self-attention. Both strategies perform better than the standard multi-head aggregation (Row 1), verifying the effectiveness of the non-linear aggregation mechanisms. Specifically, the two strategies require comparable parameters and computational speed, but EM routing achieves better performance on translation qualities. Considering the training speed and performance, _EM routing_ is used as the default multi-head aggregation method in subsequent experiments.

Effect on Attention Components(Rows 4-8) Concerning the individual attention components (Rows 4-6), we found that the encoder and decoder self-attention benefit more from the routing-based information aggregation than the encoder-decoder attention. This is consistent with the finding in Tang et al. (2018), which shows that self-attention is a strong semantic feature extractor. Encouragingly, applying EM routing in the encoder (Row 4) significantly improve the translation quality with almost no decrease in decoding speed, which matches the requirement of online MT systems. We find that this is due to the auto-regressive generation schema, modifications on the decoder influence the decoding speed more than the encoder.

Compared with individual attention components, applying routing to multiple components (Rows 7-8) marginally improves translation performance, at the cost of a significant decrease of the training and decoding speeds. Possible reasons include that the added complexity makes the model harder to train, and the benefits enjoyed by different attention components are overlapping to some extent. To balance translation performance and efficiency, we only apply EM routing to aggregate multi-head self-attention at the _encoder_ in subsequent experiments.

Encoder LayersAs shown in Row 4 of Table 2, applying EM routing to all encoder layers significantly decreases the training speed by 37.5%, which is not acceptable since Transformer is best known for both good performance and quick training. We expect applying to fewer layers can alleviate the training burden. Recent studies show that different layers of NMT encoder can capture different levels of syntax and semantic features Shi et al. (2016); Peters et al. (2018). There

\begin{table}
\begin{tabular}{c|c c|c|c||c|c c||c|c|c} \cline{2-9} \# & \multicolumn{3}{c||}{**Applying Aggregation to ...**} \\ \cline{2-10} \multirow{-2}{*}{1} & _Enc-Self_ & _Enc-Dec_ & _Dec-Self_ & \multirow{-2}{*}{**Routing**} & \multirow{-2}{*}{**\# Para.**} & \multicolumn{2}{c||}{**Speed**} & \multirow{-2}{*}{**BLEU**} & \multirow{-2}{*}{\(\triangle\)} \\ \hline
2 & ✗ & ✗ & ✗ & & n/a & 88.0M & 1.92 & 1.67 & 27.31 & – \\ \hline
3 & ✓ & ✗ & ✗ & Simple & +12.6M & 1.23 & 1.66 & 27.98 & +0.67 \\
4 & ✓ & ✗ & ✗ & EM & +12.6M & 1.20 & 1.65 & 28.28 & +0.97 \\
5 & ✗ & ✓ & ✗ & EM & +12.6M & 1.20 & 1.21 & 27.94 & +0.63 \\
6 & ✗ & ✗ & ✓ & EM & +12.6M & 1.21 & 1.21 & 28.15 & +0.84 \\
7 & ✓ & ✓ & ✗ & EM & +25.2M & 0.87 & 1.20 & 28.45 & +1.14 \\
8 & ✓ & ✓ & ✓ & EM & +37.8M & 0.66 & 0.89 & 28.47 & +1.16 \\ \end{tabular}
\end{table}
Table 2: Effect of information aggregation on different attention components, i.e., encoder self-attention (“_Enc-Self_”), encoder-decoder attention (“_Enc-Dec_”), and decoder self-attention (“_Dec-Self_”). “# Para.” denotes the number of parameters, and “Train” and “Decode” respectively denote the training speed (steps/second) and decoding speed (sentences/second).

\begin{table}
\begin{tabular}{r|c|c c|c} \# & **Layers** & **\# Para.** & **Train** & **BLEU** \\ \hline \hline
1 & None & 88.0M & 1.92 & 27.31 \\ \hline
2 & [1–6] & 100.6M & 1.20 & 28.28 \\
3 & [4–6] & 94.3M & 1.54 & 28.26 \\
4 & [1–3] & 94.3M & 1.54 & 28.27 \\
5 & [1,2] & 92.2M & 1.67 & 28.26 \\
6 & [6] & 90.1M & 1.88 & 27.68 \\
7 & [1] & 90.1M & 1.88 & 27.75 \\ \end{tabular}
\end{table}
Table 3: Evaluation of different layers in the encoder, which are implemented as multi-head self-attention with the EM routing based information aggregation. “1” denotes the bottom layer, and “6” the top layer.



fore, an investigation to study whether EM routing works for multi-head attention at different layers is highly desirable.

As shown in Table 3, we respectively employ EM routing for multi-head attention at the high-level three layers (Row 3) and low-level three layers (Row 4). The translation quality marginally drop while parameters are fewer and training speeds are quicker. This phenomena verifies that it is unnecessary to apply the proposed model to all layers. We further reduce the applied layers to low-level two (Row 5), the above phenomena still holds. However, a big drop on translation quality occurs when the number of layer is reduced to 1 (Rows 6-7).

Accordingly, to balance translation performance and efficiency, we only apply EM routing for multi-head aggregation at the _low-level two layers of the encoder_, which we term "Effective Aggregation" in the following sections.

#### 5.2.3 Main Results

In this section, we validate the proposed "Effective Aggregation" for multi-head attenion on both WMT17 Zh\(\Rightarrow\)En and WMT14 En\(\Rightarrow\)De translation tasks. The results are listed in Table 4. Our implementations of both Transformer-Base and Transformer-Big outperform the reported NMT systems on the same data and match the strong results of Transformer reported in previous works, which we believe make the evaluation convincing. Incorporating the effective aggregation consistently and significantly improves translation performance for both base and big Transformer models across language pairs, demonstrating the efficiency and universality of our proposed multi-head aggregation mechanism.

Moreover, it is encouraging to see that Transformer-Base with effective aggregation strategy even achieves comparable performances to that of Transformer-Big, with about two thirds fewer parameters, which further demonstrates that our performance gains are not simply brought by additional parameters.

## 6 Conclusion

In this work, we provide first empirical validation on the importance of information aggregation for multi-head attention. Instead of the conventional linear transformation, we propose to aggregate the partial representations learned by multiple attention heads via _routing-by-agreement_. The routing algorithm iteratively updates the proportion of how much a partial representation should be assigned to the final output representation, based on the agreement between parts and wholes.

Experimental results across 10 linguistic probing tasks reveal that our EM routing-based model indeed produces more informative representation, which benefits multi-head attention to capture more syntactic and semantic information. In addition, our approach on various machine translation tasks consistently and significantly outperforms the strong Transformer baseline. Extensive analysis further suggests that only applying EM routing to low-level two layers of the encoder can best balance the translation performance and computational efficiency 

Future work includes combining our information aggregation techniques together with other advanced information extraction models for multi-head attention (Li et al., 2018). We expect that the two kinds of approaches can complement each other to further improve the expressiveness of multi-head attention.

## Acknowledgments

Jian Li and Michael R. Lyu were supported by the Research Grants Council of the Hong Kong Special Administrative Region, China (No. CUHK 14210717 of the General Research Fund), and Microsoft Research Asia (2018 Microsoft Research Asia Collaborative Research Award). We thank the anonymous reviewers for their insightful comments and suggestions.

## References

* Ahmed et al. (2018) Karim Ahmed, Nitish Shirish Keskar, and Richard Socher. 2018. Weighted Transformer Network for Machine Translation. _arXiv preprint arXiv:1711.02132_.
* Bahdanau et al. (2015) Dzmitry Bahdanau, Kyunghyun Cho, and Yoshua Bengio. 2015. Neural Machine Translation by Jointly Learning to Align and Translate. In _ICLR_.
* Ben-Younes et al. (2017) Hedi Ben-Younes, Remi Cadene, Matthieu Cord, and Nicolas Thome. 2017. Mutan: Multimodal Tucker Fusion for Visual Question Answering. In _ICCV_.
* Chorowski et al. (2015) Jan K Chorowski, Dzmitry Bahdanau, Dmitriy Serdyuk, Kyunghyun Cho, and Yoshua Bengio. 2015. Attention-based Models for Speech Recognition. In _NIPS_.
* Conneau et al. (2018) Alexis Conneau, German Kruszewski, Guillaume Lample, Loic Barrault, and Marco Baroni. 2018. What You Can Cram into A Single SSc!.\#\(\ast\) Vector: Probing Sentence Embeddings for Linguistic Properties. In _ACL_.
* Dehghani et al. (2018) Mostafa Dehghani, Stephan Gouws, Oriol Vinyals, Jakob Uszkoreit, and Lukasz Kaiser. 2018. Universal Transformers. _arXiv preprint arXiv:1807.03819_.
* Devlin et al. (2018) Jacob Devlin, Ming-Wei Chang, Kenton Lee, and Kristina Toutanova. 2018. BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding. _arXiv preprint arXiv:1810.04805_.
* Domhan (2018) Tobias Domhan. 2018. How Much Attention Do You Need? A Granular Analysis of Neural Machine Translation Architectures. In _ACL_.
* Dou et al. (2018) Ziyi Dou, Zhaopeng Tu, Xing Wang, Shuming Shi, and Tong Zhang. 2018. Exploiting deep representations for neural machine translation. In _EMNLP_.
* Dou et al. (2019) Ziyi Dou, Zhaopeng Tu, Xing Wang, Longyue Wang, Shuming Shi, and Tong Zhang. 2019. Dynamic layer aggregation for neural machine translation. In _AAAI_.
* Fukui et al. (2016) Akira Fukui, Dong Huk Park, Daylen Yang, Anna Rohrbach, Trevor Darrell, and Marcus Rohrbach. 2016. Multimodal Compact Bilinear Pooling for Visual Question Answering and Visual Grounding. In _EMNLP_.
* Gehring et al. (2017) Jonas Gehring, Michael Auli, David Grangier, Denis Yarats, and Yann N. Dauphin. 2017. Convolutional Sequence to Sequence Learning. In _ICML_.
* Gong et al. (2018) Jingjing Gong, Xipeng Qiu, Shaojing Wang, and Xuang Huang. 2018. Information Aggregation via Dynamic Routing for Sequence Encoding. In _COLING_.
* Hassan et al. (2018)Hany Hassan, Anthony Aue, Chang Chen, Vishal Chowdhary, Jonathan Clark, Christian Federmann, Xuedong Huang, Marcin Junczys-Dowmunt, William Lewis, Mu Li, et al. 2018. Achieving Human Parity on Automatic Chinese to English News Translation. _arXiv preprint arXiv:1803.05567_.
* Hinton et al. (2011) Geoffrey E Hinton, Alex Krizhevsky, and Sida D Wang. 2011. Transforming Auto-encoders. In _ICANN_.
* Hinton et al. (2018) Geoffrey E Hinton, Sara Sabour, and Nicholas Frosst. 2018. Matrix Capsules with EM Routing. In _ICLR_.
* Koehn (2004) Philipp Koehn. 2004. Statistical Significance Tests for Machine Translation Evaluation. In _EMNLP_.
* LaLonde and Bagci (2018) Rodney LaLonde and Ulas Bagci. 2018. Capsules for object segmentation. _arXiv preprint arXiv:1804.04241_.
* Li et al. (2018) Jian Li, Zhaopeng Tu, Baosong Yang, Michael R. Lyu, and Tong Zhang. 2018. Multi-Head Attention with Disagreement Regularization. In _EMNLP_.
* Lin et al. (2017) Zhouhan Lin, Minwei Feng, Cicero Nogueira dos Santos, Mo Yu, Bing Xiang, Bowen Zhou, and Yoshua Bengio. 2017. A Structured Self-attentive Sentence Embedding. In _ICLR_.
* Luong et al. (2015) Minh-Thang Luong, Hieu Pham, and Christopher D Manning. 2015. Effective Approaches to Attention-based Neural Machine Translation. In _EMNLP_.
* Papineni et al. (2002) Kishore Papineni, Salim Roukos, Todd Ward, and Wei-Jing Zhu. 2002. BLEU: A method for Automatic Evaluation of Machine Translation. In _ACL_.
* Peters et al. (2018) Matthew E. Peters, Mark Neumann, Mohit Iyyer, Matt Gardner, Christopher Clark, Kenton Lee, and Luke Zettlemoyer. 2018. Deep Contextualized Word Representations. In _NAACL_.
* Raganato and Tiedemann (2018) Alessandro Raganato and Jorg Tiedemann. 2018. An Analysis of Encoder Representations in Transformer-Based Machine Translation. In _EMNLP Workshop BlackboxNLP: Analyzing and Interpreting Neural Networks for NLP_ 

Sara Sabour, Nicholas Frosst, and Geoffrey E Hinton. 2017. Dynamic Routing Between Capsules. In _NIPS_.
* Sennrich et al. (2016) Rico Sennrich, Barry Haddow, and Alexandra Birch. 2016. Neural Machine Translation of Rare Words with Subword Units. In _ACL_.
* Shen et al. (2018) Tao Shen, Tianyi Zhou, Guodong Long, Jing Jiang, Shirui Pan, and Chengqi Zhang. 2018. DiSAN: Directional Self-Attention Network for RNN/CNN-free Language Understanding. In _AAAI_.
* Shi et al. (2016) Xing Shi, Inkit Padhi, and Kevin Knight. 2016. Does String-based Neural MT Learn Source Syntax? In _EMNLP_.
* Strubell et al. (2018) Emma Strubell, Patrick Verga, Daniel Andor, David Weiss, and Andrew McCallum. 2018. Linguistically-Informed Self-Attention for Semantic Role Labeling. In _EMNLP_.
* Tang et al. (2018) Gongbo Tang, Mathias Muller, Annette Rios, and Rico Sennrich. 2018. Why Self-Attention? A Targeted Evaluation of Neural Machine Translation Architectures. In _EMNLP_.
* Vaswani et al. (2017) Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. 2017. Attention Is All You Need. In _NIPS_.
* Wu et al. (2016) Yonghui Wu, Mike Schuster, Zhifeng Chen, Quoc V Le, Mohammad Norouzi, Wolfgang Macherey, Maxim Krikun, Yuan Cao, Qin Gao, Klaus Macherey, et al. 2016. Google's Neural Machine Translation System: Bridging the Gap between Human and Machine Translation. _arXiv preprint arXiv:1609.08144_.
* Xi et al. (2017) Edgar Xi, Selina Bing, and Yang Jin. 2017. Capsule network performance on complex data. _arXiv preprint arXiv:1712.03480_.
* Xu et al. (2015) Kelvin Xu, Jimmy Ba, Ryan Kiros, Kyunghyun Cho, Aaron C Courville, Ruslan Salakhutdinov, Richard S Zemel, and Yoshua Bengio. 2015. Show, Attend and Tell: Neural Image Caption Generation with Visual Attention. In _ICML_.
* Yang et al. (2019) Baosong Yang, Longyue Wang, Derek F. Wong, Lidia S. Chao, and Zhaopeng Tu. 2019. Convolutional self-attention networks. In _NAACL_.
* Zhao et al. (2018) Wei Zhao, Jianbo Ye, Min Yang, Zeyang Lei, Suofei Zhang, and Zhou Zhao. 2018. Investigating Capsule Networks with Dynamic Routing for Text Classification. In _EMNLP_