SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Algorithm 2 Forward Pass for the Hash-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}$  matrices containing the hash values  $Q_{hash}\in \mathbb{R}^{N_Q\times d}$  and  $K_{hash}\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: Load current block of query hashes  $q_{hash}\gets Q_{hash}[\mathrm{start}_{m}:\mathrm{start}_{m} + B_{m}]$  5: start  $\leftarrow 0$  6: endhash  $\leftarrow 0$  7: for startn in range(0, NKV, Bm) do 8: Load block of key hashes:  $k_{hash}\gets K_{hash}[\mathrm{start}_{n}:\mathrm{start}_{n} + B_{n}]$  9: if min(khash)  $\leq$  max(qhash) then endhash  $\leftarrow$  endhash  $+1$  11: end if 12: if max(khash)  $<  \min (q_{hash})$  then 13: start  $\leftarrow$  start  $+1$  14: end if 15: end for 16: end  $\leftarrow$  endhash 17: for  $j$  in range(start, endhash) do 18: Load block of key indices:  $k_{idx} = K_{idx}[jB_n:(j + 1)B_n]$  19: if min(kidx) $\leq$ max(qidx) then end  $\leftarrow j + 1$  20: end if 21: end for 22: end for 23: for n in range(start, end) do 24: Start of current key block startn  $\leftarrow n*B_{n}$  25: Load current block of values v  $\leftarrow V[\mathrm{start}_n:\mathrm{start}_n + B_n,:]$  26: Load current block of keys k  $\leftarrow K[\mathrm{start}_n:\mathrm{start}_n + B_n,:]$  27: Load current block of key indices  $k_{idx}\gets K_{idx}[\mathrm{start}_{n}:\mathrm{start}_{n} + B_{n}]$  28: Load current block of key hashes  $k_{hash}\gets K_{hash}[\mathrm{start}_{n}:\mathrm{start}_{n} + B_{n}]$  29: Compute inner product qk  $\leftarrow \tau q.k^{\top}$  30: Apply causal and bucket mask qk  $\leftarrow$  MASK(qk,qidx,kidx,qhash,khash) 31: Update softmax statistics  $\ell ,m,o\gets$  UPDATE_STAT(S(  $\ell ,m,qk,o,v)$  32: end for 33: Write O[startm : startm + Bm,:]  $\leftarrow o$  34: Write L[startm : startm + Bm]  $\leftarrow \ell$  35: Write M[startm : startm + Bm]  $\leftarrow m$
Accumulating softmax statistics while avoiding NaNs. The following is a brief summary of how the FlashAttention algorithm (Dao et al., 2022) proposes to accumulate softmax statistics when iterating over blocks of keys. Given an input vector  $\pmb{x}$ , our goal is to compute  $\mathrm{softmax}(\pmb{x}) \triangleq e^{\pmb{x} - \max(\pmb{x})} / \sum_{i} e^{\pmb{x}_i - \max(\pmb{x})}$ . Let  $f(\pmb{x}, m) \triangleq e^{\pmb{x} - m}$ , and  $\ell(\pmb{x}, m) \triangleq \sum_{i} f(\pmb{x}, m)$ . Hence:
$$
\operatorname {s o f t m a x} (\boldsymbol {x}) = \frac {f (\boldsymbol {x} , \max  (\boldsymbol {x}))}{\ell (\boldsymbol {x} , \max  (\boldsymbol {x}))}
$$
Given a vector  $\pmb{x} \triangleq [\pmb{x}_1, \pmb{x}_2]$ , let  $m_g = \max(\pmb{x})$  (global max), and  $m_1 = \max(\pmb{x}_1)$ , we notice:
$$
\begin{array}{l} \operatorname {s o f t m a x} (x) = \frac {f (\boldsymbol {x} , m _ {g})}{\ell (\boldsymbol {x} , m _ {g})} \\ = \frac {f (\boldsymbol {x} _ {1} , m _ {1})}{e ^ {m _ {1} - m _ {g}} \ell (\boldsymbol {x} _ {1} , m _ {1}) + \ell (\boldsymbol {x} _ {2} , m _ {g})} + \frac {f (\boldsymbol {x} _ {2} , m _ {g})}{e ^ {m _ {1} - m _ {g}} \ell (\boldsymbol {x} _ {1} , m _ {1}) + \ell (\boldsymbol {x} _ {2} , m _ {g})} \\ \end{array}
$$
Therefore, if we have  $m_1, \ell(\pmb{x}_1, m_1)$ , and  $\pmb{r} = \frac{f(\pmb{x}_1, m_1)}{l(\pmb{x}_1, m_1)}$ , we can update the softmax statistics for a new block of entries  $\pmb{x}_2$  by following the following steps:
1. Compute new global max:  $m_g = \max (m_1,\max (\pmb {x}_2))$
2. Compute  $f(\pmb{x}_2, m_g) = e^{\pmb{x}_2 - m_g}$
3. Compute  $\ell (\pmb {x}_2,m_g) = \sum_if(\pmb {x}_2,m_g)$
4. Compute new  $\ell (\pmb {x},m_g)\colon \ell (\pmb {x},m_g) = e^{m_1 - m_g}\ell (\pmb {x}_1,m_1) + \ell (\pmb {x}_2,m_g)$
5. Correct running softmax result:  $r = r\frac{\ell(\pmb{x}_1,m_1)}{\ell(\pmb{x},m_g)}$
6. Add contribution from  $\pmb{x}_2$  to  $\pmb{r}$ :  $\pmb{r} = \pmb{r} + \frac{f(\pmb{x}_2, m_g)}{\ell(\pmb{x}, m_g)}$ .
In case all the keys are masked for a given query, we would have  $\max(x_2) = -\infty$ , given that the first  $m_1$  is initialized to  $-\infty$  as well (see Alg. 1 and Alg. 2) the fourth and second steps above would be undefined and result in NaN values. We solve the problem by replacing  $-\infty$  values in  $m_g$  by 0s when doing those two steps. Another potential issue is in step three: when a query has no matching key in  $x_2$  then  $\ell(x_2, m_g)$  is now 0, which generate  $+\infty$  in step five. This is an issue as we process keys by blocks, and if there are no keys for a query in the current block, we might find matching keys in following blocks. Adding  $\infty$  values to  $r$  would prevent us to accumulate statistics later on. To get our desired behaviour and have 0s when queries have no matching key we replace  $\infty$  values by 1s in  $\ell(x, m_g)$  during steps five and six. We summarize those steps in Alg. 3.
Algorithm 3 UPDATE_STATISTICS method
Require: Vector  $\ell \in \mathbb{R}^N$  , vector  $\pmb {m}\in \mathbb{R}^{N}$  , matrix of masked inner products qk  $\in \mathbb{R}^{N\times N}$  output buffer  $o\in \mathbb{R}^{B_m\times d}$  , block of values  $\pmb {v}\in \mathbb{R}^{B_m\times d}$    
1: Compute new global max (step 1) mnew  $\leftarrow$  max(rowmax(qk),m)   
2:Replace  $-\infty$  by 0:  $\hat{m}_{new}\gets \mathrm{WHERE}(m_{new} == -\infty ,0,m_{new})$    
3: Compute step 2:  $\pmb {p}\gets e^{\pmb {q}\pmb{k} - \hat{\pmb{m}}_{new}[:,None]}$  # row-wise subtraction   
4:Compute step 3:  $\ell_2\gets$  rowsum(p)   
5: Compute step 4:  $\ell_{new}\gets e^{m - \hat{m}}\ell +\ell_{2}$    
6:Compute  $z\gets \frac{1}{\ell_{new}}$    
7:Replace  $\infty$  by 1:  $z\gets$  WHERE  $(z = = \infty ,1,z)$    
8:Update  $p\gets p\times z[:,None]$  # row-wise multiplication   
9:Correct running softmax output (step 5)  $o\gets o\times (\ell z)[:,None]$  # row-wise multiplication   
10:Add contribution from current block (step 6)  $o\gets o + p.v$    
11:return  $\ell_{new},m_{new},o$
Hyperparameters. We extend the implementation of FlashAttention available in the Triton tutorial. In our benchmarks, we use a batch size  $B = 4$ , 48 heads of 64 dimensions each.
Effect of imbalanced dropping/Hashing. Our QK-sparse pattern relies on building smaller "compact" tensors containing the keys and queries that have not been dropped. In the worst case, one head can have many queries or keys dropped, while another head might not drop any. The batching would then force us to retain the entire sequence and the compact tensors would be of the same shape as the original tensors. The runtime would increase as we now have to iterate over all the blocks of keys for one of the heads. Due to the structure of FlashAttention which loops over blocks of keys, in parallel for each block of queries, the runtime is tied to the longest sequence of keys to process. A smart QK-dropping mechanism on top of our QK-sparse pattern should be cautious of dropping imbalances across heads. Concerning the Hash-sparse pattern, having all the keys and queries falling in the same
bucket would equate to the normal FlashAttention. Therefore unbalanced buckets would increase the runtime. Interestingly, in our experiments on language modeling, we do not face imbalance problems and even notice a speedup early in the training, see Fig. 14.(a) in App. C.
B.2 Runtimes in a Vacuum
Baseline implementation. We use Pytorch's FlashAttention implementation provided by the torch(nn ).functional .scaled .dot .product .attention function. To ensure fairness, we assume that all benchmarked functions receive a tensor  $Q$  of shape (BATCH, CTX_Q, H, D_HEAD), and tensors  $\pmb{K},\pmb{V}$  of shapes (BATCH, CTX_KV, H, D_HEAD), where BATCH is the batch size, CTX_Q is the number of queries, CTX_KV is the number of keys and values, H is the number of heads, and D_HEAD is number of dimensions per head. For this reason the only pre and post processing steps required are transposing the input and output tensors.
Listing 1: pytorch_fullflashattention function applying the FlashAttention algorithm on the entire sequence (no sparsity).
defpytorch_full_FLASHattention(q,k,v):   
BATCH,NCTX,H,D_HEAD  $\equiv$  q.shape   
q=q.transpose(1,2)# (BATCH,H,NCTX_Q,D_HEAD)   
k=k.transpose(1,2)# (BATCH,H,NCTX_KV,D_HEAD)   
v=v.transpose(1,2)# (BATCH,H,NCTX_KV,D_HEAD)   
y  $=$  torch(nnfunctional.scaled.dot_product attention(q,k,v,dropout_p=0.0, attn_mask=None，is_causal=True)   
return y.transpose(1,2).contiguous()
Our proposed interface. We propose the following interface to orchestrate between the Hashsparse and the QK-sparse implementations:
Listing 2: dynamic_sparseattention interface
def dynamicSparseattention(q,k,v,qidx,kidx,sm_scale=None, sparsity_mode='hash'):   
3   
q: query tensor of shape (BATCH,NCTX_Q,H,D_HEAD)   
k: key tensor of shape (BATCH,NCTX_KV,H,D_HEAD)   
v: value tensor of shape (BATCH,NCTX_KV,H,D_HEAD)   
qidx: tensor of shape (BATCH,NCTX_Q,H) for each sequence in the batch, for each query in the sequence, for each head, represents either the bucket index if sparsity_mode  $= =$  'hash' or the whether to keep that given head if sparsity_mode  $= =$  'qk'. The type should be torch.int32 if sparsity_mode  $= =$  'hash' and torch.float if sparsity_mode  $= =$  'qk'.   
k idx: tensor of shape (BATCH,NCTX_KV,H) for each sequence in the batch, for each key in the sequence, for each head, represents either the bucket index if sparsity_mode  $= =$  'hash' or the whether to keep that given head if sparsity_mode  $= =$  'qk'. The type should be torch.int32 if sparsity_mode  $= =$  'hash' and torch.float if sparsity_mode  $= =$  'qk'.   
sm_scale: normalization constant, 1/sqrt(D_HEAD) unless specified   
sparsity_mode: 'hash' to select the hash-sparse implementation and 'qk' for the qk-sparse implementation   
if sm_scale is None: sm_scale  $\equiv$  1.0 / math.sqrt(q.size(-1))   
if sparsity_mode  $= =$  'hash': return hashSparse attention(q,k,v,q_hash=qidx,k_hash=kidx, sm_scale  $\equiv$  sm_scale)   
elif sparsity_mode  $= =$  'qk': return qksparseattention(q,k,v,q_keep=qidx,k_keep=kidx, sm_scale  $\equiv$  sm_scale)   
else: raise KeyError(f"Unknown sparsity_mode: {'sparsity_mode'}, should be in ['hash', 'qk'])
Pre & post processing steps for Hash-sparse. In addition to having to transpose the  $Q, K, V$  tensors. The preprocessing steps consist in re-ordering the  $Q, K$  and  $V$  tensors based on bucket indices in q_hash and k_hash. We keep track of the original position of the queries, keys and values by storing the indices given by the sorting operations. Importantly, we use stable sorts to ensure the queries, keys and values are sorted within each bucket. The following code is showing how we implemented all those steps using Pytorch:
Listing 3: hash_sparseattention function showing pre and post processing steps for the Hash-sparse algorithm.
def hash_sparseattention(q,k,v,q_hash,k_hash,sm_scale): assert q_hash.dtype  $= =$  torch.int32 and k_hash.dtype  $= =$  torch.int32   
4 BATCH,NCTX_Q,H,D_HEAD  $= \mathbb{Q}$  .shape   
5   
6 q=q.transpose(1,2)# (BATCH,H,NCTX_Q,D_HEAD)   
7 k=k.transpose(1,2)# (BATCH,H,NCTX_KV,D_HEAD)   
8 v=v.transpose(1,2)# (BATCH,H,NCTX_KV,D_HEAD)   
9 q_hash=q_hash.transpose(1,2).contiguous() # (BATCH,H,NCTX_Q)   
10 k_hash=k_hash.transpose(1,2).contiguous() # (BATCH,H,NCTX_KV)   
11   
12 #Re-order the queries,key,values according q_hash and k_hash   
13 q_hash=q_hash.sort(dim=-1,stable=True)#q_hash.shape  $=$  (BATCH,H,NCTX_Q), stable sort to keep time ordering within a bucket   
14 k_hash=k_hash.sort(dim=-1,stable=True)#k_hash.shape  $=$  (BATCH,H,NCTX_KV)   
15   
16 qidx=q_hashindices   
17 kidx=k_hashindices   
18   
19 q_hash=q_hash.values   
20 k_hash=k_hash.values   
21   
22 qidx_EXTENDED  $= \mathbb{Q}$  idx unsqueeze(-1).expand_as(q)   
23 kidx_EXTENDED  $= \mathbb{K}$  idx unsqueeze(-1).expand_as(k)   
24   
25 q=torch.gather(q, dim=-2,index=qidxextended).contiguous()   
26 k=torch.gather(k, dim=-2,index=kidxextended).contiguous()   
27 v=torch.gather(v, dim=-2,index=kidxextended).contiguous()   
28   
29 y=hash;sparse attention_kernel(q,k,v,qidx,kidx,q_hash,k_hash,sm_scale)   
30 y=torch.zeros((BATCH,H,NCTX_Q,D_HEAD),dtype=q.dtype, device=qdevice).scatter(dim=2,index=qidxextended, src=y).transpose(1,2).contiguous()   
31 return y
Pre & post processing steps for QK-sparse. In addition to having to transpose the  $Q, K, V$  tensors. The preprocessing steps consist in removing dropped keys, values and queries from  $K$ ,  $V$  and  $Q$ . The sorting operations need to be stable to keep the original time ordering within the remaining keys and queries. Moreover, the index tensor has to be padded so our kernel can rely on those indices to compute which tile it should and shouldn't compute.
Listing 4: qk_sparseattention function showing pre and post processing steps for the QK-sparse algorithm.
def compact(keep_tensord, x, index=None):   
    '''Build a compact representation of x   
    Keyword arguments:   
    x: input tensor to compact, x.shape = (BATCH, NCTX, H, D_HEAD)   
    keep_tensord: float tensor of shape (BATCH, NCTX, H) containing a 1 when the head is kept, else 0   
    '''   
    BATCH, T, H, D_HEAD = x.shape   
    if index is None: with torch.no_grad(): indices_per_head = keep_tensord sum(dim=-2) buffer_size = indices_per_head.max().int() # first sum computes the num of non-killed elem per head, we take to max of that # sorting: it is very important that the sorting is stable, else we cannot use causal masking sorted = keep_tensord.sort(dim=-2, descending=True, stable=True) # sorted indices.shape == (BATCH x T x H), now sorted over sequence T index = sorted indices[(:, :buffer_size,:] # (BATCH x buffer_size x H) expand indices to cover all the dimensions for each heads
15 else:   
16 indices_per_head  $=$  None   
17 compact_x  $\equiv$  x.gather(dim=-3,index=index.unsqueeze(-1).expand(-1,-1,-1,D_HEAD)) # (BATCH x buffer_size x H x D_HEAD) / expand indices to cover all the dimensions for each heads   
18 return compact_x, index, indices_per_head   
19   
20   
21 @torch.no_grad()   
22 def pad_index(index, indices_per_head, padidx=-1):   
23 "" Pad the index tensor to comply with the kernel, returns a copy.   
24 Keyword arguments:   
25 index: original index tensor to pad given by 'compact', index.shape  $=$  (BATCH, buffer_size,H). For each batch and timestep, represents the head idx it's originating from.   
26 indices_per_head: of shape (BATCH,H), for each head, contains how many indices have not been dropped.   
27 ""   
28 BATCH, buffer_size, H = index.shape   
29 index_copy  $=$  torch.clone(index).type(torch.int32)   
30 mask  $=$  torch.arange(buffer_size, device=index_device).view(1,-1,1).expand(BATCH,buffer_size,H)> indices_per_head.view(BATCH,1,-1)   
31 index_copy[mask]  $=$  padidx   
32 return index_copy   
33   
34   
def qk_sparseattention(q,k,v,q_keep,k_keep,sm_scale): assert q_keep.dtype  $= =$  torch.float and k_keep.dtype  $= =$  torch.float   
35   
36   
37   
38   
39   
40 # Building compact representations   
41 q_c,qidx,iiph_q  $=$  compact(q_keep,q)#q_c.shape  $=$  (BATCH,compact_NCTX_Q,H)   
42 k_c,kidx,iiph_k  $=$  compact(k_keep,k)#k_c.shape  $=$  (BATCH,compact_NCTX_KV,H)   
43 v_c,_  $=$  compact(k_keep,v,index=kidx)#v_c.shape  $=$  (BATCH, compact_NCTX_KV,H)   
44 qidxffseted  $\equiv$  pad_index(qidx,iiph_q, padidx=-1)#(B,compact_NCTX_Q,H)   
45 kidxffseted  $\equiv$  pad_index(kidx,iiph_k, padidx=1e9)#(B,compact_NCTX_KV,H)   
46   
47 # We need to transpose everything   
48 q_c=q_c.transpose(1,2).contiguous() # (BATCH,H,compact_NCTX_Q,D_HEAD)   
49 k_c=k_c.transpose(1,2).contiguous() # (BATCH,H,compact_NCTX_KV,D_HEAD)   
50 v_c=v_c.transpose(1,2).contiguous() # (BATCH,H,compact_NCTX_KV,D_HEAD)   
51 qidxffseted  $\equiv$  qidxffseted.transpose(1,2).contiguous() # (BATCH,H, compact_NCTX_Q)   
52 kidxffseted  $\equiv$  kidxffseted.transpose(1,2).contiguous() # (BATCH,H, compact_NCTX_KV)   
53   
54 y_c=qk_sparse attentive_kernel(q_c,k_c,v_c,qidxffseted,kidxffseted, sm_scale).transpose(1,2)   
55 y  $=$  torch.zeros_like(q).scatter(dim=1, index=qidx.long().view(BATCH,-1,H,1).expand(BATCH,-1,H,D_HEAD),src=y_c)   
56 return y
B.3 Sequence Modeling Experiments
Language modeling on OpenWebText2. Our implementation is based on NanoGPT (github.com/karpathy/nanoGPT). We use the AdamW optimizer (Loshchilov & Hutter, 2019). We used bffloat16 and NVIDIA A100-40GB GPUs for all our experiments. Here is a list of hyperparameters shared by all our language models (F-LM, H-LM, and D-LM):
Weight-decay: 0.1
- Depth (number of transformer blocks): 12
Number of heads: 12
- Dropout: 0.0
- Learning rate: 0.001
- Percentage of iterations for warmup:  $2\%$ . We use a cosine learning rate scheduler.
Adam beta1: 0.9
Adam beta2: 0.95
- Tokenizer: We use the GPT2 tokenizer provided by the tiktoken library (github.com/openai/tiktoken).
- Hidden dimensions: 768
- Dimensions per head: 64
Sequential MNIST and enwik8. For the comparisons with Reformer, we use a standard GPT2 (Radford et al., 2019) implementation. For the language modeling on enwik8, the Transformer has 12 blocks with 768 hidden dimensions, 8 attention heads, and 64 dimensions per head. Dropout is set to 0.1 and the batch size is 8 with 2 gradient accumulation steps. The sequence length is 4096. For autoregressive image generation on MNIST, we use a smaller model with 8 transformer blocks and a hidden dimension of 256. Dropout is set to 0.1 and the batch size is 10. We train for 25 epochs.
C Additional Details and Analysis
Quadratic computational cost of attention in transformers. In Fig. 9 we show the runtime (forward + backward) of a transformer language model as a function of the sequence length. We separate the time taken by the attention operations from the time taken by the rest of the model. We see how the attention computation ends up dominating the runtime as the sequence length increases.
Figure 9: Quadratic computational cost of self attention dominates for larger sequences. We measure the forward and backward runtimes for a 12 layers transformer (see App. B.3 for implementation details). Except for the attention operations, the rest of the transformer runtime grows linearly with the sequence length.
Additional runtimes performances in a vacuum. In Fig. 10 and Fig. 11 we show the runtime details for the forward and backward methods separately for respectively the Hash-sparse and QK-sparse methods. We also measure runtimes of the forward and backward passes when we assume the pre and post-processing steps are free, see Fig. 12.
Linear QK-dropping scheduler. In the main paper we show results dropping keys and queries at random with a fixed pre-defined probability. In an additional experiment we start by dropping  $80\%$  of keys and queries at random and linearly decay this probability to  $20\%$ . Our intuition is earlier iterations just aim to learn contextual cues which are very redundant (and therefore quite immune to random dropping) before requiring more fine-grained representations. In Fig. 13 we show and analyse the results of that experiment.
H-LM models speeding up during training. The speed of Hash-sparse attention is conditioned on the distribution of bucket indices over keys and queries—e.g. if all the keys and queries were to fall in the same bucket then there would be no speedup over FlashAttention. Interestingly, when training on real data such as OpenWebText2, we observe our Hash-sparse based models are speeding up during training. In Fig. 14 we plot the number of iterations reached after  $x$  hours of training (normalizing the time by the number of GPUs used for training). In Fig. 14.(a) we see a speedup for H-LM models early in the training.
Different bucket sizes  $nb$  for H-LM. For sequences of 8192 tokens, we show the influence of increasing the number of buckets  $nb$  in Fig. 15. As  $nb$  increases the runtime decreases.
(a) Forward
(b) Backward
Figure 10: Forward and Backward runtimes in a vacuum for our Hash-sparse method. Fig.(a) Forward runtimes for our Hash-sparse method. Fig.(b): Backward runtimes for our Hash-sparse method.
(a) Forward
(b) Backward
Figure 11: Forward and Backward runtimes in a vacuum for our QK-sparse method. Fig.(a) Forward runtimes for our QK-sparse method. Fig.(b): Backward runtimes for our QK-sparse method.
Training H-LMs for more iterations. In Fig. 8 and Fig. 6 we show results of language models trained on OpenText2 for 15k iterations. To verify whether our finding are consistent when you train for more iterations we also try training for 50k iterations. For sequences of 8192 tokens, using the same hyperparameters described in App. B.3, we show in Fig. 16 that our findings for our Hash-sparse based models do hold when training for more iterations—we match the perplexity per iterations of the baseline model using FlashAttention over the entire sequence while being significantly faster.
Visualizing the low coverage of Reformer. In Fig. 17 we show the attention matrices for two different heads and show how the coverage—the percentage of key-query interactions actually computed vs. what should be computed according to the hash indices computed for keys and queries—of the Reformer LSH algorithm can be low.
D Limitations and Societal Impacts
Limitations. The aim of our work is to develop a method for efficiently computing attention with several sparsity structures. We don't focus on developing the best method for sparsification, although, for example, we improve the hashing-based mechanism. Moreover, as already explained, on very small sequences we incur in some constant overhead which limits our gains.
(a) QK-sparse
(b) Hash-sparse
Figure 12: Forward+Backward runtimes in a vacuum when we assume the preprocessing is free. Fig.(a) Forward+backward runtimes for our QK-sparse method when the pre and postprocessing steps are free. While the runtimes for large sequences stays relatively the same compared to Fig. 7, the offset for short sequences is now much smaller. Fig.(b): Forward+backward runtimes for our Hash-sparse method when the pre and postprocessing steps are free. While the runtimes for large sequences stays relatively the same compared to Fig. 3, the offset for short sequences is now smaller. Interestingly, unlike for the QK-sparse method, the offset for smaller sequences stays significant. We believe this might be due to the influence of the block size used (128): when the sequence length is not large enough in comparison to the block size, the block structure of the hash-sparse attention matrix cannot be efficiently leveraged and the number of tiles processed by the Hash-sparse method is larger than that of the FlashAttention method. In contrast, for the QK-sparse method the good tiles all start from 0 (we know the first tile(s) will be efficiently packed).
(c) Speed gains
Figure 13: Training Language Models (LM) on OpenWebText2 using random Query/Key dropping based sparsity (D-LM) or FlashAttention over the entire sequence (F-LM). Results are the same as in Fig 8 except for the addition of the D-LM model with linear decay of  $s$  from  $80\%$  to  $20\%$ . In Fig.(b) and Fig.(c), we observe the additional D-LM model in yellow is slower than other D-LM models we experimented with, but, as seen in Fig.(a), reaches a slightly better perplexity.
(a) Time comparison H-LM vs. F-LM
(b) Time comparison D-LM vs. F-LM