SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Amirkeivan Mohtashami
amirkeivan.mohtashami@epfl.ch
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/.
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 a form of recurrent memory inside the Transformer architecture, such as Transformer-XL [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 a knowledge base and adding them to the context. However, this approach requires a separate mechanism to identify relevant documents, called a retriever. Such retrieval models can not easily be updated to work on fresh long input data, and furthermore are also not fully compatible with the standard attention mechanism itself, and thus may fail to 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
Figure 1: An illustration comparing standard attention and our attention with landmarks. The example shows the (causal) attention given by a current token e to previous ones, illustrating our mechanism with block-size  $\ell_{\mathrm{block}} = 2$ . 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 b can have a high (green) attention score when being part of one block and a 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 at the intermediate layers.
of fixed length and introduce a special token for each block, called a landmark, which acts as a gate for attending to 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 it 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 use the landmark tokens to select the most relevant blocks and only use these blocks for computing the attention. This immediately reduces the computation cost by a factor of block length. For example, in our experiments where we use blocks of 50 tokens, this translates to almost  $50\mathrm{x}$  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, our 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 a 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 a model trained to operate at the given context length) by a substantial factor equal to the block size (e.g., 50 in our experiments).
- Compatibility with advanced data structures that can further decrease the resource requirements for operating the model with very large context lengths.
With the evolution of state-of-the-art 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.
Retrieval-Augmented Language Models. Retrieval-augmented language models use a separate module, called a retriever, to find a 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.
Memory for Transformers. Various methods have been proposed to introduce memorization capabilities to 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 a 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 to benefit from both full access to previous tokens as well as 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 flexibility 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 a random subset of previous tokens as well as several globally accessible tokens [42]. Longformer [3] further introduces dilated sliding window patterns to increase attention's receptive field and manually picks the window sizes for each layer. Linformer [39] uses a low-rank approximation of the attention matrix while Performer [8] uses a non-softmax kernel to obtain a more efficient implementation. Reformer [19] uses locality-sensitive hashing (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-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 is 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 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.
$k$ -NN Augmented Transformers.  $k$ -nearest-neighbor ( $k$ -NN) augmentation has been proposed as an alternative method for allowing transformers to access external memory. For example,  $k$ -NN-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 a nearest-neighbor search over previous keys and computes the attention to the top nearest items. However, these methods obtain the final results by interpolating between the  $k$ -NN prediction and the local attention prediction using a 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.
Context Length Extrapolation. Transformers have a 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 a 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, as 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.
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 to the non-causal case in Appendix F.
When using a Transformer to process a 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 to one, the number of keys with a large attention weight is limited even for long contexts. Thus, by retrieving only those keys with large attention scores, it is possible to 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 a block would lead to a 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 a 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 is 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 a retrieval method directly controlled by attention which can be more semantic-based. Furthermore, retrieving a 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 a landmark token after every  $\ell_{\mathrm{block}}$  tokens. Then we proceed to training the model using the standard batching method which feeds windows of  $\ell_{\mathrm{seq}}$  tokens to the model. In a list of  $\ell_{\mathrm{seq}} > \ell_{\mathrm{block}}$  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 as 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_i$ . If the last block is incomplete and does not have a landmark token, we define  $p_i := \ell_{\mathrm{seq}}$ . If the  $i$ -th token is a landmark token,  $p_i := 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 a token depends on the similarity of the query vector with both the token's key as well as with the key of its block's landmark token. To define the mechanism, we first define a generalized softmax function called Grouped Softmax. Given a vector  $\mathbf{v} \in \mathbb{R}^{\ell_{\mathrm{seq}}}$  and a group index  $\mathbf{g} \in \mathbb{N}^{\ell_{\mathrm{seq}}}$ , Grouped Softmax applies softmax separately over elements
belonging to the same group. (Using  $\mathbf{g} = \mathbb{1}_{\ell_{\mathrm{seq}}}$  recovers the standard softmax function):
$$
\sigma_ {G} (\mathbf {v}, \mathbf {g}) _ {x} := \operatorname {G r o u p e d S o f t m a x} (\mathbf {v}, \mathbf {g}) _ {x} := \frac {e ^ {\mathbf {v} _ {x}}}{\sum_ {y : \mathbf {g} _ {y} = \mathbf {g} _ {x}} e ^ {\mathbf {v} _ {y}}}. \tag {1}
$$
We replace the softmax function after computing the attention scores with Grouped Softmax. For each block, we put its regular tokens in a 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 block is ignored when computing the attention weights for the  $i$ -th token. In 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  $p_j = 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
$$
\mathbf {G} _ {i, j} := \left\{ \begin{array}{l l} p _ {j} & p _ {j} \neq j \\ - 1 & p _ {i} = j \\ p _ {i} & p _ {i} \neq j \wedge p _ {j} = j \end{array} \right. \begin{array}{l} \triangleright \text {p l a c i n g n o r m a l t o k e n s i n t h e i r o w n b l o c k s .} \\ \triangleright \text {i g n o r i n g c u r r e n t b l o c k ’ s l a n d m a r k t o k e n .} \\ \triangleright \text {p l a c i n g o t h e r l a n d m a r k s i n t h e i - t h t o k e n ’ s g r o u p .} \end{array} \tag {2}
$$
Finally, to obtain the final weights after applying GroupedSoftmax, we multiply each token'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 as its attention weight. The weight for landmark tokens is always zero. In more formal terms,
$$
\mathbf {S} _ {i, j} := \operatorname {S o f t m a x S c o r e} (\mathbf {Q}, \mathbf {K}) _ {i} \quad := \operatorname {G r o u p e d S o f t m a x} \left(\frac {\mathbf {Q} _ {i} ^ {\top} \times \mathbf {K}}{\sqrt {d _ {\text {h e a d}}}}, \mathbf {G} _ {i}\right) \tag {3}
$$
$$
\operatorname {A t t W e i g h t} (\mathbf {Q}, \mathbf {K}) _ {i, j} \quad := \left\{ \begin{array}{l l} 0 & p _ {j} = j \\ \mathbf {S} _ {i, j} & \mathbf {G} _ {i, j} = \mathbf {G} _ {i, i} \wedge p _ {j} \neq j \\ \mathbf {S} _ {i, j} \cdot \mathbf {S} _ {i, p _ {j}} & \mathbf {G} _ {i, j} \neq \mathbf {G} _ {i, i} \wedge p _ {j} \neq j \end{array} . \right. \tag {4}
$$
An example 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 (e.g. by setting  $A_{i,j}$  to  $-\infty$  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 a different grouping scheme which also trains a 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 a 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.
Similar to training, the input gets augmented by a landmark token after every  $\ell_{\mathrm{block}}$  tokens. Then, we break the input into chunks of  $\ell_{\mathrm{local}}$  length and iteratively feed chunks from the beginning to the end. To retrieve relevant blocks, each attention layer has access to a 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 to 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  $k$  highest scoring blocks and preprocess the obtained attention matrix to the local attention matrix. We finish by applying
New Input
Final Attention Position Mapping
Figure 2: Stingy position mapping: Retrieving top  $k = 2$  blocks from a memory of 5 blocks. Retrieval landmarks for the last 2 blocks are accurately mapped to sequence index positions, while previous blocks are mapped to the position of the  $(k + 1)$ -th last block. Blocks are then distributed across the prefix based on their position, with an empty block separating the last  $k$  blocks from older ones.
GroupedSoftmax to the attention matrix and computing the weighted average of the value vectors to obtain the token's representation.
Under the above scheme, each token and each head can retrieve different blocks from the cache. It is possible to limit the retrieval flexibility in order to 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  $k$  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