Improving Transformer with an Admixture of Attention Heads
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 efficient and flexible transformersthat allow the sharing of attention matrices between attention heads.
At the coreof FiSHformer is a novel finite 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 significantly improvingthe computational and memory efficiency 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 classification.
On the WikiText-103, IWSLT’14 De-En and WMT’14 EnDe, FiSHformers use much fewer floating-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 flexibility 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 ] ⊤ ∈ R N × D x of N feature vectors, self-attentiontransforms X into the output sequence H in the following two steps [Figure] L
[Figure]
R a t i o
[Figure]
A B C D
Model Dimension
[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 ∈ R D × D x , and W V ∈ 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 fori = 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 [-
-]∈ 41 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 final output.
Let H be the number of heads and W O [-
-]∈ 50 R HD × HD be the projection matrix for the output.
The multi-head attention is defined as
MultiHead ( { H } Hi =1 ) =
Concat ( H 1 , [Equation]
[Equation] H H ) W O .
1.2 Eigenvalue Analysis of the Attention Matrices
[Figure]
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 confirm 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 verifies 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 efficiency of transformers.
1.3 Contribution
Leveraging the idea of the finite admixture model (FAM) [52, 9], we propose a new class ofefficient 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 toefficiently 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 possibilitiesto design FiSHformer and GFiSHformer.3.
We empirically verify that FiSHformer and GFiSHformer achieve similar or even betteraccuracy but with much less computational cost in terms of FLOPs and smaller modelcomplexity 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 finite 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 efficiency 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 first review the finite admixture model (FAM) and derive a FAM of shared heads for the multi-headself-attention.
We then define 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 difficult 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 finite mixture of
Gaussian distributions, which can be summarized in the following lemma below.
Lemma 1.
Assume that P ∈ R D ′ 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 ′ such that we can finda mixture of M components (cid:80)
Mi
=1 p i N ( θ i , σ 2 I D ′ ) where p 1 , . .
. , p K are weight parameters andθ 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 ′ .
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 find 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 ′ j = (cid:80) M j i =1 p ij N ( θ ij , σ 2 I D ′ ) can approximate the distribution P j up to a given accuracy ϵ .
However, these approximations still involve (cid:81)
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 finite admixture models, which we are going to elaborate in the next sections.
2.1.1 Background
Finite admixture models (FAM) are extensions of finite mixture models (FMMs), which served as aworkhorse in stochastic modeling.
A finite mixture distribution of M components for a random array
X ∈ R N
× J is given by
[Equation]
where x j ∈ 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 finite 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 defined as follows:
Definition 1 (Finite Admixture Model of Shared Heads) .
The multi-head attention admits a finiteadmixture model of shared heads if the attention matrices A j at the j th head are sampled from thefollowing finite admixture model:
[Equation]
In FiSH defined 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 localattention 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 efficient 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 ) =
(cid:80) c p ( c | d ) p ( w | c ) .
Comparingthis 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 defined in Def. 1, is notdifferentiable, which poses a difficulty 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. (cid:80) Mk =1 p kj = 1 , p kj ≥ 0 , can be relaxed, and those mixingcoefficients 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 rectified 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 = (cid:80) Mk =1 ϕ ( p kj Q k K k ⊤ ) .
2.5 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 classification on ImageNet (Section 3.3).
We aim toshow that: (i) FiSHformers improve the efficiency 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.
Efficiency: 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 efficiency 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 efficiency analysis results on this language modeling task are
[Table] [Table] [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 first 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 efficientthan (G)FiSHformer.
Fig.
4 in Appendix B.1 summarizes the advantage in efficiency 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 Classification 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 classification 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 efficient 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 confirm 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]
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 efficient transformers are in the Appendix.
Applying (G)FiSH on Linear Transformers Linear transformers [34] is a class of efficient 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 for the WMT’14 En-De translationtask and obtain the BLEU score of 33.45.
This result is comparable to the SoTA result of 33.52 fromthe transformers trained with noisy back-translation but our model is more efficient.
Empirical Analysis
We study models trained for the WikiText-103 language modeling task in this section.
Efficiency Analysis
In this section, we further investigate the efficiency 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 efficiency 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 significantly more computation reduction (in both training andinference).
Furthermore, Fig.
6 shows that the efficiency 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 moresignificant 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 justifies our use of nonlinearity mapping to generate local headsfrom global heads.
Using a pre-trained model, we first compute the covariance matrix of thevectorized attention scores of the l -th layer: C l =
[Equation]
=1 (cid:80)
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 {+=
+}
λ 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, confirming 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 thetransformer 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, efficiency, 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 efficiency are summarized in Fig.
7, 8 in Appendix C, and ablationresults on representation capacity are provided in Table 6, 7.
Related Work
Efficient Transformers Efficient transformers have been developed to reduce the quadratic computational and memory cost of transformers [57].
A class of efficient transformers are sparse transformerswhich design the attention matrix to have sparse structure [50, 39, 53, 12, 7].
Another class of efficienttransformers 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 efficiency 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 efficiency 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 finite 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 justification 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 reflect 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 identifiable 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]
[Table]
Supplement
to Improving Transformer with an Admixture of Attention Heads "A Experiment Details
In this section, we provide model and training details for experiments in Section 3.
In our experiments,we consider the number of global attention matrices as a hyper parameter to finetune.
We observethat in all of our experiments, choosing the global attention matrices in FiSH to be 1/4 and 1/2 of thenumber attention heads in the original MHA results in models with good accuracy and efficiency.
Allof our experiments are conducted on a server with 4 NVIDIA A100 GPUs.
A.1 Language Modeling
Datasets and metrics WikiText-103 consists of articles from Wikipedia and is a dataset with longcontextual dependencies.
The training set is made up of about 28 K articles containing 103 M runningwords; this corresponds to text blocks of about 3600 words.
The validation and test sets are composedof 218 K and 246 K running words, respectively.
Each of them contains 60 articles and about 268 Kwords.
Our experiment follows the standard setting [43, 61] and splits the training data into L -wordindependent long segments.
For evaluation, we use a batch size of 1 and go through the text sequencewith a sliding window of size L .
We consider only the last position for computing perplexity (PPL)except in the first segment, where all positions are evaluated as in [2, 61].
A.2 Machine Translation
Datasets and metrics The dataset IWSLT’14 De-En contains about 170 K training sentence pairs,7 K validation pairs, and 7 K test pairs.
In this task, the model does the translation from German to
English.
The WMT dataset is a rich-resource English-German machine translation dataset, containingabout 4 .
5 M training sentence pairs.
Validation and test data are from the corresponding newest data.
The BLEU score [49] is used to measure the performance of the trained model.
A.3 Image Classification
Datasets
and
metrics
{+The ImageNet dataset [58] contains about 1 . 
+}
M training images and 50 Kvalidation images, the model learns to predict which one of 1000 classes an image belongs to.
Our Swin Transformer [41] experiments are based on the public code https://github.com/microsoft/
Swin-Transformer, we implemented our Hard GFiSH models with the Swin-T version.
We add ourglobal heads to the last 8 of the total 12 layers of the model, on each layer we set the number ofglobal heads to half the number of heads, which are 6 and 12 global heads for layers with 12 and 24heads, respectively.
Our experiments were conducted on a server with 1 NVIDIA RTX 3090.
We setthe batch size to 128 and the learning rate to 1.25e-4, all models are trained with single precision.
A.4 Applying (G)FiSH on linear transformers
Here we provide the detailed implementation of (G)LFiSHformer, i.e. (G)FiSH + linear transformer,discussed in Section 3.5 and Table 3 in the main text.
The linear transformer reduces the quadraticcomputational cost of self-attention to linear complexity, in terms of the sequence length, by linearizing the softmax kernel [34].
We combine (G)FiSHformer with linear transformer by generating theglobal K T V and then sampling the local K T V from the global ones, resulting in the Transformerwith a Linear Finite Admixture of Shared Heads (LFiSHformer).
Similar to (G)FiSHformer, wederive four different LFiSH-based transformers: LFiSHformer, Hard LFiSHformer, Generalized
LFiSHformer (GLFiSHformer), and Hard Generalized LFiSHformer (Hard GLFiSHformer).
Our
{+LFiSH-variants improve the performance of the linear baseline, as demonstrated in Table 
+}
B Additional Experimental Results
B.1 A Comparison on the Model Efficiency for the IWSLT14 De-En Machine Translation
Task
Fig.
4 summarizes the advantage in efficiency of 2-global-head GFiSHformer over the 4-head baselineon the IWSLT’ 14 De-En task.
B.2 Train and validation PPL of models trained for the WikiText-103 language modeling task
Figure 5 shows the train and valid PPL of 4-global-head FiSH-based models vs .
8-head MHA
Transformer trained for the WikiText-103 language modeling task.
Traning
Testing
Training
Testing
[Figure]
Dimension
R ea l T i m e R a t i o
M e m o r y R a t i o
P a r a m e t e r R a t i o
A
B
E D C
B.3 More Results to Show that FiSHformer Helps Reducing Head Redundancy
Table 10 presents the layer-average mean and variance of distances between heads in Hard FiSHformers and Hard GFiSHformers compared with those in the MHA softmax baselines.
Models are trainedfor the WikiText-103 language modeling task.
[Table]
B.4 More Results on Eigen Analysis
Table 11 presents the layer-average number of principle components for 95% variance explained ofthe covariance of attention matrices in Hard FiSHformers and Hard GFiSHformers trained for the
WikiText-103 language modeling task compared with those in the MHA softmax baselines.
Modelsare trained for the WikiText-103 language modeling task.
C Efficiency Analysis
We present the efficiency improvement of (G)FiSHformers over the MHA baseline.
We show thatthe advantage in the efficiency of (G)FiSHformers increases significantly as the model dimension Dand sequence length N grows, making (G)FiSHformers more suitable and superior for large-scale
[Table]
[Figure]
r y
[Figure]
applications.
In our analysis, we report the number of FLOPs, the number of model parameters, andmemory usage (bytes) as the measures of model efficiency.
Analysis setting We investigate the benefits of our model computation and memory reduction throughdifferent D ∈ {256, 512, 1024, 2048, 4096} and N ∈ {128, 256, 512, 1024, 2048, 4096, 8192}.
Forthe FLOP calculation, we use fvcore.
Another notice is that, for model complexity, we distinguishbetween input-embedding parameters and non-embedding parameters.
Input-embedding parametersare used to represent the inputs before sending them to the model.
Non-embedding parameters arethe parameters of the main model.
Since our method aims at reducing the size of the transformermodel, it is important to compare the reduction in non-embedding parameters.
Hence, we reportboth model’s total parameters and non-embedding parameters in this analysis.
All measurements arecalculated when running the model through data of batch size 1.
FiSH-based models significantly improves computational cost
GFiSHformer benefits computation in large-scale tasks.
By showing the FLOP ratios between 2global-head GFiSHformers with the 8-head MHA baseline across different model dimensions andsequence lengths, Figure 3 indicates that our model requires significantly less computation than thebaseline.
Especially for both training and inference, this substantially grows with the increases of Dand N (up to more than 30 % reduction), which makes our method more preferable for long-rangeand large-scale tasks.
FiSH-based variants share comparable advantages in computation saving.
For further analysis, wecompare the FLOP of FiSHformers, Hard FiSHformers, GFiSHformers, and Hard GFiSHformer.
Figure 6 shows that given a D , when N is increased, FiSH-based models have comparable trainingand inference computation (FLOP).
Since we have shown that GFiSHformer benefits computation,the FLOP comparison among FiSH-based models further confirms that our methods have a substantialadvantage in computational saving.
GFiSHformer computation efficiency rapidly increases as the number of global heads decreases Wecompare the computational-cost reduction of GFiSHformer as the global heads vary.
Figure 7 showsthe FLOP ratio of GFiSHformer with 2/4/6-global-head versus the 8-head baseline for a given Dand various sequence lengths.
As the number of heads decreases, GFiSHformers achieve significantcomputation reduction (in both training and inference)
(Figure 7).(G)FiSHformers improves memory usage and model complexity
In addition to the computationalsaving, our method achieves significant benefits in memory cost (at test time) and model complexity(total/non-embedding parameters) over the baseline.
Following the computation analysis, we first
[Figure]
[Figure]
P a r a m e t e r
R a t i o
T e s t M e m o r y R a t i o N o ne m b e dd i n g
P a r a m e t e r
[Figure]
a t i
[Figure]
present the advantage of the GFISHformers compared to the 8-head MHA Transformer.
Figure 8shows the number-of-parameter-ratio (Left) and the memory-ratio (Right) of the GFiSHformers with2/4/6 global heads and the 8-head baseline.
At a fixed sequence length N = 1024 , as we scale upthe model dimensions, our method becomes significantly more beneficial than MHA Transformer,indicating the advantage of GFiSHformer in large-scale applications.
Secondly, we show that FiSH-based models are comparable in model size and memory saving, whichfurther indicates that all variants of FiSHformers benefit space complexity.
Figure 6 shows that2-head FiSHformers, Hard FiSHformers, GFiSHformers, and Hard GFiSHformer have a comparativenumber of total/non-embedding parameters and memory usage at different model dimensions, for agiven sequence length.
Finally, we examine the model space complexity of a Fishformer variant as the number of globalheads vary.
Figure 8 shows a significant increase in total (Left) /non
-embedding parameters (Middle)and memory usage (Right) reduction of 2/6/8-global heads GFiSHformers, as we increase D .
{+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.
+}
D An Analysis on the Computational Complexity and the Number ofParameters in FiSH and the Softmax Attention
In this section, we compare the computational complexity and the number of parameters in the FiSHwith M global attention heads and H local attention heads to the H -head baseline MHA softmaxtransformer.
Following the same notation in Section 1.1 in the main text, we let D x , N , and D be theinput dimension, the input length, and the model/feature dimension, respectively.
To simplify thecomputation, we also do not take the softmax operator into account.
D.1 Computational Complexity
(i)
Softmax
H -head attention:
The number of computations needed to compute attention matricesin a softmax H -head attention is N 2 H (2 D − 1) + 2 NHD (2 D x − 1) .
Explanation: To calculate the query matrix Q and the key matrix K in Step 1 in Section 1.1 ateach head, we need 2 NDD x multiplications and 2 ND ( D x − 1) additions.
In total, these need2 ND (2 D x − 1) computations.
Next, to compute the product QK ⊤ in Eqn.
(1), we need N 2 Dmultiplications and N 2 ( D − 1) additions.
In total, computing an attention matrix in Eqn.
(1) ateach head requires 2 ND (2 D x − 1) + N 2 D
{++ N 
 ( D − 1) = N 2 (2 D − 1) + 2 ND (2 D x − 1)computations.+}
The total computation needed to compute attention matrices at H heads is then
N 2 H (2 D − 1) + 2 NHD (2 D x − 1) .
ii) FiSH with M global attention heads and H local attention heads: The number of computations
needed to compute attention matrices in a FiSH with M global attention heads and H local attention
heads
is [2( D
+
H )
M − H ] N 2 +
NMD
(2 D x − 1) .
Explanation: Similar to the above derivation, M global attention matrices need N 2 M (2 D − 1)
+2 NMD (2 D x − 1) computations.
The H local attention matrices need H ( MN 2 + ( M − 1) N 2 ) =
H (2 M − 1) N 2 .
There are also MN 2 computations needed to add noise to M global attentionmatrices.
Thus, the number of computations needed to compute attention matrices in a FiSH with
M global attention heads and H local attention heads is N 2 M (2 D − 1) + 2 NMD (2 D x − 1)
+
H (2 M − 1) N 2 + MN 2 =
2( D + H )
M − H ] N 2 + 2 NMD (2 D x − 1) .
{+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 + 
+}
( H − M ) D (2 D x − 1) N computations in a forward pass.
When N is large, this difference is significant.
D.2 The Number of Parameters
(iii) Softmax H -head attention:
The number of parameters needed to compute the attention matricesin a softmax H -head attention is 2 HDD x .Explanation: 2 HDD x parameters is from the linear projects to calculate the query matrix Q and thekey matrix K in Step 1 in Section 1.1.
(iv) FiSH with M global attention heads and H local attention heads: The number of parametersin a FiSH with M global attention heads and H local attention heads is 2 MDD x + HM + M .Explanation: 2 MDD x parameters is from the linear projects to calculate M query matrices Qand M key matrices K , which are used to compute M global attention matrices.
The extra HMparameters is from the linear mapping for computing H local attention matrices from M globalattention matrices, and the extra M parameters are the M { σ k } Mk =1 for M global heads.
Softmax H-head attention versus FiSH with M global attention heads and H local attentionheads: Given the results in (iii) and (iv), when compared to the baseline softmax H -head attention, ourFiSH with M global attention heads and H local attention heads saves 2( H − M ) DD x − HM − Mparameters.
When D is large, this saving is significant.
E Proofs
In this appendix, we provide proof for Lemma 1.
E.1 Proof of Lemma
We denote [Equation] for all x ∈ R d where f ( x ) = [Equation]
√
[Equation]
πσ
[Equation] d exp (cid:16) − ∥ x [Equation] 2 (cid:17) for given σ > 0 .
From the work of [4], thespace of Gaussian mixtures is dense in the space of continuous probability measures.
Therefore, wecan find probability distribution G 1 such that
[Equation]
To obtain the conclusion of the lemma, it is sufficient to prove that we can find a probability measure
G 2 with at most K supports where K ≤ ( C log(1 /ϵ )) d for some universal constant C such that [Equation] [Equation]
Our technique for proving the above approximation bound relies on Lemma A.1 in [28].
In particular,that lemma entails that for any k ≥ 1 there exists a probability distribution G 2 with at most (2 k − 2) dsupports such that
[Equation]
for any α = ( α 1 , α 2 , . . . , α d ) ∈ N d such that 0 ≤ | α | = (cid:80) dj =1 α j ≤ 2 k − 2 , Here, θ α = (cid:81) dj =1 θ α j j .
Now, for any M ≥ 2 a √ d , we have ∥ x − θ ∥ ≥ ∥ x ∥ − ∥ θ ∥ > M − a √ d > M/ 2 as long as ∥ x ∥ > Mand θ ∈ [ − a, a ] d . It indicates that [Equation] On the other hand, for any k ≥ 1 we also have that
[Equation]
where the final inequality stems from
[Equation]
which is due to Eqn.
To further bound the right-hand-side (RHS) of Eqn.
(15), we use the following inequality:
[Equation]
| y | k
[Equation]
k ! for any y ∈ R . Since k ! ≥ ( k/e ) k for any k ≥ 1 , the above bound can be rewritten as(cid:12)(cid:12)(cid:12)(cid:12)(cid:12)(cid:12) [Equation] [Equation]
Further simplification of Eqn.
(15) leads to [Equation] where the final inequality is based on an application of inequality (16) with y = −∥ x − θ ∥ 2 / (2 σ 2 ) . For ∥ x ∥ ≤ M and θ ∈ [ − a, a ] d , we have ∥ x − θ ∥ ≤ ∥ x
∥ + ∥ θ ∥ ≤ M + a √ d .
Therefore, we furtherhave
[Equation]
When M ≥ 2 a √ d , we have M + a √ d ≤ 3 M 2 and the above bound leads to
[Equation]
By choosing M 2 = 8 σ 2 log(1 /ϵ ′ ) for some ϵ ′ > 0 , the bounds in Eqns.
(14) and (17) become
[Equation] As long as we choose k = 9 e 2 log(1 /ϵ ′ ) and ϵ ′ ≤ 1 , we havesup [Equation] By choosing ϵ ′ =
ϵ 2max { [Equation] 2 πσ ) d , 1 } , the results from Eqns. (18) and (19) indicate that
[Equation] ≤ ϵ 2 , and sup ∥ x ∥ >M | p G 1 ( x ) − p G 2 ( x ) | ≤ ϵ 2 . Therefore, if we choose M = 8 σ 2 log (cid:18) 2max { 2 ( √ 2 πσ ) d , 1 } ϵ (cid:19) and k = 9 e 2 log (cid:18) 2max { 2 ( [Equation]
2 πσ ) d , 1 } ϵ (cid:19) , wehave
[Equation]
It indicates that we obtain the conclusion of claim (12) by choosing K = (2 k − 2) d ≤ [Equation] [Equation]
(cid:19)(cid:19) d . As a consequence, we obtain the conclusion of the lemma. F Additional Benchmark Results F.1 UEA Time Series Classification
[Table]
We compare the accuracy of the GFiSHformer and the baseline softmax transformers trained on the
UEA Time Series Classification Archive benchmark [5]. In Table 12, we show that GFiSHformersperform on par with the baselines. For each classification task in this benchmark, the number of GFiSHformer’s global heads is half the number of heads in the baseline softmax transformers.
Theexperiment setups and configurations for the baseline and our GFiSHformer are the same as in [76](for the PEMS-SF, SelfRegulationSCP2, UWaveGestureLibrary datasets) and [78] (for other tasks). [Table] F.2 Reinforcement Learning on the D4RL Benchmark [Table]
In Table
[Table]
we report the results of the GFiSHformer and the softmax transformer trained for thecontinuous control tasks from D4RL benchmark [27] to evaluate the model performance on theoffline reinforcement learning.
On average, the 2-global-heads GFiSHformers perform comparablywith the 4-head transformer baselines.
For this benchmark, we follow the architecture and trainingconfiguration from [76]
[Table]