D FSS A TTEN : D YNAMIC F INE - GRAINED S TRUCTURED S PARSE A TTENTION M ECHANISM
Anonymous authors
Paper under double-blind review
A BSTRACT
Transformers are becoming mainstream solutions for various tasks like NLP and
Computer vision.
Despite their success, the quadratic complexity of their attention mechanism hinders them from applying to latency-sensitive tasks.
Tremendous efforts have been made to alleviate this problem, and many of them successfully reduce the asymptotic complexity to linear.
Nevertheless, few of them achieve practical speedup over the original full attention, especially under the moderate sequence length.
In this paper, we present D FSS A TTEN , an attention mechanism that dynamically prunes the full attention weight matrix to the 50% fine-grained structured sparse pattern used by the sparse tensor core on NVIDIA A100 GPU.
We provide both theoretical and empirical evidence that demonstrate D FSS A TTEN is a good approximation of the full attention mechanism and can achieve speedups in wall-clock time under arbitrary sequence length.
We evaluate our method on tasks from various domains under different sequence lengths from 256 to 4096.
D FSS A TTEN achieves 1 .
27 ∼ 1 . 89 × speedups over the full-attention mechanism with no accuracy loss on A100 GPU.
I NTRODUCTION
Transformers (Vaswani et al., 2017) have achieved competitive performance across various domains like NLP (Ott et al., 2018) and Computer Vision (Dosovitskiy et al., 2021).
The key feature that sets them apart from traditional neural network architectures is the attention mechanism (Vaswani et al., 2017), which allows the transformers to gather information from the embeddings of elements in the input sequence in an adaptive and learnable manner.
Nevertheless, the high computation cost and memory footprint brought by the attention mechanism make it difﬁcult to apply transformers to latency-sensitive tasks.
Many efﬁcient attention mechanisms(Tay et al., 2020b; Zaheer et al., 2020; Beltagy et al., 2020; Tay et al., 2020a; Roy et al., 2021;
Kitaev et al., 2020) have been proposed to address this issue.
However, most of them drastically modify the original full attention mechanism and introduce a handful of hyper-parameters to tune.
Therefore, they require tremendous engineering effort to deploy and optimize.
Besides, they usually need to be trained from scratch instead of exploiting pretrained models like BERT (Devlin et al., 2019).
Some of them rely on ﬁxed sparse patterns or extremely high sparsity to achieve wall-clock time speedup.
Therefore, these methods usually require thousands of pretraining or ﬁne-tuning steps on speciﬁc tasks and toilsome tuning of several hyper-parameters to reach good accuracy.
Last but not least, previous methods usually introduce additional operators like top-k, sort that cause large overheads and offset their beneﬁts at moderate sequence length.
In this paper, we present D FSS A TTEN , a simple and effective sparse attention mechanism that address the limitations mentioned above.
D FSS A TTEN dynamically prunes the full attention score matrix using 50% fine-grained structured sparse patterns (NVIDIA, 2020).
This pattern is GPU friendly and can leverage the new sparse tensor core on NVIDIA A100 GPU (Mishra et al., 2021).
Our D FSS A TTEN offers several advantages over existing studies:
• It requires minimal changes to the original full-attention with no hyper-parameters to tune.
This makes it a drop-in replacement of the full attention that only requires to change a few lines of code.
Moreover, it can directly exploit existing pretrained models like BERT (Devlin et al., 2019) and RoBERTa (Liu et al., 2020).
• We dynamically prune the attention scores based on their magnitude under only 50% sparsity.
This allows the pruned attention matrix to reserve the important entries, achieving on par model accuracy with full attention even without ﬁne-tuning.
• Our method introduces zero overhead on existing GPU hardware.
As a result, we are able to achieve wall-clock time speedup and memory footprint reduction over the full attention in arbitrary sequence length.
To conclude, our main contributions are summarized below:
• We propose D FSS A TTEN , a dynamic sparse attention mechanism that is a drop-in proxy of the full attention mechanism.
Its effectiveness is justiﬁed by both empirical and theoretical evidence.
• We present a dedicated CUDA kernel design to completely remove the pruning overhead.
The pruning is implemented as an epilogue of the dense matrix multiplication which produces the attention score matrix.
• We evaluate D FSS A TTEN on tasks cross various domains and sequence lengths.
It achieves 1 .
27 ∼ 1 .
× speedup over the full attention with no accuracy loss.
B ACKGROUND AND M OTIVATION
We ﬁrst introduce the preliminaries, notations, and background of our paper.
2. F ULL A TTENTION M ECHANISM
Given an input sequence X =( x 1 ,.., x n ) ∈ R n × d , the full attention mechanism can be deﬁned as
[Equation]
where Q = XW q , K = XW k , and V = XW v are query, key, and value matrices.
QK T forms a full-quadratic adjacency matrix, whose edge weights are the dot-product similarity between all the elements in the sequence.
This adjacency matrix is standardized with 1 / √ d to keep the unit second moment and then normalized with softmax.
At last, the row feature vectors in V are aggregated according to the normalized adjacency matrix by multiplying them together.
In the rest of this paper, we denote A =
Softmax ( QK T / √ d ) for simplicity.
We refer QK T as the attention score matrix and A as the attention weight matrix.
2. E FFICIENT A TTENTION M ECHANISM
The high computation cost and memory footprint in the full attention mechanism come from A , whose size grows quadratically with the sequence length n .
To address this issue, various efﬁcient attention mechanisms have been proposed (Tay et al., 2020b).
Fixed Sparse Patterns .
Zaheer et al.
; Beltagy et al.
(2020) apply a set of ﬁxed sparse attention patterns on A , like global attention and sliding window attention.
These patterns are constructed from empirical observations and designed GPU-friendly to achieve wall-time speedup.
However, as these patterns are designed empirically and ﬁxed during inference, there is no guarantee that they can always capture the important entries in A or transfer easily across different tasks.
Dynamic Sparse Patterns .
Ham et al.
(2021) dynamically generate ﬁne-grained sparse attention patterns on A with low-cost binary hashing.
However, this technique requires specialized hardware to achieve speedup, so it is not available on general-purpose hardware like GPU.
Tay et al.
(2020a);
Roy et al.
; Kitaev et al.
(2020) apply various clustering methods and only compute the attention within each cluster.
Although computing full attention in each cluster is more friendly to
GPU compared with ﬁne-grained sparsity, the clustering methods contain several GPU-unfriendly operators like top-k and sorting that offsets their beneﬁts under moderate sequence length.
Low Rank / Kernel .
Wang et al.
(2020) project A from n × n to n × k with linear projection.
Choromanski et al.
(2021) introduce the FAVOR+ which approximates the softmax with kernel method.
This allows them to change the computation order and reduce the asymptotic complexity
[Figure]
to linear
However, the low-rank projection and kernel construction also introduce considerable overhead.
This makes these methods only effective under long sequence length.
Besides, the previous studies drastically change the attention mechanisms, tens of thousands pretraining or ﬁnetuning steps are required to reach a comparable performance with the origin full attention mechanism.
So they require tremendous engineering effort to deploy.
2. F INE - GRAINED S TRUCTURED S PARSITY IN NVIDIA A100 GPU
NVIDIA introduces the ﬁne-grained structured sparsity in the A100 GPU.
As shown in Figure 1, the dense input matrix is pruned with ﬁne-grained structured pruning.
If the data type is ﬂoat, 1: sparsity is used which selects the larger one in two consecutive entries.
If the data type is bﬂoat or ﬂoat16, the 2:4 sparsity is used which selects two larger ones among four consecutive elements.
After the pruning, the result is compressed to nonzeros and metadata.
The nonzeros contain the value of reserved data that is 50% smaller than the original one.
The metadata records the index of the nonzeros in the origins matrix.
It takes 4 bit metadata to record the decision of each 1:2 or 2: selection.
Therefore, the metadata is only 1/16 of the original dense matrix in terms of bits.
This compressed sparse matrix can be multiplied with a dense matrix under the support of the sparse tensor core to achieve signiﬁcant speedup.
This ﬁne-grained structured sparsity has been applied to the static weight matrices in various neural network models including transformer (Mishra et al., 2021).
It can effectively accelerate the feedforward part of the transformer up to 1 . 9 × .
However, to the best of our knowledge, no previous studies use it in the attention mechanism where the attention weight matrix is dynamically generated for each sequence.
One plausible explanation is that during pruning, GPU must read the whole dense matrix to be pruned from the memory.
Then, after selecting the elements to be reserved under the 1:2 and 2:4 pattern, it must also generate the metadata encoded in a special format such that the metadata can be used efﬁciently later.
All these overheads will offset the beneﬁt brought by the pruning if we do it on the ﬂy.
D YNAMIC F INE - GRAINED S TRUCTURED S PARSE A TTENTION M ECHANISM
In this section, we ﬁrst give an overview of our D FSS A TTEN method.
Then, we discuss the design considerations of exploring sparsity in attention and the choice of sparse granularity in our method for GPU-friendly implementation and effectiveness.
Finally, we brieﬂy introduce our GPU kernel design to remove pruning overhead.
Our proposed D FSS A TTEN mechanism is simple and effective, as illustrated in Figure 2.
Compared with the full-quadratic attention mechanism, our method dynamically prunes attention scores without incurring storage or computation overhead, while maintaining the effectiveness of attention.
More importantly, our method can achieve practical speedups of attention on existing GPU hardware with customized CUDA kernels.
Listing 1 shows all the modiﬁcations to be made to use
D FSS A TTEN .
𝑛 𝑄 𝐾 !
𝑉 𝑌 𝑛/2 𝑛/ 𝑛 𝑑 𝑑 𝑛 𝑑 𝑄𝐾 !
𝜎(𝑄𝐾 ! )
𝜎 𝑂 𝑛 𝑄 𝐾 ! 𝑉 𝑌 𝑛 𝑛 𝑑 𝑑 𝑛 𝑑 𝜎 𝑂 𝑄𝐾 !
𝜎(𝑄𝐾 ! )
[Figure]
Full Attention
[Figure]
# Full attention mechanism def full_attention(q,k,v): attn_weight
=torch.bmm(q, k.transpose(1, 2))
attn_weight = torch.nn.functional.softmax(attn_weight, -1) return torch.bmm(attn_weight, v) # DFSS attention mechanism import dspattn def dfss_attention(q,k,v):attn_weight,metadata= dspattn.bmm(q, k)
attn_weight = torch.nn.functional.softmax(attn_weight, -1) return dspattn.spmm(attn_weight, metadata, v)
Listing
: Example of using D FSS A TTEN .
The “dspattn” is the package we developed.
3. D ESIGN C ONSIDERATIONS FOR E XPLOITING A TTENTION S PARSITY
[Figure]
As illustrated in Figure 3, the attention mechanism can be considered as three stages: QK T , Softmax , and AV .
To design a sparse attention mechanism, the ﬁrst decision to make is where should we induce pruning to sparsify the attention.
If we start from 0 , all the three stages will be beneﬁted from the sparsity given effective implementation on GPU: the dense matrix multiplication between Q and K will be replaced with the sampled dense-dense matrix multiplication (SDDMM) which only computes the entries identiﬁed by the sparse pattern.
The Softmax only operates on the nonzero values in each row.
The original dense matrix multiplication between A and V will be replaced with a sparse matrix-matrix multiplication (SpMM) which multiplies a sparse matrix with a dense matrix.
However, as it is not possible to exactly know which entry in QK T has higher magnitude before computing QK T , starting from 0 usually requires some additional components to predict the location of important entries.
Starting from 1 requires us to compute a dense matrix multiplication between Q and K .
The beneﬁt is that we can explicitly select important entries from QK T without prediction.
As the softmax is a monotonically increasing function, starting from 2 does not offer any beneﬁts overbut throws away the opportunity to accelerate Softmax .
In this paper, we choose to start from 1 based on two considerations.
First, replacing the dense matrix multiplication with SDDMM at QK T offers limited speedup even at high sparsity.
Chen et al.
(2021b) show that it is difﬁcult for SDDMM to achieve speedup over its dense counterpart under 80% sparsity even with some structured design.
Second, starting from 1 allows us to keep our design simple such that it does not introduce additional overhead or hyper-parameters to tune.
3. O N THE G RANULARITY OF S PARSE A TTENTION P ATTERNS
The second decision to make is what sparse pattern to use as it will tremendously affect the latency of
SpMM as well as the overhead to encode the sparse QK T .
Existing studies exploit various sparse encoding schemes.
For instance, the compressed sparse row (CSR) is popular for encoding ﬁnegrained sparsity.
However, CSR-based SpMM requires over 95% sparsity to be on par with its dense counterpart (Chen et al., 2021b).
Block sparsity is also widely used as it can bring considerable wall-time speedup at moderate sparsity given large block size.
However, it cannot capture some ﬁne-grained attention patterns.
Moreover, these patterns require data comparisons within the entire row, which is difﬁcult to execute in parallel and unfriendly to GPUs.
We ﬁnd the ﬁne-grained structured sparsity mentioned in Section 2.3 is a good choice as long as we address the pruning overhead.
On one hand, the 1:2 and 2:4 selections are performed locally and are easy to execute in parallel.
On the other hand, the size of compressed nonzeros is half of the original dense attention matrix, so the softmax only needs half of the computations.
Powered by the
NVIDIA Sparse Tensor Core, the SpMM between the compressed A and V can also achieve 1.7 × speedup.
3. E MPIRICAL R ESULTS OF D FSS A TTEN M ECHANISM
[Table]
Empirically, we ﬁnd this pattern can well approximate the full attention mechanism.
We ﬁrst ﬁnetune a BERT-large model on SQuAD v1.1 under full attention.
Then, we directly replace the full attention with the 1:2 and 2:4 attention without additional ﬁnetuning.
The F1-scores are summarized in Table 1 under Cl = 95% .
It is obvious that the accuracy loss is only around one sigma even without ﬁnetuning.
3. R EMOVING P RUNING O VERHEAD
As mentioned in Section 2.3, the major challenge that hinders us from using the ﬁne-grained structured sparse attention is the pruning overhead.
We observe that when computing QK T , the results are ﬁrst accumulated in GPU registers and written to memory when all the computations are done.
Therefore, we can implement the pruning as an epilogue of the matrix multiplication: after the accumulation is ﬁnished, we compare the data stored in the registers, select the larger ones and generate the metadata.
Then, we only write the reserved non-zeros and metadata to memory.
This design brings two beneﬁts.
First, it completely removes the overhead caused by reading the matrix to be pruned from memory, so it has zero overhead.
Second, the memory footprint caused by the attention weight matrix is reduced from n 2 × 32 - bit to n 2 2 × 32 - bit + n 2 16 × 32 - bit as the original n × n full attention weight matrix is not written to memory.
The more detailed description of the CUDA kernel design including how to encode the metadata on the ﬂy is summarized in Appendix A.1.
T HEORETICAL R ESULTS
In this section, we provide more theoretical and empirical evidence that justify our D FSS A TTN as a good proxy of the full attention mechanism.
The strategy is to ﬁrst derive the theoretical value of 1) quality of the approximation with different sparse patterns 2) speedup can be achieved under certain sparsity.
Then, we compare the quality of different methods under the same speedup.
4. A TTENTION L OTTERY T ICKET
We borrow the lottery ticket hypothesis (Frankle & Carbin, 2019) and extend it to the attention mechanism.
The last step AV in the attention mechanism can be viewed as the aggregation in the graph neural network.
Following the Generalized Attention Mechanism (Zaheer et al., 2020), we describe it with a weighted directed graph G =
( A , X ) .
A is the adjacent matrix and A u,v >indicates that element x u attends to x v .
Inspired by the Graph Lottery Tickets (Chen et al., 2021a), we propose the Attention Lottery Ticket as follows.
Attention Lottery Ticket (ALT) .
Given a fully connected d graph G =
{ A , X } constructed from the full quadratic attention mechanism (Vaswani et al., 2017), the associated sub-graph can be deﬁned as G s = { m (cid:12) A , X } , where m is a binary mask.
If a G s has the performance matching or surpassing the original full quadratic attention mechanism, then we deﬁne the sparse attention mechanism with G s as an attention lottery ticket .
Zaheer et al.
(2020) have proved the existence of lottery tickets by showing 1) sparse attention mechanisms are universal approximators of sequence to sequence functions when being used as encoder 2) sparse encoder-decoder transformers are Turing Complete.
So the remaining problem is how to identify the winning tickets G s at runtime.
4. Q UALITY OF THE L OTTERY T ICKET
A popular strategy that empirically works well is selecting the top-k neighborhood in G based on the magnitude of edge weight.
We refer it as Top-k Sparsity .
Intuitively, this strategy is based on the hypothesis that the edges with larger edge weight are more important.
It has been widely adapted in existing studies (Frankle & Carbin, 2019;
Chen et al., 2021a; Ham et al., 2021; Wang et al., 2021) and demonstrated its ability to preserve model accuracy at a high sparsity ratio.
Following this trend of work, we deﬁne the Quality of Attention Lottery Ticket as follows:
Deﬁnition 4.1.
( L p -Quality of Attention Lottery Ticket)
The quality of attention lottery ticket G s = { m (cid:12) A , X } under density s = 1 n 2 (cid:80) nj =1 (cid:80) ni
=
m j,i is deﬁned as
[Equation]
The above deﬁnition computes the expectation of normalized L p norm in each row of the attention score matrix.
The p is a task-dependent factor that indicates how the accuracy depends on the edges with higher magnitude.
In this paper, we compare the L p -Quality of tickets yield by three types of sparse patterns: Top-K, ﬁxed, and our dynamic 1:2 and 2:4 sparse pattern.
Particularly, we have the proposition below:
Proposition 4.1.
Under the assumption that the entries in QK T / √ d follow i.i.d.
N ( µ, σ ) , we have
[Equation]
(Proof: Appendix A.2)
It is obvious that the Q ptopk achieves the upper bound of Q p under s .
Besides, the pσ is always positive.
Therefore, we also have Q p 2:4 ≥ Q p 1:2 > Q pfix | s =0 .
5 = 1 / 2 .
4. E FFICIENCY OF THE L OTTERY T ICKET
A lottery ticket with high quality does not necessarily mean that it is also efﬁcient to execute for wall-clock time speedup.
In this section, we analyze the efﬁciency of the three sparse patterns.
Top-K Sparsity .
Zhao et al.
(2019) explicitly select k neighbors in each row of A based on their magnitude.
However, as shown in their Table 4, the explicit sparse transformer has lower inference throughput despite k (cid:28) n .
On one hand, the top-k operator is difﬁcult to parallel and introduces high overhead.
On the other hand, even if an oracle top-k sparsity mask m were provided with zero overhead, it would still be difﬁcult for the explicit Top-K sparse attention to beat its dense counterpart.
We provide a theoretical upper bound for density s in Proposition 4.2.
Proposition 4.2.
Given embedding size d and the maximum tiling size supported by GPU T , the upper bound of the speedup achieved by Top-K Sparsity under density s is (Proof: Appendix A.3)
[Equation]
As typical values for the dimension d and tiling size T are d = 64 , T = 128 , s < 4 .
5% is a necessary and insufﬁcient condition to have Speedup > 1 .
Notably, this is not a strict upper bound as we did not take the overhead of identifying top-k entries into consideration.
Therefore, the strict upper bounder should be even smaller.
Fixed Sparsity .
As the ﬁxed sparse pattern are designed or learned before inference, they can be designed to be GPU-friendly and have the same tiling size with the dense matrix multiplication.
Therefore, we can derive the upper bound of the speedup under density s with the same strategy in
Proposition 4.2:
[Equation]
Dynamic 1:2 / 2:4 Sparsity .
Similarly, we can derive the theoretical speedup with 1:2 and 2: sparsity as follows
[Equation]
4. Q UALITY OF THE L OTTERY T ICKETS UNDER THE SAME E FFICIENCY
With the theoretical conclusions above, we compare the quality of the lottery ticket under our dynamic 1:2 and 2:4 sparsity with the other two methods under the same efﬁciency.
Comparison with Top-K Sparsity .
The Top-K sparsity achieves the same efﬁciency with ours at
[Equation]
With typical values T = 128 , d = 64 , we have s < 0 .
We can substitute it to Proposition 4. and get Q ptopk < Q p 1:2 when pσ < 7 .
On the other hand, when pσ > 7 , although the Top-K sparsity produces tickets with higher quality, Q p 1:2 | pσ =7 ≈ 0 . 9999996 is already very close to 1.
Comparison with Fixed Sparsity .
The ﬁxed sparsity achieves the same efﬁciency with ours when
[Equation]
With typical values T = 128 , d = 64 , we have s ≈ 0 . 63 .
On the other hand, we have theoretical value of σ ≈ 1 and p ≥ 1 . The p ≥ 1 is based on the observation that the edges with higher magnitude are more inﬂuential.
Therefore, we have pσ ≥ 1 and Q p 1:2 ≥ 0 .
76 > 0 .
63 =
Q pfix | s =0 .
To conclude, compared with both top-k sparsity and ﬁxed sparsity, our method can always yield lottery tickets with higher quality under the same efﬁciency.
To support this conclusion, we further provide some empirical studies in Appendix A.4.
Besides, we found both theoretically and empirically that our method is a good complementary to the kernel-based transformers like Performer (Choromanski et al., 2021).
We add more discussions about it in Appendix A.5.
E VALUATION
In this section, we ﬁrst evaluate the accuracy of our dynamic ﬁne-grained structured sparse attention mechanism on tasks across different domains.
Then, we proﬁle our methods on NVIDIA A100 GPU under different sequence lengths from 256 to 4096 to show that we can achieve practical speedup in arbitrary sequence length.
5. M ODEL A CCURACY
To show that our method is effective in comprehensive scenarios, we ﬁrst evaluate the model accuracy on tasks in different domains and sequence length.
For models under “bﬂoat16” data type, we ﬁrst ﬁnetune them from the pretrained model under “ﬂoat” data type as “ﬂoat” provides more precise gradient that helps convergence.
After the ﬁnetuning, we directly cast all the parameters in the model to “bﬂoat16” and test it on the test dataset.
For Question Answering and Masked Language
Modeling tasks, we report the results averaged over 8 runs under different random seeds.
Question Answering .
We evaluate the BERT-large on SQuAD v1.1 under sequence length 384.
We use the “bert-large-uncased-whole-word-masking” in Huggingface (Wolf et al., 2020) as the [Table] pretrained model and ﬁnetune it with the default conﬁguration in Huggingface 1 .
The F1-scores of “1:2(ﬂoat)” and “2:4(bﬂoat16)” without ﬁnetuning are obtained by directly using the checkpoints from “Transformer(ﬂoat)”.
The F1-scores of “Transformer (ﬂoat)” and “Transformer (bﬂoat16)” without ﬁnetuning are obtained by directly using the checkpoints from “DFSSATTEN 1:2 (ﬂoat)” and “DFSSATTEN 2:4 (bﬂoat16)”, respectively, and running inference with dense attention mechanism.
As shown in Table 2, with ﬁnetuning, our 1:2 sparsity has only 0.1 F1 score loss that is smaller than the standard deviation.
Our 2:4 sparsity even achieves a little bit of performance improvement over the dense baseline.
One plausible explanation is that while the 2:4 sparsity can keep most of the important edges, it also occasionally drops a small fraction of important edges which acts like the attention dropout technique (Zehui et al., 2019).
Besides, directly applying our methods to the dense transformer without ﬁnetuning also achieves comparable results, it justiﬁes that our method can well approximate the dense attention mechanism.
Masked Language Modeling .
We also evaluate our models on the masked modeling tasks on
Wikitext-2 and Wikitext-103 under sequence length 512.
Similar to the question-answering tasks, we choose the “roberta-large” as the pretrained model and ﬁnetune it under the default conﬁguration in Huggingface 2 .
The results are summarized in Table 3.
Similarly, the perplexities achieved by our methods are on par with the dense transformer.
[Table]
Long Range Arena .
For sequence length longer than 512, we incorporate four tasks from the Long
Range Arena (Tay et al., 2021), including ListOps, Text Classiﬁcation, Document Retrieval
[Table]
and
Image Classiﬁcation under sequence lengths 2048, 2048,
[Table]
and 1024, respectively.
We omit the
Pathﬁnder (1K) task as we cannot replicate the results, which was also reported in Lu et al.
For a fair comparison with other efﬁcient transformers, the model is trained from scratch under the default conﬁgurations.
The results are summarized in Table 4.
Our method achieves comparable accuracy on all the three benchmarks for long sequence.
5. S PEEDUP
In this section
[Table]
we demonstrate the speedup achieved by our method across different sequence lengths.
For a fair comparison, we only show the speedup achieved on the attention mechanism declared in equation 1.
This is because the end-to-end speedup can be affected by various other factors including the quantization and pruning strategies applied to the other parts of the transformer models, the embedding sizes, and the efﬁciency of the code implementation (e.g. some operators could be fused for lower latency). We present the end-to-end speedup and memory footprint reduction under different sequence length, number of heads, and hidden dimension in Appendix A.6. For models in previous studies, We also apply the PyTorch JIT script when possible in case that their implementations are not efﬁcient. The conﬁguration we use is as follows:
Each layer containsheads, the feature dimension per head is 64. The batch size is set to be large enough to keep the GPU busy. We summarize the proﬁling results in Figure 4.
We normalize the latency to the Trans-
[Table]
P e
r
f
o r m e r
R e f o r m e r
R o u t i n g
S i n k h
o r n
N y s t r o m
T r a n s f o r m e r O u r s P e r f
o r m e r R e f o r m e r R o u t i n g S i n k h
o r n
N y s t r o m
T r a n s f o r m e r O u r s
P e
r
f
o r m e r
R e f o r m e r
R o u t i n g
S i n k h
o r n
N y s t r o m T r a n s f o r m e r O u r s
P e
r
f
o r m e r
R e f o r m e r
R o u t i n g
S i n k h
o r n
N y s t r o m T r a n s f o r m e r O u r s
P e
r
f
o r m e r
R e f o r m e r
R o u t i n g
S i n k h o r n N y s t
r o m     bfloat Series Series Series
Series  0.
T r a n s f o r m e r O u r s P e r f
o r m e r R e f o r m e r
R o u t i n g S i n k h
o r n N y s t r o m
T r a n s f o r m e r O u r s
P e
r
f
o r m e r
R e f o r m e r
R o u t i n g
S i n k h
o r n
N y s t r o m T r a n s f o r m e r O u r s
P e
r
f
o r m e r
R e f o r m e r
R o u t i n g
S i n k h
o r n
N y s t r o m
T r a n s f o r m e r O u r s
P e
r
f
o r m e r
R e f o r m e r
R o u t i n g
S i n k h
o r n
N y s t r o m T r a n s f o r m e r O u r s
P e
r
f
o r m e r
R e f o r m e r
R o u t i n g
S i n k h
o r n
N y s t r o m
[Figure]
[Figure]
L a t e n c y n o r m a li z e d t o
T r a n s f o r m e r former with full attention mechanism under each conﬁguration.
We also cut the y axis off at 2 for clarity, because some methods designed for long sequence could be more than 20 × slower than the dense transformer at moderate sequence length.
First of all, our method achieves 1.27 ∼ 1.89x speedup over the transformer with full attention.
It is the only method that brings consistent speedup across different sequence lengths, while other methods from previous papers suffer from high overhead at moderate and short sequence lengths.
Second, under the ﬂoat data type, our method achieves speedup in all the three stages with zero overhead.
This accords with our arguments in Section 3.
Under the data type bﬂoat16, the QK T in our method is a little bit slower than the dense baseline.
The reason is that selecting 2 larger ones from 4 elements requires more comparisons, which results in more warp divergence.
C ONCLUSION AND D ISCUSSION
In this paper we present D FSS A TTEN , a dynamic ﬁne-grained structured sparse attention mechanism that dynamically prunes the QK T on the ﬂy to 1:2 and 2:4 structured sparsity.
As it only requires 50% sparsity, it can achieve no accuracy loss compared with the full attention mechanism across tasks in various domains and sequence lengths.
Besides, it only requires modifying a few lines of code, which makes it a drop-in replacement of the full-attention mechanism.
Moreover, powered by our customized CUDA kernels and the new sparse tensor core on Ampere GPU, we achieve 1.27 ∼ 1.89 × speedup over the full attention in arbitrary sequence length.
All of these pieces of evidence demonstrate that our method can be a good replacement for the full attention mechanism.
Our method is also orthogonal to many existing efﬁcient attention mechanisms and can potentially be applied jointly for further speedup.
We present two examples in Appendix A.
A A PPENDIX A. K ERNEL D ESIGN D ETAILS
In this section, we ﬁrst demonstrate how a dense matrix is pruned and compressed under the 50% structured sparsity on Ampere GPU in Section A.1.1.
Then, we detail the design and implementation of the SDDMM, Softmax, and SpMM kernels in Section A.1.2 and A.1.3.
A.1. S TRUCTURED P RUNING OF D ENSE M ATRIX
We ﬁrst illustrate how to dynamically prune a dense matrix with 50% structured sparsity.
Under data type ﬂoat, we select the larger one in every two consecutive elements.
If the data type is bﬂoat16, we select two larger ones in every four consecutive elements.
We compress the pruned dense matrix to nonzeros and metadata following CUTLASS et al (2021) as there are two beneﬁts.
First, it can be directly used by high-performance SpMM kernels in CUTLASS.
Second, as we will show in
Section A.1.2, it can be dynamically generated from the SDDMM kernel with neglectable overhead.
As shown in Figure 5, the basic tile size to prune is 32 × 64 -byte, this corresponds to a 32 ×block under bﬂoat16 or 32 × 16 block under ﬂoat.
There are four major steps:
Prune 50% of each consecutive 8B data, generate nonzeros and metadata; 1 Interleave the metadata rows by 8; 2 Switch the metadata along sub-diagonal.
3 Write metadata and nonzeros to global memory.
In detail, 2 out of 4 2-byte data are select based on their magnitude and a unique 4-bit metadata is assigned to each combination in 0 .
The correspondence between selection pattern and metadata is enumerated in Figure 5 (b).
Notably, with ﬂoat32 data type, each 32-bit data occupies two consecutive 2-byte slots.
Therefore, it only supports the patterns under 0x4 and 0xe .
After generating the 4-bit metadata, consecutive four of them are concatenated to a 2B metadata block.
Then, the rows of metadata are interleaved by 8 in 1 following
[Equation]
In 2 , the metadata blocks at upper right and lower left of each 2 × 2 grid are switched.
At last, in 3 , the metadata produced by 2 is written into global memory following the interleaved column-
[Figure]
Thread Block Tile
A
B
C
A
B
M
N
K K
Mtile
Ntile
Warp0 Warp
Warp2 Warp
Warp Tile 16float32bfloat   major format under stride 4-byte.
This can be realized by interpreting two consecutive metadata as an int object and then write it to DRAM in column-major.
The nonzeros are simply writen to global memory under row-major.
A.1.
[Equation]
Our strategy for dynamically pruning the attention score matrix has two steps.
First, perform a conventional dense GEMM.
Second, prune the GEMM output with procedures described in Section
A.1.1.
However, if the second step is implemented as a separate GPU kernel, we need to write the dense attention score matrix to DRAM and read it back.
This not only introduces high overhead, but also prevents us from reducing global memory footprint.
To address this issue, we implement the pruning step as an epilogue attached to the conventional GEMM kernel: the results of the dense
GEMM are stored in the registers, the epilogue processes the results and then writes nonzeros and metadata to global memory.
Dense GEMM .
The GEMM step is no different from conventional GEMM kernels, and all the existing optimizations can be used.
The tiling is shown in Figure 6: each thread block processes a Mtile × Ntiles output tile, which is further partitioned to several warp tiles.
Each warp tile is composed of a grid of 16 × 16 blocks that matches the tensor core output size.
In each thread block, all the threads jointly load Mtile × Ktile and Ktile × Ntile input tiles from matrix A and B into the shared memory.
We use the new synchronize copy feature on Ampere architecture to fully utilize the memory bandwidth and reduce register usage.
To fully annihilate shared memory bank conﬂict, we use the XOR layout.
Once the load is completed, the warps fetch their source operands from shared memory with ldmatrix and perform a (16 × 32 B ) · (32 B × 16) warp matrix multiply accumulate ( wmma ) with tensor core.
Notably, ﬂoat data will be converted to tensorﬂoat-32 before wmma .
To reduce accumulation error, we accumulate the partial sum as ﬂoat regardless of the source operand data type.
Besides, software 2-stage pipeline is used to overlap memory access and computation with double buffering (et al, 2021).
Although deeper software pipeline can be built on Ampere, we ﬁnd 2 stages is enough as the inner-product dimension K is usually very small (e.g. 64).
More detailed explanation of the above techniques can be found in this GTC 2020 talk (Kerr, 2020).
Pruning the GEMM result .
In the pruning step, the warp tile is partitioned to a grid of 32 × 64 B blocks that are processed by the warp one at a time.
Under data type ﬂoat, the register layout of the 32 × 16 block is illustrated in Figure 7 (a).
It consists of two 16 × 16 wmma blocks, so each thread has sixteen 32-bit registers to hold the results.
The registers are annotated with “T thread id { register id } ”.
As the adjacent two data are held by the same thread, we can simply compare them and the larger one is retained.
Under data type bﬂoat16, we need to select 2 larger ones from adjacent 4 entries.
However, under the naive mapping shown in Figure 8 (a), these 4 entries are held by 2 thread.
Therefore, we need additional warp shufﬂe to ﬁrst pass these 4 entries to the same thread, then compare them and obtain the 2 larger ones.
This will introduce additional overhead.
To solve this problem, we propose to interleave the columns when loading matrix B to shared memory by simply manipulating the pointer to the global memory at the beginning.
The resulted mapping to the registers is shown in
Figure 8 (b) which is equivalent with Figure 7 (a) bﬂoat16.
After the interleaving, consecutive four
[Table] [Table]
[Table] [Table] [Table]
[Table] [Table] [Table]
[Table] [Table]
[Table] [Table]
[Table] [Table]
[Table]
R\C
[Table] [Table]
[Table] [Table]
[Table] [Table]
[Table] [Table]
[Table] [Table]
[Table] [Table]
[Table] [Table]
[Table] [Table]
[Table] [Table]
[Table] [Table]
R\
[Table]
[Table] [Table]
[Table] [Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table] [Figure] [Table] [Figure]
[Table]
[Figure] [Table] [Figure] [Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure] [Table] [Figure] [Table]
[Figure]
[Table] [Figure] [Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
float
[Figure]
data are naturally held by the same thread, and we select 2 larger ones from them.
To reduce branch divergence, the selection is done by comparing the sum of any two data.
Generate Metadata and Nonzeros .
For both ﬂoat and bﬂoat16 data type, each comparison produces a 4-bit metadata.
Next, following the procedures described in Section A.1.1, we need to concatenate consecutive four metadata to a 16-bit metadata block.
This is done in two steps.
First, put the 4-bit metadata to the correct position of a int16 register with bit shift.
Second, share these int16 registers cross threads with warp shufﬂe, and concatenate them with bitwise OR.
As consecutive four metadata are held by thread 4t to 4t+3 , we put the 4-bit metadata of thread 4t+k to [ k × 4 : k × 4 + 3 ] bits in the int16 object in the ﬁrst step.
The detailed layout is shown in Figure 8 (b), where we denote each 4-bit metadata as “T thread id { register id } [ bit id ]”.
The result of the second step is shown in Figure 7 (c).
Figure 7 (d) and (e) illustrate the result after 1 and 2 in Figure 5.
Notably, these two step only change the logic mapping of the metadata and the register allocation is not affected.
So no code is required for these two steps.
At last, we need to write the metadata and nonzeros to global memory following 3 in Figure 5.
As shown in Figure 7 (e), each row is held by consecutive two int16 registers of the same thread, so we can simply reinterpret it as an =int object and write the metadata to global memory in column major.
For the nonzeros, we simply coalesce them in the shared memory and then write to global memory in row-major.
Batched Kernel .
The self-attention layer in transformer usually has multiple independent attention heads.
Instead of launching one CUDA kernel for each attention head, using a batched kernel that
[Figure] [Figure]
C
[Figure] [Figure]
for bﬂoat16.
processes all the heads can better utilize the GPU resources and reduce kernel launching overhead.
We support the batched computation by using the blockIdx.z to index the heads in the batch and update the pointers to the input and output based on the index.
Blocked-ELL Sparsity .
Under long sequence length, higher sparsity is desired to reduce computation cost and memory footprint.
Our kernel support hybrid blocked-ELL sparsity Zaheer et al.
(2020) and 50% structured sparsity.
To support this feature, we set the block size in blocked-ELL to the thread block tile size of the GEMM.
Therefore, we can simply skip those pruned blocks during the execution.
A.1. S OFTMAX AND S P MM K ERNEL
In this section, we detail the implementation of the softmax and SpMM kernels.
Softmax Kernel .
To improve numerical stability, the softmax on GPU is computed with
[Equation]
Therefore, each element in x has to be loaded for three times. 1) compute c = max ( x ) ; 2) compute s = (cid:80)
j e x j − c ; 3) compute
e x i − c /s .
Instead of loading x
i from global memory in each time, we cache it in the register when the whole row ﬁts in the register ﬁle capacity.
Besides, the ordinary softmax kernel in libraries like PyTorch can also be used.
SpMM Kernel .
As we encode the nonzeros and metadata following the CUTLASSet al (2021), we directly construct the SpMM kernels from the CUTLASS APIs.
To support the hybrid blocked-ELL and structured 50% sparsity, we modify the PredictedTileAccessIterator class in CUTLASS to skip the tiles masked out by the blocked-ELL sparsity.
A. P ROOF OF P ROPOSITION 4.
Proof.
Under the assumption that the entries in QK T / √ d follow i.i.d.
N ( µ, σ ) , we denote x i,j = e µ + σz i,j , where z ∼ i.i.d.
N (0 , 1) .
Then we can substitute it into the deﬁnition of the softmax and
[Equation]
We substitute the above equation into the deﬁnition of L p -Quality and get
[Equation]
With n → ∞ , the denominator can be approximated with
[Equation]
Top-K Sparsity .
When the sequence is long enough such that we have n → ∞
[Equation]
the numerator can be approximated with
[Equation]
Therefore, the L p -Quality of Top-K sparsity is
[Equation]
Fixed Sparsity .
Without any assumption on the distribution of important edges in A , applying a ﬁxed pattern is equivalent with uniformly sampling with probability s and we have
[Equation]
Therefore
[Equation]
the L p -Quality of the ﬁxed sparsity is Q
[Equation]
[Equation]
2-to-1 Sparsity : This sparsity pattern select the larger one in every two elements.
We denote adjacent two elements with
[Equation]
Z 1 and Z 2 are independent.
Then we have
[Equation]
We denote
[Equation]
then we have
[Equation]
With the conclusion above
[Equation]
we have
[Equation]
The L P -Quality of 1:2 sparsity can be computed with
[Equation]
2:4 Sparsity : This sparsity pattern select the largest two elements in consecutive four elements.
While it is more challenging to ﬁnd an explicit expression for Q p 4 − to −
[Equation]
a trivial lower bound can be found with
[Equation] [Equation]
where we have
[Equation]
Z 1 , ..., Z 4 are independent.
Therefore, the lower-bound of Q p 2:4 is
[Equation]
A. P ROOF OF P ROPOSITION 4.
Proof.
First of all, thanks to the Tensor Core in latest GPUs
[Equation]
the latency of matrix multiplication operations, both sparse and dense, are bounded by the memory access.
Therefore, instead of counting the number of MACs (multiply-accumulate operations), the amount of memory access is a better metric to estimate the latency.
[Equation]
[Figure]
Tiling is a basic optimization applied to optimize matrix matrix multiply on GPU.
As shown in
Figure 9 (A), the original n × n output is partitioned to independent blocks with size T × T .
When computing each block, operands with size T × r and r × T are loaded from A and V T to the fast memory, respectively.
Then, these two operand are multiplied and accumulated to the partial sum stored in the registers.
After applying the top-k, as shown in Figure 9 (B), the k elements in each row of A correspond to different rows in A .
Therefore, we can only partition the output to independent S p ee dup O v e r F u ll A tt e n t i o n
Density (s)
Top-K
Top-K Theory Fixed Fixed Theory 1: 1:2 Theory vectors with size 1 × T . During the computation, operands with size 1 × r and r × T are loaded from A and V T to the fast memory, respectively.
Then, the loaded operands are multiplied and accumulated to the partial sum stored in the registers.
With the tiling strategy mentioned above, we can summarize the amount of memory access in different attentions in the table below.
[Table]
For QK T , as we need to compute all of it before getting the top-k elements, it is a dense matrix matrix multiplication for both full and explicit top-k attention.
The Softmax needs to read the n × n
QK T in, normalizes it, and write the result A back.
As the intermediate values can be stored in registers, we only need to count reading QK T in and writing A out.
Therefore, its memory access is 2 n 2 for full attention and 2 n 2 s for explicit top-k attention.
For AV in full attention, the output size is nd .
As each output element is generated from the inner product between two vectors with length n , the total data read equals nd × 2 n .
However, with the tiling in Figure 9 (A), each operand is reused for T times.
Therefore, the total memory access for AV in full attention is nd ( 2 nT + 1) .
For AV in explicit top-k attention, as shown in Figure 9 (B), each left-hand-side data is reused for
T times while each right-hand-side data is used only for once.
Therefore, its memory access equals to nd ( snT + sn + 1) .
The theoretical speedup can be computed with
[Equation]
A. Q UALITY OF THE L OTTERY T ICKETS UNDER THE SAME E FFICIENCY
In this section, we provide more empirical evidences to support our conclusions.
We ﬁrst compare the theoretical speedup of different sparsity predicted in equation 4, equation 5, and equation 6 and the actual speedup measured on A100 GPU in Figure 10. [Figure] First of all, the Top-K sparsity is well bounded by the theoretical value, and our method achieves better speedup than the Top-K sparsity when the density s > 0 . This is because gathering top-k elements in each row of the attention weight matrix and sorting them to compressed row format introduce huge overhead.
Second, the speedup achieved by the ﬁxed sparsity is well predicted by our theoretical value.
The speedup it achieved is lower than ours when density s ≥ 0 .
63 , which accords with our theoretical conclusion.
Notably, the speedup of ﬁxed sparsity we used here is simply truncate the number of columns of the attention weight matrix based on the density.
The actual speedup will be even lower when more ﬁne-grained pattern is involved.
Our method delivers speedup a little bit higher than the theoretical value.
This is because the softmax kernel has different implementations under different sequence length.
When the sequence length is moderate, as mentioned in Appendix A.3, the data loaded from the attention score matrix can be explicitly cached in fast memory like registers or shared memory for reuse.
When sequence length too long for the fast memory to cache, it has to be implicitly reused through lower-level cache or even global memory. The second implementation is slower than the ﬁrst one as lower-level cache has longer access latency and lower throughput. As our method reduces the sequence length by half, it can use the implementation for moderate sequence length while the full attention is handled by the long sequence version. In Figure 11, we compute the theoretical value (solid line) and empirical value (box plot) of Q p over attention matrix A in BERT-Large on SQuAD v1.1.
As p is a task-dependent value that is hard to obtain, we instead sweep through several typical values.
Compared with the top-k sparsity, when p < 7 , our 1:2 and 2:4 sparsity always achieve better performance than the top-k sparsity when s < 0 .
Besides, when p = 7 , the Q p 1:2 and Q p 2:4 are [Figure] a [Figure]
very close to 1. These observations accord with our conclusion that our 1:2 and 2:4 sparsity can obtain tickets with better quality than Top-K sparsity at the same efﬁciency.
Compared with the ﬁxed sparsity, our Q p 1:2 and Q p 2:4 are also similar or better than Q pfix across different p s.
This supports our conclusion that our method achieves better performance than the ﬁxed sparsity patterns under the same efﬁciency.
To show that our Q p is a good metric to compare the performance of different sparse patterns, we plot the Q p and F1 score on BERT-large SQuAD v1.1 in Figure 12.
As we mentioned before, p is a task-speciﬁc value used to model tasks with different degree of dependency on the largest few elements. In order to identify the p for our target task, we tune the value of p until the data points from Top-K sparsity and Fixed sparsity form a monotonically increasing line. We found that p = 6 . 5 is a good choice. This large p accords our observation that the Top-K sparsity works well even under 5.4% density.
After anchored the p , we put the data points from 1:2 and 2:4 sparsity into the plot and verify if the line is still monotonically increasing.
Figure 12 shows that the data points from our 1:2/2:4 sparsity perfectly ﬁlls in the monotonically increasing line.
Oppositely, The traditional F-norm based metric cannot explain why the 1:2 sparsity has better F1-score than some
Fixed Sparsity even though it has lower score.
This demonstrates that our Q p is a better metric than existing metrices.
In this section, we add more discussions on how our method compared with kernel based transformer, i.e. Performer (Choromanski et al., 2021).
As our Deﬁnition 4.1 is designed to characterize how well the sparse pattern could reserve the important edges in A , so it is not suitable for kernelbased attention mechanisms that do not involve sparsity.
For example, an approximation of A with high positive approximation error can have Q P ≥ 1 under Deﬁnition 4.1.
Therefore, we instead compare the mean squared error (MSE) following Choromanski et al.
Given the query and two adjacent key vectors q , k , and k (cid:48) ∈ N (0 , I d ) , we denote the softmax kernel between them as SM ( q , k ) =
exp ( q T k / √ d ) .
And the softmax approximated by our dynamic 1:2 sparsity (cid:92) SM 1:2 ( q , k ) is deﬁned as
[Equation]
else
[Equation]
Then, we can compute its MSE as follows
[Equation]
Because q T k (cid:48)
= (cid:80) di
=1 q i k (cid:48)
i is the weighted sum of i.i.d variables following N (0 , 1) , we have x = q T k (cid:48) ∼ N (0 ,
||
q || 22 ) .
We can substitute it into equation 29 and get
[Equation]
With Lemma 2 and Theorem 2 in Choromanski et al.
[Equation]
the MSE of their positive softmax kernel with orthogonal random features has an upper bound as follows
[Equation]
First of all, when SM ( q , k ) → 0 , both MSE ( (cid:92) SM 1:2 ( q , k )) and MSE (cid:16) (cid:92) SM ort + m ( q , k ))
(cid:17) converge to 0.
However, for large SM ( q , k ) s that are potentially be critical for the model accuracy, the exp (cid:16)
||
q ||
+ ||
k ||
√ d (cid:17) SM 2 ( q , k ) term in the positive softmax kernel in Performer could greatly increases the MSE.
Oppositely, the 1 − erf (cid:16) √ d || q || 2 √ 2 ln
( SM ( q , k )) (cid:17) term in our method reduces the MSE.
To conclude, while both the positive softmax kernel and ours has low MSE error when approximating small edge weights, our method can better approximate the edges with high magnitude.
From the empirical perspective, as shown in Table 2 and 3, our method can achieve good accuracy even without ﬁnetuning.
Whereas the Performer still requires tens of thousands steps of ﬁnetuning (e.g. Figure 5 in Choromanski et al.
Table 4 also reveals that Performer has poor accuracy on certain tasks like byte-level document retrieval, while ours consistently achieve accuracy on par with the dense transformer.
All this observations suggest that our method can better approximate the full attention mechanism than Performer.
In terms of wall-clock time speedup, Figure 4 illustrates that the Performer can only achieve good speedup at long sequence length.
The similar phenomenon is also observed in multiple online forums 3 .
Certainly, the PyTorch JIT script does not yield the optimal implementation of the computation graph, but it reveals that tremendous engineering efforts are required for Performer to achieve good speedup under moderate sequence length.
Following Section 4.3, we also compare the theoretical speedup achieved by ours and the Performer.
[Equation]
The computation steps of Performer are listed in equation 32 where each equation denotes a sub computation graph that can potentially be fused.
Notably, this is more complex than the original [Figure] dtype=bfloat16, #head=4, hidden=     0. dtype=bfloat16, #head=4, hidden= 0.
1.0 Sequence Length 0.
E n d - t o - E n d Sp ee
d u p o v e r D e n s e T r a n s f o r m e r mathematical expression to handle the numerical instability of exp .
The total memory access can be computed with
Speedup = {
[Equation]
We have m = dln ( d ) following Theorem 4 in Choromanski et al.
We can substitute m = 266 , d = 64 , and T = 128 into equation 33 and get Speedup > 1 when n > 672 .
On the other hand, the performer achieves the same speedup with ours with n > 1002 .
To conclude, our method is a good complementary to performer.
With delicately optimized computation graph, performer can achieve good speedup and relatively good accuracy under long sequence scenario.
In contrary, our method has better speedup and accuracy under moderate and short sequence length.
Besides, our method delivers lower approximation error on important edges so it is more friendly to ﬁnetuning.
In this section, we present the end-to-end speedup achieved by our method under different conﬁgurations.
We use the 4-layer dense transformer model of Text Classiﬁcation task in Long Range Arena Tay et al.
The dimension of each head is 64.
We explore different combination of number of heads (4, 8), sequence length (512, 1024, 2048, 4096), and hidden dimension of the feed forward layer (256, 512, 1024). The end-to-end speedup over the dense transformer under different conﬁgurations are plotted in Figure 13. Our method achieves 1 . 11 ∼ 1 .
52 × and 1 .
08 ∼ 1 .
47 × end-to-end speedup over the dense transformer, it is the only method that deliver end-to-end speedup under all conﬁgurations.
Under sequence length ≤ 2048 , our method achieves higher speedup than most of the baselines.
Although
Sinkhorn transformer (Tay et al., 2020a) has higher speedup than ours at sequence length 2048, as shown in Table 4, its accuracy is less satisfying.
This result justiﬁes that our method delivers good speedup under short and moderate sequence length.
Notably, this speedup is almost a free lunch
[Table]
On one hand
[Table]
Section 5 demonstrates that our method achieves comparable accuracy across different tasks and sequence length
[Table]
so the model accuracy is not sacriﬁced.
On the other hand, unlike previous efﬁcient transformers
[Table]
our method has no hyper-parameters and only requires lightweight ﬁnetuning process.
[Table]
D e n s e O u r s D e n s e O u r s D e n s e O u r s D e n s e O u r s D e n s e O u r s D e n s e O u r s D e n s e O u r s D e n s e O u r s D e n s e O u r s
L a t e n c y N o r m a li z e d t o D e n s e # Head =
# Head =
Hidden
Seq Len
[Figure]
To study how our method contributes to the end-to-end speedup, we further break down the end-toend inference time to the attention mechanism and other components under bﬂoat16.
The results are illustrated in Figure 14.
Under moderate and short sequence length like 1024 and 512, the “Others” contributes over 70% of the total latency.
This is because the size of the matrix multiplications in the feed-forward network and query/key/value projection are comparable with the attention mechanism.
However, unlike the attention mechanism that has limited time budget for compression, the feedforward network and query/key/value projection use a static weight matrix during inference, so they can be compressed ofﬂine.
Tons of methods have been proposed in the literature to do that even before the transformers are proposed.
For instance, Mishra et al.
; Zhou et al.
(2020) show that pruning the weights to 2:4 sparsity can deliver 1 .
3 ∼ 1 . 6 × speedup and 2 × fewer parameters in the feed-forward and projection layers 4 without accuracy loss on BERT-large.
Lagunas et al.
(2021) apply structured pruning and achieve 2 .
4 × speedup on SQuAD v1.1 with 1% drop of F1.
The
MobileBERT (Sun et al., 2020), on the other hand, redesign the network architecture that reduce the hidden dimension of feed-forward network in BERT from 4096 to 512.
In terms of quantization, previous work (Zafrir et al., 2019) have shown that the linear layers can be quantized to 8 bit integer.
Besides the linear layers, there are also techniques to accelerate other components in transformers.
For instance, the MobileBERT (Sun et al., 2020) replaces the layer normalization to a simple element-wise linear transformation.
The input embedding table is also compressed with smaller embedding dimension along with an 1D convolution.
[Figure]
P e a k M e m o r y
A ll o c a t i o n N o r m a li
z e d t o D e n s e
T r a n s f o r m e r
With all these techniques in the literature, it should not be hard to achieve 2 × speedup in the nonattention part of transformer models.
Then our method could deliver 1 .
13 ∼ 1 .
41 × speedup under sequence length ≤ 1024 .
We also measure the peak memory allocation of different models and conﬁgurations, the results are summarized in Figure 15.
Our method achieves 1 .
41 ∼ 1 .
82 × memory reduction, which is comparable with or better than most existing efﬁcient transformers when sequence length ≤ 1024 .
A. C OMBINATION WITH THE E XISTING E FFICIENT T RANSFORMERS
Existing efﬁcient transformers usually sparsify the full attention mechanism to densely connected clusters (Tay et al., 2020a; Roy et al., 2021; Kitaev et al., 2020; Zaheer et al., 2020) or approximate it with low-rank projection (Wang et al., 2020).
As our method is a good approximation of the full attention mechanism and brings wall time speedup at arbtrary sequence length, it can potentially be combined with the existing efﬁcient transformers. We ﬁrst demonstrate the combination of our method with Xiong et al. Xiong et al.
(2021) propose a Nystrom-based self-attention mechanism that approximate standard self-attention with
O ( n ) complexity. The Nystromformer is illustrated in Figure 16. We observe that the computation circled in Figure 16 is identical to the standard attention mechanism, so it can be further accelerated with our method. More importantly, the two matrix multipliation involved are the two of the three largest m × n matrices.
It will be very beneﬁcial to reduce their complexity.
We report the accuracy on Image (1K) on LRA (Tay et al., 2021) in Table 6.
We ﬁrst pretrain a standard Nystromformer from the scratch for 35,000 iterations following Xiong et al.
Then, we ﬁnetune it for 3,500 iterations (1/10 of the training process) under standard Nystromformer,
Nystromformer + DFSSATTEN 1:2, and Nystromformer + DFSSATTEN 2:4.
It is obvious that by combining DFSSATTEN and Nystromformer, we can achieve higher accuracy on LRA with lightweight ﬁnetuning.
𝑋 : 𝑛×𝑑 𝐾 ! :
𝑑 " ×𝑛 𝐾’ ! :𝑑 " ×𝑚 𝑄 : 𝑛×𝑑 " 𝑉 : 𝑛×𝑑 # 𝑄+ : m ×𝑑 " × 𝑚×𝑚 n ×𝑚 m ×𝑛 n ×𝑚 𝑚×𝑑 # n ×𝑑 # 𝑂 ∶ n×𝑑 # n×𝑑 # DConvk× sMEANS sMEANS 𝑚×𝑚
pINV × × × × × +
SDDMM
SDDMM
SpMM
SpMM
[Table]
Then we provide a complexity analysis of the combination following Xiong et al.
The landmark selection with segement-means takes O ( n ) , iterative approximation of the pseudoinverse takes O ( m 3 ) .
The matrix multiplication complexity of the standard Nystromformer takes O ( nm 2 + mnd v + m 3 + nmd v ) .
{+After applying our method, it can be reduced to O ( nm 2 2 + nmd v 2 + m 
+}
+ nmd v ) .
The memory footprint can be reduced from O ( md q + nm + m 2 + nm + nd v ) to O ( md q + nm + m 2 + nd v ) .
Given n (cid:29) m > d v ≈ d p , this could be a signiﬁcant improvement that allows us to use more landmarks m to better approximate the full attention mechanism.
Besides Nystromformer, we also illustrate two possible combinations with BigBird (Zaheer et al., 2020) and Linformer Wang et al.
(2020) that can be explored in the future work.
As shown in Figure 17 (A), Zaheer et al.
(2020) use block sparsity with block size 64 and compute a full attention within each block.
We can apply the 1:2 or 2:4 sparsity within each block to bring further speedup.
Figure 17 (B) gives another example on how to combine our method with Linformer (Wang et al., 2020).
Linformer uses low-rank approximation on the attention mechanism as follows:
[Equation]
where E , F ∈ R n × k are linear projection matrices and k (cid:28) n .
We can ﬁrst prune E and F along with other weight matrices to have 1:2 or 2:4 sparsity ofﬂine following Mishra et al.
Then we compute EK and F V with Sparse Matrix-Matrix multiplication.
Next, we multiply Q and ( EK ) T and the result is pruned to 50% structured ﬁne-grained sparsity on the ﬂy.
After applying softmax to the nonzeros, we multiply it with F V .
A. V ISUALIZE A TTENTION D ISTRIBUTION
To illustrate that our DFSSATTEN can well capture the ﬁne-grained sparsity in attention, we visualize the attention weight matrices in BERT-large on SQuAD v1.1 in Figure 18. In detail, we 𝑄 𝑉 𝑌 𝑘 𝑘 𝑛 𝑑 𝑛 𝑑 𝑄𝐾 ! 𝐸 !
𝜎 𝑂 (B) Linformer: 𝐸 !
𝐾 !
𝑛 𝑑 𝐾 ! 𝐸 !
𝑘 𝑛/ 𝑛/ 𝜎(𝑄𝐾 !
𝐸 ! )
𝐹𝑉 𝐹 𝑘 𝑛/2 𝑛/
[Figure]
run inference of the same input sample in BERT-large model pretrained under dense, 1:2, and 2: settings, then collect the attention weight matrix in the ﬁrst layer.
It is obvious that the pattern in dense transformer and our DFSSATTEN are quite similar.
The magnitude of nonzero values in DFSSATTEN are a little bit higher than dense attention.
This is because the softmax normalizes the values in each row with the exponential sum of each entry.
After removing 50% smaller entries, the magnitude of remaining entries would be relatively higher.
Nevertheless, we ﬁnd that this does not inﬂuence the model accuracy, as the forthcoming normalization layers will take care of it.
[Figure]