C ONCISE M ULTI - HEAD A TTENTION M ODELS
Anonymous authors
Paper under double-blind review
A BSTRACT
Attention based Transformer architecture has enabled signiﬁcant advances in the ﬁeld 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.
Unfortunately, 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 sacriﬁcing the performance.
I NTRODUCTION
Attention based architectures, such as Transformers, have been effective for sequence modelling tasks such as machine translation (Gehring et al., 2017; Vaswani et al., 2017), question answering, sentence classiﬁcation
(Radford et al., 2018;
Devlin et al., 2018) and document generation (Liu et al., 2018).
These models have emerged as better alternatives to the recurrent models - RNNs (Sutskever et al., 2014), LSTMs (Hochreiter & Schmidhuber, 1997) and GRUs (Cho et al., 2014).
This is mainly due to their feed forward structure, which removes the sequential processing bottleneck for sequence data, making them easier to train compared to the recurrent models.
Self attention models also have found applications in vision (Wang et al., 2018), adversarial networks (Zhang et al., 2018), reinforcement learning (Zambaldi et al., 2018; Li, 2017) and speech recognition (Chiu et al., 2018).
Recent advances in using the self attention models in natural language tasks have been made by ﬁrst using a language modeling task to pre-train the models and then ﬁne tuning the learned models on speciﬁc downstream tasks.
Radford et al.
(2018) and Devlin et al.
(2018) used Transformers to pre-train a language model and showed that the ﬁne tuned model outperforms LSTMs on many natural language understanding and question answering tasks.
For example, BERT (Devlin et al., 2018), a 24 layer transformer model, is shown to achieve the state of the art performance on several
NLP tasks, including on the SQuAD dataset.
These advances, in addition to novel pre-training tasks, relied on bigger models with a larger embedding size.
BERT model uses an embedding size of(Devlin
et al., 2018); GPT-2 uses models with embedding size up to 1600 (Radford et al., 2019).
A single Transformer block consists of two key components: a multi-head self attention layer followed by a feed forward layer (Vaswani et al., 2017).
A single head in a multi-head attention layer, computes self attention between the tokens in the input sequence, which it then uses to compute a weighted average of embeddings for each token.
To keep the number of parameters ﬁxed in the attention layer, each head projects the data into a lower dimensional subspace, dimension of which scales as 1/(number of heads), and computes the self attention in this subspace.
This projection size for each head is commonly known as the head size .
Despite the advances in using Transformer models for various tasks, their functioning and design choices still remain mysterious and are not well understood.
Can the attention layer learn arbitrary contextual representations?
What is the role of the feed forward layer in the Transformer block?
Do we need such a large embedding size to capture the context of all the tokens?
Answering these questions requires an understanding of the representation power of the units in the Transformer.
In this paper we take some of the ﬁrst steps towards developing such an understanding of the
Transformer.
In particular, we focus on the representation power of the multi-head self attention layer.
[Table]
First, we analyze the representation power of a single self attention unit and show that it crucially depends on the projection sizes used to compute the dot product attention.
We next study the advantage of having multiple heads in the attention layer.
It is generally believed that increasing the number of heads helps by allowing the heads to compute context from different representation subspaces at different positions.
However, increasing the number of heads decreases the head size, decreasing the expressive power of individual heads.
We show that reducing the head size to a value below the input sequence length hurts the representation power of each head.
This is because a smaller head size introduces a rank constraint on the projection matrices in each head, and limits their representation power.
We indeed notice this effect in practice: while the performance improves with increasing the number of heads in the beginning (Devlin et al., 2018), we notice a drop in the performance once the number of heads increases beyond a certain threshold, as seen in Tableand Fig.
1 (see also Table 4(A) in Vaswani et al.
This heuristic of scaling the head size inversely with the number of heads was proposed initially in Vaswani et al.
and has become the standard way of using multi-head attention (Radford et al., 2018; Devlin et al., 2018).
In order to avoid hurting the performance, the existing Transformer models allow for multiple heads by increasing the embedding size, which in turn increases the head size.
However, larger embedding size, in addition to increasing the number of parameters, makes it expensive to use the model and the learned embeddings in downstream tasks, as the downstream model sizes scale with the embedding size of the tokens.
For example, the inference time and memory required in retrieval tasks increases linearly with the embedding size.
Based on these observations, we propose a new way to set the projection size in the attention heads, in which each head has a ﬁxed head size that is independent of both the number of heads and the embedding size of the model.
This allows us to train models with a relatively smaller embedding size without affecting the head size.
It also allows us to increase the number of heads per layer to improve the performance.
Another advantage of the ﬁxed head size Transformer is, unlike the standard Transformer, which requires the number of heads to be a factor of the embedding size, we are free to set arbitrary number of heads as required for the task.
We evaluate this ﬁxed head size Transformer on language modeling (LM1B dataset), natural language inference (MNLI dataset) and question answering tasks (SQuAD dataset).
We show that the modiﬁed Transformer trained with an embedding size of 512 can match the performance of the
BERT LARGE (Devlin et al., 2018), a Transformer with an embedding size of 1024 (see Fig. 2).
We further present experimental results evaluating the effect of different choices of the head size and the embedding size in the Section 4.
The contributions of this paper are summarized below.
• We analyze the representation power of the multi-head self attention layer and show the limitation the embedding size places on the number of heads.
• We propose a new way to set the head size, and show the proposed ﬁxed head size layers are strictly better than the standard multi-head attention layers in terms of their expressive power.
This modiﬁcation allows us to both increase the number of heads per layer and decrease the embedding size, without hurting the performance.
• We experimentally show that the ﬁxed head size Transformer can be trained with a smaller embedding size and more heads on three standard NLP tasks.
1. R ELATED W ORKS
Given the signiﬁcance of self attention models, there has been work trying to both improve the performance and speedup the computation in Transformers.
Ott et al.
and You et al.
(2019) reduce precision and use large batch training to reduce the training time of the attention models.
Child et al.
(2019) propose sparse self attention models to speed up the computation in the attention layer for long sequence data generation tasks.
They show that these sparse attention models can be trained on tasks with sequence length greater than 10k without sacriﬁcing the accuracy.
Dehghani et al.
(2018) propose a depth recurrent Transformer network that reuses the parameters across layers.
They show that this modiﬁcation makes the Transformer networks Turing complete even with ﬁnite precision weights.
Yang et al.
(2019) propose a new way to increase the effective sequence length that the Transformer attends to, by reusing the intermediate embeddings across sequences.
They show that the modiﬁed architecture performs better on tasks that require computing context over longer sequence lengths.
We note that most of these modiﬁcations 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 modiﬁcations.
Wu et al.
(2019) propose to replace the self-attention layer with lightweight dynamic convolutions and show improved performance on machine translation and language modeling.
Even though the resulting model has faster inference time, it still needs to use a large embedding size (1024), as big as the original attention models.
We believe the techniques in this paper can be combined with these results to realize both smaller embedding size and faster inference time.
Sun et al.
(2019) perform neural architecture search using evolutionary methods on sequence to sequence models and ﬁnd an evolved transformer architecture, which in addition to multi-head attention units, has convolution ﬁlter and gated linear units.
Our proposed modiﬁcations stay closer to Transformers in spirit and can be used as seed units for this architecture search. Voita et al. ; 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. ; Correia et al. (2019) impose sparsity structure on the attention layer during training to improve both interpretability and performance. Fixing the head size will in fact make it easier to learn such sparsity patterns, as a low rank constraint does not allow a head to express all possible sparsity patterns.
Combining these techniques can hence potentially enable training of sparse attention models with a smaller embedding size.
T RANSFORMER A RCHITECTURE AND A NALYSIS
In this section we present the Transformer architecture and analyze the representation power of the multi-head self attention, a key component of the Transformer block.
The input to a Transformer network is a sequence of n tokens.
Typically, each token is converted into a token embedding of dimension d by an embedding layer.
We let X ∈ R d × n be the embedding matrix corresponding to the n tokens in the input sequence.
2. S INGLE -H EAD A TTENTION
The transformer block is a combination of a self attention layer followed by a feed forward layer (Vaswani et al., 2017).
Both layers have a skip connection and use Layer Normalization (LN)
(Ba et al., 2016).
In particular, for token embeddings X , the dot product attention is computed as follows.
[Equation]
Here W q ∈ R d q × d , W k ∈ R d k × d
and W v ∈ R
d v × d represent the projection matrices associated with the query, key and value respectively in an attention unit (Vaswani et al., 2017).
For a single-head attention unit, we have d q = d k = d v = d .
In the dot-product attention (cf.
Subsequently, the output of the attention layer takes the following form.
[Equation]
where LN ( · ) represents the layer-normalization operation.
Given the attention module, as deﬁned in (1), it is natural to question its ability to represent arbitrary contexts P for a given input sequence X .
Towards this, we show that for a large enough projection dimension d , the unit has enough capacity to represent arbitrary contexts over a given input sequence.
In the following result we establish that for a large enough projection size an attention unit can represent any data pair ( X , P ) .
We also show that the model cannot represent arbitrary context when d is smaller than n .
Theorem 1 (Representation Theorem) .
If d q = d k = d ≥ n , then given any full column rank matrix
X ∈ R d × n and an arbitrary n × n positive column stochastic matrix P , there always exists d × d projection matrices W q and W k such that
[Equation]
If d q = d k = d < n , there exist X and P such that (3) does not hold for all W q and W k .
The proof of Theorem 1 is provided in the supplementary material.
This result shows that the projection dimension d q = d k = d needs to be larger than the sequence length n for the attention unit to be able to represent any desired context P .
Even though this result describes a single example sequence case, it highlights a fundamental property of the model architecture that increasing the projection size increases the capacity of the attention heads.
2. M ULTI -H EAD A TTENTION
As discussed in Section 2.1, 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.
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:
[Equation]
The output of the Multi-head attention layer then becomes
[Equation]
where W o ∈ R d × d .
For a model with h heads, the query, key and value projection matrices { W iq } , { W ik } and { W iv } are dh × d matrices.
Therefore, each head projects the input onto a dh -dimensional subspace to compute the context, and keeps the number of parameters ﬁxed per layer.
This has been observed empirically to increase the expressive power of the attention layer (Vaswani et al., 2017).
2. D EPENDENCE 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 h is larger than dn , the attention unit inside each head projects onto a dimension smaller than n , and looses its ability to represent arbitrary context vectors (cf.
Theorem 1).
Since the sequence length is ﬁxed from the data/task at hand, the only remaining way to 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 as well.
Given the widespread success of attention mechanism, this highlights the need for a modiﬁed model architecture that can leverage the advantages of MultiHead without suffering from the requirement of large embedding sizes.
C ONCISE M ULTI - HEAD T RANSFORMER
In this section we propose a different way of setting the head size of the Transformer, which allows us to 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 a ﬁxed dimension d p irrespective of the number of heads h .
This approach where d p is independent of d and h leads to the following attention mechanism.
[Equation]
Note that the projection matrices used here { V iq } , { V ik } and { V iv } are d p × d matrices.
With
V o ∈ R
d × h · d p , the output of this new multi-head attention layer takes the following form.
[Equation]
This modiﬁcation makes each attention head more similar to a hidden unit in a feed forward network or a ﬁlter 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 modiﬁcation allows us to train a model with a smaller embedding size, ultimately allowing us to reduce the total number of parameters in the model.
Choice of the head size.
Our proposed modiﬁcation introduces head size d p as a new model hyperparameter.
We choose head size to be 128 for our BERT experiments, as most of the pre-training is done with 128 sequence length data.
While we have ablation studies (cf.
Table 2(B))
showing bigger head size improves the performance, there is a tradeoff between increasing the head size vs number of heads vs layers.
We found that having sufﬁcient head size matching the pre-training sequence length, is better than having a larger embedding size (cf.
Section 4).
3. M ULTI H EAD VS . F IXED M ULTI H EAD ATTENTION
Given a MultiHead layer, we can always represent it using a FixedMultiHead layer, whenever we have the head size d p ≥ d / h .
While this shows that increasing the number of heads h beyond d / d p makes individual heads of the FixedMultiHead as expressive as the ones in the MultiHead, it is not clear if 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 FixedMultiHead layer is strictly better than the standard MultiHead layer in terms of expressive power.
Consider the standard multi-head attention units in (4).
[Equation]
We denote the collection of all parameter matrices as W .
Similarly, consider the function represented by the ﬁxed head size attention units:
[Equation]
Let V be the collection of all these parameter matrices.
We deﬁne F and G to be the class of functions f W ( · ) and g V ( · ) , respectively.
As noted above, if d p ≥ d / h , we have F ⊂ 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 a strict subset of G .
Theorem 2.
Given n ≥ 2 and d ≥ d p , let h > d / d p .
Consider a ﬁxed head size attention layer g V ( · ) with parameters that satisfy the following conditions:
[Equation]
 is full rank, and ( V ik ) T V iq =
U for all
i = 1 , . . .
, h, where U is a rank- d p matrix.
Then, for any f W ∈ F , there exists X ∈ R d × n such that f W ( X ) (cid:54) =
g V ( X ) .
Because (cid:107) f W ( X ) − g V ( X ) (cid:107) is a continuous function of X , existence of such an X implies that the integral of the norm of difference (i.e., approximation error) is 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 ﬁxed head size Transformers will increase the expressive power of the attention layers
E XPERIMENTS
[Figure] [Figure]
In this section we present our experiments on three standard NLP tasks, language modeling (LM1B), 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 modiﬁcation; 2) Decoupling the head size from embedding size allows us to train models with a smaller embedding size; and 3) Setting the head size appropriately in the
Transformers allows us to train models with a better performance scaling.
We ﬁrst describe our experimental setup followed by our results and ablation studies on the proposed modiﬁcations.
[Figure]
4. S ETUP AND D ATASETS
For the language modeling task we use the one billion word benchmark dataset (LM1B) (Chelba et al., 2013).
This dataset has around 30M training examples and around 300k examples in the test set.
We use a sub-word tokenizer with 32k vocab and cap the input to 256 sequence length.
We train a 6 layer Transformer model with the ADAM optimizer using the tensor2tensor library (Vaswani et al., 2018).
The detailed experimental setting is presented in Section C.
Multi-Genre Natural Language Inference (MNLI) is a 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 or is neutral to the premise.
We report the classiﬁcation accuracy for this task.
Stanford Question Answering Dataset (SQuAD) is a question answering dataset, where given a paragraph and a question, the goal is to predict the sequence of words in the paragraph that constitute the answer to the question (Rajpurkar et al., 2016).
This is a harder word level task, compared to the sentence classiﬁcation task.
We report both Exact
Match (EM) and F1 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 ﬁrst pre-training on a language modeling task and then ﬁne-tuning the models on the task data.
We follow the same experimental setup for both pre-training and ﬁne-tuning as BERT (Devlin et al., 2018), and use their codebase 1 .
We ﬁrst pre-train our models using the masked language model and the next sentence prediction objectives, and then ﬁne tune the pre-trained model for individual tasks (Devlin et al., 2018).
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 a complete description of the setup.
4. R ESULTS
For our ﬁrst set of experiments we want to see if the ﬁxed head size Transformer with a smaller embedding size can match the performance of standard Transformers with a 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 ﬁxed head size Transformers with a [Table] ﬁxed embedding size of 256 and a head size of 32, with an increasing number of heads from 4 to 70.
We notice that the ﬁxed head size models with an embedding size of 256 can match the performance of standard Transformers with an embedding size of 512 (see Fig. 1).
Further this provides a better performance scaling than the standard Transformers.
We repeat the similar experiment on the other two tasks, where for baseline we train BERT LARGE , a 24 layer, 16 head Transformer architecture, with embedding sizes from 512 to 1024.
We compare it with the modiﬁed model, with an embedding size of 512 and a head size of 128, with an increasing number of heads from 8 to 32.
We again notice that the ﬁxed head size model with 512 embedding size can match the performance of BERT LARGE (see Fig. 2).
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. 4).
Hence, both the head size and the number of heads needs to be set high enough for better performance.
4. A BLATION
Increasing heads.
From Table 1 and Fig.
1a we can see that increasing the number of heads hurts the performance of the Transformer after a certain number.
We repeat the same experiments with the ﬁxed head size Transformer, and present the results in Table 2(A) and Fig. 3a.
The results show that the performance of the modiﬁed model improves monotonically as 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 2(B) and Fig.
3b, we present comparisons between models with different head sizes.
This shows that the gains in the performance of the ﬁxed head size models indeed come from adjusting the head size of the query, key and value layers in the attention unit.
The table shows a clear trend of better performance with a larger head size, suggesting that it indeed is an important factor in the performance of the attention models.
C ONCLUSION
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 modiﬁed way to set the head size that allows us to increase the number of heads without increasing the embedding size.
As a consequence we are able to train Transformers with a 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 is an interesting direction for the future work
B P ROOFS
Proof of Theorem 1.
d ≥ n case .
To prove the ﬁrst part of the result, we present an explicit construction of W k and W q which allows us to generate P from X using the dot product attention.
Since X has full column rank, there exists a left inverse X † =
( X T X ) − 1 X T ∈ R n × d such that X † X
= I n .
Let W k
=
˜
W k X † and W q = ˜
W q X † .
Then
[Equation]
Now that the above choice of W q and W k has handled the dependence on X , we will choose a ˜ W kq depending on P and ﬁnish the construction.
Below we express the Softmax operation on the query and key inner products.
Note that the Softmax here is a columnwise operator computing the attention scores for each query.
By using (5), we obtain that
[Equation]
where D ˜ W kq is an n × n diagonal matrix such that
[Equation]
Hence, we can establish the desired result by showing that there always exists a ˜ W kq that satisﬁes the following ﬁxed point equation.
[Equation]
Given P , to construct such a ˜ W kq , we pick an arbitrary positive diagonal matrix D 0 , and set
[Equation]
Since P is a positive matrix, such a ˜ W kq always exists.
Next, we verify that this construction indeed satisﬁes the ﬁxed point equation (cf.
Note that
[Equation]
The last equation follows from the fact that P is a column stochastic matrix.
Now, using (7) and (8),
[Equation]
This completes the ﬁrst part of the proof.
d < n case .
Consider the case of d = 1 and n = 2 .
Then X ∈ R 1 × 2 and W q and
W k ∈ R 1 × 1 .
Let X =
Then
[Equation]
This matrix clearly cannot be used to generate P that have distinct elements in the second column , e.g., P = (cid:20) 0 .
[Table]
(cid:21) .
Proof of Theorem 2.
First let us rewrite the MultiHead and FixedMultiHead layers as follows.
The
MultiHead layer can be rewritten as
[Equation]
where W io are d × d/h matrices and W iv , W ik , and W iq are d/h × d matrices.
We denote the collection of all parameter matrices as W .
Similarly, rewrite the ﬁxed head size attention layer as
[Equation]
where V io ∈ R d × d p , and V iv , V ik , V iq ∈ R d p × d .
Let V be the collection of all these matrices.
The outline of the proof is basically a case analysis: we divide possible values of W into three categories, and show in each case that there exists a X such that f W ( X ) (cid:54) =
g V ( X ) .
Here are the three cases:
• Case 1 : (cid:80) hi =1 W io W iv (cid:54) =
(cid:80) hi
=1 V io V iv .
• Case 2 : (cid:80) hi
=
1 W io W iv = (cid:80) hi
=1 V io V iv , and there exists i ∈ { 1 , . . .
, h } such that U / (cid:112)
d p − ( W ik ) T ( W iq ) /
(cid:112) d/h is not skew-symmetric.
• Case 3 : (cid:80) hi
=
1 W io W iv = (cid:80) hi
=1 V io V iv , and all U / (cid:112)
d p − ( W ik ) T ( W iq ) /
(cid:112) d/h are skewsymmetric.
Case
In the ﬁrst case, we can choose any v such that ( (cid:80) hi
=1 W io W iv − (cid:80) hi
=1 V io V iv ) v (cid:54) = 0 .
Choose X = v1 T =
[ v v . . .
v ] .
Then, note that for any column stochastic matrix P , we have
XP = X .
Therefore, h
[Equation]
Case 2.
In cases where (cid:80) hi =1 W io W iv = (cid:80) hi
=1 V io V iv , since (cid:80) hi
=1 V io V iv is full rank by assumption and each W io W iv is at most rank d/h , it follows that all columns in W io ∈ R d × d/h must be linearly independent.
Therefore, for any v (cid:54) = 0 , { W io W iv v , i = 1 , . . .
, h } is a set of linearly independent vectors, because each W io W iv v is a linear combination of d/h column vectors of W io that are linearly independent of other column vectors in W jo , j (cid:54)
=
i .
Now consider any v ∈ R d , and X = ve T 1 , where e 1 = (1 , 0 , . .
. , 0) ∈ R n .
Deﬁne φ ( t ) = exp( t ) / (exp( t ) + n − 1) .
Then, we have g V ( X )
= h (cid:88) i =
V io V iv X · Softmax (cid:2) X T UX / √ d p (cid:3) =
h (cid:88) i =
V io V iv X · Softmax    v T Uv √ d p 0 . . .
0 . . .
  = (cid:32) h (cid:88) i =
V io V iv (cid:33) (cid:20) φ (cid:18) v T Uv √ d p (cid:19) v v n . . .
v n (cid:21) =
(cid:32) h (cid:88) i =
W io W iv (cid:33) (cid:20) φ (cid:18) v T Uv √ d p (cid:19) v v n . . .
v n (cid:21) .
Similarly, we can calculate
[Equation]
Notice that all the columns of f W ( X ) and g V ( X ) , from the second columns to the last ones, are the same.
We now compare the ﬁrst columns:
[Equation]
Recall that for any v (cid:54) = 0 , W io W iv v are linearly independent, so f W ( X ) : , 1 − g V ( X ) : , 1 =
if and only if all φ
[Equation]
are zero.
However, since there exists i ∈ { 1 , . .
, h } such that U / (cid:112)
d p − ( W ik ) T ( W iq ) /
(cid:112) d/h is not skew-symmetric, we can choose v to be one that satisﬁes v T ( W ik ) T W iq v √ d/
[Equation]
(cid:54) =
[Equation]
T Uv √
[Equation]
p , hence making φ
[Equation]
d/h (cid:19) − φ (cid:18) v T Uv √ d p (cid:19) (cid:54) = 0 , therefore f W ( X ) : , 1 − g V
( X ) : , 1 (cid:54) = 0 .
Case 3.
Now consider any X = [ vv0 . . .
0 ] , where v 1 and v 2 will be chosen later.
Deﬁne φ 1 ( t 1 , t 2 ) = exp( t 1 ) / (exp( t 1 )+exp( t 2 )+ n − 2) , φ 2 ( t 1 , t 2 ) =
exp( t 2 ) / (exp( t 1 )+exp( t 2 )+ n − 2) .
Then, we have g V ( X )
= h (cid:88) i =
V io V iv X · Softmax   v T 1 Uv 1 √ d p v T 1 Uv 2 √ d p 0 . . .v
T 2 Uv 1 √ d p v T 2 Uv 2 √ d p 0 . . .
0 . . .
 .
Therefore, the ﬁrst column of g V ( X ) can be written as g V ( X ) : , 1 = (cid:32) h (cid:88) i =
W
[Equation]
iv (cid:33)
[Equation]
φ
[Equation]
(cid:32) v T
[Equation]
Uv(cid:112) d p , v T 2 Uv 1 (cid:112) d p (cid:33) v 1 + φ(cid:32) v T 1 Uv(cid:112) d p , v T 2 Uv 1 (cid:112) d p (cid:33) v(cid:35) .
Similarly, the ﬁrst column of f W ( X ) is f W ( X ) : , 1 = h (cid:88) i =
W io W iv (cid:34)
[Equation]
(cid:33)
[Equation]
(cid:35) .
Since U / (cid:112) d p − ( W 1 k ) T ( W 1 q ) /
(cid:112) d/h is skew-symmetric by assumption, we have v T
[Equation]
U √
[Equation]
v
[Equation]
= 0 for all v 1 .
Recall that U is rank- d p by assumption, so
U /
[Equation]
d
[Equation]
(cid:112) d/h is at least rank d p − d/h ≥ 1 , so we can choose any v 1 such that (cid:18)
[Equation]
If both U √ d p v 1 and ( W 1 k ) T ( W 1 q ) √ d/h v 1 are nonzero, We can always choose ˜ v 2 such that ˜ v T(cid:18)
U √ d p (cid:19) v 1 > 0 and ˜ v T(cid:18)
( W 1 k ) T ( W 1 q ) √ d/h (cid:19) v 1 < 0 .
This means that if we choose v 2 = α ˜ v 2 and scale α → ∞ , φ(cid:32) v T 1 Uv(cid:112) d p , v T 2 Uv 1 (cid:112) d p (cid:33) → 0 , φ(cid:32) v T 1 Uv(cid:112) d p , v T 2 Uv 1 (cid:112)
d p (cid:33) → 1 , φ(cid:32) v T 1 ( W 1 k ) T W 1 q v(cid:112)
d/h , v
[Equation]
( W 1 k )
[Equation]
W 1 q v(cid:112)
[Equation]
/h (cid:33) → exp( v
[Equation]
( W 1 k )
[Equation]
W 1 q v 1 / (cid:112)
[Equation]
/h )
exp( v T 1 ( W 1 k )
[Equation]
W 1 q v 1 / (cid:112) d/h ) + n − 2 , φ(cid:32) v T 1 ( W 1 k ) T W 1 q v(cid:112)
d/h , v T 2 ( W 1 k ) T W
[Equation]
v(cid:112) d/h (cid:33) → 0 .
Then, consider the difference f W ( X ) : , 1 − g V ( X ) : , 1 .
Recall that for any v , W 1 o W 1 v v is independent of { W io W iv v , i (cid:54) = 1 } .
This means that, to show f W ( X ) : , 1 − g V ( X ) : , 1 (cid:54) = 0 , it sufﬁces to show that
[Equation]
If we scale v 2 = α ˜ v 2 with large enough α , the second term will dominate the ﬁrst term and the ﬁrst term will never be able to cancel the second one.
Thus, by choosing large enough α > 0 , we can make sure that the sum is nonzero.
Even in case where one of U √ d p v 1 and ( W 1 k ) T ( W 1 q ) √
d/h v 1 is zero (say ( W 1 k ) T
( W 1 q )
√ d/h v 1 = 0 ), we can choose ˜ v 2 = U √ d p v 1 and use a similar scaling argument.
By choosing large enough α > 0 and v 2 = α ˜ v 2 , one can show that the difference f W ( X ) : , 1 − g V ( X ) : , 1 is nonzero.
C E XPERIMENTAL SETTINGS
For our experiments with the language modeling (LM1B dataset), we train 6 layer transformer models.
We use a batch size of 4096 and train for 250k steps.
We use a learning rate of 0.1 with a linear warm up for the ﬁrst 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 ﬁx the width of the feed forward layer in the Transformer to be 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 in (Devlin et al., 2018).
We present the key details here and refer the reader to (Devlin et al., 2018).
We train with a batch size of 1024 for 450k steps with inputs of sequence length n = 128 followed by 50k steps with inputs of sequence length 512.
In contrast the BERT paper uses a 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 a learning rate of 1e-4, and a linear warmup and decay schedule as in BERT.
We use 5k warmup steps for the ﬁrst stage, and a re-warmup of 3k steps for the second stage (You et al., 2019).
Again, we 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 ﬁrst stage and 64 TPUv3 chips for the second stage.
Pre-training with this conﬁguration takes between 2 to 3 days.
We did not attempt to ﬁnd the optimal hyper-parameters for the ﬁxed head size architecture, and use the same hyper-parameters as used for training the standard Transformer.
D A DDITIONAL EXPERIMENTAL RESULTS
[Table]