DFSSATTEN:  Dynamic Fine-grained Structured Sparse Attention Mechanism

1 INTRODUCTION . Transformers ( Vaswani et al. , 2017 ) have achieved competitive performance across various domains like NLP ( Ott et al. , 2018 ) and Computer Vision ( Dosovitskiy et al. , 2021 ) . The key feature that sets them apart from traditional neural network architectures is the attention mechanism ( Vaswani et al. , 2017 ) , which allows the transformers to gather information from the embeddings of elements in the input sequence in an adaptive and learnable manner . Nevertheless , the high computation cost and memory footprint brought by the attention mechanism make it difficult to apply transformers to latency-sensitive tasks . Many efficient attention mechanisms ( Tay et al. , 2020b ; Zaheer et al. , 2020 ; Beltagy et al. , 2020 ; Tay et al. , 2020a ; Roy et al. , 2021 ; Kitaev et al. , 2020 ) have been proposed to address this issue . However , most of them drastically modify the original full attention mechanism and introduce a handful of hyper-parameters to tune . Therefore , they require tremendous engineering effort to deploy and optimize . Besides , they usually need to be trained from scratch instead of exploiting pretrained models like BERT ( Devlin et al. , 2019 ) . Some of them rely on fixed sparse patterns or 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 to reach good accuracy . Last but not least , previous methods usually introduce additional operators like top-k , sort that cause large overheads and offset their benefits at moderate sequence length . In this paper , we present DFSSATTEN , a simple and effective sparse attention mechanism that address the limitations mentioned above . DFSSATTEN 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 DFSSATTEN offers several advantages over existing studies : • It requires minimal changes to the original full-attention with no hyper-parameters to tune . This makes it a drop-in replacement of the full attention that only requires to change a few lines of code . Moreover , it can directly exploit existing pretrained models like BERT ( Devlin et al. , 2019 ) and RoBERTa ( Liu et al. , 2020 ) . • We dynamically prune the attention scores based on their magnitude under only 50 % sparsity . This allows the pruned attention matrix to reserve the important entries , achieving on par model accuracy with full attention even without fine-tuning . • Our method introduces zero overhead on existing GPU hardware . As a result , we are able to achieve wall-clock time speedup and memory footprint reduction over the full attention in arbitrary sequence length . To conclude , our main contributions are summarized below : • We propose DFSSATTEN , a dynamic sparse attention mechanism that is a drop-in proxy of the full attention mechanism . Its effectiveness is justified by both empirical and theoretical evidence . • We present a dedicated CUDA kernel design to completely remove the pruning overhead . The pruning is implemented as an epilogue of the dense matrix multiplication which produces the attention score matrix . • We evaluate DFSSATTEN on tasks cross various domains and sequence lengths . It achieves 1.27 ∼ 1.89× 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 an input sequence X= ( x1 , .. , xn ) ∈ Rn×d , the full attention mechanism can be defined as O = Softmax ( QKT / √ d ) V , ( 1 ) where Q = XWq , K = XWk , and V = XWv are query , key , and value matrices . QKT forms a full-quadratic adjacency matrix , whose edge weights are the dot-product similarity between all the elements in the sequence . This adjacency matrix is standardized with 1/ √ d to keep the unit second moment and then normalized with softmax . At last , the row feature vectors in V are aggregated according to the normalized adjacency matrix by multiplying them together . In the rest of this paper , we denote A = Softmax ( QKT / √ d ) for simplicity . We refer QKT as the attention score matrix and A as the attention weight matrix . 2.2 EFFICIENT ATTENTION MECHANISM . 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 et al. , 2020b ) . Fixed Sparse Patterns . Zaheer et al . ( 2020 ) ; Beltagy et al . ( 2020 ) 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 et 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 et al . ( 2020a ) ; Roy et al . ( 2021 ) ; Kitaev et al . ( 2020 ) apply various clustering methods and only compute the attention within each cluster . Although computing full attention in each cluster is more friendly to GPU compared with fine-grained 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 n × n to n × k with linear projection . Choromanski et al . ( 2021 ) introduce the FAVOR+ which approximates the softmax with kernel method . This allows them to change the computation order and reduce the asymptotic complexity to linear . However , the low-rank projection and kernel construction also introduce considerable overhead . This makes these methods only effective under long sequence length . Besides , the previous studies drastically change the attention mechanisms , tens of thousands pretraining or 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-grained structured sparsity in the A100 GPU . As shown in Figure 1 , the dense input matrix is pruned with fine-grained 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 float16 , the 2:4 sparsity is used which selects two larger ones among four consecutive elements . After the pruning , the result is compressed to nonzeros and metadata . The nonzeros contain the value of reserved data that is 50 % smaller than the original one . The metadata records the index of the nonzeros in the origins matrix . It takes 4 bit metadata to record the decision of each 1:2 or 2: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 a 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 feedforward part of the transformer up to 1.9× . However , to the best of our knowledge , no previous studies use it in the attention mechanism where the attention weight matrix is dynamically generated for each sequence . One plausible explanation is that during pruning , GPU must read the whole dense matrix to be pruned from the memory . Then , after selecting the elements to be reserved under the 1:2 and 2:4 pattern , it must also generate the metadata encoded in a special format such that the metadata can be used efficiently later . All these overheads will offset the benefit brought by the pruning if we do it on the fly . 3 DYNAMIC FINE-GRAINED STRUCTURED SPARSE ATTENTION MECHANISM . In this section , we first give an overview of our DFSSATTEN 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 DFSSATTEN mechanism is simple and effective , as illustrated in Figure 2 . Compared with the full-quadratic attention mechanism , our method dynamically prunes attention scores without incurring storage or computation overhead , while maintaining the effectiveness of attention . More importantly , our method can achieve practical speedups of attention on existing GPU hardware with customized CUDA kernels . Listing 1 shows all the modifications to be made to use DFSSATTEN . # Full attention mechanism def full_attention ( q , k , v ) : attn_weight = torch.bmm ( q , k.transpose ( 1 , 2 ) ) attn_weight = torch.nn.functional.softmax ( attn_weight , -1 ) return torch.bmm ( attn_weight , v ) # DFSS attention mechanism import dspattn def dfss_attention ( q , k , v ) : attn_weight , metadata = dspattn.bmm ( q , k ) attn_weight = torch.nn.functional.softmax ( attn_weight , -1 ) return dspattn.spmm ( attn_weight , metadata , v ) Listing 1 : Example of using DFSSATTEN . The “ dspattn ” is the package we developed . 3.1 DESIGN CONSIDERATIONS FOR EXPLOITING ATTENTION SPARSITY As illustrated in Figure 3 , the attention mechanism can be considered as three stages : QKT , Softmax , and AV . To design a sparse attention mechanism , the first decision to make is where should we induce pruning to sparsify the attention . If we start from 0 , all the three stages will be benefited from the sparsity given effective implementation on GPU : the dense matrix multiplication between Q and K will be replaced with the sampled dense-dense matrix multiplication ( SDDMM ) which only computes the entries identified by the sparse pattern . The Softmax only operates on the nonzero values in each row . The original dense matrix multiplication between A and V will be replaced with a sparse matrix-matrix multiplication ( SpMM ) which multiplies a sparse matrix with a dense matrix . However , as it is not possible to exactly know which entry in QKT has higher magnitude before computing QKT , starting from 0 usually requires some additional components to predict the location of important entries . Starting from 1 requires us to compute a dense matrix multiplication between Q and K. The benefit is that we can explicitly select important entries from QKT without prediction . As the softmax is a monotonically increasing function , starting from 2 does not offer any benefits over 1 but throws away the opportunity to accelerate Softmax . In this paper , we choose to start from 1 based on two considerations . First , replacing the dense matrix multiplication with SDDMM at QKT offers limited speedup even at high sparsity . Chen et al . ( 2021b ) show that it is difficult for SDDMM to achieve speedup over its dense counterpart under 80 % sparsity even with some structured design . Second , starting from 1 allows us to keep our design simple such that it does not introduce additional overhead or hyper-parameters to tune .

Summary: 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.