DFSSATTEN:  Dynamic Fine-grained Structured Sparse Attention Mechanism
Authors: Zhaodong Chen
Liu Liu
Yuying Quan
Zheng Qu
Yufei Ding
Yuan Xie
Venue: ICLR 2022 Submitted (2022)

Abstract: 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 DFSSATTEN, 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 evidences that demonstrate DFSSAT- TEN 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. DFSSATTEN achieves 1.27 ∼ 1.89× speedups over the full-attention mechanism with no accuracy loss.

Under review as a conference paper at ICLR 2022
DFSSATTEN: DYNAMIC FINE-GRAINED STRUCTURED SPARSE ATTENTION MECHANISM
Anonymous authors Paper under double-blind review
ABSTRACT
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- efforts have been made 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 DFSSATTEN, attention mechanism that dynamically prunes the full attention weight matrix to the 50% fine-grained structured sparse pattern used by the sparse tensor core NVIDIA A100 GPU. We provide both theoretical and empirical evidence that demonstrate DFSS ATTEN is 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 4096. DFSSATTER achieves 1.27 ~ 1.89 x speedups over the full-attention mechanism with no accuracy loss on A100 GPU.
tasks. Tremendous
INTRODUCTION
Transformers Vaswani et al. 2017) have achieved competitive performance across various domains like NI.P OOt et al. 2018) and Computer Vision (Dosovitskiy al.. (021). The key feature that sets them apart from traditional neural network architectures is the attention mechanism (Vaswani et al.) 2017), which allows the transformers 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 difficult to apply transformers to latency-sensitive tasks. Many efficient attention mecha- nisms(Tay et al. 2020b] Zaheer et al. 2020 Beltagy al., 2020 Tay et al., 2020a Royy 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 fixed sparse patterns extremely high sparsity to achieve wall-clock time speedup. Therefore, these methods usually require thousands of pretraining or fine-tuning steps on specific tasks and toilsome tuning of several hyper-parameters reach good accuracy. Last but not least, previous methods usually introduce additional operators like top-k, sort that cause large overheads and offset their benefits moderate sequence length.
In this paper, we present ATTEN, a simple and effective sparse attention mechanism that ad- dress the limitations mentioned above. DrSSATTEN 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 ATSSATTEN offers several advantages over existing studies:
It requires minimal changes to the original full-attention with no makes it drop-in replacement of the full attention that only requires change few lines of code. Moreover, it can directly exploit existing pretrained models like BERT (Devlin et al. 2019) and RoBERTa (Liu et al. 2020).
to tune. This

---

Under review as a conference paper at ICLR 2022
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 fine-tuning.
Our method introduces achieve wall-clock time speedup and memory footprint reduction over the full attention in arbitrary sequence length.
overhead on existing GPU hardware. As result,
able to
To conclude, our main contributions summarized below:
We propose DESSATTEN, a dynamic sparse attention mechanism that is drop-in proxy of the full attention mechanism. Its effectiveness justified by both empirical and theoretical evidence.
We present dedicated CUDA kernel design to completely remove the pruning overhead. The pruning implemented as an epilogue of the dense matrix multiplication which produces the attention score matrix.
We evaluate DESSATTEN on tasks cross various domains and sequence lengths. It achiieess 1.27 1. 1.89x speedup over the full attention with no accuracy loss.
2
BACKGROUND AND MOTIVATION
We first introduce the preliminaries, notations, and background of our paper.
2.1
FULL ATTENTION MECHANISM
Given input sequence X = (x) E Raxd the full attention mechanism can be defined
O = Softmar /Vavv
(1)
where Q XW, K= XW and V XW, are query, key, and value matrices. Q QKT forms qull-quadratic adjacency matrix, whose edge weights are the dot-product similarity between all the elements in sequence. This adjacency matrix is standardized with 1/Vd 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 /k//val /Vd) for simplicity. refer QK" as the attention score matrix and A as the attention weight matrix.
of this paper, we denote A = So Stffmaz(( (QK")
2.2 EFFICIENT ATTENTION MECUANISM
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 efficient attention mechanisms have been proposed Tay al. 2020b).
Fixed Sparse Patterns. Zaheer et al. 2020 Beltagy et al. (020) apply a set of fixed 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 fixed 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 al. (2021) dynamically generate fine-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 al. (220aa); Roy et al. (2021): Kitaev al. (2020) apply various clustering methods and only compute the at- tention within each cluster. Although computing full attention in each cluster is more friendly to GPU compared with fineinrained sparsity, the clustering methods contain several GPU-unfriendly operators like top-k and sorting that offsets their benefits under moderate sequence length.
Low Rank / Kernel. Wang et al. (2020) project A from 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
to k with linear projection.
2

---

Under review as a conference paper at ICLR 2022
Sparse Tensor Core
Selec
2:4
Dol-product
Fine-grained Structured Pruning
Compress
1:2
1:2/2:4 sparsity
110
Nonzeros
Figure A100 GPU AnneGGrinned Structured Sparsity Pruning. (NVIDIA) 2020
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 pre- training or finetuning steps are required to reach a comparable performance with the origin full attention mechanism. So they require tremendous engineering effort to deploy.
2.3 FINE-GRAINED STRUCTURED SPARSITY IN NVIDIA A100 GPU
NVIDIA introduces the fine- ine-grained structured sparsity in the A100 GPU. As shown in Figure T the dense input matrix pruned with fine-grin greeggaaied structured pruning. If the data type is float, 1:2 sparsity is used which selects the larger one in two consecutive entries. If the data type is bfloat16 or 6, the 2:4 sparsity used which selects two larger ones among four consecutive elements. After the pruning, the result is compressed 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. takes bit metadata record the decision of each 1:2 or 2:4 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 dense matrix under the support of the sparse tensor core to achieve significant speedup.
This fine-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 feed- forward part of the transformer up to 1.9x. 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, must also generate the metadata encoded in a special format such that the metadata can be used efficiently later. All these overheads will offset the benefit brought by the pruning if we do on the fly.
DYNAMIC FINE-GRAINED STRUCTURED SPARSE ATTENTION MECHANISM
In this section, we first give an overview of our DESSATTEN 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 briefly introduce our GPU kernel design to remove pruning overhead.
Our proposed DESSATTEN mechanism is simple and effective, as illustrated Figure 2 Com- pared with the full-quadratic attention mechanism, method dynamically prunes attention scores without incurring storage or computation overhead, while maintaining the effectiveness of atten- tion. More importantly, our method can achieve practical speedups of attention on existing GPU hardware with customized CUDA kernels. Listing shows all the modifications to be made to use ATSSATTEN.
3

---

Under review as a conference paper at ICLR 2022
(A) Full Attention
(B) DFSSATTEN
7
M
V
K
71
n/2 n/16
o@kk
0
Figure 2: Overview of our Dynamic Fine-grained Structured Sparse Attention Mechanism, repre- is omitted for simplicity.
sents softmax, va
Full attention mechanism def full_attentiann (q,k, v)
attn_weight corsh.bmm (q, k. tar 2))
attn_weight. torch. funne cal.ccaall............ softmax (attn_wei ghtt, -1) return terch. bam (attn weight, v)
DFSS attention mechanism import dspattn
def dfss_attention kq,k, v) v)
attn_weight, metadata dspattn. bamn (q, k)
attn_weiid Eetn_weight/ = torch. nn. functicnal, softmaxx return dspattn. spmm aattn__weight metadata, Listing 1: Example of using DESSATTEN. The "dspattn" is the package developed.
-1)
3.1 DESIGN CONSIDERATIONS FOR EXPLOITING ATTENTION SPARSITY
As illustrated in Figure 3 the attention mechanism can be considered three stages: QKT, Softmar, and AV. To design sparse atten- tion mechanism, the first decision to make is where should we induce pruning sparsify the attention.
OKT
Softmax
AV
Figure Attention Stages.
If we start from all the three stages will he benefited from the sparsity given effective imple- mentation 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 identified by the sparse pattern. The Softmaz only operates on the nonzero values in each row. The original dense matrix multiplication between A and V will be replaced with sparse matrix-matri multipli- cation (SpMM) which multiplies a sparse matrix with a dense matrix. However, as is not possible to exactly know which entry in QK² has higher magnitude before computing QKF, starting from usually requires some additional components to predict the location of important entries.
Starting from requires us to compute a dense matrix multiplication between Q and K. The benefit that we can explicitly select important entries from QK" without prediction. As the softmax is monotonically increasing function, starting from does not offer any benefits over but throws away the opportunity to accelerate Softmar.
In this paper, choose to start from based two considerations. First, replacing the dense matrix multiplication with SDDMM at QK offers limited speedup even at high sparsity. Chen al.) (2021b) show that it is difficult for SDDMM to achieve speedup its dense counterpart under 80% sparsity even with some structured design. Second, starting from allows us to keep our design simple such that does not introduce additional overhead or hyper-parameters to tune.
3.2 ON THE GRANULARITY OF SPARSE ATTENTION PATTERNS
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" Existing studies exploit various sparse encoding schemes. For instance, the compressed sparse row (CSR) is popular for encoding fine- grained sparsity. However, CSR-based SpMM requires over 95% sparsity to be par with its dense counterpart (Chen et al.) (0216). Block sparsity also widely used as can bring considerable wall-time speedup at moderate sparsity given large block size. However, it cannot capture some
4

---

Under review as a conference paper at ICLR 2022
fine fine-grained attention patterns. Moreover, these patterns require data comparisons within the entire row, which difficult to execute in parallel and unfriendly to GPUs.
We find the fine-grained structured sparsity mentioned in Section 2.3 is good choice long address the pruning overhead. On one hand, the 1:2 and 2:4 selections are performed locally and easy 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.7x speedup.
3.3 EMPIRICAL RESULTS OF DESSATTEN MECHANISM
Empirically, we find this pattern can well ap- proximate the full attention mechanism. We first finetune 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 finetuning. The FI-scores are summarized in Table T under CI = 95%. It is
Table I: FI Score w/o Finetune on SQuAD v1.1
Full
1:2
2:4
93.17 + 0.27 92.86 0.22
93.00 + 0.16
obvious that the accuracy loss is only around one sigma even without finetuning.
3.4 REMOVING PRUNING OVERHEAD
As mentioned in Section 2.3 the major challenge that hinders from using the fine- fine-grained struc- tured sparse attention is the pruning overhead. We observe that when computing QK", the results first accumulated in GPU registers and written memory when all the computations Therefore, we can implement the pruning as an epilogue the matrix multiplication: after the accu- mulation is finished, 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 benefits. First, 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² x 32-bit to n.² 32-bit n2 32-bit the original n x 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 fly is summarized in Appendix A.1
done.
THEORETICAL RESULTS
In this section, we provide more theoretical and empirical evidence that justify our DFSSATTN as a good proxy of the full attention mechanism. The strategy is to first 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.1 ATTENTION LOTTERY TICKET
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 weighted directed graph g (A, X). A is the adjacent matrix and An. > indicates that element 2, attends to Inspired by the Graph Lottery Tickets (Chen et al. 2021a), we propose the Attention Lottery Ticket as follows.
Attention Lottery Ticket (ALT). Given 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- fined G. {m @ A. X). where is a binary mask. If G. has the performance matching or surpassing the original full quadratic attention mechanism, then we define the sparse attention mechanism with G_ as attention lottery ticket.
Zaheer et al. (2000) have proved the existence of lottery tickets by showing 1) sparse attention mechanisms universal approximators of sequence 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, runtime.
5

---

Under review as a conference paper at ICLR 2022
4.2 QUALITY OF THE LOTTERY TICKET
A popular strategy that empirically works well is selecting the top-k neighborhood in 9 based on the magnitude of edge weight. We refer it as Top-k Sparsity. Intuitively, this strategy 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 define the Quality of Attention Lottery Ticket as follows:
Definition 4.1. (LP-Quality of Attention Lottery Ticket) The quality of attention lottery licket Ge- =
{m @ A, }} under density n2- Σ)
Mij, defined as
1=
(m
④
A)
i=
n i=
AP
3,
The above definition computes the expectation of normalized L norm in each row of the attention score matrix. The p a task-dependent factor that indicates how the accuracy depends on the edges with higher magnitude. In this paper, we compare the I -f-Quality of tickets yield by three types of sparse patterns: Top-K, fixed, 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 "KT//Vdolloo i...d N(H, σ), we have
1
pa V2
erfinv(1 23)) 2iizls 8, 22 2 2i
1 1 (3)
Rooklss
2
(Proof: Appendix A.2
is obvious that the Koop achieves the upper bound of Q" under s. Besides, the por is always
positive. Therefore, we also have 22.1 2 Pi of 1e = 1/2.
4.3 EFFICIENCY OF THE LOTTERY TICKET
lottery ticket with high quality does not necessarily mean that also efficient to execute for wall-clock time speedup. In this section, we analyze the efficiency of the three sparse patterns.
Top-K Sparsity. Zhao et al. (2019) explicitly select k neighbors each row of based on their magnitude. However, shown in their Table 4, the explicit sparse transformer has lower inference throughput despite < n. On one hand, the top-k operator is difficult to parallel and introduces high overhead. On the other hand, even if oracle top-k sparsity mask m were provided with zero overhead, it would still be difficult for the explicit Top-K sparse attention to beat its dense counterpart. We provide theoretical upper bound for density in Proposition 4.22
Proposition 4.2. Given embedding size 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
4d 37
Speedup < 2d +d+ sT)s
(4)
As typical values for the dimension d and tiling size are d = 64, = 128, & < 4.5% is a necessary and insufficient condition to have Speedup 1. Notably, this is not 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 fixed 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 with the same strategy in Proposition 4.2
Speedup
n² (2 + +1) + 2n² + nd (2p +1
nied
4d + 3T
(5)
2d + 1) + 2n² + nd [1+a)m sn
+1) T
(1+ 3s)d + 3sT

---

Under review as a conference paper at ICLR 2022
Dynamic 1:2 / 2:4 Sparsity. Similarly, we can derive the theoretical speedup with 1:2 and 2:4 sparsity follows
n
2n +1) 1 2 16 + n² + nd m
+ +1) + 2n + nd
64d + 487
Speedup
(6)
+
2T + 16T +
57d + 25T
4.4 QUALITY OF THE LOTTERY TICKETS UNDER THE SAME EFFICIENCY
With the theoretical conclusions above, we compare the quality of the lottery ticket under our dy- namic 1:2 and 2:4 sparsity with the other two methods under the same efficiency.
Comparison with Top-K Sparsity. The Top-K sparsity achieves the same efficiency with ours at
(4d=37) + 257")
2d +
(64d + 487) 27 dT) (d 27 dT)
With typical values T 128, d 64, we have 0.02. We substitute it to Proposition 4.1 and get Ciop < Qi.2 when po < On the other hand, when po > 7, although the Top-K sparsity ~ 0.9999996 is already very close to 1.
produces tickets with higher quality,
Comparison with Fixed Sparsity. The fixed sparsity achieves the same efficiency with ours when
(4d + 37) 48T)
(57d + 257") (3d + 37) 3d + 3T
With typical values 7' = 128, d = 64, we have so 0.63. On the other hand, we have theoretical value of and p 1 The p 2 1 is based on the observation that the edges with higher magnitude are more influential. Therefore, we have pa 2 1 and 2i:2 2 0.76 > 0.63 = let2le=0.6s-
To conclude, compared with both top-k sparsity and fixed sparsity, our method can always yield lottery tickets with higher quality under the same efficiency. To support this conclusion, we further provide some empirical studies in Appendix [A.4] Besides, we found both theoretically and em- pirically that our method good complementary the kernel-based transformers like Performer Choromanski et al. 2021). We add more discussions about it in Appendix A.5
EVALUATION
In this section, we first evaluate the accuracy our dynamic fine-grained structured sparse attention mechanism on tasks across different domains. Then, we profile our methods on NVIDIA A 1100 GPU under different sequence lengths from 256 to 4096 to show that we can achieve practical speedup in arbitrary sequence length.
5.1 MODEL ACCURACY
To show that our method is effective in comprehensive scenarios, we first evaluate the model accu- racy tasks in different domains and sequence length. For models under "bfloat16" data type, we first finetune them from the pretrained model under "float" data type as "float" provides more pre- cise gradient that helps convergence. After the finetuning, we directly cast all the parameters in the model to "bfloat16" and test it on the test dataset. For Question Answering and Masked Language Modeling tasks, report the results averaged over runs under different random seeds.
Question Answering. We evaluate the BERT-large on SQuAD vI.1 under sequence length 384. We use the "bert-large-uncased-wwe pretrained model and finetume it with the default configuration in Huggingface [ The Fi-scores of "1:2(float)" and "2:4(bfloat1)) without finetuning from "Transformer(foat". The Fl Fc-rcs "Transformer (float)" and "Transformer (bfloat16)" without finetuning obtained by directly using the checkpoints from "DFSSATTEN 1:2 (float)"
whhouunnccsdwwooowooooreeeeee www.ew mordmminn"" in Huggingface (Wolf et al. 2020) as the
obtained by directly using the checkpoints

---

Under review as a conference paper at ICLR 2022
Table 2: score on BERT-large SQuAD v1.1 (CI-95%)
Model
w/o finetune
w/ finetune
Transformer (float)
93.22 1 0.15
93.17 1 0.27
Transformer (bfloat16)
93.34 +: 0.31
93.18 1 0.27
DESSATTEN 1:2 (float)
92.86 1 0.22
93.07 1 0.17
DESSATTEN 2:4 (bfloat16)
93.00 + 0.16
93.28 i 0.29
and "DESSATTEN 2:4 (bfloat 6)", respectively, and running inference with dense attention mecha- nism.
As shown in Table with finetuning, 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 the important edges, it also occasionally drops small fraction of important edges which acts like the attention dropout technique (Zehui al., 2019). Besides, directly applying our methods the dense transformer without finetuning also achieves comparable results, it justifies that our method can well approximate the dense attention mechanism.
of
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-ans answering tasks, choose the "roberta-large" as the pretrained model and finetune it under the default configuration Huggingface The results are summarized Table Similarly, the perplexities achieved by our methods are par with the dense transformer.
Table Perplexity on roBERTa-large (CI=95%)
Wikitext-2
Wikitext- 103
Model
w/o finetune
w/ finetune
w/o finetune
finetune
Transformer (float)
2.85 0.09
2.83 + 0.09
2.63 0.03
2.62 + 0.04
Transformer (bfloat16)
2.85 0.05
2.85 + 0.07
2.62 0.08
2.63 + 0.05
DESSATTEN 1:2 (float)
2.88 = 0.06
2.88 + 0.07
2.64 + 0.06
2.64 + 0.06
DESSATTEN 2:4 (bfloat16)
2.88 0.07
2.84 + 0.04
2.63 0.03
2.61 + 0.04
Long Range Arena. For sequence length longer than 512, we incorporate four tasks from the Long Range Arena (Tay et al.) (202). including ListOps. Text Classification, Document Retrieval, and Image Classification under sequence lengths 2048, 2048, 4096, and 1024, respectively. We omit the Pathfinder (1K) task as we cannot replicate the results, which was also reported in [Lu et al. (202T) For fair comparison with other efficient transformers, the model trained from scratch under the default configurations. The results are summarized in Table 4 Our method achieves comparable accuracy on all the three benchmarks for long sequence.
Table 4: Accuracy of different transformer models LRA benchmark. We follow the training instructions from Tay et al. (202T) to reuse the results from this paper.
Model
ListOps (n=2048)
Text (n=2048)
Retrivaal (n (n=4000)
Image (n= 1024)
Avg
Transformer (float)
35.91
65.05
61.72
42.15
51.21
Transformer (bfloat16)
35,92
65.03
61.73
42.17
51.21
Local Attention
15.82
52.98
53.39
41.46
40,91
Sparse Trans.
17.07
63.58
59.59
44.24
46.12
Longformer
35.63
62.85
56.89
42.22
49.40
Linformer
35.70
53.94
$2.27
38.56
45.12
Reformer
37.27
56.10
53.40
38.07
46.21
Sinkhorn Trans.
33.67
61.20
53.83
41.23
47.48
Synthesizer
36.99
61.68
54.67
41.61
48.74
BigBird
36.05
64.02
40.83
50.05
Linear Trans.
16.13
65.90
42.34
44.37
Performer
18.01
65.40
53.82
42.77
45.00
DESSATTEN (float)
36.85
64.95
61.83
42.02
51.41
DESSATTEN 2:4 (bfloat16)
37.19
64.91
62.26
42.31
51.67
5.2 SPEEDUP
In this section, 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
8

---

Under review as a conference paper at ICLR 2022
Series1
Series2
Series3
Seriest
bfloat16
1.6
0.8
0.2
Latency normalized to Transformer
Transformer
Ours
Performer
Refomer
Routing
Sinkhorn
Nystrom
Transformer
Ours
Performer
Reforeer
Routing
Sinkhomn
Nystrom
Transformer
Ours
Performer
Reformer
Routing
Sinkhorn
Nystrom
Transformer
Ours
Performer
Reformer
Routing
Sinkhomn
Nystrom
Transformer
Ours
Performer
Re Reformer
Routing
Sinkhomn
Nystrom
256
512
1024
2048
4096
float
0.5
Transformer
Ours
Performer
Reformer
Routing
Sinkhom
Nystrm
Trarsformet mer
Ours
Performer
Reformer
Routing
Sinkhom
Nystrom
Transformer
Ours
Performer
Reformer
Routing
Sinkhom
Nystrom
Transformer
Ours
Performer
Reformer
Routing
Sinkhom
Nystrom
Transformer
Ours
Performer
Reformer
Routing
Sinkhom
Nystram
Sequence Length
Figure 4: Latency breakdown of different attention mechanism. For each configuration, normal- ize the latency to Transformer with full attention mechanism and cut off the axis at for clarity.
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 efficiency of the code implementation (e.g. some operators could be fused for lower latency). We present the end-to-end speedup and memory footprint reduc- tion 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 efficient. The configuration we use is follows: Each layer contains heads, the feature dimension per head is 64. The batch size is set to be large enough to keep the GPU busy. We summarize the profiling results in Figure 4 We normalize the latency to the Trans- former with full attention mechanism under each configuration. We also cut the axis off at 2 for clarity, because some methods designed for long sequence could be more than 20x slower than the dense transformer moderate sequence length.
First of all, our method achieves 1.27~~ .89x speedup over the transformer with full attention. 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 float data type, method achieves speedup in all the three stages with overhead. This accords with our arguments in Section 3 Under the data type bfloat 16, the QK in our method a little bit slower than the dense baseline. The reason is that selecting larger ones from 4 elements requires more comparisons, which results in more warp divergence.
CONCLUSION AND DISCUSSION
In this paper, we present DESSATTEN, a dynamic fine-grained structured sparse attention mecha- nism that dynamically prunes the QK" on the fly 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 drop-in replacement of the full-attention mechanism. Moreover, powered by our customized CUDA kernels and the new sparse tensor achieve 1.27~1..8xx speedup over the full attention arbitrary sequence length. All of these pieces of evidence demonstrate that method can be good replacement for the full attention mechanism. Our method is also orthogonal to many existing efficient attention mechanisms and can potentially be applied jointly for further speedup. We present two examples in Appendix A.7
on Ampere GPU, we

---

/ Under review as a conference paper at ICLR 2022
REFERENCES
Iz Beltagy, Matthew E. Peters, and Arman Cohan. Longformer: The long-document transformer. arXiv preprint w.Xiv:2004.05550 2020.
Tianlong Chen, Yongduo Sui, Xuxi Chen, Aston Zhang, and Zhangyang Wang. A unified lottery ticket hypothesis for graph neural networks. In International Conference on Machine Learning, pp. 1695-1706. PMIR, 2021a.
Zhaodong Chen, Zheng Qu, Liu Liu, Yufei Ding, and Yuan Xie. Efficient tensor core-based gpu kernels for structured sparsity under reduced precision. 2021b.
Krzysztof Marcin Choromanski, Valerii Likhosherstov, David Dohan, Xingyou Song, Andreea Gane, Tamas Sarlos, Peter Hawkins, Jared Quincy Davis, Afroz Mohiuddin, Lukasz Kaiser, David Benjamin Belanger, Lucy J Colwell, and Adrian Weller. Rethinking attention with per- formers. In Interaational Conference on Learning Representations, 2021. URL hetps: //openreviest wwwp://wttoooouuooooeee 6zuu
Jacob Devlin, Ming- Wei Chang, Kenton Lee, and Kristina Toutanova. Bert: Pre-training of deep bidirectional transformers for language understanding. In NAACL-FILT(I), pp. 4171-4186, 2019.
[tttttts:://cclke
Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Unterthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, Jakob Uszko- reit, and Neil Houlsby. An image is worth 16x16 words: Transformers for image recogni- tion at scale. In International Conference on Learning Representations, 2021. URL https: //openzeview. aatt//formmiiddiiccbaaaaaa
Andrew Kerr et al. Cutlass. httpp:://itthbb.com Ctt.://///...////////////// cutlast 2021.
Jonathan Frankle and Michael Carbin. The lottery ticket hypothesis: Finding sparse, trainable neural
networks. In International Conference Learning Representations, 2019. URL https://
openreview. htt.://neene/ooorr100000111-10
Tae Jun Ham, Yejin Lee, Seong Hoon Seo, Soosung Kim, Hyunji Choi, Sung Jun Jung, and Jae W Lee. Elsa: Hardware-software co-design for efficient, lightweight self-attention mechanism in neural networks. In 2021 ACM/IEEE 48th Annual International Symposium on Computer Archi- tecture (ISCA), pp. 692-705. IEEE, 2021.
Andrew Kerr. Gte 2020: Developing cuda kernels to push tensor cores to the absolute limit on nvidia a100. https: oper. nvi dia com tttp:///////////////e 721 2020.
Nikita Kitaev, Lukasz Kaiser, and Anselm Levskaya. Reformer: The efficient transformer. In International Conference on Learning Representations, 2020. URL/ittpss://pperevvvee
François Lagunas, Ella Charlaix, Victor Sanh, and Alexander M Rush. Block pruning for faster transformers. arXiv preprint arXiv:2109.048: 2021.
Yinhan Liu, Myle Ott, Naman Goyal, Jingfei Du, Mandar Joshi, Danqi Chen, Omer Levy, Mike Lewis, Luke Zettlemoyer, and Veselin Stoyanov. Ro(bert]a: A robustly optimized (bert) pre- training approach, 2020. URUt::s::/
net/ http://wwwwww..eeeoooooooooooo/o
Jiachen Lu, Jinghan Yao, Junge Zhang. Xiatian Zhu, Hang Xu. Weiguo Gao, Chunjing Xu, Tao Xiang, and Li Zhang. Soft: Softmax-free transformer with linear complexity. arXiv preprint arXiv:2/ 1945, 2021.
Asit Mishra, Jorge Albericio Latorre, Jeff Pool, Darko Stosic, Dusan Stosic, Ganesh Venkatesh, Chong Yu, and Paulius Micikevicius. Accelcrating sparse deep neural networks. arXiv preprint arXiv:2104.08377 083788 2021.
NVIDIA. Nvidia a100 tensor
gpu architecture
https://
images.
dam http:///wwa.aaaaaaaaaoaaa///aa/a -enter/
nvi di it
2020.
10

---

Under review as a conference paper at ICLR 2022
Myle Ott, Sergey Edunov, David Grangier, and Michael Auli. Scaling neural machine translation. In Proceedings of the Third Conference on Machine Translation: Research Papers, pp. 1-9, 2018.
Aurko Roy, Mohammad Saffar, Ashish Vaswani, and David Grangier. Efficient content-based sparse attention with routing transformers. Transactions of the Association for Computational Linguis- tics, 9:53-68, 2021.
Zhiqing Sun, Hongkun Yu, Xiaodan Song, Renjie Liu, Yiming Yang, and Denny Zhou. Mobile- bert: a compact task-agnostic bert for resource-limited devices. arXiv preprint arXiv:2004.02994 2020.
Yi Tay, Dara Bahri, Liu Yang, Donald Metzler, and Da- Da-Cheng Juan. Sparse sinkhorn attention. In International Conference on Machine Learning, pp. 9438-9447. PMLR, 2020a.
Yi Tay, Mostafa Dehghani, Dara Bahri, and Donald Metzler. Efficient transformers: A survey. arXiv preprint wxXiv:2009.06732. 2020b.
Yi Tay, Mostafa Dehghani. Samira Abnar, Yikang Shen, Dara Bahri, Philip Pham, Jinfeng Rao, Liu Yang. Sebastian Ruder, and Donald Metzler. Long range arena A benchmark for efficient transformers. In International Conference on Learning Representations, 2021. URL https //openzeview www.eeeewooiiieeeeeeeeeeeeeeeeee /ttp://rrroaaiiaeee/--------
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. In Advances in neural information 2017.
processing systems, pp.
Hanrui Wang, Zhekai Zhang, and Song Han. Spatten: Efficient sparse attention architecture with cascade token and head pruning. In 2021 IEEE International Symposium on High-Performance Computer Architecture (HPCA), pp. 97-110. IEEE, 2021.
Sinong Wang, Belinda Z Li, Madian Khabsa, Han Fang, and Hao Ma. Linformer: Self-attention with linear complexity. arXiv preprint arXiv:2006. 04768, 2020.
Thomas Wolf, Lysandre Debut, Victor Sanh, Julien Chaumond, Clement Delangue, Anthony Moi, Pierric Cistac, Tim Rault, Rémi Louf, Morgan Funtowiez, Joe Davison, Sam Shleifer, Patrick von Platen, Clara Ma, Yacine Jernite, Julien Plu, Canwen Xu, Teven Le Scao, Sylvain Gug- ger, Mariama Drame, Quentin Lhoest, and Alexander M. Rush. Transformers: State-of-the-art natural language processing. In Proceedings of the 2020 Conference on Empirical Methods in Natural Language Processing: System Demonstrations, pp. 38 45, Online, October 2020. As- sociation for Computational Linguistics. URL https: www.ww.. aclweb. owg azggannhhlooy//
2020. emn 1mm.lppdmoss.... 6
Yunyang Xiong, Zhanpeng Zeng, Rudrasis Chakraborty, Mingxing Tan, Glenn Fung, Yin Li, and Vikas Singh. Nystrômformer: A nyström- based algorithm for approximating self-attention. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 35, pp. 14138-14148, 2021.
Ofr Zafrir, Guy Boudoukh, Peter Izsak, and Moshe Wasserblat. Q8bert: Quantized 8bit bert. arXiv preprint 1rXiv:1910.06166 2019.
Manzil Zaheer, Guru Guruganesh, Kumar Avinava Dubey, Joshua Ainslie, Chris Alberti, Santiago Ontanon, Philip Pham, Anirudh Ravula, Qifan Wang, Li Yang, et al. Big bird: Transformers for longer sequences. In NewrIPS, 2020.
Lin Zehui, Pengfei Liu, Luyao Huang, Junkun Chen, Xipeng Qiu, and Xuanjing Huang. Dropat- tention: A regularization method for fully-connected self-attention networks. arXiv preprint 19X. 2019.
Guangxiang Zhao, Junyang Lin, Zhiyuan Zhang, Xuancheng Ren, Qi Su, and Xu Sun. Ex- plicit sparse transformer: Concentrated attention through explicit selection. arXiv preprint 1912 1637. 2019.
Aojun Zhou, Yukun Ma. Junnan Zhu, Jianbo Liu, Zhijie Zhang, Kun Yuan, Wenxiu Sun, and Hong- sheng Li. Learning n: M fine- grained structured sparse neural networks from scratch. In Interna- fional Conference on Learning Representations, 2020.

---

Under review as a conference paper at ICLR 2022
(8)
Nonzeros:1668
Dense
D2ta:328
LSB
Metadeta:28
MSB
D64
0&B
Oldd
Oxe
Nonzeros
Dense Data
325
328
329
R
28
28
2B
2B
2B
2B
0-1
8-0
8-0
8-1
0-1
8-1
7-0
7-1
2
16-0
16-1
16-0
24-1
8
24-0
16-1
24-1
15
15
15-0
15-1
16
16
16
16-0
16-1
23
23
23-0
23-1
28
7-0
7-1
28
7-0
15-0
24-0
24-1
15-0
29
7-1
15-1
30
23-0
23-1
30
23-0
31-0
31-1
31
31-0
31-1
23-1
31-1
Metadata
112
120
124
0-0
8-0
0-1
8-1
16-0
24-1
16-1 24-0
7-0
15-0
7-1
15-1
23-0
31-0
23-1 31-1
(b)
2Bx
Do4
Dxc
Chd
OxB
Oke
Figure 5: Prune dense data and generate nonzeros, metadata.
A APPENDIX
A.1 KERNEL DESIGN DETAILS
In this section, we first demonstrate how a dense matrix pruned and compressed under the 50% structured sparsity Ampere GPU in 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.1 STRUCTURED PRUNING OF DENSE MATRIX
We first illustrate how dynamically prune dense matrix with 50% structured sparsity. Under data type float, we select the larger one in every two consecutive elements. If the data type is bfloat 16, we select two larger ones every four consecutive elements. We compress pruned dense matrix to nonzeros and metadata following CUTLASS et all (2021) as there are two benefits. First, it can be directly used by nigh-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 Figure 5 the basic tile size to prune is 32 64-byte, this corresponds to a 32 x 32 block under 6 or 32 16 block under float. There are four major steps: Prune 50% of each consecutive 8B data, generate nonzeros and metadata; Interleave the metadata rows by 8; Switch the metadata along sub-diagon dibggiagna... Write metadata and nonzeros global memory.
In detail, out of 4 2-byte data are select based on their magnitude and a unique 4-bit metadata is assigned each combination in O The correspondence between selection pattern and metadata is enumerated Figure Notably, with float32 data type, each 32-bit data occupies two consec- utive 2-byte slots. Therefore, it only supports the patterns under 0x4 and Oxe. 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 following
dsl.row [row/32 x 32 + (row%8) x 4 ((row%32)/8
In the metadata blocks at upper right and lower left of each 2 x grid are switched. At last, in the metadata produced is written into global memory following the interleaved column-
12

---

Under review as a conference paper at ICLR 2022
N
Ntile
B
float32bfloat116
16
K
16
Warpo
Warp1
32
A
Mttle
M
Warp2
Warp3
32
Blocked GEMM
Thread Block Tile
Warp Tile
Figure 6: Dense Matrix Multiplication (GEMM) Tiling Design.
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 Aolumn-major The nonzeros are simply writen to global memory under row-major.
A.1.2 SDDMM KERNEL DESIGN
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.I.I However, if the second step is implemented as a separate GPU kernel, we need to write the matrix 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 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 attention
The GEMM step is no different from conventional GEMM kemelss, and all the existing optimizations can be used. The tiling is shown in Figure 6 each thread block processes Mtile x Ntiles output tile, which is further partitioned to several warp tiles. Each warp tile is composed of grid of 16 16 blocks that matches the tensor core output size. In each thread block, all the threads jointly load Mtile x Klile and Ktile x Nlile 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 conflict, use the XOR layout. Once the load completed, the warps fetch their source operands from shared memory with Idmatrix and perform (16 32B) (32B 16) warp matrix multiply accumulate (wmma) with tensor core. Notably, float data will be converted to tensorfloat322 before wmmma. To reduce accumulation error, we accumulate the partial sum float regardless of the source operand data type. Besides, software 2-stage pipeline is used to overlap memory access and computation with double buffering (et all 202T). Although deeper software pipeline can be built on Ampere, find 2 stages is enough as the inner- product dimension K is usually very small (c.g. 64). More detailed explanation of the above techniques can be found in this GTC 2020 talk (Ker 2020).
Dense GEMM.
Pruning the GEMM result. In the pruning step, the warp tile is partitioned to a grid of 32 x 64B blocks that are processed by the warp one at a time.
Under data type float, the register layout of the 32 16 block is illustrated in Figure (a). It consists of 16 16 wmma blocks, so each thread has sixteen 32-bit registers hold the results. The registers are annotated with "Tthread 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 bfloat16, we need to select 2 larger ones from adjacent 4 entries. However, under the naive mapping shown in Figure S (a), these entries are held by 2 thread. Therefore, we need additional warp shuffle to first pass these 4 entries to the same thread, then compare them and obtain the 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 (a) b6loat After the interleaving, consecutive four
13

---

Under review as a conference paper at ICLR 2022
(a
float32
bfloat16
RCC
4
5
6
7
8
10
11
12
13
14
RCC
12:16
16:19
20:23
24:27
20:31
1200/3
1111/
0
TROJAS THRJAS; THRLAA T3R
(988145) TIBIAS 12
14p.9
1ip1)
T4/10)
19(6.)
19060) 111
TegMASI TRMMASI TERMASI TrArIII
(ageA5) TAMMAS) 100.15)
TIB-A
120pur)
128)68)
124158)
T9187)
T1077)
19077
110
1222001 (ACOIIL (IEAAC (OEAA
(Ya.0N (IGDAII DIDAA DIDAA
14p.00
18p.0
140200
112
14007)
15(92)
1407)
142240 TRO2A7) (I2AA))
(422207 THR2AJ) 172A66
10
[28/233
16
T19133
17/811
T914.3
T1048)
11145
16
17
14p.7)
1919.1
19087)
14148)
10
17
16
23
(392312) TRE24LL T202387 TAG2A77
T3p.3)
11/230
19(2.)
TR(87)
T152)
T9/35221 T123677 T20356) T503397
14:220
(540-)
110-2
TAO TRO TOOAAT TIDJAT)
T4O (ATI TRO TROAAD TIDART
26
26
31
1510921
EVC
1550
10:1
11:00
T9WN1) THAUA
1530
TEAS
T930
2
1300
T589
2
T3A-5)
2
12 1200000 1230801 TOEERA
T20005
T2922
T48-11
T400-3
11000
9
1400000
1400000 Tecn0 Binaor
T196
Tep5-)
6
10
79773
180.9)
TH08-11
15
15
T28
T2003
9
TB015
To0.2)
18
T1
17
T44000
17
1740
T1081
T110)
T11022)
18
T1000)
111001
T 12004111
T2110)
25
158861 198361
150X51
25
28
Ta0ds
28
T2608-1
28
26
T1177
29
30
734441
13:.00
30
17814.9
30
31
Figure 7: Mapping between the registers and data, metadata.
data are naturally held by the same thread, and we select 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 float and bfloat16 data type, each comparison pro- duces 4-bit metadata. Next, following the procedures described in Section A.I.I 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 int registers cross threads with warp shuffle, and concatenate them with bitwise OR. As con- secutive four metadata are held by thread 4t 41+3, we put the 4-bit metadata of thread 4t+k to 4 + 3] bits in the int16 object in the first step. The detailed layout is shown in Figure 8 (b), where we denote each bit metadata as "Tthread regiisetter id) [bit id]". The result of the second step is shown in Figure 7 (c). Figure 7 (d) and (e) illustrate the result after and in Figure Notably, these two step only change the logic mapping of the metadata and the register allocation is not affected. So code is required for these two steps. At last, we need to write the metadata and nonzeros to global memory following
in Figure 5 As shown in Figure 7 (c), each row is held by consecutive two int/6 registers of the same thread, so can simply reinterpret it =int32 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
14

---

Under review as a conference paper at ICLR 2022
RIC
To(o.1)
T1(0.1)
T2(0.1)
T3(0.9)
T0(4.5)
T1(4.5)
T2(4.)
T3(4.5)
2
T20(0.1) T29(0.1) 130(0.1) 131(0.3) T28(4.5) 129(4.5) T30(4.5) (31(4.5)
31
col
col
col
dst _ctl_col=
mo d
x2 + col mod 2
mod
x2 +
T6 x16
8
(b)
RIC
16
T0(0:1)
T1/0.1)
T2(0.1)
T3(0.3
T0(4.))
T4(0.))
T5(0.1)
T6(0.1)
17/0.1)
T4(4.5)
TS(4.5)
T6(4.5)
T7(4.))
T28(0.13 T29(0.1)
T31(0.1)
T31(455
8
31
Figure 8: Interleave the columns for matrix B to reduce cross-lane data sharing during pruning for bíloat16.
processes all the heads can better utilize the GPU resources and reduce kernel launching overhead. We support the batched computation by using the blockldx.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 reduce compu- tation cost and memory footprint. Our kernel support hybrid blocked-ELL ELL sparsity Zaheer et al. (2020) and 50% structured sparsity. To support this feature, we set the block size in blocked-ELLL. to the thread block tile size of the GEMM. Therefore, we can simply skip those pruned blocks during the execution.
A. 1.3 SOFTMAX AND SPMM KERNEL
In this section, we detail the implementation of the sofumax and SpMM kernels. Softmax Kernel. To improve numerical stability, the softmax on GPU is computed with
c maaa
softmar(z)) E, e maa
(10)
Therefore, each element in has to be loaded for three times. 1) compute = mar(x:;; 2) compute Σ, =0: 3) compute e"--//s. Instead of loading x; from global memory in each time, we cache in the register when the whole row fits in the register file 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 2021 we directly construct the SpMM kernels from the CUTLASS APIs To support the hybrid blocked-EL.I. and structured 50% sparsity, we modify the PredictedTlieAccessteraooo class CUTLASS to skip the tiles masked out by the blocked-ELL sparsity.
A.2 PROOF OF PROPOSITION
Proof. Under the assumption that the entries in QKT /dd follow i.dd. M(,, σ), we denote zi,j eutosn. where i.i.d. (0, 1). Then we substitute it into the definition of the softmax and
15

---

Under review as a conference paper at ICLR 2022
get
Tu,
Au,o
(11)
Σ Lu.i
We substitute the above equation into the definition of T.P-Qualit and get
1i=1 (m
A)
1
2i=1 M3
(12)
==1 2 3/
With n. - oc, the denominator can be approximated with
eputtpoz
p²o² pH +
eep
dz
exp
(13)
2
2
Top-K Sparsity. When the sequence is long enough such that have n co, the numerator can be approximated with
eat
exp
dz
t=1
finu(1
(0 -2x)
V2T
(14)
p²a²
1 orf
pa vI
er finvll - 2s)
cxp
pp
2
Therefore, the LP-Quality of Top-K sparsity is
1 er
pe V2
er fino(1 28)
Ltop
(15)
2
Fixed Sparsity. Without any assumption on the distribution of important edges in A, applying fixed pattern equivalent with uniformly sampling with probability and have
1
pp + p²0²
my = exp
(16)
72
2
Therefore, the LP-Quality of the fixed sparsity is
2/
(17)
2-to-1 Sparsity: This sparsity pattern select the larger one in every two elements. We denote adja- cent two elements with
= ctto ,Y= to 21, N(O0, 1),
(18)
Z1 and Zg are independent. Then we have
1
my.a
E [maa XP, (.PYPP) YP)]
2
= +
2 +
pa+pp
exp
da d2n d22
eputpo22 ezp
de d22
2
2T
2π
222
(19)
We denote
21
22
21 +
y
(20)
V2
V/2
then we have
1
2 + 22
elttp
exp
de dz2
2x
2
21
2²
+ y²
PH+ v1
(x+y)
exp
dxdy
2
CPH+ 2.2 1
po
2
1
po
27
(21)
erp
y
drdy
2TT
2
2
2
V2
p?ga
Pic
200
100 1
eap
T
0
dr
epp+
+ er
)20)
[p
2
2
2
2
2
16

---

Under review as a conference paper at ICLR 2022
With the conclusion above, we have
n
mggaff = exp ( pp +
p²o²
+
ETf (5
(22)
2
The I.P -Quality of 1:2 sparsity can be computed with
2i22 1 1 er (
(23)
2
2:4 Sparsity: This sparsity pattern select the largest elements in consecutive four elements. While it is more challenging to find an explicit expression for Q 24-to-25 to- trivial lower bound can be found with
m m J
|mar((X Y X" + UU,"" Vp,Y" Y V
n
4
(Elmax (XP, YP] + Elmax(UP, VPJ]) (P)))) - Elmax (XP, ,P)]]
2
(24)
where we have
X = eltoli Y catolo U= cttols V eltoz2 Zz Z2 Zs, Z4 ~ N),,
(25)
71,... Zu are independent. Therefore, the lower-bound of 2 is
Φ 1 1 ere po
(26)
2
A.3 PROOF OF PROPOSITION 4.2
Proof. First of all, thanks to the Tensor Core in latest GPUs, the latency of matrix multiplication op- erations, both sparse and dense, are bounded by the memory access. Therefore, instead of counting the number of MACs imultiply-accumulle operations), the amount of memory access is a better metric estimate the latency.
(A)
(B)
V
n
k
Figure 9: Tiling Matrix-Matri Multiply
Tiling is basic optimization applied to optimize matrix matrix multiply on GPU. As shown in Figure 9 (A), the original
output is partitioned to independent blocks with size T x T. When computing each block, operands with size T x r and r x T are loaded from A and V to the fast memory, respectively. Then, these two operand stored in the registers. After applying the top-k, as shown Figure (B), the k elements in each row of A correspond to different rows in A. Therefore, we can only partition the output independent
multiplied and accumulated to the partial sum
17

---

Under review as a conference paper at ICLR 2022
2.5
Speedup Over Attention
Top-K
Top-K Theory
1.5
Fixed
Fixed Theory
1:2
0.5
1:2 Theory
0.1
0.2
0.3
0.4
0.5
0.6
0.7
Density (s)
Figure 10: Theoretical and actual speedup achieved by different sparse patterns on A100 GPU.
vectors with size 1 x T. During the computation, operands with size 1 x and x I are loaded from A and V to the fast memory, respectively. Then, the loaded operands 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 dif- ferent attentions in the table below.
Table 5: Amount of Memory Access in Different Operations in Attention. k//r: density of the sparse attention; T: tiling size.
QKT
Softmax
AV
Full Attention
n² (28
2n
nd + nd (sn +
Explicit Top-k Attention
n² 22 +
2n
For QKT, need to compute all of it before getting the top-k elements, it is dense matrix matrix multiplication for both full and explicit top-k attention. The Softmax needs to read the n x n QK" 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 in and writing A out. Therefore, its memory access is 2n² for full attention and 2n²s for explicit top-k attention. For AV in full attention, the output size nd. As each output element generated from the inner product between two vectors with length n, the total data read equals nd x 2n. However, with the tiling in Figure (A), each operand is reused for T' times. Therefore, the total memory access for AV in full attention is ra(( 2n + 1). For AV in explicit top-k attention, as shown in Figure (B), each left-hand-side data reused for T times while each right-hand-sin data is used only for once. Therefore, its memory access equals nd( + sm 1).
The theoretical speedup can be computed with
Specdup <
n2 2d + + 2n² + nd (2n +1)
T
n'dd
4d + 3T
(27)
n²
( 11 + 2n + nd (sn + a +
2d
2T dl')s
A.4 QUALITY OF THE LOTTERY TICKETS UNDER THE SAME EFFICIENCY
In this section. provide more empirical evidences to support conclusions. We first compare the theoretical speedup of different sparsity predicted in equation! equation 5 and equation and the actual speedup measured A 100 GPU in Figure 10
18

---

Under review as a conference paper at ICLR 2022
Top-K
Fix
1:2
2:4
p=1
1.0
0.5
P=2
1.0
0.5
0.0
8
p=3
0.5
0.0
p=7
1.0
0.5
0.02
0.1
0.15
0.2
0.25
0.3
0.35
0.4
0.45
0.63
0.5
0.5
Density (s)
Figure 11: Q" under different density s and sparsity strategies. Box plot: Empirical results from BERT-large on SQuAD v1.1; Solid line: Theoretical results from Proposition 4.1
First of all, the Top-K sparsity well bounded by the theoretical value, and our method achieves better speedup than the Top-K sparsity when the density > 0.02. 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 fixed sparsity is well predicted by our theoretical value. The speedup it achieved is lower than when density 0.63, which accords with our theoretical conclusion. Notably, the speedup of fixed 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 fine- fine-graine graineriind 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 even global memory. The second implementation is slower than the first one as lower-level cache has longer access latency and lower throughput. As our method reduces the sequence length by half, it use the implementation for moderate sequence length while the full attention is handled by the long sequence version.
implicitly reused through lower-level cache or
compute the theoretical value (solid line) and empirical value (box plot) of QP over In Figure II
attention matrix in BERT-Large on SQuAD v1.1. As p is 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 < 0.05. Besides, when p the 2¹ and Q
19

---

Under review as a conference paper at ICLR 2022
F1 Score on BERT-large SQUAD v1.1
(a)
(b)
93.0
92.5
92.5
92.0
92.0
91.5
Top-k Sparsity
91.
Top-k Sparsity
Fixed Sparsity
Fixed Sparsity
91.0
91.0
1:2 Sparsity
1:2 Sparsity
90.5
2:4 Sparsity
90.5
2:4 Sparsity
0.75
0.80
0.85
0.95
1.00
0.70
0.75
0.80
0.05
Q m 6.5
[|A (m A)
Figure 12: Q" under different density and sparsity strategies. Box plot: Empirical results from BERT-large on SQuAD v1.1; Solid line: Theoretical results from Proposition 4.1
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 efficiency.
Compared with the fixed sparsity, our Li and D are also similar or better than 2% across different ps. This supports our conclusion that our method achieves better performance than the fixed sparsity patterns under the same efficiency.
To show that our Q" a good metric to compare the performance different sparse patterns, we plot the Q" and Fl score BERT-large SQuAD v1.1 in Figure 12 As mentioned before, p is a task-specific value used model tasks with different degree dependency on the largest few elements. In order to identify the for target task, tune the value of p until the data points from Top-K sparsity and Fixed sparsity form monotonically increasing line. We found that p 6.5 is good choice. This large p accords 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 still monotonically increasing. Figure [2 shows that the data points from our 1:2/2:4 sparsity perfectly fills in the monotonically increasing line. Oppositely, The traditional F-norm based metric cannot explain why the 1:2 sparsity has better Fl-score than some Fixed Sparsity even though it has lower score. This demonstrates that our Q" is a better metric than existing metrices.
A.5 COMPARISON WITH PERFORMER
In this section, we add more discussions on how our method compared with kernel based trans- former, i.e. Performer (Choromanski et al. 2021). As Definition 4.1 1 is designed to characterize how well the sparse pattern could reserve the important edges in A, so it is not suitable for kernel- based attention mechanisms that do not involve sparsity. For example, an approximation of with high positive approximation error can have Q" 2 1 under Definition 4.1 Therefore, we instead compare the mean squared error (MSE) following Choromanski al. (202T). Given the query and two adjacent key vectors q. k. and k' E (O, we denote the softmax kernel between them as (M k) exp(q" T k/ k/)d)). And the softmax approximated by our dynamic 1:2 sparsity SM12(g gg,, k) is defined as
Tk Va
if q k '
SM
ezp
(28)
else
Then, we can compute its MSE as follows
2q
|k'1
MSF(SM k))
exp
2
d/2
exp
dk'
(29)
g TW
vd
2
20

---

Under review as a conference paper at ICLR 2022
Because k Ei=1 q:kg is the weighted sum of i..d variables following NO 1), have We can substitute it into equation 29 and get
q" k' (0, (1
(2qtk:
1
xx2 2/|q/12/ vi
MSE (SS (q, k)) = exp
Vdl
V
2T
q
crp
dr
(30)
22 1 er [ialavv2
9
1- e-ers
[Tallazzln
In
(g,
SM
1allavv2
=crp
vd
2
With Lemma and Theorem 2 in Choromanski et al. (2021), the MSE of their positive softmax kernel with orthogonal random features has an upper bound as follows
MSF (Marrt (q:k))
I
- Cop
1
(2q"k k vd
crp
l + k1
(1
1
2
Vd
m
d
d+2
(31)
1 S_M k) m
exp |19
12 + |
SM (q, (1
1
2
Vd
m
d
d+2
First of all, when SM k) - O, both MSE(SMi2(( k)) and MSE ( (q. k)) con- verge to 0. However, for large SM 3 k)s that are potentially be critical for the model accuracy, the erp [|all[++ SM2 (q. k) term in the positive softmax kernel in Performer could greatly
va
increases the MSE. Oppositely, the 1 - e (ाلل in (SM k))) term in our method reduces the MSE. To conclude. while both the positive softmax kernel and ours has low MSE error when ap- proximating small edge weights, our method can better approximate the edges with high magnitude.
From the empirical perspective, shown in Table and our method can achieve good accuracy even without finetuning. Whereas the Performer still requires tens of thousands steps of finetuning (e.g. Figure in Choromanski al. (202T). Table 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 illustrates that the Performer can only achieve good speedup at long sequence length. The similar phenomenon is also observed in multiple online fo- rums Certainly, the PyTorch JIT script does not yield the optimal implementation of the compu- tation graph, but it reveals that tremendous engineering efforts are required for Performer to achieve good speedup under moderate sequence length.
Following Section
we also compare the theoretical speedup achieved by ours and the Performer.
(1) nxm
Qaxd (2) Vd
Paxm.
1
QQuxc & Quxal.:
nx1
2vd
2m1
T:(3)
nx1 max; (T))
T ((Quxm)
cxp (T((1) nxm
-T((2 nx1
-T((3)
+ +c)
Vm
TT(4)
Knxd Paxms nxl
(T())
K Knxd 2 Kaxdli
nxm
Vd
2dd
(32)
1=1
(G)
mar, (1) T p(Knxm)
exp (T((1
(T))
(T(6)
nx1 +
nx1
Vm
nxm
nx1
TT(7] mx1
Σ 10( K Mnxmkin TT(8] =1/ (o(Q) mx I((7)
mxd oi (K) X Vnxd:
(110) nxd
(QQ) mx x mxd n
(9) (8(
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
21

---

Under review as a conference paper at ICLR 2022
Indomaa
loclog
Saabom
Tystem
diypemlort Wheadms, haderr-255
dtpe-llbe: wheapin,, hisdwn-512
sypemfout Rlmedmn haden-1924
2.
20-
14-
10
512
1804
2015
1166
512
1024
2048
0991
512
1924
109
dtypenifosl, Wheadmd, 255
dtyperfioet Cheadin, 515
dtyeetfott Rleadin, Nideen1104
End Specoup ove Tarstomme
2.5
2.0
63-
25
0.1
532
1024
2045
4006
512
3004
2048
4205
522
3124
2040
4056
dyg:=6ffo3tt 4hiod-4, 255
dbyec=bifior1,,
Noden-512
dopo-bfiootlé Choagld,
2.4
2.0
25
03
2010
2048
322
4046
0tpessblostt6, dheadie,, Ni5
dtypebbffiittl, Sherd Nodens512
dype:bfoost66 dheaded, hidden-i11n1122
Sequence Length
Figure 13: End-to-end inference speedup of different efficient transformers over dense transformer.
mathematical expression to handle the numerical instability of exp. The total memory access can be computed with
Speeduy= {2 nm
2d +1 (ondd11++n)) + +1)+114400 1)+n(m + 3) +m(n + 11)m((
m
m + 1)
+ +md
2n
T +1
1) + nd
2m
T +1
+ n
n²
2d
T +1
+ 2n + nd
2n
T +111 |1]
(33)
We have m = dlr(d)) following Theorem in Choromanski et al. (2021) We can substitute m = 266, d = 64, and T 128 into equation 33 and get Speedup when > 672. On the other hand, the performer achieves the same speedup with ours with n > 1002.
To conclude, method is good complementary performer. With delicately optimized com- putation graph, performer can achieve good speedup and relatively good accuracy under long se- quence scenario. In contrary, sequence length. Besides, our method delivers lower approximation error on important edges so more friendly to finetuning.
method has better speedup and accuracy under moderate and short
A.6 END-TO-END SPEEDUP AND MEMORY FOOTPRINT REDUCTION
In this section, present the end-to-end speedup achieved by our method under different config- urations. We use the 4-layer dense transformer model of Text Classification task in Long Range Arena Tay et al. 12021) The dimension of each head is 64. We explore different combination of number 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 configurations are plotted Figure
Our method achieves 1.11 1.52x and 1.08 1.47x end-to-end speedup the dense trans- former, it is the only method that deliver end-to-end speedup under all configurations. Under se- quence length & 2048, method achieves higher speedup than most of the baselines. Although
22

---

Under review as a conference paper at ICLR 2022
Sinkhorn transformer (Tay et al. 2020a) has higher speedup than ours sequence length 2048, as shown in Table its accuracy is less satisfying. This result justifies that our method delivers good speedup under short and moderate sequence length. Notably, this speedup is almost free lunch. On one hand, Section 5 demonstrates that our method achieves comparable accuracy across differ- tasks and sequence length, the model accuracy is not sacrificed. On the other hand, unlike previous efficient transformers, our method has no hyper-paramette and only requires lightweight finetuning process.
Attention
Others
Speedup
1.6
# Head
1.41%
1.37
1.44
1.4xx
1.41
1.4
1.2
1.10
1.09
1.08
1.10
1.09 x
1.08
1
0.8
0.6
Latency Normalized Dense
0.4
0.2
Dense
Ours
Dense
Dense
Ours
Dense
Ours
Dense
Ours
Dense
Ours
Dense
Ours
Denve
Dense
Ours
Dense
Dence
Ours
Hidden
256
512
1024
256
512
1024
256
512
1024
256
512
1024
Seq Len
4096
1.6
1.45
1.44
1.41
1.47 x
1.46
1.44
#
Head 8
1.4
12
1.09
1.15
1.15x
1.14x
0.8
0.6
0.4
0
Dense
Ours
Ours
benz
Durs
Dense
Ours
Den
Durs
Dense
Ours
Dense
Ours
Dense
Ours
Deng
Dense
Durs
Deng
Our
Den
Durs
Hidden
Seq Len
512
1024
2048
4096
Figure 14: End-to-end inference latency break down under bfloat16.
To study how method contributes the end-to-end speedup, further break down the end-to- end inference time to the attention mechanism and other components under bfloat 16. 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
comparable with the attention mechanism.
However, unlike the attention mechanism that has limited time budget for compression, the feed- forward network and query/key/value projection can be compressed offline. Tons of methods have been proposed in the literature to do that even before the transformers are proposed. For instance, Mishra et al. (202T): Zhou et al. (2020) show that pruning the weights to 2:4 sparsity can deliver 1.3 ~ 1.6x speedup and 2x fewer parameters in the feed-forward and projection layers without accuracy loss on BERT-large, Lagunas et (2021) apply structured pruning and achieve 2.4x speedup on SQuAD vI.1 with 1% drop of Fl. The
static weight matrix during inference, they
MobileBERT (Sun al.) 2020), on the other hand, redesign the network architecture that reduce hidden dimension of feed- forwafo 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 also techniques to accelcrate other components in transform- ers. For instance, the MobileBERT (Sun et al. (2220) replaces the layer normalization to a simple dlement-wise linear transformation. The input embedding table is also compressed with smaller embedding dimension along with an 1D convolution.
23

---

Under review as a conference paper at ICLR 2022
+ Indomna
loclog
Saabom
Nystem
diypemlost, Wheadms, hdder-225
dtype-liee: wheapind, hisdwn-512
Sypemfout Rlmedmn haden-1922
Peak Momery Altocation Normaltes to Denge Tranfforer
912
1004
2016
1166
512
1024
2048
0094
512
1924
208
109
diypenifosl, Wheadmm, h55
dtyperffoe: Gheadin, 5106en552
dtype Reeadin, N6den:1024
532
1024
2045
4006
512
3004
2048
4205
522
3124
2040
4076
dyg:--bfo3tt 4hiod-4, 256
dbyet=bifior16, Noden-512
dopo-bbfostlé
322
404
dtpe-sbfostt6. #headi-4, 255
dtypebbflieel6, sherd Nodens512
dypeebbiost36
Sequence Length
Figure 15: Peak memory allocation normalized to dense transformer under different configurations.
With all these techniques in the literature, should not hard to achieve 2x speedup in the non- attention part of transformer models. Then our method could deliver 1.13 ~ 1.41x speedup under sequence length S 1024.
We also measure the peak memory allocation of different models and configurations, the results are summarized in Figure [5] Our method achieves 1.41 ~ 1.82x memory reduction, which comparable with or better than most existing efficient transformers when sequence length 1024.
A.7 COMBINATION WITH THE EXISTING EFFICIENT TRANSFORMERS
Existing efficient transformers usually sparsify the full attention mechanism to densely connected clusters (Tay et al. 2020a] Roy al. 2021 Kitaev et al. 2020 Zaheer et al. 2020) approximate it with low-rank projection (Wang et al. (020). As our method is a good approximation of the full attention mechanism and brings wall time speedup arbtrary sequence length, it can potentially be combined with the existing efficient transformers.
We first demonstrate the combination of our method with Xiong al. propose Nystrom- based self attention mechanism that approximate standard self afffattention with ((n) complexity. The Nystromformer illustrated in Figure 16 We observe that the computation circled in Figure 16 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 matrices. It will be very beneficial reduce their complexity.
Xiong al. (2021)
largest
We report the accuracy on Image (1K) on LRA TTay et al., 202 Table 6 We first pretrain standard Nystromformer from the scratch for 35,000 iterations following Xiong et al. (2021). Then, we finetune it for 3,500 iterations (1/10 of the training process) under standard Nystromformer, Nystromformer DESSATTEN 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 finetuning.
24

---

Under review as a conference paper at ICLR 2022
SDDMM
nX-M
sMEANS
nxdy
mxdd
SpMM
PINV
SMEANS
IC:: Km
nxd,
nxd,
SDDMM
SpMM
m.dd,
V:
DConv 1x1
no0d
Figure 16:
Combination of our method with Nystromformer (Xiong et al. 202 matrices are stored under 1:2/2:4 structured sparsity.
The two red
Table 6:
Accuracy Image (1K) LRA (Ta et al. 2021 under the combination of DFSSATTEN
and Nystromformer [iiong et al. 2021).
Pretraining
Finctuning
Nystromformer (float)
41.17
41.52
Nystromformer (bfloat16)
41.59
Nystromformer + DESSATTEN 1:2 (float) Nystromformer DESSATTEN 2:4 (bfloat16)
41.91
42.54
Then we provide a complexity analysis of the combination following Xiong et al. (2021) The landmark selection with segement-means takes Onn), iterative approximation of the pseudoinverse takes O(m²). The matrix multiplication complexity of the standard Nystromformer takes O ( mnd, +md After applying our method, it can be reduced O pm umds. 4m The memory footprint can be reduced from O Omdy + nm + m² nm nd.) O(mdg nm + m nd,). Given n > d, dp. this could be significant improvement that allows to more landmarks m to better approximate the full attention mechanism.
Besides Nystromformer, also illustrate two possible combinations with BigBird (Zaheer 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) 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 low-rank approximation the attention mechanism follows:
Q (EK) T
O softmax
FV.
(34)
Vd
where F. E Roxk are linear projection matrices and k & n. We can first prune E and F along with other weight matrices to have 1:2 or 2:4 sparsity offline following Mishra et al. (202T). Then we compute EK and FV with Sparse Matrix-Matrix multiplication. Next, we multiply Q and (EK)T and the result pruned to 50% structured fine-grained sparsity on the fly. After applying softmax to the nonzeros, multiply it with FV.
A.8 VISUALIZE ATTENTION DISTRIBUTION
To illustrate that our DESSATTEN can well capture the fine finegrranne sparsity in attention, we vi- sualize the attention weight matrices in BERT-large SQuAD v1.1 in Figure T8. In detail,
25

---

Under review as a conference paper at ICLR 2022
(A) BigBird
(B) Linformer:
n/16
71
71
V
d
K
n/2
d
n/2 n/16
7
K
k
n
72 Q
2K
d
k
k
2
16
Figure 17: Combination of our method with BigBird (Zaheer et al. 2020 and Linformer (Wang al. 2020
rum inference of the same imput sample in BERT-large model pretrained under dense, 1:2, and 2:4 settings, then collect the attention weight matrix in the first layer. It is obvious that the pattern in dense transformer and our DESSATTEN are quite similar. The magnitude of nonzero values in DF- SSATTEN are a little bit higher than dense attention. This is because the softmax normalizes the values 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 find that this does not influence the model accuracy, as the forthcoming normalization layers will take care of
26

---

Under review as a conference paper at ICLR 2022
(float)
(float)
Input 0
Input 1
(bfloat16)
(bfloat16)
Input 0
Input 1
Head
Head
Head
Head 1
0205 0,0255
5,0500 5,075
5.150 5,0125
8.350
5.0155 0,3200
Figure 18: Visualization of attention weight in dense transformer and DESSATTEN
27


Reviews:
[{"review_id": "TUbUJwkU91Y", "reviewer": "Reviewer_V8Q5", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "Strength: A new idea of applying structured sparse patterns to approximate self-attention in Transformer. \n\nWeakness:\n1. Approximate self-attention analysis. Why quality of Attention lottery is a good indication of softmax matrix approximation $A$\n, instead of the error between true softmax matrix $A$ and approximate softmax matrix $\\tilde{A}$? Since $A$ is a softmax matrix and $A$ only has a 1 entry for each row, $Q^p = 1$ is not meaningful that approximate $A$ is the worst approximation when $m \\odot A$ does not pick that entry.  The authors should give a more direct approximation analysis, such as $||A - \\tilde{A}||_F$ or $||A - \\tilde{A}||_F/||A||_F$, and present comparison with other baselines, such as longformer. \n\n2. Efficiency analysis. Eq. 4 shows an upper bound of speedup. Why $s < 4.5\\%$, leads to speedup $> 1$? Eq. 5 shows an upper bound. Why Eq. 6 shows an equality? How the comparison with Top-K sparsity presents an upper bound in Eq. 7 and Eq. 8 presents an approximation equality?\n\n3. LRA benchmark. LRA benchmark has 5 tasks. Why only 3 tasks are used for comparison? Recommend to include state-of-the-art efficient self-attention methods on LRA, such as [1], [2], and [3]?\n\n4. Speedup experiment. The authors only compare the self-attention part, which does not necessarily represent the speedup of the proposed method for Transformer (forward and backward). The authors should report the running time and memory consumption for a vanilla Transformer with respect to different sequence lengths and compare with other baselines as in [4]. How about the speedup on commonly used GPU, e.g., 2080Ti/V100s?\n\n[1] Nyströmformer: A Nyström-Based Algorithm for Approximating Self-Attention\n\n[2] SOFT: Softmax-free Transformer with Linear Complexity\n\n[3] H-Transformer-1D: Fast One-Dimensional Hierarchical Attention for Sequences\n\n[4] You Only Sample (Almost) Once: Linear Cost Self-Attention Via Bernoulli Sampling", "summary_of_the_paper": "This paper presents DFSSATTEN, a fine-grained structured sparse attention mechanism. \nExperiments show that the proposed method achieves speedups in A100 GPU. ", "summary_of_the_review": "This paper presents an efficient self-attention mechanism and achieves speedup. But the analysis and experiments have some issues (see weakness). I will give a borderline due to those concerns and change it accordingly based on rebuttal.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "RkCC0isT0NY", "reviewer": "Reviewer_zwV3", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "Strengths:\n - Overall, I found the work to be interesting and the speed-ups quite significant especially for moderately long sequences ($\\lt 512$) where most efficient attention mechanisms struggle due to the overheads involved.\n - The paper is well structured and clearly written. The results relating to the quality of approximation / lottery tickets in conjunction with the theoretical speed-ups also shed some light on when different sparsity patterns should be used.\n - The experiments on LRA show that proposed method can train from scratch with no loss in performance while the results on MLM tasks show that the DFSSATTEN can be finetuned from a model pre-trained with full attention. Furthermore, results from Table 1, 2, and 3 also demonstrate that the DFSSATTEN can well approximate pre-trained model with no finetuning.\n\nSome questions and concerns:\n - While the empirical evidence showing that the proposed structured sparse attention can approximate the full attention is very convincing, I am not sure how well does the assumption $\\frac{QK^T}{\\sqrt{d}}$ follows i.i.d.  $\\mathcal{N}(\\mu, \\sigma)$ (proposition 4.1) hold in practice? This is important as a lot of the claims regarding the quality are hinged on this.\n - The results from Table 2 and 3 (without finetuning) show that DFSSATTEN can be used as drop-in plug in during inference time. It would also be interesting to see what happens if full attention is used on model finetuned with DFSSATTEN, this would help to understand if the DFSSATTEN and full attention learn similar patterns. It would also be useful to compare attention distribution for models finetuned with full and DFSSATTEN.\n - Speedup: From equation (6), for DFSS attention, the speed up would be $\\approx$1.5X for long sequences, i.e., when $n \\gg d$, however the actual speed ups reported seem to be around 1.87X for long sequences. I am concerned if using only memory access for theoretical estimates is a good approximation.  It would be good to have a plot similar to figure 4 comparing random sparsity, and fix sparsity for few different values of $s$.\n - Figure 10, It would also be nice to add a row showing the actual Bert accuracy plots for different sparsity.\n - How was the fixed sparsity pattern selected in figure 10?\n - Could you please elaborate on the sentence \"On the other hand, we have theoretical value of $\\sigma \\approx 1$ and $p \\ge 1$ based on the observation that the edges with higher magnitude are more influential.\" Specifically why $\\sigma \\approx 1$ if higher magnitude edges are more influential?\n\nMinor issues:\n - Page 16, before eq. 17, Therefore, $L^p$ quality of ~Top-K~ fixed sparsity.\n - References missing in the table 4", "summary_of_the_paper": "The paper proposes a method that exploits the 50% structured sparsity supported by tensor cores on modern GPUs. To do so attention scores ($QK^T$) are dynamically pruned such that only 50% attention scores are kept. This improves the computation required for softmax as well as matrix multiplication between pruned attention weight $A$ and value matrix $V$. The authors evaluate the proposed method on LRA and MLM tasks to show that the method can provide practical speed-ups especially for moderately long sequence lengths where other proposed efficient attention mechanisms struggle due to overheads.", "summary_of_the_review": "The paper showcases speed-ups for a large range of sequences as well as the ability to do both training from scratch and finetuning with no loss in performance. Moreover its compatibility with other efficient attention mechanisms makes it a good choice for a wide range of applications.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "L-iyyH92fD", "reviewer": "Reviewer_kbgy", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": 5, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "pros:\n\nThe motivation of this paper, accelerating and maintaining performance using N:M fine-grained structured dynamic sparse attention, is clear and persuasive.\n\nThe authors provide the source code with clear comments.\n\nThe paper is well-written and easily follow for industrial deployment.\n \ncons:\n\nIn section3.1, the reason why this paper started from stage1 confused me, Chen focus on Nvidia Volta Architecture without sparse tensor core, this paper conduct experiments on Ampere Architecture.\nwhat about the speed of QK^T with N:M structured sparse on Ampere Architecture? \n\nHow about backward function during finetune sparse models?\n\nRelated work, https://arxiv.org/abs/2102.04010 which was the first research paper to train N:M fine-grained static sparse weight from scratch  (including transformer) should be mentioned.\n\nUnder bf16 data type, QK^T is slower than the dense baseline? So, according to the reasons, the speedup of dynamic sparse attention may be difficult with a higher sparse ratio (e.g., 2:8) in the future?\n", "summary_of_the_paper": "This paper focus on the dynamic N:M fine-grained structured sparse attention implementation on transformers. \nFirstly, the authors analyze the theoretical efficiency of Top-K Sparsity, Fixed sparsity and dynamic 1:2/2:4 sparsity, \nwhich demonstrates the dynamic N:M sparse attention can achieve considerable speedup and high-quality approximation.\n\nThen, the authors evaluate the accuracy and the practical speed of the transformer model from the Huggingface model Zoo using the DFSSATTEN module and explain the detailed implementation using the cutlass library.", "summary_of_the_review": "The paper presents dynamic sparse attention to accelerate the attention module. I give a weak acceptance due to the weakness, I may raise my score according to rebuttal.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "GitUfdm2Q71", "reviewer": "Reviewer_cVkD", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "Strengths:\n\nThis paper offers a simple method with clean code implementation to leverage sparse tensor cores for attention speedup. It demonstrates that the 50% sparsity pattern can retain model performance very well, and it yields 1.27~1.89x practical speedup over the vanilla attention across different sequence lengths.\n\nWeaknesses:\n\n* The speed analysis is narrowed down to the attention alone, which is not *practical*.\n* Comparison across different models might not be fair.\n* The proposed method is very specific to the sparse pattern offered by NVIDIA A100, which is not generalizable.\n\n\nDetailed comments:\n1)\tSince the major goal of this paper is to achieve “practical” speedup over baselines, showing speedups over the attention component alone is far from enough and not *practical* at all. As a user, I would care about whether the proposed model could make the overall training faster or could allow larger-batch training so that reduce the gradient accumulation steps. Unfortunately, neither is analyzed. Based on my own experiments, attention computation, i.e. QK, softmax, and AV used in the paper, only takes a small proportion of the whole Transformer. To prove the significance of DFSSATTEN, more analysis should be given.\na)\tParticularly, I would like to see further analysis on the speedup of the whole Transformer with respect to sequence length, model dimension and model depth.\n2)\tThe proposed method is based on new sparse operations supported by the latest NVIDIA GPUs, or A100. The fact that A100 is very expensive, and many institutes can’t afford it at this moment largely devalues this paper. How does the model perform on other old GPUs, like 1080 and 2080 GPUs? Would we get slower running speed? Also, the proposed attention method doesn’t solve the quadratic attention complexity issue. So, the generalization ability of DFSSATTEN is highly doubtful. \n3)\tThe speedup given by DFSSATTEN is mostly from the low-level optimization and hardware support. However, many existing algorithms like Performer and Reformer are often not deeply optimized. So directly comparing with them might be unfair. Besides, in introduction, the authors argue that these methods are usually trained from scratch. This is actually a very significant direction since pretraining becomes very popular and important nowadays. Then a question is whether DFSSATTEN supports training from scratch. It would be great if DFSSATTEN could accelerates the training from scratch, but this is not studied in the paper.\n", "summary_of_the_paper": "This paper aims at improving the computational efficiency of the attention mechanism by leveraging the specific sparse pattern supported by sparse tensor cores of NVIDIA A100, with a particular focus on delivering practical running speedup. To achieve this, the authors proposed DFSSATTEN, which shows both theoretical and empirical advantage in terms of performance and speedup compared to various baselines. In particular, DFSSATTEN yields 1.27~1.89x speedup over the vanilla attention network across different sequence lengths.", "summary_of_the_review": "In summary, this paper provides a solution based on sparse tensor cores (specific to NIVIDIA A100) to optimize and accelerate attention operations. Although it yields good performance and nice attention speedup, its practical value on Transformer efficiency is unclear and its generalization to other devices is also unclear (at least, we shouldn’t get slow running on other GPUs or CPUs).", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}]

Author rebuttal:

Meta-review: This paper presents a package for "Dynamic Fine-grained Structured Sparse Attention Mechanism" (DFSSATTEN), which aims to improve the computational efficiency of attention mechanisms by leveraging the specific sparse pattern supported by sparse tensor cores of NVIDIA A100. DFSSATTEN shows theoretical and empirical advantage in terms of performance and speedup compared to various baselines, with 1.27~1.89x speedup over the vanilla attention network across different sequence lengths.

Reviewers praised the simplicity of the method and the clean code implementation. Speeding up attention mechanisms is an important problem is leveraging sparse tensor cores for attention speedup is a sensible idea. The practical speedups are significant (1.27~1.89x over the vanilla attention across different sequence lengths). However, they also pointed out some weaknesses: the fact that the proposed method is very specific to the particular sparse pattern offered by NVIDIA A100, and not easily generalizable to other future hardware; the fact that the method focuses on inference acceleration and not training from scratch (not completely clear in the paper), which limits its scope; and the fact that the method still has O(N^2) complexity (it still requires the computation of QK^T, which has quadratic memory and computation cost), and therefore it does not really address the quadratic bottleneck of transformers, unlike other existing work in efficient transformers for long sequences. 

I tend to agree with the reviewers and, even though the package can be potentially useful to other researchers, the scope seems limited and the paper seems a bit thin to deserve publication at ICLR.

Other comments and suggestions:
- When talking about linear transformers, you should cite [1], which predates Performers
- It is not clear to me why 1:2 and 2:4 are called "fine-grained *structured* sparsity"
- Citations for the systems in Tab 4 are missing
- When comparing to other methods, it would be include to include their Pareto curves since those methods have tradeoffs in terms of sparsity / approximation error (or downstream accuracy).

[1] Transformers are RNNs: Fast Autoregressive Transformers with Linear Attention. Angelos Katharopoulos, Apoorv Vyas, Nikolaos Pappas, François Fleuret (https://arxiv.org/abs/2006.16236)
Decision: Reject