SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Runtime performances in a vacuum. We test our implementation with different numbers of buckets  $nb$  and random keys, queries, and values. In these tests, we assume a hash bucket is provided for free for each head of each key and query (they are sampled uniformly at random (torch.randint(0, nb))). In practice, runtime experiments on sequence modeling tasks show that obtaining the buckets can be cheap and in no way prevents us from improving the attention runtime (see Fig. 6). We compare with causal FlashAttention over the entire sequence. Importantly, to ensure a fair comparison, we take into account pre-processing and post-processing steps required to reshape the tensors for both methods. For our method this includes stable sorting by bucket index and transposing tensors, for the baseline only the transposition is required, see App. B.2 for detailed code. Fig. 3.b summarises our findings. We observe large improvements in runtime as the number of buckets  $nb$  and the sequence length increases.
Language modeling on OpenWebText2. For sequences of length  $T = 8192$  and  $T = 16384$  we train transformer language models using FlashAttention (F-LM), and identical models replacing only the FlashAttention by our Hash-based sparse attention (H-LM) using  $nb = 16$  hash buckets. In Fig. 6 we see that it takes the same number of iterations for H-LM and F-LM to reach a given perplexity. However, H-LM iterations are  $1.8 \times$  and  $2.3 \times$  faster for respectively  $T = 8192$  and  $T = 16384$ . As a result, H-LM models reach a given perplexity much faster than their F-LM counterpart. Interestingly, we observe the H-LM models gain speed during training, see App. C for additional details.
Comparison with Reformer. We compare the speed and performance of our hash-sparse implementation with the Reformer hashed attention. For all comparisons, we always equalize the average bucket size. Results are summarized in Fig. 4. Benchmarks with random inputs show that both our hash-sparse implementation and the Reformer, as expected, are linear with respect to the sequence length (Fig. 4.a). However, we still achieve a significant speedup thanks to our more efficient kernel. More importantly, Fig. 4.b shows that the fixed attention structure imposed by the Reformer does not allow to capture all of the hash collisions, with the coverage decreasing steeply as the sequence length increases. On the contrary, our method is exact and covers every bucket collision in the attention
(a) Iter to reach perplexity
(b) Time to reach perplexity
(c) Speed gains
Figure 6: Training Language Models (LM) on OpenWebText2 (Gao et al., 2020) using our hash-based sparsity (H-LM) or FlashAttention over the entire sequence (F-LM). We train on sequences of length 8192 and 16384 and use 16 buckets for all of our H-LM models. We show that it is possible to use our proposed hash-based sparsity to significantly gain in speed while not compromising the perplexity. In Fig.(a) we see for both sequence lengths the perplexity decreasing similarly as a function of the iteration. In fact, H-LM even slightly outperform the baseline. Fig.(b): H-LM reach lower perplexity much faster than their F-LM counterpart. Fig.(b) and (c): H-LM models are significantly faster than F-LM models for a given sequence length. The gap widens as the sequence length increases.
matrix. This is reflected in Table 5: our hash-sparse attention layer outperforms the Reformer attention even for shorter sequences.
4.3 Query/Key-Dropping Based Attention
Q/K-dropping mechanism used. We show that naively dropping heads for each key and query at random can already yield competitive results while significantly improving the runtime. While better dropping schemes could be devised, they are outside of the scope of this work.
Runtime performances in a vacuum. We test our implementation with different sparsity ratios, corresponding to the probability of dropping some head associated to a given key or query. We assume that the tensors indicating the dropping of each head of each query and key are given for free, along with some random key, query, and value tensors. To ensure a fair comparison, we take into account pre-processing and post-processing steps required to reshape the tensors for both methods, see App. B for more details. For our approach, we hope reducing the size of the key, query and value tensors and computing the attention on those would be faster than using FlashAttention over the entire sequence. For this, the time gained by computing the attention on smaller tensors should be larger than the overhead of re-ordering tensors to build those smaller tensors. In Fig. 7.a, we show a naive implementation using existing PyTorch functionalities only starts to provide a speedup when dropping more than  $70\%$  of the keys and queries. Fig. 7.b shows that using our proposed implementation provides significant speedups even at relatively low sparsity levels. The linearly increasing cost of reshaping the tensors is rapidly compensated by large gains over the quadratic cost of self-attention.
Language modeling on OpenWebText2. For sequences of length  $T = 8192$ , we train transformer language models using FlashAttention (F-LM), as well as identical models replacing only the FlashAttention by our Q/K-dropping sparse attention (D-LM). We train with several sparsity ratios, dropping  $30\%$ ,  $50\%$ , and  $70\%$  of the heads of keys and queries at random. In Fig. 8 we observe that while high sparsity can negatively affect the perplexity, lower sparsity D-LM models are matching F-LM models in perplexity per iterations, while training nearly twice as fast. Importantly, the dropping pattern is not static. An interesting approach similar to curriculum learning in which we start the training with very large sparsity and reduce it linearly during training is studied in App. C.
(a) Naive implementation
(b) Our implementation
Figure 7: Runtimes of the full Flash-attention of Dao et al. (2022) and several implementations of Query/Key dropping based sparsity. For this figure we show total times for the forward and backward passes. For sparse methods, we drop at random a percentage of keys and queries, this percentage is indicated on the right of each curve. Fig.(a): A naive implementation consisting in creating compact representations of the key, value, and query tensors by removing dropped keys and queries. As a result, the attention matrix is no longer triangular (see Fig. 1). We call the PyTorch scaled.dot_productattention method with a custom but still causal mask. The non-triangular mask prevents FlashAttention to be used and only dropping more than  $70\%$  of the keys and queries seems to improve the runtime over attending the entire sequence using FlashAttention. Fig.(b): Our modification of FlashAttention allows to improve over the runtime. Similar to the naive implementation, reshaping the tensor induce an overhead which compensates the speed gain for shorter sequences. However, this offset is compensated by a strong margin as the sequence length increases. Our implementation allows significant gains over FlashAttention even for low levels of sparsity. The detailed runtimes for the forward and backward passes can be found in App. C.
(a) Iter to reach perplexity
Figure 8: Training Language Models (LM) on OpenWebText2 (Gao et al., 2020) using random Query/Key dropping based sparsity (D-LM) or FlashAttention over the entire sequence (F-LM). Dropping keys and queries randomly is naive and our point here is not to show that this approach is a good way to use the proposed Q/K-sparsity attention, rather we want to demonstrate that it is possible to significantly gain in speed while not losing too much in perplexity—even with a naive approach, and in a very dynamic way (two sequences are allowed to have completely different dropping patterns). For all methods we train over sequences of 8192 tokens. Fig.(a): While dropping large portions of keys and queries slows down the decrease of perplexity per iteration, dropping  $30\%$  seems to match the baseline F-LM. Fig.(b) Our method is significantly faster to reach a given perplexity. Interestingly, more sparsity does not necessarily mean decreasing the perplexity faster. Fig.(b) and (c): Using our Q/K-sparse implementation we train significantly faster than the baseline method.
(c) Speed gains
We develop and validate an efficient kernel that can make sparse attention based on dynamic patterns very fast. We hope that our contribution will inspire the community to research dynamic attention patterns in a way that is less constrained by a tight computational budget.
The computational cost of large attention models remains both a practical issue in scaling up to very large contexts, and a fundamental research question to close the gap between the energy usage of biological systems to that of GPU systems able to run very large models. Dynamically modulating the computation is an obvious direction to address this challenge.
The authors acknowledge support from the Swiss National Science Foundation under grant number CR-SII5-193716 - "Robust Deep Density Models for High-Energy Particle Physics and Solar Flare Analysis (RODEM)". We also thank Igor Krawczuk for interesting discussions and suggesting using Triton.
Andoni, A., Indyk, P., Laarhoven, T., Razenshteyn, I., and Schmidt, L. Practical and optimal lsh for angular distance, 2015.
Behnke, M. and Heafield, K. Losing heads in the lottery: Pruning transformer attention in neural machine translation. In EMNLP (1), pp. 2664-2674. Association for Computational Linguistics, 2020.
Beltagy, I., Peters, M. E., and Cohan, A. Longformer: The long-document transformer. CoRR, abs/2004.05150, 2020. URL https://arxiv.org/abs/2004.05150.
Bender, E. M., Gebru, T., McMillan-Major, A., and Shmitchell, S. On the dangers of stochastic parrots: Can language models be too big? In Proceedings of the 2021 ACM Conference on Fairness, Accountability, and Transparency, FAccT '21, pp. 610-623, New York, NY, USA, 2021. Association for Computing Machinery. ISBN 9781450383097. doi: 10.1145/3442188.3445922. URL https://doi.org/10.1145/3442188.3445922.
Borgeaud, S., Mensch, A., Hoffmann, J., Cai, T., Rutherford, E., Millican, K., van den Driessche, G., Lespiau, J., Damoc, B., Clark, A., de Las Casas, D., Guy, A., Menick, J., Ring, R., Hennigan, T., Huang, S., Maggiore, L., Jones, C., Cassirer, A., Brock, A., Paganini, M., Irving, G., Vinyals, O., Osindero, S., Simonyan, K., Rae, J. W., Elsen, E., and Sifre, L. Improving language models by retrieving from trillions of tokens. CoRR, abs/2112.04426, 2021. URL https://arxiv.org/abs/2112.04426.
Child, R., Gray, S., Radford, A., and Sutskever, I. Generating long sequences with sparse transformers. CoRR, abs/1904.10509, 2019. URL http://arxiv.org/abs/1904.10509.
Choromanski, K., Likhosherstov, V., Dohan, D., Song, X., Gane, A., Sarlós, T., Hawkins, P., Davis, J., Mohiuddin, A., Kaiser, L., Belanger, D., Colwell, L. J., and Weller, A. Rethinking attention with performers. CoRR, abs/2009.14794, 2020. URL https://arxiv.org/abs/2009.14794.
Dao, T., Fu, D. Y., Ermon, S., Rudra, A., and Ré, C. Flashattention: Fast and memory-efficient exact attention with io-awareness. In NeurIPS, 2022.
Gao, L., Biderman, S., Black, S., Golding, L., Hoppe, T., Foster, C., Phang, J., He, H., Thite, A., Nabeshima, N., Presser, S., and Leahy, C. OpenWebText2 dataset, as part of 'the Pile: An 800gb dataset of diverse text for language modeling'. arXiv preprint arXiv:2101.00027, 2020.
Goyal, S., Choudhury, A. R., Raje, S., Chakaravarthy, V. T., Sabharwal, Y., and Verma, A. Power-bert: Accelerating BERT inference via progressive word-vector elimination. In ICML, volume 119 of Proceedings of Machine Learning Research, pp. 3690-3699. PMLR, 2020.
Hutter, M. The human knowledge compression contest. URL http://prize.hutter1.net, 6, 2012.
Wang, H., Zhang, Z., and Han, S. Spatten: Efficient sparse attention architecture with cascade token and head pruning. In HPCA, pp. 97-110. IEEE, 2021.
Wang, S., Li, B. Z., Khabsa, M., Fang, H., and Ma, H. Linformer: Self-attention with linear complexity. CoRR, abs/2006.04768, 2020.
Weidinger, L., Mellor, J., Rauh, M., Griffin, C., Uesato, J., Huang, P., Cheng, M., Glaese, M., Balle, B., Kasirzadeh, A., Kenton, Z., Brown, S., Hawkins, W., Stepleton, T., Biles, C., Birhane, A., Haas, J., Rimell, L., Hendricks, L. A., Isaac, W., Legassick, S., Irving, G., and Gabriel, I. Ethical and social risks of harm from language models. CoRR, abs/2112.04359, 2021.
Wu, Y., Rabe, M. N., Hutchins, D., and Szegedy, C. Memorizing transformers, 2022.
Xia, M., Zhong, Z., and Chen, D. Structured pruning learns compact and accurate models. In ACL (1), pp. 1513-1528. Association for Computational Linguistics, 2022.
Zaheer, M., Guruganesh, G., Dubey, K. A., Ainslie, J., Alberti, C., Ontaño, S., Pham, P., Ravula, A., Wang, Q., Yang, L., and Ahmed, A. Big bird: Transformers for longer sequences. In Larochelle, H., Ranzato, M., Hadsell, R., Balcan, M., and Lin, H. (eds.), Advances in Neural Information Processing Systems 33: Annual Conference on Neural Information Processing Systems 2020, NeurIPS 2020, December 6-12, 2020, virtual, 2020. URL https://proceedings.neurips.cc/paper/2020/black/c8512d142a2d849725f31a9a7a361ab9-Abtract.html.
Zheng, L., Wang, C., and Kong, L. Linear complexity randomized self-attention mechanism. In ICML, volume 162 of Proceedings of Machine Learning Research, pp. 27011-27041. PMLR, 2022.
Zhou, W., Ge, T., Wei, F., Zhou, M., and Xu, K. Scheduled drophead: A regularization method for transformer models. In EMNLP (Findings), volume EMNLP 2020 of Findings of ACL, pp. 1971-1980. Association for Computational Linguistics, 2020.
A.1 Structured attention
With a focus on performances on downstream applications, Raganato et al. (2020) show that it is possible to replace all but one attention head with static attentive patterns only relying on the position, e.g. attending to the previous token. Similarly, Tay et al. (2021a) investigate removing entirely the dot product attention and instead let the input token alone predict its attention pattern over the entire sequence, or use a random attention pattern. While they show their approach can be competitive on certain downstream task, results on language modeling seems to indicate those attention patterns are not expressive enough to model natural language. Peng et al. (2021) propose Random Feature Attention (RFA) which relies on kernel methods to approximate the softmax attention and achieve a linear computational complexity w.r.t. the sequence length. A follow up work by Zheng et al. (2022) augments RFA with randomized attention to build an unbiased estimator of the softmax attention. Instead of looking for a linear approximation of the softmax, Qin et al. (2022) propose to replace the softmax by a linear projection kernel and a cos-based re-weighting mechanism which scales linearly with the sequence length. While kernel based method seem like a good compromise in terms of speed vs. performance, they have been shown to underperform on certain downstream tasks Tay et al. (2021b). Compared to those methods, we are not trying to replace or find an estimator of the softmax attention, we instead provide an efficient way to leverage different forms of sparsity of the softmax attention matrix. Our speed gains do not come from a conceptually different way to compute the attention but simply deciding not to compute certain regions of the softmax attention matrix.
A.2 Pruning
The large body of works on pruning mostly focus on faster inference on downstream tasks. As a result, many methods are modifying the training process to facilitate pruning, often resulting in slower training.
Pruning heads in transformers. Many works have investigated dropping entire attention heads in transformer architectures. Let the entire attention matrix  $\mathbf{A}$  be of shape  $B \times H \times T \times T$  with  $B$  the batch size,  $H$  the number of heads, and  $T$  the sequence length. Dropping entire heads imposes an implicit structure over  $\mathbf{A}$  which is now of shape  $B \times H' \times T \times T$ , with  $H'$  the number of non-dropped heads. Michel et al. (2019) and Voita et al. (2019) both observe—in the context of Neural Machine Translation (NMT)—that a large fraction of heads can be dropped without significantly impacting performance. While their goal is not primarily to speed up training, many of methods following on this insight try to incorporate a sparsifying mechanism during training which facilitate dropping heads at test time (Behnke & Heafield, 2020; Peng et al., 2020; Xia et al., 2022; Li et al., 2021). Still limited to NMT and downstream classification tasks, Zhou et al. (2020) found a regularizing effect of dropping heads during training. In comparison, less works have investigated dropping head in the context of language modeling, with text generations applications in mind. In comparison to those works, we propose a way to take advantage of dynamic sparsity structures, which is much more general but still includes head dropping. We show the potential of our work to enable speedup during training. Our work is orthogonal to head-dropping mechanisms and could be used in addition to those.
Pruning tokens in transformers. While dropping heads can be seen as dropping the model parameters generating keys, queries and values for those heads, another more recent line of work looked into dropping tokens. Dropping entire tokens imposes an implicit structure over the attention matrix  $A$  which is now of shape  $B \times H \times T_{Q} \times T_{KV}$ , with  $T_{Q}$  and  $T_{KV}$  the number remaining queries and keys. Goyal et al. (2020) obtain faster inference on downstream tasks by using an attention based scoring mechanism to eliminate redundant input vectors. Wang et al. (2021) develop a joint algorithm-architecture framework which speeds up inference for downstream classifications tasks and language generation tasks. Their method includes head and token pruning along with specific hardware optimizations. Our work can be used to implement those approaches and does not require custom hardware accelerators. We moreover allow dropping individual heads instead of entire tokens.
B Details On The Implementation
We here give implementation details for experiments of § 4. First, in App. B.1, we describe in more details the implementations of our custom Triton kernels introduced in § 3. Secondly in App. B.2 we provide the python code used in our runtime benchmarks of § 4, including pre and post processing steps required to reshape and re-order tensors. Lastly in § B.3 we give more details on the hyperparameters used in our sequence modeling experiments of § 4. The code for all our experiments can be found via the following link: https://github.com/epfml/dynamic-sparse-flash-attention.
B.1 Triton Kernels
QK-Sparse Triton Kernel. In Alg. 1 we detail the core of the QK-sparse algorithm from § 3.1. This algorithm is computing the softmax attention result only for one block of queries, corresponding to one head. In practice this algorithm would be run in parallel for all blocks of queries for all heads. We find the index of the last relevant tile to compute by iterating over block of key indices and comparing them with the largest query index for the current block of queries. This works as  $q_{idx}$  and  $k_{idx}$  have a monotonic structure thanks to the stable sort used when reshaping the tensors (see the pre-processing code in App. B.2 for more details). We apply causal masking locally by looking at  $q_{idx}$  and  $k_{idx}$ , q query can only attend to past keys:  $\text{mask} = q\_idx[:, None] >= k\_idx[None, :]$ . The backward pass relies exactly on the same trick, we first iterate over query indices to find the starting block of queries.
Algorithm 1 Forward Pass for the QK-sparse kernel
Require: Matrix  $Q\in \mathbb{R}^{N_Q\times d}$  matrices  $K,V\in \mathbb{R}^{N_{KV}\times d}$  index matrices  $Q_{idx}\in \mathbb{R}^{N_Q\times d}$  and  $K_{idx}\in \mathbb{R}^{N_{KV}\times d}$  ,softmax scaling constant  $\tau \in \mathbb{R}$  , softmax statistics vectors  $M\in \mathbb{R}^{N_Q}$  and  $L\in \mathbb{R}^{N_Q}$  output tensor  $O\in \mathbb{R}^{N_Q\times d}$  query block size  $B_{m}$  ,key block size  $B_{n}$  starting query index startm. 1: Initialize  $\pmb {o}\gets (0)_{B_m\times d}\in \mathbb{R}^{B_m\times d},\ell \gets (0)_{B_m}\in \mathbb{R}^{B_m},\pmb {m}\gets (-\infty)_{B_m}\in \mathbb{R}^{B_m}$  2: Load current block of queries  $\pmb {q}\leftarrow \pmb {Q}[\mathrm{start}_m:\mathrm{start}_m + B_m,:]$  3: Load current block of query indices  $q_{idx}\gets Q_{idx}[\mathrm{start}_{m}:\mathrm{start}_{m} + B_{m}]$  4: end  $\leftarrow 0$  5: for startn in range(0,  $N_{KV},B_m)$  do 6: Load block of key indices:  $k_{idx}\gets K_{idx}[\mathrm{start}_{n}:\mathrm{start}_{n} + B_{n}]$  7: if min  $(k_{idx})\leq \max (q_{idx})$  then 8: end  $\leftarrow$  end+1 9: end if 10: end for 11: for n in range(0, end) do 12: Start of current key block startn  $\leftarrow n*B_n$  13: Load current block of values  $\pmb {v}\gets V[\mathrm{start}_n:\mathrm{start}_n + B_n,:]$  14: Load current block of keys  $\pmb {k}\gets \pmb {K}[\mathrm{start}_n:\mathrm{start}_n + B_n,:]$  15: Load current block of key indices  $k_{idx}\gets K_{idx}[\mathrm{start}_{n}:\mathrm{start}_{n} + B_{n}]$  16: Compute inner product qk  $\leftarrow \tau q.k^{\top}$  17: Apply causal mask qk  $\leftarrow$  MASK(qk,qidx,kidx) 18: Update softmax statistics  $\ell ,m,o\gets$  UPDATE_STAT(Sl,m,qk,o,v) 19: end for 20: Write O[startm : startm + Bm,:]  $\leftarrow$  o 21: Write L[startm : startm + Bm]  $\leftarrow$  l 22: Write M[startm : startm + Bm]  $\leftarrow$  m
Hash-Sparse Triton Kernel. In Alg. 2 we detail the core of the Hash-sparse algorithm from § 3.2. This algorithm is computing the softmax attention result only for one block of queries, corresponding to one head. In practice this algorithm would be run in parallel for all blocks of queries for all heads. We find the index of the first and last relevant tiles to compute by iterating over block of key hashes and comparing them with the largest and smallest query hashes for the current block of queries. This works as we reshaped our  $Q, K, V$  tensors by sorting them by their hash values (see the pre-processing code in App. B.2 for more details). In addition to causal masking, we also
enforce attention to happen within the same bucket: mask = (q_IDX[ :, None] >= k_IDX[ None, :]) & (q_hash[ :, None] == k_hash[ None, :]). In our experiments we often replace >= by > to prevent a query to attend to itself as in the Reformer. As a side note, for most application it would also be fine to only enforce causal masking and allow attention across buckets within a tile. While this could add some serendipity in the attention computation, some applications might require masking based on hash. The backward pass relies exactly on the same trick, we first iterate over query indices to find the starting and end blocks of queries.