Random-Access Infinite Context Length for Transformers
Authors: Amirkeivan Mohtashami
Martin Jaggi
Venue: NeurIPS 2023 poster (2023)
Keywords: large language models
memory
context length

Abstract: While Transformers have shown remarkable success in natural language processing, their attention mechanism's large memory requirements have limited their ability to handle longer contexts. Prior approaches, such as recurrent memory or retrieval-based augmentation, have either compromised the random-access flexibility of attention (i.e., the capability to select any token in the entire context) or relied on separate mechanisms for relevant context retrieval, which may not be compatible with the model's attention. In this paper, we present a novel approach that allows access to the complete context while retaining random-access flexibility, closely resembling running attention on the entire context. Our method uses a landmark token to represent each block of the input and trains the attention to use it for selecting relevant blocks, enabling retrieval of blocks directly through the attention mechanism instead of by relying on a separate mechanism. Our approach seamlessly integrates with specialized data structures and the system's memory hierarchy, enabling processing of arbitrarily long context lengths. We demonstrate that our method can obtain comparable performance with Transformer-XL while significantly reducing the number of retrieved tokens in each step. Finally, we show that fine-tuning LLaMA 7B with our method successfully extends its context length capacity to over 32k tokens, allowing for inference at the context lengths of GPT-4. We release the implementation of landmark attention and the code to reproduce our experiments at https://github.com/epfml/landmark-attention/.

Random-Access Infinite Context Length for Transformers
Amirkeivan Mohtashami EPFL amirkeivan www.tahaamieeplll.... ch
Martin Jaggi EPFL martin.jaggieepll...
Abstract
While Transformers have shown remarkable success natural language processing, their attention mechanism's 's large memory requirements have limited their ability to handle longer contexts. Prior approaches, such as recurrent memory retrieval-based augmentation, have either compromised the random-access flexibility of attention (i.c., the capability to select any token the entire context) relied separate mechanisms for relevant context retrieval, which may not be compatible with the model's s attention. In this paper, we present a novel approach that allows access the complete context while retaining random-access flexibility, closely resembling running attention on the entire context. Our method uses a landmark token to represent each block of the input and trains the attention it for selecting relevant blocks, enabling retrieval of blocks directly through the attention mechanism instead of by relying seamlessly integrates with specialized data structures and the system's 's memory hierarchy, enabling processing of arbitrarily long context lengths. We demonstrate that our method can obtain comparable performance with Transformer-XL while significantly reducing the number retrieved tokens in each step. Finally, we show that fine-tuning LLaMA 7B with our method successfully extends its context length capacity over 32k tokens, allowing for inference the context lengths of GPT-44 We release the implementation of landmark attention and the code to reproduce our experiments https: //ggps:://iie /ggtttb.......gg www.co/////ceeee epfml/ ceimn/laandenn Landmark attention/.
separate mechanism. Our approach
Introduction
Large transformers have revolutionized language modeling and demonstrated remarkable abilities to perform various tasks with zero or few examples [4]. This success can be largely attributed to the attention mechanism, which allows each token to access the representation of any other token in each layer. However, this flexibility comes with quadratic computational cost and highly problematic memory footprint, limiting the number of tokens that can be attended to, and thus the context length.
To overcome this limitation, researchers have proposed various solutions, including incorporating form of recurrent memory inside the Transformer architecture, such as Transformer-XI. [9]. However, these approaches often sacrifice the random access flexibility of attention.
An alternative approach to overcome the context length limit is to use retrieval- based methods that incorporate additional static knowledge by searching for relevant documents in knowledge base and adding them to the context. However, this approach requires a separate mechanism to identify relevant documents, called retriever. Such retrieval models can not easily be updated work on fresh long input data, and furthermore are also not fully compatible with the standard attention mechanism itself, and thus may fail mimic attention over long documents.
In this work, we propose a novel approach for overcoming the context length limit by allowing earlier blocks of the input to be directly incorporated into the attention itself. We break the input into blocks
37th Conference on Neural Information Processing Systems (NcurIPS 2023).

---

b
G
Standard Attention
b
b
c
Attention with Landmarks
e
Figure 1: An illustration comparing standard attention and attention with landmarks. The example shows the (causal) attention given by a current token e to previous ones, illustrating our mechanism with block-size folock =. The attention scores rely on the similarity of query vector with the key vector, and in our case also with the landmark vector corresponding to the block. This is why the same token can have a high (green) attention score when being part of one block and low (red) attention score when being in other one, despite having the same representative vector in both cases. Landmark tokens (same as regular tokens) have the same vector representation at the first layer. However, this changes as they are updated though depth, leading to the illustrated behavior of attention the intermediate layers.
of fixed length and introduce a special token for each block, called a landmark, which acts as a gate for attending its corresponding block. The gating mechanism is controlled by the attention score to the landmark token. At inference time, the attention scores on the landmarks allow us to retrieve any previous block and integrate with standard attention. The idea is illustrated in Figure 1. Our proposed approach maintains the random-access flexibility of attention and offers an alternative solution to the recurrent memory approaches.
Our model can process any context length at inference time regardless of the context length used at training time. To achieve this, we split the input into chunks and feed the chunks sequentially to the model, while maintaining a cache of previous tokens, usually referred to as KV cache. When processing each chunk, we first the landmark tokens select the most relevant blocks and only these blocks for computing the attention. This immediately reduces the computation cost by factor of block length. For example, in our experiments where we use blocks of 50 tokens, this translates to almost 50x reduction of computation. We note that to the overhead of computing the attention for the retrieved blocks does not depend on the input length and becomes negligible for very large inputs. Furthermore, it is possible to obtain the same reduction in memory usage since all tokens in a block (except the landmark itself) can be swapped out and only loaded when the corresponding landmark token is activated (see Appendix G). We also point out that this reduction can be further improved by using special data structures designed for retrieving closest neighbors, such as FAISS [15].
We demonstrate the efficacy of our method in practice by applying our training method both for training models from scratch and for fine-tuning pre-trained models. In both cases, model effectively utilizes landmark tokens to retrieve relevant blocks from memory, enabling inference at arbitrary context lengths much longer than those encountered during training. As result, our model obtains comparable performance with Transformer XL trained to use recurrence on a much larger window. More importantly, we demonstrate that using our method to fine-tune LLaMA 7B [38], a large language model, allows it to retrieve relevant information from contexts with over 32k tokens, which is the context length of GPT-4 [25].
The primary advantages of our method can be summarized as follows:
Enabling inference at any context length, irrespective of the context length utilized during training, without incurring additional training costs.
Instantly reducing inference time and memory usage (compared to model trained to operate at the given context length) by substantial factor equal to the block size (e.g., 50 in our experiments).
Compatibility with advanced data structures that can further decrease the resource require- ments for operating the model with very large context lengths.

---

2
Related Work
With the evolution of state-of-the-an commercial models and their applications towards very long context window lengths, such as 32k tokens (GPT-4 [25]) or even 100k (Claude [2]), the research question of efficient while accurate long context models is receiving increased attention.
Gatrieval-Augmenter Language Models. Retrieval-augmentee module, called a retriever, to find set of relevant documents in the knowledge base, which are then prepended to the input. The augmented input is then fed into the main model, called the reader. Various methods have been proposed for training retrievers and readers [16]. For example, REALM [11] jointly trains the reader and retriever, where both components are transformers. Atlas [13] further investigates the effect of various losses on the performance of the retriever. Previous work has also looked into using the attention in the reader to build a retriever but used manually crafted rules to reduce the token scores to document scores [14, 18, 31]. In contrast, our landmark approach is trained to directly produce meaningful landmark embeddings on the fly, without needing any notion of corpus for retrieval.
language models use a separate
Memory for Transformers. Various methods have been proposed to introduce memorization capabilities Transformers through recurrence [5, 40]. Transformer-XL [9] feeds the input to the model in windows of a fixed length and allows each token to attend to tokens in the current window as well as the preceding window. Memory Transformers [6] introduce special memory tokens that are prepended to the input, and their representation at the final layer of the model is used for the next input. Infinite Memory Transformers [23] map the input to continuous space and then sample points to be used for memory in the next step according to the probability distribution defined by the attention mechanism. However, while these methods improve upon the memory-less variants, they do not allow for attending to specific tokens in the past, as the model only has access to a compressed version of this information. In fact, Mu et al. [24] in simultaneous work propose adding special "gist" tokens which are trained to summarize the prompt so far, and find that the model is incapable of remembering specific details that should be copied into the output. Furthermore, the decision about whether to keep or discard a piece of information needs to be made without knowledge of future tokens, which makes it likely that aspects of information will be lost, especially if the topic changes. In contrast, using our method, the model always has the possibility of retrieving and attending to any tokens in the past. Nonetheless, we note that these methods can be combined with ours, allowing the model benefit from both full access previous tokens as well access to a summarized version in terms of the recurrent memory state.
Approximate and Sparse Attention. Various methods have also been proposed to reduce the memory footprint of attention. However, similar to recurrent memories, these approximations significantly reduce the ficxibility of attention in attending to arbitrary individual tokens. For example, Child et al. [7] limit the attention to a local window around each token, while BigBird additionally suggests attending to random subset of previous tokens as well several globally accessible tokens [42]. Longformer [3] further introduces dilated sliding window patterns to increase attention's 's receptive field and manually picks the window sizes for each layer. Linformer [39] uses low-rank approximation of the attention matrix while Performer [8] uses a non-softmax kernel to obtain a more efficient implementation. Reformer [19] uses Docaliy-sensitivv--aal
(LSH) to retrieve the closest key vectors which should account for the highest scores in the attention matrix. Combiner [30] utilizes a hierarchical attention mechanism and heuristic reduction techniques, such as max-pout max-pooling, to derive key and query vectors for input blocks. The block weight is determined based on the pooled key vector, while the weight of each token within the block determined by the pooled query vector. However, this approach limits the control of the current token over the weights of the tokens inside the block, resulting in reduced flexibility of attention. In contrast, our proposed method enables the current token's 's query vector to control the weight for each token, and the gating mechanism is learned through the attention process instead of relying on heuristic reductions.
KNN Augmented Transformers. k-nearest-neighhbrr (kNN) augmentation has been proposed as an alternative method for allowing transformers to access external memory. For example, ANN-LM [17] stores the hidden representation of tokens in memory and uses the distribution of the next token among the stored vectors that are closest to the current token to predict the next token. Memorizing Transformer [41] performs nearest- neighbo search over previous keys and computes the attention to the top nearest items. However, these methods obtain the final results by interpolating between the ANN prediction and the local attention prediction using tuned parameter as interpolation weight. Therefore, the interpolation does not depend on the current input and does not consider whether the memory contains relevant information.
3

---

Context Length Extrapolation. Transformers have well-known limitation in extrapolating to contexts longer than what was observed during training [27], even when relative positional encoding is used [36]. Current solutions to address this problem often result in weakened attention scores for long-range tokens, which undermines the benefits of longer context (27, 36). Moreover, these methods only work when combined with windowed attention, which restricts direct attention to long-range tokens [36]. We hypothesize that the limitation may partially stem from the model's learning of its own positional encoding with the use of causal masking, demonstrated in [12]. This limitation poses a challenge as our goal is to enable access to long-range tokens during inference at distances that were not observed during training. We discuss solutions in Section 3.2.1.
3
Methodology
In this paper, we mainly focus on the causal language modeling where each token can only attend to previous tokens in the input. We briefly discuss the extension of our method the non-causal case in Appendix F.
When using Transformer to process long input, the ideal case would be to allow each token to attend to all previous tokens. However, this becomes computationally infeasible as the input length increases. Nevertheless, since the attention scores always sum one, the number of keys with large attention weight is limited even for long contexts. Thus, by retrieving only those keys with large attention scores, it is possible closely emulate the ideal case. In this work, we propose a method to find these keys by dividing a long input into blocks of consecutive tokens and using the attention to retrieve relevant blocks.
More particularly, we assign a representative vector to each block such that a high attention score to any token inside block would lead to high attention score to the block's ' representative vector. Therefore, we can directly retrieve blocks based on the attention score of their representative vector.
To obtain the representative vector of a block, we introduce a new special token to the vocabulary, called the landmark token. We insert landmark token after the last token of each block and train the model such that the key vector for this token becomes the representative vector we seek. The process illustrated in Figure 1.
We will first describe the method we use to train the landmark tokens in Section 3.1 and then describe the inference process in Section 3.2.
We note that an alternative for directly finding a candidate set of keys with high attention score is using a data structure that allows finding nearest neighbors of the query vectors efficiently such as FAISS [15]. In comparison, our method provides retrieval method directly controlled by attention which can be more semantic-based, Furthermore, retrieving block instead of a single token allows the attention to also access the local context around the token which may be more accommodating of observed classic attention patterns [20]. Finally, we point out the aforementioned data structures can also be applied on top of our method to search for relevant blocks.
3.1 Training Landmark Tokens
In order to train the landmark tokens, we first go over the text corpus and add landmark token after every falock tokens. Then we proceed to training the model using the standard batching method which feeds windows of Esey tokens to the model. In a list of fueq > bpteek tokens, we train the model such that each landmark token represents the block consisting of all previous tokens until the previous landmark token (or the beginning of the input if no previous landmark token exists). The token is passed through the transformer any other token while its representation is updated using the self-attention mechanism. Let us denote the index (token position) of the landmark corresponding to the i-th token's block by P;- If the last block is incomplete and does not have landmark token, we define p. := Esey- the i-th token a landmark token, p; := i.
In order to train the transformer to make use of landmark tokens, we alter the standard attention mechanism such that the attention weight for token depends on the similarity of the query vector with both the token's 's key as well as with the key of its block's landmark token. To define the mechanism, we first define generalized softmax function called Grouped Softmax. Given a vector fM and group index g e Nd,, Grouped Softmax applies softmax separately over elements

---

belonging to the same group. (Using En recovers the standard softmax function):
oG(v.g) GroupedSoftmax((..gm g). = EvR.=RE eV
(I)
We replace the softmax function after computing the attention scores with Grouped Softmax. For each block, we put its regular tokens in separate group, ensuring that all regular tokens within the same block share the same group, while tokens outside the block are assigned to different groups. When computing the attention weights for the i-th token, landmark tokens for other blocks are placed in the same group as the i-th token. The landmark token for the i-token's 's block is ignored when computing the attention weights for the i-th token. other words, the landmark token for each block is only used by tokens in other blocks. This is intuitive as the landmark token should only be accessed when tokens in other blocks require to retrieve information from the landmark's corresponding block. Building on the fact that Pj- = j only holds when the j-th token is a landmark token, we can define the grouping used for the i-th token more formally as
Pj
Pj j
placing normal tokens in their own blocks.
G
p: p=jj pi P D placing other landmarks in the i-th token's group.
> ignoring current block's ' landmark token.
Finally, to obtain the final weights after applying GroupedSoftnax, we multiply each token's 's softmax output with the softmax output for its block's landmark token. For the tokens in the same group as the i-th token, we directly use the softmax output its attention weight. The weight for landmark tokens is always zero. In more formal terms,
Q x K
S.j == SoftmaxScore( K);
GroupedSoftmax
Gi)
(3)
Vhend
0
Pj G. = G
AttWeight Q, K)
SJ
Si.j Siny
G # G. APp $
An exmaple illustration of various values defined above is given in Appendix A. Note that under this scheme, the attention weights sum to one as is the case for the standard softmax function. More importantly, attending to tokens in other blocks is gated by the attention score to the landmark token as expected. Since tokens in the same block and the landmark tokens share the softmax group. the model has to choose between attending to other blocks and current tokens. Thus, the intuition behind the grouping is to force the model to only attend to relevant blocks due to this trade-off.
We note that attention masks can be applied normally by ignoring the masked elements in the softmax (c.g. by setting A.,, -00 on the masked elements in practice). Indeed we focus our experiments on the causal language modeling. We also point out that the grouping scheme can be further extended to introduce additional hierarchy for retrieval. For example, we refer the interested reader to Appendix D, where we briefly discuss different grouping scheme which also trains global retrieval gate token that controls whether retrieval from an earlier block needs to be performed. At inference, this gate can be used to decide whether memory call is needed or the model already has the information it needs in the context. We leave further investigations of this setting for future work.
3.2 Inference
Similar to training, the input gets augmented by landmark token after every bbink tokens. Then, we break the input into chunks of Gocal length and iteratively feed chunks from the beginning to the end. To retrieve relevant blocks, each attention layer has access to cache of previous blocks. The cache stores the key-value vectors for all tokens of those blocks, including the landmark token. Since the retrieval only requires access the landmarks, we can reduce the memory usage significantly by swapping out (for example to CPU memory or even to disk) all regular tokens' cached key-value vectors, and then swapping them back in only if their corresponding block is retrieved by the attention.
We start by discussing the most permissive retrieval scheme. When processing each chunk at each layer, we first compute the attention score of each token with the landmark tokens currently in the cache. For each token, we compute the attention score for all the tokens in the highest scoring blocks and prepend the obtained attention matrix to the local attention matrix. We finish by applying
5

---

Memory
9
b
9
Q
f
Retroival Position Mapping
h
9
5
5
P
9
Input
Final Position Mapping
Figure 2: Stingy position mapping: Retrieving top k= blocks from memory of 5 blocks. Retrieval landmarks for the last blocks are accurately mapped to sequence index positions, while previous blocks mapped to the position of the 1-th last block. Blocks are then distributed across the prefix based on their position, with an empty block separating the last & blocks from older ones.
GroupedSoftmax to the attention matrix and computing the weighted average of the value vectors to obtaín the token's 's representation.
Under the above scheme, each token and each head can retrieve different blocks from the cache. is possible limit the retrieval flexibility in order improve efficiency. For example, it is possible to merge the scores across heads by taking a maximum over the landmark attention scores (after applying softmax) of each head. Under this scheme, the same set of blocks is retrieved for all heads. It is also possible to take the maximum over different tokens, retrieving only A blocks per head for all the tokens in the current window combined. We study the effect of these limitations at the end of Section 4.1. Unless otherwise stated, we experiment using the permissive scheme described above.
3.2.1 Positional Encoding
When computing the attention scores to cache elements (both landmark and normal tokens), it is important correctly incorporate positional information. The transformer model is sensitive to positional information. It also intuitive that the model would rely on position information some cases. For example tokens right after the last memory block do have access to any context and are unable to select memory blocks based on semantics. Instead, they need to rely on positional information 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 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.
window of tokens. We decide against using
We allocate segment with length (k+1) (folock +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 & blocks the corresponding place within the last A 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 decide which blocks to retrieve, 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.
map those among the latest & blocks to the right of
Note that we found out that when k == 1 mapping memory blocks to segment of at least 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 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 is possible that the importance of pre-allocating the additional block decreases k grows, adapt this scheme so that the attention would be at least as flexible of simply keeping the last & blocks.
We point out that the above approximation relies the ability to add position information when performing the retrieval. 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 GroupedSoft- max. In particular, our method does not require maintaining 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 (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 infer- ence time by computing the attention score over smaller set of token pairs. For instance, during auto-regressive generation, where tokens are generated sequentially by using the previously generated token input for obtaining the next one, the traditional approach involves computing attention across all preceding tokens when generating the next token. contrast, our method allows computing atten- tion 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 every folock + 1 tokens. This immediately reduces the number of operations by a factor of block length blookk- For example, when using block = 50 as in our experiments, this can lead to a 50x boost. Importantly, the same reduction can be obtained in terms of memory. In the standard Transformer, a cache of all previous key and values (K V-cachee) is needed to perform the generation efficiently. In contrast, we only need immediate access to the landmark tokens and can offload the blocks slow memory (e.g. CPU), loading only the retrieved blocks. The reduction in memory and compute can be further improved the search for retrieving blocks is performed using more advanced data structures such FAISS [15].
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 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].
Experiments
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 Transformer-Xl. while reducing FL.OPs. In contrast with Transformer-XLL using our method, the information retrieval 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 GPT-2 [28] architecture: a 12-layer decoder-only transformer with heads each layer, each with dimension 128 (embedding dimension 1024), and hidden FFN size of 4096. We trained our model using AdamW [22] with 31 = 0.9 and 82 = 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 cosine scheduler with minimum (final)
7

---

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-XI19) When using landmarks, the column " Blocks" shows the maximum number of blocks stored in memory. Each block contains (block = 50 normal tokens and one landmark token. Due to computation limitations we only report results for Transformer-Xe XL on PG-19 as this method takes longer train in our implementation.
Eval. Length
foeal
XL cache
Blocks
k
Attention Size
PG19
arXiv
512
None
None
512
16.12
4.01
Baseline
512
360
None
None
360
16.76
4.31
250
None
10
2
360
16.23
4.01
Ours
256
256
None
512
14.72
[9]
250
None
40
2
360
15.14
3.43
350
None
40
2
460
15.07
3,41
2048
Ours
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
256
256
None
512
14,55
[9]
4096
250
None
40
460
14.79
3.19
None
80
2
370
15.00
3,29
Ours
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 mixeddppreissio at most 4 Nvidia A A100 GPUs For our method, we train the model on each dataset for 240K steps with context length leeq 512. We train Transformer-Xi XL with 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 method which translates performing 60K steps.
training with bfloat16 over
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 hoaat, represents the local context accessible without any memory. Table 1 presents the perplexity of the trained models under various inference settings. Notably, by using local context length of 250 and retrieving the top k: 2 most relevant blocks, achieve comparable performance with 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 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 significantly smaller number 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-2 XL, which 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 model still outperforms the baseline even with only 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-XI.
8

---

Table 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
14.92
4096
4
250
8
14.72
2048
2
8
15.48
X
2048
4
8
8
15.10
4096
4
8
4
14.95
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 inference, for increased system throughput. In this
possible to further limit this granularity section we evaluate the effect of maintaining the same set of retrieved blocks across tokens or across heads. The results are presented in Table 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, 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 less flexible selection scheme, further discuss the distribution of the retrieved blocks in Appendix C.
4.2 Fiie-Tuuiig Pre-Trained Models
We demonstrate the possibility of fine-tuning large language model using landmark's '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' 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 certain lengths, 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. 32K), we use additional techniques to reduce the memory usage by offloading the KV cache (exeept 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 trained on. In this work, we provide a special indexing method which can be combined with landmark tokens to bypass this issue. However, as result, the model can only attend 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)
hatg://iihhccommllggsutcooeeeeem Data

---

There is important info hidden inside a lot of irrelevant text. Find it and memorize them. will quiz you about the important information there.
100
80
<prefix filler by continuously repeating:
Accuracy
60
LLaMA 78
The grass is green. The sky is blue.
Landmark Finetuning
The sun is yellow. Here we go. There and back again.>
40-
The paas key is <PASS KEY>. Remember it.
<PASS KEY> is the pass key. <suffix filler>
What is the pass key? The pass key is
5000
10000
20000 25000 39000 Prompt (Tokens)
15000
(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 red cross represent cases where the model ran out of memory. Points marked with a green star use 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 and 50000), which each time is located at random position in the full-length prompt. The space before and after the pass key 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 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. 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 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 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 incor- porating the cache during training. We leave investigation of such training variants as a future work.
Conclusion
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-XX 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.
10

---

Acknowledgment
The authors would like to express their gratitude to Matteo Pagliardini for insightful discussions and Olivia Simin Fan for their invaluable feedback initial drafts of this paper. This project was supported by SNSF grant number 200020 200342.
References
Reza Yazdani Aminabadi, Samyam Rajbhandari, Minjia Zhang, Ammar Ahmad Awan, Cheng Li. Du Li, Elton Zheng, Jeff Rasley, Shaden Smith, Olatunji Ruwase, and Yuxiong He. DeepSpeed Inference: Enabling Efficient Inference of Transformer Models at Unprecedented Scale, June 2022. arXiv:2207.00033 [es].
[2] Anthropic.
claude-vi v1. 3-10kk, Blog post 'Introducing 100K
Context
Windows'.
htps://www.nntropicccommm
http:pppp...coomneeeeooooooooooo contee context-windoww, 2023.
[3] Iz Beltagy, Matthew E. Peters, and Arman Cohan. Longformer: The dongmeeemmntt transformer.
[4] Tom B. Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, Sandhini Agarwal, Ariel Herbert-Voss Gretchen Krueger, Tom Henighan, Rewon Child, Aditya Ramesh, Daniel M. Ziegler, Jeffrey Wu, Clemens Winter, Christopher Hesse, Mark Chen, Eric Sigler, Mateusz Litwin, Scott Gray, Benjamin Chess. Jack Clark, Christopher Berner,
Radford, Ilya Sutskever. Dario Amodei. Language models are few-shot learners. In Hugo Larochelle, Marc' Aurelio Ranzato, Raia Hadsell, Maria-Florin Balcan, and Hsuan-Tien Lin, editors. Advances in Neural Information Process- ing Systems 33: Annual Conference Neural Information Processing Systems 2020, NeurIPS 2020, December 6- 12, 2020, virtual, 2020. URL https://tstt httpss////cccee/e/////// neessinnrispp.. wttppss.cccppe/e/eeeeeeeeppeeeet /2020/ hash/ 1457 (0060000bbbbb990 41116667718bbbaacccc 421 64a Abstract serrratthtmll htmll.
McCandlish,
[5] Aydar Bulatov, Yuri Kuratov, and Mikhail S. Burtsev. Recurrent memory transformer.
[6] Mikhail S. Burtsev, Kuratov, Anton Peganov, and Grigory Sapunov. Memory transformer.
[7] Rewon Child, Scott Gray, Alec Radford, and llya Sutskever. Generating Long Sequences with Sparse Transformers, 2019. URL htt.ttt/tta/iiii.....a org/ /bte///aaatt/a /abb///994...4 10509.
[8] Krzysztof Choromanski, Valerii Likhosherstov, David Dohan, Xingyou Song, Andreea Gane, Tamas Sarlos, Peter Hawkins, Jared Davis, Afroz Mohiuddin, Lukasz Kaiser, David Belanger, Lucy Colwell, and Adrian Weller. Rethinking Attention with Performers, November 2022. arXiv:2009. 14794 [cs, stat].
[9] Zihang Dai, Zhilin Yang, Yiming Yang, Jaime Carbonell, Quoe Le, and Ruslan Salakhutdinov. Transformer- XL: Attentive language models beyond fixed-length context. In Proceedings of the 57th Annual Meeting of the Association for Computational Linguistics, pages 2978 2988, Florence, Italy, 2019, Associa- tion for Computational Linguistics. doi: 10.18653/v1//19 1285. URL https::/aclaathholooo org/ P19 1285.
[10] Tri Dao, Daniel Y. Fu, Stefano Ermon, Atri Rudra, and Christopher Ré. FlashAttention: Fast and Memory-Fificient Exact Attention with IO-Awareness, June 2022.
[11] Kelvin Guu, Kenton Lee, Zora Tung, Panupong Pasupat, and Ming-Wo: Chang. REALM: Retrieval- augmented language model pre-training.
[12] Adi Haviv, Ori Ram, Ofir Press, Peter Izsak, and Omer Levy. Transformer Language Models without Positional Encodings Still Learn Positional Information, 2022. URL https: http///axvvvv//////eee///e wwwaaiiivgg.....e..eeeeeee abs/ 16634.
[13] Gautier Izacard, Patrick Lewis, Maria Lomeli, Lucas Hosseini, Fabio Petroni, Timo Schick, Jane Dwivedi- Yu, Armand Joulin, Sebastian Riedel, and Edouard Grave. Atlas: Few-shot learning with retrieval augmented language models.
[14] Zhengbao Jiang, Luyu Gao, Jun Araki, Haibo Ding, Zhiruo Wang, Jamie Callan, and Graham Neubig. Retrieval as attention: End-to-end learning of retrieval and reading within single transformer.
[15] Jeff Johnson, Matthijs Douze, and Hervé Jégou. sclle similarity search with GPUs.

---

[16] Vladimir Karpukhin, Barlas Oguz, Sewon Min, Patrick Lewis. Ledell Wu, Sergey Edunov, Danqi Chen, and Wen-tau Yih. Dense passage retrieval for open-domain question answering. In Proceedings of the 2020 Conference on Empirical Methods in Natural Language Processing (EMNLP), pages 6769-6781, Online, 2020, Association for Computational Linguisties. doi: 10. 10.186631122022mmmmmmmmmm mainn........ URL https: hatpss//aaaatthhlle org/2020 emmlp- main.550.
[17] Urvashi Khandelwal, Omer Levy. Dan Jurafsky. Luke Zettlemoyer, and Mike Lewis. Generalization through memorization; Nearest neighbor language models. In 8th International Conference on Learning Representations, ICIR 2020, Addis Ababa, Ethiopia, April 26-30, 2020. OpenReview.net, 2020. URL. https://openreriis net/ oormforrmm/tteeemmm htt///ormmmmdd111// HB-HHBBEC
[18] Omar Khattab, Christopher Potts, and Matei Zaharia. Relevance-guided supervision for OpenQA with ColBERT. Tramsactions of the Association for Computational Linguistics, 9:929-944, 2021. doi: 10.1162/ a a_0_000 URL https: http///alaatttthh anthology wwwproggg.rgggoooe 2rg/2021.2 tacl-1.55
[19] Nikita Kitaev, Lukasz Kaiser, and Anselm Levskaya.
Reformer: The efficient transformer. In 8th International Conference Learning Representations, ICLR 2020, Addix Ababa, Ethiopia, April 26-30,
2020. OpenReview.net, 2020. URL https:/ http::::///pppeeeiiieeee wetteeoreiiddeeeeeeeeeeeettteeet id-kggikkttt...
[20] Olga Kovaleva, Alexey Romanov, Anna Rogers, and Anna Rumshisky. Revealing the dark secrets BERT. In Proceedings of the 2019 Conference on Empirical Methods in Natural Language Processing and the 9th
International Joint Conference on Natural Language Processing (EMNLP (IINNP) pages 4365 4374, Hong Kong. China, 2019. Association for Computational Linguistics. doi: 10. 10.18653/v1/D19/ 1445. URL https: /acl/a/hoool org/D1g- 1445.
[21] Alex Krizhevsky, Ilya Sutskever, and Geoffrey E Hinton. ImageNct Classification with Deep Convolutional Neural Networks. In Advances in Neural Information Processing Systems, volume 25, Curran Associates, Inc., 2012.
[22] Ilya Loshchilov and Frank Hutter. Decoupled weight decay regularization. In 7th International Conference on Learning Representations, ICLR 2019, New Orleans, IA,, USA, May 6-9, 2019. OpenReview.net, 2019. URL https: wwwpeerrevieeeeeefoooooeeee 1d-Bkg6Ri
[23] Pedro Henrique Martins, Zita Marinho, and André F. T. Martins. SkinftyS-fommer: Infinite memory transformer.
[24] Jesse Mu, Xiang Lisa Li, and Noah Goodman. Learning to Compress Prompts with Gist Tokens, 2023.
URL https: / //wxxvv.vvv/////////////////wtv. org/ abo//304..0 08467.
[25] OpenAL. GPT-4 Technical Report. arXiv, 2023.
[26] Reiner Pope, Sholto Douglas, Aakanksha Chowdhery, Jacob Devlin, James Bradbury, Anselm Levskaya, Jonathan Heek. Kefan Xiao, Shivani Agrawal, and Jeff Dean. Efficiently Scaling Transformer Inference. November 2022. arXiv:221 100510022000 [cs].
[27] Ofir Press, Noah A. Smith, and Mike Lewis. Train short, test long: Attention with linear biases enables input length extrapolation. In The Tenth International Conference on Learning Representations, ICLR 2022, Virtual Event, April 29, 2022. OpenReview.net, 2022. URL htths: / hpwpreeeeevvveeeeee veew..eeee.. nwt//fermmii R8sQPPGCvO.
[28] Alec Radford, Jeff Wu, Rewon Child, David Luan, Dario Amodei, and llya Sutskever. Language models unsupervised multitask learners. 2019.
[29] Jack W. Rac, Anna Potapenko, Siddhant M. Jayakumar, Chloe Hillier, and Timothy P. Lillicrap. Com- pressive transformers for long-range sequence modelling. In 8th International Conference on Learning Representations, ICLR 2020, Addis Ababa, Ethiopia, April 26-30, 2020. OpenReview.net, 2020. URL attpa://openrevier http:/nnee//fneieeeeeeeee
[30] Hongyu Ren, Hanjun Dai, Zihang Dai, Mengjiao Yang, Jure Leskovec, Dale Schuurmans, and Bo Dai. Combiner: Full Attention Transformer with Sparse Computation Cost, October 2021. arXiv:2107.05768 [es].
[31] Keshav Santhanam, Omar Khattab, Jon Saad-Falcon, Christopher Potts, and Matei Zaharia. ColBERT Effective and efficient retrieval via lightweight late interaction. Proceedings of the 2022 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, pages 3715 3734, Seattle, United States, 2022. Association for Computational Linguistics. doi: 10.18653/v11200 1/2022. www.c-aain.227777 272. URL https: http:///aaaaaallcccoooo 2022 naacl-ma maaac-mminn... 272.
12

---

[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. Hlightthruuhhut Generative Inference Large Language Models with a Single GPU, March 2023. arXiv:2303, 068665 [es].
[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 Transformers, August 2019. arXiv:1905 07799 (cs, stat].
[35] Simeng Sun, Kalpesh Krishna, Andrew Mattarella-Micke and Mohit lyyer. 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.18663/1/2021.mmmll http://22...mmmmmmmmmm2222 URL https:// aclanthology , org/2021. cannp-maie 62.
[36] Yutao Sun, Li Dong, Barun Patra, Shuming Ma, Shaohan Huang, Alon Benhaim, Vishrav Chaudhary. Xia Song, and Furu Wei. cength-Extrapolataba Transformer, 2022. URL httpp://aaxii..t/
2212. 10554.
[37] Philippe Tillet, - Kung, David Cox. Triton: an intermediate language and compiler for tiled neural network computations. 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, Maric Marie-Anne Lachaux, Timothéc 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/ /2tp//23.2..222 13971.
[39] Sinong Wang, Belinda Z. Li, Madian Khabsa, Han Fang, and Hao Ma. Linformer: Self-attention with linear complexity. arXiv preprint 047668,000...788, 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 Learning Representations, ICLR 2022, Virtual Event, April 25-29, 2022. OpenReview.nel, 2022. URL httpp:://penreevies.... wett//eruriid/eeeeeeeeeeee www.riidzrrrr............
[42] Manzil Zaheer, Guru Guruganesh, Kumar Avinava Dubey, Joshua Ainslic, Chris Alberti, Santiago Ontañón. Philip Pham, Anirudh Ravula, Qifan Wang. Li Yang. and Amr Ahmed. Big bird: Trans- formers for longer sequences. In Hugo Larochelle, Mare' Aurelio Ranzato, Raia Hadsell, Maria- Florina Balcan, and Hsuan-Tien Lin, editors. Advances in Neural Information Processing Systems 33: Annual Conference on Neural Information Processing Systems 2020, NeurIPS 2020, Decem- ber 6-12, 2020, viruual, 2020. URL https: https://proon http////ooeeen...///////// neurips wwtpiip...pppppee wttppcppp///////// /ataa/er/020//1 /ashh/ c8512d 22221177722224771111133
Abstract. html.
13

---

A Grouped Softmax Example
0
2
3
4
5
6
7
8
0 K
mask +
1
1
1
1
mask
mask
V head
pi
2
5
5
8
8
G.,j
2
8
3
8
8
S
0.5
0.33
0.33
ign
Wo.i
0.167
0.167
0
0.167
0.167
0
0.33
0
Figure 4: Example illustration of applying the method used for applying attention through landmark tokens. the above example, a causal mask applied as well. Tokens at indices 2, 5, and are landmarks in the example can be seen by checking the condition y; i. The example is based on the assumption that the dot product of the query vector of the token at index and other key vectors an all-one vector. It can be seen that when computing the attention for this token, the landmark position is ignored while the other two landmark token, act as gate for attending to tokens their block. The attention weight 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
arxiv Math We use the cleaned arXiv math subset of proof-pile dataset. The dataset contains all the TEX files for submissions under math category arXiv. Files with encodings other than (tff-/16// 1/1622222 or latini 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, subsections. The final training dataset contains around 5.6B tokens.
PG-19 PG-19 is 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 length 2048 (in chunks of 250 tokens). During this process, we retrieved the top = blocks from the cache and kept track of the number of unique blocks retrieved in the penultimate chunk. Note that 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, the last chunk may have 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 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 each layer, indicated by the spike in the last bin. Interestingly, the additional spike in the flexible retrieval corresponds to a very low number of blocks
14

---

0.30
Token & Head
Per Head
0.25
Per Token
0.20
Density
0.15
0.05
0.00
10
20
25
30
35
# of Different Retrieved Blocks
Figure 5: The distribution of the number 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 attributed to the retrieval pattern the initial layer. This finding aligns with previous research suggesting a locality attention in the earlier layers [34].
Moreover, the results demonstrate a significant decrease in the number unique retrieved blocks when only allowing the set of retrieved blocks to vary across heads. Although there 8 heads, each retrieving 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), can significantly improve performance by facilitating offloading slower devices through 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
increase the number of retrieved blocks. This be observed by comparing the results between retrieving the top k = 2 blocks at the highest level of flexibility and retrieving the top blocks with reduced flexibility, as shown 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 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.
tokens but not heads, it is still possible
D Context Miss Token
In this section, we demonstrate how to create additional hierarchy in the landmark structure by simply changing the grouping. particular, we use new grouping scheme to train 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 tokens in their respective blocks. In particular, let denote LCMT the set of landmark tokens that are controlled by CMT. Then, using the same notation as in Section 1. we use the following grouping during training (For brevity, the below formulation assume that the negative of all preceding conditions hold for each case):
15

---

Pi
j = -1
placing CMT in the i-th token's ' group.
Pj
Pj +
placing normal tokens in their own blocks.
Gij=
-1 Pi =
ignoring current block's landmark token.
(5)
Pi j & LCMT
placing free landmarks in the i-th token's 's group.
-2 j E LCMT
> placing CMT-controlled landmarks in a separate group.
Using the above grouping the attention weights can be calculated as:
0
= =-jjv== =-1
S
G.j G
AttWeight(Q.. K).
(6)
S
St.s
j LCMT
S.. Sipp S
j LCMT
To build intuition, we note that using an empty LOMT: 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 LCMT is not empty the above formulation ensures that the total sum of attention weights is equal to one. During training, we suggest to choose LCMT randomly by selecting each landmark token with probability PeMr = 0.3.
Intuitively, the above formulation should teach the network to attend to the CMT token whenever 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 beacon for signaling the context lacks enough information. That is why we draw 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 a lower (more granular) level whenever it 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 model with CMT on PG19. At inference we set LCMT 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 retrieval is necessary or not. For simplicity, in our implementation we always perform the retrieval but set S to if it 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 particular, around 50% 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.
still utilizing the memory.
E Positional Augmentation
In our main method, we use a special position mapping scheme (See Section 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 settings since one of our main goals is to allow attention to any token, including those that are very far. While development extrapolate is outside the scope of this work, propose a possible solution and provide results of an early investigation. We leave a more thorough assessment and development such schemes for future work.
a positional encoding scheme which
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 positional
16

---

Table 3: Performance on PG19 dataset with 2048 evaluation length and 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 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%
40
Standard
fyup 100
Perplexity
30
25-
500
1000
1500
2000
2500
3000
3500
4000
Prompt Length (Tokens)
Figure 6: Comparison of perplexity on PG19 for different context lengths. The evaluation does use the landmark cache and feeds the whole input in 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 Eseq where is the length of the input. In particular, instead of assigning positions from 1 to T, where I is the length of the input, tokens by a random integer between 1 and Pjump after each landmark token. We refer to these increases making positional jumps. When Pjump 1, no augmentation is applied and the standard positions are recovered.
propose to increase the positions of all subsequent
To evaluate our proposal, we train the same model that we used in Section 4.1 on PG19 with Pjump = 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. Pjump 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 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 expected and justified since the model learning harder task and therefore may require additional steps to reach the same performance.
17

---

k = 2
3.25
k 4
3.00
k
2.75
2.50
2.25
2.00
1.50
20 40
60
80 Block Size
100
120
Figure 7: Time needed to process 4870 tokens (excluding the landmarks) using Landmark Attention for different number of retrieved blocks & and different block sizes. The efficient implementation of landmark attention combined with FlashAttention is used. The retrieval flexibility is reduced that all tokens in each chunk use the same set retrieved blocks.
F Additional Extensions and Details
Masked Language Modeling In this work, we focus on next-word prediction tasks. this section, briefly discuss how to adapt 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 be used train 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 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 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 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 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).
kept and adjusted each step.
Trade-off between number 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 blcck size Thus increasing block size reduces memory usage while also reducing the time needed to find relevant blocks since there are less blocks. However, the number of retrieved blocks & 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 's maximum allowed context length, i.e. the train context length. Therefore, increasing the block size increases the number of chunks the input has be broken into, which slows down the inference. This trade-off can be in the Figure 7, showing the inference time for different block sizes and different values k.
18

---

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 address this bottleneck. However, using landmark attention allows to reduce memory usage further using 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 of retrieved blocks to vary across heads but 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 selecting the maximum, softmax is applied to ensure all scores are on the same scale.
blocks for retrieval. Prior to
Using the above method, we evaluated the fine-tuned LLaMA with context length of 32070 tokens (the exact number of tokens may vary slightly due to random generation, but the difference less than 10 tokens). Across 50 randomly generated prompts, achieved an encouraging accuracy of 98% in retrieving the pass key. This clearly demonstrates the model 's capability to effectively operate context lengths comparable to those supported by GPT-4 [25].
19


Reviews:
[{"review_id": "IzGglcJnAx", "reviewer": "Reviewer_DxW4", "summary": "The authors propose a new transformer architecture, the Landmark Transformer, which uses a novel approach to the self-attention mechanism to allow the model to handle longer sequences.\n\nThe Landmark Transformer introduces a new type of token called a \"landmark token\" which acts as a gateway to a block of tokens. The blocks are assigned to these landmark tokens in a way that allows the model to handle longer sequences while maintaining computational efficiency. The landmark tokens are distributed across the sequence and form a part of the model's input. They are processed in a particular order and are used to control the flow of attention in the transformer model. In some ways, this could also be viewed as introducing hierarchy (and or tree-structures) into the context window. They provide a higher level abstraction over local sets of tokens.\n\nThe authors propose a \"Grouped Softmax\" operation, which is used to calculate attention scores in the self-attention mechanism. This operation allows the transformer to regulate its attention across different blocks of tokens. The Grouped Softmax operation is applied to both the query and key vectors in the self-attention mechanism. The authors also propose an efficient method for calculating these Grouped Softmax operations.\n\nThe authors evaluate their proposed models on the PG-19 dataset and an arXiv math dataset. The results demonstrate that their models can handle longer sequences while maintaining good performance when compared to the baseline .", "questions": "There is a typo on line 185: exmaple -> example\n\nwhere these 80GB or 40GB version of the a100?\n\ncould you provide a graph of how memory/compute/time scales with block size and length?\n\nExperiments to demonstrate the memory offloading capabilities of this approach would improve this work.\n\nIt would alse be good to benchmark against the memorizing transformer and similar competing appraoches.\n\nI will improve my score if these questions are answered and the weaknesses raised are addressed.", "limitations": "There is no discussion of the social impact of this work, it should be added though.", "rating": 8, "confidence": 5, "soundness": 4, "presentation": 3, "contribution": 4, "strengths": "The paper offers a novel and simpple approach to extending the context length of transformer models.\n\nThe multi-resolution nature of this approach allows for exact retention of fine-grained details and is very amenable to model interpertability approaches.\n\nThis method can be applied post-hoc, is extensible to tools/frameworks that can better take advantage of CPU processing and system RAM.\n\nThis approach can be used to directly build document retrievers without any further training required.\n\nThe approach is also quite efficient, requiring at most 4 A100 GPUs to achieve the results presented in this work.\n\nThis approach has the potential to enable heirarchical scaling of the context window, which when combined with an external cache could enable quadratic/exponential growth of the context window.\n\nThe illustrations in the article help alot in conveying the core ideas.\n\nOverall, a solid, timely and well-justified piece of work.", "weaknesses": "There should be more diversity in the models and datasets benchmarked against in this paper.\n\n\nSections 3.1 (presenting grouped softmax) and the middle paragraph os section 3.2 are confusing and less clear than the rest of the article.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Dear Reviewer DxW4,\n\nWe make the following comments to address your questions and concerns:\n\n1. The language modeling experiments were done using 4 A100 40GB GPUs. LLaMA fine-tuning was done using 8 A100 80GB GPUs.\n\n2. We have successfully used offloading parts of KV cache to CPU to perform inference at 32k context lengths using LLaMA. We have described this result in Appendix G. \n\n3. Here is an overview of the effects of increasing block size or context length during training and inference.\n    * During training, the block size has almost no effect on time or memory usage since we keep the context size (including the landmarks) fixed. Since we use the standard training procedure for transformers, computation time scales quadratically with the context size (during training). We note that, a future work can look into using a similar retrieval method used at inference during training to reduce the training computation by block size as well.\n    * 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 1/(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 1 of the PDF attached to our global rebuttal, showing the inference time for different block sizes and different values of $k$.\n    * When increasing the context length, the standard transformer’s memory usage and operation time grows quadratically since it needs to compute the full attention matrix. When using landmark attention, the attention matrix is only computed for a fixed number of blocks, which means the total memory usage and operation time for computing the attention matrix increases only linearly. However the bottleneck becomes finding the set of relevant blocks. As we mentioned it is possible to use a kNN data structure to do this operation efficiently as well but we leave exploration of such implementation for future work. With the direct (not kNN) implementation we used for our experiments, the memory usage and operation time of this operation is also quadratic in context length but is reduced with the noticeable factor of block size (e.g. 50x reduction). The implementation still needs to be done efficiently for this improvement to be observable. Finally, we note that using flash attention, it is possible to reduce the quadratic memory usage for both our method and standard transformer to linear. We have already implemented this version of our method in Triton and have seen its success. Please note that the reduction of operation time by the noticeable block size factor remains even with flash attention.\n\n4. This unique characteristic of our method sets it apart from previous approaches like Memorizing Transformers, which require expensive training with large context lengths. As a baseline, we implemented Transformer-XL to showcase that we can achieve comparable performance in utilizing long contexts. Comparing against additional baselines becomes a challenging task considering the increased training costs and the complexity of implementing such methods (e.g., Memorizing Transformers rely on FAISS during training). \n\n\nWe hope that the above replies adequately clarify our contributions and that you would consider raising your score. We would also like to ask that you read the general reply we have provided which addresses common concerns raised by the reviewers. We remain at your disposal to answer any additional questions or comments.\n"}, {"review_id": "4fBNRsK9IB", "reviewer": "Reviewer_Hd8A", "summary": "The paper proposes a new attention mechanism that uses \"landmark\" tokens to allow access to long contexts while retaining the flexibility of standard attention. The landmark token represents each block of the input context, and the attention mechanism is trained to use landmark tokens to select relevant blocks. This allows retrieving blocks directly through attention instead of separate retrieval mechanisms. Experiments including language modeling and fine-tuning LLaMA 7B show landmark tokens' effectiveness.", "questions": "- In Table 2, for different levels of retrieval flexibility, the change of perplexity is small. Does this mean the choice of $k$ is not that important?\n- What is the effect of doing memory-offloading for unused KV cache? To what extent it can save memory consumption?\n- Can you explain and compare the choice of pre-training the landmark LM from scratch and finetuning a landmark LM based on the existing LM like LLaMA.\n- The line spacing between line 339-340, and line 347-348 is too narrow.\n", "limitations": "The authors have not addressed their limitations officially in the conclusion but admitted the computation limitations in the middle of the paper and also stated their future work.", "rating": 5, "confidence": 5, "soundness": 2, "presentation": 3, "contribution": 3, "strengths": "- The landmark design is interesting, which maintains the random-access flexibility of standard attention, unlike other retrieve-based approaches. It enables the possibility of processing arbitrarily long contexts and practically extends the context length of large LMs like LLaMA by 5x.\n- This method is able to reduce memory footprint since only landmark tokens need to be stored in GPU and some KV cache can be offloaded to CPU. Besides, the landmark approach needs fewer retrieved tokens per step, compared with Transformer-XL.\n", "weaknesses": "- One of the advantages of the landmark approach in the paper, is to reduce memory usage by swapping out (to CPU memory or to disk) all regular tokens’ cached key-value vectors (line205). However, this is not tested in experiments, as stated in line 56: \n>For simplicity, we focus our experiments on storing everything in GPU memory, but we note that the above techniques can be directly applied in large-scale settings.\nIt is better to add the related experiments about memory saving to validate the effectiveness of this approach.\n- The language modeling experiment is not that persuasive. The perplexity of the landmark approach is not better than Transformer-XL, although it uses a smaller attention size. It is still not sure whether the language modeling ability as well as the memory consumption is supreme to transformer-XL.\n- The experiment of finetuning LLaMA 7B is not that persuasive, LLaMA-2k is not good at extrapolation, and this baseline is really weak. Besides: i) only test Retrieval tasks, lack of other common long document understanding/generation tasks (e.g. scrolls benchmark or many-shot in-context learning); ii) there are many other approaches targeting long document processing (listed in Related work), lack of experiments to compare some of these baselines.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Dear Reviewer Hd8A,\n\nWe make the following comments to address your questions and concerns:\n\n1. In Appendix G, we demonstrate how offloading the KV cache to the CPU, along with reduced retrieval flexibility, enables us to extend LLaMA's context length to 32K. As a result, we achieve 98% accuracy in retrieving the pass key. Without our method, performing inference at such long context lengths would be impossible, as the KV cache grows and cannot fit in a single A100 80GB GPU.\n\n2. We present Transformer XL as a baseline to show that landmark attention can perform comparably in utilizing long contexts. However, despite performing slightly better, Transformer XL has inherent limitations as it cannot directly access earlier tokens, hindering its ability to perform certain tasks, such as retrieval. In contrast, our LLaMA experiments demonstrate that Landmark Attention successfully retrieves and attends to early tokens. \n\n3. As you mentioned the original LLaMA 7B was trained on 2k tokens and can only perform inference up to this length. We want to demonstrate that fine-tuning using landmark attention (at 512 context length) makes the model capable of handling arbitrarily long context lengths. Our experiments clearly verify this for 32k tokens. Please note that all existing methods for handling long contexts usually require training at the target inference length which increases the training cost or sacrifices the ability to directly attend to any tokens. \n\n4. We point out that most existing methods for retrieval use a separate embedding model to perform the retrieval which needs to be trained separately, increasing the training cost. In contrast, our method allows using the model itself to perform the retrieval. Furthermore, many of the proposed retrieval-based methods usually focus on retrieving whole documents rather than finding relevant parts of the input. Extending these methods to be capable of identifying small relevant blocks of the input is beyond the scope of our work and requires further research. Other existing methods for augmenting transformers with memory or long context capabilities either remove the capability to directly attend to each token , which is unlikely to perform well for retrieving an exact pass key, or involve a costly training method which is challenging to perform for a model as large as LLaMA. We have also left further evaluation of the fine-tuned model on other tasks such as summarization as future work given the resource limitation.\n\n5. When the set of retrieved blocks is allowed to change for each token and each head, increasing k beyond certain values can have a negligible effect since softmax is applied over the landmarks of the retrieved blocks. For example tokens in the 10-th top block can only have a weight of at most 0.1. However, in a less flexible regime where the retrieval set remains the same across all tokens (but can vary across heads), k can have a more noticeable effect. For example, as shown in Table 2, increasing k from 2 to 4 in this setting improves performance from 15.48 to 15.10. \n\nWe hope that the above replies adequately clarify our contributions and that you would consider raising your score. We would also like to ask that you read the general reply we have provided which addresses common concerns raised by the reviewers. We remain at your disposal to answer any additional questions or comments. "}, {"review_id": "Z0M17dIAod", "reviewer": "Reviewer_fLCX", "summary": "This works proposed a hierarchical structure to organize previous context in blocks and represent them via novel landmark tokens at the end of each block. Additionally, a novel GroupedSofxmax mechanism is proposed to replace the original softmax to enable current token to attend on both local tokens and retrieved previous contexts. Also, a novel positional encoding method is proposed to avoid the positional confusion due to retrieved previous contexts. The author demonstrate the effectiveness of proposed method on several language model datasets.", "questions": "The presentation of the whole paper is perfect. I think I have understand the details of the paper well thanks to the great presentation. The only issue regarding the paper is that most of references lack of publication journal/conference and the publication year.", "limitations": "N/A", "rating": 5, "confidence": 4, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "1. Different from the retrieval-based attention on all previous cached contexts in Memorizing Transformer, the proposed method proposes a novel way to split the previous context into blocks and attend on the well-organized blocks based on GroupedSoftmax on the landmark token representation of each token. The whole design is intuitive and feasible.\n\n2. I do appreciate the related work section of this paper. This section discuss almost all feasible solutions on extending context length of Language Models. ", "weaknesses": "1. My major concern towards the proposed method is that it is too trivial to implement, thus leading to the lower training stability and higher computational timecost. Now the authors store everything like the landmark token representations, previous contextual representations in GPU memory for simple implementation, and thus for each attention layer the stored vectors will take up a lot of GPU memory. The authors talk about the future work regarding storing cached vectors in CPU memory or even disks. However, the computation for the proposed GroupedSoftmax is still deployed on GPU. The major efficiency bottleneck will be the loading and offloading process from GPU to CPU or from CPU to GPU, as this process is several times longer than that of attention computation timecost. Additionally, implementing these cached index on disks seems like impossible as the data transfer timecost from disk to CPU and then to GPU is unacceptable for any real-time applications. Thus, analysis and statistics towards the introduced additional timecost and GPU memory cost will be appreciated.\n\n2. My second concern is that the proposed method introduces too many hyperparameters to control the granularity of the block architecture, including the local context length, the block length, number of blocks, number of retrieved blocks, and attention size. The author should make a constant decision on which group of hyperparameters to use in all experiments or propose a brute-force strategy to select such group of hyperparameters.\n\n3. You only follow Memorizing Transformer to compare the proposed method to Transformer-XL. However, the memorizing transformer, as the most important baseline to your methods, should be also compared with. The memorizing transformer can be regarded as a special case of your model when number of block is 1, local context length is eval length, number of retrieved block is 32, block length is 1 token, attention size is the local context length plus the number of retrieved blocks, and original softmax and attention is applied. ", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Dear Reviewer fLCX,\n\nWe make the following comments to address your questions and concerns:\n\n1. We did not observe any training instability due to using landmark attention. Furthermore, as mentioned in the paper, the additional computational costs are negligible especially when combined with Flash Attention. This is because our method facilitates performing inference at any length regardless of the train context length, removing the necessity of a costly training procedure for long context inference. Choosing the alternative approach of training at the target inference length would be much more expensive and scale quadratically with the length during both training and inference.\n\n2. In Appendix G, we demonstrate how offloading the KV cache to the CPU, along with reduced retrieval flexibility, enables us to extend LLaMA's context length to 32K. As a result, we achieve 98% accuracy in retrieving the pass key. Without our method, performing inference at this context length would be impossible, as the KV cache cannot fit in a single A100 GPU.\n\n3. The local context length and the attention size are directly determined based on the model’s training context length (which already is a hyperparameter that has to be chosen and depends on the available resources) and the number of retrieved blocks k. These values are provided in Table 1 for simplicity and better comparison. Therefore, the method only introduces two new hyperparameters which are the block size and the number of retrieved blocks. We note that only the block size needs to be chosen at training time and even this can be mitigated by providing the model with blocks of different sizes at training. In early stages of our work, we briefly experimented with dropping some of the landmark tokens, thus merging consecutive blocks, and did not find a noticeable effect on performance. For simplicity we decided to avoid doing this for the final experiments. Overall we do not expect much tuning to be required for the choice of the block size and this value can be chosen based on common sense (keeping the block small enough so that retrieving a single block does not fill the context with unimportant tokens and it can be summarized effectively, but keeping it large enough to use the computational benefits). Since the optimal k can be chosen at inference, it is easy to tune. Our results demonstrate the trade off between using a larger and smaller k. Finally, increasing k beyond a certain point has negligible effect since softmax is applied over the landmarks of the retrieved blocks. For example tokens in the 10-th top block can only have a weight of at most 0.1. \n\n4. We would like to emphasize that a very important difference between Memorizing transformers and our method is the need to train the model at the large context length which either increases training cost significantly or makes the implementation more complicated since it requires connecting the model to FAISS data structure. The increased training cost makes it challenging to train such a model for large contexts. In contrast, our method allows inference at arbitrary context lengths regardless of the training context length. As a side note, we do not expect the stingy mapping to work for extremely small blocks (such as block size 1) as in that case mapping earlier blocks all to the same position in the stage of finding relevant blocks resembles searching in a bag of words. \n\nWe hope that the above replies adequately clarify our contributions and that you would consider raising your score. We would also like to ask that you read the general reply we have provided which addresses common concerns raised by the reviewers. We remain at your disposal to answer any additional questions or comments. \n\n"}, {"review_id": "KEaQCcydwm", "reviewer": "Reviewer_g7s5", "summary": "This paper presents a novel approach that enables the Transformer language models to process much longer sequences. Specifically, the authors propose to group the input sequence into multiple blocks, each of which is represented with a landmark token. The attention scores are calculated regularly among all tokens, but are normalized within each block and multiplied by the landmark scores. In this way, the attention mechanism is trained to identify and select relevant token blocks given each query. As a result, during inference, the model could dynamically select relevant context blocks for each incoming query without loading all previous contexts into memory. Combined with a careful design of positional encodings, the model is allowed to process arbitrarily long contexts and demonstrates its effectiveness in language modeling tasks across different model scales.", "questions": "1. Regarding LLAMA fine-tuning experiments, it is known that LLAMA performs poorly when it comes to length extrapolation, possibly due to positional encodings. Consequently, it remains unclear whether the improved results depicted in Figure 3(b) are solely attributed to the proposed model's enhanced ability to handle longer contexts or are primarily a consequence of the utilization of more robust positional encodings. This issue might require further investigation.\n2. During inference, token blocks within the memory are frequently retrieved and replaced due to the attention landmark selection. Intuitively, such swapping-in/out processes can be computationally expensive, especially in the case of per-head and per-token selection. Does this case hold true in practical settings and are the benefits of the proposed approach only apparent in the long sequence regime?\n3. How well does the model scale with longer sequences during training in terms of evaluation performance (e.g., language modeling perplexity)?", "limitations": "The authors have addressed the limitations; my suggestions can be found above.", "rating": 5, "confidence": 5, "soundness": 2, "presentation": 2, "contribution": 3, "strengths": "- The idea is novel and interesting in that it associates each contiguous block of tokens with a pointer; each incoming query is then compared against these pointer vectors (as well as tokens within the local neighborhood) and retrieves past token blocks only when the corresponding pointer is semantically relevant. A notable advantage of this approach is that the representation of these pointer vectors is learned through the attention mechanism, thereby reducing the need for heuristic reductions and enhancing the overall simplicity of the model.\n- Furthermore, the exploration of positional encodings in this paper is a valuable contribution. While it is not new to investigate the impact of positional embeddings on length extrapolation, this research offers a comprehensive analysis of this phenomenon and proposes several techniques to alleviate related challenges. The empirical findings presented here not only provide insights specific to the context of this study but may also hold relevance for the general language modeling community.", "weaknesses": "- The paper is not that easy to follow. For example, the formulation of Equation (1) introduces confusion. The subscript $i$ in Equation (1) actually indexes the $i$-th element of vectors $v$ and $g$, rather than representing the $i$-th query as in Equation (2,3,4). I would suggest using a different subscript other than $i$ to enhance clarity. In addition, it seems that $(k+1)\\cdot \\ell_{\\text{block}}$ in L234 should be $(k+1) \\cdot (\\ell_{\\text{block}} + 1)$, according to the specification in L200 that “augmented by a landmark token after every $\\ell_{\\text{block}}$ tokens”.\n- The proposed model is not evaluated sufficiently in the following aspects:\n    - 1) Some baselines are missing. In my opinion, it would be beneficial to compare against other long-range methods such as Combiner (as mentioned in the paper), which processes the distant context with heuristic pooling operations.\n    - 2) Ablation studies are not sufficiently extensive. For instance, this work does not examine the impact of positional encodings and block sizes in modeling long contexts;\n    - 3) The conducted experiments do not provide qualitative or quantitative results supporting the claim such as “In contrast with Transformer-XL, using our method, the information retrieval is interpretable”?\n- One notable advantage of this proposed model is the ability to efficiently process long contexts by enabling random access to tokens. However, this work does not demonstrate the efficiency gains (e.g., memory usage or decoding-time speedup) achieved in dealing with long contexts. Providing evidence of these gains would further strengthen the paper's contributions in this regard.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Dear Reviewer g7s5,\n\nWe make the following comments to address your questions and concerns:\n\n1. We thank you for the valuable suggestions to increase clarity (and also for pointing out the typo) and will apply them in the final revision.\n\n2. Please note that all existing methods for handling long contexts usually require training at the target inference length, which increases the training cost, or sacrifices the ability to directly attend to any tokens, which are unlikely to perform well for retrieving an exact pass key. \n\n3. We want to emphasize that throughout the paper we always use the standard rotary positional encoding. The stingy positional mapping is solely applied during inference and only affects how tokens are indexed, not how they are encoded for the model. Furthermore, this mapping can only be applied in combination with a method such as landmarks which identifies relevant previous blocks. It is not a separate method that can be directly used with LLaMA or other models. Therefore, the study of the effect of different positional encodings on context length is out of the scope of our work though we discuss the limitations of the existing methods at the final paragraph of the related works. \n\n4. Overall we do not expect our method to be very sensitive to the choice of the block size. Instead, this value can be chosen based on common sense (keeping the block small enough so that retrieving a single block does not fill the context with unimportant tokens and it can be summarized effectively, but keeping it large enough to use the computational benefits). We note that the need to choose a fixed block size also can be mitigated by providing the model with blocks of different sizes at training. In early stages of our work, we briefly experimented with dropping some of the landmark tokens, thus merging consecutive blocks, and did not find a noticeable effect on performance. For simplicity we decided to avoid doing this for the final experiments.\n\n5. Please note that the information in Transformer XL has to be passed through recurrence which prevents finding the tokens the model attends to. In contrast, in our method, the model directly attends to the tokens. In this sense, our model is more 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, i.e. blocks with highest scoring landmarks.  We will clarify this further in the final revision.\n\n6. While swapping out and swapping in can lead to a slowdown, we show in Appendix G that only allowing the retrieved set to change across heads and not across tokens allows us to mitigate this problem and successfully perform inference with LLaMA at 32k inference length with 98% accuracy. \n\n7. We acknowledge that without the reduced flexibility, i.e. in per token and per head setting, the CPU-GPU communication could become the bottleneck. However, we note that due to limitations, we did not implement caching of the retrieved blocks in GPU to avoid double communication which we expect to significantly reduce the load on CPU-GPU connection. Regardless, we emphasize that our method still improves over previous methods and allows the model to operate at arbitrary inference context lengths (though slowly) despite being trained at a much smaller context length. Therefore, using landmarks, inference at larger context lengths becomes a question of speed or having additional resources at inference time instead of the need to re-train the model.\n\n8. In terms of perplexity, we expect our method to behave the same as the standard transformer as the context length grows. While this is not included in the paper, we can confirm for example that LLaMA 7B can also be successfully trained at 2k context length using a more efficient implementation of our method (combined with Flash Attention).  Finally, a similar method used at inference can also be applied at training to allow direct training at larger context lengths as well while reducing the quadratic computational cost by the block size factor (though we have not yet implemented this).\n\nWe hope that the above replies adequately clarify our contributions and that you would consider raising your score. We would also like to ask that you read the general reply we have provided which addresses common concerns raised by the reviewers. We remain at your disposal to answer any additional questions or comments. \n"}, {"review_id": "6NfSrbsJiM", "reviewer": "Reviewer_tGXa", "summary": "The paper presents a new architecture for long-range decoder-transformers: a new \"landmark token\" is inserted into the input within constant strides (that is, after every $k$ tokens). Every such \"landmark token\" is thus the \"representative\" of a block of tokens.\nThe attention score to this \"landmark token\" is treated as a multiplicative gate for attending the tokens within its block.\n\nThat is, during training: tokens in the local proximity are attended to as usual; the attention score of far away tokens is their usual token-attention probability, multiplied by their landmark's attention probability.\n\nDuring inference, the model attends to the local proximity tokens (the last k tokens) and to the previous landmarks, chooses the top-k landmarks, and attends only to the local block's tokens (the most recent tokens) and to tokens within these top-k blocks.\nIn other words, the model always attends to the most recent tokens, but performs a kind of \"hierarchical attention\" to the past tokens.\n\nThis allows, after fine-tuning, extending LLaMa 7B from 2k to 10k inputs.", "questions": "## Questions\n1. What is the metric in Table 1? perplexity?\n2. Am I right to conclude that the proposed approach outperforms only the vanilla base transformer, but not outperforming Transformer-XL?\n3. The paper clearly discusses relevant long-range transformers in the Related Work section. Isn't any of them a relevant baseline?\n4. The only baseline is Transformer-XL, which came out in 2019. Can Memorizing Transformers and TRIME$_{\\mathrm{long}}$ be relevant baselines?\n5. In the LLaMa experiments in Section 4.2, the paper mentions that LLaMA was fine-tuned with a context length of 512 tokens. For fairness, was the proposed model trained with a context of 512 as well?\n\n## Additional Comments\n6. Possibly related paper [Bertsch et al., \"Unlimiformer: Long-Range Transformers with\nUnlimited Length Input\"](https://arxiv.org/pdf/2305.01625.pdf). I refer to Unlimiformer as concurrent work since it appeared on arxiv 2 weeks before the NeurIPS deadline, so a direct comparison is not required.\n7. I think that the explanation of GroupedSoftmax and the entire Section 3.1 can be simplified, as it explains the operator from the implementation point of view, rather than from the conceptual idea. I think that it could have been explained more simply by just saying that \"every token attends only to the tokens within the same block, and to the other landmarks of other blocks\", without bothering the reader with the $\\mathbf{g}$ indices. Papers such as [Longformer](https://arxiv.org/pdf/2004.05150.pdf) also implement attention with masking and summing over indices, but the explanation in the paper is much more intuitive (see Figure 2 in [Longformer](https://arxiv.org/pdf/2004.05150.pdf)).\n5. Equation (2) uses $\\mathbf{G}$ rather than $\\mathbf{g}$, unless $\\mathbf{G}$ is something else; in that case, $\\mathbf{G}$ was not defined.\n\n## Summary\nI like the overall approach and I think that it can be very useful in practice. However, there are a few significant issues: \n1. Overclaiming, and claiming abilities that were not demonstrated, but only \"theoretically possible\"\n2. Lack of baselines - the only baseline is Transformer-XL, which performs better (!) than the proposed approach\n\nI am thus voting for only a borderline acceptance at this time. I would increase my score if the above questions, issues and weaknesses would be resolved.", "limitations": "NA", "rating": 5, "confidence": 4, "soundness": 2, "presentation": 3, "contribution": 3, "strengths": "## Strengths\n* The approach allows finetuning the newest models (such as LLaMA) to extend their input length significantly\n* The discussion of related work is thorough, albeit with some inaccuracies regarding Memorizing Transformers (see Weaknesses below).\n", "weaknesses": "## Weaknesses\n- The authors claim \"infinite context length\", but demonstrate only on ~10k tokens. The authors write that \"it is theoretically possible for the model to access any token in the entire past\" - but this was not demonstrated. Even if the authors argue that this is only a matter of missing engineering, it needs to be demonstrated that the model can actually leverage the attention to the entire past in an effective way to claim \"infinite context length\" (although I believe that even solely the engineering is not as simple as the authors try to make it sound).\n- The paper notes that a model that was trained this way can be used as a general-purpose document retriever, but this was not demonstrated as well.\n- The authors make some claims about the most related work, [Memorizing Transformers (Wu et al., ICLR 2022)](https://arxiv.org/pdf/2203.08913.pdf), that are inaccurate:\n> Memory Transformers [33] ... However, while these methods improve upon the memory-less variants, they do not allow for attending to *specific* tokens in the past, as the model only has access to a compressed version of this information\n\nThis is incorrect: Memorizing Transformers [33] do allow for attending to specific tokens in the past, by indexing and retrieving them in a kNN index. The main difference may be that in Memorizing Transformers, the trained gate is applied to the entire attention head, rather than to the local \"landmark\". Thus, I also think that the claim about having access only to a compressed version of the information is incorrect - later in the Related Work, the paper says: \n> Memorizing Transformers [33] performs... However, these methods obtain the final results by interpolating between the kNN prediction and the local attention prediction using a tuned parameter as interpolation weight\n\nThis is again incorrect. Memorizing Transformers does not perform interpolation with a tuned parameter, it attends to the retrieved kNN tokens with a learned attention, as in standard transformers (the claim in the paper is correct regarding kNN-LM, but not to Memorizing Transformers)\n\n* The explanation about positional encoding (Section 3.2.1) is very unclear. What does \"allocate a segment\" mean? where is this allocated and why? What does \"*we index the current chunk starting after this segment*\" mean? I could not parse sentences such as \"*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*\", although I am very familiar with the long-range transformers literature and with Rotary positional embeddings.\n* Evaluation\n    * Table 1 shows the language modeling evaluation. \n        * What metric is this? The table only shows numbers without mentioning the metric. is this perplexity?\n        * From what I am able to understand, in the longest evaluation length of 4096, Transformer-XL [7] performs better than the proposed approach? (achieving perplexity of 14.55 on PG19 which is lower than \"Ours\"). \n        * In evaluation length of 2048, Transformer-XL [7] also performs better than \"Ours\", achieving perplexity of 14.72.\n        * Am I right to conclude that the proposed approach outperforms only the vanilla base transformer, but not outperforming Transformer-XL?\n        * The authors argue that \"*In contrast with Transformer-XL, using our method, the information retrieval is interpretable. Particularly, it is possible to understand which parts of the text was recovered to generate a certain answer*\" - however this interpretability was not demonstrated, so I did not take it as an advantage.\n    * The experiments with LLaMA-7B were demonstrated only on a synthetic task\n    * baselines: \n        * The only baseline is Transformer-XL, which came out in 2019.\n        * [Zhong et al., \"Training Language Models with Memory Augmentation\", EMNLP '2022](https://arxiv.org/pdf/2205.12674.pdf) discusses several retrieval scenarios - retrieving from an external datastore, and using kNN retrieval for long inputs. I believe that their TRIME-long is a very relevant baseline. I agree that TRIME-long is conceptually more limited because it uses a tuned interpolation hyperparameter as in kNN-LM, but the insight there of training with retrieval *from other examples in the same batch* may compensate for it.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Dear Reviewer tGXa,\n\nWe make the following comments:\n\n1. Our argument on infinite context length relies on the model's capabilities. While practical demonstration of infinite context length isn't feasible, our method's efficacy is shown in Appendix G at 32k context length. The original transformer remains crippled by the inability of positional encodings to fully generalize to arbitrary length. Current solutions sacrifice direct attention to early tokens to address this. However, our approach, using stingy positional mapping, eliminates this and reduces quadratic complexity by the block size factor.  While inevitably going to larger lengths requires more resources, we show that this can be further alleviated by offloading the cache to CPU which is what we use in Appendix G. Our experiments reveal the model's ability in key retrieval despite stingy positional mapping, justifying the term \"infinite.\" We welcome any suggestions to make our argument more concise.\n\n2. We only note the ability to use for retrieval as a future prospect. We will clarify this in the final revision.\n\n3. Memory Transformer is different from Memorizing Transformer. By mistake, we have referenced the wrong paper in Line 79. We will fix this in the final revision and cite Burtsev et. al. (2020). Indeed Memory Transformers do not allow for attending to specific tokens in the past while Memorizng Transformers do. \n\n4. We refer to Section 3.1 of the Memorizing Transformer paper where the interpolation is discussed. The knn-augmented layer in memorizing transformers similar to kNN-LM performs an interpolation between the memory and the LM output to obtain the final output.\n\n5. Please note that we still use standard Rotary Positional Encoding in our model. Our approach solely modifies token indexing (not the encoding) to eliminate the need for extrapolation to indices unseen at training. Here is an alternative description: We feed the input to the model in chunks; as in the paper. So the total number of tokens in the chunk and in the retrieved blocks should be less than or equal to the training context length $l_{\\text{train}}$. So we introduce $l_{\\text{train}}$ placeholders corresponding to indices 1 to $l_{\\text{train}}$ to be filled with tokens from either the chunk or retrieved blocks, as shown in Figure 2. By allocating a segment in the beginning of the input we mean reserving some of the early placeholders. The placeholders beyond the reserved segment are filled with the chunk tokens left to right. Note that the token to placeholder (index) allocation differs between finding relevant blocks and attending to them. When attending to relevant blocks, those that are among the last $k$ blocks fed to the model are right-aligned in the reserved segment while others are left-aligned. Given that the reserved segment spans $k + 1$ blocks, this leaves one empty block between the two groups. We refer to Figure 2 for a visual representation of this process.\n\n6. As specified in Line 282, the numbers in Table 1 are perplexity scores. We will make it clear in the caption of the Table in the final revision. \n\n7. We present Transformer XL as a baseline to show that landmark attention can perform comparably in utilizing long contexts. However, Transformer XL has inherent limitations as it cannot directly access earlier tokens, hindering its ability to perform certain tasks, such as retrieval. In contrast, our LLaMA experiments demonstrate that Landmark Attention successfully retrieves and attends to early tokens.\n\n8. Please note that the information in Transformer XL has to be passed through recurrence which prevents identifying the tokens the model attends to. In this sense, our model is more 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, i.e. blocks with highest scoring landmarks.  We will clarify this further in the final revision.\n\n9. We want to strongly emphasize that **in our method each token can attend to any other token**. In our method, a token attends to specific tokens in other blocks as well. However, the attention scores to those tokens are gated by their block's landmark token's score (through multiplication as in Eq 4). We provide high level intuition for the method in the introduction and Fig 1 but formally define our method in Sec. 3. \n\n10. We do not train a separate model in LLaMA experiments and only compare a fine-tuning of LLaMA using our method at length 512 with Meta's vanilla LLaMA (not fine-tuned) at length 2048, our baseline. If we have misunderstood your question, please let us know.\n\n11. We emphasize a key distinction between Memorizing transformers and our approach: Memorizing transformers require training with a large context length. This increases training cost or complicates implementation due to the need for linking the model to FAISS data structure. Due to the increased training cost it is challenging to train such models for large contexts. In contrast, our method allows inference at arbitrary context lengths regardless of the training context length. We tried training our model using the TRIME-Long method. Unfortunately, the training collapsed after 140k steps. At that point, Landmark Attention noticeably outperforms TRIME on PG-19 dataset in addition to being cheaper to train. The perplexity with 4096 tokens is 15.54 for landmarks and 17.43 for TRIME. We also note that the Trime-Long method is inherently limited since the model can not directly attend to the past context and intuitively it can only match close repetitions. \n\n12. $g$ is any vector and denotes one of the inputs of GroupedSoftmax in Equation 1. $G$ is the specific grouping matrix defined in Equation 2. \n\nWe hope that the above replies adequately addresses your concerns and that you would consider raising your score. Please also read our general reply. We remain at your disposal to answer any additional questions or comments."}]

Author rebuttal:
Dear Reviewers,

We would like to address some of the common concerns raised by the reviewers with the following comments:

* We have successfully used offloading parts of KV cache to CPU to perform inference at 32k context lengths using LLaMA. We have described this result in Appendix G. Note that while offloading to CPU might make the inference a bit slower, without our method performing inference at such long context lengths becomes impossible on a single A100 80GB GPU since it is not possible to store the KV cache.

* One of the major advantages of Landmark Attention is its ability to decouple the training and inference context lengths. This means that the model can be trained at a limited context length, thus avoiding increased training costs. The trained model can then be effectively used for making inferences at any context length. We have extensively demonstrated this feature in our language modeling experiments where we train our models at context length 512 and successfully infer at much larger context lengths, such as 4096. Similarly, with LLaMA, we fine-tune at context length 512 using a model originally trained at 2k context length, and we can perform inference at a much larger 32k context length with excellent results. 

* This unique characteristic of our method sets it apart from previous approaches like Memorizing Transformers, which require expensive training with large context lengths. As a baseline, we implemented Transformer-XL to showcase that we can achieve comparable performance in utilizing long contexts. Comparing against additional baselines becomes a challenging task considering the increased training costs and the complexity of implementing such methods (e.g., Memorizing Transformers rely on FAISS during training).

* We note that while our method performs comparably (but not better than) Transformer-XL it has the advantage of being able to directly attend to all previous tokens. This is important to utilize the full power of attention, for example allowing our method to excel at retrieving tasks as demonstrated in our LLaMA experiment.

In addition to these general comments, we have provided individual responses to the reviewers to address their specific concerns. We hope that these explanations further clarify the contributions of our work. We are available to address any further comments or questions you may have.


Meta-review: Meta Review for Random-Access Infinite Context Length for Transformers

As reviewer DxW4 had summarized, this work proposes a new transformer architecture, the Landmark Transformer, which uses a novel approach to the self-attention mechanism to allow the model to handle longer sequences. The Landmark Transformer introduces a new type of token called a "landmark token" which acts as a gateway to a block of tokens. The blocks are assigned to these landmark tokens in a way that allows the model to handle longer sequences while maintaining computational efficiency.

Several innovations had been proposed, including the grouped softmax operation and its efficient calculation. This operation allows the transformer to regulate its attention across different blocks of tokens. The Grouped Softmax operation is applied to both the query and key vectors in the self-attention mechanism.

This work clearly adds value to the community, especially in the context of improving generative AI. There are no clear negative issues or weaknesses. Looking at the feedback, 4 reviewers gave borderline accept (5), and one high confidence review gave it a strong accept (8). The reviewer's feedback has been incorporated into the manuscript, and I believe it will be a paper that will offer good value and interest to the NeurIPS community.
Decision: Accept (poster)