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 efficient and flexible transformers that allow the sharing of attention matrices between attention heads.
At the core of FiSHformer is a novel finite 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 significantly improving the computational and memory efficiency 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 classification.
On the WikiText-103, IWSLT’14 De-En and WMT’
EnDe, FiSHformers use much fewer floating-point operations per second (FLOPs), memory and parameters compared to the baseline transformers.
[Figure]
R a t i o
[Figure]
B
C
D
Model
[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 flexibility 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 ] ⊤ ∈ 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 ∈ 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 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 final 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 defined as
MultiHead ( { H } Hi =1 ) =
Concat ( H 1 , . .
, H H
[Equation]
W O .
1.2 Eigenvalue Analysis of the Attention Matrices
[Figure]
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 confirm 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 verifies 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 efficiency of transformers.
1.3 Contribution
Leveraging the idea of the finite admixture model (FAM) [57, 9], we propose a new class of efficient 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 efficiently 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 finite 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 efficiency 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 first review the finite admixture model (FAM) and derive a FAM of shared heads for the multi-head self-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 for all j .
Since the distributions P 1 , . . . , P H can have complex forms and be difficult 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 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 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 ′ such that we can find a 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 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 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 a workhorse 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 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 finite 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 defined as follows:
Definition 1 (Finite Admixture Model of Shared Heads) .
The multi-head attention admits a finite admixture model of shared heads if the attention matrices A j at the j th head are sampled from the following 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 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 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 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 )
=
(cid:80) 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 defined in Def. 1, is not differentiable, which poses a difficulty 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. (cid:80) Mk =1 p kj = 1 , p kj ≥ 0 , can be relaxed, and those mixing coefficients 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]
k
[Equation] [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 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, we formulate local attention matrices
A j in the Transformer with a Generalized Hard Finite Admixture of Shared Heads (Hard GFiSHformer) as A j = (cid:80) Mk =1 ϕ ( p kj Q k K k ⊤ ) .
[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]
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 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 2( 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 classification on ImageNet (Section 3.3), time series classification 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 efficiency 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] [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.
Efficiency:
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, 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 efficiency 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 efficiency 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 first 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 efficient than (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 [44], a state-of-the-art vision transformer architecture, for the image classification 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
[Table] [Table]
[Table] [Table]
slightly more efficient 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 Classification
We compare the accuracy of the GFiSHformer and the baseline softmax transformers trained on the
UEA Time Series Classification Archive benchmark [5].
In Table 6, we show that GFiSHformers perform 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.
The experiment setups and configurations 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 benefits 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 offline 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 configuration 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 confirm 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 efficient transformers are in the Appendix.
Applying (G)FiSH on Linear Transformers Linear transformers [37] is a class of efficient 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 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
and compare the efficiency 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 significantly more computation reduction (in both training and inference).
Furthermore, Fig.
6 shows that the efficiency 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 significant 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 justifies our use of nonlinearity mapping to generate local heads from global heads.
Using a pre-trained model, we first compute the covariance matrix of the vectorized attention scores of the l -th layer: C l = 1 M · H (cid:80) Mm =1 (cid:80)
Hh =
( A l,hm )( A l,hm ) ⊤ .
We use spectral decomposition to represent C l in terms of eigenvalues and eigenvectors, namely,[-C l = (cid:80) n 2 i =
-]
[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, confirming the advantage of (G)FiSH over MHA.
Table 13 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 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, efficiency, 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 efficiency are summarized in Fig.
7, 8 in Appendix C
[Table]
and ablation results on representation capacity are provided in Table 8, 9.
Related Work
Efficient Transformers Efficient transformers have been developed to reduce the quadratic computational and memory cost of transformers [62].
A class of efficient transformers are sparse transformers which design the attention matrix to have sparse structure [55, 42, 58, 12, 7].
Another class of efficient 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 efficiency 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 efficiency 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 finite 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 efficiency 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.
Checklist
The checklist follows the references.
Please read the checklist guidelines carefully for information on how 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 referencing the 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 instructions block 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] 2.
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] 3.
If you ran experiments... (a) Did you include the code, data, and instructions needed to reproduce the main experimental 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 experiments 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] 4.
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] 5.
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 applicable?
[N/A] (b) Did you describe any potential participant risks, with links to Institutional Review Board (IRB) approvals, if applicable?
[N/A] (c) Did you include the estimated hourly wage paid to participants and the total amount spent on participant compensation?
[N/A]
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 observe that in all of our experiments, choosing the global attention matrices in FiSH to be 1/4 and 1/2 of the number attention heads in the original MHA results in models with good accuracy and efficiency.
All of 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 long contextual dependencies.
The training set is made up of about 28 K articles containing 103 M running words; this corresponds to text blocks of about 3600 words.
The validation and test sets are composed of 218 K and 246 K running words, respectively.
Each of them contains 60 articles and about 268 K words.
Our experiment follows the standard setting [46, 66] and splits the training data into L -word independent long segments.
For evaluation, we use a batch size of 1 and go through the text sequence with a sliding window of size L .
We consider only the last position for computing perplexity (PPL) except in the first segment, where all positions are evaluated as in [2, 66].
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, containing about 4 .
5 M training sentence pairs.
Validation and test data are from the corresponding newest data.
The BLEU score [54] is used to measure the performance of the trained model.
A.3 Image Classification
Datasets and metrics The ImageNet dataset [63] contains about 1 .
M training images and 50 K validation images, the model learns to predict which one of 1000 classes an image belongs to.
Our Swin Transformer [44] 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 our global heads to the last 8 of the total 12 layers of the model, on each layer we set the number of global heads to half the number of heads, which are 6 and 12 global heads for layers with 12 andheads, respectively.
Our experiments were conducted on a server with 1 NVIDIA RTX 3090.
We set the batch size to 128 and the learning rate to 1.25e-4, all models are trained with single precision.
A.4 UEA Time Series Classification
Datasets and metrics There are 30 datasets in the benchmark [5].
Following [81], to evaluate our models on temporal sequences, we choose 10 datasets, which vary in input sequence lengths, the number of classes, and dimensionality.
We report the test accuracy as an evaluation for the benchmark.
Models and baseline The experiment setups and configurations for the softmax and our models are the same as in [81] 1 (for the PEMS-SF, SelfRegulationSCP2, UWaveGestureLibrary datasets) and [83] 2 (for the other tasks).
In all models, the number of heads is 8, whereas the model dimension and number of transformer layers are varied.
Our GFiSHformer uses 4 global heads to producelocal heads.
A.5 Reinforcement learning on the D4RL benchmark
Datasets and metrics The D4RL benchmark [29] contains the continuous control tasks for offline reinforcement learning.
We adapt the selection from [81], including HalfCheetah, Hopper, and Walker as experiment environments and Medium-Expert, Medium, and Medium-Replay as behavior policies.
Models and baseline The softmax baseline trained on this benchmark adopts the configuration from [81], with 3 transformer layers and 4 heads per layer.
The 2-global-head GFiSHformer also shares the same configuration and training set-ups.
A.6 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.7 and Table 3 in the main text.
The linear transformer reduces the quadratic computational cost of self-attention to linear complexity, in terms of the sequence length, by linearizing the softmax kernel [37].
We combine (G)FiSHformer with linear transformer by generating the global K T V and then sampling the local K T V from the global ones, resulting in the Transformer with a Linear Finite Admixture of Shared Heads (LFiSHformer).
Similar to (G)FiSHformer, we derive 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 baseline on the IWSLT’ 14 De-En task.
[Figure]
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
[Figure]
t
i
[Figure]
[Table] [Table]
[Table]
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.
B.3 More Results to Show that FiSHformer Helps Reducing Head Redundancy
Table 12 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 trained for the WikiText-103 language modeling task.
[Table] [Table]
B.4 More Results on Eigen Analysis
Table 13 presents the layer-average number of principle components for 95% variance explained of the 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.
Models are trained for the WikiText-103 language modeling task.
C Efficiency Analysis
In this section, we present the efficiency improvement of (G)FiSHformers over the MHA baseline on the WikiText-103 language modeling task.
We show that the advantage in the efficiency of (G)FiSHformers increases significantly as the model dimension D and sequence length N grows, making (G)FiSHformers more suitable and superior for large-scale applications.
In our analysis, we report the number of FLOPs, the number of model parameters, and memory usage (bytes) as the measures of model efficiency.
Analysis setting We investigate the benefits of our model computation and memory reduction through different D ∈ {256, 512, 1024, 2048, 4096} and N ∈ {128, 256, 512, 1024, 2048, 4096, 8192}.
For the FLOP calculation, we use fvcore.
Another notice is that, for model complexity, we distinguish between input-embedding parameters and non-embedding parameters.
Input-embedding parameters are used to represent the inputs before sending them to the model.
Non-embedding parameters are the parameters of the main model.
Since our method aims at reducing the size of the transformer model, it is important to compare the reduction in non-embedding parameters.
Hence, we report both model’s total parameters and non-embedding parameters in this analysis.
All measurements are calculated 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 and sequence lengths, Figure 3 indicates that our model requires significantly less computation than the baseline.
Especially for both training and inference, this substantially grows with the increases of D and N (up to more than 30 % reduction), which makes our method more preferable for long-range and large-scale tasks.
FiSH-based variants share comparable advantages in computation saving.
For further analysis, we compare 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 training and 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 substantial advantage in computational saving.
F L O
PS
Training
Testing
[Figure]
T e s t
M e m o r y
( G
B ) P a r a m e t e r s ( M )
[Figure]
C
D
[Figure] [Figure]
GFiSHformer computation efficiency rapidly increases as the number of global heads decreases We compare the computational-cost reduction of GFiSHformer as the global heads vary.
Figure 7 shows the FLOP ratio of GFiSHformer with 2/4/6-global-head versus the 8-head baseline for a given D and various sequence lengths.
As the number of heads decreases, GFiSHformers achieve significant computation reduction (in both training and inference)
(Figure 7).
(G)FiSHformers improves memory usage and model complexity
In addition to the computational saving, 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 present the advantage of the GFISHformers compared to the 8-head MHA Transformer.
Figureshows the number-of-parameter-ratio (Left) and the memory-ratio (Right) of the GFiSHformers with 2/4/6 global heads and the 8-head baseline.
At a fixed sequence length N = 1024 , as we scale up the 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, which further indicates that all variants of FiSHformers benefit space complexity.
Figure 6 shows that 2-head FiSHformers, Hard FiSHformers, GFiSHformers, and Hard GFiSHformer have a comparative number of total/non-embedding parameters and memory usage at different model dimensions, for a given sequence length.
Finally, we examine the model space complexity of a Fishformer variant as the number of global heads 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 .
While only having the fewest heads, 2-head GFiSHformer is the most efficient model (lowest metric ratio
Model Dimension
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 on - e m bedd i ng
P a r a m e t e r R a t i o  6 Number of Global Heads
[Figure]
with the 8-head MHA baseline) and achieves comparable performance with the baseline, for model dimension, and sequence length are 128 and 256 respectively, as indicated in Table 3.
D An Analysis on the Computational Complexity and the Number of Parameters in FiSH and the Softmax Attention
In this section, we compare the computational complexity and the number of parameters in the FiSH with M global attention heads and H local attention heads to the H -head baseline MHA softmax transformer.
Following the same notation in Section 1.1 in the main text, we let D x , N , and D be the input dimension, the input length, and the model/feature dimension, respectively.
To simplify the computation, 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 matrices in 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 at each head, we need 2 NDD x multiplications and 2 ND ( D x − 1) additions.
In total, these need 2 ND (2 D x − 1) computations.
Next, to compute the product QK ⊤ in Eqn.
(1), we need N 2 D multiplications and N 2 ( D − 1) additions.
In total, computing an attention matrix in Eqn.
(1) at each 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
−
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 attention matrices.
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
[Equation] [Equation]
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 matrices in 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 the key matrix K in Step 1 in Section 1.1.
iv) FiSH with M global attention heads and H local attention heads:
The
number of parameters in 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 Q and M key matrices K , which are used to compute M global attention matrices.
The extra HM parameters is from the linear mapping for computing H local attention matrices from M global attention 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 attention heads:
Given the results in (iii) and (iv), 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 ) DD x − HM − M parameters.
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]
(cid:17) for given σ > 0 .
From the work of [4], the space of Gaussian mixtures is dense in the space of continuous probability measures.
Therefore, we can 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]
Our technique for proving the above approximation bound relies on Lemma A.1 in [30].
In particular, that lemma entails that for any k ≥ 1 there exists a probability distribution G 2 with at most (2 k − 2)
d supports 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
[Equation]
we have ∥ x − θ ∥ ≥ ∥ x ∥ − ∥ θ ∥ > M − a √ d > M/ 2 as long as ∥ x ∥ > M and θ ∈ [ − 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.
[Equation]
we use the following inequality:
[Equation]
for any y ∈ R .
Since k ! ≥ ( k/e ) k for any k ≥
[Equation]
the above bound can be rewritten as
[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 further have
[Equation]
When M ≥ 2 a √ d , we have M + a √ d ≤ 3 M 2 and the above bound leads to
[Equation]
[Equation]
By choosing M 2 = 8 σ 2 log(1 /ϵ ′ ) for some ϵ ′ > 0 , the bounds in Eqns.
(14) and (17) become
[Equation]
k k .
[Equation]
As long as we choose k = 9 e 2 log(1 /ϵ ′ ) and ϵ ′ ≤ 1 , we have
[Equation]
By choosing ϵ ′ =
ϵ 2max { 2 ( √ 2 πσ ) d , 1 } , the results from Eqns.
(18) and (19) indicate that sup
[Equation]
and sup
[Equation]
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 ( √ 2 πσ ) d , 1 } ϵ (cid:19) , we have
[Equation]
It indicates that we obtain the conclusion of claim (12) by choosing K =
(2 k − 2) d ≤ (cid:18) 18 e 2 log (cid:18) 2max {
[Equation]
[Equation]
ϵ (cid:19)(cid:19) d .
As a consequence, we obtain the conclusion of the lemma.