SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Szymon Tworkowski $^{1,3*}$  Konrad Staniszewski $^{1,3*}$  Mikołaj Pacek $^{1,3*}$  Yuhuai Wu $^{6\dagger}$
Henryk Michalewski $^{3,4}$  Piotr Miłos $^{1,2,5}$
$^{2}$ Institute of Mathematics, Polish Academy of Sciences
<sup>3</sup>University of Warsaw
4Google DeepMind
$^{5}$ deepsense.ai
$^{6}$ XAI
Large language models have an exceptional capability to incorporate new information in a contextual manner. However, the full potential of such an approach is often restrained due to a limitation in the effective context length. One solution to this issue is to endow an attention layer with access to an additional context, which comprises of (key, value) pairs. Yet, as the number of documents increases, the proportion of relevant keys to irrelevant ones decreases, leading the model to focus more on the irrelevant keys. We identify a significant challenge, dubbed the distraction issue, where keys linked to different semantic values might overlap, making them hard to distinguish. To tackle this problem, we introduce the Focused Transformer (FoT), a technique that employs a training process inspired by contrastive learning. This novel approach enhances the structure of the (key, value) space, enabling an extension of the context length. Our method allows for fine-tuning pre-existing, large-scale models to lengthen their effective context. This is demonstrated by our fine-tuning of  $3B$  and  $7B$  OpenLLaMA checkpoints. The resulting models, which we name LONGLLAMA $^2$ , exhibit advancements in tasks requiring a long context. We further illustrate that our LONGLLAMA models adeptly manage a  $256k$  context length for passkey retrieval.
Language models have served as a catalyst for substantial advancements in several areas, including natural language processing [Radford et al., 2019, Brown et al., 2020], code generation [Chen et al., 2021, Li et al., 2022], quantitative reasoning [Lewkowycz et al., 2022] and theorem proving [Polu and Sutskever, 2020, Jiang et al., 2022, Mikula et al., 2023]. One of the central challenges with language models is the effective incorporation of extensive new knowledge. The common practice of fine-tuning the model is not only resource-intensive and complex to manage, but it also does not always clearly indicate how to incorporate new knowledge. For example, fine-tuning on a text such as "Alice in Wonderland" does not equip the model to answer questions about the story itself, but rather
*Equal contribution† Work done while at Google Research.
2We release the checkpoints and source code of LONGLLAMA, see also our colabs.
Figure 1: Accuracy of LONGLLAMA  $3B$  on passkey retrieval compared to the original OpenLLaMA model. Our method extrapolates beyond the training length, achieving  $94.5\%$  accuracy at a context length of  $100k$  and  $73\%$  at  $256k$  tokens, while the baseline is unable to handle context longer than its training length  $(2k)$ .
it trains the model to predict the next token or complete masked sentences. A promising alternative – integrating the new knowledge within the context – doesn’t require training but is considerably restricted by the model’s effective context length. For this method to work with large knowledge databases (like large code repositories), the model needs to manage a context length extending to millions of tokens.
In this research, we highlight one of the primary obstacles in augmenting the context length: as the number of documents increases, the ratio of pertinent to irrelevant tokens diminishes. The standard training procedure frequently results in overlaps between keys connected with irrelevant values and those related to relevant ones, exacerbating the model's task of differentiating between them. We term this challenge the distraction issue.
We propose the Focused Transformer (FOT), an innovative technique developed explicitly to address this issue. The Focused Transformer permits a subset of attention layers to access an additional context of (key, value) pairs through the k-nearest neighbors (kNN) algorithm, akin to the method used in [Wu et al., 2022]. This mechanism effectively extends the total context length. The distinctive aspect of the Focused Transformer is its training procedure, drawing from contrastive learning. This method addresses the distraction issue and facilitates larger context capacities. Specifically, during the training phase, we deliberately expose the chosen subset of attention layers to both relevant and irrelevant keys (like negative samples from unrelated documents). This strategy incentivizes the model to differentiate keys connected with semantically diverse values, thereby enhancing their structure.
We introduce and make available LONGLLAMAs (→), fine-tuned OpenLLaMA models with FOT, demonstrating that our method does not require long context during training and can be applied to existing models. Notably, LONGLLAMAs show significant improvements on tasks necessitating long-context modeling. In particular, they can manage a 256k context length on the passkey retrieval task [Mootashami and Jaggi, 2023].
Our research contributions are the following:
1. We pinpoint the distraction issue as a significant challenge and a primary obstacle to scaling up the context length in Transformer models, particularly in multi-document scenarios.
2. We develop the Focused Transformer (FoT), designed to alleviate the distraction issue. FoT includes a unique training objective that improves the (key, value) structure, enabling the use of extensive additional context and k-nearest neighbors lookup to scale the context length.
3. Our method is simple to implement, and it provides the benefit of extending model context without modifying the architecture, facilitated by cost-effective fine-tuning. We demonstrate this on the  $3B$  and  $7B$  OpenLLaMA checkpoints. The resulting models, named LONGLLAMAs, display enhancements on tasks that benefit from increasing the number of few-shot demonstrations in the extended context, such as TREC [Li and Roth, 2002, Hovy et al., 2001] and WebQS [Berant et al.,
[2013]. We also prove that for passkey retrieval Mohtashami and Jaggi [2023], our LONGLLAMA models successfully handle a  $256k$  context length.
4. We further scrutinize FoT's capabilities across various datasets and model sizes. We show that a FoT trained with a total context of 512 tokens can extrapolate to 16 million tokens in a benchmark dictionary lookup task. We also assess FoT on long-context language modeling tasks such as books (PG-19), mathematics (arXiv), code (GitHub), and formal proofs (Isabelle), where it exhibits improvements in perplexity over baselines.
Long-context transformer architectures A multitude of approaches have been developed to increase the context length of transformers, mostly focusing on alleviating the quadratic complexity of the attention computation. For instance, Transformer-XL [Dai et al., 2019] caches the previous context and enables the linear extension of context with the number of layers. Longformer [Beltagy et al., 2020] employs an attention mechanism that allows tokens to attend to distant tokens sparsely, reducing the computational complexity. BigBird [Zaheer et al., 2020], LongT5 [Guo et al., 2021], and [Dao et al., 2022] also use sparse attention to handle long sequences. Different efficiency considerations have been studied in [Kaddour et al., 2023], showing that they lead to limited gains. Hierarchical transformers [Nawrot et al., 2021, 2023] downsample activations in intermediate layers to reduce computation and enable longer contexts. COLT5 [Ainslie et al., 2023] proposes conditional computation to save memory and enable larger contexts. Memorizing Transformer [Wu et al., 2022] uses kNN lookup to pick up the most relevant tokens, which might also be seen as a way to reduce the computational complexity of attention. Our work adheres to this approach and aims to train a key space that handles longer attention context length (e.g., by mitigating the distraction issue) and, thus, has better long-context capabilities.
Fine-tuning LLMs for longer retrieval Prior works such as RETRO [Borgeaud et al., 2022] (RETROfitting) and Memorizing Transformer [Wu et al., 2022] have demonstrated a promising path for fine-tuning existing LMs to add new capabilities without the need to retrain the entire model. In contrast to those approaches our method is not framed as a retrieval but as a way of extending the context of the model. In contrast to RETRO, we propose a single-stage method for context extension instead of a two-stage retrieve-then-embed approach. We provide a more detailed comparison with the Memorizing Transformer in Appendix C.3. More recently, a number of works have explored fine-tuning LLaMA to extend its context length. Landmark attention [Mohtashami and Jaggi, 2023] proposes a compression scheme of LLM's context into landmarks, increasing the context length of LLaMA-7B to  $32K$ . Position Interpolation (PI, [Chen et al., 2023] and [kaiokendev, 2023]) introduces a modification to the rotary positional encoding scheme that enables fine-tuning for  $32K$  context. In contrast to this work, our method does not rely on positional encodings, following the findings from [Haviv et al., 2022]. Removing positional encoding in additional context allows us to extrapolate to  $256k$  tokens, although the model was only trained on sequences up to  $8K$ , yielding theoretically unbounded context length.
Zero-shot methods KNN-LM [Khandelwal et al., 2019] shows that one can improve the performance of a LLM by combining two probability distributions. One created by a pre-trained model, and one based on the similarity between the embedding of the currently processed token and the embeddings of tokens retrieved from a large database. Meanwhile, we extend the model context in a subset of attention layers, potentially allowing for reasoning within this extended context. Parallel Context Windows for Large Language Models [Ratner et al., 2023] introduces a method for extending the context of language models without training. They achieve this by embedding several context windows independently in parallel and allowing only a subset of tokens to attend to all windows. On the other hand, we fine-tune existing models and allow all tokens to attend to all previous tokens but only in a subset of layers. Additionally, our method allows us to improve the structure of the key-value space of the existing models.
Contrastive learning. Contrastive learning aims to learn good representations by comparing positive and negative examples. CLIP [Radford et al., 2021] and SimCLR [Chen et al., 2020] are two popular contrastive learning methods that have achieved state-of-the-art performance in the image domain. During contrastive pre-training, negative examples are kept in the same batch to learn to distinguish
Figure 2: The Focused Transformer overview. During inference, a memory attention layer (green) uses additional context of (key, value) pairs via kNN lookup, which effectively extends its context length. This layer is trained using crossbatch. Namely, the tokens from the current context  $C_{curr}$  attend in a differentiable way (Att + ∇) to the previous context  $C_{prev}$  of the same document and, importantly,  $d - 1$  contexts of other documents. The latter serve as 'negative' examples intended to better shape the (key, value) space.
them from positive examples. Scaling the batch size in contrastive learning has been demonstrated to enhance the quality of representations, as shown in [Gao et al., 2021b]. It has been suggested [Gao et al., 2019] that the embedding space in language modeling suffers from degeneracy, where embeddings are tightly packed in a narrow cone, making it difficult to distinguish between them. TRIME [Zhong et al., 2022] proposes a training approach designed for training LMs with memory augmentation, which uses negatives to improve the quality of representations. The main difference between this and our approach is that we incorporate negatives into the chosen subset of attention layers instead of interpolating in the output layer and use the standard language modeling loss. TRIME [Zhong et al., 2022] also focuses on retrieval from large databases, whereas we focus on extending the context of the model. ContraCLM [Jain et al., 2023] applies contrastive losses at both the token and sequence levels during training to promote more uniformly distributed, isotropic representations. It is shown to enhance the discrimination of representations on textual semantic similarity benchmarks. While ContraCLM focuses on improving the general expressiveness of representations, our work introduces contrastive-inspired techniques designed specifically for training the attention mechanism to handle longer context lengths. Nonetheless, exploring other contrastive learning objectives could be beneficial for further improving the key structure in future work.
3 FoT: Focused Transformer
Our method, the Focused Transformer (FOT), is a simple plug-and-play extension of transformer models and can be used both to train new models or fine-tune existing, possibly large, models with longer context. To this end, FoT uses memory attention layers and the crossbatch training procedure. Memory attention layers enable the model to retrieve information from the additional context at inference time, effectively extending the context. The crossbatch training procedure biases the model to learn (key, value) representations, which are easy to use by a memory attention layer. See Figure 2 for an overview of the FOT architecture and Appendix L for pseudocode.
3.1 Memory attention layers
Memory attention layers  $\mathcal{L}$  are endowed with access to an additional context during inference. Namely, each query in  $\ell \in \mathcal{L}$  attends to preceding keys from the local context and the top  $k$  most matching keys (i.e. having the largest inner product with the query) from memory. The memory keys are ranked by the inner product with the query and retrieved using the kNN search algorithm. We use the exact kNN search implemented in FAISS [Johnson et al., 2017]. The memory is populated incrementally with  $(key, value)$  pairs processed by  $\ell$  beforehand. Our memory attention layer design is closely related to [Wu et al., 2022], we follow most of its design choices, except for the gating, which we replace with a simpler mechanism, which turns out to be more effective in our applications. See details in Section C.3 and Appendix B.2. We remove positional encodings in memory layers
in all our models except LONGLLAMAs. This allows LONGLLAMA checkpoints to be a drop-in replacement for LLaMA checkpoints. We treat the kNN search algorithm as an approximation of full dense attention, which opens the doors for future speed-ups.
3.2 Crossbatch training procedure
Our training procedure is a novel way of training (or fine-tuning) transformer-based architectures in order to improve the structure of the (key, value) space. The main motivation is to shape this space so that a memory attention layer  $\ell \in \mathcal{L}$  can easily focus on relevant information. The key idea, inspired by contrastive learning, is to expose  $\ell$  to (key, value) pairs from the current and previous local context of the given document (positives) and  $d - 1$  contexts from unrelated documents (negatives). Importantly, this is done in a differentiable way.
To achieve this, we use a data pipeline in which each element of the batch corresponds to a different document. We embed the previous  $(C_{\mathrm{prev}})$  and the current  $(C_{\mathrm{curr}})$  local context for each of the processed documents. The overview of our procedure can be found in Figure 2. Specifically for each document  $\delta$  in  $C_{\mathrm{curr}}$  we create a set  $\{p_i^\delta\}_{i = \{1,\dots,d\}}$  consisting of the (key, value) pairs from the previous local context of  $\delta$  (positives), along with pairs from  $d - 1$  other contexts coming from  $C_{\mathrm{prev}}$  (negatives). We also experiment with varying the number of previous contexts and negatives for different batch elements.
The operation is fully differentiable, and thus, we improve all the (key, value) pairs in  $p^\delta$ . Two, the procedure is easy to implement; it does not require any additional loss (i.e., uses the standard transformer training objective) and is done on the level of the data loading pipeline and a minor self-attention change. The only new hyperparameter is  $d$ , which prescribes the ratio of positive to negative samples. Typically, we find it beneficial to start with small  $d \leq 8$  (otherwise, the model tends to ignore the previous local context) and later switch to bigger values, say  $d \geq 64$ . Appendix B.3 provides more details about the method. Listing 1 outlines an implementation of the crossbatch.
3.3 The distraction issue
In this section, we conceptualize what we call the distraction issue and hypothesize it is one of the key problems in dealing with long multi-document contexts (like large code repositories). Namely, during the standard training, the model is not incentivized to distinguish the keys from different documents. We measure that the attention mass is evenly spread on the related and unrelated documents; see Figure 3. More precisely, for a document  $\delta$ , let  $w_{ij}$  be the softmax weights related to  $p_{ij}^{\delta}$  constructed as described in Section 3.2. We define the positive attention mass as  $r_d := \sum_j w_{1j} / \sum_{i=1}^d \sum_j w_{ij}$ . We observe that  $r_d \approx 1/d$ , which can be interpreted as the fact that the attention is equally distracted by the positive (coming from the current document at  $i = 1$ ) and negative keys. This is an undesirable prop
erty since when scaling the memory, the attention becomes increasingly distracted. We show that the crossbatch mostly alleviates the distraction issue, resulting in a focused attention. More information can be found in Appendix B.4. In Section 5.3, we also show that the distraction issue has a harmful effect on metrics like perplexity.
Figure 3: Distraction issue. We compare FoT trained with different values of parameter  $d$  to the standard Transformer baseline. During the evaluation, both models see the previous local context and some contexts from other documents in the chosen layer (as in crossbatch training procedure). For a document  $\delta$  we measure the distribution of attention mass on  $p^{\delta}$ . Scale  $x$ : the number of contexts from documents that the model can see. Scale  $y$ : avg attention mass to the previous local context of the current document.
4 LONGLLAMA: extending LLaMA's context length with FoT
One of the promises of our work is that FoT can be used to fine-tune already existing large models to extend their context length. In this section, we show that this is indeed the case. We use OpenLLaMA-3B and OpenLLaMA-7B models trained for  $1T$  tokens as starting points and fine-tune them with FoT. We show that the resulting models, which we call LONGLLAMAs, are capable of extrapolating beyond their training context length (even up to  $256K$ ) and retain the performance on short-context tasks. We release the inference code on GitHub: https://github.com/CStanKonrad/long_1lama and the LONGLLAMA-3B checkpoint on Hugging Face: https://huggingface.co/syzymon/long_1lama_3b. We note that our checkpoint is backward compatible, i.e. can be used with any existing LLaMA inference code (both in Hugging Face and other implementations), albeit without long-context capabilities.
The architecture of the models is the same as OpenLLaMAs, see Geng and Liu [2023] and Appendix A.1. We use  $\mathcal{L} = \{6, 12, 18\}$  (resp.  $\mathcal{L} = \{8, 16, 24\}$ ) as the memory layers for  $3B$  (resp.  $7B$ ) LongLLAMA model. We fine-tune the models on  $10B$  (resp.  $3B$ ) tokens using FOT,  $8k$  context length and our dataset mixture based on RedPajama [TogetherComputer, 2023], see Appendix A.3.
There are three minor differences from the standard FoT procedure. First, we retain the positional encodings in the local context of the memory layers (this is not necessary for FoT, but makes our checkpoints fully compatible with any existing LLaMA inference codebase). To be more precise, queries and keys from the local context (up to  $2K$  tokens) receive the standard LLaMA rotary positional encoding, whereas memory keys are encoded as if they had position 0 in the local context window. Second, we use dense attention instead of the kNN retrieval, as we found only marginal performance differences, and it is simpler to implement. Third, we modify the crossbatch training procedure to have more fine-grained control over the number of additional contexts and the ratio of positive to negative samples. All these differences are detailed in Appendix A.2.
4.2 Context length extrapolation on the passkey retrieval task
We first measure the effective context length of LONGLLAMA, namely the distance for which tokens can effectively attend each other. We use passkey retrieval introduced in [Mohtashami and Jaggi, 2023], a synthetic task designed to measure this property. In this task, the model has to retrieve a passkey placed randomly in a long prompt. Results are shown in Figure 1 - importantly, our  $3B$  model is capable of solving this task much beyond its training context length  $8K$ , achieving  $94.5\%$  accuracy for prompts of length  $100k$  and  $73\%$  for  $256k$ .
4.3 Question answering over research papers
In Table 6 we present the performance on the validation set of Qasper [Dasigi et al., 2021] from SCROLLS [Shaham et al., 2022] and compare our results to LongChat 7B [Ma and Zhang, 2023] and two baseline short-context models. We note that our model shows gains from increased context length.
4.4 Improving few-shot learning accuracy with longer context
We measure long-context capabilities of these models on two downstream tasks, TREC question classification [Li and Roth, 2002, Hovy et al., 2001] and WebQS question answering [Berant et al., 2013]. We follow the experimental setup of [Hao et al., 2022]. Namely, we few-shot prompt the models with as many demonstration examples as possible up to the given context length. We do not use structured prompting like in [Hao et al., 2022] - instead, we directly provide all demonstrations in context.
We observe significant accuracy gains from longer contexts on TREC and some improvements on WebQS (see Table 1). The TREC dataset consists of 50 classes. A model is tasked to predict the class label given in-context examples. Only 100 examples fit the standard context length  $(2K)$ ; it is not unusual that no class example is present for a given question, making the task impossible. Increasing the context length and the number of examples mitigates this risk. Moreover, having more demonstrations of the given class is also likely to be beneficial.
Table 1: Few-shot in-context learning performance of LONGLLAMA; accuracy on TREC and WebQS. We see significant gains from the additional context on the TREC dataset. To calculate the results, we average over 20 trials for sampling in-context demonstrations from the train set; the resulting confidence intervals for TREC and WebQS are smaller than  $1\%$  and  $0.1\%$ , respectively.
Dataset	TREC	WebQS
Context	LONGLLAMA 3B	LONGLLAMA 7B	LONGLLAMA 3B	LONGLLAMA 7B
2K	67.0	63.2	21.2	25.5
4K	71.6	72.7	21.4	26.4
6K	72.9	74.9	22.2	27.2
8K	73.3	75.9	22.4	27.7
Table 2: Few-shot in-context learning performance comparison between standard fine-tuning on  $4K$  context (baseline) and FoT fine-tuning on the same context length for  $1B$  tokens. On TREC, FoT is able to utilize additional examples beyond its training context length to achieve higher accuracy at  $8K$  context length, which is not possible for the baseline since its context is bounded to  $4K$ .
Dataset	TREC	WebQS
Context	baseline	FoT (ours)	baseline	FoT (ours)
2K	52.8	55.6	20.7	20.8
4K	57.2	60.9	18.7	21.0
6K	-	61.7	-	21.2
8K	-	62.5	-	20.7
4.5 Comparison to standard long-context fine-tuning