Improving Transformer with an Admixture of Attention Heads
Abstract
Transformers with multi-head self-attention have achieved remarkable success in sequence modeling and beyond.
However, they suffer from high computational and memory complexities for computing the attention matrix at each head.
Recently, it has been shown that those attention matrices lie on a low-dimensional manifold and, thus, are redundant.
We propose the Transformer with a Finite Admixture of
Shared Heads (FiSHformers), a novel class of efﬁcient and ﬂexible transformers that allow the sharing of attention matrices between attention heads.
At the core of FiSHformer is a novel ﬁnite admixture model of shared heads (FiSH) that samples attention matrices from a set of global attention matrices.
The number of global attention matrices is much smaller than the number of local attention matrices generated.
FiSHformers directly learn these global attention matrices rather than the local ones as in other transformers, thus signiﬁcantly improving the computational and memory efﬁciency of the model.
We empirically verify the advantages of the FiSHformer over the baseline transformers in a wide range of practical applications including language modeling, machine translation, and image 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.
[Figure]
o
[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
Introduction
Transformers have become the state-of-the-art model for solving many challenging problems in natural language processing [75, 1, 17, 80, 21, 10, 35, 61] and computer vision [19, 68, 22, 73].
Transformers learn from unlabeled data effectively and take advantage of the pre-trained models on downstream tasks that involve different data modalities with limited supervision [59, 60, 21, 82, 43].
The success of transformer is credited to the multi-head self-attention (MHA) mechanism as their fundamental building block.
For each token in the sequence, self-attention in transformers aggregates information from other tokens by computing a weighted average of their feature representations with a weight proportional to a similarity score between the representations.
This attention mechanism allows arbitrary input-dependent interaction between tokens in the sequence where a token can pay attention to other tokens and attain a contextual representation [6, 75, 38].
Multi-head selfattention captures multiple such contextual representations, one at each head, thereby increasing the representation capacity of the self-attention.
It has been argued that the representation capacity of the attention mechanism [72] and its ﬂexibility in capturing diverse syntactic and semantic relationships [72, 76, 15, 77, 32] 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-attention transforms X into the output sequence H in the following two steps:
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 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 compute the 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 structure that the attention matrix A learns from training determines the ability of the self-attention to capture contextual representation for each token.
Multi-head Attention (MHA)
Each output sequence H forms an attention head.
In MHA, multiple heads 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
[Equation]
1.2 Eigenvalue Analysis of the Attention Matrices
The multi-head mechanism allows transformers to capture more diverse attention patterns and increase the capacity of the model.
However, in many practical tasks, transformers learn redundant heads [47, 78], whose learned attention matrices lie on a low-dimensional manifold [8].
To conﬁrm this claim, in Figure 2, we follow the eigenvalue analysis in [8] and investigate the eigenvalues of the covariance matrix of vectorized attention matrices aggregated over each layer of a transformer model trained on the WikiTextdataset for language modeling.
We observe that this covariance matrix is low rank with top(1.2%) eigenvalues capturing more than 90% of the energy.
This result veriﬁes that the variability of learned attention matrices in transformers can be explained by a relatively small number of principal components, and those attention matrices lie on a low-dimensional manifold.
Therefore, in multi-head attention, the effective number of heads is much smaller than the actual number of heads, and a more effective 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) [57, 9], we propose a new class of efﬁ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 an admixture of a small number of global attention matrices.
This sharing mechanism between heads helps reduce the computational complexity and the model size compared to the MHA softmax transformer.
Our contribution is three-fold:
We construct an admixture model for shared attention matrices between heads and propose FiSHformer, 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.
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 model of shared heads and then present our FiSHformer and its extensions.
In Section 2.5, we analyze the reduction in model complexity and computational cost from FiSH.
In Section 3 and 4, we validate and empirically analyze the efﬁciency and accuracy of FiSHformer, as well as conducting ablation studies on the model.
We discuss related works in Section 5.
The paper ends up with concluding remarks.
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-head self-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 for all j .
Since the distributions P 1 , . . . , P H can have complex forms and be difﬁcult to compute, our approach is to consider approximated distributions of P j and these approximated distributions have simple 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 and admits differentiable and bounded density estimation p .
Then, for any scale parameter  > 0 and for any ✏ > 0 , there exists universal constant C and M  ( C log(1 / ✏ ))
D 0 such that we can ﬁnd a mixture of M components
[Equation]
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 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 be computationally expensive.
To overcome this issue, we assume that M 1 =
M 2 = . . .
=
M H and the location parameters ( ✓ 1 j , . .
. , ✓ M j j )
=
( ✓ 1 , . .
. , ✓ M ) for all j , i.e., these approximated mixtures share a similar set of location parameters.
This sharing information of location parameters has a deep connection 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 a workhorse 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 chosen probability 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 mixture distributions 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), in which A j follows ﬁnite admixture distribution of M components given by
[Equation] [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 means and coveriances, respectively.
FiSH is then deﬁned as follows:
Deﬁnition 1 (Finite Admixture Model of Shared Heads) .
The multi-head attention admits a ﬁnite admixture model of shared heads if the attention matrices A j at the j th head are sampled from the following ﬁ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 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 length of the input sequence.
In contrast, FiSH only need to compute M < H , e.g. M = 2 , global attention matrices, each of which also requires O ( N 2 ) computational and memory costs.
Then FiSH combines those global attention matrices to form a FAM from which H , e.g. H = 8 , local attention matrices are sampled as in Eqn.
This second step of sampling local attention matrices from a set of global attention 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 the word w whose topic is c , pLSA models the occurrence of the word w in the document d as a mixture of conditionally independent Multinomial distributions p ( w | d ) =
P c p ( c | d ) p ( w | c ) .
Comparing 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 the local attention matrices in FiSH are words sampled from those topics.
It is interesting to note that pLSA is equivalent to the famous Latent Dirichlet Allocation model under a uniform Dirichlet prior on 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 not differentiable, which poses a difﬁculty in training FiSHformers.
Applying the reparameterization trick [39], 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 convex combination condition of p kj , i.e.
P
[Equation]
p
[Equation]
= 1 , p kj  0 , can be relaxed, and those mixing coefﬁ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 shared heads (MiSHformer) can be used to reduce the amount of computation with the cost of accuracy reduction.
MiSH is a special case of FiSH when the mixture weights p kj are the same for all j .
The local 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 in learning representation by replacing the linear mapping in Eqns.
(8) and (9) by a nonlinear mapping such as 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, we formulate local attention matrices
A j in the Transformer with a Generalized Hard Finite Admixture of Shared Heads (Hard GFiSHformer) as A j =
[Equation]
[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
A
B
C
D
[Figure]
￼￼
Training
Testing
[Figure]
[Table]
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 heads saves [2( H  M ) D  2 MH )] N 2 + 2( H  M ) D (2 D x  1) N FLOPs in a forward pass and
( 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), image classiﬁcation on ImageNet (Section 3.3), time series classiﬁcation on the UEA benchmark (Section 3.4), and reinforcement learning on the D4RL Benchmark (Section 3.5).
We aim to show that: (i) FiSHformers improve the efﬁciency and accuracy upon the
MHA baseline; (ii) FiSH is a universal method that can be applied to state-of-the-art transformer models to improve their performance on large-scale applications.
In Section 4, we also show that
FiSH helps reduce the redundancy between attention heads.
We compare FiSHformers, Hard FiSHformers, GFiSHformers, and Hard GFiSHformers with the baseline MHA softmax transformers.
In our experiments, we apply the discriminative relaxation explained 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 followed by a linear neural network.
All of our results are averaged over 5 runs with different seeds.
More details on datasets, models, and training are provided in Appendix A. Our PyTorch code with documentation can be found at https://github.com/minhtannguyen/FishFormer.
3.1 WikiText-103 Language Modeling
Models and baselines We compare the 2 and 4-global-head FiSHformers with the 8-head softmax transformers [75].
Each model has 16 layers, and our training follows the setting from [66].
Results Perplexity: Table 3.1 demonstrates that our 2/4-global-head (G)FiSHformers and their hard versions obtain comparable or better PPLs than the corresponding 8 head MHA baseline on
[Table]
[Table] [Table]
WikiText-103.
Interestingly, the 2-global-head (G)FiSHformers perform on par with the 8-head baseline 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 of local attention matrices in (G)FiSH, though being generated from only 2 global bases, is comparable to those in the 8-head MHA
[Table]
Efﬁciency: In Fig.
3A and 3B
[Table]
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 model dimension D and sequence length N .
In Fig.
3C and 3D
[Table]
we show the reduction ratio of model size and GPU memory usage at test time, respectively, of our 2-global-head GFiSHformer vs. the same baseline.
We observe that the efﬁciency advantage of GFiSHformer over the baseline grows with D and N , making it more suitable and superior for large-scale applications.
Note that the model size in terms of the number of parameters does not depend on the sequence length N , and from our experiments, we observe that the GPU memory usage reduction ratio is almost the same for different sequence lengths.
More efﬁciency analysis results on this language modeling task are 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 translation task, an important task in natural language processing in which the sequence lengths of the input are not the same.
We ﬁrst compare (G)FiSHformers with the baselines MHA softmax transformers on the IWSLT’ 14 De-En [11] and then scale up our experiments to the WMT’14 En-De [45].
On these tasks, we calculate the BLEU scores for evaluation.
Models and baselines For the IWSLT’ 14 De-En task, we compare 2-global-heads (G)FiSHformers with the baseline 4-head softmax transformer.
Each model consists of 12 layers, 6 layers for an encoder 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 and model setting are the same as those in [53].
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 attention matrices generated by (G)FiSH.
Furthermore, (G)FiSHformer outperforming Hard (G)FiSHformer in all settings suggests the positive value of adding noise into the models to turn them into a proper probabilistic model.
Nevertheless, it is worth noticing that Hard (G)FiSHformer is more efﬁcient than (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 [44], a state-of-the-art vision transformer architecture, for the image classiﬁcation task on the ImageNet dataset [20].
The baseline Swin-T we use has a total of 12 layers, across 4 stages of transformer blocks with 3, 6, 12, and 24 heads each.
Our GFiSH Swin-T uses 6 and 12 global attention heads at the last two stages.
The model and training follow the settings in [44].
We summarize our results in Table 5.
Our GFiSHformer is only translation.
Our methods obtain comparable/better results than
[Table] [Table]
[Table] [Table]
slightly more efﬁcient than the baseline in this case because the sequence length N per window for this task is small, i.e. N = 49 .
However, like with the previous language tasks and as pointed by formula of computational cost and model complexity reduction in Section 2.5, these advantages grow with larger D and N .
3.4 UEA Time Series Classiﬁcation
We compare the accuracy of the GFiSHformer and the baseline softmax transformers trained on the
UEA Time Series Classiﬁcation Archive benchmark [5].
In Table 6, we show that GFiSHformers perform on par with the baselines.
For each classiﬁcation task in this benchmark, the number of
GFiSHformer’s global heads is half the number of heads in the baseline softmax transformers.
The experiment setups and conﬁgurations for the baseline and our GFiSHformer are the same as in [81] (for the PEMS-SF, SelfRegulationSCP2, UWaveGestureLibrary datasets) and [83] (for other tasks).
3.5 Reinforcement Learning on the D4RL Benchmark
We also demonstrate the beneﬁts of GFiSHformers on reinforcement learning.
In Table 7, we report the results of the GFiSHformer and the softmax transformer trained for the continuous control tasks from D4RL benchmark [29] to evaluate the model performance on the ofﬂine reinforcement learning.
On average, the 2-global-heads GFiSHformers perform comparably with the 4-head transformer baselines.
For this benchmark, we follow the architecture and training conﬁguration from [81].
3.6 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 the adaptively 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
[Table] [Table]
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 whileattention heads are used in each AST layer.
3.7 Beyond Multi-Head Softmax Transformers
We show that (G)FiSH can be applied on top of many transformer architectures to improve their performance including the linear transformers [37] and the SoTA transformer with noisy backtranslation [26].
More results of combining (G)FiSH with efﬁcient transformers are in the Appendix.
Applying (G)FiSH on Linear Transformers Linear transformers [37] is a class of efﬁcient transformers that linearize the softmax kernel in Eqns.
1 and 2 when computing attention matrices.
We apply (G)FiSH on linear transformers trained for the IWSLT14 De-En machine translation task and summarize the results in Table 3.
The empirical results verify that applying (G)FiSH using only 2-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 [26] for the WMT’14 En-De translation task and obtain the BLEU score of 33.45.
This result is comparable to the SoTA result of 33.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
and compare the efﬁciency of our FiSH-based models in Fig.
Fig.
6, 7, 8 and details on our setting are provided in Appendix C. From Fig.
7 and 8, we observe that when using fewer number of global heads, GFiSHformers achieve signiﬁcantly more computation reduction (in both training and inference).
Furthermore, Fig.
6 shows that the efﬁciency measures, i.e. FLOPs, model size, and GPU memory 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 more signiﬁcant distances between heads than the baseline.
Thus, our models capture more diverse patterns across heads than the baseline.
For a given pre-trained model, we compute the pair-wise L 2 distances between heads in the same layer.
We show the layer-average mean and variance of distances between heads in GFiSHformers compared with those in the MHA softmax baselines in Table 8.
We provide additional results for Hard GFiSHformers and Hard FiSHformer in Table 12 in the Appendix.
Eigen Analysis We show that heads in GFiSHformer lie on a higher-dimensional subspace compared to those in FiSHformer.
This justiﬁes our use of nonlinearity mapping to generate local heads from global heads.
Using a pre-trained model, we ﬁrst compute the covariance matrix of the vectorized attention scores of the l -th layer: C l =
[Equation]
A l,hm ) > .
We use spectral decomposition to represent C l in terms of eigenvalues and eigenvectors, namely, C l =
[Equation]
Without losing generality, we assume that eigenvalues are sorted in descending order.
We illustrate the layer-average number of principle components that are needed to explain 95% variance in Table 9.
Interestingly, Table 9 shows that attention matrices in all of our proposed FiSH-based models lie on higher-dimensional subspace than those in the baseline MHA softmax transformers, which indicates that our models achieve better representational capacity than the baseline, conﬁrming the advantage of (G)FiSH over MHA.
Table 13 in the Appendix provide additional results for Hard (G)FiSHformer.
[Table] [Table]
[Table]
Admixture vs. Mixture of Heads We compare the transformer with a mixture of heads and the transformer with an admixture of heads.
We show that the transformer with a mixture of heads yields worse accuracy.
We summarize our results on the WikiText-103 language modeling task in Table
[Table]
Ablation Study on the Effect of the Number of Global Heads on FiSH-based Models We investigate the accuracy, efﬁciency, and representation capacity of FiSH-based models under different numbers of global heads on the WikiText-103 language modeling task.
Since GFiSH obtains the best accuracy on this task, we use GFiSH as a study case in this section and report our results on accuracy in Table 11.
Ablation results on efﬁciency are summarized in Fig.
7, 8 in Appendix C
[Table]
and ablation results on representation capacity are provided in Table 8, 9.
Related Work
Efﬁcient Transformers Efﬁcient transformers have been developed to reduce the quadratic computational and memory cost of transformers [62].
A class of efﬁcient transformers are sparse transformers which design the attention matrix to have sparse structure [55, 42, 58, 12, 7].
Another class of efﬁcient transformers are models that integrate different access patterns for better coverage [12, 33].
These access patterns can also be learned from the data [40, 62, 71].
In other works, a side memory module is used to access multiple tokens simultaneously [41, 69, 3, 7].
Recently, low-rank and kernelization methods have been proposed to improve the computational and memory efﬁciency of computing self-attention [74, 79, 37, 14, 50, 67, 52, 56].
Our (G)FiSHformers are complementary to these methods.
Redundancy in Transformers Most of the neurons and heads in the pre-trained transformer are redundant and can be pruned when applied on downstream tasks [18, 47, 23].
The contextualized embeddings in pre-trained networks under this redundancy have also been studied to demonstrate that the representations learned within these models are highly anisotropic [48, 27].
Knowledge distillation and sparse approximation have also been used to enhance the efﬁciency of transformers, including [65, 70, 78, 64].
Our FiSH-based approach are complementary to these methods
Mixture Models for Transformers Recently, mixture models have been employed to study and enhance transformers.
Among these works is switch transformers [28] that uses the routing algorithm in Mixture of Experts (MoE) to decrease the communication and computational costs in transformers.
[49] derives a GMM for each attention head.
Other works that combine mixture models with transformers include [51, 13, 31, 36].
Concluding Remarks
In this paper, we proposed the FiSHformer, a class of transformers that samples local attention matrices from a ﬁnite admixture of global attention matrices.
FiSHformers and their generalized version GFiSHformers attain better computational cost and model complexity than their baseline
MHA softmax transformer counterparts.
Furthermore, (G)FiSHformers help increase the diversity between 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, and this is a limitation of our model.
It is interesting to impose additional structures such as low-rank and sparsity into the global attention matrices.
Finally, establishing theoretical guarantee for optimization algorithms [25, 24, 34] for solving FiSHformers is also an important future direction to improve the efﬁciency of FiSHformers.
Acknowledgements
This material is based on research sponsored by the AFOSR MURI FA9550-18-1-0502, the ONR grant N00014-20-1-2093, the MURI N00014-20-1-2787, and the NSF under Grant# 2030859 to the
Computing Research Association for the CIFellows Project (CIF2020-UCLA-38).
NH acknowledges support from the NSF IFML 2019844 and the NSF AI Institute for Foundations of Machine Learning.