M ULTI -H EAD A TTENTION : C OLLABORATE I NSTEAD OF C ONCATENATE
Anonymous authors
Paper under double-blind review
A BSTRACT
Attention layers are widely used in natural language processing (NLP) and are beginning to inﬂuence computer vision architectures.
Training very large transformer models allowed signiﬁcan improvement in both ﬁelds, but once trained, these networks show symptoms of over-parameterization.
For instance, it is known that many attention heads can be pruned without impacting accuracy.
This work aims to enhance current understanding on how multiple heads interact.
Motivated by the observation that trained attention heads share common key/query projections, we propose a collaborative multi-head attention layer that enables heads to learn shared projections.
Our scheme decreases the number of parameters in an attention layer and can be used as a drop-in replacement in any transformer architecture.
For instance, by allowing heads to collaborate on a neural machine translation task, we can reduce the key dimension by 4 × without any loss in performance.
We also show that it is possible to re-parametrize a pre-trained multi-head attention layer into our collaborative attention layer.
Even without retraining, collaborative multi-head attention manages to reduce the size of the key and query projections by half without sacriﬁcing accuracy.
Our code is public.
I NTRODUCTION
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.
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 ampliﬁed 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, speciﬁcally requiring one head per pixel in the ﬁlter’s receptive ﬁeld.
Beyond the number of heads, ﬁnding the adequate head dimension is also an open question.
Bhojanapalli et al.
(2020) ﬁnds 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 ﬁxed 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.
We discover that some key/query projected dimensions are redundant, as trained concatenated heads tend to compute their attention patterns on common features.
Our ﬁnding 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 signiﬁcantly decreases the number of parameters of the attention layer without sacriﬁcing 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 efﬁcient.
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 efﬁcient 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.
M ULTI -H EAD A TTENTION
We ﬁrst review standard multi-head attention introduced by Vaswani et al.
2. A TTENTION
Let X ∈ R T × D in and Y ∈ R T (cid:48)
× D in be two input matrices consisting of respectively T and T (cid:48) tokens of D in dimensions each.
An attention layer maps each of the T query token from D in to D out dimensions as follows:
[Equation]
(cid:18) QK [Equation] d
[Equation]
(cid:19)
V , with Q = XW Q , K = Y W K , V = Y W V (1)
The layer is parametrized by a query matrix W Q ∈ R D in × D k , a key matrix W K ∈ R D in × D k and a value matrix W V ∈ R D in × D out .
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. C ONTENT VS . C ONTEXT
Some re-implementations of the original transformer architecture 2 use biases in the linear layers.
This differs from the attention operator deﬁned in eq.
(1) where the biases b Q and b K ∈ R D k are ommited.
Key and query projections are computed as K = XW K + 1 T × 1 b K and Q =
Y W Q + 1 T × 1 b Q , respectively, where 1 a × b is an all one matrix of dimension a × b .
The exact computation of the (unscaled) attention scores can be decomposed as follows:
[Equation]
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 ﬁrst two terms have a clear meaning: XW Q W (cid:62) K Y (cid:62) considers the relation between keys and query pairs, whereas 1 T × 1 b (cid:62) Q W (cid:62) K Y (cid:62) computes attention solely based on key content.
[Figure] c a
[Figure]
t u r e d v a r i a n
[Figure]
The above ﬁndings suggest that the bias b K of the key layer can be always be disabled without any consequence.
Moreover, the query biases b Q 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-S YNTHESIZER (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 sufﬁce.
2. M ULTI -H EAD A TTENTION
Traditionally, the attention mechanism is replicated by concatenation to obtain multi-head attention deﬁned for N h heads as:
[Equation]
where distinct parameter matrices W ( i ) Q , W ( i ) K ∈ R D in × d k and W ( i ) V ∈ R D in × d out are learned for each head
i ∈
[ N h ] and the extra parameter matrix W O ∈ R N h d out × D out projects the concatenation of the
N h head outputs (each in R d out ) to the output space R D out .
In the multi-head setting, we call d k the dimension of each head and D k
= N h
d k the total dimension of the query/key space.
I MPROVING THE M ULTI -H EAD M ECHANISM
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. H OW 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 sufﬁce to look at the similarity between query (or key) matrices { W ( i ) Q } i ∈
[ N h ] 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 ∈ R
d k × d k such that
[Equation]
Even though the two heads are computing identical attention scores, i.e. W (1) Q RR (cid:62) W (1) (cid:62) K =
W (1) Q W (1) (cid:62) K , they can have orthogonal column-spaces and the concatenation [ W (1) Q , W (2) Q ] ∈
R D in × 2 d k can be full rank.
[Figure] To disregard artiﬁcial differences due to common rotations or scaling of the key/query spaces, we study the similarity of the product W ( i ) Q W ( i ) (cid:62) K ∈ R D in × D in 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 ) (cid:62) 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 sufﬁces to capture almost all the energy of W Q W (cid:62) K , which suggests that there is inefﬁciency in the way multi-head attention currently operate.
3. C OLLABORATIVE M ULTI -H EAD A TTENTION
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 deﬁned as follows:
[Equation]
The main difference with standard multi-head attention deﬁned in eq.
(5) is that we do not duplicate the key and query matrices for each head.
Instead, each head learns a mixing vector m i ∈ R ˜
D k that deﬁnes a custom dot product over the ˜ D k projected dimensions of the shared matrices ˜ W Q and ˜ W K of dimension D in × ˜ 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 efﬁcient 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
= N h d k ).
The left of Figure
[Equation]
laying out the mixing vectors m i as rows.
In the concatenated MHA, the mixing vector m i for the i -th head is a vector with ones aligned with the d k dimensions allocated to the i -th head among the
D k = N h d k total dimensions.
Some alternative collaborative schema can be seen on the right side of Figure 2.
By learning the mixing vectors { m i } i ∈
[ N h ] instead of ﬁxing 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 d k 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.
3. H EAD C OLLABORATION AS T ENSOR D ECOMPOSITION
As we show next, there is a simple way to convert any standard attention layer to collaborative attention without retraining.
To this end, we must extract the common dimensions between query/key matrices { [Equation] i
[Equation]
[ N h ] across the different heads.
This can be solved using the
Tucker tensor decomposition (Tucker, 1966) of the 3rd-order tensor
[Equation]
Following the notation 3 of Kolda & Bader (2009), the Tucker decomposition of a tensor T ∈ R
I × J × K is written as
[Equation]
with A ∈ R I × P , B ∈ R J × Q , and C ∈ R K × R being factor matrices, whereas G ∈ R P × Q × R is the core tensor.
Intuitively, the core entry g pqr =
G p,q,r quantiﬁes the level of interaction between the components a p , b q , and c r .
In the case of attention, it sufﬁces to consider the dot product of the aligned key/query components of the Q and K matrices, which means that the core tensor is super-diagonal (i.e. g pqr (cid:54) =
only if q = r ).
We further simplify the Tucker decomposition by setting the factors dimensions P, Q and R to ˜ D k , a single interpretable hyperparameter equal to the dimension of the shared key/query space that controls the amount of compression of the decomposition into collaborative heads.
These changes lead to a special case of Tucker decomposition called the canonical decomposition, also known as CP or PARAFAC (Harshman, 1970) in the literature (Kolda & Bader, 2009).
Fix any positive rank R .
The decomposition yields:
[Equation]
with A ∈ R I × R , B ∈ R J × R and C ∈
R K × R .
What is remarkable is that the above can be used to express any (trained) attention layer parametrized by { W ( i ) Q , b ( i ) Q , W ( i ) K , b ( i ) K }
i ∈
[ N h ] as a collaborative layer.
In particular, if we apply the decomposition to the stacked heads W QK we obtain the three matrices (cid:74) M , ˜ W Q , ˜ W K (cid:75) that deﬁne a collaborative attention layer: the mixing matrix M ∈
R N h ×
˜ D k , as well as the key and query projection matrices ˜ W Q , ˜ W K ∈ R D in × ˜ D k .
On the other hand, biases can be easily dealt with based on the content/context decomposition of eq.
(3), by storing for each head the vector
[Equation]
With this in place, the computation of the (unscaled) attention score for the i -th head is given by:
[Equation]
where m i is the i -th row of M .
If ˜ D k ≥ D k the decomposition is exact (eq.
(11) is an equality) and our collaborative heads layer can express any concatenation-based attention layer.
We also note that the proposed re-parametrization can be applied to the attention layers of many transformer architectures, such as the ones proposed by Devlin et al.
; Sanh et al.
; Lan et al.
3. P ARAMETER AND C OMPUTATION E FFICIENCY
[Table]
Collaborative MHA introduces weight sharing across the key/query projections and decreases the number of parameters and FLOPS.
While the size of the heads in the standard attention layer is set to d k = 64 and the key/query layers project into a space of dimension D k
= N h d k , the shared key/query dimension
˜
D
k of collaborative MHA can be set freely.
According to our experiments in Section 4 (summarized in Table 1), a good rule of thumb when transforming a trained MHA layer to collaborative is to set ˜
D k to half or one third of D k .
When training from scratch,
˜ D k can even be set to 1 / 4 -th of D k .
Parameters.
Collaborative heads use (2 D in + N h )
˜
D k parameters, as compared to 2 D in D k in the standard case (ignoring biases).
Hence, the compression ratio is ≈ D k / ˜ D k , controlled by the shared key dimension ˜ D k .
The collaborative factorization introduces a new matrix M of dimension
N h × ˜ D k .
Nevertheless, as the number of heads is small compared to the hidden dimension (in
BERT-base N h = 12 whereas D in = 768 ), the extra parameter matrix yields a negligible increase as compared to the size of the query/key/values matrices of dimension D in × D k .
Computational cost.
Our layer decomposes two matrices into three, of modulable dimensions.
To compute the attention scores between T tokens for all the N h heads, collaborative MHA requires 2 T ( D in + N h )
˜
D
k + T 2 N h ˜ D k FLOPS, while the concatenation-based MHA uses 2 TD in D k + T 2 D k
FLOPS.
Assuming that D in (cid:29) N h
= O (1) (as is common in most implementations), we obtain a theoretical speedup of Θ( D k / ˜ D k ) .
However in practice, having two matrix multiplications instead of a larger one makes our implementation slightly slower, if larger multiplications are supported by the hardware.
E XPERIMENTS
The goal of our experimental section is two-fold.
First, we show that concatenation-based MHA is a drop-in replacement for collaborative MHA in transformer architectures.
We obtain a signiﬁcant reduction in the number of parameters and number of FLOPS without sacriﬁcing performance on a
Neural Machine Translation (NMT) task with an encoder-decoder transformer.
Secondly, we verify that our tensor decomposition allows one to reparametrize pre-trained transformers, such as BERT (Devlin et al., 2019) and its variants.
To this end, we show that collaborative MHA performs on par with its concatenation-based counter-part on the GLUE benchmark (Wang et al., 2018) for Natural
Language Understanding (NLU) tasks, even without retraining.
The NMT experiments are based on the FairSeq (Ott et al., 2019) implementation of transformer-base by Vaswani et al.
For the NLU experiments, we implemented the collaborative MHA layer as an extension of the Transformers library (Wolf et al., 2019).
The ﬂexibility of our layer allows it to be applied to most of the existing transformer architectures, either at pre-training or after ﬁne-tuning using tensor decomposition.
We use the tensor decomposition library Tensorly (Kossaiﬁ et al., 2019) with the PyTorch backend (Paszke et al., 2017) to reparameterize pre-trained attention layers.
Our code and datasets are publicly available
and all hyperparameters are speciﬁed in the Appendix.
[Figure]
4. C OLLABORATIVE MHA FOR N EURAL M ACHINE T RANSLATION
We replace the concatenation-based MHA layers of an encoder-decoder transformer by our collaborative MHA and evaluate it on the WMT14 English-to-German translation task.
Following (Vaswani et al., 2017), we train on the WMT16 train corpus, apply checkpoint averaging and report compound split tokenized BLEU.
We use the same hyperparameters as the baseline for all our runs.
Results are shown in Figure 3.
Our run of the original base transformer with N h = 8 heads and D k =key/query total dimensions achieves 27.40 BLUE (instead of 27.30).
As observed in the original paper by Vaswani et al.
(2017), decreasing the key/query head size d k degrades the performance ( × in Figure 3).
However, with collaborative heads (+ in Figure 3), the shared key/query dimension can be reduced by 4 × without decreasing the BLEU score.
As feed-forward layers and embeddings are left untouched, this translates to a 10% decrease in number of parameters for a slight increase in training time.
When setting a total key/query dimension of
D k = 64 , corresponding to d k = 8 dimensions per head, the classic MHA model suffers a drop of 0.6 BLEU points, meanwhile the collaborative MHA stays within 0.1 point of the transformer-base model using concatenation.
We conclude that sharing key/query projections across heads allows attention features to be learned and stored only once.
This weight sharing enables decreasing D k without sacriﬁcing expressiveness.
4. R E - PARAMETRIZE A P RE - TRAINED MHA INTO C OLLABORATIVE MHA
[Figure]
We turn to experiments on Natural Language Understanding (NLU) tasks, where transformers have been decisive in improving the state-of-the-art.
As pre-training on large text corpora remains an expensive task, we leverage the post-hoc reparametrization introduced in Section 3.3 to cast already pretrained models into their collaborative form.
We proceed in 3 steps for each GLUE task (Wang et al., 2018).
First, we take a pre-trained transformer and ﬁne-tune it on each task individually.
Secondly, we replace all the attention layers by our collaborative MHA using tensor decomposition to compute ˜ W Q , ˜ W K and M and re-parametrize the biases into v .
This step only takes a few minutes as shown in Figure 4.
Finally, we ﬁne-tune the compressed model again and evaluate its performance.
We experiment with a pre-trained BERT-base model (Devlin et al., 2019).
We also repurpose two variants of BERT designed to be more parameter efﬁcient: ALBERT (Lan et al., 2020), an improved transformer with a single layer unrolled, and DistilBERT (Sanh et al., 2019) a smaller version of
BERT trained with distillation.
We report in Table 2 the median performance of 3 independent runs of the models on the GLUE benchmark (Wang et al., 2018).
We ﬁrst verify that tensor decomposition without compression ( ˜ D k
= D k
= 768 ) does not alter performance. As shown in Table 2, both BERT-base and its decomposition performs similarly with an average score of 83.0% and 83.2% respectively.
We should clarify that, for consistency, we opted to
[Table]
re-ﬁnetune the model in all cases (even when ˜ D
k = D k ), and that the slight score variation disappears without re-ﬁnetuning.
Nevertheless, even with re-ﬁnetuning, reparametrizing the attention layers into collaborative form is beneﬁcial in 4 out of the 8 tasks, as well as in terms of the average score.
We then experiment with compressed decomposition using a smaller ˜ D k .
Comparing the original models with their well-performing compressed counterpart (gray rows) shows that the key/query dimension of BERT and DistilBERT can be reduced by 2 × and 3 × respectively without sacriﬁcing more than 1.5% of performance.
This is especially remarkable given that DistilBERT was designed to be a parameter-efﬁcient version of BERT.
It seems that ALBERT suffers more from compression, but the dimension can be reduced by a factor 1.5 × with minor performance degradation.
We suspect that unrolling the same attention layer over the depth of the transformer forces the heads to use different projections and decreases their overlap, decreasing the opportunity for weight-sharing.
Our hypothesis is that better performance may be obtained by pre-training the whole BERT architecture variants from scratch.
Recovering from compression with ﬁne-tuning.
We further investigate the necessity of the second ﬁne-tuning—step 3 of our experimental protocol—after the model compression.
Figure 5 shows the performance of BERT-base on 3 GLUE tasks for different compression parameters ˜ D
k with and without the second ﬁne-tuning.
We ﬁnd that for compression up to 1.5 × (from D k = 768 to ˜
D k = 512 ), the re-parametrization is accurate and performance is maintained without ﬁne-tuning again.
Further compressing the model starts to affect performance.
Nevertheless, for compression by up to 3 × (to ˜ D k = 256 ), this loss can readily be recovered by a second ﬁne-tuning ( in orange ).
[Figure]
Repeating ﬁne-tuning after compression can make the model recover the original performance when compression was drastic.
The
GLUE baseline gives a reference for catastrophic failure.
C ONCLUSION
This work showed that trained concatenated heads in multi-head attention models can extract redundant query/key representations.
To mitigate this issue, we propose to replace concatenation-based
MHA by collaborative MHA.
When our layer is used as a replacement for standard MHA in encoder/decoder transformers for Neural Machine Translation, it enables the decrease of effective individual head size from d k = 64 to 8 without impacting performance.
Further, without pre-training from scratch, switching a MHA layer to collaborative halves the number of FLOPS and parameters needed to compute the attentions score affecting the GLUE score by less than 1.5%.
Our model can impact every transformer architecture and our code (publicly available) provides post-hoc compression of already trained networks.
We believe that using collaborative MHA in models pre-trained from scratch could force heads to extract meaningful shared query/key features.
We are curious if this would translate to faster pre-training, better performance on downstream tasks and improved interpretability of the attention mechanism.
Supplementary Material A H YPERPARAMETERS FOR N EURAL M ACHINE T RANSLATION E XPERIMENTS
Our implementation is based on Fairseq implementation Ott et al.
We report in the following tables the speciﬁcation of the architecture.
We used the default hyperparameters if they are not speciﬁed below.
[Table]
B H YPERPARAMETERS FOR N ATURAL L ANGUAGE U NDERSTANDING E XPERIMENTS
We use standard models downloadable from HuggingFace repository along with their conﬁguration.
Models
BERT-base
Devlin et al.
(2019) bert-base-cased
DistilBERT
Sanh et al.
(2019) distilbert-base-cased
ALBERT
Lan et al.
(2020) albert-base-v
We use HuggingFace default hyperparameters for GLUE ﬁne-tuning in all our runs.
We train with a learning rate of 2 · 10 − 5 for 3 epochs for all datasets except SST-2 and RTE where we train for 10 epochs.
In preliminary experiments, we tried to tune the tensor decomposition tolerance hyperparameter among
[Table]