Concise Multi-head Attention Models
Authors: Srinadh Bhojanapalli
Chulhee Yun
Ankit Singh Rawat
Sashank Reddi
Sanjiv Kumar
Venue:  (2020)
Keywords: Transformers
Attention
Multihead
expressive power
embedding size

Abstract: Attention based Transformer architecture has enabled significant advances in the field of natural language processing. In addition to new pre-training techniques, recent improvements crucially rely on working with a relatively larger embedding dimension for tokens. This leads to models that are prohibitively large to be employed in the downstream tasks. In this paper we identify one of the important factors contributing to the large embedding size requirement. In particular, our analysis highlights that the scaling between the number of heads and the size of each head in the existing architectures gives rise to this limitation, which we further validate with our  experiments. As a solution, we propose a new way to set the projection size in attention heads that allows us to train models with a relatively smaller embedding dimension, without sacrificing the performance.

Under review as a conference paper at ICLR 2020
CONCISE MULTI-HEAD ATTENTION MODELS
Anonymous authors Paper under double-blind review
ABSTRACT
Attention based Transformer architecture has enabled significant advances in the field of natural language processing. addition to new pre-training techniques, recent improvements crucially rely on working with relatively larger embedding dimension for tokens. Unfortunately, this leads to models that are prohibitively large to be employed in the downstream tasks. In this paper identify one of the important factors contributing to the large embedding size requirement. In particular, our analysis highlights that the scaling between the number of heads and the size of each head in the existing architectures gives rise to this limitation, which we further validate with our experiments. As a solution, we propose a new way to set the projection size in attention heads that allows us to train models with a relatively smaller embedding dimension, without sacrificing the performance.
INTRODUCTION
Attention based architectures, such as Transformers, have been effective for sequence modelling tasks such as machine translation (Gehring al. 2017 Vaswani al. 2017), question answering, sentence classification (Radford et al. [2118] Devlin et al. 2018)) and document generation (Liu et al. 2018)). These models have emerged as better alternatives 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:
2017) and speech recognition (Chiu al. 2018
Recent advances in using the self attention models in natural language tasks have been made by first using 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 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). 24 layer transformer model, 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 larger embedding size. BERT model uses embedding size of 1024 (Devlin et al. 2018); GPT-2 uses models with embedding size up to 1600 (Radford et al. 2019).
single Transformer block consists of two key components: mult 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 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/(umber 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 partieular, focus the representation power of the multi-head self attention layer.

---

Under review as a conference paper at ICLR 2020
# heads
16
32
# params
336M
336M
336M
SQuAD FI
90.8910.15
90.61=0.14
90.45+0.08
SQuAD - EM MNLI
84.1+:034
83.75.10.27
83.4
8520.2
84.5=0.4
84.410.2
Table 1: Performance of BERTLARGE (Devlin et al. 2018 3. 24 layer Transformer with an embedding size of 1024, does not improve with the increasing number of heads after 8 heads.
First, we analyze the representation power of a single self attention unit and show that it crucially depends the projection sizes used to compute the dot product attention.
We next study the advantage 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). notice a drop in the performance once the number of heads increases beyond a certain threshold, as seen in Table and Fig. |
(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..._2013 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, 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 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 factor of the embedding size, we are free to set arbitrary number of heads as required for the task.
propose
way to set the projection size in the attention heads,
We evaluate this fixed head size Transformer on language modeling (LMIB dataset). natural lan- guage 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 LERTLARGE (Devlin et al., 2018). a Transformer with an embedding size of 1024 (see Fig. (). We further present experimental results evaluating the effect of different choices of the head size and the embedding size 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 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 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 three standard NLP tasks.
2

---

Under review as a conference paper at ICLR 2020
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, [Ot al. (2018) and You et al. (2019) reduce precision and use large batch training to reduce the training time of the attention models.
Child al. (2119) propose sparse self attention models 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." (018)) propose 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 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 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 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 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. (209)) impose sparsity structure on the attention layer during training to improve both interpretability and performance. Fixing the head size will in fact make it casier to learn such sparsity patterns, low rank constraint does not allow head to express all possible sparsity patterns. Combining these techniques can hence potentially enable training of sparse attention models with smaller embedding size.
TRANSFORMER ARCHITECTURE AND ANALYSIS
In this section 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 Transformer network is sequence of tokens. Typically, each token is converted into a token embedding of dimension d by an embedding layer. We let X E [Rdxn be the embedding matrix corresponding to the n tokens in the input sequence.
2.1
SINGLE-HEES HEAD ATTENTION
The transformer block is combination of a self attention layer followed by feed forward layer (Vaswani et al. 2017). Both layers have a skip connection and use Layer Normalization (LN) (Ba) al. 2016). In particular, for token embeddings X, the dot product attention is computed as follows.
(W&X)"(W,X) ddg
Attention X)= W.,X Softmax
W.X.P.
Here Wg E Rd WA E 4 and W. E Rd represent the projection matrices associated with the query, key and value respectively in an attention unit (Vaswani et al., (0177). For single-heat -ead attention unit, we have d, = d₂ = d, = In the dot-produee attention (cf. (I).. P aims to capture the context of the input for given token based the remaining tokens in the input sequence.
3

---

Under review as a conference paper at ICLR 2020
Subsequently, the output of the attention layer takes the following form.
LN (X + W. Attention(X)),
(2)
where LN(-) represents the layer- aayer-normalization operation. Given the attention module, as defined in ①T. it is natural to question its ability to represent arbitrary contexts P for a given input sequence X.
Towards this, we show that for 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 cannot represent arbitrary context when d is smaller than n.
Theorem (Representation Theorem). df = d2 2 n, then given any full column rank matrix X E Rdxn and an arbitrary n X n positive column stochastic matrix P, there always exists d projection matrices Wg q and W such that
(WxX)" =
Softmax
(3)
Vdk If d d; - < there exist X and P such that (3) does not hold for all W and Wt-
The proof of Theorem D is provided in the supplementary material. This result shows that the projection dimension d, dy = 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.
2.2 MULTI-HEAD ATTENTION
As discussed Section an attention unit updates the embedding of an input token based on a weighted average of the embeddings of all the tokens in the sequence, using the context P (cf. (I). Vaswani et al. 2017) proposed Multi-Head attention mechanism that increases the representation power of an attention layer, where multiple attention units operate on different low dimensional projections of the input, with each attention unit being referred to as a head. This is followed by concatenation of the outputs from different heads. In particular, the computation inside a Multi-Head attention with h heads takes the following form:
head X): = W Softmax (W:X) E E
MultiHead X) (oncat(head(()):
head X)n] Rdxn
The output of the Multi-head attention layer then becomes
Z LN (X + W. (ul))
(4)
where W. E Rdxd For model with h heads, the query, key and value projection matrices { (Wi) and { [ are x d matrices. Therefore, each head projects the input onto a d subspace to compute the context, and keeps the number of parameters fixed per layer. This has been observed empirically to increase the expressive power of the attention layer (Vaswani et al. 2017)
- 4ddimensional
2.3 DEPENDENCE OF THE NUMBER OF HEADS ON THE EMBEDDING SIZE
While increasing the number of heads seemingly gives the model more expressive power, at the same time we are reducing the head size, which can decrease the expressive power. When the number of heads is larger than á the attention unit inside each head projects onto dimension smaller than n, and looses its ability to represent arbitrary context vectors (cf. Theorem I Since the sequence length fixed from the data/task hand, the only remaining way increase the number of heads, without loosing the expressiveness, is by increasing the embedding size d. This corresponds to a fundamental limitation of the model architecture that one needs to increase the embedding size in order to support more heads.
Unfortunately, increasing the embedding size leads to higher computation and memory requirements to train and store the model. Further, since it is common to use learned embeddings from Transformer based models for downstream tasks (Devlin et al. 2018) larger embedding size increases the model size and computation required for all the downstream tasks well. Given the widespread success of attention mechanism, this highlights the need for a modified model architecture that can leverage the advantages of MultiHead without suffering from the requirement of large embedding sizes.
4

---

Under review as a conference paper at ICLR 2020
3 CONCISE MULTI-HEAD TRANSFORMER
In this section we propose a different way of setting the head size the Transformer, which allows us enjoy the advantage of higher expressive power of multiple heads without requiring the embedding size to be large. The key is to decouple the dependency between the projection size in a head and the embedding size of the model. The projection matrices now project onto subspaces of fixed dimension dp irrespective of the number of heads h. This approach where dp is independent of d and leads to the following attention mechanism.
fixedhead X) = V Softmax
E R24xx
FixedMulttiHea(()) Concat (iieehefixedeead X) 1),,' fixedhead Rd
Note that the projection matrices used here {V+), {Vi+ and {V+} are dp x d matrices. With Vo Rdxhed, the output of this new multi-head attention layer takes the following form.
z LN (X V. FixeMMuliHHadd())
This modification makes each attention head more similar to a hidden unit in feed forward network or filter in a convolutional network, and allows us to vary the number of heads without the worry of reducing the representation power per head. The downside is, unlike the standard MultiHead, the number of parameters per layer increase with the number of heads. However, this modification allows us to train a model with smaller embedding size, ultimately allowing us to reduce the total number of parameters in the model.
Choice of the head size. Our proposed modification introduces head size d, as a new model hyper- parameter. We choose head size to be 128 for our BERT experiments, most of the pre-training is done with 128 sequence length data. While we have ablation studies (ef. Table[[(B) showing bigger head size improves the performance, there is tradeoff between increasing the head size number of heads vs layers. We found that having sufficient head size matching the pre-training sequence length, is better than having larger embedding size (ef. Section
3.1 MULTIHEAD VS. FIXEDMULTHHEAD ATTENTION
Given MultiHead layer, we can always represent it using a FixedMultHeaad layer, whenever we have the head size dp 2 4/h. While this shows that increasing the number of heads h beyond d/d, makes individual heads of the FixedMultiHeal as expressive as the ones in the MultiHead, it is not clear FixedMultiHead is strictly more expressive. Can the FixedMultiHead layer represent functions that the standard MultiHead layer can not represent? In this subsection we show that indeed, in the multi-head regime, the FixedMultiHaad layer is strictly better than the standard MultiHead layer in terms of expressive power.
Consider the standard multi-head attention units
Jw(X) = W. MultiHead(X)..
We denote the collection of all parameter matrices as W. Similarly, consider the function represented by the fixed head size attention units:
gv = Vo FixedMuliHHaad(())
Let V be the collection of all these parameter matrices. We define F and G to be the class of functions fw(-) and gv(-), respectively. As noted above, dp 2/h,, we have F C G.
The following theorem shows that even for simple examples in G, functions in F fail to approximate them beyond certain accuracy; this shows that F is strict subset of G.
Theorem 2. Given n 2 2 and d 2 dp. let h > 4/d,. Consider a fixed head size attention layer gv with parameters that satisfy the following conditions:
gv(+) (-)
V is full rank, and (Vi) = for all I,...,,,, where U is a rank-d, matrix.
Vh
Then, for any Sw E F. there exists X E Rdxn such that fw (X) & (X).
5

---

Under review as a conference paper at ICLR 2020
Because fl gv ox]]] is continuous function of X, existence of such an X implies that the integral the norm of difference (i.e., approximation error) strictly positive.
This theorem shows that the expressive power of the FixedMultiHead attention function class is strictly superior to the standard MultiHead attention function class. Hence the heuristic of reducing the head size with the number of heads is limiting the expressive power of MultiHead, whereas using the fixed head size Transformers will increase the expressive power of the attention layers.
4 EXPERIMENTS
41
d/h,
d d/h,
h 16
-
16
h
Perplexity
Perplexity
d 266, de
32
d= 266, 32
37
36
36
35
35
34-
15
22
20
30
250
300
350
400
430
300
Parameters millions
Embedding size (d)
(a) LMIB
(b) LMIB
Figure 1: Performance of the standard Transformer (d, d/h) compared with the fixed head size (d, = 32) models on a language modeling task (LMIB) on the test set. We vary the embedding size of the standard Transformer from 256 512. We train the fixed head size models with fixed embedding size of 256 and head size of 32, and vary the number of heads from 4 to 70, while matching the number of parameters. The plots clearly indicate that fixing the head size allows us to train Transformers with a smaller embedding size (plot (b)), and with better scaling of performance (plot (a)). Note that for perplexity lower values are better.
hn:00
MNLI accuracy
EM
e-1024
14.9
0.3
44432
0.5
BRRTTaaaa d, d/b.b- 18
BRRTesnme dy h/
019
BRRTT d, d/b.b
FisslBeed d-- 412.4,
FlodBHed
FlolBBed d-
163
1:3
Parameters millions
Parameters millions
Parameters millions
(a) FQ
(b) SQUAD FM
(c) MILI
24 layer standard Transformer model BERTLARGE with the fixed head size model on the SQuAD and MNLI dev sets. We vary the embedding size of the BERT models from 512 to 1024. We train the fixed head size models with a fixed embedding size 512 and head size of 128, with varying number of heads from 32, while matching the number of parameters. Fixing the head size allows us to train models with an embedding size 512 with better performance.
Figure 2: Comparison
In this section we present our experiments three standard NLP tasks, language modeling (LMIB). question answering (SQuAD), and sentence entailment (MNLI), to demonstrate: 1) Increasing the number of heads beyond certain point hurts the performance of the standard Transformer, but always helps with our proposed modification; 2) Decoupling the head size from embedding size allows us to train models with smaller embedding size; and 3) Setting the head size appropriately in the Transformers allows us to train models with a better performance scaling. We first describe our experimental setup followed by our results and ablation studies on the proposed modifications.

---

Under review as a conference paper at ICLR 2020
46
dp
d/h.
h
dp
Fex
h
16
40
44-
16
39-
32
42-
Perplexity
40-
37
30-
88
36-
34-
33-
2
2
2
2040 Width of the Feedforward layer
Head size (dg)
(b)
Figure 3: Ablation studies on LMIB: (a) We fix the embedding size of all the models 256 and vary the capacity of the standard Transformers by increasing the size of the feedforward layers. For the modified models we the head size to 32, so head modified model is the same as the head standard Transformer. We notice that again in Transformers increasing the number of heads beyond 16 hurts the performance, whereas with a fixed head size increasing the number of heads monotonically improves the performance. (b) We show the effect of head size the performance with different number of heads. Both plots clearly show the advantage in having an additional way to tune the capacity of Transformers with fixed embedding size.
4.1 SETUP AND DATASETS
For the language modeling task we use the one billion word benchmark dataset (LMIB) (Chelba et al. 2013). This dataset has around 30M training examples and around 300k examples in the test set. We use sub-word tokenizer with 32k vocab and cap the input to 256 sequence length. We train 6 layer Transformer model with the ADAM optimizer using the tensor2tensor library (Vaswani) ct al. 2018). The detailed experimental setting is presented in Section
Menti-Genre Natural Language Inference (MNLI) is sentence level entailment task, designed to test natural language understanding [Williams et al. 2018). Given a premise sentence and a hypothesis sentence, the goal is to predict whether hypothesis entails, contradicts is neutral to the premise. We report the classification accuracy for this task. Stanford Question Answering Dataset (SQuAD) is question answering dataset, where given paragraph and question, the goal is predict the sequence of words in the paragraph that constitute the answer to the question (Rajpurkar et al.) 2016). This is harder word level task, compared the sentence classification task. We report both Exact Match (EM) and Fl scores for this task. All results in this section are reported on the Dev set, which has not been used in any experimental choices in this paper.
For these latter two tasks, we follow the two stage approach of first pre-training on a language modeling task and then fine-tuning the models on the task data. We follow the same experimental setup for both pre-train vre-training and fine-tuning as BERT (Dellnn et al. 2018), and use their codebase"] We first pre-train our models using the masked language model and the next sentence prediction objectives, and then fine tune the pre-trained model for individual tasks (Devlin et al.) (018). For pre-training we use English Wikipedia and BooksCorpus dataset (Zhu et al. 2015). The input to the models is tokenized using the WordPiece representation with 30000 tokens in the vocabulary. We present the key experiment choices in Section C. and refer the reader to Devlin et al. (2018) for complete description of the setup.
4.2 RESULTS
For our first set of experiments we want to see if the fixed head size Transformer with a smaller embedding size can match the performance of standard Transformers with larger embedding size. As a baseline for the language modeling task, we train Transformers with the embedding size increasing from 256 to 512 with different number of heads. We train the fixed head size Transformers with
7

---

Under review as a conference paper at ICLR 2020
heads
12
16
32
# params
168M
193M
218M
319M
SQuAD FI
89.6:1.17
90.25=1.221
90.43=0.14
90.95:10.14
SQuAD EM MNLI
82.73=0.21
83.18:0.24
83,59=0.06
84.4:10.29
83.5
84.2.10.2
83.9=002
84,9
(A) Increasing number of heads
head size
32
64
128
256
# params
130M
142M
168M
218M
SQuAD F1
88.53=0.006
89.51=0.15
89.6=t..7
90.33+0.23
SQuAD EM MNLI
81.19=0.21
82.41=0.32
82.73=0.21
83.3610.48
82.5
83.4+0.3
83.5=0.2
83.9
(B) Increasing head size
Table (A): 24 layer modified Transformer with a fixed head size of 128 and 512 embedding size shows an improvement in the accuracy with the increasing number of heads. (B) The fixed head size model with 512 embedding size and heads shows an improvement in accuracy with the increasing head size. This shows that indeed head size is an important capacity controlling parameter in the self attention architecture.
fixed embedding size of 256 and head size of 32, with an increasing number of heads from 4 to 70. We notice that the fixed head size models with an embedding size of 256 can match the performance of standard Transformers with an embedding size of 512 (see Fig. [T]. Further this provides better performance scaling than the standard Transformers. We repeat the similar experiment two tasks, where for baseline we train BERTLARGE a 24 layer, 16 head Transformer architecture, with embedding sizes from 512 to 1024. We compare it with the modified model, with an embedding size 512 and a head size 128, with an increasing number heads from to 32. We again notice that the fixed head size model with 512 embedding size can match the performance of BERTLARGS (see Fig. 2.
the other
Note that simply trying to increase the head size of the standard Transformers by decreasing the number of heads, does not improve the performance, as decreasing the number of heads reduces the expressive power of the model (see Fig. 44. Hence, both the head size and the number of heads needs to be set high enough for better performance.
4.3 ABLATION
Increasing heads, From Table and and Fig. La we can see that increasing the number of heads hurts the performance of the Transformer after certain number. We repeat the same experiments with the fixed head size Transformer, and present the results in Table A) and Fig. [3a The results show that the performance of the modified model improves monotonically the number of heads increase. This is because the model capacity (a function of the head size) is no longer reduced with the increasing number of heads.
Increasing head size. In Table ((B) and Fig. 3b we present comparisons between models with different head sizes. This shows that the gains the performance of the fixed head size models indeed come from adjusting the head size of the query, key and value layers in the attention unit. The table shows clear trend of better performance with a larger head size, suggesting that indeed an important factor in the performance of the attention models.
CONCLUSION
In this paper we studied the representation power of the multi-head self attention models and showed that the larger embedding size used in the current models is a consequence of the limitations of the current multi-head attention formulation. We propose a modified way to set the head size that allows us increase the number of heads without increasing the embedding size. As consequence
8

---

Under review as a conference paper at ICLR 2020
able to train Transformers with smaller embedding size and fewer parameters, without hurting the performance. In the future, it will be interesting to experiment with varying head sizes within an attention block and across layers. This requires further understanding of the role of each layer in computing the context, which an interesting direction for the future work.
REFERENCES
Jimmy Lei Ba, Jamie Ryan Kiros, and Geoffrey E Hinton. Layer normalization. arXiv preprint arXiv: (oo::660005555 2016.
Ciprian Chelba, Tomas Mikolov, Mike Schuster, Qi Ge, Thorsten Brants, and Phillipp Koehn. One billion word benchmark for measuring progress in statistical language modeling. CoRR, abs/1312.3005,, 2013. URL http:://aaxx...r//aa////////////
3005
Rewon Child, Scott Gray, Alec Radford, and Ilya Sutskever. Generating long sequences with sparse transformers. arXiv preprint arXiv: 10509, 2019.
Chung-Cheng Chiu, Tara N Sainath, Yonghui Wu, Rohit Prabhavalkar, Patrick Nguyen, Zhifeng Chen, Anjuli Kannan, Ron Weiss, Kanishka Rao, Ekaterina Gonina, al. State-of-the-an speech recognition with requence-to-sequenne models. In 2018 IEEE International Conference on Acoustics, Speech and Signal Processing (ICASSP), pp. 4774-4778. IEEE, 2018.
Kyunghyun Cho, Bart van Merrienboer, Dzmitry Bahdanau, and Yoshua Bengio. On the properties of neural machine translation: Encoder-decoder approaches. In Proceedings of SSST-8, Eighth Workshop on Syntax, Semantics and Structure Statistical Translation, pp. 111-111, 2014.
Gonçalo M Correia, Vlad Niculae, and André FT Martins. Adaptively sparse transformers. In Proceedings of the 2019 Conference on Empirical Methods in Natural Language Processing and the 9th International Joint Conference on Natural Language Processing (EMNLP-IJCNLP),, pp- 2174-216 2184, 2019.
Mostafa Dehghani, Stephan Gouws, Oriol Vinyals, Jakob Uszkoreit, and Eukasz Kaiser. Universal transformers. arXiv preprint 1887vv1807.038111 2018
Jacob Devlin, Ming-Was Chang, Kenton Lee, and Kristina Toutanova. Bert: Pre-training of deep bidirectional transformers for language understanding. arXiv preprint arXiv: 1rxiv:1810.04805, 2018.
Jonas Gehring, Michael Auli, David Grangier, Denis Yarats, and Yann N Dauphin. Convolutional sequence to sequence learning. In Proceedings of the 34th International Conference on Machine Volume 70, pp. 1243-1252. JMLR. org, 2017.
Sepp Hochreiter and Jürgen Schmidhuber. Long short-term memory. Neural computation, 9(8):
1735-1780, 1997.
Yuxi Li. Deep reinforcement learning: An overview. arXiv preprint arXiv: 2017.
Peter Liu, Mohammad Salch, Etienne Pot, Ben Goodrich, Ryan Sepassi, Lukasz Kaiser, and Noam Shazeer. Generating wikipedia by summarizing long sequences. arXiv preprint 1011991801.10198. 2018.
Paul Michel, Omer Levy, and Graham Neubig. Are sixteen heads really better than one? arXiv preprint 1905 2019.
Myle Ott. Sergey Edunov, David Grangier, and Michael Auli. Scaling neural machine translation. In Proceedings of the Third Conference on Machine Translation: Research Papers, pp. 1-9, 2018.
Alec Radford, Karthik Narasimhan, Tim Salimans, and Ilya Sutskever. Improving language under- standing by generative pre-training. Technical report, OpenAl, 2018.
Alec Radford, Jeffrey Wu, Rewon Child, David Luan, Dario Amodei, and Ilya Sutskever. Language models are unsupervised multitask learners. Technical report, OpenAI, 2019.
Pranav Rajpurkar, Jian Zhang, Konstantin Lopyrev, and Percy Liang. Squad: 100,000+ questions for machine comprehension of text. arXiv preprint 1roiiv/606.0555 2016.

---

Under review as a conference paper at ICLR 2020
Hao Sun, Xu Tan, Jun-Wei Gan, Hongzhi Liu, Sheng Zhao, Tao Qin, and Tie- Yan Liu. Token-level ensemble distillation for grapheme-to- to-ooenem-- conversion. arXiv preprint 2019.
Iya Sutskever, Oriol Vinyals, and Quoc Le. Sequence to sequence learning with neural networks. In Advances in neural information processing systems, pp. 3104-3112, 2014.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. Advances in Neural Information Processing Systems, pp. 5998-6008, 2017.
Ashish Vaswani, Samy Bengio, Eugene Brevdo, Francois Chollet, Aidan N. Gomez, Stephan Gouws, Llion Jones, Łukasz Kaiser, Nal Kalchbrenner, Niki Parmar, Ryan Sepassi, Noam Shazeer, and Jakob Uszkoreit. Tensor2tensor for neural machine translation. CoRR, abs/1803 (bs//103,,7 2018.
Elena Voita, David Talbot, Fedor Moiseev, Rico Sennrich, and Ivan Titov. Analyzing multi-head self-attention: Specialized heads do the heavy lifting, the rest can be pruned. arXiv preprint arXiv:1905.09418,, 2019.
Xiaolong Wang, Ross Girshick, Abhinav Gupta, and Kaiming He. Non-local neural networks. In Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition, pp. 7794-7803, 2018.
Adina Williams, Nikita Nangia, and Samuel Bowman. A broad-coverage challenge corpus for sentence understanding through inference. In Proceedings of the 2018 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technolo- gies, Volume / (Long Papers), pp. 1112-1122. Association for Computational Linguistics, 2018. 8. org anttoooogo//n-----
URL
Felix Wu, Angela Fan, Alexei Baevski, Yann N Dauphin, and Michael Auli. Pay less attention with lightweight and dynamic convolutions. arXiv preprint 1eXiiv1901.100 2019.
Zhilin Yang, Zihang Dai, Yiming Yang, Jaime Carbonell, Ruslan Salakhutdinov, and Quoc V Le. Xlnet: Generalized autoregressive pretraining for language understanding. arXiv preprint arXiv:1906.08233 2019.
Yang You, Jing Li, Jonathan Hseu, Xiaodan Song, James Demmel, and Cho-Jui Hsich. Reducing bert pre-training time from 3 days to 76 minutes. arXiv preprint arXiv: (00iv::994.00000 909400962, 2019.
Vinicius Zambaldi, David Raposo, Adam Santoro, Victor Bapst, Yujia Li, Igor Babuschkin, Karl Tuyls, David Reichert, Timothy Lillicrap, Edward Lockhart, et al. Relational deep reinforcement learning. arXiv preprint arXiv: 008000 2018.
Han Zhang, Ian Goodfellow, Dimitris Metaxas, and Augustus Odena. Self-attention generative adversarial networks. arXiv preprint 180..v1805.0833 2018.
Yukun Zhu, Ryan Kiros, Rich Zemel, Ruslan Salakhutdinov, Raquel Urtasun, Antonio Torralba, and Sanja Fidler. Aligning books and movies: Towards story-like visual explanations by watching movies and reading books. In Proceedings of the IEEE international conference computer vision, pp. 19-27, 2015.
10

---

Under review as a conference paper at ICLR 2020
A NOTATION
Embedding size
d
Number of layers
Number of heads
Sequence length
Vocab size
Head size
dp
B PROOFS
Proof of Theorem d 2 n case. To prove the first part of the result, we present an explicit construc- tion of W and W which allows us to generate P from X using the dot product attention. Since X has full column rank, there exists left inverse X = (X"X) X X E Raxd such that X = I Let W W and W, W X Then
X X"W,WW, W = (X WT W2X'X II W I, W W
Wky kg
(5)
Now that the above choice of W, and We & has handled the dependence on X, we will choose a Way depending on P and finish the construction. Below express the Softmax operation on the query and key inner products. Note that the Softmax here is columnwise operator computing the attention scores for each query. By using (5), we obtain that
Softmax
(W&X)"
Softmax
W kg Vok
= exp
Wkg Vdg
kg
D
1 Wey
vdk
where Wa. ko is an n x n diagonal matrix such that
(Wkg)a
-(1T
(W
kg)
D Dwe,le )u
exp
exp
Vdle
Hence, we can establish the desired result by showing that there always exists a W kg that satisfies the following fixed point equation.
W
exp
P. DWDi.
(6)
ddz
Given P, to construct such W kg. we pick arbitrary positive diagonal matrix Do, and set
W = Vd log (P- PD))
Since is a positive matrix, such W kig always exists. Next, we verify that this construction indeed satisfies the fixed point equation (cf. (6)). Note that
(W kw))) Vdg
)
D D Wi..
= Diag
1" cxp
Diag (1"'-Do) Do-
(8)
The last equation follows from the fact that is column stochastic matrix. Now, using 7 and
W kgg
exp
P. D. = PP-D..
This completes the first part of the proof.
< n case.
Consider the case of d = and n = 2. Then X E Rix2 and W, and
W₁ E B1x1 Let X = [1, 0]. Then
(W x (W,X)
Softmax [1 W W W [1, 0]
W W
0
Softmax
Softmax
Vdk
Vd
This matrix clearly cannot be used to generate P that have distinct elements in the second column
[0.5
0.75
egg., P
0.5
0.25
11

---

Under review as a conference paper at ICLR 2020
First let rewrite the MultiHead and FixedMultiHead layers follows. The MultiHead layer can be rewritten as
Proof of Theorem 2
[w = W. MultiHead()) =
Σ W W Softmax [wix)]
(wix)"((xx)) (/:x) \
where W are d x d/h matrices and W e² W and W are d/h x matrices. We denote the collection of all parameter matrices W.
Similarly, rewrite the fixed head size attention layer as
gv gv(X) (X) = Vo FixedMultiHead
Σ V ==1
Softmax (VVX)
where V E Rd and V k. V E Rdxd Let V be the collection of all these matrices.
The outline of the proof is basically case analysis: we divide possible values of W into three categories, and show in each case that there exists a X such that (X(X) * gv(X). Here are the three cases:
Case 1: EE W: + Σ V.V,
Case 2: E W: E V and there exists E {l...., ,h} such that U/ ad (Wi) /// d/h is not syew-symmetric.
Case 3: Eh w = Ehe V. uy and all U/ (W)) /d/h are skew- symmetric.
In the first case, we can choose any such that (E W: - Σ V.V. Choose X = -1"==1 v v]. Then, note that for any column stochastic matrix P, we have
Case
XP = X. Therefore,
W Softmax (w:)))(w;x///// h
h
Σ
v Softmax [VXx) VT,)
1=1
W Σ V.
w
v. #
since Σか 1 V is full rank by assump- tion and each W::: is at most rank d/h. it follows that all columns in W E gdxd/h must be linearly independent. Therefore, for any v + 0, (W: W.W.v. i 1....,h h) set of linearly independent vectors, because each W W v is a linear combination of d/h column vectors of W that are linearly independent of other column vectors in W j #
Case 2. In cases where E W = V.
Now consider any v E Rd, and X = vet, where exp(t)/(exp(I) + 1). Then, we have
(1, 0) E R". Define @(t)
Uv
0
dp
(v Σ 1=1
V V Softmax xTUx\
Σ V v Softmax
0
V
-T dp
W Wi
/dp
12

---

Under review as a conference paper at ICLR 2020
Similarly, calculate
fw(X)= Σ W W Softmax (wix)/
/a/h
h
(
v" ('"())""W"
= Σ
w
W::
Va/h Va/h
Notice that all the columns of fw and gv (X), from the second columns to the last ones, are the same. We now compare the first columns:
h Σ ==1
(W)) W
(v'Uv] w'wir.
fw(X).1 -
Vd/h ④
ddp
Recall that for any v + O. W W' are linearly independent, fw - 0 if and
only if all Φ v" "WW""W's
Uv zero. However, since there exists (I,...,h)
Vath
such that U/ U/ (Wi) (W d/h not skew-symmetric, we can choose v to one that
satisfies "² (W
Vd/h * UF
hence making & v (WW))) WWw á
* 0, therefore
Vajh
ddp
fw( (V) * O.
Case Now consider any [v1 Define Φ1 (2),,t)= n - 2). Then, we have
V2
0). where v, and V2 will be chosen later.
2), t2)
v UVI
Uv2
0
ddz
dp
Uv:
Uv2
h gv(X) Σ v.vi Softmax
0
dly
:0:0
Therefore, the first column of gv (X) can be written
vi Uv ddp
Uvi
Vi
SV (X),,, Σ W:w:
Φ1
ddp
vi +
(Ti
Vip
ddz
V2
Similarly, the first column of (w)Xiix is
Wiv (WW)) WW. vI (W:)
fw(X).,1
W Wi
vit
Va/h
Vd/h
VF(W))" VI vs (W
Φ
V2
Vd/h
Vd/h
Since U/ (W)) is skew-symmetric by assumption, we have
U (WL)T(W)
vT
vi 0 for all Vi. Recall that U is rank-d, by assumption, so d/h is at least rank dp d/h 2 1, so we can choose any v such that
Vath
U/
dd (WH)"(W))
U (WLTT(W)
* 0.
44
Va/h
13

---

Under review as a conference paper at ICLR 2020
(W)"(wW)) ddth
both
vd, Vi and
V1 are nonzero, We can always choose V2 such that v
v
vd
(W))"(W") Vdjh
0 and V/2
V1 < O. This means that if we choose v2 = au² and scale a - oc,
Ui v. vv
Φ
0 02
++
ddp Vdp
dd
ddp
( v2 (WW)T
exp(vf (WL) W / Va/h)
Φ1
Vath
Vath
explvf (W))T
V (W()) W V
2
Va/h
Vd/h
Then, consider the difference (x),)., gV (X): ). Recall that for any v. W! is independent of (W: v, + 13. This means that, to show fw (X):, 1- (x o O, suffices to show that
T TW, (WW)
ViUv Uv
W!
Φ1
Va/h
Vafh
Vdz Vdp
vi (WW)" FVI
V
(W W
viTUv
vV
2
02
W W! + 0.
Va/h
Va/h
Vdp
vdp
we scale V2 = av2 with large enough a, the second term will dominate the first term and the first term will never be able to cancel the second one. Thus, by choosing large enough a >> 0, we can make sure that the sum is nonzero.
Even in case where one of U and (WL) is zero (say (W (WL) Vdp choose 22 U Vi and use a similar scaling argument. By choosing large enough a > 0 and
0), we can
Va/h
Vayk
ddp
V2 aV,,, one can show that the difference /w(X).,1 - (X),, nonzero.
C EXPERIMENTAL SETTINGS
For experiments with the language modeling (LMIB dataset), train 6 layer transformer models. We use a batch size of 4096 and train for 250k steps. We use learning rate of 0. with a linear warm up for the first 10k steps. We decay the learning rate with the square root of the number of steps. We train the standard transformers with the embedding dimension varying from 256 to 512. We fix the width of the feed forward layer in the Transformer to 1024. In addition, we use weight decay of 0.01 and dropout with probability of 0.1 on all the layers.
For our experiments with BERT, we follow the same experimental settings as (Devlin et 2018). We present the key details here and refer the reader to (Devlin et al. (0188). We train with batch size of 1024 for 450k steps with inputs of sequence length
128 follower by 50k steps with inputs of sequence length 512. In contrast the BERT paper uses batch size of 512, and does the pre-training for 900K steps with 128 sequence length inputs and 100k steps with 512 sequence length inputs. We train using ADAM with learning rate of le-4, and linear warmup and decay schedule as in BERT. We use 5k warmup steps for the first stage, and re-warmup 3k steps for the second stage Youu al. 2019). Again, use weight decay of 0.01 and dropout with probability of 0.1 on all the layers.
For the language modeling task, training is performed on 4 TPUv2 chips for a couple of hours. For BERT models training is performed on 16 TPUv3 chips in the first stage and 64 TPUv3 chips for the second stage. Pre-training with this configuration takes between to 3 days. We did not attempt find the optimal hyper-parameters for the fixed head size architecture, and use the same hyper-parameters as used for training the standard Transformer.
D ADDITIONAL EXPERIMENTAL RESULTS
14

---

Under review as a conference paper at ICLR 2020
41
d=256, h=8
dp = dd = 32
40
d d-256, dp 32
39
Perplexity
38
3884 h=12
d=448, h=14
36
d=256,
h=16
35
d=256,
h=32
d=512, h=16
34
d=256,
h=48
d=256,
h=60
d=256,
n070
33
16
18
20
22
24
26
28
30
# Parameters in millions
Figure Performance of the standard Transformer training compared with the fixed head size (dp) models for language modeling task (LMIB) the test set. Unlike Fig [ we vary both the embedding size and the number heads of the standard Transformer to keep its head size fixed to 32, We train the fixed head size models with fixed embedding size of 256 and a head size of 32, and vary the number of heads from 4 to 70, while matching the number of parameters. The plot again clearly indicates the advantage of the fixed head size models. The main issue with standard Transformers is that fixing the head size to 32, forces the number of heads to be small when the dimension is small. Reducing the number of heads below certain extent hurts the performance of the Transformer.
# heads
12
16
20
# params
214M
252M
290M
327M
SQuAD Fl
90.35:1.114
90.48.10.09
90.92-1014
90.8910.08
SQuAD EM
83.37:0.12
83.6740.03
84.1620.35
84.2910.16
MNLI
84.4+0.2
84.4+0.2
84.7-0.1
85.1
(A) Increasing number of heads
Table 3: (A): 24 layer modified Transformer with fixed head size of 128 and embedding size of 768 shows an improvement in the accuracy with the increasing number heads.
15


Reviews:
[{"review_id": "SJe97j5FKr", "reviewer": "AnonReviewer3", "summary": "", "questions": "", "limitations": "", "rating": 1, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "This paper proposes a “concise” version of the well-established Transformer model. The main proposal is to explicitly set the head size in the Transformer model instead of having to divide (share) representation ability amongst heads. \n\nThis paper is poorly written and after an entire 2-page long-winded introduction, the reader is left wondering what is the main contribution of this work. The term “concise” is also not well-defined and left vague to readers. I re-read this paper multiple times and the only concluding finding I have is that this paper proposes an explicit way of setting the projection dimension regardless of the number of heads. \n\nAfter many mathematical formulations, theorems (seemingly ornamental, or handwavy actually), the final contribution seems to be to set the head size of BERT (size of each head) to 128. This is really trivial. The authors kept teasing a “different” way to do this, but this left the reader completely unsatisfied when the different way refers to explicitly setting each head to 128 and using a smaller model overall. \n\nThe value 128 is derived from a theorem derived by the authors, which suggests that each head should at least be greater or equal than the sequence length (the sequence length here stated by the authors is 128). I’m not very convinced by the argument. While it is intuitive that each head has to be sufficiently large, being under-sized can be made up for with multiple heads. It is also not clear why every X and P must be expressed with transforms W_q and W_k. P here represents the affinity matrix between tokens in a sequence.  It does not make any sense to me to ensure that every variation of P can be expressed because P is literally the pairwise scores between every token in the fully-connected attention graph. \n\nWhile I did not have the luxury of time to parse the Appendix to validate the legitimacy of the proof, I think the overall shortcomings of the paper (highly non-readable, bad presentation and perhaps a fair attempt at masking the lack of contribution) warrants a clear reject from me. \n", "extra_scores": {"experience_assessment": "I have published in this field for several years.", "review_assessment:_checking_correctness_of_derivations_and_theory": "I assessed the sensibility of the derivations and theory.", "review_assessment:_checking_correctness_of_experiments": "I assessed the sensibility of the experiments.", "review_assessment:_thoroughness_in_paper_reading": "I read the paper at least twice and used my best judgement in assessing the paper."}, "extra_text": {}}, {"review_id": "S1xczKie9S", "reviewer": "AnonReviewer1", "summary": "", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "This work discusses how to set the projection size for each head (head size) in multi-head attention module, especially Transformer. Theorem 1 is interesting, which points out a lower bound for the head size. The proposed method is to decouple the dependency between the head size and the embedding size. The experiments show that the proposed method is able to achieve comparable performance to BERT with fewer training cost.\n\nThe lower bound for the head size is a valuable result. However, the novelty is very limited. To decouple the dependency between the head size and the embedding size is not a novel point. In BERT/Transformer, it is set d_q=d_k=d_v=d, which is not a strict constraint. The only constraint in attention is to have d_q=d_k to allow dot product. Therefore, the proposed method is more like a tuning of hyper-parameters.", "extra_scores": {"experience_assessment": "I have published one or two papers in this area.", "review_assessment:_thoroughness_in_paper_reading": "I read the paper at least twice and used my best judgement in assessing the paper.", "review_assessment:_checking_correctness_of_experiments": "I assessed the sensibility of the experiments.", "review_assessment:_checking_correctness_of_derivations_and_theory": "I assessed the sensibility of the derivations and theory."}, "extra_text": {}}, {"review_id": "BJxY2KeTtH", "reviewer": "AnonReviewer2", "summary": "", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "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. \n\nPros:\n- The arguments on head size and head number tradeoff could be inspiring to future works.\n- A simple approach that proves strong in several NLP tasks.\n\nCons:\n- The theoretical discussion imposes too strong assumptions that might make it less interesting in practice.\n- No NMT experiments.\n- The takeaway seems a bit trivial.\n\nDetails:\n- Theorem 1 presents a rank-based view of each attention head's capacity, which is nice. Yet it is still unclear whether it is the case that the more expressive are the heads the better. For example, several recent works argues for specialized attention heads, i.e., each head has specific \"job,\" which may not require it being very expressive [1, 2]. Further, other works shows that a low-rank P matrix could be beneficial [3, 4, 5], which contradicts the argument in this work. It would be nice if the authors and discuss this in the revision\n\n(To be clear, I do believe this is still an open question, and do not think presenting a different view from previous works hurts the contribution of this work in any way.)\n\n- Theorem 2. I didn't carefully check the proof. Why is it required that the V matrices for each head have the same product. For both Thm.1 and 2, it would be nice to see some discussion on how they translate into the models in practice.\n\n- Can the authors compare the training/inference speed? It probably will be the same as standard transformers, but it would be nice to confirm.\n\n- Figure 1: the caption says trying out embedding sizes from 256 to 512. But it seems that only 4 values are tried. Can the authors comment on this? Also, it is a bit awkward to plot a line chart out of 4 points. Same for Figure 2.\n\n- It would be nice to see some NMT experiments. \n\n- The proposed method is so straightforward that I'm actually very surprised that this paper is the first trying this. The authors might need justify the technical contribution more.\n\n(I'm on the fence for this one, but the system doesn't allow me to. I'm happy to revise the score if the authors can address my concerns.)\n\n\n[1] Analyzing Multi-Head Self-Attention: Specialized Heads Do the Heavy Lifting, the Rest Can Be Pruned. https://arxiv.org/abs/1905.09418\n\n[2] Are Sixteen Heads Really Better than One? https://arxiv.org/abs/1905.10650\n\n[3] Generating Long Sequences with Sparse Transformers. https://arxiv.org/abs/1904.10509\n\n[4] Generating Long Sequences with Sparse Transformers. https://arxiv.org/pdf/1904.10509.pdf.\n\n[5] Adaptively Sparse Transformers. https://arxiv.org/abs/1909.00015.", "extra_scores": {"experience_assessment": "I have published in this field for several years.", "review_assessment:_thoroughness_in_paper_reading": "I read the paper at least twice and used my best judgement in assessing the paper.", "review_assessment:_checking_correctness_of_experiments": "I assessed the sensibility of the experiments.", "review_assessment:_checking_correctness_of_derivations_and_theory": "I assessed the sensibility of the derivations and theory."}, "extra_text": {}}]

Author rebuttal:

Meta-review: This paper studies tradeoffs in the design of attention-based architectures. It argues and formally establishes that the expressivity of an attention head is determined by its dimension and that fixing the head dimension, one gains additional expressive power by using more heads.

Reviewers were generally positive about the question under study here, but raised important concerns about the significance of the results and the take-home message in the current manuscript. The AC shares these concerns, and recommends rejection, while encouraging the authors to address the concerns raised during this discussion. 
Decision: Reject