SlowGuess/NeurIPS_2023_OCR huggingface dataset:
arXiv Math We use the cleaned arXiv math subset of proof-pile  $^2$  dataset. The dataset contains all the  $\mathrm{TEX}$  files for submissions under math category in arXiv. Files with encodings other than utf-8/16/32 or latin1 were removed. Additionally, the dataset curation includes certain heuristic to only keep well-structured papers. An example is removing files that do not contain any chapter, sections, or subsections. The final training dataset contains around 5.6B tokens.
PG-19 PG-19 is a large dataset (3.7B tokens in the training dataset) of English books that were published before 1919 and were retrieved from the Project Gutenberg archive. This dataset is widely used for evaluating model capabilities in utilizing long-range token interactions [35, 41].
C Number of Unique Retrieved Blocks
We feed the validation set of PG19 in segments of at length 2048 (in chunks of 250 tokens). During this process, we retrieved the top  $k = 4$  blocks from the cache and kept track of the number of unique blocks retrieved in the penultimate chunk. Note that we count the same block retrieved by two different tokens only once. To ensure consistent statistics, we utilize the penultimate chunk instead of the last one, as the last chunk may have a lower number of tokens that could disrupt the analysis. We record this number separately for each batch element and each layer and plot the distribution at different retrieval flexibility levels in Figure 5. In the penultimate chunk, the cache contained 35 blocks. As depicted in the figure, at the most flexible level, it can be observed that in many cases, all the blocks are accessed at least once at each layer, as indicated by the spike in the last bin. Interestingly, the additional spike in the flexible retrieval corresponds to a very low number of blocks
$^{2}$ https://github.com/zhangir-azerbayev/proof-pile/
Figure 5: The distribution of the number of unique retrieved blocks by the 255 tokens (including the landmark tokens) in the penultimate chunk when evaluating over 2048 token inputs from PG19.
being retrieved and is attributed to the retrieval pattern in the initial layer. This finding aligns with previous research suggesting a locality of attention in the earlier layers [34].
Moreover, the results demonstrate a significant decrease in the number of unique retrieved blocks when only allowing the set of retrieved blocks to vary across heads. Although there are 8 heads, each retrieving 4 blocks from the cache, theoretically allowing access to 32 different blocks, we typically observe the number of unique blocks to be below 10. While this reduction does come at the cost of perplexity (as observed in Table 2), it can significantly improve performance by facilitating offloading to slower devices through a reduction in the required bandwidth for loading retrieved blocks from the cache. It is worth noting that one approach to mitigate the negative impact of reduced flexibility on perplexity is to increase the number of retrieved blocks. This can be observed by comparing the results between retrieving the top  $k = 2$  blocks at the highest level of flexibility and retrieving the top  $k = 4$  blocks with reduced flexibility, as shown in Table 2. By retrieving more blocks, we can partially compensate for the reduced flexibility and its effect on perplexity.
Finally, when allowing the set of retrieved blocks to vary across tokens but not heads, it is still possible to observe improvement in the number of unique retrieved blocks which can reduce the bandwidth needed to load blocks from the cache. At the same time, better perplexities can be achieved in this setting.
D Context Miss Token
In this section, we demonstrate how to create additional hierarchy in the landmark structure by simply changing the grouping. In particular, we use a new grouping scheme to train a new special token, called context miss token (CMT). CMT is always placed at the beginning of the input (at the special position  $-1$ ) and is used to signal the need to retrieve landmarked blocks from the memory. To train this token, we change the grouping scheme so that attention to some of the blocks is regulated by CMT, similar to how landmark tokens act as gateways to the tokens in their respective blocks. In particular, let us denote  $L_{\mathrm{CMT}}$  as the set of landmark tokens that are controlled by CMT. Then, using the same notation as in Section 3.1, we use the following grouping during training (For brevity, in the below formulation we assume that the negative of all preceding conditions hold for each case):
$$
\mathbf {G} _ {i, j} := \left\{ \begin{array}{l l} p _ {i} & j = - 1 \\ p _ {j} & p _ {j} \neq j \\ - 1 & p _ {i} = j \\ p _ {i} & j \notin L _ {\mathrm {C M T}} \\ - 2 & j \in L _ {\mathrm {C M T}} \end{array} \right. \quad \triangleright \text {p l a c i n g C M T i n t h e i - t h t o k e n ' s g r o u p .} \quad \triangleright \text {p l a c i n g n o r m a l t o k e n s i n t h e i r o w b l o c k s .} \quad \triangleright \text {i n g o r i n g c u r r e n t b l o c k ' s l a n d m a r k t o k e n .} \quad \triangleright \text {p l a c i n g f r e e l a n d m a r k s i n t h e i - t h t o k e n ' s g r o u p .} \quad \triangleright \text {p l a c i n g C M T - c o n t r o l l e d l a n d m a r k s i n a s e p a r a t e g r o u p .} \tag {5}
$$
Using the above grouping the attention weights can be calculated as:
$$
\operatorname {A t t W e i g h t} (\mathbf {Q}, \mathbf {K}) _ {i, j} := \left\{ \begin{array}{l l} 0 & p _ {j} = j \vee j = - 1 \\ \mathbf {S} _ {i, j} & \mathbf {G} _ {i, j} = \mathbf {G} _ {i, i} \\ \mathbf {S} _ {i, j} \cdot \mathbf {S} _ {i, p _ {j}} & j \notin L _ {\mathrm {C M T}} \\ \mathbf {S} _ {i, j} \cdot \mathbf {S} _ {i, p _ {j}} \cdot \mathbf {S} _ {i, - 1} & j \in L _ {\mathrm {C M T}} \end{array} . \right. \tag {6}
$$
To build intuition, we note that using an empty  $L_{\mathrm{CMT}}$ , almost recovers the original formulation (with the exception of the existence of CMT token which can absorb part of the attention score). Note that assuming  $L_{\mathrm{CMT}}$  is not empty the above formulation ensures that the total sum of attention weights is equal to one. During training, we suggest to choose  $L_{\mathrm{CMT}}$  randomly by selecting each landmark token with probability  $\mathbb{P}_{\mathrm{CMT}} = 0.5$ .
Intuitively, the above formulation should teach the network to attend to the CMT token whenever it does not have enough information within the free landmark tokens (landmark tokens that are not controlled by CMT). Furthermore, we note that since CMT is the first token, it can not attend to any token. As such, its representation in each layer does not depend on the input. As such, the token acts as a beacon for signaling the context lacks enough information. That is why we draw a parallel to the notion of a cache miss, naming it a context miss token. We note that CMT can be used to create different levels of landmark block storage, going to a lower (more granular) level whenever it is activated at each level, similar to a cache miss signal. We leave a thorough investigation of such designs to future work.
To evaluate our method, we follow the configuration used in Section 4.1 and train a model with CMT on PG19. At inference we set  $L_{\mathrm{CMT}}$  to be the set of landmarks that are in the memory so the in-context landmarks are not controlled by the CMT. Using this setting, the CMT can be used to check whether a retrieval is necessary or not. For simplicity, in our implementation we always perform the retrieval but set  $\mathbf{S}_{i,-1}$  to 0 if it is below a cutoff threshold. This exactly emulates not performing the retrieval at all. We report the perplexity on PG19 with different cutoff thresholds and the ratio of CMT scores below the cutoff threshold that were subsequently dropped in Table 3. Our results show that it is possible to use CMT to significantly reduce the number of retrieval calls while still utilizing the memory. In particular, around  $50\%$  of the retrievals can be dropped with minor effect on perplexity. Note that the difference observed between training without CMT (baseline) and training with CMT (no cutoff) is reasonable since the model is learning a harder task. We conjecture that this difference can be alleviated by simply training the model longer.
E Positional Augmentation
In our main method, we use a special position mapping scheme (See Section 3.2.1), called Stingy mapping. The need for Stingy mapping is due to the model's inability to extrapolate to positions that were not observed during training [36]. Existing methods allow the model to handle longer inputs by manually dampening or completely limiting attention to far-away tokens [27, 36]. These techniques are not applicable to our settings since one of our main goals is to allow attention to any token, including those that are very far. While development of a positional encoding scheme which can extrapolate is outside the scope of this work, we propose a possible solution and provide results of an early investigation. We leave a more thorough assessment and development of such schemes for future work.
Data augmentation has been used in various settings to allow the model to generalize to additional settings such as reflections of the same image [21]. We propose to apply augmentation on positional
Table 3: Performance on PG19 dataset with 2048 evaluation length and  $k = 4$  for different CMT cutoff thresholds. When the GroupedSoftmax for CMT is below the cutoff threshold, it is set to zero to emulate not performing a retrieval. The drop rate column shows the ratio of CMT scores below the cutoff threshold. Baseline refers to the model with landmarks but without CMT. The models are trained for 60K steps.
Cutoff	Perplexity	Drop Rate
Baseline	16.28	0%
0.0	16.38	0%
0.1	16.38	23%
0.3	16.43	57%
0.5	16.86	84%
1.0	19.49	100%
Figure 6: Comparison of perplexity on PG19 for different context lengths. The evaluation does not use the landmark cache and feeds the whole input in a single iteration to the model. Still, the landmark tokens are inserted every 50 tokens. The context length does not include the additionally inserted landmark tokens.
information in Transformers to allow them to extrapolate to longer contexts. In the standard positional encoding, the positions are increased by one at each token, leading to the tokens being assigned positions 1 to  $\ell_{\mathrm{seq}}$  where  $T$  is the length of the input. In particular, instead of assigning positions from 1 to  $T$ , where  $T$  is the length of the input, we propose to increase the positions of all subsequent tokens by a random integer between 1 and  $p_{\mathrm{jump}}$  after each landmark token. We refer to these increases as making positional jumps. When  $p_{\mathrm{jump}} = 1$ , no augmentation is applied and the standard positions are recovered.
To evaluate our proposal, we train the same model that we used in Section 4.1 on PG19 with  $p_{\mathrm{jump}} = 100$ . Since we use a context length of 512 for training, each input can has between 10 and 11 landmark tokens. Theoretically, this should allow the model to extrapolate to context lengths as long as  $1100 + 512 = 1612$ . We plot the performance of the model on different context lengths as well as the performance of the model with standard positional encoding (i.e.  $p_{\mathrm{jump}} = 1$ ) after 60K steps in Figure 6. Note that we do not use a cache in this section and feed the whole input at once.
We can see that using the augmentation, the model becomes capable of utilizing longer contexts. This is evident by the fact that we observe reduction in perplexity as we increase context lengths until 1400 tokens which is close to the theoretical estimate of model's extrapolation capacity. In contrast, the decreasing trend stops for the standard model before reaching 1024 tokens. We can also observe that when evaluating at training context lengths, i.e. 512 tokens, the performance when making positional jumps is lower than standard training. However, this can be expected and justified since the model is learning a harder task and therefore may require additional steps to reach the same performance.
Figure 7: Time needed to process 4870 tokens (excluding the landmarks) using Landmark Attention for different number of retrieved blocks  $k$  and different block sizes. The efficient implementation of landmark attention combined with FlashAttention is used. The retrieval flexibility is reduced so that all tokens in each chunk use the same set of retrieved blocks.
F Additional Extensions and Details
Masked Language Modeling In this work, we focus on next-word prediction tasks. In this section, we briefly discuss how to adapt our scheme for masked language modeling. Our discussion here is conceptual, and experimental verification of our proposal is left for future work. We note that the scheme described in Section 3.1 can be used to train a model with masked language modeling. However, certain changes need to be made to the inference method. Specifically, the input cannot be given to the model sequentially from left to right because each token should be able to access the subsequent tokens as well. Thus, we propose processing the entire input layer by layer. In each layer, the input is broken into chunks, and each chunk is processed separately, with each token retrieving the top- $k$  landmarked blocks from other chunks.
Combination with Flash Attention Flash attention computes the attention matrix in blocks [10] with the block size chosen to optimize memory accesses. Landmark Attention can be naturally combined with flash attention by choosing the frequency of the landmark tokens to be equal to the aforementioned block size. Note that in this case, the calculation of the GroupedSoftmax incurs negligible overhead (both in terms of memory and compute) because only the softmax values for the currently processed block and the landmark block need to be kept and adjusted at each step. Combination with more advanced extensions is also possible. For example, the landmark token can naturally be used to decide on dropping blocks when using block sparse flash attention. Finally, when reducing the retrieval flexibility and enforcing the same set of retrieved blocks for all tokens in the chunk, Flash Attention can be directly applied at inference as well. We implement this combination using Triton and show that it can be used to fine-tune LLaMA7B with 2048 context length (instead of the maximum 512 tokens possible without flash attention).
Trade-off between number of retrieved blocks and block size During inference, when using landmark attention, landmark tokens need to be also stored in memory which slightly increases the memory usage and compute time by a factor of  $\frac{1}{\text{block size}}$ . Thus increasing block size reduces memory usage while also reducing the time needed to find relevant blocks since there are less blocks. However, if the number of retrieved blocks  $k$  is kept constant, the length of the chunks have to become smaller so the longer retrieved blocks and the current chunk can fit into the model's maximum allowed context length, i.e. the train context length. Therefore, increasing the block size increases the number of chunks the input has to be broken into, which slows down the inference. This trade-off can be seen in the Figure 7, showing the inference time for different block sizes and different values of  $k$ .
G Offloading KV Cache to CPU
Previous research has proposed various techniques to reduce the size of the KV cache [1, 26, 32], which can be directly applied to address this bottleneck. However, using landmark attention allows to reduce memory usage further using a simpler approach: When using landmarks, it is possible to offload the cache to the CPU and only load back the retrieved blocks in GPU at each step, while keeping all landmarks in GPU memory. While for inference with very high context lengths the techniques proposed by previous research would still be necessary, offloading to CPU allows us to perform inference with LLaMA 7B with more than 32K tokens.
Although the aforementioned technique works well, it introduces significant CPU-GPU traffic, resulting in slow inference. To mitigate this issue, we limit the number of retrieved blocks by reducing retrieval flexibility, allowing the set of retrieved blocks to vary across heads but not across tokens. We use the same reduction mechanism as in Section 4.1. Specifically, we calculate the maximum score for each block across all tokens and select the top  $k = 5$  blocks for retrieval. Prior to selecting the maximum, softmax is applied to ensure all scores are on the same scale.
Using the above method, we evaluated the fine-tuned LLaMA with a context length of 32070 tokens (the exact number of tokens may vary slightly due to random generation, but the difference is less than 10 tokens). Across 50 randomly generated prompts, we achieved an encouraging accuracy of  $98\%$  in retrieving the pass key. This clearly demonstrates the model's capability to effectively operate at context lengths comparable to those supported by GPT-4 [25].