SlowGuess/NeurIPS_2023_OCR huggingface dataset:
When computing the attention scores to cache elements (both landmark and normal tokens), it is important to correctly incorporate positional information. The transformer model is sensitive to positional information. It is also intuitive that the model would rely on position information in some cases. For example tokens right after the last memory block do not have access to any context and are unable to select memory blocks based on semantics. Instead, they need to rely on positional information to select the last memory block.
Optimally, the tokens are encoded with their actual position index. However, a known flaw of Transformers is their inability to extrapolate to lengths not observed during training. Various methods proposed to alleviate this condition also do not fully resolve the problem unless they are combined with block attention which only allows attending to a window of tokens. We decide against using block attention since our main goal is to facilitate attending to large contexts. In Appendix E we propose an alteration of the training method which allows using the actual position index. However, to avoid overlapping changes, we use the following approximation for most of our experiments.
We allocate a segment with length  $(k + 1) \cdot (\ell_{\mathrm{block}} + 1)$  in the beginning of the context. We index the current chunk starting after this segment. For the retrieved blocks, we map the index for any of the latest  $k$  blocks to the corresponding place within the last  $k$  blocks in the allocated prefix segment. Other blocks in the cache have their position mapped to the first block in the allocated prefix segment. Once we decide which blocks to retrieve, we map those among the latest  $k$  blocks to the right of pre-allocated segment and map the rest of the blocks to the left of the pre-allocated segment, while respecting the order of the blocks. We call this scheme stingy position mapping which is further illustrated in Figure 2.
Note that we found out that when  $k = 1$  mapping memory blocks to a segment of at least 2 blocks is crucial. Using only a single block, all memory blocks are mapped to the same index. However, as we discussed at the beginning, it is essential to at least retrieve the last memory block solely based on position information which is impossible unless this block is mapped to a different index. While it is possible that the importance of pre-allocating the additional block decreases as  $k$  grows, we adapt this scheme so that the attention would be at least as flexible of simply keeping the last  $k$  blocks.
We point out that the above approximation relies on the ability to add position information when performing the retrieval. In our experiments, we use Transformer models with Rotary positional encoding [33] which adds the position information to the key and query vectors just before computing
the attention. Thus, we can store the key vectors without position information in the cache and add the position information when performing the retrieval according to the following scheme.
3.3 Memory & Computation
During training, our method has only a negligible overhead due to the computation of GroupedSoftmax. In particular, our method does not require maintaining a cache of previous values at training time. Furthermore, we decouple the training context length from the inference context length since it is possible to perform inference at any context length using the method described in Section 3.2 regardless of the train context length. As such, when comparing training time in terms of inference context length, we offer constant training time  $(\mathcal{O}(1))$  whereas training time for a standard transformer scales quadratically with the operational (inference) context length.
Furthermore, in comparison with standard attention over the whole input, our method reduces inference time by computing the attention score over a smaller set of token pairs. For instance, during auto-regressive generation, where tokens are generated sequentially by using the previously generated token as input for obtaining the next one, the traditional approach involves computing attention across all preceding tokens when generating the next token. In contrast, our method allows computing attention solely over landmark tokens, followed by attention computation over the retrieved blocks. Given the constant size of the blocks, the cost of computing the attention over the retrieved blocks remains constant regardless of the total context length. While the cost of finding the most relevant landmark tokens increases linearly with the context length, the rate of increase is 1 every  $\ell_{\mathrm{block}} + 1$  tokens. This immediately reduces the number of operations by a factor of block length  $\ell_{\mathrm{block}}$ . For example, when using  $\ell_{\mathrm{block}} = 50$  as in our experiments, this can lead to a  $50\mathrm{x}$  boost. Importantly, the same reduction can be obtained in terms of memory. In the standard Transformer, a cache of all previous key and values (KV-cache) is needed to perform the generation efficiently. In contrast, we only need immediate access to the landmark tokens and can offload the blocks to slow memory (e.g. CPU), loading only the retrieved blocks. The reduction in memory and compute can be further improved if the search for retrieving blocks is performed using more advanced data structures such as FAISS [15].
It is worth noting that the additional computational overhead introduced by performing two matrix multiplications (one for block selection and another for attention to the retrieved blocks) instead of a single matrix multiplication in the standard setting becomes relatively negligible, especially when dealing with larger inputs.
Finally, we point out that our method can be naturally combined with flash attention [10], reducing the overhead further. We discuss this in Appendix F. In this work, to reduce the complexity and allow flexibility in experiments, we use a high-level implementation (not combined with Flash Attention). However, we also publish a version of efficient implementation in Triton [37].
4.1 Language Modeling
We first evaluate the efficacy of retrieving earlier blocks on two language modeling tasks which can be expected to have long-range token interactions: English language books (PG-19) [29] (3.7B tokens), and math papers from arXiv (5.6B tokens). We provide additional details about the datasets in Appendix B. Our results show that models trained with landmark tokens can retrieve relevant blocks, obtaining comparable perplexity as a Transformer-XL while reducing FLOPs. In contrast with Transformer-XL, using our method, the information retrieval is interpretable since the exact tokens attended to by the model can be identified by looking at the attention scores or looking at the set of retrieved blocks. Particularly, it is possible to understand which parts of the text was recovered to generate a certain answer, which can for example be useful to remove inaccurate information. Our results also demonstrate that using the inference mechanism described in Section 3.2, our models can be used at much longer context than the one used for training.
Model & Training. We use a GPT-2 [28]-like architecture: a 12-layer decoder-only transformer with 8 heads in each layer, each with dimension 128 (embedding dimension 1024), and hidden FFN size of 4096. We trained our model using AdamW [22] with  $\beta_{1} = 0.9$  and  $\beta_{2} = 0.95$ . We applied weight decay with factor 0.001. We used base learning rate 0.002 for all our experiments with a warmup stage that was  $2\%$  of the whole training and applied a cosine scheduler with minimum (final)
Table 1: Performance of different training and inference settings in terms of language modeling perplexity. The column XL cache shows the size of the XL cache available both during training and inference which was only used when training Transformer-XL[9]. When using landmarks, the column "☑ Blocks" shows the maximum number of blocks stored in memory. Each block contains  $\ell_{\mathrm{block}} = 50$  normal tokens and one landmark token. Due to computation limitations we only report results for Transformer-XL on PG-19 as this method takes longer to train in our implementation.
Eval. Length	\( \ell_{local} \)	XL cache	\( ♀Blocks \)	k	Attention Size	PG19	arXiv
512	512	None	None	-	512	16.12	4.01	Baseline
360	None	None	-	360	16.76	4.31
250	None	10	2	360	16.23	4.01	Ours
2048	256	256	None	-	512	14.72	-	[9]
250	None	40	2	360	15.14	3.43	Ours
350	None	40	2	460	15.07	3.41
300	None	40	3	460	14.97	3.36
250	None	20	4	460	15.02	3.37
250	None	40	4	460	14.92	3.35
4096	256	256	None	-	512	14.55	-	[9]
250	None	40	4	460	14.79	3.19	Ours
250	None	80	2	370	15.00	3.29
250	None	80	4	470	14.72	3.18
learning rate being 0.0004. We used GPT-2's [28] tokenizer. When using landmark tokens, the tokens were added to the dataset and stored as part of the train dataset, leaving the batching mechanism unchanged. We used gradient accumulation as well as data-parallel training across four nodes to maintain an effective total batch size of 128. We used mixed-precision training with bf10at16 over at most 4 Nvidia A100 GPUs. For our method, we train the model on each dataset for 240K steps with context length  $\ell_{\mathrm{seq}} = 512$ . We train Transformer-XL with a window size of 256 (i.e. effective context size 512) over segments of length 2048. We train Transformer-XL to observe the same number of tokens during training as our method which translates to performing 60K steps.
Results. To evaluate our model's performance with different context lengths, we divide the validation data into equally sized segments, referred to as evaluation lengths. Each segment is separately inputted into our model, which is further divided into chunks using the method described in Section 3.2. The chunk size, denoted as  $\ell_{\mathrm{local}}$ , represents the local context accessible without any memory. Table 1 presents the perplexity of the trained models under various inference settings. Notably, by using a local context length of 250 and retrieving the top  $k = 2$  most relevant blocks, we achieve a comparable performance with a context length of 512. This corresponds to attending to 360 tokens, including 250 tokens from the local context, 10 landmark tokens, and 100 tokens from the retrieved blocks. The effectiveness of using landmark tokens with retrieval becomes even more evident when comparing it to standard inference with an attention length of 360. Our results demonstrate that intelligently recovering relevant blocks enables attending to a significantly smaller number of tokens while maintaining performance.
Furthermore, our results highlight that landmark tokens enable the model to operate with larger context lengths than those encountered during training. The improvement in perplexity clearly indicates that the retrieved blocks contribute to the model's performance, making the results comparable to a Transformer-XL trained with segments of length 2048. However, unlike Transformer-XL, which can only leverage past information through recurrence, our method allows the model to attend to any token from the past, facilitating both the retention of exact fine-grained details and the interpretability of information utilization mechanisms.
Finally, the number of retrieved blocks and the number of blocks stored in memory can be adjusted during inference. While reducing the number of retrieved blocks  $k$  adversely affects performance, our results demonstrate that the model still outperforms the baseline even with only 2 retrieved blocks at context lengths of 2048 and 4096. Notably, when keeping only the last 40 blocks in memory, the model performs better at an evaluation length of 4096 compared to 2048. This suggests that the model is also learning to utilize recurrent mechanisms similar to those in Transformer-XL.
Table 2: Performance on PG19 dataset for different levels of retrieval flexibility. The blocks column shows the theoretical total number of blocks that can be accessed from the memory when feeding the input in windows of length 250 to the model.
Per Head	Per Token	Eval. Length	k	Blocks	Perplexity
✓	2048	2	250·8·2	15.14
2048	4	250·8·4	14.92
4096	4	250·8·4	14.72
✓	2048	2	8·2	15.48
2048	4	8·4	15.10
4096	4	8·4	14.95
X	✓	2048	2	250·2	15.44
2048	4	250·4	15.04
4096	4	250·4	14.89
Granularity of Cache Block Retrieval. Block retrieval can be performed on different levels of granularity. At the most granular level, the set of retrieved blocks can be different for each head and each token. This setting is the same as the model experiences during training. However, it is possible to further limit this granularity at inference, for increased system throughput. In this section we evaluate the effect of maintaining the same set of retrieved blocks across tokens or across heads. The results are presented in Table 2 which also shows the total number of retrieved block, with the same block retrieved by different token or head counted multiple times. While reducing the flexibility has a noticeable adverse effect on performance, the model still improves over the baseline. In particular, we note that it is possible to retrieve the same set of blocks for all tokens (which varies across heads) while only suffering 0.23 points in perplexity. To provide further insights into the expected improvement in speed gained from using a less flexible selection scheme, we further discuss the distribution of the retrieved blocks in Appendix C.
4.2 Fine-Tuning Pre-Trained Models
We demonstrate the possibility of fine-tuning a large language model using landmark's token and therefore extending the model's context length. Namely, we fine-tune LLaMA 7B [38] for 15000 steps using our method. To reduce computation, we fine-tune the model with context length 512. We use the sample subset of RedPajama<sup>1</sup> for the fine-tuning which closely follows the dataset curation process used for training LLaMA.
We evaluate the efficacy of our method by comparing model's ability to recover a hidden pass phrase inside a text segment. In particular, we use randomly generated prompts of the format described in Figure 3a and compute the accuracy of generating the correct pass key (as the first integer within the first 100 generated tokens). The result is plotted in Figure 3b for different context lengths. We observe that the base model is capable of finding the pass phrase until a certain length, even slightly exceeding its default training context length of 2048 (the area shared in grey). However, the base model completely fails at the task for larger contexts. In contrast, our landmark version can always retrieve the pass phrase with high accuracy, even for significantly larger context lengths. We point out that when evaluating our model with very large inputs (e.g.  $32\mathrm{K}$ ), we use additional techniques to reduce the memory usage by offloading the KV cache (except the landmarks) to CPU. We discuss this in more detail in Appendix G.
5 Future Work
Extrapolating Positional Encoding. One of the obstacles in attaining infinite context length is the inability of models to attend to context lengths much larger than those they were trained on. In this work, we provide a special indexing method which can be combined with landmark tokens to bypass this issue. However, as a result, the model can only attend to tokens that are too far based on their semantic (and not their position). While this is an important improvement and facilitates extrapolation to large context lengths, it can be expected that the performance would be further improved if the exact indexing method can be used. Unfortunately, existing proposals limit (or completely disable)
$^{1}$ https://github.com/togethercomputer/RedPajama-Data
There is an important info hidden inside a lot of irrelevant text. Find it and memorize them. I will quiz you about the important information there. <prefix filler by continuously repeating: The grass is green. The sky is blue. The sun is yellow. Here we go. There and back again.> The pass key is <PASS KEY>. Remember it. <PASS KEY> is the pass key. <suffix filler> What is the pass key? The pass key is
(a) Prompt Format
(b) Retrieval Accuracy
Figure 3: Prompt format used for comparing retrieval accuracy of the vanilla LLaMA 7B and its counterpart fine-tuned with landmarks. The points marked with a red cross represent cases where the model ran out of memory. Points marked with a green star use a more efficient inference mechanism (see Section G). Inference is done by feeding the segment in windows of length 250 tokens (excluding the inserted landmark tokens). The top  $k = 4$  landmarked blocks are retrieved. Retrieval accuracy is measured for a fixed total prompt length, by using the suffix and prefix filler. Results are averaged over 50 random generation of the pass key (a random number between 1 and 50000), which each time is located at a random position in the full-length prompt. The space before and after the pass key is filled accordingly by the suffix and prefix filler. The gray box marks the region where the prompt length is within lengths used during original LLaMA training.
attention to far tokens which defeats our purpose. While we briefly discuss a possible solution for models with landmark tokens in Appendix E, we leave a more thorough investigation as future work. We note that once such method is developed, it can be directly combined with landmark tokens, yielding inference capabilities at any length.
Hierarchical Landmarks. In large-scale settings, the landmark tokens can be stored in k-nearest neighbor data structures to improve retrieval performance and reduce memory usage. However, an alternative is to introduce hierarchy with higher level landmark tokens controlling the attention to lower level landmarks. In Appendix D, we investigate adding a special token which acts as a gate to all landmark tokens. This token can for example be used to decide whether a retrieval is necessary. Similarly, this token can be used at different memory cache levels where high attention to this token would constitute a cache miss, leading to lookup in lower-level (and slower) caches. We leave exploration of possible hierarchical landmark tokens as a future direction.
Training with Cache. For simplicity, in this work we focus on using the standard training procedure. While we expect the standard softmax mechanism to closely resemble the retrieval at inference, given the special indexing scheme, it is possible that the model would gain additional benefit from incorporating the cache during training. We leave investigation of such training variants as a future work.
In conclusion, this work presents a novel method for training attention to retrieve relevant blocks from memory. Unlike previous methods that rely on recurrence to create memory, our approach enables direct access to previous tokens, ensuring accurate information retrieval without the problem of slowly forgetting past data. We have demonstrated that our method achieves comparable performance to recurrent methods such as Transformer-XL while utilizing less computational resources. Additionally, our attention-based retrieval process allows for tracking and interpretability, providing insights into the information used to generate the output. Importantly, our results highlight the ability of our approach to handle significantly longer context lengths than those encountered during training. Moreover, we have shown that this capability can efficiently be incorporated into existing pre-trained models through fine-tuning, showcasing improved retrieval capabilities in the LLaMA 7B language model. Overall, our method enables efficient inference with arbitrary context lengths, making it suitable for accessing large inputs and processing fine-grained information within the large context.
The authors would like to express their gratitude to Matteo Pagliardini for insightful discussions and Olivia Simin Fan for their invaluable feedback on initial drafts of this paper. This project was supported by SNSF grant number 200020_200342.
[32] Ying Sheng, Lianmin Zheng, Binhang Yuan, Zhuohan Li, Max Ryabinin, Daniel Y. Fu, Zhiqiang Xie, Beidi Chen, Clark Barrett, Joseph E. Gonzalez, Percy Liang, Christopher Ré, Ion Stoica, and Ce Zhang. High-throughput Generative Inference of Large Language Models with a Single GPU, March 2023. arXiv:2303.06865 [cs].
[33] Jianlin Su, Yu Lu, Shengfeng Pan, Ahmed Murtadha, Bo Wen, and Yunfeng Liu. RoFormer: Enhanced transformer with rotary position embedding.
[34] Sainbayar Sukhbaatar, Edouard Grave, Piotr Bojanowski, and Armand Joulin. Adaptive Attention Span in Transformers, August 2019. arXiv:1905.07799 [cs, stat].
[35] Simeng Sun, Kalpesh Krishna, Andrew Mattarella-Micke, and Mohit Iyyer. Do long-range language models actually use long-range context? In Proceedings of the 2021 Conference on Empirical Methods in Natural Language Processing, pages 807-822, Online and Punta Cana, Dominican Republic, 2021. Association for Computational Linguistics. doi: 10.18653/v1/2021.emnlp-main.62. URL https://aclanthology.org/2021.emnlp-main.62.
[36] Yutao Sun, Li Dong, Barun Patra, Shuming Ma, Shaohan Huang, Alon Benhaim, Vishrav Chaudhary, Xia Song, and Furu Wei. A Length-Extrapolatable Transformer, 2022. URL https://arxiv.org/abs/2212.10554.
[37] Philippe Tillet, Hsiang-Tsung Kung, and David Cox. Triton: an intermediate language and compiler for tiled neural network computations. In Proceedings of the 3rd ACM SIGPLAN International Workshop on Machine Learning and Programming Languages, pages 10-19, 2019.
[38] Hugo Touvron, Thibaut Lavril, Gautier Izacard, Xavier Martinet, Marie-Anne Lachaux, Timothée Lacroix, Baptiste Rozière, Naman Goyal, Eric Hambro, Faisal Azhar, Aurelien Rodriguez, Armand Joulin, Edouard Grave, and Guillaume Lample. LLaMA: Open and Efficient Foundation Language Models, 2023. URL https://arxiv.org/abs/2302.13971.
[39] Sinong Wang, Belinda Z Li, Madian Khabsa, Han Fang, and Hao Ma. Linformer: Self-attention with linear complexity. arXiv preprint arXiv:2006.04768, 2020.
[40] Qingyang Wu, Zhenzhong Lan, Kun Qian, Jing Gu, Alborz Geramifard, and Zhou Yu. Memformer: A memory-augmented transformer for sequence modeling.
[41] Yuhuai Wu, Markus Norman Rabe, DeLesley Hutchins, and Christian Szegedy. Memorizing transformers. In The Tenth International Conference on Learning Representations, ICLR 2022, Virtual Event, April 25-29, 2022. OpenReview.net, 2022. URL https://openreview.net/forum?id=TrjbxzRcnf-.
[42] Manzil Zaheer, Guru Guruganesh, Kumar Avinava Dubey, Joshua Ainslie, Chris Alberti, Santiago Ontañón, Philip Pham, Anirudh Ravula, Qifan Wang, Li Yang, and Amr Ahmed. Big bird: Transformers for longer sequences. In Hugo Larochelle, Marc'Aurelio Ranzato, Raia Hadsell, MariaFlorina Balcan, and Hsuan-Tien Lin, editors, 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/hash/c8512d142a2d849725f31a9a7a361ab9-Abstract.html.
A Grouped Softmax Example
Figure 4: Example illustration of applying the method used for applying attention through landmark tokens. In the above example, a causal mask is applied as well. Tokens at indices 2, 5, and 8 are landmarks in the example as can be seen by checking the condition  $p_i = i$ . The example is based on the assumption that the dot product of the query vector of the token at index 6 and other key vectors is an all-one vector. It can be seen that when computing the attention for this token, the landmark at position 8 is ignored while the other two landmark tokens act as gate for attending to tokens in their block. The attention weight to landmark tokens is distributed among normal token in their block. Therefore, the final attention weight to the landmark tokens themselves is always zero.
B Dataset Description