Talking-Heads Attention

2020·Arxiv

Abstract

Abstract

We introduce "talking-heads attention" - a variation on multi-head attention which includes linear projections across the attention-heads dimension, immediately before and after the softmax operation. While inserting only a small number of additional parameters and a moderate amount of additional computation, talking-heads attention leads to better perplexities on masked language modeling tasks, as well as better quality when transfer-learning to language comprehension and question answering tasks.

1 Introduction

Neural Attention was introduced by [Bahdanau et al., 2014] as a way of extracting information from variablelength representations. The Transformer model [Vaswani et al., 2017] uses "multi-head" attention, consisting of multiple attention layers ("heads") in parallel, each with different projections on its inputs and outputs. By using a dimensionality reduction in the input projections, the computational cost is kept similar to that of basic attention. Quality is improved, presumably due to the ability to attend to multiple positions simultaneously based on multiple different types of relationships.

As noted in [Vaswani et al., 2017], taking this process to the extreme (more attention heads projected to lower dimensionality) becomes counterproductive. We believe that this is due to the fact that the query-vectors and key-vectors become so low-dimensional that their dot product can no longer constitute an informative matching function.

In this paper, we introduce a new variant, "talking-heads attention", that addresses this problem by inserting a learned linear projection across the attention-heads dimension of the attention-logits tensor. This allows each attention function to depend on all of the keys and queries. We also insert a second such projection immediately following the softmax.

We show experimentally that inserting these "talking-heads" projections leads to better perplexities on masked language modeling tasks, as well as better quality when transfer-learning to language comprehension and question answering tasks.

2 Notation

In our pseudocode, we use capital letters to represent tensors and lower-case letters to represent their dimensions. Each tensor is followed by a dimension list in brackets. For example, a 4-dimensional image-

tensor with (batch, height, width, channels) dimensions would be written as:

We use einsum notation for generalized contractions between tensors of arbitrary dimension. The computation is numerically equivalent to broadcasting each input to have the union of all dimensions, multiplying component-wise, and summing across all dimensions not in the output. Rather than identifying the dimensions by an equation, as in TensorFlow and numpy, the dimensions are indentified by the dimension-list annotations on the arguments and on the result. For example, multiplying two matrices would be expressed as:

3 Review of Attention Algorithms

3.1 Dot-Product Attention

Simple dot-product attention can be described by the pseudocode below. The logits L are computed as the dot-products of the query-vectors and the memory-vectors. For each query, the logits are passed through a softmax function to produce weights, and the different memory-vectors are averaged together, weighted by those weights. In this code, we show the case where there are n different queries all attending to the same m memory-vectors. If there is only one query, the code is identical except that the "n" dimension is removed from all tensors.

3.2 Dot-Product Attention With Projections

[Vaswani et al., 2017] propose a dimensionality-reduction to reduce the computational complexity of the attention algorithm. In this version, instead of computing the attention algorithm directly on the inputs X and M, we first project the inputs using the learned linear projections , to produce lower-dimensional query-vectors, key-vectors and value-vectors Q, K and V . We use a fourth learned linear projection, , to produce the output.

3.3 Multi-Head Attention

The multi-head attention described in [Vaswani et al., 2017] consists of the sum of multiple parallel attention layers. This can be represented by adding a "heads" dimension h to the above computation.

The pseudo-code above illustrates the practical step-by-step computation of multi-head attention. The costs of the einsum operations (the number of multiplications in a naive implementation) are shown in the comments. The equivalent pseudo-code below uses multi-way einsums and is more concise:

Note: [Vaswani et al., 2017] include a constant scaling factor on the logits. We omit this in our code, as it can be folded into the linear projections or .

4 Talking-Heads Attention

In multi-head attention, the different attention heads perform separate computations, which are then summed at the end. Our new variation, which we call "Talking-Heads Attention" breaks that separation. We insert two additional learned linear projections, , which transform the attention-logits and the attentionweights respectively, moving information across attention heads. Instead of one "heads" dimension h across the whole computation, we now have three separate heads dimensions: , and , which can optionally differ in size (number of "heads"). refers to the number of attention heads for the keys and the queries. h refers to the number of attention heads for the logits and the weights, and refers to the number of attention heads for the values. The algorithm is shown by the pseudo-code below. The costs of the einsum operations are shown in the comments.

Again, we can write this more concisely using multi-way einsum operations:

5 Complexity Analysis

If we assume that , then the number of scalar multiplications in multi-head attention is:

The number of scalar multiplications in talking-heads attention is:

The first term in this expression matches up with the cost of multi-head attention. The second term is due to the talking-heads projections. If and , the the costs of the new talking-heads projections, and are less than the existing terms and , respectively.

In practice, the talking-heads projections may be expensive on some neural-network accelerators due to the small dimension sizes involved.

6 One More Way To Look At It

Mathematically, one can view multi-head attention and talking-heads attention as two special cases of the same general function, which we will call "general bilinear multihead attention" (GBMA). GBMA uses two three-dimensional parameter tensors, as defined in the pseudocode below. Due to its high computational cost, GBMA may have no practical use. Multi-head attention is mathematically equivalent to a version of GBMA where each of the two parameter tensors is expressed as the product of two factors, as shown below. Talking-heads attention is mathematically equivalent to a version of GBMA where each of the two parameter

tensors is expressed as the product of three factors, as shown below.

P[d_X , d_M , h] = einsum(P_q[d_X , d_k , h_k], P_k[d_M , d_k , h_k], P_l[h_k , h]) Q[d_M , d_Y , h] = einsum(P_v[d_M , d_v , h_v], P_o[d_Y , d_v , h_v], P_w[h, h_v ]) return GeneralBilinearMultiheadAttention (X, M, P, Q)

7 Experiments

7.1 Text-to-Text Transfer Transformer (T5)

We test various configurations of multi-head attention and talking-heads attention on the transfer-learning setup from [Raffel et al., 2019]. An encoder-decoder transformer model [Vaswani et al., 2017] is pre-trained on a denoising objective of predicting missing text segments (average span length 3) from the C4 dataset [Raffel et al., 2019] , and subsequently fine-tuned on various language understanding tasks. We use the same code base and model architecture as the base model from [Raffel et al., 2019]. The encoder and decoder each consist of 12 layers, with and each decoder layer contains a multi-head self-attention layer and a multi-head attention-over-encoder layer. For their base model, [Raffel et al., 2019] follow [Devlin et al., 2018] and others, using h = 12 and multi-head and talking-heads attention, as detailed in table 1.

Similar to [Raffel et al., 2019], we pre-train our models for 524288 steps. Each training batch consists of 128 examples, each of which has an input of 512 tokens and an output of 114 tokens, the output containing multiple spans of tokens which were deleted from the input. Similarly to [Raffel et al., 2019], we use the Adafactor optimizer [Shazeer and Stern, 2018] and an inverse-square-root learning-rate schedule. We also decay the learning rate linearly for the final 10 percent of the training steps. Our main departure from [Raffel et al., 2019] is that we, as suggested by [Lan et al., 2019], use no dropout during pre-training. We find this to produce superior results. We compute the log-perplexity on the training objective on a held-out shard of C4, which we believe to be a good indicator of model quality. For each configuration, we train one model for the "full" 524288 steps and four models for a shorter time (65536 steps) to measure inter-run variability. The results are listed in table 1.

We then fine-tune each of the models on an examples-proportional mixture of SQUAD [Rajpurkar et al., 2016], GLUE [Wang et al., 2018] and SuperGlue [Wang et al., 2019]. Fine-tuning consists of 131072 additional steps with a learning rate of 10. Following [Raffel et al., 2019], we use a dropout rate 0.1 on the layer outputs, feed-forward hidden-layers and attention weights. The embedding matrix (also used as the projection in the final classifier layer) is fixed during fine-tuning. Tables 1, 2, 3 and 4 include results for SQUAD and MNLI-m. Results for all other tasks are listed in the appendix.

7.1.1 Multi-Head vs Talking-Heads Attention

In table 1, we compare multi-head attention to talking-heads attention. For each of the two algorithms, we test versions with 6, 12, 24 and 48 heads. Following [Vaswani et al., 2017], as we increase the number of

Table 1: Multi-Head vs. Talking-Heads attention on T5

heads, we decrease the key/value dimensionality , so as to keep the number of parameters constant. For each number of heads, talking-heads attention improves over multi-head attention on all quality metrics.

Additionally, multi-head attention gets worse as we increase the number of heads from 24 to 48 and decrease the key and value dimensionalty from 32 to 16, while talking-heads attention gets better. We presume that this is due to the keys being too short to produce a good matching signal.

For additional comparison, we include in table 1 two models with significantly more parameters and computation in the attention layers. In the first, we double the number of heads in our baseline model from 12 to 24 without reducing and , resulting in a multi-head attention layer with double the parameters and double the computation. In the second, we use "general bilinear multihead attention", as described in section 6.

We also list the results from [Raffel et al., 2019]. We believe that their results are worse due to their use of dropout during pre-training.

7.1.2 Varying the Heads-Dimensions Separately

Table 2: Talking-heads attention has three "heads" dimensions that can vary independently.

In table 2, we experiment with independently varying the sizes of the three heads-dimensions. From the results, it appears that all three are good to increase, but that the softmax-heads dimension h is particularly important.

7.1.3 Logits-Projection Only and Weights-Projection Only

In the middle two experiments of table 3, we examine hybrids of multi-head attention and talking-heads attention, where there is a projection on one but not both of the logits and the weights.

7.1.4 Encoder vs. Decoder

The transformer model contains three types of attention layers - self-attention in the encoder, self-attention in the decoder, and attention-over-encoder in the decoder. In each of the middle three experiments of table 4,

Table 3: The logits-projection and the weights-projection can be employed separately.

Table 4: In each of the middle three experiments, talking-heads attention is employed in only one of the three types of attention layers in the model.

we employ talking-heads attention in only one of these types of attention layers, and multi-head attention in the others. We find that modifying the encoder-self-attention layers has the biggest effect on the downstream language-understanding tasks. This is unsurprising, given that these tasks have more to do with analyzing the input than with generating output.

7.2 ALBERT

[Lan et al., 2019] introduce ALBERT, a variation on BERT [Devlin et al., 2018]. The main difference between the ALBERT and BERT architectures is that ALBERT shares layer parameters among all layers, significantly reducing the number of parameters. For example, a 12-layer ALBERT model has about 1/12 the number of parameters in the attention and feed-forward layers as a similar BERT model. Another difference is that the ALBERT model factorizes the word embedding as the product of two matrices with smaller bases, again significantly reducing the parameter count. This makes ALBERT appealing for memory limited devices such as mobiles. Besides above architecture differences, ALBERT also uses sentence order prediction (SOP) to replace next sentence prediction (NSP) in BERT.

We report here experiments done with the base setting for ALBERT: a Transformer network with 12 layers of attention, the hidden and embedding size set to 768. The pre-training and fine-tuning hyperparameter settings are also exactly the same as in [Lan et al., 2019]. We use the English Wikipedia and book corpus datasets [Devlin et al., 2018] to pre-train various models with different head sizes and talking-heads configurations. We evaluate the resulting representations by using them as a starting point to finetune for the SQuAD task (SQuAD1.1, SQuAD2.0 dev set) and various tasks (MNLI, SST-2, RACE) from the GLUE benchmark. Results are in Table 5.

Table 5: Multi-Head vs. Talking-Heads attention on ALBERT.

We find that as the number of heads increases beyond 12 and the dimensionality of the attention-keys and attention-values decreases below 64, the performance of multi-head attention decays. On the other hand, the performance of talking-head attention keeps improving.

In addition, we also compare the logits projection and the weight projection separately with multi-head and talking-heads attention. The results are shown in Table 6. Similar to our observation in T5 experiments, only applying either the logits projection or the weight projection does not result in significant improvement compared to without them. These results again confirm the importance of having both projections.

Table 6: The logits-projection and the weights-projection can be employed separately.

7.3 BERT

We test various configurations of talking-heads attention based on [Devlin et al., 2018]. All of our experiments use the simplified relative position embeddings [Raffel et al., 2019] instead of fixed position embedding. We first pre-train a 12 Transformer layers using the same dataset as [Devlin et al., 2018]. And then we finetune for the SQuAD1.1 task and the MNLI from the GLUE dataset. Our experiments show that quality continues to improve when we grow the number of heads up to 768 and decrease the key and value dimensionality down to 1

Table 7: Talking-Heads attention on BERT.

7.4 Visualizing the Projection Matrices of Talking-Heads

To illustrate how different heads exchange information with each other, we visualize the projection matrices (and ) of a 12 layer BERT with 12 talking-heads in figure 1. Since is applied after (although there is a softmax non-linearity in between), we also visualize the combined transformation 1. As can be observed, the main diagonals of the projection matrices do not have significant greater values than other entries. This is expected because with talking-heads, a pair of query and key do not corresponds to any specific value-vector. All keys and queries jointly decide how the values in each head interchange data. Additionally, all projection matrices are well conditioned (magnitude of determinant above 10smallest eigenvalue above 10), indicating that no significant approximation can be achieved.

Figure 1: Visualization of the learned projection matrices , and the multiplication . All entries in each matrix are normalized to [-1, 1]. The fact that these matrices are not diagonal or sparse at all, shows that there are significant data exchange across different attention heads.

8 Conclusions and Future Work

We have proposed talking-heads attention and shown some promising results. One potential challenge is speed on modern deep-learning accelerators, which are optimized for large-dimension matrix multiplications. We imagine that this will be an area of future work. One approach is to build hardware which is better at smalldimension matrix-multiplication. Another potential approach is to decrease the number of memory-positions considered for each query-position - for example, by using the local-attention and memory-compressed-attention approaches described in [Liu et al., 2018]. We look forward to more applications of talking-heads attention, as well as to further architectural improvements.

References

Dzmitry Bahdanau, Kyunghyun Cho, and Yoshua Bengio. Neural machine translation by jointly learning to align and translate, 2014. 1

Jacob Devlin, Ming-Wei Chang, Kenton Lee, and Kristina Toutanova. Bert: Pre-training of deep bidirectional transformers for language understanding. arXiv preprint arXiv:1810.04805, 2018. 5, 7, 8

Zhenzhong Lan, Mingda Chen, Sebastian Goodman, Kevin Gimpel, Piyush Sharma, and Radu Soricut. Albert: A lite bert for self-supervised learning of language representations, 2019. 5, 7

Peter J Liu, Mohammad Saleh, Etienne Pot, Ben Goodrich, Ryan Sepassi, Lukasz Kaiser, and Noam Shazeer. Generating wikipedia by summarizing long sequences. In Proceedings of the International Conference on Learning Representations, 2018. 9

Colin Raffel, Noam Shazeer, Adam Roberts, Katherine Lee, Sharan Narang, Michael Matena, Yanqi Zhou, Wei Li, and Peter Liu. Exploring the limits of transfer learning with a unified text-to-text transformer. arXiv e-prints, 2019. 5, 6, 8, 11, 13, 14, 15

Pranav Rajpurkar, Jian Zhang, Konstantin Lopyrev, and Percy Liang. Squad: 100,000+ questions for machine comprehension of text. arXiv preprint arXiv:1606.05250, 2016. 5, 12, 15

Noam Shazeer and Mitchell Stern. Adafactor: Adaptive learning rates with sublinear memory cost. arXiv preprint arXiv:1804.04235, 2018. 5

Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N. Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. In NIPS, 2017. 1, 2, 3, 5

Alex Wang, Amapreet Singh, Julian Michael, Felix Hill, Omer Levy, and Samuel R. Bowman. GLUE: A multi- task benchmark and analysis platform for natural language understanding. arXiv preprint arXiv:1804.07461, 2018. 5, 12, 13

Alex Wang, Yada Pruksachatkun, Nikita Nangia, Amanpreet Singh, Julian Michael, Felix Hill, Omer Levy, and Samuel R. Bowman. Superglue: A stickier benchmark for general-purpose language understanding systems. arXiv preprint arXiv:1905.00537, 2019. 5, 12, 14

A Variation: Dynamic Projections

In the basic talking-heads attention algorithm described in section 4, the talking-heads projections are represented by two learned weight matrices ] and ]. In an additional wrinkle, we can make these projections matrices themselves input-dependent, adding terms to the projection matrices that are themselves learned linear projections of the inputs X and M. The algorithm is described by the pseudo-code below.

We observed that the model only trained well if we initialized the projection-generating parameter matrices () to contain small enough values. We used normal initializers with standard deviations of 0, 0, 0, and 0, respectively.

Table 8: Dynamic-Projections Results on T5

A.1 Experiments

We evaluate talking-heads attention with dynamic projections on T5 [Raffel et al., 2019] in a set of experiments similar to those described in section 7.1.

Table 8 compares multi-head attention, taking-heads attention with static projections, and talking-heads attention with dynamic projections. The dynamic projections reduce perplexity on the pre-training task. However, in our experiments, we did not see an improvement on the downstream tasks.

Table 9: In each of the middle four experiments, only one of the dynamic projections is employed.

A.1.1 Comparing the Four Dynamic Projections

Table 9 examines the effects of the four dynamic projections employed individually. The middle four rows represent experiments where only one of the four dynamic projections were employed. These are compared to static projections (top row) and all four dynamic projections together (bottom row).

Table 10: Effects of applying talking-heads attention (with or without dynamic projections) in the encoder only.

A.1.2 Talking-Heads in Encoder Only

In section 7.1.4 we saw that talking heads were particularly useful in the encoder part of the model. Table 10 presents a set of experiments where the decoder uses only multi-head attention, while the encoder uses either multi-head attention (top row), talking-heads attention with static projections (middle row), or talking-heads attention with dynamic projections (bottom row). We observe that in this case, the dynamic projections do not appear to degrade performance on the downstream tasks.

B T5 Fine-Tuning Full Results

Tables 11, 12 and 13 present the results of fine-tuning the models in section 7.1 and appendix A on the GLUE [Wang et al., 2018] and SuperGlue [Wang et al., 2019], and Stanford Question-Answering Dataset (SQuAD) [Rajpurkar et al., 2016] benchmarks.

Table 11: T5 on GLUE Language-Understanding Benchmark [Wang et al., 2018] (dev). Experiments described in Section 7.1 and appendix A.

Table 12: T5 on SuperGLUE Language-Understanding Benchmark [Wang et al., 2019] (dev). Experiments described in Section 7.1 and appendix A.

Table 13: T5 on SQuAD [Rajpurkar et al., 2016] v1.1 (dev). Experiments described in Section 7.1 and appendix A.