On the diversity of multi-head attention

Abstract

Multi-head attention is appealing for the ability to jointly attend to information from different representation subspaces at different positions. In this work, we propose two approaches to better exploit such diversity for multi-head attention, which are complementary to each other. First, we introduce a disagreement regularization to explicitly encourage the diversity among multiple attention heads. Specifically, we propose three types of disagreement regularization, which respectively encourage the subspace, the attended positions, and the output representation associated with each attention head to be different from other heads. Second, we propose to better capture the diverse information distributed in the extracted partial-representations with the routing-by-agreement algorithm. 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 the machine translation, sentence encoding and logical inference tasks demonstrate the effectiveness and universality of the proposed approaches, which indicate the necessity of better exploiting the diversity for multi-head attention. While the two strategies individually boost performance, combining them together can further improve the model performance.

Introduction

Attention model becomes a standard component of the deep learning networks, contributing to impressive results in machine translation [1], [2], image captioning [3], speech recognition [4], among many other applications. Recently, the performance of attention is further improved by multi-head mechanism [5], which concurrently performs the attention functions on different representation subspaces of the input sequence. Consequently, different attention heads are able to capture distinct properties of the input, which are embedded in different subspaces [6]. Subsequently, a linear transformation is generally employed to aggregate the partial representations extracted by different attention heads [5], [7], producing the final output representation.

However, the conventional multi-head mechanism may not fully exploit the diversity among attention heads. First, one strong point of multi-head attention is the ability to jointly attend to information from different representation subspaces at different positions. But currently there is no mechanism to guarantee that different attention heads indeed capture distinct information. Second, we believe that information extraction and information aggregation are both important to produce an informative representation. We argue that the straightforward linear transformation are not expressive enough to fully capture the rich information distributed in the extracted partial-representations. In this work, we propose two strategies to better exploit the diversity of multi-head attention, namely disagreement regularization and advanced information aggregation.

In response to the first problem, we introduce a disagreement regularization term to explicitly encourage the diversity among multiple attention heads. The disagreement regularization serves as an auxiliary objective to guide the training of the related attention component. Specifically, we propose three types of disagreement regularization, which are applied to the three key components that refer to the calculation of information vector using multi-head attention. Two regularization terms are respectively to maximize cosine distances of the input subspaces and output representations, while the last one is to disperse the positions attended by multiple heads with element-wise multiplication of the corresponding attention matrices. The three regularization terms can be either used individually or in combination.

To address the second problem, we replace the standard linear transformation in conventional multi-head attention [5] with an advanced routing-by-agreement algorithm, to better aggregate the diverse information distributed in the extracted partial-representations. Specifically, we cast information aggregation as the assigning-parts-to-wholes problem [8], and investigate the effectiveness of the routing-by-agreement algorithm, which is an appealing alternative to solving this problem [9], [10]. 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.

In addition, it is natural to combine the two types of approaches and apply them simultaneously, since the former focuses on extracting more diverse information while the latter aims to better aggregate the extracted information. We apply them simultaneously by modifying both the training objective and network architecture.

We evaluate the performance of the proposed approaches on three representative NLP tasks: machine translation, sentence encoding, and logical inference tasks. For machine translation, we validate our approaches on top of the advanced Transformer model [5] on both WMT14 English$\Rightarrow$German and WMT17 Chinese$\Rightarrow$English data. Experimental results show that our approaches consistently improve the translation performance across language pairs while keeping the computational efficiency. For sentence encoding, we evaluate with the linguistic probing tasks [11], which consist of 10 classification problems to study what linguistic properties are captured by input encoding representations. Probing analysis shows that our approaches indeed produce more informative representation, which embeds more syntactic and semantic information. Experiments on logical inference further demonstrate the ability of modeling hierarchical structure. Precisely, our study reveals that:

• Directly applying disagreement regularization on the output representations of multiple attention heads is most effective.

• The EM routing algorithm shows its superiority on information aggregation over the standard linear transformation and other aggregation algorithms.

• Disagreement regularization and advanced information aggregation are complementary to each other, as indicated from analyses in machine translation and sentence encoding.

This paper combines and extends results presented at the 2018 Conference on Empirical Methods in Natural Language Processing (entitled “Multi-Head Attention with Disagreement Regularization” [12]) and at the 2019 Annual Conference of the North American Chapter of the Association for Computational Linguistics (entitled “Information Aggregation for Multi-Head Attention with Routing-by-Agreement” [13]). The extensions include

• We further refine our proposed model by combining the two sorts of strategies and exploiting the advantages of simultaneously applying them (Section 3.3). We demonstrate the effectiveness of the combined method in experiments (Table 5).

• We carry out more experiments and in-depth analyses to validate the effectiveness of our approaches on more tasks, including linguistic probing tasks (Section 4.2) and and logical inference tasks (Section 4.3). Results on linguistic probing tasks prove the superiority of our approach on capturing surface, syntactic and semantic information. Results on logical inference tasks show that the proposed approach performs better at modeling hierarchical structure.

• We present a more comprehensive description of the proposed models and algorithms (Section 3).

• For reproducibility, we release the source code, preprocessed data, and trained models, which make it easy to reproduce the experiments in this work.¹