Year: 2020

Abstract:
Attention layers are widely used in natural language processing (NLP) and are beginning to influence computer vision architectures. Training very large transformer models allowed significan improvement in both fields, 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 sacrificing accuracy. Our code is public.1

heading: 1 INTRODUCTION
text: 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.

heading: 2 MULTI-HEAD ATTENTION
text: We first review standard multi-head attention introduced by Vaswani et al. (2017).

heading: 2.1 ATTENTION
text: 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.

heading: 2.2 CONTENT VS. CONTEXT
text: 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.

heading: 2.3 MULTI-HEAD ATTENTION
text: 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.

heading: 3 IMPROVING THE MULTI-HEAD MECHANISM
text: 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.

heading: 3.1 HOW MUCH DO HEADS HAVE IN COMMON?
text: 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.

heading: 3.2 COLLABORATIVE MULTI-HEAD ATTENTION
text: 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.

heading: 3.3 HEAD COLLABORATION AS TENSOR DECOMPOSITION
text: 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 {W (i)Q W (i)> K ∈ RDin×Din}i∈[Nh] across the different heads. This can be solved using the Tucker tensor decomposition (Tucker, 1966) of the 3rd-order tensor
WQK := stack i∈[Nh]
[ W
(i) Q W (i)> K
] ∈ RNh×Din×Din . (9)
Following the notation3 of Kolda & Bader (2009), the Tucker decomposition of a tensor T ∈ RI×J×K is written as
T ≈ G×1 A×2 B ×3 C = P∑
p=1 Q∑ q=1 R∑ r=1 gpqr ap ◦ bq ◦ cr =: q G;A,B,C y , (10)
with A ∈ RI×P , B ∈ RJ×Q, and C ∈ RK×R being factor matrices, whereas G ∈ RP×Q×R is the core tensor. Intuitively, the core entry gpqr = Gp,q,r quantifies the level of interaction between the components ap, bq , and cr.
In the case of attention, it suffices 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. gpqr 6= 0 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:
T ≈ R∑
r=1
ar ◦ br ◦ cr =: q A,B,C y , (11)
with A ∈ RI×R, B ∈ RJ×R and C ∈ RK×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∈[Nh] as a collaborative layer. In particular, if we apply the decomposition to the stacked heads WQK we obtain the three matrices JM , W̃Q, W̃KK that define a collaborative attention layer: the mixing matrix M ∈ RNh×D̃k , as well as the key and query projection matrices W̃Q, W̃K ∈ RDin×D̃k .
3◦ represents the vector outer product
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
vi = W (i) K b (i) Q ∈ R Din . (12)
With this in place, the computation of the (unscaled) attention score for the i-th head is given by:( XW
(i) Q + 1T×1b > Q
)( Y W
(i) K + 1T×1b > K )> ≈XW̃Q diag(mi)W̃>K Y > + 1T×1v>i Y >, (13)
where mi is the i-th row of M . If D̃k ≥ Dk 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. (2019); Sanh et al. (2019); Lan et al. (2020).
3.4 PARAMETER AND COMPUTATION EFFICIENCY Table 1: Comparison of a layer of concatenate vs. collaborative MHA with chosen D̃k to give negligible performance difference. T=128.
train re-param. FairSeq §4.1 HuggingFace §4.2
concat. collab. concat. collab.
Dk → D̃k 512→ 128 768→ 256 Params (×106) 1.05 0.66 2.36 1.58 FLOPS (×108) 1.51 1.09 3.27 2.65 inference (ms) 0.99 0.81 1.71 1.65 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 dk = 64 and the key/query layers project into a space of dimension Dk = Nhdk, 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 Dk. When training from scratch, D̃k can even be set to 1/4-th of Dk.
Parameters. Collaborative heads use (2Din + Nh)D̃k parameters, as compared to 2DinDk in the standard case (ignoring biases). Hence, the compression ratio is ≈ Dk/D̃k, controlled by the shared key dimension D̃k. The collaborative factorization introduces a new matrix M of dimension Nh × D̃k. Nevertheless, as the number of heads is small compared to the hidden dimension (in BERT-base Nh = 12 whereas Din = 768), the extra parameter matrix yields a negligible increase as compared to the size of the query/key/values matrices of dimension Din ×Dk.
Computational cost. Our layer decomposes two matrices into three, of modulable dimensions. To compute the attention scores between T tokens for all the Nh heads, collaborative MHA requires 2T (Din+Nh)D̃k +T
2NhD̃k FLOPS, while the concatenation-based MHA uses 2TDinDk +T 2Dk FLOPS. Assuming that Din Nh = O(1) (as is common in most implementations), we obtain a theoretical speedup of Θ(Dk/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.

heading: 4 EXPERIMENTS
text: 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 significant reduction in the number of parameters and number of FLOPS without sacrificing 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. (2017). For the NLU experiments, we implemented the collaborative MHA layer as an extension of the Transformers library (Wolf et al., 2019). The flexibility of our layer allows it to be applied to most of the existing transformer architectures, either at pre-training or after fine-tuning
BLEU params (×106) time (h) Dk concat. collab. concat. collab. concat. collab.

heading: 512 27.40 27.58 60.9 61.0 18.0 21.0

heading: 256 27.10 27.41 56.2 56.2 17.3 19.0

heading: 128 26.89 27.40 53.8 53.8 17.3 18.4

heading: 64 26.77 27.31 52.6 52.7 16.9 17.9
text: 54 56 58 60 Number of parameters ×106
26.8
27.0
27.2
27.4
27.6
BL EU
256
512
64 128 512
128
256
64 collab. concat.
17 18 19 20 21 Training time on 4xV100 (hours)
26.8
27.0
27.2
27.4
27.6
BL EU
256
512
64 128
512
128
256
64 collab. concat.
Figure 3: Comparison of the BLEU score on WMT14 EN-DE translation task for an encoder-decoder transformer-base (Vaswani et al., 2017) using collaborate vs. concatenate heads with key/query dimension Dk. We visualize performence as a function of number of parameters (middle) and training time (right). Collaborative attention consistently improves BLEU score, Dk can be decreased by a factor of 4 without drop in performance.
using tensor decomposition. We use the tensor decomposition library Tensorly (Kossaifi et al., 2019) with the PyTorch backend (Paszke et al., 2017) to reparameterize pre-trained attention layers. Our code and datasets are publicly available4 and all hyperparameters are specified in the Appendix.

heading: 4.1 COLLABORATIVE MHA FOR NEURAL MACHINE TRANSLATION
text: 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 Nh = 8 heads and Dk = 512 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 dk 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 Dk = 64, corresponding to dk = 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 Dk without sacrificing expressiveness.

heading: 4.2 RE-PARAMETRIZE A PRE-TRAINED MHA INTO COLLABORATIVE MHA
text: 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 fine-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 fine-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 efficient: ALBERT (Lan et al., 2020), an improved
4https://github.com/...
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 first verify that tensor decomposition without compression (D̃k = Dk = 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 re-finetune the model in all cases (even when D̃k = Dk), and that the slight score variation disappears without re-finetuning. Nevertheless, even with re-finetuning, reparametrizing the attention layers into collaborative form is beneficial 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 sacrificing more than 1.5% of performance. This is especially remarkable given that DistilBERT was designed to be a parameter-efficient 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.
decomposed with collaborative heads of compressed dimension D̃k (horizontal axis). Repeating fine-tuning after compression can make the model recover the original performance when compression was drastic. The GLUE baseline gives a reference for catastrophic failure.
Recovering from compression with fine-tuning. We further investigate the necessity of the second fine-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 fine-tuning. We find that for compression up to 1.5× (from Dk = 768 to D̃k = 512), the re-parametrization is accurate and performance is maintained without fine-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 fine-tuning (in orange).

heading: 5 CONCLUSION
text: 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 dk = 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.

heading: A HYPERPARAMETERS FOR NEURAL MACHINE TRANSLATION EXPERIMENTS
text: Our implementation is based on Fairseq implementation Ott et al. (2019). We report in the following tables the specification of the architecture. We used the default hyperparameters if they are not specified below.

heading: B HYPERPARAMETERS FOR NATURAL LANGUAGE UNDERSTANDING EXPERIMENTS
text: We use standard models downloadable from HuggingFace repository along with their configuration.
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-v2
We use HuggingFace default hyperparameters for GLUE fine-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 {10−6, 10−7, 10−8} but did not see significant improvement and kept the default 10−6 for all our experiments.
GLUE fine-tuning hyperparameters
Number of epochs 3 for all tasks but 10 for SST-2 and RTE Batch size 32 Learning rate 2e-5 Adam 1e-8 Max gradient norm 1 Weight decay 0 Decomposition tolerance 1e-6

Review summaries:
This paper analyzes the multi-head attention in transformers and suggests to use collaboration instead of concatenation of multiple heads. Empirical results on WMT’16 English-German demonstrates that the proposed approach reduces the of parameters without sacrificing performance. Further experiments on pre-trained BERT models also demonstrate its efficiency.  Overall, the paper is well motivated and provides a deep analysis of redundancy of the multi-head attention. 

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.