Transformer with a Mixture of Gaussian Keys
Authors: Tam Minh Nguyen
Tan Minh Nguyen
Dung Duy Le
Nguyen Duy Khuong
Viet Anh TRAN
Richard Baraniuk
Nhat Ho
Stanley Osher
Venue: ICLR 2022 Submitted (2022)
Keywords: transformer
gaussian mixture model
attention heads
attention keys

Abstract: Multi-head attention is a driving force behind state-of-the-art transformers which achieve remarkable performance across a variety of natural language processing (NLP) and computer vision tasks. It has been observed that for many applications, those attention heads learn redundant embedding, and most of them can be removed without degrading the performance of the model. Inspired by this observation, we propose Transformer with a Mixture of Gaussian Keys (Transformer-MGK), a novel transformer architecture that replaces redundant heads in transformers with a mixture of keys at each head. These mixtures of keys follow a Gaussian mixture model and allow each attention head to focus on different parts of the input sequence efficiently. Compared to its conventional transformer counterpart, Transformer-MGK accelerates training and inference, has fewer parameters, and requires fewer FLOPs to compute while achieving comparable or better accuracy across tasks. Transformer-MGK can also be easily extended to use with linear attention. We empirically demonstrate the advantage of Transformer-MGK in a range of practical applications including language modeling and tasks that involve very long sequences. On the Wikitext-103 and Long Range Arena benchmark, Transformer-MGKs with 4 heads attain comparable or better performance to the baseline transformers with 8 heads.

Under review as a conference paper at ICLR 2022
TRANSFORMER WITH A MIXTURE OF GAUSSIAN KEYS
Anonymous authors Paper under double-blind review
ABSTRACT
Multi-head attention is a driving force behind state-of-the-art transformers which achieve remarkable performance across variety of natural language processing (NLP) and computer vision tasks. It has been observed that for many applications, those attention heads learn redundant embedding, and most of them can be removed without degrading the performance of the model. Inspired by this observation, we propose Transformer with Mixture of Gaussian Keys (Transformer-MGin a novel transformer architecture that replaces redundant heads in transformers with mixture of keys at each head. These mixtures of keys follow a Gaussian mixture model and allow each attention head to focus on different parts the input sequence efficiently. Compared to its conventional transformer counterpart, Transformer- MGKK accelerates training and inference, has fewer parameters, and requires less FLOPs to compute while achieving comparable better accuracy across tasks. Transformer-MGK can also be easily extended to use with linear attentions. We empirically demonstrate the advantage of Transformer-MGK in range of practical applications including language modeling and tasks that involve very long sequences. On the Wikitext-103 and Long Range Arena benchmark, Transformer-Min MGKs with 4 heads attain comparable or better performance to the baseline transformers with 8 heads.
INTRODUCTION
Transformers (Vaswani et al., 2017) have become the state-of-the-art model for sequence processing tasks, solving many challenging problems in natural language processing and computer vision (Al- Rfou et al., 2019; Dai al., 2019; Williams al., 2018; Devlin al., 2018; Brown & al., 2020; Howard & Ruder, 2018 8; Rajpurkar et al., 2016; Dehghani et al., 2018; So et al., 2019; Dosovitskiy al., 2020; Touvron et al., 2020). These models can also transfer the learned knowledge from a pre-trained model to task that involves different data modalities and has limited supervision (Radford et al., 2018; 2019; Devlin et al., 2018; Yang et al., 2019; Liu et al., 2019). The success of transformers is rooted in the self-attention mechanism their fundamental building blocks for modeling (Cho et al., 2014; Parikh et al., 2016; Lin et al., 2017). For each token, self-attention computes a weighted average of the feature representations of other tokens where the weight is proportional to similarity score between each pair of tokens. This mechanism allows token to pay attention to other tokens in the sequence and attain contextual representation (Bahdanau et al., 2014; Vaswani et al., 2017; Kim et al., 2017). It has been shown that the representation capacity of the attention mechanism (Tenney al., 2019) and its capability of capturing diverse syntactic and semantic relationships (Tenney et al., 2019; Vig & Belinkov, 2019; Clark et al., 2019; Voita et al., 2019a; Hewitt & Liang, 2019) is key to the impressive performance of transformers in practice.
1.1 SELF-ATTENTION
For given input sequence X := [...... IN transforms X into the output sequence H in the following two steps:
NXD, of N feature vectors, self-attention
Step 1. The input sequence X is projected into the query matrix Q, the key matrix K, and the value matrix V via three linear transformations
Q XW Tj K K. XWT
XWI.
where W WK K E R DxD, and Wv E D.xD2 are the weight matrices. We denote Q := q
AN K = [k,.... kx] and == " where the vectors qus k, U, for 1, 1,... N are the query, key, and value vectors, respectively.

---

Under review as a conference paper at ICLR 2022
Step 2. The output sequence H := [h1:-"
hyl
is then computed as follows
H softmax
QK V
(I)
where the softmax function is applied to each row of the matrix (QK \/VD. For each query vector qi for i = 1,.......... . N, an equivalent form of Eqn. (1) to compute the output vector h, given by
N
h = softmax q k.kky kg/VDD Uj
Uz
The matrix A E RRXN and its component dij for i, 1, N the attention matrix and attention scores, respectively. The self-attention computed by Eqn. (1) and (2) is called the scaled dot-pre dot-product attention or softmax attention. In our paper, we call a transformer that uses this attention the softmax transformer. The structure that the attention matrix A learns from training determines the ability of the self-attention capture contextual representation for each token.
Multi-head Attention Each output sequence H forms an attention head. In multi-head attention, multiple heads are concatenated to compute the final output. Let H be the number of heads and W E IR xD,xxD, be the projection matrix for the output. The multi-head attention is defined as
MultiHead((Q) [K) [V) = Concat(H,, H H W
Even though multi-head attention extends single-head attention to capture diverse attention patterns and improve the performance of transformers, has been shown that transformers for practical tasks including sequence classification and language modeling learn redundant heads (Michel et al., 2019). These redundant heads compute similar attention mappings. Having many of them in the model limits the representation capacity of the transformer while wasting parameters, memory and computation, impeding the application of transformers to many important scale tasks.
1.2 CONTRIBUTION
We establish the correspondence between self-attention in transformer and a Gaussian mixture model (GMM) and propose Transformer with a Mixture of Gaussian Keys (Transformer- of transformers that can avoid the head redundancy. the attention key kj in each head by a GMM to allow the query qt. as well its associated token, to attend to more diverse positions in the input sequence, thereby increasing the representation of each attention head and reducing the chance of learning redundant heads. Our contribution is four-fold:
the core of Transformer-MGK novel class replacing
1. We construct a GMM and show that attention scores in self-attention match posterior distribu- tion in our model, providing probabilistic framework to study self-attention in transformers.
2. Under our probabilistic framework for self-attention, we introduce an additional mixture of Gaussian to model each attention key. We empirically show that this mixture of Gaussian keys (MGK) can capture a diversity of attention patterns, thus alleviating head redundancy.
3. We extend our MGK to use with linear attentions and propose the mixture of linear keys (MLK) for efficient computation and better memory footprint.
4. We empirically show that Transformer-MGK and Transformer-MLL are comparable or better than the corresponding baseline transformers with softmax and linear attentions while only using half the number of attention heads and reducing both model complexity measured by the number parameters and computational cost terms of FLOPs.
Organization: We structure this paper follows: In Section 2, we establish the connection be- tween GMM and self-attention and then present our Transformer-MGK and its extensions including Transformer-MLL In Section 3, we validate and empirically analyze the efficiency and accuracy of Mansfformer-GGKKMLL We discuss related works in Section 4. The paper ends up with concluding remarks. More experimental details are provided in the Appendix.
2
TRANSFORMER WITH A MIXTURE OF GAUSSIAN KEYS
2.1
ATTENTION SCORE AS A POSTERIOR DISTRIBUTION
We first consider a query 4: E Q and a key k, E K. Let be a N-dimensional binary random variable having a 1-of-N representation in which a particular element tj is equal and all other
2

---

Under review as a conference paper at ICLR 2022
elements equal to O. We use t; to indicate the position j of the key k,. In particular, let I be the identity matrix, we model the distribution p(q.) by the following GMM:
n Ny p(q.) Σ 1xplalt==))=>>>>>
(4)
j=1
where n is the prior p(t 1). Given the query qi, how likely qs matches the key ky is given by posterior p(tj = -l This posterior is computed as follows
TyN(qi 1ky,
Ty exp (- q1 1;,11// 22117/200)
p(t; = 1/q,) = Ey (q: kg, Ej Ty exp (-l1q - ky 12,3///20/
Ty exp [- (11q:11² + ||k; (12) /203] exp (q:k3 /03)
(5)
Σi Ty exp
[- (11::1'' 112 b
M (1²) /20 exp (a.k)//03))
We further assume that the query qi and the key ky are normalized, and the prior тÿ is uniform. We will justify these assumptions in our Remarks at the end of this section. We also let of o².
j = ..2...., K. Then the posterior p(t, = 1/q:) can be written as
plt; = 1/1:) = exp kg /0²)/ E exp (0/))
(6)
The right-hand side of Eqn. (6) matches the attention score given in Eqn. (2) when 0² vD. Thus, we show that under right assumptions, the attention score between the query qi and the key k, in an attention unit of a transformer the posterior p(t, 1/ which indicates the responsibility that the key ky takes for 'explaining' the query q, which in turn decide, for example, how much a token position pays attention token at position j in the input sequence.
Remark The assumption that the query qi and the key kj are normalized is realistic and not artificial. In many applications, those two vectors are normalized. Schlag et al. (2021) points out that such normalization is to avoid instability occurring during the training.
Remark 2 In practice, the prior is chosen to be uniform when there is prior knowledge available.
2.2 TRANSFORMER WITH A MIXTURE OF GAUSSIAN KEYS: EACH KEY IS AGAIN A GAUSSIAN MIXTURE MODEL
As we have seen from Eqn. (6), the key k; is used to explain the query q via the posterior P(t, 1 (g:)) Via this simple connection, each query gi is treated to be as a sample from the mixture of N keys Σ NnNV(gil [[j Ta). However, the assumption that the distribution P(gilt, = 1) at each subpopulation is Gaussian in Eqn. (4) can be quite strong as there is no guarantee that this assumption is valid in practice. In particular, it may happen that the distribution of each subpopulation is asymmetric or skewed or even multimodal. Therefore, using the Gaussian distribution for each subpopulation can potentially limit the explanation power and diversity of each subpopulation/ key. indicates that we should use more expressive distributions to represent P(g. L = 1).
Ta).
Mixture of Gaussian keys: the representation of each attention head, and reduce the chance of learning redundant heads, would like to model it as mixture of Gaussian distributions. We refer to this model as Transformer In particular, in Transformer-MC
To improve the explanation power of each key k, potentially increase
with Mixture of Gaussian Keys (Transformer-MGrr each key k, at position j as a mixture of M Gaussians ((krr T). = 1,2,..., M.. Here we are overloading the notation a little bit and use ky, and - to denote the mean and covariance matrix of the th Gaussian at position j. Let be a M-dimensional binary random variable having a 1-of-M representation. We use z, to indicate the th Gaussian in the mixture. Let Tjr = P(z, = 1|1 = 1), MGK can be written as
we model
P(q.lt, 1) Σ P(z, 1|t, = ()P z, 1, 1) Σ Nq: Kyrz oz
Our motivation of using mixture of Gaussian distributions to represent the distribution of each subpopulation in Eqn. (7) stems from the following important approximation result:
3

---

Under review as a conference paper at ICLR 2022
Theorem Assume that P is probability distribution on I-a, for some 0 and admits density function p such that p is differentiable and bounded. Then, for any given variance a > 0 and for any > O, there exists a mixture of K components E Ti== T,N(0,, 1² where K S (Cloo(//))) for some universal constant C such that
K Σ 0 Se,
sup Ip(x)
mERd
where ((x(0, 621) the density function of multivariate Gaussian distribution with meun and covariance matrix I
The proof of Theorem is in Appendix C. The result of Theorem suggests that regardless of use finite mixture of Gaussian distributions to approximate P(q.ltz = 1]. It allows us to have richer approximation of P'qlit = 1) than by using a simple Gaussian distribution in Eqn. (4). Similar the derivation above, the posterior p(t l/q.) in
the real form of P(gg(t = 1),
Transformer-MGK can be written as
P(C = 110:) Σ. exp (q.k)./03))
Tyr Tyr exp kk (1,3),/0))))
Furthermore, in Transformer-MGG we relax the assumption that the queries and keys are normalized. Thus, when computing P(t, l.q:), we compute the Gaussian kernels between the queries and keys instead of their dot products. The posterior P(t, 1/g)) in TraK then given by
E, orr exp (- 1-11q: k5r /2/03)
P(t = 1/q.)) = Eg Er Tyr exp (-11g: ky 10/111//20/
As proven in Section 2.1, this posterior corresponds to the attention score. Thus, Eqn. (9) is the formula for computing the attention score in Transformer-MGK. We compute the output vector h, of the self-attention in Transformer-MGK as follows
Σ, Thr exp (- 11q. 1o--kjn kgr 11/22//203
hi = Σ 3 Ej Ex Tyr Tyr exp
V/.
(10)
(-11q: k2//11//203.
2.3 INFERENCE AND LEARNING VIA THE EXPECTATION MAXIMIZATION ALGORITHM
Let Yrr P(z, 1/q,,t, 1). in MGK, apply the E-step inference in the Expectation- Maximization (EM) algorithm to estimate this posterior given the query q. The posterior Yrr is also known the responsibility that the component, N(k of I) takes to account for the observation, which in MGK is the query q: Below we propose two approaches to estimate this responsibility.
Soft E-step Using soft E-step inference, the EM algorithm makes soft assignment, in which each query is associated with all clusters. The responsibilities are then given by
Tir exp
(-1lg:
11g:
k 10.11/2200 /20
(11)
Σn Tjr exp
(-11g: kgr (0rl1/2/00 12/2033
At learning time, the responsibilities estimated by Eqn. (11) are used to update the prior Tjr, i.e.
Tyr NN,,/N, where N is the number of queries and Nyr- = >N -i=1 Mr. These updated priors Tyr are then used in Eqn. (9) to compute attention scores.
Hard E-step Hard E-step performs a hard assignment of queries to key clusters, in which each query is associated uniquely with one cluster. This is similar to the K-means algorithm (Lloyd, 1982) and corresponds the MGK at the limit when the variance parameter of, goes to O. Following the derivation of K -emaass from GMM in (Bishop, 2006), Eqn. (9) becomes
max, exp
(-llq: k5 )2,))
P(t, 1/q.) Eg max, exp (-1ly: k
(12)
4

---

Under review as a conference paper at ICLR 2022
The hard E-step inference allows the attention score to be computed more efficiently because the priors jr no longer play an active role in the algorithm and can be completely ignored. Learning via Stochastic Gradient Descent (SGD) In order increase the efficiency of the model, in MGK, we fix the variance parameter or to be vD as in the standard softmax attention and make the cluster means, i.. the keys, j,r learnable parameters. We also make the prior Tjr learnable parameters as one of the design options. In that case, both kj+ and 7rr are learned via SGD. This update via SGD can be considered as a generalized M-step (Bishop, 2006).
Remark 3
Design Options for Keys (Option A) We follow the standard setting in the softmax transformer and make the keys k linear projection of the input æj, i.c. kj. = 2, W K,, where I, R1xDs WK. K. E RDDxr and r = 1,2, 1,2..... ... M. (Option B) Alternatively, we also make the keys k shifted version of each other to save computation, i.c. kyr = 2, W K + by, where WK E DDD
W
2.4 TRANSFORMER WITH A MIXTURE OF LINEAR KEYS
The MGK can be easily extended to use with linear attentions. We call that model Transformer with a Mixture of Linear Keys (Transformer-MLKK) In this section, we adopt the formulation of linear attentions from (Katharopoulos et al., 2020) to derive Transformer-MLK Similar approach can be taken to derive Transformer-We -MLK when using with other linear attentions such as those in performers (Choromanski et al., 2021) and fast-weight transformers (Schlag et al., 2021). In Transformer-MLL the Gaussian kernel in Eqn. (10) is linearized as the product of feature maps Φ(-) on the vectors qi and ky. The associative property of matrix multiplication is then utilized to derive the following efficient computation of the attention map
h: = E,E, T (q:) To(k)))) ((q.) T TE,,,,,,,.. TT (k))))) ΣjE, mg) (.) (kkr))
(13)
o(q:) E2 T T (kjr)
Replacing Ej E, (T:)(a:) T )o, with ((q:) E, TT. JuJ as in linear transformers, reduces the memory and computational cost of computing the attention map in Transformer-Min
from O(N2) to O(N), making Transformer-MLLK scalable to very long sequences.
EXPERIMENTAL RESULTS
In this section, we numerically justify the efficiency of Transformer-MGK/MLL and empirically study the advantage of using mixture of keys on various benchmarks, including different tasks in the Long Range Arena (LRA) (Tay et al., 2021) (Section 3.1) and language modeling on Wikitext- 103 (Merity et al., 2017) (Section 3.2). We aim to show that: (i) HoansfommerMGGK/Mee heads comparable or better than the baseline softmax and linear transformers with the full number of heads while being more efficient in both computational cost and memory footprints; (ii) Mixture of keys helps reduce the redundancy in multi-head transformers and benefits learning of the long-term dependency in long input sequences; (iii) Using the same number of heads, Transformer-MGKKMMM significantly outperforms the baseline softmax and linear transformers. Especially in the case of Transformer-MH MLK,, it helps reduce the performance gap between softmax and linear transformers.
with half the number of
Throughout this section, we compare Transformer-MGK/MM formers that have the same or double the number of attention heads. In all experiments, for our Saitaffomee-MGGGAA
with the softmax and linear trans-
models, we set M=2 where M is the number of Gaussians, i.e. keys, at each timestep. Among the design options for Transformer- MGK mentioned in Section 2.3, we the one with Soft-E step but make the parameter jrr and kjr learnable and fix the variance of to be constants. We study both implementations for keys: (A) kjr is linear projection of the input Øj, i.e. kjr- = 2, WT. and (B) k are shifted version of each other, i.e., kgr = az W + by.
In this section, we refer to the Meaeffomer-MMGKMMLL and whose keys are implemented by (B) as Transformer-MMGKAMLLX We empirically compare these models with other design options for Transformer-MGE in Section 3.4. Details on datasets, models, and training are provided in Appendix A.1.
Seaeffrmer-MGGKMMLL whose keys are implemented by (A) as
3.1 LONG RANGE ARENA (LRA) BENCHMARK
Models and baselines We compare our 1-head, 2-head, 4-head Hennfforer-MGGK and MLK with the baseline softmax (Vaswani et al., 2017) and linear transformers (Katharopoulos et al., 2020) that have head, 2 heads, 4 heads, and 8 heads. Each model consists of two layers, and we adopt the model and training setting from (Xiong et al., 2021) in experiments.
5

---

Under review as a conference paper at ICLR 2022
Table Test Accuracy (%) of Transformer-Mer compared with the bascline softmax transformer on the LRA benchmark, Our Transform-MGHt outperform softmax transformers while using half the number of heads, having less parameters, and requiring less FI,OPs (see Figure 3 for details). Results averaged over 5 runs.
Model
ListOps
Text
Retrieval
Average
Softmax 12 heads
36.64
65.62
82.18
61.48
Softmax 8 heads
37.03
65.71
81.74
61.49
Transformer-sMGK 4 heads
37.25
65.51
82.79
61.85
Transformer-GGK 4 heads
36.98
65.69
82.23
61.63
Suftmax 4 heads
36.89
65.26
81.54
61.23
Transformer=MMGK3 2 heads
37.35
65.17
82.20
61.57
Transformer-Men MGK heads
36.88
65.37
81.83
61.36
Softmax 2 heads
36.76
64.90
79.1
60.25
Transformer-MMGK head
37.31
65.04
81.23
61.19
Transformer-MGH head
37.13
65.40
80.63
61.05
Softmax head
36.81
64.48
77.9
59.73
Table 2: Test Accuracy (%) of Transformer-MLK compared with the linear transformer on LRA. Our
Transform-Mal achieve comparable/better accuracy than the baselines while using half the number of heads, having less parameters, and requiring less FLOPs (see Figure 3 for details). Results are averaged over 5 runs.
Model
ListOps
Text
Retrieval
Average
Linear 12 heads
20.26
65.87
81.97
56.03
Linear heads
19.17
65.85
81.18
55.40
Transformer-sML. heads
20.11
65.74
81.53
55.79
Transformer-ML.. 4 heads
20.06
65.7
81.34
55.7
Linear 4 heads
19.37
65.81
81.65
55.61
Transformer-sMI. 2 heads
19.88
65.61
81,66
55.71
Transformer-MI.K heads
20.12
65.72
80.80
55.54
Linear heads
18.35
65.94
80,94
55.07
Transformer-sMI head
18.87
65.57
80.37
54.93
Transformer-ML. head
18.34
65.70
81.09
55.04
Linear head
18.60
65.70
80.6
54.96
Tenetemennnel head
Tam heads
Tre Trenffrmmeee... heat
Trneformeeosux X heads
50 60
Figure 1: Training loss and test accuracy of SwaissfonnenMMGGMMM vs. softmax/linear transformer on the retricval task, which has the longest average scquence-lenee length and attention span among the LRA tasks (Tay et al., 2021). The impressive performance of TransfoomerMGGK/MLKK on this challenging task validates the capability of our models capture long-rang rongrrage dependencies via learning diversity of attention patterns.
Results We summarize our results in Table 1. Transformer-MGH. with half the number of heads consistently achieve better test accuracy than the baseline softmax attention across tasks. Since fewer heads are needed, transformer- MGKKs use less parameters and need less FLOPS to compute than the baselines. We provide a detailed efficiency analysis for Transfommer-MGK. in Figure 3. More interestingly, these efficiency advantages of Transformer-MGG. over the baseline become more significant as the number of heads in the baseline model grows. When using the same number of heads as the baseline models, Transformer-MGi. further improve over those baselines. Among the models, Transformer-sMGK performs the best across LRA tasks.
We also compare the performance of Transformer-MLKK with the baseline linear transformers in Table 2. Like Monsformer-GGK. Transformer-ML.K yields comparable or better results than the baseline using only half the number of heads with less parameters and FLOPs. When using the same number of heads, Transformer-Me MLK helps improve the linear transformer further.
We provide results the 12-head baselines in Table and 2 for reference. It is interesting to notice from Table 1 and that even our 2-head Transformer-o -Metmer- models achieve better equivalent results to the 12-head and 8-head baselines. A comparison between the 12-head baselines with our 6-head Transformer MaMieerrMMKKMKKKK models on the retrieval task is provided in Table 8 in Appendix A.9.
In Figure 1, we compare the training loss and test accuracy curves of our 1-head and 2-head with the 2-head softmax and 2-head linear transformers the document

---

Under review as a conference paper at ICLR 2022
Table 3: Perplexity (PPL) on WikiText- 103 of Transformer MGK and MLK compared to the baselines. Both Transformer-MGHt and MLK achieve comparable or better PPL. than the baselines while using only half the number of heads, When using the same number of heads, our models significantly improve the haselines.
Method
Valid PPI.
Test PPI.
Softmax 8 heads (small)
33.15
34.29
Transformer-MGi heads (small)
33.28
34.21
Transformer-sMCG 8 heads (small)
32.92
33,99
TrKnKformer-MGK heads (small)
32.74
33,93
Softmax 4 heads (small)
34,80
35.85
Linear 8 heads (small)
38.07
39.08
Transformer- MLLK 4 heads (small)
38.49
39.45
Transformer-Me MLLK 8 heads (small)
37.78
38.99
Linear 4 heads (small)
39.32
40.17
Softmax 8 heads (medium)
27.90 (Schlag et al., 2021) 29.60 (Schlag et al., 2021)
TraK heads (medium)
27.58
28.86
retrieval task. This retrieval task has the longest average sequence-length and attention span among the LRA tasks (Tay et al., 2021). On this task, shown in Figure 1, our Seaaformer-MGKs/MM are always better than the baseline models throughout the training. This observation corroborates our models's 's capability of capturing long-range dependencies in very long input sequences.
3.2 LANGUAGE MODELING ON WIKITEXT-103
Next we confirm the advantage of our models on large-scale application. We consider the word-level language modeling task on WikiText- 103 (Merity et al., 2017) for our experiments this section.
Models and baselines We compare 4 and 8-head Transformeer-MGKs/MM... with 8-head soft- max (Vaswani al., 2017) and linear transformers (Katharopoulos et al., 2020). Each model consists of 16 layers. Our experiments follow the setting for small/medium models from (Schlag et al., 2021).
Results As shown in Table 3, our Transformer-MGK. outperform the baseline softmax transformers. Even when using half the number of attention heads (i.e., 4 vs. 8 heads as in the baselines), the Transformer-MGH. still achieves better test perplexities (PPL) than the baseline. Adding more heads into Transformer-MGE. improves their performance. Similarly, Transformer-MLKK attain comparable test/validation PPL to the baseline linear transformers when using half the number of attention heads. When using the same number of attention heads in the baseline, Transformer- MLKs consistently achieve better performance. Note that reducing the number of heads from 8 to 4 in the baseline models significantly decreases their performance with more than 1.5 reduction in test/validation PPL for the softmax transformer and more than 1.0 reduction in test/validation PPL for the linear transformer. Our proposed Transformer-MGE and Transformer-MLL helps close this gap.
To further examine the scalability of our models, we apply the MGK on a stronger baseline, which is the 8-head medium softmax transformer (Schlag et al., 2021). This model has 90M parameters, 16 layers, 8 attention heads per layer, and hidden size of 256. The size of our baseline model is close to BERT Base (Devlin al., 2019), which has 110M parameters, 12 layers, 12 attention heads per layer, and hidden size of 768. Applying our MGK on top of this baseline and using only 4 heads instead of 8, we significantly improve the test PPL from 29.60 9.60 to 28.86 while reducing the model size and computational cost, demonstrating the advantages of our scaled models.
3.3 NEURAL MACHINE TRANSLATION ON IWSLT' 14 GERMAN TO ENGLISH
We further examine the advantages of our methods the IWSLT' 14 German-English machine translation task (Cettolo et al., 2014), Table 9 in Appendix A. shows that the 2-head Transformer- MGK and sMGK models achieve the BLEU score of 34.34 and 34.69, respectively, which are comparable to and better than the BLUE score of 34.42 of the 4-head softmax transformer baseline.
3.4 EMPIRICAL ANALYSIS
We conduct empirical analysis based on the Transformer-MGK trained for the document retrieval tasks. Results for Transformer-MLKK and the WikiTeex- 103 task provided in the Appendix.
Transformer-MGH helps avoid learning redundant heads We visually compare attention matrices leamned by Transformer-GGe and the baseline softmax transformer the document retrieval task in Figure 2. In particular, we randomly select an attention matrix at each head in each layer and visualize that attention matrix for each model in comparison. Figure 2(Left) shows that the queries
7

---

Under review as a conference paper at ICLR 2022
PSC Feac
3000
Figure 2: (Left) Visualization of attention matrices in the baseline 4-head softmax transformer (left), 4-head Transformer- MGK (middle). and 2-head Transformer-M MGGK (right) trained on the document retrieval task. Attention matrices from our Transforme--GGH MGKKs have more diverse pattern than those from the baseline softmax transformer, reducing the risk of learning redundant heads. (Right) Rank distribution of attention matrices shows that attention matrices in Transformer-MGi have higher rank than those in the softmax transformer and thus can capture more diverse attention patterns.
in Transformer-MGKS can attend to variety of keys and equivalently to other tokens at different positions in the input sequence. This diversity in attention pattern helps reduce the chance that the model learns similar and redundant attention matrices at different heads significantly.
Another metric to measure the representation capacity of an attention matrix is its rank. Attention matrices with high ranks can capture more diverse attention patterns compared to those with low ranks (Nguyen et al., 2021). We study the rank of the attention matrix from the Transformer-MGK and the softmax transformer trained for the retrieval task. In particular, we randomly select 1000 different attention matrices at each layer from each model. Then, we perform singular value decomposition (SVD) compute the rank of each matrix and threshold the singular values smaller than 10 6 Figure 2(Right) presents the distribution of the rank of attention matrices at each layer of the Transformer-MGGi MGK and the softmax transformer. We observe that attention matrices in Transformer- MGK has higher rank than those in the softmax transformer. Thus, our attention with MGK is capable of capturing more diverse and complex attention patterns than the baseline softmax attention.
Transformer-MGKKML. reduces model complexity and computational cost Figure 3 compares the computational cost, measured in FLOPS, and model complexity, measured in the number of parameters, between our Maunsforme--GG//MM softmax/linear transformer. The more heads being used, the more advantage Transfommer-MGKAMM has over the softmax/linear transformer. For much larger transformer models, this saving is significant.
that has half the number of heads and the full-head
Retrieall
Hamber Tranaterme---i XVGX Heands
Listnga
Text
Retrovall
Number Thansformann MLX Huade
Lirlops
Text
FLOPG
1.28
1.40
Parnndors (ee))
Aepamettnn (160)
Number Softree Thenebyrrer
Namberr Transormer
Softmac Tonsformer
Tonsformer -- -Tanslommrrte
Figure 3: Computational cost (FLOPs) and the number of parameters of Transformer-MGG vs the baseline softmax transformer (Left) and Transformer-Mii MLLK the baseline linear transformer (Right). The efficiency advantage of Transformer Transformr-MGKKMLLK MGK/MLK over the baselines both metries grows with the number of head.
Comparing different inference and learning techniques Table 4 compares the performance of Transformer-MGGK using different design options mentioned in Section 2.3 on the LRA benchmark. In particular, we consider the following three design options: A) Soft-E step, parameters Tjr and kjr are learnable via SGD, and variance dz. are constants, B) Soft-E step, parameter Tjr is updated according to the M-step update, ky, are learnable via SGD, and variance o) are constants, and C) Hard-E step, jjr and j are learnable via SGD, and variance 03, are constants. Note that Sinnsforer-MGK. with setting A are the default models use in all experiments above. In Table 4, Transformer-MCC MGK + Hard-E is the Transformer- MGK with setting C, Transformer-MGG the Transformer-MGG with setting B, and -GK only is the Transformer-MGK with
Soft-E is
8

---

Under review as a conference paper at ICLR 2022
Table 4: Performance of Transformer- MGK using different inference//carming techniques on LRA benchmark.
Model
ListOps
Text
Retrieval
Average
Transformer::MGG Hard-E head
37.25
64.7
81.29
61.08
Transformer-sMGGK Soft-E 1 head
37.05
64.68
81.44
61.05
Transformer-sMGG head
37.31
65.04
81.23
61.19
Transtormer- Hard-E I head
19.40
65,40
80.72
55.17
Transformer- MGG Soft-E 1 head
33.85
65.25
80.73
59,94
Transformer-Mor MGK head
37.13
65.40
80.63
61.05
setting A. It is worth noting that Transformer-s + Hard-E obtains comparable results to the models with the best performance in each task even though it is the most efficient model in our study,
RELATED WORK
Efficient Transformers Efficient transformers can be classified into several categories, as summa- rized in (Roy et al., 2021). Among these categories are models which design the attention matrix to have sparse structure (Parmar et al., 2018; Liu et al., 2018; Qiu et al., 2019; Child et al., 2019; Beltagy et al., 2020). Another category includes models that combine two or more different access patterns to improve the coverage (Child et al., 2019; Ho et al., 2019). Access patterss can also be made learnable in a data-driven fashion (Kitaev et al., 2020; Roy et al., 2021; Tay et al., 2020). Other efficient transformers take advantage of a side memory module to access multiple tokens at once (Lee et al., 2019; Sukhbaatar et al., 2019; Asai & Choi, 2020; Beltagy et al., 2020). Finally, low-rank and kernel- ization approximation are utilized to enhance the memory and computational efficiency of computing self-attention (Tsai et al., 2019; Wang et al.. 2020: Katharopoulos et al., 2020; Choromanski et al., 2021; Shen et al., 2021; Nguyen et al., 2021; Peng et al., 2021). In addition to the aforementioned efficient transformers, mulli--uery attention that shares keys and values between different attention heads (Shazeer, 2019) has also been studied to reduce the memory-bandwide speed for incremental transformer inference (see Appendix AII Last but not least, methods such as using auxiliary losses (Al-Rfou et al., 2019) and adaptive input embedding (Baevski & Auli, 2019) have been explored to speed up the convergence of training transformers. Our MGK/MLK can be easily incorporated into the methods above to further improve their accuracy and efficiency.
cost and increase the
Redundancy in Transformers Latest works suggest that most of the neurons and heads in the pre-trained transformer are redundant and can be removed when optimzing towards a downstream task (Dalvi al., 2020; Michel al., 2019; Durrani et al., 2020). Other works also study the contextualized embeddings in pretrained networks under this redundancy due to overparameterization and show that the representations learned within these models are highly anisotropic (Mu & Viswanath, 2018; Ethayarajh, 2019). An emerging body of work is proposed to distill and prune the model, including (Sanh al., 2019; Sun al., 2019; Voita al., 2019b: Sajjad et al., 2020). Our MGK/MLK approch can be combined with these distilling and pruning methods to improve their performance.
Mixture Models for Transformers Several works have used mixture models study and enhance transformers. Switch transformers (Fedus et al., 2021) employ the routing algorithm in Mixture of Experts (MoE) to reduce the communication and computational costs in transformers. (Nguyen et al., 2018; Patel et al., 2016) derive probablistic framework based on GMMs for deep neural networks that can be extended to study transformers and attention- based architectures. Other works that use mixture models with transformers include (Cho et al., 2020; Guo et al., 2019; Jiang et al., 2020).
5 CONCLUDING REMARKS
In this paper, we proposed Transformer-MGr a class of transformers that use Gaussian mixture model to represent the key vectors in self-attention. Transformer-MGK reduces the redundancy among heads in transformer. Furthermore, attention heads in the TrKKsformer-MGK have better representation capability than those in the baseline, allowing the Transformer-MGG to achieve comparable or better performance than the baseline softmax transformer while using only half of the number of heads. Comparing to the baseline, the Transformer-MG. uses fewer parameters and requires less FL.PPs to compute. We extend the TrKK into the Transformer- MLLK to use linear attentions for better efficiency. We empirically validate the advantage of the Transformer- MGK/MLK over the baseline softmax and linear transformers on various benchmarks including tasks in the LRA benchmark, WikiText-103 language modeling, and IWSLT' 14 machine translation. In our work, we make the means and the variance of the cluster learnable variables and constants, respectively. It is interesting to explore how to leverage the M-step update in the EM algorithm to update those parameters. Furthermore, we leave the application of Trunsfommer-GGK/MN improving the vision transformer (Dosovitskiy et al., 2020; Touvron et al., 2020) future work.
for

---

Under review as a conference paper at ICLR 2022
Reproducibility Statement: Source codes for experiments are provided in the supplementary materials of the paper. The details of our experimental settings and computational infrastructure are given in Section 3 and the Appendix. All datasets that we used in the paper published, and they are easy to find in the Internet.
Ethics Statement: Given the nature of the work, do not foresee any negative societal and ethical impacts of our work.
REFERENCES
Rami Al-Rfou, Dokook Choe, Noah Constant, Mandy Guo, and Llion Jones. Character-level language modeling with deeper self-attention. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 33, pp. 3159-3166 2019.
Akari Asai and Eunsol Choi. Challenges in information seeking qa: Unanswerable questions and paragraph retrieval. arXiv preprint arXiv:2010.11911 2020.
A. G. Bacharoglou. Approximation of probability distributions by convex mixtures of Gaussian measures. Proceedings of the American Mathematical Society, 138:2619-2628, 2010.
Alexei Baevski and Michael Auli. Adaptive input representations for neural language modeling. In International Conference on Learning Representations, 2019. URL https///peemeevie.ee//ffuuu id-ByxZX20qFQ.
Dzmitry Bahdanau, Kyunghyun Cho, and Yoshua Bengio. Neural machine translation by jointly learning align and translate. arXiv preprini 1rXiv::409.04 2014.
Iz Beltagy, Matthew E Peters, and Arman Cohan. Longformer: The long-document transformer.
arXiv preprint arXiv:2004.0511 050, 2020.
Christopher M Bishop. Pattern recognition. Machine learning, 128(9), 2006.
Tom Brown and et al. Language models are few-shot learners. In H. Larochelle, M. Ranzato, R. Hadsell, M. Balcan, and H. Lin (eds.), Advances in Neural Information Processing Sys- tems, volume 33, pp. 1877-1901, 2020. URL
1457c0666fbb9974488Bb883 142f64a Pape.aa-a.
Mauro Cettolo, Jan Nichues, Sebastian Stüker, Luisa Bentivogli, and Marcello Federico. Report on the 11th iwslt evaluation campaign, iwslt 2014. In Proceedings of the International Workshop on Spoken Language Translation, Hanoi, Vietnam, volume 57, 2014.
Rewon Child, Scott Gray, Alec Radford, and Ilya Sutskever. Generating long sequences with sparse transformers. arXiv preprint PxXiv:/904. 10509, 2019.
Kyunghyun Cho, Bart van Merrienboer, Caglar Gulcehre, Dzmitry Bahdanau, Fethi Bougares, Holger Schwenk, and Yoshua Bengio. Learning phrase representations using RNN encoder- decoder for statistical machine translation. In Proceedings of the 2014 Conference on Empirical Methods in Natural Language Processing (EMNLP), pp. 1724-1734 Doha, Qatar, October 2014. Association for Computational Linguistics. doi: 10.3115/v1/D14-1177 URL https://www.aclweb.
org/anthology/DI4- 1179.
Sung Min Cho, Eunhyeok Park, and Sungjoo Yoo. Meantime: Mixture attention mechanisms with multi-temporal embeddings for sequential recommendation. In Fourteenth ACM Conference on Recommender Systems, pp. 515-520, 2020.
Krzysztof Marcin Choromanski, Valerii Likhosherstov, David Dohan, Xingyou Song, Andreea Gane, Tamas Sarlos, Peter Hawkins, Jared Quincy Davis, Afroz Mohiuddin, Lukasz Kaiser, David Benjamin Belanger, Lucy J Colwell, and Adrian Weller. Rethinking attention with performers. In International Conference on Learning Representations, 2021. URL hitps: wttprr/wwwee.ttttommmmmmdouuaaaa
Kevin Clark, Urvashi Khandelwal, Omer Levy, and Christopher D. Manning, What does BERT look at? an analysis of BERT"s attention. In Proceedings of the 2019 ACL Workshop BlackboxNLP: Analyzing and Interpreting Neural Networks for NLP, pp. 276-286, Florence, Italy, August 2019. Association for Computational Linguistics. doi: 10.18653/v1/W194440 //ww..dlwe..rgaathllgg/111..
URL
https:
10

---

Under review as a conference paper at ICLR 2022
Zihang Dai, Zhilin Yang, Yiming Yang, Jaime Carbonell, Quoc V Le, and Ruslan Salakhutdi- nov. Transformer-x:: Attentive language models beyond a fixed-length context. arXiv preprint arXiv:1901.02860, 2019.
Fahim Dalvi, Hassan Sajjad, Nadir Durrani, and Yonatan Belinkov. Analyzing redundancy in pretrained transformer models. arXiv preprint (rxXxx:2004.400000 2020.
Mostafa Dehghani, Stephan Gouws, Oriol Vinyals, Jakob Uszkoreit, and Lukasz Kaiser. Universal transformers. arXiv preprint FrxXi::807.038899 2018.
Jacob Devlin, Ming-Wer Chang, Kenton Lee, and Kristina Toutanova. Bert: Pre-training of deep bidirectional transformers for language understanding. arXiv preprint arXhv:1810.0480) 2018.
Jacob Devlin, Ming- Wei Chang, Kenton Lee, and Kristina Toutanova. BERT: Pre-training of deep bidirectional transformers for language understanding. In Proceedings of the 2019 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, Volune I (Long and Short Papers), pp. 4171-4186, Minncapolis, Minnesota, June 2019. Association for Computational Linguistics. doi: 10. 10.18653//1/N19- 1423. V/cclanhologyyorrgNN33 1422444433
URL https:
Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Unterthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, et al. An image is worth 16x16 words: Transformers for image recognition at scale. arXiv preprins arXiv:2010.11992 11929, 2020.
Nadir Durrani, Hassan Sajjad, Fahim Dalvi, and Yonatan Belinkov. Analyzing individual neurons in pre-trained language models. arXiv preprint arXiv:2010.02695, 2020.
Kawin Ethayarajb. How contextual are contextualized word representations? comparing the geometry of bert, elmo, and gpt-2 embeddings. arXiv preprint /9X.:: 2019.
William Fedus, Barret Zoph, and Noam Shazeer. Switch transformers: Scaling to trillion parameter models with simple and efficient sparsity. arXiv preprint arXiv:2/01.03961 2021.
S. Ghosal and A. van der Vaart. Entropics and rates of convergence for maximum likelihood and bayes estimation for mixtures of normal densities. Ann. Statist., 29: 1033-122-,,, 2001.
Maosheng Guo, Yu Zhang, and Ting Liu. Gaussian transformer: a lightweight approach for natural language inference. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 33, pp. 6489-6496, 2019.
Stephen José Hanson. A stochastic version of the delta rule. Physica D: Nonlinear Phenomena, 42 (1-3):265-272, 1990.
John Hewitt and Percy Liang. Designing and interpreting probes with control tasks. In Proceedings of the 2019 Conference on Empirical Methods in Natural Language Processing and the 9th International Joint Conference Natural Language Processing (EMNLP-IJCNLP), pp. 2733- 2743, Hong Kong, China, November 2019. Association for Computational Linguistics. doi: 10.18653/v 11116533/1/DD9-... 1275-1275. URL https:/wwww..clbbbbarr//tttbtb/p 1275.
Jonathan Ho, Nal Kalchbrenner, Dirk Weissenborn, and Tim Salimans. Axial attention in multidi- mensional transformers. arXiv preprint arXiv: 2019.
Jeremy Howard and Sebastian Ruder. Universal language model fine-tuning for text classification. In Proceedings of the 56th Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers), pp. 228-339, Melbourne, Australia, July 2018. Association for Computational Linguistics. doi: 10. 1/t.//////1//1.. 1
URL
http://wwwwwwww.ww./////toooooot 1031
Junyan Jiang, Gus G Xia, Dave B Carlton, Chris N Anderson, and Ryan H Miyakawa. Transformer vae: A hierarchical model for structure-aware and interpretable music representation learning. In ICASSP 2020-2020 IEEE International Conference on Acoustics, Speech and Signal Processing (ICASSP), pp. 516-520. IEFE, 2020.
11

---

Under review as a conference paper at ICLR 2022
Angelos Katharopoulos, Apoorv Vyas, Nikolaos Pappas, and François Fleuret. Transformers are mns:: Fast autoregressive transformers with linear attention. In International Conference on Machine Learning, pp. 5156-51655 PMLR, 2020.
Yoon Kim, Carl Denton, Luong Hoang, and Alexander M Rush. Structured attention networks. arXiv preprint orXiv:1702.00885 2017.
Diederik P Kingma and Jimmy Ba. Adam: A method for stochastic optimization. arXiv preprint arXiv: 1412.6980 2014.
Nikita Kitaev, Lukasz Kaiser, and Anselm Levskaya. Reformer: The efficient transformer. arXiv preprint arXiv:2001.0445/ 2020.
Juho Lee, Yoonho Lee, Jungtack Kim, Adam Kosiorek, Seungjin Choi, and Yee Whye Teh. Set trans- former: A framework for attention- based permutation-invariaa neural networks. In International Conference on Machine Learning, pp. 3744-3753 PMLR, 2019.
Zhouhan Lin, Minwei Feng, Cícero Nogueira dos Santos, Mo Yu, Bing Xiang, Bowen Zhou, and Yoshua Bengio. A structured self-attentive sentence embedding. CoRR, abs/1703.03130, 2017.
URL http://arxiv.org/abb// wttp://wwwwr.abb///.....33 1703.03130.
Peter Liu, Mohammad Saleh, Etienne Pot, Ben Goodrich, Ryan Sepassi, Lukasz Kaiser, and Noam Shazeer. Generating wikipedia by summarizing long sequences. arXiv preprint arXiv: 2018.
Yinhan Liu, Myle Ott, Naman Goyal, Jingfei Du, Mandar Joshi, Danqi Chen, Omer Levy, Mike Lewis, Luke Zettlemoyer, and Veselin Stoyanov. Roberta: A robustly optimized bert pretraining approach. arXiv preprint arXiv: 1907. 1692, 2019.
Stuart Lloyd. Least squares quantization in pem. IEEE transactions on information theory, 28(2): 139-137, 1982.
Andrew L. Maas, Raymond E. Daly, Peter T. Pham, Dan Huang, Andrew Y. Ng, and Christopher Potts. Learning word vectors for sentiment analysis. In Proceedings of the 49th Annual Meeting of the Association for Computational Linguisties: Human Language Technologies, pp. 142-150, Portland, Oregon, USA, June 2011. Association for Computational Linguistics. URL hitps: //ww.adwee.og//attooog///Pl
Stephen Merity, Caiming Xiong, James Bradbury, and Richard Socher. Pointer sentinel mixture models. In 5th International Conference on Learning Representations, ICLR 2017, Toulon, France, April 24-26, 2017, Conference Truck Proceedings. OpenReview.net, 2017. URL https: Vlopenneview.nee/fooum htter/oooummmmdByyooee Bidddyyyyyxxee
Paul Michel, Omer Levy. and Graham Neubig. Are sixteen heads really better than one? In H. Wallach, H. Larochelle, A. Beygelzimer. F. d''lché-Buee E. Fox, and R. Garett (eds.), Advances in Neural Information Processing Systems, volume 32. Curran Associates, Inc., 2019. URL
httpss//cccedinns.eerrrsccpppppp wttp://www......ooooooooooooooee w0227996bbb2200/0///40000------0
Jiaqi Mu and Pramod Viswanath. All-but-the-top: Simple and effective postprocessing for word
representations. In International Conference on Learning Representations, 2018. URL https:
Nikita Nangia and Samuel Bowman. ListOps: A diagnostic dataset for latent tree learning. In
Proceedings of the 2018 Conference of the North American Chapter of the Association for Computational Linguistics: Student Research Workshop, pp. 92-99. New Orleans, Louisiana, USA, June 2018. Association for Computational Linguistics. doi: 10.18653/1/N118-40133 URL
http::/www..eleeee.gggnntgggg//
Tan Nguyen, Nhat Ho, Ankit Patel, Anima Anandkumar, Michael Jordan, and Richard G Baraniuk. A Bayesian perspective of convolutional neural networks through deconvolutional generative model, arXiv preprint arXiv: /8/ 1./8/1101 02265, 2018.
12

---

Under review as a conference paper at ICLR 2022
Tan M. Nguyen, Vai Suliafu, Stanley J. Osher, Long Chen, and Bao Wang. Fmmformer: Effi- cient and flexible transformer via decomposed near-field and far-field attention. arXiv preprint 20Xx::2008.02347,, 2021.
Myle Ott, Sergey Edunov, Alexei Baevski, Angela Fan, Sam Gross, Nathan Ng, David Grangier, and Michael Auli. fairseq: A fast, extensible toolkit for sequence modeling. In Proceedings of NAACL-IILT 2019: Demonstrations, 2019.
Ankur Parikh, Oscar Täckström, Dipanjan Das, and Jakob Uszkoreit. A decomposable attention model for natural language inference. In Proceedings of the 2016 Conference on Empirical Methods in Natural Language Processing, pp. 1249-2255 Austin, Texas, November 2016. Association for Computational Linguistics. doi: 10. (0.18653/v1DD16 1244. URL http://wwwwwwwww...///n//nbo////
D16- 1244.
Niki Parmar, Ashish Vaswani, Jakob Uszkoreit, Lukasz Kaiser, Noam Shazeer, Alexander Ku, and Dustin Tran. Image transformer. In Jennifer Dy and Andreas Krause (eds.), P Proceedings of the 35th International Conference on Muchine Learning, volume 80 Proceedings of Machine Learning Research, pp. 4055-4064. PMLR, 10-15 Jul 2018. URL http://proceedings..les wwl.piigg.......s..............s v80/parmar 8a.html.
Ankit B Patel, Minh T Nguyen, and Richard Baraniuk. A probabilistic framework for deep learning. Advances in neural information processing systems, 29:2558 02-2558-2266, 2016.
Hao Peng, Nikolaos Pappas, Dani Yogatama, Roy Schwartz, Noah Smith, and Lingpeng Kong. Random feature attention. In International Conference on Learning Representations, 2021. URL http:://pwweweeeeeee/////odddddd wittoornnnnniittooooooooota
Jiezhong Qiu, Hao Ma. Omer Levy, Scott Wen-tau Yih, Sinong Wang, and Jie Tang. Blockwise self-attention for long document understanding, arXiv preprint 1xXiv:9911.02922 2019.
Dragomir R Radev, Pradeep Muthukrishnan, Vahed Qazvinian, and Amjad Abu-Jbara. The acl anthology network corpus. Language Resources and Evaluation, 47(4):919-944 2013.
Alec Radford, Karthik Narasimhan, Tim Salimans, and Ilya Sutskever. Improving language under- standing by generative pre-training, OpenAl report, 2018.
Alec Radford, Jeffrey Wu, Rewon Child, David Luan, Dario Amodei, and Ilya Sutskever. Language models are unsupervised multitask learners. OpenAl blog, 1(8):9, 2019.
Pranav Rajpurkar, Jian Zhang, Konstantin Lopyrev, and Percy Liang. SQuAD: 100,000+ questions for machine comprehension of text. In Proceedings of the 2016 Conference Empirical Methods in Natural Language Processing, pp. 2383-2392 Austin, Texas, November 2016. Association for Computational Linguistics. doi: 10.18653/11/D16- 1264. URL https://www..llle..og/aabbbobbo. D16- 1264.
Aurko Roy, Mohammad Saffar, Ashish Vaswani, and David Grangier. Efficient content-based sparse attention with routing transformers. Transactions of the Association for Computational Linguistics, 9:53-68, 2021. doi: 16.1/aa/ a 00353. URL
1.4.
Hassan Sajjad. Fahim Dalvi, Nadir Durrani, and Preslav Nakov. Poor man's bert: Smaller and faster transformer models. arXiv e-prints, pp. arXiv-2004, 2020.
Victor Sanh, Lysandre Debut, Julien Chaumond, and Thomas Wolf. Distilbert, a distilled version of bert: smaller, faster, cheaper and lighter. arXiv preprini 1910v:1910.01100 2019.
Imanol Schlag, Kazuki Irie, and Jürgen Schmidhuber. Linear transformers are secretly fast weight programmers. In International Conference on Machine Learning, pp. 9355-9366. PMLR, 2021.
Noam Shazeer. Fast transformer decoding: One write-head is all you need. arXiv preprint arXiv:1911.02150 2019.
Zhuoran Shen, Mingyuan Zhang, Haiyu Zhao, Shuai Yi, and Hongsheng Li. Efficient attention: Attention with linear complexities. In Proceedings of the IEEE/CVF Winter Conference on Applications of Computer Vision, pp. pp. 351133333333 2021.
13

---

Under review as a conference paper at ICLR 2022
David R So, Chen Liang, and Quoc V Le. The evolved transformer. arXiv preprint arXiv: 2019.
Nitish Srivastava, Geoffrey Hinton, Alex Krizhevsky, Ilya Sutskever, and Ruslan Salakhutdinov. Dropout: simple way prevent neural networks from overfitting The journal of machine learning research, ((()):199-1955) 2014.
Sainbayar Sukhbaatar, Edouard Grave, Guillaume Lample, Herve Jegou, and Armand Joulin. Aug- menting self-attention with persistent memory. arXiv preprint 1OXivv1907.0144 2019.
Siqi Sun, Yu Cheng, Zhe Gan, and Jingjing Liu. Patient knowledge distillation for bert model compression. arXiv preprint arXiv:/088.08 09355, 2019.
Yi Tay, Dara Bahri, Liu Yang, Donald Metzler, and Da-Cheng Juan. Sparse Sinkhorn attention. In Hal Daumé III and Aarti Singh (eds.), Proceedings of the 37th International Conference on Machine Learning, volume 119 of Proceedings of Machine Learning Research, pp. 9438-9447. PMLR, 13-18 Jul 2020. URL
wtwp:reeeeedingg....rr..... Salvv119/tta 19/tay0a.html..
Yi Tay, Mostafa Dehghani, Samira Abnar, Yikang Shen, Dara Bahri, Philip Pham, Jinfeng Rao, Liu Yang, Sebastian Ruder, and Donald Metzler. Long range arena : A benchmark for efficient transformers. In International Conference Learning Representations, 2021. URL https: /lopeneeieew.nefforummm didiuummiiiiiiiia
-rCkk.
lan Tenney, Dipanjan Das, and Ellie Pavlick. BERT rediscovers the classical NLP pipeline. In Proceedings of the 57th Annual Meeting of the Association for Computational Linguistics, pp. 4593-4601, Florence, Italy, July 2019. Association for Computational Linguistics. doi: 10.18653/ v1/P19- 1452. URL http::/www..alllw....////na/////
1452.
Hugo Touvron, Matthieu Cord, Matthijs Douze, Francisco Massa, Alexandre Sablayrolles, and Hervé Jégou. Training data-efficient image transformers & distillation through attention. arXiv preprint arXiv:2012.1287,, 2020.
Yao-Hung Hubert Tsai, Shaojie Bai, Makoto Yamada, Louis-Philippe Morency, and Ruslan Salakhut- dinov. Transformer dissection: An unified understanding for transformer's attention via the lens of kernel. arXiv preprint arXiv: 11775, 2019.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention all you need. In Advances neural information processing systems, pp. 5998-6008, 2017.
Jesse Vig and Yonatan Belinkov. Analyzing the structure of attention in transformer language model. In Proceedings of the 2019 ACL Workshop BlackboxNLP: Analyzing and Interpreting Neural Networks for NLP. pp. 63-76, Florence, Italy, August 2019. Association for Computational 19-4808.
Linguistics. doi: 10. N0.18653/v1// 1/W1/ 488888...
URL
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. In Proceedings of the 57th Annual Meeting of the Association for Computational Linguistics, pp. 5797-5808, Florence, Italy, July 2019a. Association for Computational Linguistics. doi: 10.1 16535v11P19-11555/ URI.
hwtp://wwwww.gggaoooo/////////// 1580.
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: 2019b.
Sinong Wang, Belinda Li, Madian Khabsa, Han Fang. and Hao Ma. Linformer: Self-attention with linear complexity. arXiv preprint arXiv:2006.04700 2020.
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, June 2018. doi: 10.18653/v1/N18 18- 1101. URL http::/www...lvveebor///ttlggg//
14

---

Under review as a conference paper at ICLR 2022
Yunyang Xiong, Zhanpeng Zeng, Rudrasis Chakraborty, Mingxing Tan, Glenn Fung, Yin Li, and Vikas Singh. Nystrómformer: A Nysttömmbasse Algorithm for Approximating Self-Attentio.. 2021.
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.08237, 2019.
15

---

Under review as a conference paper at ICLR 2022
Supplement to "Transformer with a Mixture of Gaussian Keys"
In this supplementary material, we provide experimental details and additional experiments of Transformer-MGis and Transformer-MLKK
A
ADDITIONAL EXPERIMENTS
A.1
EXPERIMENT DETAILS
In this section, we provide model and training details for experiments in Section 3. All our experiments conducted on a server with 4 NVIDIA A100 GPUs.
A.1.1 LONG RANGE ARENA BENCHMARK
Datasets and metries We consider the following tasks in the LRA benchmark: Listops (Nangia & Bowman, 2018), byte-level IMDb reviews text classification (Maas et al., 201 1), and byte-level document retrieval (Radev al., 2013). These tasks involve long sequences of length 2K, 4K, and 4K, respectielly. We follow the setup/evaluation protocol in (Tay et al., 2021) and report the test accuracy for individual task and the average result across all tasks.
Models and baselines We use the softmax transformer (Vaswani et al., 2017) and linear trans- former (Katharopoulos et al., 2020) as our baselines. All models have layers, 64 embedding dimension, and 128 hidden dimension. The number of heads in each layer are set to 1, 2, 4, and 8. For Beeeffomer-MGG/MLL..
and their shifted versions, we share ār for all position j and learn for each head. The initial value for each Tjr is set to 0.5. For Transformer-MMGG. we learn b, and initialize its elements from standard normal distribution. Each jr is a constant with value vD where is the dimension of each head, which is the same as in the baselines models
Details about the Long Range Arena (LRA) benchmarks can be found in the original pa- per (Tay et al., 2021). Our implementation is based on the public code by (Xiong et al., 2021), and follow their training procedures. The training setting and additional baseline model details are provided in the configuration file used in (Xiong et al., 2021) and available www.////wppppaoyytooooeeer/bbbbe RA/cc....
A.1.2 LANGUAGE MODELING ON WIKITEXT-100
Datasets and metrics WikiText- 103 consists of articles from Wikipedia and is dataset with long contextual dependencies. The training set is made up of about 28K articles containing 103M running words; this corresponds to text blocks of about 3600 words. The validation and test sets are composed of 218K and 246K running words, respectively. Each of them contains 60 articles and about 268 K words. Our experiment follows the standard setting (Merity al., 2017: Schlag et al., 2021) and splits the training data into L-word independent long segments. For evaluation, we use a batch size of and go through the text sequence with a sliding window of size L. We consider only the last position for computing perplexity (PPL) except in the first segment, where all positions are evaluated as in (Al-Rfou et al., 2019; Schlag et al., 2021).
Models and baselines Our language modeling implementation is based on the public code by (Schlag al., 2021). We use their small and medium model configurations for models in our experiments. In particular, for small models, we set the key, value, and query dimension to 128, and the training and evaluation context length to 256. For medium models, we set the key, value, and query dimension to 256, and the training and evaluation context length to 384. In both configurations, we set the number of heads to 8, the -orward layer dimension to 2048, and the number of layers to 16. For Transformer- PrunsforerrMGGKMML 8-head versions to compare with the 8-head baselines. Tir, b, and jjr for this task follow our setings for LRA experiments. Other than those, our language modeling share the same configurations the baselines.
4 and
We train our models for language modeling on 2 A100, 40GB each with a batch size of 96, and each model is trained for 120 epochs. We apply 10% dropout (Hanson, 1990; Srivastava et al., 2014) and use the Adam optimizer (Kingma & Ba, 2014) with an initial learning rate of 0.00025 and 2000 steps for learning rate warm-up.
A.1.3 NEURAL MACHINE TRANSLATION ON IWSLT' 14 GERMAN TO ENGLISH
Datasets and metries The IWSLT14 German German-English dataset (Cettolo et al., 2014) contains 153K training, 7K validation, and 7K test TED-talks scripts
translated sentences. We
16

---

Under review as a conference paper at ICLR 2022
follow the same preprocessing steps in (Ott al., 2019). We use the BiLingual Evaluation Understudy (BLEU) score as our evaluation metric for this task. All trained models are evaluated using the evaluation protocol in (Ott et al., 2019).
Models and baselines The bascline model we use in this machine translation task is an encoder- decoder transformer with six encoder/decoder layers, four attention heads per layer. The embedding dimension and the hidden size are 512 and 1024, respectively, for both encoder and decoder. These architecture configurations are the same for Transfome=-GGK/MMGK except for the number of attention heads per layer, which is reduced by half. We share jrr across all position j and learn it for each head. The initial value for each Tyr is set to 0.5. For Transformer-sMGK we learn b, and initialize its elements from a standard normal distribution. Each orr is a constant with the value
VD where D is the dimension of each head, which is the same in the baselines models. Our experiments follow the settings from (Ott et al., 2019), and our implementation is based on the public code
A.2 MORE EMPIRICAL ANALYSIS OF TRRANSFORMER-MGKE
TRAINED FOR
LANGUAGE MODELING
In Table in the main text, we show the improvements in valid and test perplexity of our Transformer- MGKs/MLKs compared with the baseline softmax and linear transformers. In particular, Transformer- MGK/MLKs with the same number of heads as the baselines, e.g. 8 heads, significantly improve the baselines during training while MGGK//L with 4 head achieve comparable or better performance than the 8-head baselines. In this section, we provide more empirical analysis to shed light on those results. Figure 4 shows the validation perplexity curve of our models versus the softmax and linear transformers.
Figure visualizes the attention matrices from a randomly selected sample for the trained softmax transformer with 8 heads and Maanformer-MGKK with 8 heads and 4 heads. These visualizations attend to more diverse positions in all heads and layers than the softmax attention baseline. We also compare the rank distributions of these attention matrices computed from 1000 samples at each layer in the model. Figure presents the rank histograms of the 8-head softmax attention and 4-head and 8-head MGK attentions for the 1s and 5thh layer. It is clear that attention matrices from the Saansfommr-MGGi have higher ranks than those in the softmax transformer, which implies that Transformer-MGK can attend to more diverse regions than the softmax transformer without the need of using more attention heads.
show that Transformer-MGG
Softmaz Tansformer Il Heads Thenformer-le NCC trach
Tre MDC & heads foreee---ett beats
Lineat mransformer haads Trnnstrmm-//LK Heads
Taanformme heads
Sa PR
60
Flaning
Figure 4: Validation perplexity of Transformer- MGK vs. the softmax transformer (Left) and Transformer-MLK vs. the linear transformer (Right) for language modeling on WikiText-103.
A.3 ADDITIONAL TRAINING RESULTS FOR LRA
In this section, we provide additional experimental results on the LRA benchmark. Figure 7 compares the computational cost measured by FLOPS and model complexity in term of the number of parameters of different inference and learning methods for Transf The computational costs of TransformerMGKK/MMM
and Transformer-MGG Hard-E/Soft-E are par with each other, while Transformer-sMGr uses fewer parameters than the other without trade-off in performance 4 for all tasks. The naming is as explained in Section 3.4 in the main text. In addition, Figure 8 visualizes the attention matrices in the 4-head linear transformer baseline, 4-head Transformer MLMoomrrrMMM.....- and 2-head Transformer-Mo MLKK trained on the document retrieval task.
17

---

Under review as a conference paper at ICLR 2022
Softmax Heads
MGK Heads
MGK 4 Heads
Layer
Layer
Layer
Layer
Layer
Layer
Layer
Layer
Layer
Head
Head 3 Head
Head
Head Head 6 Head
Head
Figure 5: Visualization of attention matrices from 8-head softmax transformer (Left), 8-head Transformer-MGK (Middle), and 4-head Transformer-Me MGK (Right) trained on WikiTexx-102 language modeling. Here, we plot attention matrices for all heads and layers in the models. The sequence length is 256 and the size of each matrix is 256.256.
Softmax 8 Heads
MBK B Heads
MGK Heads
100
30
0--40
320
240
100
330
100
Rank
Figure 6: Rank distributions of attention matrices from 8-head softmax transformer (Left), 8-head Transformer- MGK (Middle), and 4-head Transformer-MGG (Right) trained WikiText- 103 language modeling. The rank histograms computed from 1000 attention matrices each layer. The attention matrices from Transformer- MGKs have higher ranks than those from softmax transformers. This implies more diverse attention patterns, which allows us to reduce the number of heads in Transformer-MGK.
Transformer-MGK. have
A.4 ABLATION STUDY ON THE IMPACT OF THE MIXTURE OF KEYS, THE GAUSSIAN DISTANCE, AND THE KEY SHIFTING
In this section, we conduct an ablation study of the Transformer-MGK on the LRA retrieval task to investigate where the performance improvement from. In particular, we would like to understand the impact of the following factors on the performance of Transformer-MGK: 1) the mixture of keys, 2) the Gaussian distance, and 3) the key shifting. We summarize our empirical results Table and discuss the impact of 2 and 3 below.
Impact of the Mixture of Keys We apply our mixture of keys (MGK) approach to the softmax transformer using the dot product between queries and keys instead of the Gaussian distance as in our paper. We name this model Softmax MGK. We compare the Softmax MGK that has head (Sofmax MGK 1 head in Table 5) with the baseline softmax transformers that use I and 2 heads (Softmax 1 head and Softmax heads in Table 5). Results in Table show that the Softmax MGK head outperforms both the baseline sofumax transformers of and heads. Note that our Softmax MGK
18

---

Under review as a conference paper at ICLR 2022
Retrieval
Listops
Text
33
FLOPs
25
255
23
18.
11
Parameters (1e5)
20
42
35
35
Number Heads
Transttmmns sMGK Transtormor MGK
Transtormer
Transformar Soft
Figure 7: Model complexity (Top) and computational cost (Bottom) different inference and learning methods for Transformer-Me MMG trained on the document retrieval task. While computational costs are almost the same, Transformer-sMCG has more advantage in model size, comparing to Transformer-MGr MGK,, Transformer-Mer Hard-E, and Soft-E. The naming is explained in Section 3.4 in the main text.
Linear 4 Heads
MLK 4 Heads
MLK 2 Heads
layer 1
layer
Inyer
Inyer
Inyer
Iayer
1K
3K
2.54
2.5c
2.5
0.
235
255
4K
3X 2.35%
1K
16
2.33
2.34
2,5g
236
25
25
0.8
2.35
4%
3K. 2.55
48
35.
2.5
46
3K
23K
4X
1K
5X. 2.30
15
2.50
40
10
14
3k
1K
9.
1.38
2,5g
2.5k
dx
1K
2.54
6K
3K
7.5c
<c
1K
2.5
4K
5.5
1K
1..4
2,5c
2.55
235 46
Figure 8: Visualization attention matrices in the 4-head linear transformer haseline (Left), 4-head Transformer- MLK (Middle), and 2-head Transformer MLKK (Right) trained on the document retrieval task. Here, the sequence length is 4000, and the size of each matrix is 4000
1 head is more efficient than the baseline of 2 heads in terms of the number of parameters and the number FLOPs. These results confirm the benefit of using MGK.
Impact of using the Gaussian distance Next, we compare the softmax MGK with the Gaussian MGK. Here the Gaussian MGK is the Transformer-Me MGK proposed and discussed in our paper, which computes the attention scores using the MGK approach and the Gaussian distance between the queries and keys. Results Table 5 suggest that the Gaussian MGK head improves over the Softmax MGK 1 head (80.63% vs. 79.23%) This result justifies the advantage of using Gaussian distance over dot product to compute the attention scores.
Impact of key shifting Finally, apply key shifting to both Softmax MGK and Gaussian MGK (Softmax sMGK and Gaussian sMGK Table 5). From Table 5, we observe that the Softmax sMGK 1 head and Gaussian sMGK 1 head outperform the Softmax MGK head and Gaussian MGK head, respectively. These results, again, corroborate the benefit of using key shifting.
We also include the result for the Gaussian I head model in Table 5. This Gaussian I head model is similar to the Softmax head model but uses the Gaussian distance to compute the attention scores. Comparing the results of the Gaussian I head model, the Softmax 1 head model, and the Gaussian
19

---

Under review as a conference paper at ICLR 2022
Table 5: Ablation study on the impact of the mixture of keys, the Gaussian distance, and the key shifting on the LRA retrieval task. We denote the softmax Transformer by Softmax. All Softmax models (i.e., Softmax 2 heads, Softmax I head, Softmax MGK head, and Softmax sMGK head) use dot product to compute the attention scores. All Gaussian models (i.c., Gaussian MGK head, Gaussian sMGK head, and Gaussian head) use Gaussian distance to compute the attention scores. We denote MGK with key shifting by sMGK. Here MGK used to denote our approach of using a mixture of keys at each timestep.
Method
Accuracy (%)
Softmax 2 heads
79.10
Softmax sMGK head
79.81
Softmax MGK head
79.23
Softmax head
77.90
Gaussian sMCK head
81.23
Gaussian MGK head
80.63
Gaussian / head
80.38
Table 6: Comparing the GPU memory footprint and computational time overhead (seconds/iteration) hetween our 4-head Sonsformer-MGK<MMl retrieval task at test time. Our TransformerMGGK//MLL MGKs/MLK, save much more memory and have significantly smaller wall-clock time compared to the baselines. Here we use a batch size of 32 for each iteration.
and the 8-head softmax/linear transformer baselines trained on the LRA
Method
Memory (Gb) Time overhcad (scconds/itcration)
Softmax 8 heads
58.00
0.357
Transformer- MGK 4 heads
53.58
0.278
Transformer-sMGG. 4 heads
45.50
0.227
Linear 8 heads
3.38
0.055
Transformer MLK 4 heads
2.90
0.043
Transformer-MM.. 4 heads
2.83
0.042
Trarsformer SMGK head Varsformer MGK I head
Transforme: heeds cocosforme-MMGee 2 heads
Transformer SMLK head
Trensformr MLK heads
Trarsformer-LL MLK heed
Teensfrreessk heads
Softmax Wransformer heads
Linear Wansfarmer heads
Logs (10 ()
Ancunacy (10-1)
Loss (10'
Arruracy L10
0
00
63
50
10C
25
40
14
100
20
60
150
20
62
107
10C
40
Training Epochs
Taining Epachs
Figure 9: Training and test loss/accuracy of Transfomer-MGGK softmax transformer (Left) and Transformer- ML.K vs. linear Transformer (Right) on the retrieval task, which has the longest average requence-lengtt and attention span among the LRA tasks (Tay et al., 2021), training, we apply carly stopping to avoid overfitting. That explains why the training loss and accuracy curves stop carly. The test loss/accuracy curves already converge. The impressive performance of Minsformer-MGK/MM on this challenging task validates the capability of our models capture long-range dependencies via learning diversity of attention patterns.
MGK head model reported in Table 5 further confirms the advantage of using MGK and Gaussian distance.
A.5 TIME OVERHEAD AND MEMORY FOOTPRINT ANALYSIS
Table 6 compares the GPU memory footprint and computational time overhead (seconds/iteration) of our 4-head Transformer-MGKK/MM
with those of the 8-head softmax/linear transformer baselines at test time. All models are trained on the LRA retrieval task, and we use batch size of 32 for each iteration. Our MGK/MLK models save memory and reduce wall-clock time significantly compared to the baselines. Using key shifting in our models, i.c. Satnsffomer-sMGKK//see the efficiency further.
helps improve
A.6 LEARNING CURVES SHOWING CONVERGENCE
In this section, we replot Figure and 4 to show the convergence of our trainings. In Figure 1, the results do not seem to converge since we plot the loss and the accuracy in log-scale. In Figure 4, the
20

---

Under review as a conference paper at ICLR 2022
Seftmex Transformer
Lincor Transformer Hoads
Transforme--MCK Heads
Transforme-MLK Heads
Transformer MGK heads
Tranooormer--L
Transtomer-sMäKK heads
100
120
Vpiidaation PR)
1.40
100
120
80
25 50 75
Training opochs
100
Training epochs
100
Figure 10: Validation perplexity of the Transformer-MGr vs. the softmax transformer (Left) and the Transformer-MLKK vs. the linear transformer (Right) for language modeling WikiText- 103. Training converges on this task after 500000 iterations, equivalent to 115 epochs.
Softmax 2 Heads
MGK 1 Head
sMGK
Head
Layer
Layer 2
Layer 1
Layer
Layer
Layer
10
10
Weight Key
Welght Key1
0.75
8
20
20
20
Head
30
0.50
40
80
40
025
40
60
20
0.00
10
10
Weight Key
-0.2
Head
30
30
-0
20
Figure 11: Weight matrices W for computing the keys, for all heads and layers, in the 2-head softmax transformer baseline (Lefft), the 1-head Transformer-MGr with 2 keys (Middle), and the 1-head Transformer- sMGK with keys (Right) trained on the LRA retrieval task. Here, the dimension of each head = 32 and that of input is D. 64. Hence, each weight matrix has the shape of (64, 32).
Table 7: The learned mixing coefficient Typ of all heads and layers in the I-head Transformer-Men MGGKs trained on the LRA retrieval task. Here we use the same Tyi, 7yz for all time step j N.
Method
Layer 1
Layer 2
7y1
7y2
Ty1
7y2
TraKsformer:MGG heads
0.488
0.512
0.500
0.500
Transformer-MGK heads
0,502
0.498
0.497
0.503
results do not seem to converge since we zoom into the specific range on the and x-axes. Figure and are the replotted versions of Figure and 4, respectively. Figure the training loss/accuracy curves stop early because use early stopping to avoid overfitting. The test loss/accuracy curves in this figure already converge.
A.7 WEIGHT MATRICES OF THE KEYS, KEYS AND MIXING COEFFICIENT
In this section, we analyze the learned Kjr. kjr, and K,, j = I,...,N N and r I,..., M in the Transformer-MGK trained the LRA retrieval task. In all of our experiments, we set M 2. In Figure 11 and 12, we visualize the weight matrices Wx that computes the keys and the keys K, respectively, for all heads and layers in the 2-head softmax transformer baseline. the 1-head Transformer-Mer with 2 keys, and the I-head Transformer-sMi MMG with keys trained on the LRA retrieval task. Note that for the keys, only plot the first 100 tokens. Also, Table 7 summarizes the
21

---

Under review as a conference paper at ICLR 2022
Softmax 2 heads
MGK
Head
SMGK 1
Head
Layer
Layer 2
Layer
Laver 2
Layer 1
Laver 2
1.00
20
0.75
Head 1
40
Key
40
Key
0.50
0.25
0.00
30-
20
-0.25
Head 2
Key
Key 2
40
40
80
-075
Figure 12: Key embeddings K for all heads and layers of the 2-head softmax transformer baseline (Left), the I-head Transforme-Mo MeKK with keys (Middle). and the 1-head Transformer- on the LRA retrieval task. Here the dimension D each head is 32, and we plot the key embeddings of the first 100 tokens in a randomly chosen sequence. Hence, each key matrix has the shape of (100, 32).
with 2 keys (Right) trained
Retrieval
Listops Number Transformer- Heads
Text
Retrieval
Listops
Text
Number Transformer Heads
ELOP
10
Number Softmme Transforme Heachs
Number of Lincar Transformrr Heads
Softmax Transformer --- Traneforre--MGK
Linear Transtomer
Transformor
Computational cost (FLOPs) for each training iteration of Transformer MGK the baseline softmax transformer (Left) and Transformer-Min MLK vs. the baseline linear transformer (Right) on the LRA retrieval task. The efficiency advantage of Transforme-MGKKMM... over the baselines grows with the number of heads. Here, the batch size of each training iteration is 32.
Figure 13:
learned mixing coefficient jjr of all heads and layers in the 1-head Transformer-MGH trained on the same retrieval task. Here we use the same Ty,, T2 for all time step j = I....,N.
A.8 ADDITIONAL COMPUTATIONAL COMPLEXITY (FLOPS) ANALYSIS AT TRAINING TIME
Figure 13 demonstrates the computational cost (FLOPs) for each training iteration of the Transformer- MGK vs. the baseline softmax transformer (Left) and the Transformer-LL. vs. the baseline linear transformer (Right) on the LRA retrieval task. The efficiency advantage of Transformer-MGK/ML... over the baselines grows with the number of heads.
Figure 14 shows the computational cost per training iteration (measured in FLOPs) of different infer- ence and learning methods for T'rasformmr-MM MGK trained on the document retrieval task. Transformer- sMGK, Minnffrmer-MGK. Transformer-MGK Hard-E, and Soft-E have similar computational costs.
A.9 SCALING TO 12-HEAD BASELINE MODELS FOR THE RETRIEVAL TASK.
To further study the scalability of our model, in this section, we investigate the performance of our 6-head Transformer-MGKs/MM transformers on the retrieval task. Table indicates that our 6-head significantly outperform 12-head softmax/linear transformers, respectively. Moreover, comparing these results to those in Table and Table 2, although the 12-head softmax/lincar transformers improve over the 8-head ones, their accuracies are still worse than or only equivalent to those of our 4-head and even 2-head Transformer-MGK/MLL.
in comparison with the 12-head baseline softmax/linear
models.
A. NEURAL MACHINE TRANSLATION ON IWSLT' 14 GERMAN TO ENGLISH
Table 9 shows that the 2- head Transformer-MGK/SMGGG MGK/ models have comparable or better BLEU scores than the head softmax transformer baseline.
22

---

Under review as a conference paper at ICLR 2022
Retrieval
Listops
Text
16
80
60
8-
12
FLOPs
50
40
30
20
Number Heads
T--nnnnoorr--MMMCS Taanfforer-MGG Transformer-Me Hard-E
Transformer MGK Soft-E
Figure 14: Computational cost (measured in FLOPs) per training iteration of different inference and learning methods for Transformer-MGG trained on the LRA retrieval task, Computational costs are almost the same for Transformer-MGGK, TransformerMGGK, Transformer-MGK Hard-F, and Soft-E. The naming explained in Section 3,4 in the main text.
Table 8: Test Accuracy (%) of 6-head MietoormeerMGKK/MLLL.. compared with the baseline 12-head softmax linear transformers on the retrieval task. 6-head Transformer-Me Sanssormer-MGKs/MLL. significantly outperform softmax and linear transformers, respectively, while being more efficient in terms of computational cost, model size, and memory usage.
Method
Accuracy (%)
Safimax 12 heads Transformer sMGK head
82.18
83.31
Transformer MGK 6 head
83.05
Linear sMGK 12 head
81.97
Transformer sMLK 6 head Transformer MLK 6 hcad
82.80
82.11
Table Machine translation BLEU scores of 2-head Transformer-oe MGGKs on the IWSLT14 De-En dataset better than or equivalent to that the 4-head baseline.
Method
BLEU score
Sofimax heads
34.42
Transformer sMGK 2 head
34.69
Transformer MGK 2 head
34.34
A.11 COMPARISON TO MULTI- MULTI-QUERY ATTENTION
In this section, compare MGK approach to the multi-query attention (Shazeer, 2019). The multi-query attention shares the same set keys and values different heads to reduce the memory- bandwidth cost during incremental inference, which allows faster inference since the size of the reloaded "keys" and "values" tensors are significantly reduced. On another hand, our Transformer- MGK models the key each head a Gaussian mixture model, which leads the of multiple keys at each head and allows us to decrease the number of attention heads. This helps reduce the computations and parameters needed calculate additional queries and values. If using key shifting (see option (B) in paragraph Design Options for Keys in Section 2.3), the MGK approach also helps reduce the computations and parameters needed calculate additional keys. The advantages of Transformer-Mo MrKK hold in both training and inference, including incremental inference. We have provided a detailed analysis on the computational complexity and the number of parameters of the Transformer-MGE in comparison with the corresponding softmax transformer in Appendix B. Combining the multi-query attention and our MGK approach is interesting since each method has its own advantage and they complementary to each other. In particular, model share the same set of values and mixtures of keys at different heads. This approach can potentially have the advantages of both multi-query attention and MGK.
can let the transformer
23

---

Under review as a conference paper at ICLR 2022
B AN ANALYSIS ON THE COMPUTATIONAL COMPLEXITY AND THE NUMBER OF PARAMETERS IN TRANSFORMER -GGK AND THE SOFTMAX TRANSFORMER
In this section, we compare the computational complexity and the number of parameters in transformer- MGK with M keys each timestep and //M heads to the baseline softmax transformer that has 1 key at each timestep and H heads. Here we choose M such that H is multiple of M and use keys design option (A) that make the key kjr linear projection
the input aj, i.e. j = 2; WK, K. where 2 E IRIDD=, WK. E DxD= and r = 1, ....., M (see Design Options for Keys in Section 2.3 for more details). To simplify notation and without loss of generalization, let M as in our experiments and assume that D, = D, i.e. the values have the same feature dimension as the queries and the keys. To simplify the computation, we also do not take the softmax operator into account since this operator yields similar costs when applied in Transformer-Mor and softmax transformer.
B.1 COMPUTATIONAL COMPLEXITY
(i) Softmax H-head attention: The number of computations in softmax H-head attention is H - 1) + (6D, + 2HD - 5).
Explanation: To calculate the query matrix Q, the key matrix K. and the value matrix V in Step in Section 1.1 at each head, we need 3NDD, multiplications and 3N 3ND( D(D, DD=11) additions. In total, these need 3N (2 1) computations. Next, to compute the product QK" in Eqn. (1), we need N D multiplications and 2 1) additions. Similarly, the product AV requires 22 D multiplications and (N 1) additions. In total, computing the output sequence H in Eqn. (I) at each head requires 3/V D(2D, - +)++++++++ + (2 1DD-10 ++++ N2 D
- (N = N -
computations. The total computation for all H heads is then
T(N - (D( 4)) NHD 1)
= H - 1) + NHD( + 2HD - 5),
where the extra NHH - is from the linear projection by W
(if) Mixture of Gaussian keys attention with 1/2-head: The number of computations a mixture of 2 Gaussian keys attention with H/2-head is N H 0.5) DHD( + HD - 4).
Similar to the above derivation, in a Mixture of M Gaussian keys attention, to compute the output sequence H we need ( + 2) - 1) + (NM + (2D, 1) 1) computations. Note that, to compute the Gaussian distances between the queries q, and the keys k, as in the mixture of M Gaussian keys attention and to compute their dot product in the softmax attention, need the similar number of computations. Therefore, the total computation for all H/M heads is then
Explanation:
(H/M) ((2 + +2)D 1) ND((M + +2)(2D, 1) 1)- 1)) + (HD(((H/MM))))))
= =N2H 2( + +1) + NHD
M (-M
2(M + 22
2 HD
3M +
D.
M
M
M
Let M = 2, then the total computation of the mixture of 2 Gaussian keys attention is given by 22 H33 0.5) + NHDD(D + HD 4)..
Soft-max H-head attention versus mixture of Gaussian keys attention with H/2-head: Given the results in (i) and (ii), when compared to the baseline softmax H-head attention, our mixture of 2 Gaussian keys attention with H/2 /22-head saves
2 0.5) + NHD(2D, + HD 1)
computations. When N is large, this difference is significant. In conclusion, the mixture of Gaussian keys attention with H/2-heads has cheaper computational complexity than that of soft-max H-head attention.
B.2 THE NUMBER PARAMETERS
(iii) Softmax H-head attention: The number of parameters in a softmax heaa attention is 3HDD, + (HD)
24

---

Under review as a conference paper at ICLR 2022
Explanation: 3HDD, is from the linear projects to calculate the query matrix Q, the key matrix K, and the value matrix V in Step I in Section 1.1. (HD)² is from the linear project to compute the final output in Eqn. (3).
(iv) Mixture of 2 Gaussian Keys attention with H //2-hhad: The number of parameters in a Mixture of Gaussian Keys attention with H/ /22-he is 2HDD, + D) + H.
Explanation: The linear projections to calculate Q and V contribute HDD parameters each. The linear projection to calculate K contributes HDD,, The linear project compute the final output has dimension D x HD, so it contributes 0.5( D)² parameters. H more parameters is from the prior Tjr-- These priors contribute parameters each head since we make (mzI, T22) for all
j 1; NV the same.
Soft-max H-head attention versus mixture of Gaussian keys attention with H/2-head: Given the results in (iii) and (iv), when compared to the baseline softmax HH- attention, our mixture of 2 Gaussian keys attention with H /22hhead saves HDD, + OD5(H H parameters. When H and
D are large, this saving is significant.
C PROOFS OF MAIN RESULTS
In this appendix, we provide proofs for the main results in the paper.
C.I PROOF OF THEOREM 1
To ease the presentation of the proof, for any probability distribution G, we denote
pa(x) = / f(x ( ((d=()) = f ((x)), (T))G()))
for all x E Rd where f(x) (v2road exp
|=ll for given a > 0. means that pc is the 2
convolution of f and the probability distribution G. Since the space of Gaussian mixtures is dense in the space of continuous probability measures (Bacharoglou, 2010), indicates that there exists probability distribution G1 such that
sup Ip(x) Pa; (x) E 2
(14)
ZER
Our next step is to prove that there exists a probability measure G2 with at most K supports where K & (C')og(1/e)) for some universal constant C' such that
sup |pc, - (P) 2' zERd
(15)
Indeed, from Lemma A. in (Ghosal & van der Vaart, 2001), for any 2 there exists a probability distribution G2 with most (2k 2)d supports such that
f (0d(Cc G2) = 0.
(16)
for any (a1, .2,...,,ad) No such that 0 0 |a| Σラ= aj S 2k - 2, Here, lr II 0g:
Now, for any M > 2avd, we have || - 0ll 2 |||||| - -100 > M - avd > M/2 as long as ||||| > M
and 0 [-a, a)d.. It indicates that
/ f(0 [(x O)a(( -))(()))))(())
sup |pG, (x) PG, (=)) ||>||>
sup
[|]]
| 012
< sup
exp
d(G, (2)
(V2TTo)d
20
||| >M
M2
exp
(17)
( (VTT))d
80
25

---

Under review as a conference paper at ICLR 2022
On the other hand, for any & 2 1 we also have that
/ M( - [( 00 (2)
sup |
|PG, (ar) PGG (x))
sup MM
/ f(z 00)
k-1 1) | 01
0) -
Σ
d(G -2)
sup M
(v2T)
dd
1=0
(-1)"11a 0
+ sup
== (V2r)dod+23)) d(G) G2)(0)
k-1 1 | 0
= sup |
f(x f(0- 00 (V27)
Σ
d - G2)
1=0
(18)
where the final equality stems from
(-17|=-0 - d(G: G2) = 0,
(V2m)dgd+2131
which is due to Eqn. (16).
To further bound the right-hand-side (RHS) of Eqn. (18), we use the following inequality:
exp(v) Σ(0) /5!
/
0
for any y E R. Since 2 (k/e)* for any le 2 1, the above bound can be rewritten as
k-1
lyel* h
exp(v)
Σ
Cu)0/1 /5!
(19)
Further simplification of Equ. (18) leads to
/
1-1) 01133
sup |PG, (x) - Pa. (x)] |
sup
f(x 0)
d(G1 G2)(
== 2n 2d
||||| SM
1-1)| -1 01
2
sup
[x 0)
|| Mr
(
j=0
|| 1 013
11 || - 01123
sup
exp
| SM 2. (V2noo)
20
ct ||x
0 2k
sup
2k
(2k)*
||| SM
a.
where the final inequality is based on an application of inequality (19) with y -11 01 (2o2)
For || s M and E [-a.a]] a)d, we have | have
01 < |||||| + ||0|| < M avd. Therefore, we further
e 11 1 -
e (M + av 22 (2k))
sup Ipa, (x) pas (x:) 5 [|]] SM
sup || M.llla.aaa
a 2k (2k)*
When M 2 2dyd, we have M + aadd & 2
3M and the above bound leads to
(9e)* * 2
sup |
|pG, (x) pa, (x)) 5
(20)
k80
26

---

Under review as a conference paper at ICLR 2022
By choosing M 802 log(1e') for some è' 0, the bounds in Eqns. (17) and (20) become
sup 1
IpG, (x) - -PG2(x)) (V
(9e) (log((()))))))
sup IPG, (x) (P))(2))= |
(21)
hb
As long as we choose k = 9e² log(1/e') and E' s 1, we have
sup |
Iaa, (2)) - (x)( S e e 96" loR(1/e)) = (e)) Sé.
(22)
By choosing
the results from Eqns. (21) and (22) indicate that (vine
2max
sup |
|PG, (2) PG, (x) S 2'
and sup |
|PG, (z) PG., (x)) S
mmxx (Viroja 1)
and h: = 96² log 2 max
Therefore, if we choose M = 80 log
(VITa)a
have
sup PG,, (x) - Pc, (x) S C zER4
2
indicates that we obtain the conclusion of claim (15) by choosing K (2k 2)d s
2max (yzen .1) ))
18c² log
Combining the results from Eqns. (14) and (15), we have sup |Pa, (æ) Pa (x)) S
sup pp(z) Paa (x) S sup (p(z) Pa, (x))
.ER1
zERd
zER*
As consequence, we obtain the conclusion of the theorem.
27


Reviews:
[{"review_id": "tsZ4LzAj8NG", "reviewer": "Reviewer_gyah", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": 4, "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": "Strong points\n \n1. The proposed method is empirically efficient in terms of the number of parameters and computational cost in terms of FLOPS.\n\n2. Although there is previous work for the RBF kernel interpretation of self-attention [1], the proposed probabilistic interpretation of self-attention provides an interesting direction for future research.\n\n3. This paper provides experimental results on LRA benchmark and a language modeling benchmark to show the effectiveness of the proposed method.\n\n4. The paper also introduces the linear-attention variant of the proposed method.\n \n[1] Yao-Hung Hubert Tsai et al. Transformer Dissection: An Unified Understanding for Transformer’s Attention via the Lens of Kernel.  EMNLP-IJCNLP 2019\n\n##########################################################################\n\nWeak points\n\n1.  Although the authors argue that Transformer-MGK is introduced to improve the explanation power of each key, increase the representation of each attention head, and reduce the chance of learning redundant heads, this paper provides only empirical results and no theoretical explanation for these improvements.\n(1) It is valuable to describe the computational/space complexity of the proposed method and compare them with those of conventional methods.\n(2) Although I can somehow expect that the mixture of Gaussian increases the representation power of keys, I cannot imagine why the proposed method reduces redundant heads and promote the diversity in attention pattern.\n\n2. Since pre-training and finetuning paradigm is out of scope of LRA benchmark, I suggest the authors to conduct experiments for standard NLP benchmarks, such as GLUE, SQuAD, SuperGLUE, or Winogrande.\n\n \n", "summary_of_the_paper": "The paper addresses the efficiency of Transformer models.\nThe paper provides a probabilistic view of self-attention in transformers and introduces a novel transformer model, Transformer-MGK, and its extensions.\nIn Transformer-MGK, each key of self-attention is modeled as a mixture of Gaussians whose mixture weights are estimated by the E-step of EM algorithm.\nIn the experiments, they show Transformer-MGK is comparable or better than Transformers on LRA benchmark and a language modeling benchmark.\nIn addition, the authors empirically show Transformer-MGK alleviates the head redundancy of conventional Transformer.", "summary_of_the_review": "Overall, I vote for weak acceptance. I like the probabilistic interpretation of self-attention and the empirical results are strong. \nMy major concern is about the clarity of the paper (please see weak point #1 below).  \n", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "n_v9x3NbqA-", "reviewer": "Reviewer_EtfP", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": 4, "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": "\n- First, I like the Gaussian mixture perspective of multihead attention in general, but it could lead to cases that can be a bit awkward. For example, in causal attention, the number of components of the mixture changes over time steps. This could be avoided by, e.g., treating only the multihead part as a mixture, not the keys over timesteps. \n- Second, it can be confusing to weave into the narrative something that is actually not used at all, e.g., the norm constraints around Eq 5, and the M step around Eq 13. I suggest reworking these part and probably presenting them as connections to other models. \n- Third, some modeling details need to be clarified. Optional (A) for key’s parameterization, the notation W_{k_{jr}} seems to suggest it needs NM number of projection matrices. This seems  much more expensive than standard multihead attention. Similarly in option (B), why would it need N number of matrices at all? Further about the parameterizations of queries and values. The notations seem to suggest that they are never “multiheaded” as keys are. Could the authors confirm this? If this is true, it is worth clarifying.\n- Adding onto the above, what does a MGK with 4 heads mean? Does it have 4 x seq_len number of Gaussian components for the keys?\n- Additionally, it would be nice to include a time overhead comparison.\n- Lastly about baselines and further experiments: the WikiText103 baseline seems pretty far behind, e.g., Baeski and Auli (2018). Could the authors explore MGK on top of a stronger baseline? Besides, [1] seems pretty related and worth comparing to. Given the weak LM baseline, I suggest the authors including other sequence modeling tasks such as machine translation.\n\nTypos:\n- Above 1.1, keyed -> key.\n- Above Eq 4, unit matrix -> identity\n- Above Eq 7, abusing -> overloading\n- Below Eq 11, at learning -> at learning time.\n- Start of Section 3, “full the number of heads” doesn’t parse for me.\n\nReferences:\n\n[1] https://arxiv.org/abs/1911.02150\n\n[2] https://arxiv.org/abs/2103.02143", "summary_of_the_paper": "This paper proposes a mixture-model perspective of multihead attention. Each key within each attention head is treated as a Gaussian component; with some assumptions, the softmax-based attention weights can be recovered as the posterior probability of the Gaussian mixture, conditioning on the query. The framework leads to a new parameterization of multihead attention, which is explored in transformers. It can be applied to several linear transformer variants, too. Experiments with text classification, language modeling show that the proposed method achieves same or better accuracy controlling the number of attention heads, and reduces computational overhead. Analysis suggests that it reduces the redundancy of multihead attention", "summary_of_the_review": "Strengths:\n- An interesting mixture model perspective of multihead attention\n- It is nice to reduce compute cost with the same or better accuracy\n\nWeaknesses:\n- It can be confusing to introduce concepts that are never used in the model\n- Several key details of the implementation need to be clarified.\n- Weak LM baseline. Would be great to include more experiments. ", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "TFcncGYXls_", "reviewer": "Reviewer_r2tu", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 5, "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": "Strengths\n-------------\n\n- The paper provides a nice probabilistic view on the attention matrix that could be a new tool to study the transformer architecture.\n- Language modeling experiments besides LRA are greatly appreciated\n- The paper is well written and easy to follow\n\nWeaknesses\n------------------\n\nThe most significant weakness is that the provided experiments do not provide adequate evidence regarding any benefit of the proposed method.\n1. In both LRA and Wikitext-103 all scores way too close to be able to judge whether any of the methods performs better.\n2. Even if they do perform better (the only one could be sMGK on LRA-Retrieval) the model is sufficiently changed that more tests would be required. For instance, is the performance improvement coming from using the distance instead of the dot product? Is it coming from the key shifting? \n3. In all experiments, what is M? It is quite hard to compare the performance with different number of heads without having M for MGK and MLK.\n4. Moreover, it would be interesting to see the performance of softmax and linear transformers with 4 heads on Wikitext-103. It is quite interesting in tables 1 and 2 that reducing the number of heads has minimal impact in performance.\n\nFinally, it would be interesting and beneficial to the paper's main point to have a proposition regarding why parameterizing the attention matrix in this way would improve the diversity of the attention heads. The rank argument does not show diversity between attention heads given that the attention matrices can all be the same but have high rank. Moreover, when thinking about how diversity can be achieved with this formulation, it can also be measured and enforced with some type of loss or regularizer.\n\nSuggestions\n-----------------\n\n- It would be much more informative to exchange FLOPS with wall-clock time and although number of parameters is important for large transformer models, so is the memory footprint in the GPU during training.", "summary_of_the_paper": "The paper constructs a Gaussian mixture model with uniform prior and one Gaussian centered at each key and shows that the posterior distribution ie. the probability that a given query is generated by a given key matches the attention scores in the transformer architecture. In addition, it is proposed to describe each key as an additional Gaussian mixture and it is shown that the posterior is easy to compute and resembles the multi-head attention (given that the query remains the same for all heads). The paper proposes that this formulation can learn more diverse attention-heads and reduce the computation by replacing some heads with the Gaussian mixture heads which are easier to compute. Moreover, the authors show that it is straightforward to transfer the aforementioned formulation to linear transformers. Finally, the authors perform experiments on the LRA benchmark and language modelling experiments on Wikitext-103.", "summary_of_the_review": "As mentioned in the top of the weaknesses section, the main reason for my recommendation is that there is not sufficient evidence showing an improvement in performance with the proposed method.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "1: The contributions are neither significant nor novel."}, "extra_text": {}}, {"review_id": "3Wu-y_Vr3Om", "reviewer": "Reviewer_fA4U", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 4, "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": "Strengths:\nThe authors use the probabilistic framework to establish the connection between self-attention and gaussian distribution. Even though there are some assumptions: (i) the query and the key are normalized; (ii) all the standard derivations are identity; (iii) the prior distribution is uniform, I think these assumptions are easy to follow. Based on the connection, a natural extension is the mixture Gaussian model and they apply the EM algorithm to learn the posterior distribution. The experiment results are comprehensive.\n\nWeaknesses:\n1. In Table 1 and Table 2, the authors discover the influence of the number of heads. However, [1] claims that one head is indeed sufficient, which is not consistent with your results. I wonder whether the reason is the experiment settings.\n2. In Table 1, Transformer-sMGK is better than Transformer-MGK in some cases but not in all cases. Can you provide some explanation about the reason? Also the situation in Table 2.\n3. In Table 1 and Table 2, the authors start from 8 heads. Can you provide the results from the vanilla Transformer (12 heads) for better illustration?\n\n[1] Michel, Paul, Omer Levy, and Graham Neubig. \"Are Sixteen Heads Really Better than One?.\" Advances in Neural Information Processing Systems 32 (2019): 14014-14024.\n", "summary_of_the_paper": "This paper proposes Transformer with a Mixture of Gaussian Keys (Transformer-MGK), which replaces the multi-head self-attention module with a Gaussian mixture model. Specifically, the authors build the connection between self-attention and probabilistic framework. Based on their findings, they extend the self-attention to a Gaussian mixture attention. Compared with vanilla Transformer, Transformer-MGK has fewer parameters and accelerates training and inference with comparable performance. Various experiment results illustrate the advantage of Transformer-MGK.", "summary_of_the_review": "This paper builds the connection between self-attention and Gaussian distribution with some assumptions, which is novel for Transfomer model. The extension from original self-attention to mixture Gaussian model is natural and convincing to me. The experiment results are sufficient to illustrate the improvement of their models.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "0NOZ2KoxvhG", "reviewer": "Reviewer_qyjK", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 3, "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": " Strong points\n\n- The idea of introducing the mixture of Gaussian keys to increase the expression capacity of key embeddings and to allow the query to attend to more diverse positions is interesting.\n- Overall, the paper is well written. Especially, in section 2, the formula explanations have a neat flow so that it is easy to follow.\n- The proposed model shows comparable results to general softmax-based transformers with fewer parameters and FLOPs.\n- In the experiment analysis, the visualization of the attention matrix and rank distribution supports that the paper's hypothesis that Transformer-GMK would learn diverse attention patterns.\n\nWeak points\n\n- The main concern about this paper is the experimental improvements are marginal. Moreover, although the method shows comparable to original transformers with fewer parameters and FLOPs, I suggest the authors to discuss whether there is any difficulty in the learning process using EM-based algorithms for practical usability; for example, how many learning epochs are required, whether it is sensitive to initial values, and an increase of computation required to calculate the E-step during the training process.\n- Beyond the attention matrix analysis, in order to demonstrate the keys are well modeled as Gaussian mixtures, the learned $\\pi_{jr}$,  $k_{jr}$, or $W^k_r$ which is the weight matrix projecting to $k_{jr}$ seems to be required to be analyzed.\n- In the experiment section, the number of Gaussian mixtures is not noted. Also, comparing with baseline methods (Softmax and Linear heads) and proposing methods with half heads is, for me, unfair without considering the number of the mixtures.\n-  I'm not convinced with the experimental results are reported at their convergence since the learning curves in Figures 1 and 4 end before converging.\n\nQuestions\n\n- How does the proposed model result with the increasing number of mixtures?\n- How many epochs are needed to converge via the EM algorithm?\n- Also FLOPs of training time including the E-steps should be compared.\n- Can this proposed approach be applied to deeper and large Transformers such as BERT and GPTs and show its effectiveness?\n\nComments\n\n- I would like to recommend comparing the model complexity and computational cost with notations and formulas rather than empirical figures.\n- Typo in Figure 3. The label of Softmax Transformer is wrong as Liner Transformer in the left figure.\n- Typo in Table 4. The lower rows group: Transformer-MGK → Transformer-MLK", "summary_of_the_paper": "This paper proposes a transformer with a mixture of Gaussian keys (Transformer-MGK) in order to avoid training redundant attention heads in the transformer modules. The authors formulate the attention mechanism as a posterior distribution $p(t_j = 1 | q_i)$ (i.e., given the query $q_i$, how likely $q_i$ matches the key $k_j$), and extend this as the mixture of Gaussian keys where each key $k_j$ is modeled as a mixture of $M$ Gaussians with several assumptions. The paper introduces inference methods using the EM algorithm: Soft E-step and Hard E-step. The authors also propose a liner variant, Transformer-MLK, which can be more computationally efficient to long sequences. Empirical studies are performed to show the effectiveness of Transformer-MGK over the softmax Transformers. Results are shown on the long-range arena benchmark and language model on WikiText-103 corpus. Also, comparisons of the model complexity and computational cost and ablation study depending on different inference and learning techniques are also provided.", "summary_of_the_review": "Overall, I vote for \"marginally below the acceptance threshold\". I think the extension of attention keys to the mixture of Gaussian and exploring the learning methods with EM algorithms are somewhat technically sound. However, experimentally demonstrating its practical usefulness is required further. Hopefully, the authors can address my concern in the rebuttal period.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}]

Author rebuttal:

Meta-review: The paper proposes using a mixture of Gaussian models for transformer keys (MGK) so that the posterior distribution of key given query matches the attention scores in the transformer architecture under some assumption. Similarly but in reverse, the query given key under a MoG also matches transformer attention score. The paper proposes that this formulation can learn more diverse attention-heads and reduce the computation by replacing some heads with the Gaussian mixture heads which are easier to compute. Moreover, the authors show that it is straightforward to transfer their formulation to linear transformers. The authors perform experiments on the Long Range Arena benchmark and on Wikitext-103 where they show some improvements with less attention heads, while using up to 20% less FLOPs for softmax transformers and around mostly 20% but up to 80% less parameters for the worse linear transformer.

Reviewers generally find that the interpretation of transformers as mixture of Gaussian interesting, although it is based on normalization assumptions. The main objections from both positive and negative reviewers is the weakness of the empirical results. In summary, all the improvements in perplexity are less than 0.5, and accuracy improvements are less than 1. In the discussions, the authors claim that their comparisons are using fewer heads vs baselines with more heads, which is how the tables are presented, but upon closer look I find this unconvincing since the same head count comparisons can be reconstructed and still show very weak results. For example in table 1, acc(MGK4)=61.85 vs. acc(baseline4)=61.23, or in table 2 acc(LinearMG4)=55.7 vs acc(LK4-baseline)=55.61. To begin with, it would be better to compare the same number of heads for more head-to-head comparisons, so I find the whole argument to be misleading. The claims about FLOPS and memory has a similar flavor where the differences are small but the claims were big. A visual evaluation of figure 3 show that most reductions are around 10% in FLOPs or parameters for the better softmax model (note the axis does not start at 0). The retrieval task for the worse linear transformer is the only case where the FLOPs reduction seems significant. Given that EM is now required during training, I'd interpret the results as a negative for MGK given the marginal improvements achieved.

Several reviewers are unhappy with the strength of the empirical results while most reviewers gave favorable scores after the discussion where the authors insisted that some valid points are not valid. The authors requested the AC to look further into the one negative review due to the lack of a reviewer's response. After taking some time to look at the substance of the paper and the review, I recommend rejection due to the weakness of the results, the misleading presentation and discussions.
Decision: Reject