Multi-Head Attention: Collaborate Instead of Concatenate

1 INTRODUCTION . Since the invention of attention ( Bahdanau et al. , 2014 ) and its popularization in the transformer architecture ( Vaswani et al. , 2017 ) , multi-head attention ( MHA ) has become the de facto architecture for natural language understanding tasks ( Devlin et al. , 2019 ) and neural machine translation . Attention mechanisms have also gained traction in computer vision following the work of Ramachandran et al . ( 2019 ) and Bello et al . ( 2019 ) . Nevertheless , despite their wide adoption , we currently lack solid theoretical understanding of how transformers operate . In fact , many of their modules and hyperparameters are derived from empirical evidences that are possibly circumstantial . The uncertainty is amplified in multi-head attention , where both the roles and interactions between heads are still poorly understood . Empirically , it is well known that using multiple heads can improve model accuracy . However , not all heads are equally informative , and it has been shown that certain heads can be pruned without impacting model performance . For instance , Voita et al . ( 2019 ) present a method to quantify head utility and prune redundant members . Michel et al . ( 2019 ) go further to question the utility of multiple heads by testing the effect of heavy pruning in several settings . On the other hand , Cordonnier et al . ( 2020 ) prove that multiple heads are needed for self-attention to perform convolution , specifically requiring one head per pixel in the filter ’ s receptive field . Beyond the number of heads , finding the adequate head dimension is also an open question . Bhojanapalli et al . ( 2020 ) finds that the division of the key/query projection between heads gives rise to a low-rank bottleneck for each attention head expressivity that can be fixed by increasing the head sizes . In contrast , our approach increases heads expressivity by leveraging the low-rankness accross heads to share common query/key dimensions . This work aims to better detect and quantify head redundancy by asking whether independent heads learn overlapping or distinct concepts . This relates to the work on CNN compression that factorizes common filters in a trained convolutional network ( Kim et al. , 2016 ) using Tucker decomposition . In attention models , we discover that some key/query projected dimensions are redundant , as trained 1https : //github.com/ ... concatenated heads tend to compute their attention patterns on common features . Our finding implies that MHA can be re-parametrized with better weight sharing for these common projections and a lower number of parameters . This differs from concurrent work ( Shazeer et al. , 2020 ) that orchestrate collaboration between heads on top of the dot product attention scores . Contribution 1 : Introducing the collaborative multi-head attention layer . Section 3 describes a collaborative attention layer that allows heads to learn shared key and query features . The proposed re-parametrization significantly decreases the number of parameters of the attention layer without sacrificing performance . Our Neural Machine Translation experiments in Section 4 show that the number of FLOPS and parameters to compute the attention scores can be divided by 4 without affecting the BLEU score on the WMT14 English-to-German task . Contribution 2 : Re-parametrizing pre-trained models into a collaborative form renders them more efficient . Pre-training large language models has been central to the latest NLP developments . But pre-training transformers from scratch remains daunting for its computational cost even when using more efficient training tasks such as ( Clark et al. , 2020 ) . Interestingly , our changes to the MHA layers can be applied post-hoc on pre-trained transformers , as a drop-in replacement of classic attention layers . To achieve this , we compute the weights of the re-parametrized layer using canonical tensor decomposition of the query and key matrices in the original layer . Our experiments in Section 4 show that the key/query dimensions can be divided by 3 without any degradation in performance . As a side contribution , we identify a discrepancy between the theory and some implementations of attention layers and show that by correctly modeling the biases of key and query layers , we can clearly differentiate between context and content-based attention . 2 MULTI-HEAD ATTENTION . We first review standard multi-head attention introduced by Vaswani et al . ( 2017 ) . 2.1 ATTENTION . Let X ∈ RT×Din and Y ∈ RT ′×Din be two input matrices consisting of respectively T and T ′ tokens of Din dimensions each . An attention layer maps each of the T query token from Din to Dout dimensions as follows : Attention ( Q , K , V ) = softmax ( QK > √ dk ) V , with Q = XWQ , K = Y WK , V = Y WV ( 1 ) The layer is parametrized by a query matrix WQ ∈ RDin×Dk , a key matrix WK ∈ RDin×Dk and a value matrix WV ∈ RDin×Dout . Using attention on the same sequence ( i.e . X = Y ) is known as self-attention and is the basic building block of the transformer architecture . 2.2 CONTENT VS . CONTEXT . Some re-implementations of the original transformer architecture2 use biases in the linear layers . This differs from the attention operator defined in eq . ( 1 ) where the biases bQ and bK ∈ RDk are ommited . Key and query projections are computed as K = XWK + 1T×1bK and Q = Y WQ + 1T×1bQ , respectively , where 1a×b is an all one matrix of dimension a × b . The exact computation of the ( unscaled ) attention scores can be decomposed as follows : QK > = ( XWQ + 1T×1b > Q ) ( Y WK + 1T×1b > K ) > ( 2 ) = XWQW > K Y > ︸ ︷︷ ︸ context +1T×1b > QW > K Y > ︸ ︷︷ ︸ content +XWQbK11×T + 1T×T b > QbK ( 3 ) As the last two terms of eq . ( 3 ) have a constant contribution over all entries of the same row , they do not contribute to the computed attention probabilities ( softmax is shift invariant and softmax ( x + c ) = softmax ( x ) , ∀c ) . On the other hand , the first two terms have a clear meaning : XWQW > K Y > 2For instance : the BERT orignal implementation , its HuggingFace re-implementation and FairSeq encoderdecoder transformer . considers the relation between keys and query pairs , whereas 1T×1b > QW > K Y > computes attention solely based on key content . The above findings suggest that the bias bK of the key layer can be always be disabled without any consequence . Moreover , the query biases bQ play an additional role : they allow for attention scores that are content-based , rather than solely depending on key-query interactions . This could provide an explanation for the recent success of the Dense-SYNTHESIZER ( Tay et al. , 2020 ) , a method that ignores context and computes attention scores solely as a function of individual tokens . That is , perhaps context is not always crucial for attention scores , and content can suffice . 2.3 MULTI-HEAD ATTENTION . Traditionally , the attention mechanism is replicated by concatenation to obtain multi-head attention defined for Nh heads as : MultiHead ( X , Y ) = concat i∈ [ Nh ] [ H ( i ) ] WO ( 4 ) H ( i ) = Attention ( XW ( i ) Q , Y W ( i ) K , Y W ( i ) V ) , ( 5 ) where distinct parameter matrices W ( i ) Q , W ( i ) K ∈ RDin×dk and W ( i ) V ∈ RDin×dout are learned for each head i ∈ [ Nh ] and the extra parameter matrix WO ∈ RNhdout×Dout projects the concatenation of the Nh head outputs ( each in Rdout ) to the output space RDout . In the multi-head setting , we call dk the dimension of each head and Dk = Nhdk the total dimension of the query/key space . 3 IMPROVING THE MULTI-HEAD MECHANISM . Head concatenation is a simple and remarkably practical setup that gives empirical improvements . However , we show that another path could have been taken instead of concatenation . As the multiple heads are inherently solving similar tasks , they can collaborate instead of being independent . 3.1 HOW MUCH DO HEADS HAVE IN COMMON ? . We hypothesize that some heads might attend on similar features in the input space , for example computing high attention on the verb of a sentence or extracting some dimensions of the positional encoding . To verify this hypothesis , it does not suffice to look at the similarity between query ( or key ) matrices { W ( i ) Q } i∈ [ Nh ] of different heads . To illustrate this issue , consider the case where two heads are computing the same key/query representations up to a unitary matrix R ∈ Rdk×dk such that W ( 2 ) Q = W ( 1 ) Q R and W ( 2 ) K = W ( 1 ) K R. Even though the two heads are computing identical attention scores , i.e . W ( 1 ) Q RR > W ( 1 ) > K = W ( 1 ) Q W ( 1 ) > K , they can have orthogonal column-spaces and the concatenation [ W ( 1 ) Q , W ( 2 ) Q ] ∈ RDin×2dk can be full rank . To disregard artificial differences due to common rotations or scaling of the key/query spaces , we study the similarity of the product W ( i ) Q W ( i ) > K ∈ RDin×Din across heads . Figure 1 shows the captured energy by the principal components of the key , query matrices and their product . It can be seen on the left that single head key/query matrices W ( i ) Q W ( i ) > K are not low rank on average . However , as seen on the right , even if parameter matrices taken separately are not low rank , their concatenation is indeed low rank . This means that heads , though acting independently , learn to focus on the same subspaces . The phenomenon is quite pronounced : one third of the dimensions suffices to capture almost all the energy of WQW > K , which suggests that there is inefficiency in the way multi-head attention currently operate . 3.2 COLLABORATIVE MULTI-HEAD ATTENTION . Following the observation that heads ’ key/query projections learn redundant projections , we propose to learn key/query projections for all heads at once and to let each head use a re-weighting of these projections . Our collaborative head attention is defined as follows : CollabHead ( X , Y ) = concat i∈ [ Nh ] [ H ( i ) ] WO ( 6 ) H ( i ) = Attention ( XW̃Q diag ( mi ) , Y W̃K , Y W ( i ) V ) . ( 7 ) The main difference with standard multi-head attention defined in eq . ( 5 ) is that we do not duplicate the key and query matrices for each head . Instead , each head learns a mixing vector mi ∈ RD̃k that defines a custom dot product over the D̃k projected dimensions of the shared matrices W̃Q and W̃K of dimension Din × D̃k . This approach leads to : ( i ) adaptive head expressiveness , with heads being able to use more or fewer dimensions according to attention pattern complexity ; ( ii ) parameter efficient representation , as learned projections are shared between heads , hence stored and learned only once . It is instructive to observe how standard multi-head attention ( where heads are simply concatenated ) can be seen as a special case of our collaborative framework ( with D̃k = Nhdk ) . The left of Figure 2 displays the standard attention computed between xn and ym input vectors with the mixing matrix M : = concat i∈ [ Nh ] [ mi ] ∈ RNh×D̃k , ( 8 ) laying out the mixing vectors mi as rows . In the concatenated MHA , the mixing vector mi for the i-th head is a vector with ones aligned with the dk dimensions allocated to the i-th head among the Dk = Nhdk total dimensions . Some alternative collaborative schema can be seen on the right side of Figure 2 . By learning the mixing vectors { mi } i∈ [ Nh ] instead of fixing them to this “ blocks-of-1 ” structure , we increase the expressive power of each head for a negligible increase in the number of parameters . The size dk of each head , arbitrarily set to 64 in most implementations , is now adaptive and the heads can attend to a smaller or bigger subspace if needed .

Summary: The paper investigates the over-parameterization of attention heads in Transformer’s multi-head attention. The authors show that query-key projections are redundant because trained concatenated heads tend to compute their attention patterns on common features. They propose a reparameterization of multi-head attention allowing the parameters of queries and keys to be shared between heads: this is called “collaborative attention”. This attention can be applied either instead of the standard attention during training, or as a drop-in replacement for an already trained model. To use as a drop-in replacement, the method requires to use tensor decomposition and subsequent model fine-tuning.