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