SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Figure 14: Training speed as a function of time for H-LM, D-LM and F-LM. All models trained on sequences of 8192 tokens. In Fig.(b) we observe D-LM models having a fixed speed throughout the experiments except for the one corresponding to the linear decaying of the sparsity ratio  $s$  which slows down during training. In Fig.(a) H-LM models are speeding up early during training and then seem to keep a constant speed.
Figure 15: Influence of the number of buckets  $nb$  on H-LM training speed for  $T = 8192$  and  $T = 16384$ . In accordance with Fig. 3 increasing  $nb$  speeds up the training but the speed gain from doubling  $nb$  decreases as  $nb$  is increasing.
(a) Iter to reach perplexity
(b) Speed gains
Figure 16: Comparison between H-LM and F-LM when training for  $50\mathrm{k}$  iterations on sequences of 8192 tokens. Those curves are averaged over two seeds. In Fig.(a) we zoom in on the perplexity per iteration for the baseline F-LM and two of our H-LM models using our Hash-sparse attention with 8 and 16 buckets. We observe that our methods are matching the perplexity of the F-LM model which uses FlashAttention over the entire sequence. Fig.(b): We match the perplexity but gain in speed, our methods are  $1.9\times$  and  $2.1\times$  faster than the baseline for respectively 8 and 16 buckets.
Figure 17: Attention pattern of Reformer. The two images show the attention pattern for the Reformer LSH attention after reordering queries and keys according to the bucket. The red squares are the positions for which attention is computed in the Reformer with bucket size  $= 32$ . The back squares are positions for which the queries and keys map to the same hash bucket. Queries are sampled from a normal distribution with  $\mu = 3$  and  $\sqrt{\sigma} = 5$ . Each image refers to a different attention head. Black regions which are not covered by red tiles are key-query interactions which should be computed but are missed by the Reformer LSH attention.
Societal impacts. The attention mechanism is central to Large Language Models (LLMs). Moreover, efficient attention mechanisms can make these models more powerful, by making them faster to train and by increasing their context length. The social impact and risks associated with our work, therefore, are included in the risks associated with the deployment of such systems (Bender et al., 2021; Weidinger et al., 2021).