Improving Transformer with an Admixture of Attention Heads
Anonymous Author(s)
AfﬁliationAddress email
Abstract
Transformers with multi-head self-attention have achieved remarkable success insequence modeling and beyond.
However, they suffer from high computational andmemory complexities for computing the attention matrix at each head.
Recently, ithas been shown that those attention matrices lie on a low-dimensional manifoldand, thus, are redundant.
We propose the Transformer with a Finite Admixture of
Shared Heads (FiSHformers), a novel class of efﬁcient and ﬂexible transformersthat allow the sharing of attention matrices between attention heads.
At the coreof FiSHformer is a novel ﬁnite admixture model of shared heads (FiSH) thatsamples attention matrices from a set of global attention matrices.
The numberof global attention matrices is much smaller than the number of local attentionmatrices generated.
FiSHformers directly learn these global attention matricesrather than the local ones as in other transformers, thus signiﬁcantly improvingthe computational and memory efﬁciency of the model.
We empirically verifythe advantages of the FiSHformer over the baseline transformers in a wide rangeof practical applications including language modeling, machine translation, andimage classiﬁcation.
On the WikiText-103, IWSLT’14 De-En and WMT’14 EnDe, FiSHformers use much fewer ﬂoating-point operations per second (FLOPs),memory, and parameters compared to the baseline transformers.
Introduction
Transformers have become the state-of-the-art model for solving many challenging problems innatural language processing [70, 1, 17, 75, 21, 10, 32, 56] and computer vision [19, 63, 22, 68].
Transformers learn from unlabeled data effectively and take advantage of the pre-trained models ondownstream tasks that involve different data modalities with limited supervision [54, 55, 21, 77, 40].
The success of transformer is credited to the multi-head self-attention (MHA) mechanism as theirfundamental building block.
For each token in the sequence, self-attention in transformers aggregatesinformation from other tokens by computing a weighted average of their feature representations witha weight proportional to a similarity score between the representations.
This attention mechanismallows arbitrary input-dependent interaction between tokens in the sequence where a token canpay attention to other tokens and attain a contextual representation [6, 70, 35].
Multi-head selfattention captures multiple such contextual representations, one at each head, thereby increasing therepresentation capacity of the self-attention.
It has been argued that the representation capacity of theattention mechanism [67] and its ﬂexibility in capturing diverse syntactic and semantic relationships[67, 71, 15, 72, 30] account for the impressive performance of transformers in practice.
1.1 Background: Self-Attention
For a given input sequence X :=
[ x 1 , · · · , x N ] > 2 R N ⇥ D x of N feature vectors, self-attentiontransforms X into the output sequence H in the following two steps:
[Figure]
a t i o
[Figure]
B
[Figure]
Testing
P a r a m e t e r s
R a t i o
T e s t M e m o r y R a t i o
Step 1.
The input sequence X is projected into the query matrix Q , the key matrix K , and the value matrix V via three linear transformations
[Equation]
where W Q , W K 2 R D ⇥ D x , and W V 2 R D v ⇥ D x are the weight matrices.
We denote Q :=[ q 1 , · · · , q N ] > , K :=
[ k 1 , · · · , k N ] > , and V :=
[ v 1 , · · · , v N ] > , where the vectors q i , k i , v i for 38 i = 1 , · · · , N are the query, key, and value vectors, respectively.
Step 2.
The output sequence H :=
[ h 1 , · · · , h N ] > is then computed as follows
[Equation]
where the softmax function is applied to each row of the matrix
A
=
( QK > ) .
This matrix A {+
+} R N ⇥ N and its component a ij for i, j = 1 , · · · , N are called the attention matrix and attention scores,respectively.
For each query vector q i for i = 1 , · · · , N , an equivalent form of Eqn.
(1) to computethe output vector h i is given by
[Equation]
The self-attention computed by Eqn.
(1) and (2) is called the scaled dot-product or softmax attention.
In our paper, we call a transformer that uses this attention the softmax transformer.
The structurethat the attention matrix A learns from training determines the ability of the self-attention to capturecontextual representation for each token.
Multi-head Attention (MHA)
Each output sequence H forms an attention head.
In MHA, multipleheads are concatenated to compute the ﬁnal output.
Let H be the number of heads and W O {+
+} R HD ⇥ HD be the projection matrix for the output.
The multi-head attention is deﬁned as
MultiHead
[Equation]
1.2 Eigenvalue Analysis of the Attention Matrices
The multi-head mechanism allows transformersto capture more diverse attention patterns andincrease the capacity of the model.
However,in many practical tasks, transformers learn redundant heads [44, 73], whose learned attentionmatrices lie on a low-dimensional manifold [8].
To conﬁrm this claim, in Figure 2, we followthe eigenvalue analysis in [8] and investigate theeigenvalues of the covariance matrix of vectorized attention matrices aggregated over each layer of a transformer model trained on the WikiText-103 dataset for language modeling.
We observe that this covariance matrix is low rank with top 200(1.2%) eigenvalues capturing more than 90% of the energy.
This result veriﬁes that the variability oflearned attention matrices in transformers can be explained by a relatively small number of principalcomponents, and those attention matrices lie on a low-dimensional manifold.
Therefore, in multi-headattention, the effective number of heads is much smaller than the actual number of heads, and a moreeffective way to compute multi-head attention is needed to improve the efﬁciency of transformers.
1.3 Contribution
Leveraging the idea of the ﬁnite admixture model (FAM) [52, 9], we propose a new class ofefﬁcient transformer architectures, namely the Transformer with a Finite Admixture of Shared
Heads (FiSHformer).
At the core of FiSHformer is to sample local attention matrices from anadmixture of a small number of global attention matrices.
This sharing mechanism between headshelps reduce the computational complexity and the model size compared to the MHA softmaxtransformer.
Our contribution is three-fold:
We construct an admixture model for shared attention matrices between heads and proposeFiSHformer, a novel class of transformers that take advantage of this admixture model to efﬁciently compute multi-head attention.
We introduce a nonlinearity mapping from global heads to local heads into FiSH and
propose the Generalized FiSHformer (GFiSHformer).
We then explore different possibilities to design FiSHformer and GFiSHformer.3.
We empirically verify that FiSHformer and GFiSHformer achieve similar or even better accuracy but with much less computational cost in terms of FLOPs and smaller model complexity measured by the number of parameters.
The advantages of our methods grow
with the model/feature dimension D and the input sequence length N .
We also show that FiSHformer-based models help reduce head redundancy in our experiments.
Organization: We structure this paper as follows:
In Section 2, we develop a ﬁnite admixture modelof shared heads and then present our FiSHformer and its extensions.
In Section 2.5, we analyze thereduction in model complexity and computational cost from FiSH.
In Section 3 and 4, we validateand empirically analyze the efﬁciency and accuracy of FiSHformer, as well as conducting ablationstudies on the model.
We discuss related works in Section 5.
The paper ends up with concludingremarks.
More results and details are provided in the Appendix.
Transformer with a Finite Admixture of Shared Heads
We ﬁrst review the ﬁnite admixture model (FAM) and derive a FAM of shared heads for the multi-headself-attention.
We then deﬁne Transformer with a Finite Admixture of Shared Heads (FiSHformer).
2.1 A Probabilistic Viewpoint of Attention Matrices
Let A j denote the attention matrix at the j th head, j = 1 , 2 , . . .
, H .
From a probabilistic viewpoint,to have diversity among A 1 , . .
A H , we can assume that A j comes from a distribution P j forall j .
Since the distributions P 1 , . . . , P H can have complex forms and be difﬁcult to compute, ourapproach is to consider approximated distributions of P j and these approximated distributions havesimple forms.
A natural choice for each of these approximated distributions is via a ﬁnite mixture of
Gaussian distributions, which can be summarized in the following lemma below.
Lemma 1.
Assume that P 2 R D 0 is a probability distribution supported on some compact set andadmits differentiable and bounded density estimation p .
Then, for any scale parameter  > 0 andfor any ✏ > 0 , there exists universal constant C and M  ( C log(1 / ✏ ))
D 0 such that we can ﬁnda mixture of M components
P
[Equation] p i N ( ✓ i ,  2 I D 0 ) where p 1 , . .
. , p K are weight parameters and 106 ✓ 1 , . .
. , ✓ K are location parameters that satisfy the following inequality
[Equation]
where  ( .
| ✓ ,  2 I ) is Gaussian density function with location parameter ✓ and covariance matrix 2 I D 0 .
The proof of Lemma 1 is in Appendix E.
In light of Lemma 1, for each scale parameter > 0 and for each distribution P j , we can ﬁnd the corresponding number of components M j , weight parameters p 1 j , . . .
, p M j j , and location parameters ✓ 1 j , . .
✓ M j j such that the mixtures
P 0 j = P M j i =1 p ij N ( ✓ ij ,  2 I D 0 ) can approximate the distribution P j up to a given accuracy ✏ .
However, these approximations still involve
Q
Hj =
1 M j number of location parameters, which can becomputationally expensive.
To overcome this issue, we assume that M 1 =
M 2 = . . .
=
M H andthe location parameters ( ✓ 1 j , . .
. , ✓ M j j ) =
( ✓ 1 , . .
. , ✓ M ) for all j , i.e., these approximated mixturesshare a similar set of location parameters.
This sharing information of location parameters has a deepconnection to ﬁnite admixture models, which we are going to elaborate in the next sections.
2.1.1 Background
Finite admixture models (FAM) are extensions of ﬁnite mixture models (FMMs), which served as aworkhorse in stochastic modeling.
A ﬁnite mixture distribution of M components for a random array
X 2 R N ⇥
J is given by
[Equation]
where x j 2 R N is the j -th row of X randomly sampled from the mixture distribution, f is a chosenprobability measure, such as a Gaussian distribution, p = { p 1 , p 2 , . .
. , p M } are mixture weights, and✓ k denotes the parameter values for the k -th component.
A FAM is a generalization of a FMM where rows x j , j = 1 , . . .
, H , are drawn from different mixturedistributions that share the components f ( x ; ✓ k ) , k = 1 , . . .
M but with different mixture weights
[Equation]
2.2 Multi-head as a Finite Admixture Model of Shared Heads (FiSH)
As demonstrated in Section 2.1, we propose a Finite Admixture Model of Shared Heads (FiSH), inwhich A j follows ﬁnite admixture distribution of M components given by
[Equation]
Here M < H and f ( A ; ✓ k ) are chosen probability measures.
In particular, we choose f
( A ; ✓ k ) to be
Gaussian distributions N ( A ; G k , ⌃ k ) , where G k = Q k K k > and ⌃ k =  2 k
I are the cluster meansand coveriances, respectively.
FiSH is then deﬁned as follows:
Deﬁnition 1 (Finite Admixture Model of Shared Heads) .
The multi-head attention admits a ﬁniteadmixture model of shared heads if the attention matrices A j at the j th head are sampled from thefollowing ﬁnite admixture model:
[Equation]
In FiSH deﬁned in Def. 1, we call { G k = Q k K k > } k =1 ,...,M global attention matrices and{ A j } j =1 ,...,H local attention matrices.
FiSH computes M global attention matrices G k , and H local 138 attention matrices
A j are sampled from FiSH as in Eqn.
7 with M < H .
Remark 1 (FiSH vs. Baseline MHA) .
The baseline MHA with H heads need to compute H , e.g.
H = 8 , attention matrices, each of which requires O ( N 2 ) computational costs where N is the lengthof the input sequence.
In contrast, FiSH only need to compute M < H , e.g. M = 2 , global attentionmatrices, each of which also requires O ( N 2 ) computational and memory costs.
Then FiSH combinesthose global attention matrices to form a FAM from which H , e.g. H = 8 , local attention matricesare sampled as in Eqn. 7.
This second step of sampling local attention matrices from a set of globalattention matrices in FiSH requires very few computations.
Thus, FiSH is more efﬁcient than MHA.
Remark 2 (Connection to Topic Models) .
FiSH can be interpreted as a Probabilistic Latent Semantic Analysis (pLSA) model for topic modeling.
Considering the document d that contains theword w whose topic is c , pLSA models the occurrence of the word w in the document d as a mixtureof conditionally independent Multinomial distributions p ( w | d ) =
P c p ( c | d ) p ( w | c ) .
Comparing 150 this pLSA with Eqn.
7 of FiSH, we can associate the mixture weights p kj and the distribution
N ( A ; Q k K k > ,  2 k I ) in FiSH with the distributions p ( c | d ) and p ( w | c ) in pLSA, respectively.
Therefore, it can be interpreted that the global attention matrices in FiSH play the role of topics, and thelocal attention matrices in FiSH are words sampled from those topics.
It is interesting to note thatpLSA is equivalent to the famous Latent Dirichlet Allocation model under a uniform Dirichlet prioron the per-document topic distribution p ( c | d ) .
2.3 Transformer with a Finite Admixture of Shared Heads
FiSHformers are transformers that use FiSH instead of MHA.
FiSH, as deﬁned in Def. 1, is notdifferentiable, which poses a difﬁculty in training FiSHformers.
Applying the reparameterizationtrick [36], the attention matrices A j can be written in a differentiable form as follows:
[Equation]
FiSHformers use the formulation of local attention heads in Eqn. 8 to implement FiSH.
Transformer with a Hard Finite Admixture of Shared Heads (Hard FiSHformer)
Hard FiSHformer takes the zero-noise limit of Eqn.
8 to reduce the computational cost.
The attention matrices
A j in Hard FiSHformer are then calculated as
[Equation]
Remark 3 (Discriminative Relaxation) .
To take the advantage of learning from data, the convexcombination condition of p kj , i.e. P [Equation] =1 p kj = 1 , p kj  0 , can be relaxed, and those mixingcoefﬁcients are made learnable parameters that are learned from data during training
Remark 4 (Transformers with a Mixture of Shared Heads) .
A transformer with a mixture of sharedheads (MiSHformer) can be used to reduce the amount of computation with the cost of accuracyreduction.
MiSH is a special case of FiSH when the mixture weights p kj are the same for all j .
Thelocal attention matrices
A j in MiSHformer are given by
[Equation]
An empirical comparison between FiSHformer and MiSHformer is provided in Section 4.
2.4 Transformer with a Generalized Finite Admixture of Shared Heads
In order to increase the representation capacity of attention heads, we follow a common approach inlearning representation by replacing the linear mapping in Eqn. 8 and 9 by a nonlinear mapping suchas a neural network with the rectiﬁed linear units (ReLU).
The Transformer with a Generalized Finite
Admixture of Shared Heads (GFiSHformer) is then formulated as
[Equation]
where  is a nonlinear mapping and p kj are relaxed to be learnable parameters.
Similarly, weformulate local attention matrices
A j in the Transformer with a Generalized Hard Finite Admixtureof Shared Heads (Hard GFiSHformer) as A j = P [Equation] [Equation]
Reduction in Model Complexity and Computational Cost from FiSH
Compared to its H -head MHA counterpart, a FiSH attention of M global heads and H local headssaves [2( H  M ) D  2 MH )] N 2 +
( H  M ) D (2 D x  1) N FLOPs in a forward pass and2( H  M ) DD x  HM  M parameters.
Detailed derivations are provided in Appendix D.
Experimental Results
In this section, we empirically study the advantages of FiSHformer on various tasks and benchmarks,including language modeling on WikiText-103 (Section 3.1), machine translation on IWSLT’
De-En and WMT’14 (Section 3.2), and image classiﬁcation on ImageNet (Section 3.3).
We aim toshow that: (i) FiSHformers improve the efﬁciency and accuracy upon the MHA baseline; (ii) FiSHis a universal method that can be applied on state-of-the-art transformer models to improve their
[Figure]
B
C
D
[Figure]
P a r a m e t e r s
R a t i o
T e s t M e m o r y
R a t i o
[Table]
performance on large-scale applications.
In Section 4, we also show that FiSH helps reduce theredundancy between attention heads
[Table]
We compare FiSHformers, Hard FiSHformers, GFiSHformers
[Table]
and Hard GFiSHformers with thebaseline MHA softmax transformers.
In our experiments, we apply the discriminative relaxationexplained in Remark 3 on our FiSHformers to make the mixture weights p kj learnable parameters.
For GFiSHformers/Hard GFiSHformers, we choose the nonlinear mapping  to be a ReLU followedby a linear neural network.
All of our results are averaged over 5 runs with different seeds.
Moredetails on datasets, models, and training are provided in Appendix A.
3.1 WikiText-103 Language Modeling
Models and baselines We compare the 2 and 4-global-head FiSHformers with the 8-head softmaxtransformers [70].
Each model has 16 layers, and our training follows the setting from [61].
Results Perplexity: Table 3.1 demonstrates that our 2/4-global-head (G)FiSHformers and theirhard versions obtain comparable or better PPLs than the corresponding 8 head MHA baseline on
WikiText-103.
Interestingly, the 2-global-head (G)FiSHformers perform on par with the 8-headbaseline even though only 2 global attention matrices are used to span all local attention matrices,indicating that the attention matrices in MHA are indeed redundant and the representation capacity oflocal attention matrices in (G)FiSH, though being generated from only 2 global bases, is comparableto those in the 8-head MHA.
Efﬁciency: In Fig.
3A and 3B, we presents the reduction ratio of train and test FLOPS, respectively,of our 2-global-head GFiSHformer vs. the baseline 8-head MHA transformer as functions of modeldimension D and sequence length N .
In Fig.
3C and 3D, we show the reduction ratio of modelsize and GPU memory usage at test time, respectively, of our 2-global-head GFiSHformer vs. thesame baseline.
We observe that the efﬁciency advantage of GFiSHformer over the baseline growswith D and N , making it more suitable and superior for large-scale applications.
Note that themodel size in terms of the number of parameters does not depend on the sequence length N , andfrom our experiments, we observe that the GPU memory usage reduction ratio is almost the samefor different sequence lengths.
More efﬁciency analysis results on this language modeling task are
[Table] [Table]
[Table]
provided in Section 4 and Appendix C.
Also, Figure 5 in Appendix B.2 shows the train and test
PPL of (G)FiSHformers and the MHA softmax transformers.
In our experiment, 4-global-head
GFiSHformers obtain the best validation PPL.
3.2 Machine Translation
In this section, we examine the performance of (G)FiSHformer on the neural machine translationtask, an important task in natural language processing in which the sequence lengths of the input arenot the same.
We ﬁrst compare (G)FiSHformers with the baselines MHA softmax transformers onthe IWSLT’ 14 De-En [11] and then scale up our experiments to the WMT’14 En-De [42].
On thesetasks, we calculate the BLEU scores for evaluation.
Models and baselines For the IWSLT’ 14 De-En task, we compare 2-global-heads (G)FiSHformerswith the baseline 4-head softmax transformer.
Each model consists of 12 layers, 6 layers for anencoder and the other 6 layers for a decoder.
Our experiments follow the setting on fairseq.
For the
WMT’14 En-De task, we use similar models as in the IWSLT’ 14 De-En task.
However, we compare(G)FiSHformers of 8 and 4 global heads with the 16-head MHA softmax baseline.
Our training andmodel setting are the same as those in [48].
Results As shown in Table 2 and 4, (G)FiSHformers outperform or at least are on par with the baseline
MHA softmax transformers.
Again, these results indicate rich representations of the local attentionmatrices generated by (G)FiSH.
Furthermore, (G)FiSHformer outperforming Hard (G)FiSHformerin all settings suggests the positive value of adding noise into the models to turn them into a properprobabilistic model.
Nevertheless, it is worth noticing that Hard (G)FiSHformer is more efﬁcientthan (G)FiSHformer.
Fig.
4 in Appendix B.1 summarizes the advantage in efﬁciency of 2-globalhead GFiSHformer over the 4-head baseline on the IWSLT’ 14 De-En task.
These advantages of
GFiSHformer grow with the model dimension D .
3.3 Image Classiﬁcation on ImageNet
The advantages of (G)FiSHformers also hold across different data modalities.
To illustrate this point,in this section, we apply (G)FiSH to Swin transformer [41], a state-of-the-art vision transformerarchitecture, for the image classiﬁcation task on the ImageNet dataset [20].
The baseline Swin-Twe use has a total of 12 layers, across 4 stages of transformer blocks with 3, 6, 12, and 24 headseach.
Our GFiSH Swin-T uses 6 and 12 global attention heads at the last two stages.
The model andtraining follow the settings in [41].
We summarize our results in Table 5.
Our GFiSHformer is onlyslightly more efﬁcient than the baseline in this case because the sequence length N per window forthis task is small, i.e. N = 49 .
However, like with the previous language tasks and as pointed byformula of computational cost and model complexity reduction in Section 2.5, these advantages growwith larger D and N .
3.4 FiSHformer is more effective than other methods for head-redundancy reduction
To futher demonstrate the effectiveness of our method, we compare FiSHformers against the headredundancy reduction method in [16] on the WMT’14 machine translation task.
[16] proposes theadaptively sparse transformer (AST), reducing redundancy within each head by zeroing out lowattention scores.
In comparison, the results further conﬁrm the effectiveness of our method since the
BLEU score of FiSHformer and GFiSHformer are 27.26 and 27.67, respectively, better than that of
AST, which is 26.93.
All models share the same architecture, with 12 transformer layers, 6 encoder,and 6 decoder layers.
Our FiSHformers have 4 global heads and 8 local heads per layer while 8attention heads are used in each AST layer. [Table] [Table]
[Table] [Table] [Table] [Table]
[Table]
3.5 Beyond Multi-Head Softmax Transformers
We show that (G)FiSH can be applied on top of many transformer architectures to improve theirperformance including the linear transformers [34] and the SoTA transformer with noisy backtranslation [24].
More results of combining (G)FiSH with efﬁcient transformers are in the Appendix.
(G)FiSH on Linear Transformers Linear transformers [34] is a class of efﬁcient transformers that linearize the softmax kernel in Eqns.
1 and 2 when computing attention matrices.
Weapply (G)FiSH on linear transformers trained for the IWSLT14 De-En machine translation task andsummarize the results in Table 3.
The empirical results verify that applying (G)FiSH using only2-global heads on a 4-head linear transformer improves the accuracy of the baseline model.(G)FiSH Improves the State-of-the-Art Noisy Back-Translation We apply an 8-global-head
Hard GFiSH on the transformers trained with noisy back-translation [24] for the WMT’14 En-Detranslation task and obtain the BLEU score of 33.45.
This result is comparable to the SoTA result of33.52 from the transformers trained with noisy back-translation but our model is more efﬁcient.
Empirical Analysis
We study models trained for the WikiText-103 language modeling task in this section.
Efﬁciency Analysis
In this section, we further investigate the efﬁciency reduction of 2-global-head
GFiSHformers over the 8-head baseline as a function of the number of heads in Fig.
7 and 8 andcompare the efﬁciency of our FiSH-based models in Fig.
Fig.
6, 7, 8 and details on our settingare provided in Appendix C. From Fig.
7 and 8, we observe that when using fewer number ofglobal heads, GFiSHformers achieve signiﬁcantly more computation reduction (in both training andinference).
Furthermore, Fig.
6 shows that the efﬁciency measures, i.e. FLOPs, model size, and GPUmemory usage, of all FiSH-based models we study in this paper follow similar patterns.
FiSHformer Helps Reducing Head Redundancy We show that (G)FiSHformers attain moresigniﬁcant distances between heads than the baseline.
Thus, our models capture more diverse patternsacross heads than the baseline.
For a given pre-trained model, we compute the pair-wise L 2 distancesbetween heads in the same layer.
We show the layer-average mean and variance of distances betweenheads in GFiSHformers compared with those in the MHA softmax baselines in Table 6.
We provideadditional results for Hard GFiSHformers and Hard FiSHformer in Table 10 in the Appendix.
Eigen Analysis We show that heads in GFiSHformer lie on a higher-dimensional subspace comparedto those in FiSHformer.
This justiﬁes our use of nonlinearity mapping to generate local headsfrom global heads.
Using a pre-trained model, we ﬁrst compute the covariance matrix of thevectorized attention scores of the l -th layer: C l = 1 M · H
P Mm =
P
Hh =
( A l,hm )( A l,hm ) > .
We use spectraldecomposition to represent C l in terms of eigenvalues and eigenvectors, namely, C l =[-(cid:80) n 2 i =
-]
P n 2 i =1  i v i v > i .
Without losing generality, we assume that eigenvalues are sorted in descending order.
We illustratethe layer-average number of principle components that are needed to explain 95% variance in Table 7.
Interestingly, Table 7 shows that attention matrices in all of our proposed FiSH-based models lie onhigher-dimensional subspace than those in the baseline MHA softmax transformers, which indicatesthat our models achieve better representational capacity than the baseline, conﬁrming the advantageof (G)FiSH over MHA.
Table 11 in the Appendix provide additional results for Hard (G)FiSHformer.
[Table]
[Table] [Table]
[Table]
Admixture vs. Mixture of Heads We compare the transformer with a mixture of heads and the
[Table]
transformer with an admixture of heads.
We show that the transformer with a mixture of heads yieldsworse accuracy.
We summarize our results on the WikiText-103 language modeling task in Table 8.
Ablation Study on the Effect of the Number of Global Heads on FiSH-based Models Weinvestigate the accuracy, efﬁciency, and representation capacity of FiSH-based models under differentnumbers of global heads on the WikiText-103 language modeling task.
Since GFiSH obtains the bestaccuracy on this task, we use GFiSH as a study case in this section and report our results on accuracyin Table 9.
Ablation results on efﬁciency are summarized in Fig.
7, 8 in Appendix C, and ablationresults on representation capacity are provided in Table 6, 7.
Related Work
Efﬁcient Transformers Efﬁcient transformers have been developed to reduce the quadratic computational and memory cost of transformers [57].
A class of efﬁcient transformers are sparse transformerswhich design the attention matrix to have sparse structure [50, 39, 53, 12, 7].
Another class of efﬁcienttransformers are models that integrate different access patterns for better coverage [12, 31].
Theseaccess patterns can also be learned from the data [37, 57, 66].
In other works, a side memory moduleis used to access multiple tokens simultaneously [38, 64, 3, 7].
Recently, low-rank and kernelizationmethods have been proposed to improve the computational and memory efﬁciency of computingself-attention [69, 74, 34, 14, 62, 47, 51].
Our (G)FiSHformers are complementary to these methods.
Redundancy in Transformers Most of the neurons and heads in the pre-trained transformer areredundant and can be pruned when applied on downstream tasks [18, 44, 23].
The contextualizedembeddings in pre-trained networks under this redundancy have also been studied to demonstratethat the representations learned within these models are highly anisotropic [45, 25].
Knowledgedistillation and sparse approximation have also been used to enhance the efﬁciency of transformers,including [60, 65, 73, 59].
Our FiSH-based approach are complementary to these methods
Mixture Models for Transformers Recently, mixture models have been employed to study andenhance transformers.
Among these works is switch transformers [26] that uses the routing algorithmin Mixture of Experts (MoE) to decrease the communication and computational costs in transformers.[46] derives a GMM for each attention head.
Other works that combine mixture models withtransformers include [13, 29, 33].
Concluding Remarks
In this paper, we proposed the FiSHformer, a class of transformers that samples local attentionmatrices from a ﬁnite admixture of global attention matrices.
FiSHformers and their generalizedversion GFiSHformers attain better computational cost and model complexity than their baseline
MHA softmax transformer counterparts.
Furthermore, (G)FiSHformers help increase the diversitybetween attention heads.
It is worth noting that there is no potential negative societal impacts of
FiSHformers.
Also, global attention matrices in FiSHformers currenly do not have any structure, andthis is a limitation of our model.
It is interesting to impose additional structures such as low-rank andsparsity into the global attention matrices.
Checklist
The checklist follows the references.
Please read the checklist guidelines carefully for information onhow to answer these questions.
For each question, change the default [TODO] to [Yes] , [No] , or[N/A] .
You are strongly encouraged to include a justiﬁcation to your answer , either by referencingthe appropriate section of your paper or providing a brief inline description.
For example:
• Did you include the license to the code and datasets?
[Yes] See Section ?? .
• Did you include the license to the code and datasets?
[No] The code and the data are
proprietary.
• Did you include the license to the code and datasets?
[N/A]
Please do not modify the questions and only use the provided macros for your answers.
Note that the
Checklist section does not count towards the page limit.
In your paper, please delete this instructionsblock and only keep the Checklist section heading above along with the questions/answers below.
For all authors...
(a) Do the main claims made in the abstract and introduction accurately reﬂect the paper’s
contributions and scope?
[Yes]
(b) Did you describe the limitations of your work?
[Yes] See Section 6.
(c) Did you discuss any potential negative societal impacts of your work?
[Yes] See
Section 6.
(d) Have you read the ethics review guidelines and ensured that your paper conforms to
them?
[Yes]
If you are including theoretical results...
(a) Did you state the full set of assumptions of all theoretical results?
[Yes]
(b) Did you include complete proofs of all theoretical results?
[Yes]
If you ran experiments...
(a) Did you include the code, data, and instructions needed to reproduce the main experi-
mental results (either in the supplemental material or as a URL)?
[Yes]
(b) Did you specify all the training details (e.g., data splits, hyperparameters,
how they
were chosen)?
[Yes]
(c) Did you report error bars (e.g., with respect to the random seed after running experi-
ments multiple times)?
[Yes]
(d) Did you include the total amount of compute and the type of resources used (e
g., type
of GPUs, internal cluster, or cloud provider)?
[Yes]
If you are using existing assets (e
g., code, data, models) or curating/releasing new assets...
(a) If your work uses existing assets, did you cite the creators?
[N/A]
(b) Did you mention the license of the assets?
[N/A]
(c) Did you include any new assets either in the supplemental material or as a URL?
[N/A]
(d) Did you discuss whether and how consent was obtained from people whose data you’re
using/curating?
[N/A]
(e) Did you discuss whether the data you are using/curating contains personally identiﬁable
information or offensive content?
[N/A]
If you used crowdsourcing or conducted research with human subjects...
(a) Did you include the full text of instructions given to participants and screenshots, if
[Table]
applicable?
[N/A]
[Table]
(b) Did you describe any potential participant risks, with links to Institutional Review
[Table]
Board (IRB) approvals, if applicable?
[N/A]
[Table]
(c) Did you include the estimated hourly wage paid to participants and the total amount
[Table]
spent on participant compensation?
[N/A]
[-The ImageNet dataset [58] contains about 1 . 
-]
[-LFiSH-variants improve the performance of the linear baseline, as demonstrated in Table 
-]
[Table]
[-Whileonly having the fewest heads, 2-head GFiSHformer is the most efficient model (lowest metric ratiowith the 8-head MHA baseline) and achieves comparable performance with the baseline, for modeldimension, and sequence length are 128 and 256 respectively, as indicated in Table 3.
-]
[-+ N 
-]
[-Soft-max H-head attention versus FiSH with M global attention heads and H local attention heads: Given the results in (i) and (ii), when compared to the baseline softmax H -head attention, our FiSH with M global attention heads and H local attention heads saves [2( H − M ) D − 2 MH )] N 2 + 
-]