Concise Multi-head Attention Models

1 INTRODUCTION . Attention based architectures , such as Transformers , have been effective for sequence modelling tasks such as machine translation ( Gehring et al. , 2017 ; Vaswani et al. , 2017 ) , question answering , sentence classification ( Radford et al. , 2018 ; Devlin et al. , 2018 ) and document generation ( Liu et al. , 2018 ) . These models have emerged as better alternatives to the recurrent models - RNNs ( Sutskever et al. , 2014 ) , LSTMs ( Hochreiter & Schmidhuber , 1997 ) and GRUs ( Cho et al. , 2014 ) . This is mainly due to their feed forward structure , which removes the sequential processing bottleneck for sequence data , making them easier to train compared to the recurrent models . Self attention models also have found applications in vision ( Wang et al. , 2018 ) , adversarial networks ( Zhang et al. , 2018 ) , reinforcement learning ( Zambaldi et al. , 2018 ; Li , 2017 ) and speech recognition ( Chiu et al. , 2018 ) . Recent advances in using the self attention models in natural language tasks have been made by first using a language modeling task to pre-train the models and then fine tuning the learned models on specific downstream tasks . Radford et al . ( 2018 ) and Devlin et al . ( 2018 ) used Transformers to pre-train a language model and showed that the fine tuned model outperforms LSTMs on many natural language understanding and question answering tasks . For example , BERT ( Devlin et al. , 2018 ) , a 24 layer transformer model , is shown to achieve the state of the art performance on several NLP tasks , including on the SQuAD dataset . These advances , in addition to novel pre-training tasks , relied on bigger models with a larger embedding size . BERT model uses an embedding size of 1024 ( Devlin et al. , 2018 ) ; GPT-2 uses models with embedding size up to 1600 ( Radford et al. , 2019 ) . A single Transformer block consists of two key components : a multi-head self attention layer followed by a feed forward layer ( Vaswani et al. , 2017 ) . A single head in a multi-head attention layer , computes self attention between the tokens in the input sequence , which it then uses to compute a weighted average of embeddings for each token . To keep the number of parameters fixed in the attention layer , each head projects the data into a lower dimensional subspace , dimension of which scales as 1/ ( number of heads ) , and computes the self attention in this subspace . This projection size for each head is commonly known as the head size . Despite the advances in using Transformer models for various tasks , their functioning and design choices still remain mysterious and are not well understood . Can the attention layer learn arbitrary contextual representations ? What is the role of the feed forward layer in the Transformer block ? Do we need such a large embedding size to capture the context of all the tokens ? Answering these questions requires an understanding of the representation power of the units in the Transformer . In this paper we take some of the first steps towards developing such an understanding of the Transformer . In particular , we focus on the representation power of the multi-head self attention layer . First , we analyze the representation power of a single self attention unit and show that it crucially depends on the projection sizes used to compute the dot product attention . We next study the advantage of having multiple heads in the attention layer . It is generally believed that increasing the number of heads helps by allowing the heads to compute context from different representation subspaces at different positions . However , increasing the number of heads decreases the head size , decreasing the expressive power of individual heads . We show that reducing the head size to a value below the input sequence length hurts the representation power of each head . This is because a smaller head size introduces a rank constraint on the projection matrices in each head , and limits their representation power . We indeed notice this effect in practice : while the performance improves with increasing the number of heads in the beginning ( Devlin et al. , 2018 ) , we notice a drop in the performance once the number of heads increases beyond a certain threshold , as seen in Table 1 and Fig . 1 ( see also Table 4 ( A ) in Vaswani et al . ( 2017 ) ) . This heuristic of scaling the head size inversely with the number of heads was proposed initially in Vaswani et al . ( 2017 ) and has become the standard way of using multi-head attention ( Radford et al. , 2018 ; Devlin et al. , 2018 ) . In order to avoid hurting the performance , the existing Transformer models allow for multiple heads by increasing the embedding size , which in turn increases the head size . However , larger embedding size , in addition to increasing the number of parameters , makes it expensive to use the model and the learned embeddings in downstream tasks , as the downstream model sizes scale with the embedding size of the tokens . For example , the inference time and memory required in retrieval tasks increases linearly with the embedding size . Based on these observations , we propose a new way to set the projection size in the attention heads , in which each head has a fixed head size that is independent of both the number of heads and the embedding size of the model . This allows us to train models with a relatively smaller embedding size without affecting the head size . It also allows us to increase the number of heads per layer to improve the performance . Another advantage of the fixed head size Transformer is , unlike the standard Transformer , which requires the number of heads to be a factor of the embedding size , we are free to set arbitrary number of heads as required for the task . We evaluate this fixed head size Transformer on language modeling ( LM1B dataset ) , natural language inference ( MNLI dataset ) and question answering tasks ( SQuAD dataset ) . We show that the modified Transformer trained with an embedding size of 512 can match the performance of the BERTLARGE ( Devlin et al. , 2018 ) , a Transformer with an embedding size of 1024 ( see Fig . 2 ) . We further present experimental results evaluating the effect of different choices of the head size and the embedding size in the Section 4 . The contributions of this paper are summarized below . • We analyze the representation power of the multi-head self attention layer and show the limitation the embedding size places on the number of heads . • We propose a new way to set the head size , and show the proposed fixed head size layers are strictly better than the standard multi-head attention layers in terms of their expressive power . This modification allows us to both increase the number of heads per layer and decrease the embedding size , without hurting the performance . • We experimentally show that the fixed head size Transformer can be trained with a smaller embedding size and more heads on three standard NLP tasks . 1.1 RELATED WORKS . Given the significance of self attention models , there has been work trying to both improve the performance and speedup the computation in Transformers . Ott et al . ( 2018 ) and You et al . ( 2019 ) reduce precision and use large batch training to reduce the training time of the attention models . Child et al . ( 2019 ) propose sparse self attention models to speed up the computation in the attention layer for long sequence data generation tasks . They show that these sparse attention models can be trained on tasks with sequence length greater than 10k without sacrificing the accuracy . Dehghani et al . ( 2018 ) propose a depth recurrent Transformer network that reuses the parameters across layers . They show that this modification makes the Transformer networks Turing complete even with finite precision weights . Yang et al . ( 2019 ) propose a new way to increase the effective sequence length that the Transformer attends to , by reusing the intermediate embeddings across sequences . They show that the modified architecture performs better on tasks that require computing context over longer sequence lengths . We note that most of these modifications rely on the multi-head self attention , the same building block of the Transformers . Our work is studying this basic multi-head attention layer , and suggesting a new way to set the head size , which can be easily used along with any of the above architectural modifications . Wu et al . ( 2019 ) propose to replace the self-attention layer with lightweight dynamic convolutions and show improved performance on machine translation and language modeling . Even though the resulting model has faster inference time , it still needs to use a large embedding size ( 1024 ) , as big as the original attention models . We believe the techniques in this paper can be combined with these results to realize both smaller embedding size and faster inference time . Sun et al . ( 2019 ) perform neural architecture search using evolutionary methods on sequence to sequence models and find an evolved transformer architecture , which in addition to multi-head attention units , has convolution filter and gated linear units . Our proposed modifications stay closer to Transformers in spirit and can be used as seed units for this architecture search . Voita et al . ( 2019 ) ; Michel et al . ( 2019 ) study the importance of different heads in an attention layer . They observe that , during inference , many of the heads in each layer can be pruned away with a little effect on the prediction . However , they still need multiple heads during the training . Child et al . ( 2019 ) ; Correia et al . ( 2019 ) impose sparsity structure on the attention layer during training to improve both interpretability and performance . Fixing the head size will in fact make it easier to learn such sparsity patterns , as a low rank constraint does not allow a head to express all possible sparsity patterns . Combining these techniques can hence potentially enable training of sparse attention models with a smaller embedding size . 2 TRANSFORMER ARCHITECTURE AND ANALYSIS . In this section we present the Transformer architecture and analyze the representation power of the multi-head self attention , a key component of the Transformer block . The input to a Transformer network is a sequence of n tokens . Typically , each token is converted into a token embedding of dimension d by an embedding layer . We let X ∈ Rd×n be the embedding matrix corresponding to the n tokens in the input sequence . 2.1 SINGLE-HEAD ATTENTION . The transformer block is a combination of a self attention layer followed by a feed forward layer ( Vaswani et al. , 2017 ) . Both layers have a skip connection and use Layer Normalization ( LN ) ( Ba et al. , 2016 ) . In particular , for token embeddings X , the dot product attention is computed as follows . Attention ( X ) = WvX · Softmax [ ( WkX ) T ( WqX ) √ dk ] = WvX ·P . ( 1 ) Here Wq ∈ Rdq×d , Wk ∈ Rdk×d and Wv ∈ Rdv×d represent the projection matrices associated with the query , key and value respectively in an attention unit ( Vaswani et al. , 2017 ) . For a single-head attention unit , we have dq = dk = dv = d. In the dot-product attention ( cf . ( 1 ) ) , P aims to capture the context of the input for a given token based on the remaining tokens in the input sequence . Subsequently , the output of the attention layer takes the following form . LN ( X+Wo · Attention ( X ) ) , ( 2 ) where LN ( · ) represents the layer-normalization operation . Given the attention module , as defined in ( 1 ) , it is natural to question its ability to represent arbitrary contexts P for a given input sequence X . Towards this , we show that for a large enough projection dimension d , the unit has enough capacity to represent arbitrary contexts over a given input sequence . In the following result we establish that for a large enough projection size an attention unit can represent any data pair ( X , P ) . We also show that the model can not represent arbitrary context when d is smaller than n. Theorem 1 ( Representation Theorem ) . If dq = dk = d ≥ n , then given any full column rank matrix X ∈ Rd×n and an arbitrary n× n positive column stochastic matrix P , there always exists d× d projection matrices Wq and Wk such that Softmax [ ( WkX ) T ( WqX ) √ dk ] = P. ( 3 ) If dq = dk = d < n , there exist X and P such that ( 3 ) does not hold for all Wq and Wk . The proof of Theorem 1 is provided in the supplementary material . This result shows that the projection dimension dq = dk = d needs to be larger than the sequence length n for the attention unit to be able to represent any desired context P. Even though this result describes a single example sequence case , it highlights a fundamental property of the model architecture that increasing the projection size increases the capacity of the attention heads .

Summary: This work studies the head size <--> head number tradeoff in multihead attention. It argues and formally establishes that (1) the expressivity of an attention head is determined by its dimension and (b) fixing the head dimension, one gains additional expressive power by using more heads. In response to such observations, the paper proposes Fixed Multihead Attention, where the constraint that `head_size * number_of_heads = embedding_size` in standard multihead attention is lifted; and it allows for using more attention heads without making each head smaller. One can control the total amount of parameters by using smaller embedding sizes, making it comparable (in terms of #parameters) to standard multihead attention. Empirical results on language modeling and NLI tasks confirms the arguments.