Original: T RANSFORMER WITH A M IXTURE OF G AUSSIAN K EYS
Revised: T RANSFORMER WITH A M IXTURE OF G AUSSIAN K EYS

Original: Anonymous authors
Revised: Anonymous authors

Original: Paper under double-blind review
Revised: Paper under double-blind review

Original: A BSTRACT
Revised: A BSTRACT

Original: 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.
Revised: 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.

Original: 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.
Revised: 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.

Original: 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.
Revised: 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.

Original: These mixtures of keys follow a Gaussian mixture model and allow each attention head to focus on different parts of the input sequence efﬁciently.
Revised: These mixtures of keys follow a Gaussian mixture model and allow each attention head to focus on different parts of the input sequence efﬁciently.

Original: Compared to its conventional transformer counterpart,
Revised: Compared to its conventional transformer counterpart,

Original: Transformer-MGK accelerates training and inference, has fewer parameters, and requires less FLOPs to compute while achieving comparable or better accuracy across tasks.
Revised: Transformer-MGK accelerates training and inference, has fewer parameters, and requires less FLOPs to compute while achieving comparable or better accuracy across tasks.

Original: Transformer-MGK can also be easily extended to use with linear attentions.
Revised: Transformer-MGK can also be easily extended to use with linear attentions.

Original: We empirically demonstrate the advantage of Transformer-MGK in a range of practical applications including language modeling and tasks that involve very long sequences.
Revised: We empirically demonstrate the advantage of Transformer-MGK in a range of practical applications including language modeling and tasks that involve very long sequences.

Original: On the Wikitext-103 and Long Range Arena benchmark,
Revised: On the Wikitext-103 and Long Range Arena benchmark,

Original: Transformer-MGKs with 4 heads attain comparable or better performance to the baseline transformers with 8 heads.
Revised: Transformer-MGKs with 4 heads attain comparable or better performance to the baseline transformers with 8 heads.

Original:  I NTRODUCTION
Revised:  I NTRODUCTION

Original: 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 (AlRfou et al., 2019; Dai et al., 2019; Williams et al., 2018;
Revised: 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 (AlRfou et al., 2019; Dai et al., 2019; Williams et al., 2018;

Original: Devlin et al., 2018; Brown & et al., 2020;
Revised: Devlin et al., 2018; Brown & et al., 2020;

Original: Howard & Ruder, 2018; Rajpurkar et al., 2016; Dehghani et al., 2018;
Revised: Howard & Ruder, 2018; Rajpurkar et al., 2016; Dehghani et al., 2018;

Original: So et al., 2019; Dosovitskiy et al., 2020; Touvron et al., 2020).
Revised: So et al., 2019; Dosovitskiy et al., 2020; Touvron et al., 2020).

Original: 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;
Revised: 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;

Original: Yang et al., 2019; Liu et al., 2019).
Revised: Yang et al., 2019; Liu et al., 2019).

Original: The success of transformers is rooted in the self-attention mechanism as their fundamental building blocks for modeling (Cho et al., 2014; Parikh et al., 2016; Lin et al., 2017).
Revised: The success of transformers is rooted in the self-attention mechanism as their fundamental building blocks for modeling (Cho et al., 2014; Parikh et al., 2016; Lin et al., 2017).

Original: For each token, self-attention computes a weighted average of the feature representations of other tokens where the weight is proportional to a similarity score between each pair of tokens.
Revised: For each token, self-attention computes a weighted average of the feature representations of other tokens where the weight is proportional to a similarity score between each pair of tokens.

Original: This mechanism allows a token to pay attention to other tokens in the sequence and attain a contextual representation (Bahdanau et al., 2014; Vaswani et al., 2017;
Revised: This mechanism allows a token to pay attention to other tokens in the sequence and attain a contextual representation (Bahdanau et al., 2014; Vaswani et al., 2017;

Original: Kim et al., 2017).
Revised: Kim et al., 2017).

Original: It has been shown that the representation capacity of the attention mechanism (Tenney et 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;
Revised: It has been shown that the representation capacity of the attention mechanism (Tenney et 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;

Original: Hewitt & Liang, 2019) is key to the impressive performance of transformers in practice.
Revised: Hewitt & Liang, 2019) is key to the impressive performance of transformers in practice.

Original: 1. S ELF -A TTENTION
Revised: 1. S ELF -A TTENTION

Original: For a given input sequence X :=
Revised: For a given input sequence X :=

Original: [ x 1 , · · · , x N ] (cid:62) ∈ R N × D x of N feature vectors, self-attention transforms X into the output sequence H in the following two steps:
Revised: [ x 1 , · · · , x N ] (cid:62) ∈ R N × D x of N feature vectors, self-attention transforms X into the output sequence H in the following two steps:

Original: Step 1.
Revised: Step 1.

Original: The input sequence X is projected into the query matrix Q , the key matrix K , and the value matrix V via three linear transformations
Revised: The input sequence X is projected into the query matrix Q , the key matrix K , and the value matrix V via three linear transformations

Original: [Equation]
Revised: [Equation]

Original: where W Q , W K ∈ R D × D x , and W V ∈ R D v × D x are the weight matrices.
Revised: where W Q , W K ∈ R D × D x , and W V ∈ R D v × D x are the weight matrices.

Original: We denote Q :=
Revised: We denote Q :=

Original: [ q 1 , · · · , q N ] (cid:62) , K :=
Revised: [ q 1 , · · · , q N ] (cid:62) , K :=

Original: [ k 1 , · · · , k N ] (cid:62) , and V :=
Revised: [ k 1 , · · · , k N ] (cid:62) , and V :=

Original: [ v 1 , · · · , v N ] (cid:62) , where the vectors q i , k i , v i for i = 1 , · · · , N are the query, key, and value vectors, respectively.
Revised: [ v 1 , · · · , v N ] (cid:62) , where the vectors q i , k i , v i for i = 1 , · · · , N are the query, key, and value vectors, respectively.

Original: Step 2.
Revised: Step 2.

Original: The output sequence H :=
Revised: The output sequence H :=

Original: [ h 1 , · · · , h N ] (cid:62) is then computed as follows
Revised: [ h 1 , · · · , h N ] (cid:62) is then computed as follows

Original: [Equation]
Revised: [Equation]

Original: where the softmax function is applied to each row of the matrix ( QK (cid:62) ) / √ D .
Revised: where the softmax function is applied to each row of the matrix ( QK (cid:62) ) / √ D .

Original: For each query vector q i for i = 1 , · · · , N , an equivalent form of Eqn.
Revised: For each query vector q i for i = 1 , · · · , N , an equivalent form of Eqn.

Original: (1) to compute the output vector
Revised: (1) to compute the output vector

Original: h
Revised: h

Original: i is given by
Revised: i is given by

Original: [Equation]
Revised: [Equation]

Original: The matrix A ∈ R N × N and its component a ij for i, j = 1 , · · · , N are the attention matrix and attention scores, respectively.
Revised: The matrix A ∈ R N × N and its component a ij for i, j = 1 , · · · , N are the attention matrix and attention scores, respectively.

Original: The self-attention computed by Eqn.
Revised: The self-attention computed by Eqn.

Original: (1) and (2) is called the scaled dot-product attention or softmax attention.
Revised: (1) and (2) is called the scaled dot-product attention or softmax attention.

Original: In our paper, we call a transformer that uses this attention the softmax transformer.
Revised: In our paper, we call a transformer that uses this attention the softmax transformer.

Original: The structure that the attention matrix A learns from training determines the ability of the self-attention to capture contextual representation for each token.
Revised: The structure that the attention matrix A learns from training determines the ability of the self-attention to capture contextual representation for each token.

Original: Multi-head Attention
Revised: Multi-head Attention

Original: Each output sequence H forms an attention head.
Revised: Each output sequence H forms an attention head.

Original: In multi-head attention, multiple heads are concatenated to compute the ﬁnal output.
Revised: In multi-head attention, multiple heads are concatenated to compute the ﬁnal output.

Original: Let H be the number of heads and
Revised: Let H be the number of heads and

Original: W O ∈ R HD v × HD v be the projection matrix for the output.
Revised: W O ∈ R HD v × HD v be the projection matrix for the output.

Original: The multi-head attention is deﬁned as
Revised: The multi-head attention is deﬁned as

Original: [Equation]
Revised: [Equation]

Original: Even though multi-head attention extends single-head attention to capture diverse attention patterns and improve the performance of transformers, it has been shown that transformers for practical tasks including sequence classiﬁcation and language modeling learn redundant heads (Michel et al., 2019).
Revised: Even though multi-head attention extends single-head attention to capture diverse attention patterns and improve the performance of transformers, it has been shown that transformers for practical tasks including sequence classiﬁcation and language modeling learn redundant heads (Michel et al., 2019).

Original: These redundant heads compute similar attention mappings.
Revised: These redundant heads compute similar attention mappings.

Original: 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 large-scale tasks.
Revised: 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 large-scale tasks.

Original: 1. C ONTRIBUTION
Revised: 1. C ONTRIBUTION

Original: 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-MGK), a novel class of transformers that can avoid the head redundancy.
Revised: 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-MGK), a novel class of transformers that can avoid the head redundancy.

Original: At the core of Transformer-MGK is replacing the attention key k j in each head by a GMM to allow the query q i , as well as 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.
Revised: At the core of Transformer-MGK is replacing the attention key k j in each head by a GMM to allow the query q i , as well as 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.

Original: Our contribution is four-fold:
Revised: Our contribution is four-fold:

Original: We construct a GMM and show that attention scores in self-attention match posterior distribution in our model, providing a probabilistic framework to study self-attention in transformers.
Revised: We construct a GMM and show that attention scores in self-attention match posterior distribution in our model, providing a probabilistic framework to study self-attention in transformers.

Original: Under our probabilistic framework for self-attention, we introduce an additional mixture of Gaussian to model each attention key.
Revised: Under our probabilistic framework for self-attention, we introduce an additional mixture of Gaussian to model each attention key.

Original: We empirically show that this mixture of Gaussian keys (MGK) can capture a diversity of attention patterns, thus alleviating head redundancy.
Revised: We empirically show that this mixture of Gaussian keys (MGK) can capture a diversity of attention patterns, thus alleviating head redundancy.

Original: We extend our MGK to use with linear attentions and propose the mixture of linear keys (MLK) for efﬁcient computation and better memory footprint.
Revised: We extend our MGK to use with linear attentions and propose the mixture of linear keys (MLK) for efﬁcient computation and better memory footprint.

Original: We empirically show that Transformer-MGK and Transformer-MLK 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 of parameters and computational cost in terms of FLOPs.
Revised: We empirically show that Transformer-MGK and Transformer-MLK 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 of parameters and computational cost in terms of FLOPs.

Original: Organization: We structure this paper as follows:
Revised: Organization: We structure this paper as follows:

Original: In Section 2, we establish the connection between GMM and self-attention and then present our Transformer-MGK and its extensions including
Revised: In Section 2, we establish the connection between GMM and self-attention and then present our Transformer-MGK and its extensions including

Original: Transformer-MLK.
Revised: Transformer-MLK.

Original: In Section 3, we validate and empirically analyze the efﬁciency and accuracy of
Revised: In Section 3, we validate and empirically analyze the efﬁciency and accuracy of

Original: Transformer-MGK/MLK.
Revised: Transformer-MGK/MLK.

Original: We discuss related works in Section 4.
Revised: We discuss related works in Section 4.

Original: The paper ends up with concluding remarks.
Revised: The paper ends up with concluding remarks.

Original: More experimental details are provided in the Appendix.
Revised: More experimental details are provided in the Appendix.

Original:  T RANSFORMER WITH A M IXTURE OF G AUSSIAN K EYS 2. A TTENTION S CORE AS A P OSTERIOR D ISTRIBUTION
Revised:  T RANSFORMER WITH A M IXTURE OF G AUSSIAN K EYS 2. A TTENTION S CORE AS A P OSTERIOR D ISTRIBUTION

Original: We ﬁrst consider a query q i ∈ Q and a key k j ∈ K .
Revised: We ﬁrst consider a query q i ∈ Q and a key k j ∈ K .

Original: Let t be a N -dimensional binary random variable having a 1-of- N representation in which a particular element t j is equal to 1 and all other elements are equal to 0.
Revised: Let t be a N -dimensional binary random variable having a 1-of- N representation in which a particular element t j is equal to 1 and all other elements are equal to 0.

Original: We use t j to indicate the position j of the key k j .
Revised: We use t j to indicate the position j of the key k j .

Original: In particular, let I be the identity matrix, we model the distribution p ( q i ) by the following GMM:
Revised: In particular, let I be the identity matrix, we model the distribution p ( q i ) by the following GMM:

Original: [Equation]
Revised: [Equation]

Original: where π j is the prior p ( t j = 1) .
Revised: where π j is the prior p ( t j = 1) .

Original: Given the query q i
Revised: Given the query q i

Original: [Equation]
Revised: [Equation]

Original: how likely q i matches the key k j is given by posterior p ( t j = 

Revised: how likely q i matches the key k j is given by posterior p ( t j =

Original: | q i ) .
Revised: | q i ) .

Original: This posterior is computed as follows
Revised: This posterior is computed as follows

Original: [Equation]
Revised: [Equation]

Original: We further assume that the query q i and the key k j are normalized, and the prior π j is uniform.
Revised: We further assume that the query q i and the key k j are normalized, and the prior π j is uniform.

Original: We will justify these assumptions in our Remarks at the end of this section.
Revised: We will justify these assumptions in our Remarks at the end of this section.

Original: We also let σ 2 j = σ 2 , j = 1 , 2 , . . .
Revised: We also let σ 2 j = σ 2 , j = 1 , 2 , . . .

Original: , K .
Revised: , K .

Original: Then the posterior p ( t j = 1 | q i ) can be written as
Revised: Then the posterior p ( t j = 1 | q i ) can be written as

Original: [Equation]
Revised: [Equation]

Original: The right-hand side of Eqn.
Revised: The right-hand side of Eqn.

Original: (6) matches the attention score given in Eqn.
Revised: (6) matches the attention score given in Eqn.

Original: (2) when σ 2 = √ D .
Revised: (2) when σ 2 = √ D .

Original: Thus, we show that under right assumptions, the attention score between the query q i and the key k j in an attention unit of a transformer is the posterior p ( t j = 1 | q i ) , which indicates the responsibility that the key k j takes for ‘explaining’ the query q i , which in turn decide, for example, how much a token at position i pays attention to a token at position j in the input sequence.
Revised: Thus, we show that under right assumptions, the attention score between the query q i and the key k j in an attention unit of a transformer is the posterior p ( t j = 1 | q i ) , which indicates the responsibility that the key k j takes for ‘explaining’ the query q i , which in turn decide, for example, how much a token at position i pays attention to a token at position j in the input sequence.

Original: Remark 1 The assumption that the query q i and the key k j are normalized is realistic and not artiﬁcial.
Revised: Remark 1 The assumption that the query q i and the key k j are normalized is realistic and not artiﬁcial.

Original: In many applications, those two vectors are normalized.
Revised: In many applications, those two vectors are normalized.

Original: Schlag et al.
Revised: Schlag et al.

Original: (2021) points out that such normalization is to avoid instability occurring during the training.
Revised: (2021) points out that such normalization is to avoid instability occurring during the training.

Original: Remark 

Revised: Remark

Original: In practice, the prior is chosen to be uniform when there is no prior knowledge available.
Revised: In practice, the prior is chosen to be uniform when there is no prior knowledge available.

Original: 2. T RANSFORMER WITH A M IXTURE OF G AUSSIAN K EYS : E ACH K EY IS A GAIN A G AUSSIAN M IXTURE M ODEL
Revised: 2. T RANSFORMER WITH A M IXTURE OF G AUSSIAN K EYS : E ACH K EY IS A GAIN A G AUSSIAN M IXTURE M ODEL

Original: As we have seen from Eqn.
Revised: As we have seen from Eqn.

Original: (6), the key k j is used to explain the query q i via the posterior P ( t j = 1 | q i ) .
Revised: (6), the key k j is used to explain the query q i via the posterior P ( t j = 1 | q i ) .

Original: Via this simple connection, each query q i is treated to be as a sample from the mixture of N keys (cid:80) Nj =1 π j N ( q i | k j , σ 2 j I d ) .
Revised: Via this simple connection, each query q i is treated to be as a sample from the mixture of N keys (cid:80) Nj =1 π j N ( q i | k j , σ 2 j I d ) .

Original: However, the assumption that the distribution P ( q i | t j = 1) at each subpopulation is Gaussian in Eqn.
Revised: However, the assumption that the distribution P ( q i | t j = 1) at each subpopulation is Gaussian in Eqn.

Original: (4) can be quite strong as there is no guarantee that this assumption is valid in practice.
Revised: (4) can be quite strong as there is no guarantee that this assumption is valid in practice.

Original: In particular, it may happen that the distribution of each subpopulation is asymmetric or skewed or even multimodal.
Revised: In particular, it may happen that the distribution of each subpopulation is asymmetric or skewed or even multimodal.

Original: Therefore, using the Gaussian distribution for each subpopulation can potentially limit the explanation power and diversity of each subpopulation/ key.
Revised: Therefore, using the Gaussian distribution for each subpopulation can potentially limit the explanation power and diversity of each subpopulation/ key.

Original: It indicates that we should use more expressive distributions to represent P ( q i | t j = 1) .
Revised: It indicates that we should use more expressive distributions to represent P ( q i | t j = 1) .

Original: Mixture of Gaussian keys:
Revised: Mixture of Gaussian keys:

Original: To improve the explanation power of each key k j , potentially increase the representation of each attention head, and reduce the chance of learning redundant heads, we would like to model it as mixture of Gaussian distributions.
Revised: To improve the explanation power of each key k j , potentially increase the representation of each attention head, and reduce the chance of learning redundant heads, we would like to model it as mixture of Gaussian distributions.

Original: We refer to this model as Transformer with a Mixture of Gaussian Keys (Transformer-MGK).
Revised: We refer to this model as Transformer with a Mixture of Gaussian Keys (Transformer-MGK).

Original: In particular, in Transformer-MGK we model each key k j at position j as a mixture of M Gaussians N ( k jr , σ 2 jr I ) , r = 1 , 2 , . . .
Revised: In particular, in Transformer-MGK we model each key k j at position j as a mixture of M Gaussians N ( k jr , σ 2 jr I ) , r = 1 , 2 , . . .

Original: , M .
Revised: , M .

Original: Here we are overloading the notation a little bit and use k jr and σ 2 jr I to denote the mean and covariance matrix of the r th Gaussian at position j .
Revised: Here we are overloading the notation a little bit and use k jr and σ 2 jr I to denote the mean and covariance matrix of the r th Gaussian at position j .

Original: Let z be a M -dimensional binary random variable having a 1-of- M representation.
Revised: Let z be a M -dimensional binary random variable having a 1-of- M representation.

Original: We use z r to indicate the r th Gaussian in the mixture.
Revised: We use z r to indicate the r th Gaussian in the mixture.

Original: Let π jr ≡ P ( z r = 1 | t j = 1) , our MGK can be written as
Revised: Let π jr ≡ P ( z r = 1 | t j = 1) , our MGK can be written as

Original: [Equation]
Revised: [Equation]

Original: Our motivation of using mixture of Gaussian distributions to represent the distribution of each subpopulation in Eqn.
Revised: Our motivation of using mixture of Gaussian distributions to represent the distribution of each subpopulation in Eqn.

Original: (7) stems from the following important approximation result:
Revised: (7) stems from the following important approximation result:

Original: Theorem 1 Assume that P is probability distribution on [ − a, a ] d for some a > 0 and admits density function p such that p is differentiable and bounded.
Revised: Theorem 1 Assume that P is probability distribution on [ − a, a ] d for some a > 0 and admits density function p such that p is differentiable and bounded.

Original: Then, for any given variance σ > 0 and for any (cid:15) > 0 , there exists a mixture of K components (cid:80) Ki
Revised: Then, for any given variance σ > 0 and for any (cid:15) > 0 , there exists a mixture of K components (cid:80) Ki

Original: =
Revised: =

Original: π i N ( θ i , σ 2 I ) where K ≤ ( C log(1 /(cid:15) ))
Revised: π i N ( θ i , σ 2 I ) where K ≤ ( C log(1 /(cid:15) ))

Original: d for some universal constant C such that
Revised: d for some universal constant C such that

Original: [Equation]
Revised: [Equation]

Original: where φ ( x | θ, σ 2 I ) is the density function of multivariate Gaussian distribution with mean θ and covariance matrix σ 2 I .
Revised: where φ ( x | θ, σ 2 I ) is the density function of multivariate Gaussian distribution with mean θ and covariance matrix σ 2 I .

Original: The proof of Theorem 1 is in Appendix C.
Revised: The proof of Theorem 1 is in Appendix C.

Original: The result of Theorem 1 suggests that regardless of the real form of P ( q i | t j = 1) , we can use ﬁnite mixture of Gaussian distributions to approximate P ( q i | t j = 1) .
Revised: The result of Theorem 1 suggests that regardless of the real form of P ( q i | t j = 1) , we can use ﬁnite mixture of Gaussian distributions to approximate P ( q i | t j = 1) .

Original: It allows us to have richer approximation of P ( q i | t j = 1) than by using a simple
Revised: It allows us to have richer approximation of P ( q i | t j = 1) than by using a simple

Original: Gaussian distribution in Eqn.
Revised: Gaussian distribution in Eqn.

Original: Similar to the derivation above, the posterior p ( t j = 1 | q i ) in
Revised: Similar to the derivation above, the posterior p ( t j = 1 | q i ) in

Original: Transformer-MGK can be written as
Revised: Transformer-MGK can be written as

Original: [Equation]
Revised: [Equation]

Original: Furthermore, in Transformer-MGK, we relax the assumption that the queries and keys are normalized.
Revised: Furthermore, in Transformer-MGK, we relax the assumption that the queries and keys are normalized.

Original: Thus, when computing P ( t j = 1 | q i ) , we compute the Gaussian kernels between the queries and keys instead of their dot products.
Revised: Thus, when computing P ( t j = 1 | q i ) , we compute the Gaussian kernels between the queries and keys instead of their dot products.

Original: The posterior P ( t j = 1 |
Revised: The posterior P ( t j = 1 |

Original: q
Revised: q

Original: i ) in Transformer-MGK is then given by
Revised: i ) in Transformer-MGK is then given by

Original: [Equation]
Revised: [Equation]

Original: As proven in Section 2.
Revised: As proven in Section 2.

Original: [Equation]
Revised: [Equation]

Original: this posterior corresponds to the attention score.
Revised: this posterior corresponds to the attention score.

Original: Thus, Eqn.
Revised: Thus, Eqn.

Original: (9) is the formula for computing the attention score in Transformer-MGK.
Revised: (9) is the formula for computing the attention score in Transformer-MGK.

Original: We compute the output vector h i of the self-attention in Transformer-MGK as follows
Revised: We compute the output vector h i of the self-attention in Transformer-MGK as follows

Original: [Equation]
Revised: [Equation]

Original: 2. I NFERENCE AND L EARNING VIA THE E XPECTATION M AXIMIZATION A LGORITHM
Revised: 2. I NFERENCE AND L EARNING VIA THE E XPECTATION M AXIMIZATION A LGORITHM

Original: Let γ ir ≡ P ( z r = 1 | q i , t j = 1) , in MGK, we apply the E-step inference in the ExpectationMaximization (EM) algorithm to estimate this posterior given the query q i .
Revised: Let γ ir ≡ P ( z r = 1 | q i , t j = 1) , in MGK, we apply the E-step inference in the ExpectationMaximization (EM) algorithm to estimate this posterior given the query q i .

Original: The posterior γ ir is also known as the responsibility that the component N ( k jr , σ 2 jr I ) takes to account for the observation, which in MGK is the query q i .
Revised: The posterior γ ir is also known as the responsibility that the component N ( k jr , σ 2 jr I ) takes to account for the observation, which in MGK is the query q i .

Original: Below we propose two approaches to estimate this responsibility.
Revised: Below we propose two approaches to estimate this responsibility.

Original: Soft E-step Using soft E-step inference, the EM algorithm makes a soft assignment, in which each query is associated with all clusters.
Revised: Soft E-step Using soft E-step inference, the EM algorithm makes a soft assignment, in which each query is associated with all clusters.

Original: The responsibilities are then given by
Revised: The responsibilities are then given by

Original: [Equation]
Revised: [Equation]

Original: At learning time, the responsibilities estimated by Eqn.
Revised: At learning time, the responsibilities estimated by Eqn.

Original: (11) are used to update the prior π jr , i.e. π jr = N jr /N , where N is the number of queries and N jr = (cid:80) Ni
Revised: (11) are used to update the prior π jr , i.e. π jr = N jr /N , where N is the number of queries and N jr = (cid:80) Ni

Original: =1 γ ir .
Revised: =1 γ ir .

Original: These updated priors π jr are then used in Eqn.
Revised: These updated priors π jr are then used in Eqn.

Original: (9) to compute attention scores.
Revised: (9) to compute attention scores.

Original: 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.
Revised: 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.

Original: This is similar to the K -means algorithm
Revised: This is similar to the K -means algorithm

Original: (Lloyd, 1982) and corresponds to the MGK at the limit when the variance parameter σ 2 jr goes to 0.
Revised: (Lloyd, 1982) and corresponds to the MGK at the limit when the variance parameter σ 2 jr goes to 0.

Original: Following the derivation of K -means from a GMM in (Bishop, 2006), Eqn.
Revised: Following the derivation of K -means from a GMM in (Bishop, 2006), Eqn.

Original: (9) becomes
Revised: (9) becomes

Original: [Equation]
Revised: [Equation]

Original: Remark 

Revised: Remark

Original: The hard E-step inference allows the attention score to be computed more efﬁciently because the priors π jr no longer play an active role in the algorithm and can be completely ignored.
Revised: The hard E-step inference allows the attention score to be computed more efﬁciently because the priors π jr no longer play an active role in the algorithm and can be completely ignored.

Original: Learning via Stochastic Gradient Descent (SGD)
Revised: Learning via Stochastic Gradient Descent (SGD)

Original: In order to increase the efﬁciency of the model, in MGK, we ﬁx the variance parameter σ 2 jr to be √ D as in the standard softmax attention and make the cluster means, i.e. the keys, k jr learnable parameters.
Revised: In order to increase the efﬁciency of the model, in MGK, we ﬁx the variance parameter σ 2 jr to be √ D as in the standard softmax attention and make the cluster means, i.e. the keys, k jr learnable parameters.

Original: We also make the prior π jr learnable parameters as one of the design options.
Revised: We also make the prior π jr learnable parameters as one of the design options.

Original: In that case, both k jr and π jr are learned via SGD.
Revised: In that case, both k jr and π jr are learned via SGD.

Original: This update via SGD can be considered as a generalized M-step (Bishop, 2006).
Revised: This update via SGD can be considered as a generalized M-step (Bishop, 2006).

Original: Design Options for Keys (Option A)
Revised: Design Options for Keys (Option A)

Original: We follow the standard setting in the softmax transformer and make the keys k jr a linear projection of the input x j , i.e. k jr = x j W (cid:62) K r , where x j ∈ R 1 × D x ,
Revised: We follow the standard setting in the softmax transformer and make the keys k jr a linear projection of the input x j , i.e. k jr = x j W (cid:62) K r , where x j ∈ R 1 × D x ,

Original: W K r ∈ R D × D x and r = 1 , 2 , . .
Revised: W K r ∈ R D × D x and r = 1 , 2 , . .

Original: , M .
Revised: , M .

Original: (Option B) Alternatively, we also make the keys k jr shifted version of each other to save computation, i.e. k jr = x j W (cid:62) K + b r , where W K ∈ R D
Revised: (Option B) Alternatively, we also make the keys k jr shifted version of each other to save computation, i.e. k jr = x j W (cid:62) K + b r , where W K ∈ R D

Original: × D x .
Revised: × D x .

Original: 2. T RANSFORMER WITH A M IXTURE OF L INEAR K EYS
Revised: 2. T RANSFORMER WITH A M IXTURE OF L INEAR K EYS

Original: The MGK can be easily extended to use with linear attentions.
Revised: The MGK can be easily extended to use with linear attentions.

Original: We call that model Transformer with a Mixture of Linear Keys
Revised: We call that model Transformer with a Mixture of Linear Keys

Original: (Transformer-MLK).
Revised: (Transformer-MLK).

Original: In this section, we adopt the formulation of linear attentions from (Katharopoulos et al., 2020) to derive Transformer-MLK.
Revised: In this section, we adopt the formulation of linear attentions from (Katharopoulos et al., 2020) to derive Transformer-MLK.

Original: Similar approach can be taken to derive Transformer-MLK when using with other linear attentions such as those in performers (Choromanski et al., 2021) and fast-weight transformers (Schlag et al., 2021).
Revised: Similar approach can be taken to derive Transformer-MLK when using with other linear attentions such as those in performers (Choromanski et al., 2021) and fast-weight transformers (Schlag et al., 2021).

Original: In
Revised: In

Original: Transformer-MLK, the Gaussian kernel in Eqn.
Revised: Transformer-MLK, the Gaussian kernel in Eqn.

Original: (10) is linearized as the product of feature maps φ ( · ) on the vectors q i and k j .
Revised: (10) is linearized as the product of feature maps φ ( · ) on the vectors q i and k j .

Original: The associative property of matrix multiplication is then utilized to derive the following efﬁcient computation of the attention map
Revised: The associative property of matrix multiplication is then utilized to derive the following efﬁcient computation of the attention map

Original: [Equation]
Revised: [Equation]

Original: Replacing (cid:80) j (cid:80) r π jr φ ( q i ) (cid:62) φ ( k jr ) v j with φ ( q i ) (cid:62) (cid:80) j (cid:80) r π jr φ ( k jr ) v (cid:62) j , as in linear transformers, reduces the memory and computational cost of computing the attention map in Transformer-MLK from O ( N 2 ) to O ( N ) , making Transformer-MLK scalable to very long sequences.
Revised: Replacing (cid:80) j (cid:80) r π jr φ ( q i ) (cid:62) φ ( k jr ) v j with φ ( q i ) (cid:62) (cid:80) j (cid:80) r π jr φ ( k jr ) v (cid:62) j , as in linear transformers, reduces the memory and computational cost of computing the attention map in Transformer-MLK from O ( N 2 ) to O ( N ) , making Transformer-MLK scalable to very long sequences.

Original:  E XPERIMENTAL R ESULTS
Revised:  E XPERIMENTAL R ESULTS

Original: In this section, we numerically justify the efﬁciency of Transformer-MGK/MLK and empirically study the advantage of using mixture of keys on various benchmarks, including different tasks in the Long
Revised: In this section, we numerically justify the efﬁciency of Transformer-MGK/MLK and empirically study the advantage of using mixture of keys on various benchmarks, including different tasks in the Long

Original: Range Arena (LRA) (Tay et al., 2021)
Revised: Range Arena (LRA) (Tay et al., 2021)

Original: (Section 3.1) and language modeling on Wikitext-103 (Merity et al., 2017)
Revised: (Section 3.1) and language modeling on Wikitext-103 (Merity et al., 2017)

Original: (Section 3.2).
Revised: (Section 3.2).

Original: We aim to show that: (i) Transformer-MGK/MLK with half the number of heads is comparable or better than the baseline softmax and linear transformers with the full number of heads while being more efﬁcient in both computational cost and memory footprints; (ii) Mixture of keys helps reduce the redundancy in multi-head transformers and beneﬁts learning of the long-term dependency in long input sequences; (iii) Using the same number of heads, Transformer-MGK/MLK signiﬁcantly outperforms the baseline softmax and linear transformers.
Revised: We aim to show that: (i) Transformer-MGK/MLK with half the number of heads is comparable or better than the baseline softmax and linear transformers with the full number of heads while being more efﬁcient in both computational cost and memory footprints; (ii) Mixture of keys helps reduce the redundancy in multi-head transformers and beneﬁts learning of the long-term dependency in long input sequences; (iii) Using the same number of heads, Transformer-MGK/MLK signiﬁcantly outperforms the baseline softmax and linear transformers.

Original: Especially in the case of
Revised: Especially in the case of

Original: Transformer-MLK, it helps reduce the performance gap between softmax and linear transformers while still maintaining linear memory and computational complexities.
Revised: Transformer-MLK, it helps reduce the performance gap between softmax and linear transformers.

Original: Throughout this section, we compare Transformer-MGK/MLK with the softmax and linear transformers that have the same or double the number of attention heads.
Revised: Throughout this section, we compare Transformer-MGK/MLK with the softmax and linear transformers that have the same or double the number of attention heads.

Original: In all experiments, for our
Revised: In all experiments, for our

Original: Transformer-MGK/MLK models, we set M=2 where M is the number of Gaussians, i.e. keys, at each timestep.
Revised: Transformer-MGK/MLK models, we set M=2 where M is the number of Gaussians, i.e. keys, at each timestep.

Original: Among the design options for Transformer-MGK mentioned in Section 2.3, we use the one with Soft-E step but make the parameter π jr and k jr learnable and ﬁx the variance σ 2 jr to be constants.
Revised: Among the design options for Transformer-MGK mentioned in Section 2.3, we use the one with Soft-E step but make the parameter π jr and k jr learnable and ﬁx the variance σ 2 jr to be constants.

Original: We study both implementations for keys: (A) k jr is a linear projection of the input x j , i.e., k jr = x j W (cid:62) K r and (B) k jr are shifted version of each other, i.e., k jr = x j W (cid:62) K + b r .
Revised: We study both implementations for keys: (A) k jr is a linear projection of the input x j , i.e., k jr = x j W (cid:62) K r and (B) k jr are shifted version of each other, i.e., k jr = x j W (cid:62) K + b r .

Original: In this section, we refer to the Transformer-MGK/MLK whose keys are implemented by (A) as
Revised: In this section, we refer to the Transformer-MGK/MLK whose keys are implemented by (A) as

Original: Transformer-MGK/MLK, and whose keys are implemented by (B) as Transformer-sMGK/
Revised: Transformer-MGK/MLK, and whose keys are implemented by (B) as Transformer-sMGK/

Original: sMLK.
Revised: sMLK.

Original: We empirically compare these models with other design options for Transformer-MGK in Section 3.4.
Revised: We empirically compare these models with other design options for Transformer-MGK in Section 3.4.

Original: All experiments are conducted on a server with 4 NVIDIA A100 GPUs.

Original: Details on datasets, models, and training are provided in Appendix A.1.
Revised: Details on datasets, models, and training are provided in Appendix A.1.

Original: 3. L ONG R ANGE A RENA (LRA) B ENCHMARK
Revised: 3. L ONG R ANGE A RENA (LRA) B ENCHMARK

Original: Models and baselines We compare our 1-head, 2-head, 4-head Transformer-MGK and MLK with the baseline softmax (Vaswani et al., 2017) and linear transformers (Katharopoulos et al., 2020) that 
Revised: Models and baselines We compare our 1-head, 2-head, 4-head Transformer-MGK and MLK with the baseline softmax (Vaswani et al., 2017) and linear transformers (Katharopoulos et al., 2020) that have 1 head, 2 heads, 4 heads, and 8 heads.

Revised: Each model consists of two layers, and we adopt the model and training setting from (Xiong et al., 2021) in our experiments.

Original: [Table]
Revised: [Table]

Revised: [Table]

Original: [Table]
Revised: [Table]

Original: have 1 head, 2 heads, 4 heads, and 8 heads.
Revised: Linear

Original: Each model consists of two layers, and we adopt the model and training setting from (Xiong et al., 2021) in our experiments.
Revised: [Table]

Original: Results We summarize our results in Table 1.
Revised: heads [Table] [Table] [Table]

Revised: [Table]

Revised: [Table]

Revised: [Table]

Revised: [Table]

Revised: Results We summarize our results in Table 1.

Original: Transformer-MGKs with half the number of heads consistently achieve better test accuracy than the baseline softmax attention across tasks.
Revised: Transformer-MGKs with half the number of heads consistently achieve better test accuracy than the baseline softmax attention across tasks.

Original: Since fewer heads are needed, transformer-MGKs use less parameters and need less FLOPs to compute than the baselines.
Revised: Since fewer heads are needed, transformer-MGKs use less parameters and need less FLOPs to compute than the baselines.

Original: We provide a detailed efﬁciency analysis for Transformer-MGKs in Figure 3.
Revised: We provide a detailed efﬁciency analysis for Transformer-MGKs in Figure 3.

Original: More interestingly, these efﬁciency advantages of Transformer-MGK over the baseline become more signiﬁcant as the number of heads in the baseline model grows.
Revised: More interestingly, these efﬁciency advantages of Transformer-MGK over the baseline become more signiﬁcant as the number of heads in the baseline model grows.

Original: When using the same number of heads as the baseline models, Transformer-MGKs further improve over those baselines.
Revised: When using the same number of heads as the baseline models, Transformer-MGKs further improve over those baselines.

Original: Among the models, Transformer-sMGK performs the best across LRA tasks.
Revised: Among the models, Transformer-sMGK performs the best across LRA tasks.

Original: We also compare the performance of Transformer-MLK with the baseline linear transformers in
Revised: We also compare the performance of Transformer-MLK with the baseline linear transformers in

Original: Table 2.
Revised: Table 2.

Original: Like Transformer-MGK, Transformer-MLK yields comparable or better results than the baseline using only half the number of heads with less parameters and FLOPs.
Revised: Like Transformer-MGK, Transformer-MLK yields comparable or better results than the baseline using only half the number of heads with less parameters and FLOPs.

Original: When using the same number of heads, Transformer-MLK helps improve the linear transformer further.
Revised: When using the same number of heads, Transformer-MLK helps improve the linear transformer further.

Original: We provide results of the 12-head baselines in Table 1 and 2 for reference.
Revised: We provide results of the 12-head baselines in Table 1 and 2 for reference.

Original: It is interesting to notice from Table 1 and 2 that even our 2-head Transformer-MGK/MLK models achieve better or equivalent results to the 12-head and 8-head baselines.
Revised: It is interesting to notice from Table 1 and 2 that even our 2-head Transformer-MGK/MLK models achieve better or equivalent results to the 12-head and 8-head baselines.

Original: A comparison between the 12-head baselines with our 6-head Transformer-MGK/MLK models on the retrieval task is provided in Table 8 in Appendix A.9.[Table]
Revised: A comparison between the 12-head baselines with our 6-head Transformer-MGK/MLK models on the retrieval task is provided in Table 8 in Appendix A.9.

Original: In Figure [Table] we compare the training loss and test accuracy curves of our 1-head and 2-head
Revised: In Figure 1, we compare the training loss and test accuracy curves of our 1-head and 2-head

Original: Transformer-MGK/MLK with the 2-head softmax and 2-head linear transformers on the document  retrieval task.
Revised: Transformer-MGK/MLK with the 2-head softmax and 2-head linear transformers on the document [Table] retrieval task.

Original: This retrieval task has the longest average sequence-length and attention span among the LRA tasks (Tay et al., 2021).
Revised: This retrieval task has the longest average sequence-length and attention span among the LRA tasks (Tay et al., 2021).

Original: On this task, as shown in Figure 1, our Transformer-MGKs/MLKs are always better than the baseline models throughout the training.
Revised: On this task, as shown in Figure [Table] our Transformer-MGKs/MLKs are always better than the baseline models throughout the training.

Original: This observation corroborates our models’s capability of capturing long-range dependencies in very long input sequences.
Revised: This observation corroborates our models’s capability of capturing long-range dependencies in very long input sequences.

Original: 3. L ANGUAGE M ODELING ON W IKI T EXT - 
Revised: 3. L ANGUAGE M ODELING ON W IKI T EXT -

Original: Next we conﬁrm the advantage of our models on a large-scale application.
Revised: Next we conﬁrm the advantage of our models on a large-scale application.

Original: We consider the word-level language modeling task on WikiText-103 (Merity et al., 2017) for our experiments in this section.
Revised: We consider the word-level language modeling task on WikiText-103 (Merity et al., 2017) for our experiments in this section.

Original: Models and baselines We compare 4 and 8-head Transformer-MGKs/MLKs with 8-head softmax (Vaswani et al., 2017) and linear transformers (Katharopoulos et al., 2020).
Revised: Models and baselines We compare 4 and 8-head Transformer-MGKs/MLKs with 8-head softmax (Vaswani et al., 2017) and linear transformers (Katharopoulos et al., 2020).

Original: Each model consists of 16 layers.
Revised: Each model consists of 16 layers.

Original: Our experiments follow the setting for small/medium models from (Schlag et al., 2021).
Revised: Our experiments follow the setting for small/medium models from (Schlag et al., 2021).

Original: Results As shown in Table 3, our Transformer-MGKs outperform the baseline softmax transformers.
Revised: Results As shown in Table 3, our Transformer-MGKs outperform the baseline softmax transformers.

Original: Even when usingonly half the number of attention heads (i.e., 4heads vs. 8 heads as in the baselines), the Transformer-MGK still achieves better test perplexities  than the baseline.
Revised: Even when using half the number of attention heads (i.e., 4 vs. 8 heads as in the baselines), the Transformer-MGK still achieves better test perplexities (PPL) than the baseline.

Original: Adding more heads intoour Transformer-MGKs improves their performance.
Revised: Adding more heads into Transformer-MGKs improves their performance.

Original: Similarly, Transformer-MLKs attain comparable test and validation perplexities to the baseline linear transformers when using half the number of attention heads.
Revised: Similarly, Transformer-MLKs attain comparable test/validation PPL to the baseline linear transformers when using half the number of attention heads.

Original: When using the same number of attention heads as in the baseline, Transformer-MLKs consistently achieve better performance.
Revised: When using the same number of attention heads as in the baseline, TransformerMLKs consistently achieve better performance.

Original: Note that reducing the number of heads from 8 to 4 in the baseline models signiﬁcantly decreases the performanceof these models with more than 1.5 reduction in validation and test PPL for the softmax transformer and more than 1. reduction in validation and test PPL for the linear transformer.
Revised: Note that reducing the number of heads from 8 to 4 in the baseline models signiﬁcantly 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.

Original: Our proposed Transformer-MGK and Transformer-MLK helps close this gap.
Revised: Our proposed Transformer-MGK and Transformer-MLK helps close this gap.

Original: To further examine the scalability of our models, we apply the MGK on a stronger baseline, which is the 8-head medium softmax transformer in (Schlag et al., 2021).
Revised: To further examine the scalability of our models, we apply the MGK on a stronger baseline, which is the 8-head medium softmax transformer in (Schlag et al., 2021).

Original: This model has 90M parameters,layers, 8 attention heads per layer, and hidden size of 256.
Revised: This model has 90M parameters,layers, 8 attention heads per layer, and hidden size of 256.

Original: The size of our baseline model is close to
Revised: The size of our baseline model is close to

Original: BERT Base (Devlin et al., 2019), which has 110M parameters, 12 layers, 12 attention heads per layer, and hidden size of 768.
Revised: BERT Base (Devlin et al., 2019), which has 110M parameters, 12 layers, 12 attention heads per layer, and hidden size of 768.

Original: Applying our MGK on top of this baseline and using only 4 heads instead of 8, we signiﬁcantly improve the test PPL from 29.60 to 28.86 while reducing the model size and computational cost, demonstrating the advantages of our scaled models.
Revised: Applying our MGK on top of this baseline and using only 4 heads instead of 8, we signiﬁcantly improve the test PPL from 29.60 to 28.86 while reducing the model size and computational cost, demonstrating the advantages of our scaled models.

Original: 3. N EURAL M ACHINE T RANSLATION ON IWSLT’14 G ERMAN TO E NGLISH
Revised: 3. N EURAL M ACHINE T RANSLATION ON IWSLT’14 G ERMAN TO E NGLISH

Original: We further examine the advantages of our methods on the IWSLT’14 German-English machine translation task (Cettolo et al., 2014).
Revised: We further examine the advantages of our methods on the IWSLT’14 German-English machine translation task (Cettolo et al., 2014).

Original: Table 9 in Appendix A.10 shows that the 2-head TransformerMGK 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.
Revised: Table 9 in Appendix A.10 shows that the 2-head TransformerMGK 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.

Original: 3. E MPIRICAL A NALYSIS
Revised: 3. E MPIRICAL A NALYSIS

Original: We conduct empirical analysis based on the Transformer-MGK trained for the document retrieval tasks.
Revised: We conduct empirical analysis based on the Transformer-MGK trained for the document retrieval tasks.

Original: Results for Transformer-MLKs and the WikiText-103 task are provided in the Appendix.
Revised: Results for Transformer-MLKs and the WikiText-103 task are provided in the Appendix.

Original: Transformer-MGK helps avoid learning redundant heads We visually compare attention matrices learned by Transformer-MGKs and the baseline softmax transformer on the document retrieval task 
Revised: Transformer-MGK helps avoid learning redundant heads We visually compare attention matrices learned by Transformer-MGKs and the baseline softmax transformer on the document retrieval task in Figure 2.

Revised: In particular, we randomly select an attention matrix at each head in each layer and visualize that attention matrix for each model in comparison.

Revised: Figure 2(Left) shows that the queries

Original: C o un t
Revised: C o un t

Original: L a y e r L a y e r
Revised: L a y e r L a y e r

Original: [Figure]
Revised: [Figure]

Original: in Figure 2.

Original: In particular, we randomly select an attention matrix at each head in each layer and visualize that attention matrix for each model in comparison.

Original: Figure 2(Left) shows that the queries in Transformer-MGKs can attend to a variety of keys and equivalently to other tokens at different positions in the input sequence.
Revised: in Transformer-MGKs can attend to a variety of keys and equivalently to other tokens at different positions in the input sequence.

Original: This diversity in attention pattern helps reduce the chance that the model learns similar and redundant attention matrices at different heads signiﬁcantly.
Revised: This diversity in attention pattern helps reduce the chance that the model learns similar and redundant attention matrices at different heads signiﬁcantly.

Original: Another metric to measure the representation capacity of an attention matrix is its rank.
Revised: Another metric to measure the representation capacity of an attention matrix is its rank.

Original: Attention matrices with high ranks can capture more diverse attention patterns compared to those with low ranks (Nguyen et al., 2021).
Revised: Attention matrices with high ranks can capture more diverse attention patterns compared to those with low ranks (Nguyen et al., 2021).

Original: We study the rank of the attention matrix from the Transformer-MGK and the softmax transformer trained for the retrieval task.
Revised: We study the rank of the attention matrix from the Transformer-MGK and the softmax transformer trained for the retrieval task.

Original: In particular, we randomly select 1000 different attention matrices at each layer from each model.
Revised: In particular, we randomly select 1000 different attention matrices at each layer from each model.

Original: Then, we perform singular value decomposition (SVD) to compute the rank of each matrix and threshold the singular values smaller than 10 − 6 .
Revised: Then, we perform singular value decomposition (SVD) to compute the rank of each matrix and threshold the singular values smaller than 10 − 6 .

Original: Figure 2(Right) presents the distribution of the rank of attention matrices at each layer of the
Revised: Figure 2(Right) presents the distribution of the rank of attention matrices at each layer of the

Original: Transformer-MGK and the softmax transformer.
Revised: Transformer-MGK and the softmax transformer.

Original: We observe that attention matrices in TransformerMGK has higher rank than those in the softmax transformer.
Revised: We observe that attention matrices in TransformerMGK has higher rank than those in the softmax transformer.

Original: Thus, our attention with MGK is capable of capturing more diverse and complex attention patterns than the baseline softmax attention.
Revised: Thus, our attention with MGK is capable of capturing more diverse and complex attention patterns than the baseline softmax attention.

Original: Transformer-MGK/MLK 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 Transformer-MGK/MLK that has half the number of heads and the full-head softmax/linear transformer.
Revised: Transformer-MGK/MLK 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 Transformer-MGK/MLK that has half the number of heads and the full-head softmax/linear transformer.

Original: The more heads being used, the more advantage Transformer-MGK/MLK has over the softmax/linear transformer.
Revised: The more heads being used, the more advantage Transformer-MGK/MLK has over the softmax/linear transformer.

Original: For much larger transformer models, this saving is signiﬁcant.
Revised: For much larger transformer models, this saving is signiﬁcant.

Original: Comparing different inference and learning techniques Table 4 compares the performance of
Revised: Comparing different inference and learning techniques Table 4 compares the performance of

Original: Transformer-MGKs using different design options mentioned in Section 2.3 on the LRA benchmark.
Revised: Transformer-MGKs using different design options mentioned in Section 2.3 on the LRA benchmark.

Original: In particular, we consider the following three design options: A) Soft-E step, parameters π jr and k jr are learnable via SGD, and variance σ 2 jr are constants, B) Soft-E step, parameter π jr is updated according to the M-step update, k jr are learnable via SGD, and variance σ 2 jr are constants, and
Revised: In particular, we consider the following three design options: A) Soft-E step, parameters π jr and k jr are learnable via SGD, and variance σ 2 jr are constants, B) Soft-E step, parameter π jr is updated according to the M-step update, k jr are learnable via SGD, and variance σ 2 jr are constants, and

Original: C) Hard-E step, π jr and k jr are learnable via SGD, and variance σ 2 jr are constants.
Revised: C) Hard-E step, π jr and k jr are learnable via SGD, and variance σ 2 jr are constants.

Original: Note that
Revised: Note that

Original: Transformer-MGKs with setting A are the default models we use in all experiments above.
Revised: Transformer-MGKs with setting A are the default models we use in all experiments above.

Original: In Table 4,
Revised: In Table 4,

Original: Transformer-MGK + Hard-E is the Transformer-MGK with setting C, Transformer-MGK + Soft-E is 
Revised: Transformer-MGK + Hard-E is the Transformer-MGK with setting C, Transformer-MGK + Soft-E is the Transformer-MGK with setting B, and Transformer-MGK only is the Transformer-MGK with

Original: [Table]
Revised: [Table]

Original: the Transformer-MGK with setting B, and Transformer-MGK only is the Transformer-MGK with setting A.
Revised: setting A.

Original: It is worth noting that Transformer-sMGK + Hard-E obtains comparable results to the models with the best performance in each task even though it is the most efﬁcient model in our study.
Revised: It is worth noting that Transformer-sMGK + Hard-E obtains comparable results to the models with the best performance in each task even though it is the most efﬁcient model in our study.

Original:  R ELATED W ORK
Revised:  R ELATED W ORK

Original: Efﬁcient Transformers Efﬁcient transformers can be classiﬁed into several categories, as summarized in (Roy et al., 2021).
Revised: Efﬁcient Transformers Efﬁcient transformers can be classiﬁed into several categories, as summarized in (Roy et al., 2021).

Original: 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).
Revised: 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).

Original: Another category includes models that combine two or more different access patterns to improve the coverage (Child et al., 2019; Ho et al., 2019).
Revised: Another category includes models that combine two or more different access patterns to improve the coverage (Child et al., 2019; Ho et al., 2019).

Original: Access patterns can also be made learnable in a data-driven fashion (Kitaev et al., 2020; Roy et al., 2021; Tay et al., 2020).
Revised: Access patterns can also be made learnable in a data-driven fashion (Kitaev et al., 2020; Roy et al., 2021; Tay et al., 2020).

Original: Other efﬁcient transformers take advantage of a side memory module to access multiple tokens at once (Lee et al., 2019;
Revised: Other efﬁcient transformers take advantage of a side memory module to access multiple tokens at once (Lee et al., 2019;

Original: Sukhbaatar et al., 2019; Asai & Choi, 2020; Beltagy et al., 2020).
Revised: Sukhbaatar et al., 2019; Asai & Choi, 2020; Beltagy et al., 2020).

Original: Finally, low-rank and kernelization approximation are utilized to enhance the memory and computational efﬁciency of computing self-attention (Tsai et al., 2019; Wang et al., 2020; Katharopoulos et al., 2020; Choromanski et al., 2021;
Revised: Finally, low-rank and kernelization approximation are utilized to enhance the memory and computational efﬁciency of computing self-attention (Tsai et al., 2019; Wang et al., 2020; Katharopoulos et al., 2020; Choromanski et al., 2021;

Original: Shen et al., 2021; Nguyen et al.,  2021).
Revised: Shen et al., 2021; Nguyen et al., 2021; Peng et al., 2021).

Revised: In addition to the aforementioned efﬁcient transformers, multi-query attention that shares keys and values between different attention heads (Shazeer, 2019) has also been studied to reduce the memory-bandwidth cost and increase the speed for incremental transformer inference (see Appendix A.11).

Original: In addition to the aforementioned efﬁcient transformers, 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.
Revised: 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.

Original: Our MGK/MLKare complementery to those efﬁcient transformers and training methods and can be easily incorporated into them to further improve their accuracy and efﬁciency.
Revised: Our MGK/MLK can be easily incorporated into the methods above to further improve their accuracy and efﬁciency.

Original: 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 et al., 2020; Michel et al., 2019; Durrani et al., 2020).
Revised: 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 et al., 2020; Michel et al., 2019; Durrani et al., 2020).

Original: 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).
Revised: 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).

Original: An emerging body of work is proposed to distill and prune the model, including (Sanh et al., 2019;
Revised: An emerging body of work is proposed to distill and prune the model, including (Sanh et al., 2019;

Original: Sun et al., 2019; Voita et al., 2019b; Sajjad et al., 2020).
Revised: Sun et al., 2019; Voita et al., 2019b; Sajjad et al., 2020).

Original: Our MGK/MLK approch can be combined with these distilling and pruning methods to improve their performance.
Revised: Our MGK/MLK approch can be combined with these distilling and pruning methods to improve their performance.

Original: Mixture Models for Transformers Several works have used mixture models to study and enhance transformers.
Revised: Mixture Models for Transformers Several works have used mixture models to study and enhance transformers.

Original: Switch transformers (Fedus et al., 2021) employ the routing algorithm in Mixture of
Revised: Switch transformers (Fedus et al., 2021) employ the routing algorithm in Mixture of

Original: Experts (MoE) to reduce the communication and computational costs in transformers.
Revised: Experts (MoE) to reduce the communication and computational costs in transformers.

Original: (Nguyen et al., 2018; Patel et al., 2016) derive a probablistic framework based on GMMs for deep neural networks that can be extended to study transformers and attention-based architectures.
Revised: (Nguyen et al., 2018; Patel et al., 2016) derive a probablistic framework based on GMMs for deep neural networks that can be extended to study transformers and attention-based architectures.

Original: Other works that use mixture models with transformers include (Cho et al., 2020; Guo et al., 2019; Jiang et al., 2020).
Revised: Other works that use mixture models with transformers include (Cho et al., 2020; Guo et al., 2019; Jiang et al., 2020).

Original:  C ONCLUDING R EMARKS
Revised:  C ONCLUDING R EMARKS

Original: In this paper, we proposed Transformer-MGK, a class of transformers that use Gaussian mixture model to represent the key vectors in self-attention.
Revised: In this paper, we proposed Transformer-MGK, a class of transformers that use Gaussian mixture model to represent the key vectors in self-attention.

Original: Transformer-MGK reduces the redundancy among heads in transformer.
Revised: Transformer-MGK reduces the redundancy among heads in transformer.

Original: Furthermore, attention heads in the Transformer-MGK have better representation capability than those in the baseline, allowing the Transformer-MGK to achieve comparable or better performance than the baseline softmax transformer while using only half of the number of heads.
Revised: Furthermore, attention heads in the Transformer-MGK have better representation capability than those in the baseline, allowing the Transformer-MGK to achieve comparable or better performance than the baseline softmax transformer while using only half of the number of heads.

Original: Comparing to the baseline, the Transformer-MGK uses fewer parameters and requires less FLOPs to compute.
Revised: Comparing to the baseline, the Transformer-MGK uses fewer parameters and requires less FLOPs to compute.

Original: We extend the Transformer-MGK into the Transformer-MLK to use linear attentions for better efﬁciency.
Revised: We extend the Transformer-MGK into the Transformer-MLK to use linear attentions for better efﬁciency.

Original: We empirically validate the advantage of the TransformerMGK/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.
Revised: We empirically validate the advantage of the TransformerMGK/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.

Original: In our work, we make the means and the variance of the cluster learnable variables and constants, respectively.
Revised: In our work, we make the means and the variance of the cluster learnable variables and constants, respectively.

Original: It is interesting to explore how to leverage the M-step update in the EM algorithm to update those parameters.
Revised: It is interesting to explore how to leverage the M-step update in the EM algorithm to update those parameters.

Original: Furthermore, we leave the application of Transformer-MGK/MLK for improving the vision transformer (Dosovitskiy et al., 2020; Touvron et al., 2020) as future work.
Revised: Furthermore, we leave the application of Transformer-MGK/MLK for improving the vision transformer (Dosovitskiy et al., 2020; Touvron et al., 2020) as future work.

Original: Reproducibility Statement: Source codes for our experiments are provided in the supplementary materials of the paper.
Revised: Reproducibility Statement: Source codes for our experiments are provided in the supplementary materials of the paper.

Original: The details of our experimental settings and computational infrastructure are given in Section 3 and the Appendix.
Revised: The details of our experimental settings and computational infrastructure are given in Section 3 and the Appendix.

Original: All datasets that we used in the paper are published, and they are easy to ﬁnd in the Internet.
Revised: All datasets that we used in the paper are published, and they are easy to ﬁnd in the Internet.

Original: Ethics Statement: Given the nature of the work, we do not foresee any negative societal and ethical impacts of our work.
Revised: Ethics Statement: Given the nature of the work, we do not foresee any negative societal and ethical impacts of our work.

Original: Supplement to “Transformer 
Revised: Supplement to “Transformer with a Mixture

Original: with a Mixture of Gaussian Keys”
Revised: of Gaussian Keys”

Original: In this supplementary material, we provide experimental details and additional experiments of
Revised: In this supplementary material, we provide experimental details and additional experiments of

Original: Transformer-MGK and Transformer-MLK.
Revised: Transformer-MGK and Transformer-MLK.

Original: A A DDITIONAL E XPERIMENTS A. E XPERIMENT DETAILS
Revised: A A DDITIONAL E XPERIMENTS A. E XPERIMENT DETAILS

Original: In this section, we provide model and training details for experiments in Section 3.
Revised: In this section, we provide model and training details for experiments in Section 3.

Revised: All our experiments are conducted on a server with 4 NVIDIA A100 GPUs.

Original: A.1. L ONG R ANGE A RENA B ENCHMARK
Revised: A.1. L ONG R ANGE A RENA B ENCHMARK

Original: Datasets and metrics We consider the following tasks in the LRA benchmark: Listops (Nangia & Bowman, 2018), byte-level IMDb reviews text classiﬁcation (Maas et al., 2011), and byte-level document retrieval (Radev et al., 2013).
Revised: Datasets and metrics We consider the following tasks in the LRA benchmark: Listops (Nangia & Bowman, 2018), byte-level IMDb reviews text classiﬁcation (Maas et al., 2011), and byte-level document retrieval (Radev et al., 2013).

Original: These tasks involve long sequences of length 2 K , 4 K , and 4 K , respectively.
Revised: These tasks involve long sequences of length 2 K , 4 K , and 4 K , respectively.

Original: 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.
Revised: 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.

Original: Models and baselines We use the softmax transformer (Vaswani et al., 2017) and linear transformer (Katharopoulos et al., 2020) as our baselines.
Revised: Models and baselines We use the softmax transformer (Vaswani et al., 2017) and linear transformer (Katharopoulos et al., 2020) as our baselines.

Original: All models have 2 layers, 64 embedding dimension, and 128 hidden dimension.
Revised: All models have 2 layers, 64 embedding dimension, and 128 hidden dimension.

Original: The number of heads in each layer are set to 1, 2, 4, and 8.
Revised: The number of heads in each layer are set to 1, 2, 4, and 8.

Original: For Transformer-MGK/MLKs and their shifted versions, we share π jr for all position j and learn it for each head.
Revised: For Transformer-MGK/MLKs and their shifted versions, we share π jr for all position j and learn it for each head.

Original: The initial value for each π jr is set to 0.5.
Revised: The initial value for each π jr is set to 0.5.

Original: For Transformer-sMGK, we learn b r and initialize its elements from a standard normal distribution.
Revised: For Transformer-sMGK, we learn b r and initialize its elements from a standard normal distribution.

Original: Each σ jr is a constant with value √ D where D is the dimension of each head, which is the same as in the baselines models
Revised: Each σ jr is a constant with value √ D where D is the dimension of each head, which is the same as in the baselines models

Original: Details about the Long Range Arena (LRA) benchmarks can be found in the original paper (Tay et al., 2021).
Revised: Details about the Long Range Arena (LRA) benchmarks can be found in the original paper (Tay et al., 2021).

Original: Our implementation is based on the public code by (Xiong et al., 2021), and we follow their training procedures.
Revised: Our implementation is based on the public code by (Xiong et al., 2021), and we follow their training procedures.

Original: The training setting and additional baseline model details are provided in the conﬁguration ﬁle used in (Xiong et al., 2021) and available at https:/
Revised: The training setting and additional baseline model details are provided in the conﬁguration ﬁle used in (Xiong et al., 2021) and available at https://github.com/mlpen/Nystromformer/blob/main/LRA/code.

Original: A.1. L ANGUAGE M ODELING ON W IKI T EXT - 
Revised: A.1. L ANGUAGE M ODELING ON W IKI T EXT -

Original: Datasets and metrics WikiText-103 consists of articles from Wikipedia and is a dataset with long contextual dependencies.
Revised: Datasets and metrics WikiText-103 consists of articles from Wikipedia and is a dataset with long contextual dependencies.

Original: The training set is made up of about 28 K articles containing 103 M running words; this corresponds to text blocks of about 3600 words.
Revised: The training set is made up of about 28 K articles containing 103 M running words; this corresponds to text blocks of about 3600 words.

Original: The validation and test sets are composed of 218 K and 246 K running words, respectively.
Revised: The validation and test sets are composed of 218 K and 246 K running words, respectively.

Original: Each of them contains 60 articles and about 268 K words.
Revised: Each of them contains 60 articles and about 268 K words.

Original: Our experiment follows the standard setting (Merity et al., 2017; Schlag et al., 2021) and splits the training data into L -word independent long segments.
Revised: Our experiment follows the standard setting (Merity et al., 2017; Schlag et al., 2021) and splits the training data into L -word independent long segments.

Original: For evaluation, we use a batch size of 1, and go through the text sequence with a sliding window of size L .
Revised: For evaluation, we use a batch size of 1, and go through the text sequence with a sliding window of size L .

Original: We consider only the last position for computing perplexity (PPL) except in the ﬁrst segment, where all positions are evaluated as in (Al-Rfou et al., 2019; Schlag et al., 2021).
Revised: We consider only the last position for computing perplexity (PPL) except in the ﬁrst segment, where all positions are evaluated as in (Al-Rfou et al., 2019; Schlag et al., 2021).

Original: Models and baselines Our language modeling implementation is based on the public code https://github.com/IDSIA/lmtool-fwp by (Schlag et al., 2021).
Revised: Models and baselines Our language modeling implementation is based on the public code https://github.com/IDSIA/lmtool-fwp by (Schlag et al., 2021).

Original: We use their small and medium model conﬁgurations for models in our experiments.
Revised: We use their small and medium model conﬁgurations for models in our experiments.

Original: In particular, for small models, we set the key, value, and query dimension to 128, and the training and evaluation context length to 256.
Revised: In particular, for small models, we set the key, value, and query dimension to 128, and the training and evaluation context length to 256.

Original: For medium models, we set the key, value, and query dimension to 256, and the training and evaluation context length to 384.
Revised: For medium models, we set the key, value, and query dimension to 256, and the training and evaluation context length to 384.

Original: In both conﬁgurations, we set the number of heads to 8, the feed-forward layer dimension to 2048, and the number of layers to 16.
Revised: In both conﬁgurations, we set the number of heads to 8, the feed-forward layer dimension to 2048, and the number of layers to 16.

Original: For Transformer-MGK/MLK, we use our 4 and 8-head versions to compare with the 8-head baselines.
Revised: For Transformer-MGK/MLK, we use our 4 and 8-head versions to compare with the 8-head baselines.

Original: π ir , b r and σ jr for this task follow our setings for LRA experiments.
Revised: π ir , b r and σ jr for this task follow our setings for LRA experiments.

Original: Other than those, our language modeling share the same conﬁgurations as the baselines.
Revised: Other than those, our language modeling share the same conﬁgurations as the baselines.

Original: 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.
Revised: 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.

Original: 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.
Revised: 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.

Original: A.1. N EURAL M ACHINE T RANSLATION ON IWSLT’14 G ERMAN TO E NGLISHDatasets and metrics
Revised: A.1. N EURAL M ACHINE T RANSLATION ON IWSLT’14 G ERMAN TO E NGLISH

Revised: Datasets and metrics

Original: The IWSLT14 German-English dataset (Cettolo et al., 2014) contains 153 K training, 7 K validation, and 7 K test TED-talks scripts German-English translated sentences.
Revised: The IWSLT14 German-English dataset (Cettolo et al., 2014) contains 153 K training, 7 K validation, and 7 K test TED-talks scripts German-English translated sentences.

Original: We follow the same preprocessing steps as in (Ott et al., 2019).
Revised: We follow the same preprocessing steps as in (Ott et al., 2019).

Original: We use the BiLingual Evaluation
Revised: We use the BiLingual Evaluation

Original: Understudy (BLEU) score as our evaluation metric for this task.
Revised: Understudy (BLEU) score as our evaluation metric for this task.

Original: All trained models are evaluated using the evaluation protocol in (Ott et al., 2019).
Revised: All trained models are evaluated using the evaluation protocol in (Ott et al., 2019).

Original: Models and baselines The baseline model we use in this machine translation task is an encoderdecoder transformer with six encoder/decoder layers, four attention heads per layer.
Revised: Models and baselines The baseline model we use in this machine translation task is an encoderdecoder transformer with six encoder/decoder layers, four attention heads per layer.

Original: The embedding dimension and the hidden size are 512 and 1024, respectively, for both encoder and decoder.
Revised: The embedding dimension and the hidden size are 512 and 1024, respectively, for both encoder and decoder.

Original: These architecture conﬁgurations are the same for Transformer-MGK/sMGK except for the number of attention heads per layer, which is reduced by half.
Revised: These architecture conﬁgurations are the same for Transformer-MGK/sMGK except for the number of attention heads per layer, which is reduced by half.

Original: We share π jr across all position j and learn it for each head.
Revised: We share π jr across all position j and learn it for each head.

Original: The initial value for each π jr is set to 0.5.
Revised: The initial value for each π jr is set to 0.5.

Original: For Transformer-sMGK, we learn b r and initialize its elements from a standard normal distribution.
Revised: For Transformer-sMGK, we learn b r and initialize its elements from a standard normal distribution.

Original: Each σ jr is a constant with the value √
Revised: Each σ jr is a constant with the value √

Original: D where D is the dimension of each head, which is the same as in the baselines models.
Revised: D where D is the dimension of each head, which is the same as in the baselines models.

Original: Our experiments follow the settings from (Ott et al., 2019), and our implementation is based on the public code https://github.com/pytorch/fairseq/tree/main/examples/translation.
Revised: Our experiments follow the settings from (Ott et al., 2019), and our implementation is based on the public code https://github.com/pytorch/fairseq/tree/main/examples/translation.

Original: A. M ORE E MPIRICAL A NALYSIS OF T RANSFORMER -MGK S /MLK S TRAINED FOR L ANGUAGE M ODELING
Revised: A. M ORE E MPIRICAL A NALYSIS OF T RANSFORMER -MGK S /MLK S TRAINED FOR L ANGUAGE M ODELING

Original: In Table 3 in the main text, we show the improvements in valid and test perplexity of our TransformerMGKs/MLKs compared with the baseline softmax and linear transformers.
Revised: In Table 3 in the main text, we show the improvements in valid and test perplexity of our TransformerMGKs/MLKs compared with the baseline softmax and linear transformers.

Original: In particular, TransformerMGK/MLKs with the same number of heads as the baselines, e.g. 8 heads, signiﬁcantly improve the baselines during training while Transformer-MGK/MLKs with 4 head achieve comparable or better performance than the 8-head baselines.
Revised: In particular, TransformerMGK/MLKs with the same number of heads as the baselines, e.g. 8 heads, signiﬁcantly improve the baselines during training while Transformer-MGK/MLKs with 4 head achieve comparable or better performance than the 8-head baselines.

Original: In this section, we provide more empirical analysis to shed light on those results.
Revised: In this section, we provide more empirical analysis to shed light on those results.

Original: Figure 4 shows the validation perplexity curve of our models versus the softmax and linear transformers.
Revised: Figure 4 shows the validation perplexity curve of our models versus the softmax and linear transformers.

Original: Figure 5 visualizes the attention matrices from a randomly selected sample for the trained softmax transformer with 8 heads and Transformer-MGKs with 8 heads and 4 heads.
Revised: Figure 5 visualizes the attention matrices from a randomly selected sample for the trained softmax transformer with 8 heads and Transformer-MGKs with 8 heads and 4 heads.

Original: These visualizations show that Transformer-MGKs attend to more diverse positions in all heads and layers than the softmax attention baseline.
Revised: These visualizations show that Transformer-MGKs attend to more diverse positions in all heads and layers than the softmax attention baseline.

Original: We also compare the rank distributions of these attention matrices computed from 1000 samples at each layer in the model.
Revised: We also compare the rank distributions of these attention matrices computed from 1000 samples at each layer in the model.

Original: Figure 6 presents the rank histograms of the 8-head softmax attention and 4-head and 8-head MGK attentions for the 1 st and 5 th layer.
Revised: Figure 6 presents the rank histograms of the 8-head softmax attention and 4-head and 8-head MGK attentions for the 1 st and 5 th layer.

Original: It is clear that attention matrices from the Transformer-MGKs 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.
Revised: It is clear that attention matrices from the Transformer-MGKs 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.

Original: . A. A DDITIONAL TRAINING RESULTS FOR LRA
Revised: . A. A DDITIONAL TRAINING RESULTS FOR LRA

Original: In this section, we provide additional experimental results on the LRA benchmark.
Revised: In this section, we provide additional experimental results on the LRA benchmark.

Original: 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 Transformer-MGK.
Revised: 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 Transformer-MGK.

Original: The computational costs of
Revised: The computational costs of

Original: Transformer-MGK/sMGK and Transformer-MGK Hard-E/Soft-E are on par with each other, while
Revised: Transformer-MGK/sMGK and Transformer-MGK Hard-E/Soft-E are on par with each other, while

Original: Transformer-sMGK uses fewer parameters than the other without trade-off in performance 4 for all tasks.
Revised: Transformer-sMGK uses fewer parameters than the other without trade-off in performance 4 for all tasks.

Original: The naming is as explained in Section 3.4 in the main text.
Revised: The naming is as explained in Section 3.4 in the main text.

Original: In addition, Figure 8 visualizes the attention matrices in the 4-head linear transformer baseline, 4-head Transformer-MLK, and 2-head
Revised: In addition, Figure 8 visualizes the attention matrices in the 4-head linear transformer baseline, 4-head Transformer-MLK, and 2-head

Original: Transformer-MLK trained on the document retrieval task.
Revised: Transformer-MLK trained on the document retrieval task.

Original: [Figure] [Figure]
Revised: [Figure] [Figure]

Original: A. A BLATION S TUDY ON THE I MPACT OF THE M IXTURE OF K EYS , THE G AUSSIAN D ISTANCE , AND THE K EY S HIFTING
Revised: A. A BLATION S TUDY ON THE I MPACT OF THE M IXTURE OF K EYS , THE G AUSSIAN D ISTANCE , AND THE K EY S HIFTING

Original: In this section we conduct an ablation study of the Transformer-MGK on the LRA retrieval task to investigate where the performance improvement is from.
Revised: In this section we conduct an ablation study of the Transformer-MGK on the LRA retrieval task to investigate where the performance improvement is from.

Original: 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.
Revised: 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.

Original: We summarize our empirical results in Table 5 and discuss the impact of 1, 2 and 3 below.
Revised: We summarize our empirical results in Table 5 and discuss the impact of 1, 2 and 3 below.

Original: 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.
Revised: 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.

Original: We name this model Softmax MGK.
Revised: We name this model Softmax MGK.

Original: We compare the Softmax MGK that has 1 head (Sofmax
Revised: We compare the Softmax MGK that has 1 head (Sofmax

Original: MGK 1 head in Table 5) with the baseline softmax transformers that use 1 and 2 heads (Softmax 1 head and Softmax 2 heads in Table 5).
Revised: MGK 1 head in Table 5) with the baseline softmax transformers that use 1 and 2 heads (Softmax 1 head and Softmax 2 heads in Table 5).

Original: Results in Table 5 show that the Softmax MGK 1 head outperforms both the baseline softmax transformers of 1 and 2 heads.
Revised: Results in Table 5 show that the Softmax MGK 1 head outperforms both the baseline softmax transformers of 1 and 2 heads.

Original: Note that our Softmax MGK
Revised: Note that our Softmax MGK

Original: [Figure]
Revised: [Figure]

Original: [Figure]
Revised: [Figure]

Original: 5k 4K 1K 2.5K 4K 1K 2.5k 

Revised: 5k 4K 1K 2.5K 4K 1K 2.5k

Original: K h e a d1K 2.5K 4K 1K
Revised: K h e a d1K 2.5K 4K 1K

Original: 2.5k 4K 1K 2.5K 4K 1K 2.5k 

Revised: 2.5k 4K 1K 2.5K 4K 1K 2.5k

Original: K h e a d1K 2.5K 4K 1K 2.5k 4K 0.
Revised: K h e a d1K 2.5K 4K 1K 2.5k 4K 0.

Original: [Figure]
Revised: [Figure]

Original: 1 head is more efﬁcient than the baseline of 2 heads in terms of the number of parameters and the number of FLOPs.
Revised: 1 head is more efﬁcient than the baseline of 2 heads in terms of the number of parameters and the number of FLOPs.

Original: These results conﬁrm the beneﬁt of using MGK.
Revised: These results conﬁrm the beneﬁt of using MGK.

Original: Impact of using the Gaussian distance Next, we compare the softmax MGK with the Gaussian
Revised: Impact of using the Gaussian distance Next, we compare the softmax MGK with the Gaussian

Original: MGK.
Revised: MGK.

Original: Here the Gaussian MGK is the Transformer-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.
Revised: Here the Gaussian MGK is the Transformer-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.

Original: Results in Table 5 suggest that the Gaussian MGK 1 head improves over the Softmax MGK 1 head (80.63% vs. 79.23%).
Revised: Results in Table 5 suggest that the Gaussian MGK 1 head improves over the Softmax MGK 1 head (80.63% vs. 79.23%).

Original: This result justiﬁes the advantage of using Gaussian distance over dot product to compute the attention scores.
Revised: This result justiﬁes the advantage of using Gaussian distance over dot product to compute the attention scores.

Original: Impact of key shifting Finally, we apply key shifting to both Softmax MGK and Gaussian MGK (Softmax sMGK and Gaussian sMGK in Table 5).
Revised: Impact of key shifting Finally, we apply key shifting to both Softmax MGK and Gaussian MGK (Softmax sMGK and Gaussian sMGK in Table 5).

Original: From Table 5, we observe that the Softmax sMGK 1 head and Gaussian sMGK 1 head outperform the Softmax MGK 1 head and Gaussian MGK 1 head, respectively.
Revised: From Table 5, we observe that the Softmax sMGK 1 head and Gaussian sMGK 1 head outperform the Softmax MGK 1 head and Gaussian MGK 1 head, respectively.

Original: These results, again, corroborate the beneﬁt of using key shifting.
Revised: These results, again, corroborate the beneﬁt of using key shifting.

Original: We also include the result for the Gaussian 1 head model in Table 5.
Revised: We also include the result for the Gaussian 1 head model in Table 5.

Original: This Gaussian  is similar to the Softmax 1 head model but uses the Gaussian distance to compute the attention scores.
Revised: This Gaussian 1 head model is similar to the Softmax 1 head model but uses the Gaussian distance to compute the attention scores.

Original: Comparing the results of the Gaussian 1 head model, the Softmax 1 head model, and the Gaussian
Revised: Comparing the results of the Gaussian 1 head model, the Softmax 1 head model, and the Gaussian

Original: [Table] [Table]
Revised: [Table] [Table]

Original: which has the longest average sequence-length and attention span among the LRA tasks (Tay et al., 2021).
Revised: which has the longest average sequence-length and attention span among the LRA tasks (Tay et al., 2021).

Original: In training, we apply early stopping to avoid overﬁtting.
Revised: In training, we apply early stopping to avoid overﬁtting.

Original: That explains why the training loss and accuracy curves stop early.
Revised: That explains why the training loss and accuracy curves stop early.

Original: The test loss/accuracy curves already converge.
Revised: The test loss/accuracy curves already converge.

Original: The impressive performance of Transformer-MGK/MLK on this challenging task validates the capability of our models to capture long-range dependencies via learning a diversity of attention patterns.
Revised: The impressive performance of Transformer-MGK/MLK on this challenging task validates the capability of our models to capture long-range dependencies via learning a diversity of attention patterns.

Original: MGK 1 head model reported in Table 5 further conﬁrms the advantage of using MGK and Gaussian distance.
Revised: MGK 1 head model reported in Table 5 further conﬁrms the advantage of using MGK and Gaussian distance.

Original: A. T IME O VERHEAD AND M EMORY F OOTPRINT A NALYSIS
Revised: A. T IME O VERHEAD AND M EMORY F OOTPRINT A NALYSIS

Original: Table 6 compares the GPU memory footprint and computational time overhead (seconds/iteration) of our 4-head Transformer-MGKs/MLKs with those of the 8-head softmax/linear transformer baselines at test time.
Revised: Table 6 compares the GPU memory footprint and computational time overhead (seconds/iteration) of our 4-head Transformer-MGKs/MLKs with those of the 8-head softmax/linear transformer baselines at test time.

Original: All models are trained on the LRA retrieval task, and we use a batch size of 32 for each iteration.
Revised: All models are trained on the LRA retrieval task, and we use a batch size of 32 for each iteration.

Original: Our MGK/MLK models save memory and reduce wall-clock time signiﬁcantly compared to the baselines.
Revised: Our MGK/MLK models save memory and reduce wall-clock time signiﬁcantly compared to the baselines.

Original: Using key shifting in our models, i.e. Transformer-sMGKs/sMLKs, helps improve the efﬁciency further.
Revised: Using key shifting in our models, i.e. Transformer-sMGKs/sMLKs, helps improve the efﬁciency further.

Original: A. L EARNING C URVES S HOWING C ONVERGENCE
Revised: A. L EARNING C URVES S HOWING C ONVERGENCE

Original: In this section, we replot Figure 1 and 4 to show the convergence of our trainings.
Revised: In this section, we replot Figure 1 and 4 to show the convergence of our trainings.

Original: In Figure 1, the results do not seem to converge since we plot the loss and the accuracy in log-scale.
Revised: In Figure 1, the results do not seem to converge since we plot the loss and the accuracy in log-scale.

Original: In Figure 4, the
Revised: In Figure 4, the

Original: [Figure]
Revised: [Figure]

Original: H ea d H ea d
Revised: H ea d H ea d

Original: [Figure]
Revised: [Figure]

Original: e i g h t
Revised: e i g h t

Original: K e y
Revised: K e y

Original: W e i g h t
Revised: W e i g h t

Original: K e y
Revised: K e y

Original: [Figure]
Revised: [Figure]

Original: i
Revised: i

Original: [Figure]
Revised: [Figure]

Original: y
Revised: y

Original: [Figure]
Revised: [Figure]

Original: [Table]
Revised: [Table]

Original: [Figure]
Revised: [Figure]

Original: [Table]
Revised: [Table]

Original: results do not seem to converge since we zoom into the speciﬁc range on the y and x-axes.
Revised: results do not seem to converge since we zoom into the speciﬁc range on the y and x-axes.

Original: Figureand 10 are the replotted versions of Figure 1 and 4, respectively.
Revised: Figureand 10 are the replotted versions of Figure 1 and 4, respectively.

Original: In Figure 9, the training loss/accuracy curves stop early because we use early stopping to avoid overﬁtting.
Revised: In Figure 9, the training loss/accuracy curves stop early because we use early stopping to avoid overﬁtting.

Original: The test loss/accuracy curves in this ﬁgure already converge.
Revised: The test loss/accuracy curves in this ﬁgure already converge.

Original: A. W EIGHT M ATRICES OF THE K EYS , K EYS AND M IXING C OEFFICIENT
Revised: A. W EIGHT M ATRICES OF THE K EYS , K EYS AND M IXING C OEFFICIENT

Original: In this section, we analyze the learned π jr , k jr , and W K r , j = 1 , . . . , N and r = 1 , . . .
Revised: In this section, we analyze the learned π jr , k jr , and W K r , j = 1 , . . . , N and r = 1 , . . .

Original: , M in the Transformer-MGK trained on the LRA retrieval task.
Revised: , M in the Transformer-MGK trained on the LRA retrieval task.

Original: In all of our experiments, we set M = 2 .
Revised: In all of our experiments, we set M = 2 .

Original: In Figure 11 and 12, we visualize the weight matrices W K that computes the keys and the keys
Revised: In Figure 11 and 12, we visualize the weight matrices W K that computes the keys and the keys

Original: K , respectively, for all heads and layers in the 2-head softmax transformer baseline, the 1-head
Revised: K , respectively, for all heads and layers in the 2-head softmax transformer baseline, the 1-head

Original: Transformer-MGK with 2 keys, and the 1-head Transformer-sMGK with 2 keys trained on the LRA retrieval task.
Revised: Transformer-MGK with 2 keys, and the 1-head Transformer-sMGK with 2 keys trained on the LRA retrieval task.

Original: Note that for the keys, we only plot the ﬁrst 100 tokens.
Revised: Note that for the keys, we only plot the ﬁrst 100 tokens.

Original: Also, Table 7 summarizes the
Revised: Also, Table 7 summarizes the

Original: Layer Layer
Revised: Layer Layer

Original: Layer
Revised: Layer

Original: Softmax 2 heads
Revised: Softmax 2 heads

Original: H ea d
Revised: H ea d

Original: H ea d
Revised: H ea d

Original: MGK 1 Head sMGK 1 Head
Revised: MGK 1 Head sMGK 1 Head

Original: K e y
Revised: K e y

Original: K e y K e y
Revised: K e y K e y

Original: Layer Layer
Revised: Layer Layer

Original: Layer
Revised: Layer

Original: K e y learned mixing coefﬁcient π jr of all heads and layers in the 1-head Transformer-MGK trained on the same retrieval task.
Revised: K e y learned mixing coefﬁcient π jr of all heads and layers in the 1-head Transformer-MGK trained on the same retrieval task.

Original: Here we use the same π j 1 , π j 2 for all time step j = 1 , . .
Revised: Here we use the same π j 1 , π j 2 for all time step j = 1 , . .

Original: . , N .
Revised: . , N .

Original: A. A DDITIONAL C OMPUTATIONAL C OMPLEXITY (FLOP S ) A NALYSIS AT T RAINING T IME
Revised: A. A DDITIONAL C OMPUTATIONAL C OMPLEXITY (FLOP S ) A NALYSIS AT T RAINING T IME

Original: Figure 13 demonstrates the computational cost (FLOPs) for each training iteration of the TransformerMGK vs. the baseline softmax transformer (Left) and the Transformer-MLK vs. the baseline linear transformer (Right) on the LRA retrieval task.
Revised: Figure 13 demonstrates the computational cost (FLOPs) for each training iteration of the TransformerMGK vs. the baseline softmax transformer (Left) and the Transformer-MLK vs. the baseline linear transformer (Right) on the LRA retrieval task.

Original: The efﬁciency advantage of Transformer-MGK/MLK over the baselines grows with the number of heads.
Revised: The efﬁciency advantage of Transformer-MGK/MLK over the baselines grows with the number of heads.

Original: Figure 14 shows the computational cost per training iteration (measured in FLOPs) of different inference and learning methods for Transformer-MGK trained on the document retrieval task.
Revised: Figure 14 shows the computational cost per training iteration (measured in FLOPs) of different inference and learning methods for Transformer-MGK trained on the document retrieval task.

Original: TransformersMGK, Transformer-MGK, Transformer-MGK Hard-E, and Soft-E have similar computational costs.
Revised: TransformersMGK, Transformer-MGK, Transformer-MGK Hard-E, and Soft-E have similar computational costs.

Original: A. S CALING TO 12-H EAD B ASELINE M ODELS FOR THE R ETRIEVAL TASK .
Revised: A. S CALING TO 12-H EAD B ASELINE M ODELS FOR THE R ETRIEVAL TASK .

Original: To further study the scalability of our model, in this section, we investigate the performance of our 6-head Transformer-MGKs/MLKs in comparison with the 12-head baseline softmax/linear transformers on the retrieval task.
Revised: To further study the scalability of our model, in this section, we investigate the performance of our 6-head Transformer-MGKs/MLKs in comparison with the 12-head baseline softmax/linear transformers on the retrieval task.

Original: Table 8 indicates that our 6-head Transform-MGKs/MLKs signiﬁcantly outperform 12-head softmax/linear transformers, respectively.
Revised: Table 8 indicates that our 6-head Transform-MGKs/MLKs signiﬁcantly outperform 12-head softmax/linear transformers, respectively.

Original: Moreover, comparing these results to those in Table 1 and Table 2, although the 12-head softmax/linear 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/MLK models.
Revised: Moreover, comparing these results to those in Table 1 and Table 2, although the 12-head softmax/linear 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/MLK models.

Original: A. N EURAL M ACHINE T RANSLATION ON IWSLT’14 G ERMAN TO E NGLISH
Revised: A. N EURAL M ACHINE T RANSLATION ON IWSLT’14 G ERMAN TO E NGLISH

Original: Table 9 shows that the 2-head Transformer-MGK/sMGK models have comparable or better BLEU scores than the 4-head softmax transformer baseline.
Revised: Table 9 shows that the 2-head Transformer-MGK/sMGK models have comparable or better BLEU scores than the 4-head softmax transformer baseline.

Original: [Figure]
Revised: [Figure]

Revised: [Table]

Original: [Table]
Revised: [Table]

Revised: A. C OMPARISON TO M ULTI - QUERY A TTENTION

Revised: In this section

Original: [Table]
Revised: [Table]

Original: B A N A NALYSIS ON THE C OMPUTATIONAL C OMPLEXITY AND THE N UMBER OF P ARAMETERS IN T RANSFORMER -MGK AND THE S OFTMAX T RANSFORMER
Revised: we compare our MGK approach to the multi-query attention (Shazeer, 2019). The multi-query attention shares the same set of keys and values at different heads to reduce the memorybandwidth cost during incremental inference, which allows faster inference since the size of the reloaded ”keys” and ”values” tensors are signiﬁcantly reduced. On another hand, our TransformerMGK models the key at each head as a Gaussian mixture model, which leads to the use of multiple keys at each head and allows us to decrease the number of attention heads. This helps reduce the computations and parameters needed to calculate additional queries and values.

Revised: 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 to calculate additional keys.

Revised: The advantages of

Revised: Transformer-MGK hold in both training and inference, including incremental inference.

Revised: We have provided a detailed analysis on the computational complexity and the number of parameters of the Transformer-MGK in comparison with the corresponding softmax transformer in Appendix B.

Revised: Combining the multi-query attention and our MGK approach is interesting since each method has its own advantage and they are complementary to each other.

Revised: In particular, we can let the transformer model share the same set of values and mixtures of keys at different heads.

Revised: This approach can potentially have the advantages of both multi-query attention and MGK.

Revised: B A N A NALYSIS ON THE C OMPUTATIONAL C OMPLEXITY AND THE N UMBER OF P ARAMETERS IN T RANSFORMER -MGK AND THE S OFTMAX T RANSFORMER

Original: In this section, we compare the computational complexity and the number of parameters in transformer-MGK with M keys at each timestep and H/M heads to the baseline softmax transformer that has 1 key at each timestep and H heads.
Revised: In this section, we compare the computational complexity and the number of parameters in transformer-MGK with M keys at each timestep and H/M heads to the baseline softmax transformer that has 1 key at each timestep and H heads.

Original: Here we choose M such that H is a multiple of
Revised: Here we choose M such that H is a multiple of

Original: M and use keys design option (A) that make the key k jr a linear projection of the input x j , i.e. k jr = x j W (cid:62) K r , where x j ∈ R 1 × D x , W K r ∈ R D × D x and r = 1 , 2 , . .
Revised: M and use keys design option (A) that make the key k jr a linear projection of the input x j , i.e. k jr = x j W (cid:62) K r , where x j ∈ R 1 × D x , W K r ∈ R D × D x and r = 1 , 2 , . .

Original: , M (see Design Options for
Revised: , M (see Design Options for

Original: Keys in Section 2.3 for more details).
Revised: Keys in Section 2.3 for more details).

Original: To simplify notation and without loss of generalization, we let M = 2 as in our experiments and assume that D v = D , i.e., the values have the same feature dimension as the queries and the keys.
Revised: To simplify notation and without loss of generalization, we let M = 2 as in our experiments and assume that D v = D , i.e., the values have the same feature dimension as the queries and the keys.

Original: To simplify the computation, we also do not take the softmax operator into account since this operator yields similar costs when applied in Transformer-MGK and softmax transformer.
Revised: To simplify the computation, we also do not take the softmax operator into account since this operator yields similar costs when applied in Transformer-MGK and softmax transformer.

Original: B. C OMPUTATIONAL C OMPLEXITY
Revised: B. C OMPUTATIONAL C OMPLEXITY

Original: (i) Softmax H -head attention:
Revised: (i) Softmax H -head attention:

Original: The

Original:  number of computations in a softmax H -head attention is N 2 H (4 D − 1)
Revised: The number of computations in a softmax H -head attention is N 2 H (4 D − 1)

Original: + NHD (6 D x + 2 HD − 5) .
Revised: + NHD (6 D x + 2 HD − 5) .

Original: Explanation: To calculate the query matrix Q , the key matrix K , and the value matrix V in Step 1 in
Revised: Explanation: To calculate the query matrix Q , the key matrix K , and the value matrix V in Step 1 in

Original: Section 1.1 at each head, we need 3 NDD x multiplications and 3 ND ( D x − 1) additions.
Revised: Section 1.1 at each head, we need 3 NDD x multiplications and 3 ND ( D x − 1) additions.

Original: In total, these need 3 ND (2 D x − 1) computations.
Revised: In total, these need 3 ND (2 D x − 1) computations.

Original: Next, to compute the product QK (cid:62) in Eqn.
Revised: Next, to compute the product QK (cid:62) in Eqn.

Original: (1), we need N 2 D multiplications and N 2 ( D − 1) additions.
Revised: (1), we need N 2 D multiplications and N 2 ( D − 1) additions.

Original: Similarly, the product AV requires N 2 D multiplications and N ( N − 1) D additions.
Revised: Similarly, the product AV requires N 2 D multiplications and N ( N − 1) D additions.

Original: In total, computing the output sequence H in Eqn.
Revised: In total, computing the output sequence H in Eqn.

Original: (1) at each head requires 3 ND (2 D x − 1)+ N 2 D + N 2 ( D − 1)+ N 2 D + N ( N − 1) D
Revised: (1) at each head requires 3 ND (2 D x − 1)+ N 2 D + N 2 ( D − 1)+ N 2 D + N ( N − 1) D

Original: = N 2 (4 D − 1)+ ND (
Revised: = N 2 (4 D − 1)+ ND (

Original: [Equation]
Revised: [Equation]

Original: D x − 4) computations.
Revised: D x − 4) computations.

Original: The total computation for all H heads is then
Revised: The total computation for all H heads is then

Original: [Equation]
Revised: [Equation]

Original: where the extra NHD (2 HD − 1) is from the linear projection by W O .
Revised: where the extra NHD (2 HD − 1) is from the linear projection by W O .

Original: (ii) Mixture of 

Revised: (ii) Mixture of

Original: Gaussian keys attention with H/ 2 -head:
Revised: Gaussian keys attention with H/ 2 -head:

Original: The number of computations in a mixture of 2 Gaussian keys attention with H/ 2 -head is N 2 H (3 D − 0 .
Revised: The number of computations in a mixture of 2 Gaussian keys attention with H/ 2 -head is N 2 H (3 D − 0 .

Original: + NHD (4 D x + HD − 4) .
Revised: + NHD (4 D x + HD − 4) .

Original: Explanation: Similar to the above derivation, in a Mixture of M Gaussian keys attention, to compute the output sequence H we need N 2 ((2 M + 2)
Revised: Explanation: Similar to the above derivation, in a Mixture of M Gaussian keys attention, to compute the output sequence H we need N 2 ((2 M + 2)

Original: D − 1) + ND (( M + 2)(2 D x − 1) − 1) computations.
Revised: D − 1) + ND (( M + 2)(2 D x − 1) − 1) computations.

Original: Note that, to compute the Gaussian distances between the queries q i and the keys k j as in the mixture of M Gaussian keys attention and to compute their dot product as in the softmax attention, we need the similar number of computations.
Revised: Note that, to compute the Gaussian distances between the queries q i and the keys k j as in the mixture of M Gaussian keys attention and to compute their dot product as in the softmax attention, we need the similar number of computations.

Original: Therefore, the total computation for all H/M heads is then
Revised: Therefore, the total computation for all H/M heads is then

Original: [Equation]
Revised: [Equation]

Original: Let M = 2 , then the total computation of the mixture of 2 Gaussian keys attention is given by
Revised: Let M = 2 , then the total computation of the mixture of 2 Gaussian keys attention is given by

Original: N 2 H (3 D − 0 . 
Revised: N [Equation] D − 0 . 5)

Original: + NHD (4 D x + HD − 4) .
Revised: + NHD (4 D x + HD − 4) .

Original: Soft-max H-head attention versus mixture of 2 Gaussian keys attention with H/2-head:
Revised: Soft-max H-head attention versus mixture of 2 Gaussian keys attention with H/2-head:

Original: Given the results in (i) and (ii), when compared to the baseline softmax H -head attention, our mixture of
Revised: Given the results in (i) and (ii), when compared to the baseline softmax H -head attention, our mixture of

Original: Gaussian keys attention with H/ 2 -head saves
Revised: Gaussian keys attention with H/ 2 -head saves

Original: [Equation]
Revised: [Equation]

Original: computations.
Revised: computations.

Original: When N is large, this difference is signiﬁcant.
Revised: When N is large, this difference is signiﬁcant.

Original: In conclusion, the mixture of 2 Gaussian keys attention with H/2-heads has cheaper computational complexity than that of soft-max H-head attention.
Revised: In conclusion, the mixture of 2 Gaussian keys attention with H/2-heads has cheaper computational complexity than that of soft-max H-head attention.

Original: B. T HE N UMBER OF P ARAMETERS
Revised: B. T HE N UMBER OF P ARAMETERS

Original: (iii) Softmax H -head attention
Revised: (iii) Softmax H -head attention:

Revised: The number of parameters

Original: : The
Revised: in

Original: number of parameters in a softmax H -head attention is 3 HDD x + ( HD ) 2 .
Revised: a softmax H - head attention is 3 HDD x + ( HD ) 2 .

Original: Explanation: 3 HDD x is from the linear projects to calculate the query matrix Q , the key matrix K , and the value matrix V in Step 1 in Section 1.1.
Revised: Explanation: 3 HDD x is from the linear projects to calculate the query matrix Q , the key matrix K , and the value matrix V in Step 1 in Section 1.1.

Original: ( HD ) 2 is from the linear project to compute the ﬁnal output as in Eqn.
Revised: ( HD ) 2 is from the linear project to compute the ﬁnal output as in Eqn.

Original: (iv) Mixture of 2 Gaussian Keys attention with H/ 2 -head:
Revised: iv) Mixture of 2 Gaussian Keys attention with H/ 2 -head

Original: The number of parameters in a Mixture of 2 Gaussian Keys attention with H/ 2 -head is 2 HDD x + 0 .
Revised: The number of parameters in a Mixture of 2 Gaussian Keys attention with H/ 2 -head is 2 HDD x + 0 .

Original: 5( HD ) 

Revised: 5( HD )

Original: + H
Revised: + H

Original: Explanation: The linear projections to calculate Q and V contribute HDD x / 2 parameters each.
Revised: Explanation: The linear projections to calculate Q and V contribute HDD x / 2 parameters each.

Original: The linear projection to calculate K contributes HDD x .
Revised: The linear projection to calculate K contributes HDD x .

Original: The linear project to compute the ﬁnal output has dimension ( H/ 2) D × HD , so it contributes 0 .
Revised: The linear project to compute the ﬁnal output has dimension ( H/ 2) D × HD , so it contributes 0 .

Original: 5( HD ) 2 parameters.
Revised: 5( HD ) 2 parameters.

Original: H more parameters is from the prior π jr .
Revised: H more parameters is from the prior π jr .

Original: These priors contribute 2 parameters at each head since we make { π j 1 , π j 2 } for all j = 1 , . . .
Revised: These priors contribute 2 parameters at each head since we make { π j 1 , π j 2 } for all j = 1 , . . .

Original: , N the same.
Revised: , N the same.

Original: Soft-max H-head attention versus mixture of 2 Gaussian keys attention with H/2-head:
Revised: Soft-max H-head attention versus mixture of 2 Gaussian keys attention with H/2-head:

Original: Given the results in (iii) and (iv), when compared to the baseline softmax H -head attention, our mixture of 2 Gaussian keys attention with H/ 2 -head saves HDD x + 0 .
Revised: Given the results in (iii) and (iv), when compared to the baseline softmax H -head attention, our mixture of 2 Gaussian keys attention with H/ 2 -head saves HDD x + 0 .

Original: 5( HD ) 

Revised: 5( HD )

Original: − H parameters.
Revised: − H parameters.

Original: When H and
Revised: When H and

Original: D are large, this saving is signiﬁcant.
Revised: D are large, this saving is signiﬁcant.

Original: C P ROOFS OF MAIN RESULTS
Revised: C P ROOFS OF MAIN RESULTS

Original: In this appendix, we provide proofs for the main results in the paper.
Revised: In this appendix, we provide proofs for the main results in the paper.

Original: C. P ROOF OF T HEOREM
Revised: C. P ROOF OF T HEOREM

Original: To ease the presentation of the proof, for any probability distribution G , we denote
Revised: To ease the presentation of the proof, for any probability distribution G , we denote

Original: [Equation]
Revised: [Equation]

Original:  for all x ∈ R d where f ( x ) = 
Revised: , for all x ∈ R d where f ( x ) = 1 (

Original: [Equation]
Revised: [Equation]

Original:  − [Equation] for given σ > 0 .
Revised: 2 πσ ) d exp (cid:16) − (cid:107) x (cid:107) 2 2 σ 2 (cid:17) for given σ > 0 .

Original: It means that p G is the convolution of f and the probability distribution G .
Revised: It means that p G is the convolution of f and the probability distribution G .

Original: Since the space of Gaussian mixtures is dense in the space of continuous probability measures (Bacharoglou, 2010), it indicates that there exists probability distribution G 1 such that
Revised: Since the space of Gaussian mixtures is dense in the space of continuous probability measures (Bacharoglou, 2010), it indicates that there exists probability distribution G 1 such that

Original: [Equation]
Revised: [Equation]

Original: Our next step is to prove that there exists a probability measure G 2 with at most K supports where
Revised: Our next step is to prove that there exists a probability measure G 2 with at most K supports where

Original: K ≤ ( C log [Equation] [Equation]
Revised: K ≤ ( C log(1 /(cid:15) ))

Original: d for some universal constant C such that
Revised: d for some universal constant C such that

Original: [Equation]
Revised: [Equation]

Original:  Indeed, from Lemma A.1 in (Ghosal & van der Vaart, 2001), for any k ≥ 1 there exists a probability distribution G 2 with at most (2 k − 2)
Revised: [Equation] Indeed, from Lemma A.1 in (Ghosal & van der Vaart, 2001), for any k ≥ 1 there exists a probability distribution G 2 with at most (2 k − 2)

Original: d supports such that
Revised: d supports such that

Original: [Equation]
Revised: [Equation]

Original: (16) for any α = ( α 1 , α 2 , . . 
Revised: [Equation] for any α = ( α 1 , α 2 , . . .

Original: . , α d ) ∈ N d such that 0 ≤ | α | = (cid:80) dj
Revised: , α d ) ∈ N d such that 0 ≤ | α | = (cid:80) dj

Original: =1 α j ≤ 2 k − 2 , Here, θ α = (cid:81) dj
Revised: =1 α j ≤ 2 k − 2 , Here, θ α = (cid:81) dj

Original: =1 θ α j j .
Revised: =1 θ α j j .

Original: Now, for any M ≥ 2 a √ d , we have (cid:107) x − θ (cid:107) ≥ (cid:107) x (cid:107) − (cid:107) θ (cid:107) >
Revised: Now, for any M ≥ 2 a √ d , we have (cid:107) x − θ (cid:107) ≥ (cid:107) x (cid:107) − (cid:107) θ (cid:107) >

Original: M − a √ d > M/ 2 as long as (cid:107) x
Revised: M − a √ d > M/ 2 as long as (cid:107) x

Original: (cid:107) >
Revised: (cid:107) >

Original: M and θ ∈ [ − a, a ] d .
Revised: M and θ ∈ [ − a, a ] d .

Original: It indicates that
Revised: It indicates that

Original: [Equation]

Original: [Equation]
Revised: [Equation]

Original: On the other hand [Equation] for any k ≥ 1 we also have that
Revised: On the other hand, for any k ≥ 1 we also have that

Original: [Equation]
Revised: [Equation]

Original: where the ﬁnal equality is stems from
Revised: where the ﬁnal equality is stems from

Original: [Equation]
Revised: [Equation]

Original: which is due to Eqn.
Revised: which is due to Eqn.

Revised: [Equation]

Original: To further bound the right-hand-side (RHS) of Eqn.
Revised: To further bound the right-hand-side (RHS) of Eqn.

Original: (18), we use the following inequality:
Revised: [Equation] we use the following inequality:

Original: [Equation]
Revised: [Equation]

Original: for any y ∈ R .
Revised: for any y ∈ R .

Original: Since k ! ≥ ( k/e ) k for any k ≥ 1 , the above bound can be rewritten as
Revised: Since k ! ≥ ( k/e ) k for any k ≥ [Equation] the above bound can be rewritten as

Original: [Equation]
Revised: [Equation]

Original: Further simpliﬁcation of Eqn.
Revised: Further simpliﬁcation of Eqn.

Original: (18) leads to
Revised: (18) leads to

Original: [Equation]
Revised: [Equation]

Original: where the ﬁnal inequality is based on an application of inequality (19) with y= −(cid:107) x − θ (cid:107) 2 / (2 σ 2 ) .
Revised: where the ﬁnal inequality is based on an application of inequality (19) with y

Original:  For (cid:107) x (cid:107) ≤ M and θ ∈ [ − a, a ] d , we have (cid:107) x − θ (cid:107) ≤ (cid:107) x (cid:107) + (cid:107) θ (cid:107) ≤ M + a √ d .
Revised: [Equation] For (cid:107) x (cid:107) ≤ M and θ ∈ [ − a, a ] d , we have (cid:107) x − θ (cid:107) ≤ (cid:107) x (cid:107) + (cid:107) θ (cid:107) ≤ M + a √ d .

Original: Therefore, we further have
Revised: Therefore, we further have

Original: [Equation]
Revised: [Equation]

Original: When M ≥ 2 a √ d , we have M + a √ d ≤ 3 M 2 and the above bound leads to
Revised: When M ≥ 2 a √ d , we have M + a √ d ≤ 3 M 2 and the above bound leads to

Original: [Equation]
Revised: [Equation]

Original: By choosing M 2 = 8 σ 2 log(1 /(cid:15) (cid:48) ) for some (cid:15) (cid:48) > 0 , the bounds in Eqns.
Revised: By choosing M 2 = 8 σ 2 log(1 /(cid:15) (cid:48) ) for some (cid:15) (cid:48) > 0 , the bounds in Eqns.

Original: (17) and (20) become
Revised: (17) and (20) become

Original: [Equation]
Revised: [Equation]

Original: As long as we choose k = 9 e 2 log(1 /(cid:15) (cid:48) ) and (cid:15) (cid:48) ≤ 1 , we have
Revised: As long as we choose k = 9 e 2 log(1 /(cid:15) (cid:48) ) and (cid:15) (cid:48) ≤ 1 , we have

Original: [Equation]
Revised: [Equation]

Original: By choosing (cid:15) (cid:48)
Revised: By choosing (cid:15) (cid:48)

Original: =
Revised: =

Original: (cid:15) 2max { 2 ( √ 2 πσ ) d , 1 } , the results from Eqns.
Revised: (cid:15) 2max { [Equation] ) d , 1 } , the results from Eqns.

Original: (21) and (22) indicate that
Revised: (21) and (22) indicate that

Original: [Equation]
Revised: [Equation]

Original: , and sup
Revised: , and sup

Original: [Equation]
Revised: [Equation]

Original: Therefore, if we choose M = 8 σ 2 log (cid:18) 2max {
Revised: Therefore, if we choose M = 8 σ 2 log (cid:18) 2max {

Original: [Equation]
Revised: [Equation]

Original: d , 1 } (cid:15) (cid:19) and k = 9 e 2 log (cid:18) 2max {
Revised: (cid:19) and k = 9 e 2 log (cid:18) 2max {

Original: [Equation]
Revised: [Equation]

Original: (cid:19) , we have
Revised: (cid:19) , we have

Original: [Equation]
Revised: [Equation]

Original: It indicates that we obtain the conclusion of claim (15) by choosing K =
Revised: It indicates that we obtain the conclusion of claim (15) by choosing K =

Original: (2 k − 2) d ≤ (cid:18)18 e 2 log (cid:18)
Revised: (2 k − 2) d ≤ (cid:18)

Original: [Equation]
Revised: [Equation]

Revised: log

Original: [Equation]
Revised: [Equation]

Original: d .
Revised: d .

Original: Combining the results from Eqns.
Revised: Combining the results from Eqns.

Original: (14) and (15), we have
Revised: (14) and (15), we have

Original: [Equation]
Revised: [Equation]

Original: As a consequence, we obtain the conclusion of the theorem.
Revised: As a consequence, we obtain the conclusion of the theorem.