Focused Transformer: Contrastive Training for Context Scaling
Authors: Szymon Tworkowski
Konrad Staniszewski
Mikołaj Pacek
Yuhuai Wu
Henryk Michalewski
Piotr Miłoś
Venue: NeurIPS 2023 poster (2023)
Keywords: Transformers
Language Models
Natural Language Processing
TL;DR: We show how to alleviate distraction issues in scaling the context length of language models to multiple documents.

Abstract: 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 $3 B$ and $7 B$ OpenLLaMA checkpoints. The resulting models, which we name LongLLaMA, exhibit advancements in tasks requiring a long context. We further illustrate that our LongLLaMA models adeptly manage a $256 k$ context length for passkey retrieval.

Focused Transformer: Contrastive Training for Context Scaling
Szymon Tworkowskil.3. 1.3+
Konrad Staniszewski 1.3+
Mikolaj Pacek 1.++
Yuhuai Wu
Henryk Michalewski%.4
Piotr Mitog
IDEAS NCBR
Institute of Mathematics, Polish Academy of Sciences
University of Warsaw
GGoole DeepMind
deepsense.aa xAI
Abstract
Large language models have an exceptional capability to incorporate new infor- mation in a contextual manner. However, the full potential such an approach is often restrained due to 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 the irrelevant keys. We identify significant challenge, dubbed the distraction issue, where keys linked to different semantic values might over- lap. making them hard distinguish. To tackle this problem, introduce the Focused Transformer (FOT), technique that employs 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 lengthen their effective context. This is demonstrated by our fine-tuning of 3B and 7B OpenLLaMA checkpoints. The resulting models, which name LONGLLAMA exhibit advancements in tasks requiring a long context. We further illustrate that our LONGLLAMA models adeptly manage 256k context length for passkey retrieval.
Introduction
Language models have served 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 al. 2022). quantitative reasoning Lewkowycz et al. 2022 and theorem proving [Polu
and Sutskever 2020 Jiang et al., 2022 Mikula et al. 2023]. One the central challenges with knowledge. The common practice of fine-tuning the model is not only sesource-intensive and complex to manage, but it also does not always clearly indicate how incorporate new knowledge. For example, fine-tuning
language models is the effective incorporation of extensive
text such as "Alice in Wonderland" does not equip the model to answer questions about the story itself, but rather
'Equal contribution "ork done while Google Research.
We release the checkpoints and source of LONGLLAMA
also our colabs.
37th Conference on Neural Information Processing Systems (NcurIPS 2023).

---

Long
100
OpenL LamLLaMM training context
Long Lagg training context
OpenLLaMA 3B
80
Accuracy (%6)
60
40
20
2k
8k
16k
100k
256k
Prompt length (#tokens)
Figure I: Accuracy of LONGLLAMA 3B on passkey retrieval compared the original OpenL model. Our method extrapolates heyond training length, achieving 94.5% accuracy a context length 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 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 augmenting the context length: as the number of documents increases, the ratio of pertinent to irrelevant tokens diminishes. The standard training procedure frequently results 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 attention layers access an additional context of (key, value) pairs through the k-nearest neighbors (kNN) algorithm, akin to the method used in [Wu al. [222]. 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 incentives the model to differentiate keys connected with semantically diverse values, thereby enhancing their structure.
We introduce and make available LONGLLAMAs (T)), 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 tasks necessitating long-context modeling. In particular, they can manage a 256k context length on the passkey retrieval task (Mohtashami 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 multi-document scenarios.
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 3/ and 7B OpenI.LMA checkpoints. The resulting models, named LONGLLAMA, display enhancements on tasks that benefit from increasing the number of few-shot demonstrations in the extended context, such 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 2564 context length.
4. We further scrutinize FoT's capabilities across various datasets and model sizes. We show that a FoT trained with total context of 512 tokens extrapolate to 16 million tokens in 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 proof's (Isabelle), where exhibits improvements in perplexity over baselines.
2 Related work
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 al. 2019] caches the previous
context and enables the linear extension of context with the number of layers. Longformer [Beltagy
al. 2020] employs an attention mechanism that allows tokens to attend to distant tokens sparsely. reducing the computational complexity. BigBird (Zaheer et al.) 2020). LongTS [Guo et al. [02 ]
and (Dao et al.) [022] also use sparse attention to handle long sequences. Different efficincyy considerations have been studied in [Kaddour et al. [202)), showing that they lead to limited gains. Hicrarchical transformers [Nawrot et al. [02 2023] downsample activations in intermediate layers to reduce computation and enable longer contexts. COLTS [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 way to reduce the computational complexity of attention. Our work adheres to this approach and aims to train key space that handles longer attention context length (e.g., by mitigating the distraction issue) and, thus, has better long-conted capabilities.
Fine-tuning LLMs for longer retrieval Prior works such as RETRO [Borgeaud et al. 2022
for (RETROfitting) and Memorizing Transformer [Wu et al. 2022] have demonstrated a promising path 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 stage method for context extension instead of a two-stage retrieve-then-emmeed approach. We provide a more detailed comparison with the Memorizing Transformer in Appendix C.33 More recently, a number of works have explored fine-tuning I. LLM to extend its context length. Landmark attention |Mohtashami and Jaggi 2023
proposes a compression scheme of L.M'' 's context into landmarks, increasing the context lenghh of LLaMA-7B to 32K. Position Interpolation (PI, [Chen et al.) 2023] and [kaiokendev] 2023
introduces 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 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 perfor- mance 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
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 RRaffrd et al. 2021 and SimCLR [Chen et al. 2020 are two popular contrastive learning methods that have achieved state-of-the-arr performance in the image domain. During contrastive pre-training, negative examples are kept in the same batch to learn to distinguish
3

---

Train
CrossBatch
Decument
Documint
d=2
Dreumeed It
Dozumen
Oprey
CourT
Oorev
Court
Inference
The Focused Transformer overview. During infer- ence, memory attention layer (green) uses additional context (key, value) pairs via kNN lookup, which effectively ex- tends its context length. This layer trained using crosshatch. Namely, the tokens from the current context Ceurr attend in differentiable way (Aut V) to the previous context Cyrea of the same document and, importantly d contexts of other documents. The latter serve as 'negative" examples intended to better shape the (key, value) space.
Figure 2:
Naighbors
Retrieve
Alt
:
Store Keys
Memory
and Volues
Crur
them from positive examples. Scaling the batch size in contrastive learning has been demonstrated to enhance the quality of representations, as shown [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 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 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 [Zhog 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-inspires 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.
FOT: Focused Transformer
Our method, the Focused Transformer (FoT), is simple plug-and-play extension of transformer models and can be used both to train new models 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 inference time, effectively extending the context. The crossbatch training procedure biases the model learn (key, value) representations, which are easy to use by memory attention layer. See Figure for an overview of the FoT' architecture and Appendix [ for pseudocode.
3.1 Memory attention layers
Memory attention layers C are endowed with access to an additional context during inference. Namely, each query in é L attends to preceding keys from the local context and the top & 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 Pohnson et al. [017]. The memory is populated incrementally with (key, value) pairs processed by & 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 simpler mechanism, which turns out to more effective in our applications.
See details in Section C33 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 -nop-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 novel way of training (or fine-tuning) transformer-basee architectures in order to improve the structure of the (key.aluu)) value) space. The main motivation to shape this space so that a memory attention layer & L can easily focus on relevant information. The key idea, inspired by contrastive learning, is to expose i to (key, value) pairs from the current and previous local context of the given document (positives) and d - contexts from unrelated documents (negatives). Importantly, this done differentiable way.
To achieve this, we use data pipeline in which each element of the batch corresponds to a different document. We embed the previous prev) and the current (Ceurr) local context for each of the processed documents. The overview of our procedure can be found in Figure 2 Specifically for each document 8 in Ceurr we create a set (i consisting of the (key, value) pairs from the previous local context of 5 (positives). along with pairs from d 1 other contexts coming from Cprev (negatives). We also experiment with varying the number of previous contexts and negatives for different batch elements.
The operation is fully differentiable, and thus, improve all the (key, value) pairs in p². Two, the procedure easy to implement: 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, find it beneficial start with small d < 8 (otherwise, the model tends to ignore the previous local context) and later switch to bigger values, say 2 64. Appendix B.3 provides more details about the method. Listing outlines an implementation of the crossbatch.
3.3 The distraction issue
In this section, 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 in- centivized distinguish the keys from dif-
052-064
Positive Attention Mass
FoT deB
0.8
FoT Transformer
-//d
0.6
0.4
ferent documents. We measure that the attention mass evenly spread on the re- lated and unrelated documents; see Fig- ure 3 More precisely, for a document d, let waj be the softmax weights related to
0.2
0
Number of Documents
py, constructed as described in Section 3.2 We deñne the positive attention mass as
Figure 3: Distraction issue. We compare FoT trained with different values 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 erossbatch training procedure). For document 8 we measure p². Scale z: 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.
Ta : -,, w/5/ Ei= Σ3 Wij- We ob- serve that тa & 1/d, which can be in- terpreted 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 we also show that the distraction issue has a harmful
distribution of attention mass on
can be found in Appendix B.4] Section 5.3 effect on metrics like perplexity.
5

---

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 IT tokens as start- ing points and fine-tune them with FoT. We show that the resulting models, which call LONGLLAMAS, are capable of extrapolating beyond their training context length (even up to 256K) and retain the performance short-context tasks. We release the inference code GitHub: bettps//giithn
cowm/SStankonrra//oorgg.. 1lama and the LONGLLAMA-3 check- point Hugging Face: htttp::/huggingfaae co/ syzymo//y//////o long 1lama_maaaaa 3b We note that our checkpoint is backward compatible, i.c. can be used with any existing LLaMA inference code (both in Hugging Face and other implementations), albeit without long-content capabilities.
4.1 Experimental setup
The architecture of the models the same as SapenLLMMAs, see Geng and Liu [023 and Ap-
pendix A.1 We use L = {6, 12, 18} (resp. 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, 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 encoded as if they had position in the local context window. Second, we use dense attention instead of the kNN retrieval, as we found only marginal performance differences, and 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 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), synthetic task designed to measure this property. In this task, the model has to retrieve passkey placed randomly in a long prompt. Results are shown in Figure [ importantly, our 3B model 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 [M 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 clas- sification [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 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): 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 Few-shot in-context learning performance 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 in-context demonstrations from the train set; the resulting confidence intervals for TREC and WebQS 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 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 achieve higher accuracy 8K context length, which not possible for the baseline since its context is bounded 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
In this section, compare FoT to standard long-context fine-tuning, showing that it already achieves better performance for the context length used for fine-tuning and, importantly, that can extrapolate beyond this context length, which is not the case for the baseline.
For comparisons, we fine-tune two models, one trained with FoT and another one (baseline) with standard fine-tuning (done similarly to [MosaicML 2023 Nijkamp et al. 2023). In both cases, we use 3B models fine-tuned on BB tokens using the 4 context length. evaluate both models on number of few-shot downstream tasks in the setting described in Section 4.4
In most cases, Table observe accuracy improvements when more few-shot demonstrations are provided in the extended context (from 2K used by OpenLLaMA to 4K used in our fine-tuning). On TREC, the gains from additional context are significant for both models, while on WebQS, the standard fine-tuning baseline does not provide any improvement from extended context. Notably, the model fine-tuned with FoT enjoys further accuracy gains when evaluated with context lengths beyond its training length (6K and 8K). This shows extrapolation capabilities of FoT, which are not present in the baseline (see e.g. Figure
4.6 Performance on short-context tasks
Fine-tuning for longer contexts could hurt performance on the original context length (2K), as the training data distribution changes. We show that this the for the LANGLLAMA models by evaluating them using the LM Evaluation Harness library [Gao et al. 2021a]. On most tasks, the performance is kept intact; Appendix A.4 for details. This also confirms that LONGLLAMAS could be used as drop-in replacement of LLaMA models as they are compatible with the original LLaMA inference code.
Analysis of FoT
In this section, we perform extensive experiments smaller models to analyze and further validate our approach. In particular, we answer the following questions: (1) How does FoT perform when scaling the context length at inference time? (2) Can FoT be used to extend the context length of
7

---

an existing, pre-trained model? (3) How effectively can it handle distractions, and how does this capability translate enhanced performance in long-context language modeling tasks? Moreover, we provide ablation studies of our method and additional analysis.
5.1
Experimental setup
For experiments described in this section we use decoder-only Transformer [Vaswani et al. 2017] models with 12 layers and 184M parameters (unless stated otherwise). Following Wu et al. 2022|:
Architecture
: we pick & = 8 as the memory attention layer. We tune I = 128, the number of top keys start training with small crossbatch dimension S 8 and switch to d 2 64 after some training. For more details about the architecture and hyperparameters,
retrieved by KNN. In most experiments,
see Appendix B and Appendix
Evaluation We distinguish two evaluation settings: single-document (abbreviated to single-doe) and multi-document (abbreviated to multi-doc). The single-doc setting is typically used for evaluating models that process long contexts. Here, we clear the memory for each new document, ensuring that only the current document is available in the context. The multi-doc setting retains memory across multiple documents without resets, This scenario tests whether the model can ignore irrelevant information and focus on the relevant data, which can be useful in setups like repository-level code generation.
Datasets We evaluate on the following long-context language modeling datasets: PG-19 (English books), Xrviv (mathematical papers), GitHub (code), and Isabelle (formal proofs). PG-19 |Rae et al. 2019 is large dataset of English-language books published prior to 1919, sourced from the Project Gutenberg archive. This dataset is well-established benehmark for evaluating Long-context language models [Sun et al. 2021]. The arXiv dataset contains ETEX source of papers labeled as "Mathematics" that were obtained by downloading articles through the arXiv Bulk Data Access. The token count per paper in this dataset is comparable to that of book in PG19. For details on the remaining datasets, refer to Appendix
5.2 FoT fine-tuning and context length extrapolation
FoT is minimal modification to the standard transformer architecture; therefore, it is possible to fine-tune existing models to endow them with a longer context length via the memory attention layer, we already demonstrated in Section 4 In this section, we deepen this analysis (on smaller model) by studying perplexity improvements on various datasets.
As base model, we use a standard transformer model pre-trained for 100k steps with context of 1K tokens using the standard objective and fine-tune with the FoT objective (i.e. crossbatch). The data used for both fine-tur tuning and pre- pre-training is the C4 dataset Raffel et al. [019a] (we omit documents shorter than 2K tokens). The fine-tuning phase takes 10k steps. We use the crossbatch dimension d 128 and local context of 1K tokens (context is 2K during training). We evaluate models in zero-shot way on 4 language modeling datasets, which require long context: arXiv, PG-19, GitHub and Isabelle, see Section 5.1 and Appendix E for details.
In Table we observe that FoT enjoys steady perplexity gains up to 64K tokens, although it was fine-tuned only with the 2K total differentiable context length. We compare the model perplexity to the following baselines: Memorizing Transformer (MT) [Wu al. 2022 fine-tuned with the local context of 1 K and memory size of 16K, and Transformer-X XI. [Dai et al. 2019) fine-tuned with both local context and window length of 1K. To ensure fair comparison, all three models are fine-tuned from the same base checkpoint. When evaluated with a context of 2K, our method achieves results on par with the Transformer-X XL baseline, which has access to the previous context in all layers, unlike MT and FOT. Compared to the MT baseline, we achieve better scaling when evaluated with 64K context length and significantly better perplexity values. Unlike MT, our method does not require training long sequences, which is reflected by the lower perplexities of FoT when evaluated in the zero-shot setting. For more details, see Appendix G
We also confirm the context extrapolation abilities using synthetic dictionary lookup task. In this task, the model first provided with k; 6 mappings and then asked what value is associated with a particular key. We train 37M parameter models using documents of length 512. Figure 10 shows that
8

---

Table 3: Perplexity for different context lengths after fine-tuning standard transformer model. The model fine-tuned using the FoT objective (i.e., crossbatch) on C4 and evaluated zero-shot varying the context size. Transformer-XI XL Dai et al. 2019 and Memorizing Transformer Wu ct al. 2022] fine-tuned in the same setting are used as baselines.
Method
Context Length
GitHub
Isabelle
arXiv
PG-19
2K
6.72
5.63
8.17
23.74
4K
5.88
4.93
7.44
23.25
FoT
16K
5.43
4.51
6.94
22.85
64K
5.32
4,44
6.81
22.65
Transformer-XI.
2K
6.85
5.76
8.21
23.57
2K
8.10
7.34
9.39
24.03
4K
7.55
6.93
8.95
23.62
Memorizing Transformer
16K
7.27
6.66
8.66
23.32
64K
7.26
6.64
8.60
23.24
FoT, after 5k steps of training, can effectively utilize memory consisting of 16M tokens achieving accuracy above 92%. Details can be found in Appendix F
5.3 Handling distractions in language modeling tasks
In this section, we measure how handling distractions the multi-document setting helps in language modeling. We pick the PG-19 dataset [Rae et al., [219]] and measure the perplexity of the next token prediction (language modeling task) when varying the size of multi-doc memory (in this case consisting of books). Intuitively, the memory tokens corresponding the current book might be beneficial (which is also confirmed in [Wu et al., 2022)), while the ones from the other books are unlikely to be useful and thus are distractions.
We observe, see Figure 8 that higher values the crossbatch dimension d lead to better perplexity. This aligns with the observations in Section 3.3 indicating that by mitigating the distraction issue, we experience benefits language modeling.
Moreover, all versions of FoT are able to utilize memory and achieve much better perplexity than the standard Transformer (no memory). Unsurprisingly, perplexity increases with memory size, but we stress that this happens gracefully. In the standard variant of FoT (bold line). the perplexity increases only by 0818 when scaling to > 500k tokens. Importantly, the perplexity of FoT close to this of Memorizing Transformer with the single-doc memory, which we treat as soft lower bound since it not exposed to distractions from unrelated books.
5.4 Context length extrapolation in single-doc
The original motivation behind FoT is improve the multi-doc setting performance by handling distractions. Interestingly, our method also helps to extrapolate to longer contexts, even when evaluated in the single-doc setting.
To study this, we perform FoT fine-tuning (as in Section 5.2 and evaluate the perplexity of the resulting model on the PG-19 dataset with different context lengths in the zero-shot fashion. To deepen the analysis, we introduce an additional parameter w (the number of previous contexts used in cross batch training procedure). We provide results for I (the standard setting for FoT, that corresponds to the total differentiable context being 2 1024) and w = 2 (corresponding to the total differentiable context 3- 1024).
observe, see Figure improvements when context grows, even far beyond the training context length, which reaffirms the hypothesis that FoT helps with extrapolation to longer contexts. Moreover, is significantly better than 1. When comparing d and w to d and w 1, we observe that the former is slightly better. This is natural, as the former has longer training context.

---

5.5 Ablations and design choices
In Appendix C we present ablations our design choices. In particular, we note the importance of differentiability and the inclusion of negatives. We also discuss the relation to Memorizing Transformer. We note that due to the limited resources we have followed the Memorizing Transformer the choice of memory layers.
6 Limitations and future work
Our research opens few avenues for future work. We list them as well as challenges and limitations.
Scaling up context This is by far the most important future research direction. The challenges start from purely engineering, storing more than 16M (key, value) pairs will require a distributed multi-node system. In experiments, we the exact kNN search, which is not scalable to large memory. Using approximate kNN search will require lot of engineering effort, as well as careful evaluation of the impact of the approximation the model performance.
Scaling up crossbatch We observed that increasing d is beneficial. In our experiments, we used d = 64 or d = 128, which is the maximum value that fits into the memory of single TPUv3/TPUv2 machine, see also Appendix[[ [ In future work, we want to further increase as well as test on devices with bigger memory or utilize multi-node training. We also note that crossbatch increases the training cost, but only in subset layers.
Exploring contrastive learning The FoT training is inspired by rather basic contrastive learning (CL) techniques. We show that this improves the key structure so that the distraction issue mitigated. We expect that other CL methods could be beneficial, for example, hard negative mining to utilize larger memory during training (see [Lindgren et al. 2021]). We leave this for future work.
Combining with other methods Developing long-come long-context methods is an active research field, see Section 2 We believe that some of these methods could be combined with FoT, resulting in mutually beneficial interactions.
Listing 1: Possible implementation of cross-batch. To simplify the code we assume that each document occupies two consecutive elements of the batch. A more detailed version is in Appendix
keys from other contexts will be encoded as if they were at the beginning of the local context fsey_fst pos_encode_ass_frsstt (xk=key)
local contert keys encoded in the standard way pquery, pkey = pos_encode (xq=query, xk=key)
for each element of the batch calculate indices of the batch that will used in cross- cross_batch_rel_. = arange -nun_attentions, -1)
rreshape -1)
batch_ids jup. arange batch_size), reshhppe-.. 1)
bross_batch_selectto
cross_batch_ren
batch_ids
here we want other contexts cross_ batch keys pkey fat [cross beatchheeeetorrrr...oooaaaeeeoee [:, 1:]1
here we concatenate local context with other contexts
attention_ (apematee(ppeyeee [[pkey[:, None], bross_ axis=1) jnp concatenate
cb ab_attn_weights
jup. einsum ("bqhd, brahdd>bhqche pquery, Ectention_keys preciaionssrecision))
10

---

Acknowledgments and Disclosure of Funding
We gratefully acknowledge the TPU Research Cloud program, which was instrumental our research by providing significant computational resources. Parts of the project were realized using the resources of Poznańskie Centrum Superkomputerowo- Sieciowe. We would also like to thank Markus Rabe for reviewing the initial manuscript and Christian Szegedy, Charles Staats, and DeLesley Hutchins for helpful discussions. We are also grateful to Xinyang Geng and Hao Liu for releasing OpenLLaMA checkpoints and the EasyLM library [Geng] 2023]. allowing for training these models, which significantly accelerated our research. Piotr Milos was supported by the Polish National Science Centre grant 201 2019/55/0//76/00444.. Henryk Michalewski was supported by the Polish
National Science Center grant UMM0-201899////7600003
References
Joshua Ainslie, Tao Lei, Michiel de Jong. Santiago Ontañón, Siddhartha Brahma, Yury Zemlyanskiy, David C. Uthus, Mandy Guo. James Lee-Ther Thorp, Yi Tay, Yun-Hsuan Sung, and Sumit Sanghai. Colt5: Faster long-range transformers with conditional computation. CoRR, abs/2303.09700 097522. 2023, 2303 09752
doi: 10.48550/arXi.22333.0999
URL http:htttt//att http://oti.o/g///////// wttp://////1...///rrrr..././ 48550/arXiv
Iz Beltagy. Matthew E Peters, and Arman Cohan. Longformer: The long-documen transformer. arXiv preprint arXiv:2004.05115 2020.
Jonathan Berant, Andrew Chou, Roy Frostig, and Percy Liang. Semantic parsing on Freebase from question-ansee answer pairs. In Proceedings of the 2013 Conference on Empirical Methods in Natural Language Processing, pages 1533 1544, Seattle, Washington, USA, October 2013. Association for Computational Linguistics. URL https: //tp:://wwwww wwww://wwwwww.wwb.////./om//.... aclweb. org/ htt///ntthooyyy// DD3- 1160
Sebastian Borgeaud, Arthur Mensch, Jordan Hoffmann, Trevor Cai, Eliza Rutherford, Katie Mil- lican, George van den Driessche, Jean- -aptittt Lespiau, Bogdan Damoc, Aidan Clark, Diego de Las Casas, Aurelia Guy. Jacob Menick, Roman Ring, Tom Hennigan, Saffron Huang, Loren Maggiore, Chris Jones, Albin Cassirer, Andy Brock, Michela Paganini, Geoffrey Irving, Oriol Vinyals, Simon Osindero, Karen Simonyan, Jack W. Rae, Erich Elsen, and Laurent Sifre. Im- proving language models by retrieving from trillions of tokens. In Kamalika Chaudhuri, Stefanie Jegelka, Le Song, Csaba Szepesvári, Gang Niu, and Sivan Sabato, editors, International Confer- ence on Machine Learning. ICML 2022, 17-23 July 2022, Baltimore, Maryland, USA, volume 162 of Proceedings of Machine Learning Research, pages 2206-2240. PMLR, 2022. URL
httpps:/pproceeding mlr. wwwpppresssviee bttp/////brreaa22a..oo html
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, Sam McCandlish, Alec Radford, Ilya Sutskever, and Dario Amodei. Language models are few-shot learners. CoRR, abs/2005. 14165, 2020. URL htL/sttps:// //tttss///axxx. http:://axxxvv.oooaaa/ao bbrg//bbs/in absa /205 14165
Mark Chen, Jerry Tworek, Heewoo Jun, Qiming Yuan, Henrique Ponde de Oliveira Pinto, Jared Kaplan, Harrison Edwards, Yuri Burda, Nicholas Joseph, Greg Brockman, Alex Ray, Raul Puri, Gretchen Knueger, Michael Petrov, Heidy Khlaaf, Girish Sastry, Pamela Mishkin, Brooke Chan, Scott Gray, Nick Ryder, Mikhail Pavlov, Alethea Power, Lukasz Kaiser, Mohammad Bavarian, Clemens Winter, Philippe Tillet, Felipe Petroski Such, Dave Cummings, Matthias Plappert, Fotios Chantzis, Elizabeth Barnes, Ariel Herbert- Voss. William Hebgen Guss, Alex Nichol, Alex Paino, Nikolas Tezak, Jie Tang, Igor Babuschkin, Suchir Balaji, Shantanu Jain, William Saunders, Christopher Hesse, Andrew N. Carr, Jan Leike, Joshua Achiam, Vedant Misra, Evan Morikawa, Alec Radford, Matthew Knight, Miles Brundage, Mira Murati, Katie Mayer, Peter Welinder, Bob McGrew, Dario Amodei, Sam McCandlish, Ilya Sutskever, and Wojeiech Zaremba. Evaluating large language models trained on code. CoRR, abs/2107.03374 2021. URL https://axxxi oxrg/abs//107003374
Shouyuan Chen, Sherman Wong, Liangjian Chen, and Yuandong Tian. Extending context window of large language models via positional interpolation, 2023.
11

---

Ting Chen, Simon Kornblith, Mohammad Norouzi, and Geoffrey Hinton. A simple framework for contrastive learning of visual representations. International conference on machine learning. pages 1597-1607. PMLR, 2020.
Zihang Dai, Zhilin Yang, Yiming Yang, Jaime Carbonell, Quoc Le, and Ruslan Salakhutdi- nov. Transformer-xl: Attentive language models beyond fixed-length context. arXiv preprint arXiv: 1901002000000 2019.
Tri Dao, Daniel Y. Fu, Stefano Ermon, Atri Rudra, and Christopher Ré.
Flashat-
tention: Fast and memory-efficient exact attention with io-awareness.
In
NeurlPS, 2022. URL http://papers: mips. cc/ wwwpappr_rpppeeeeeeee paper_ffles files/ httpr//sssss/ap//2/2222//// paper/2022/hash/ 77657339200000733322111166666360 http://eeeeeooobbbboaoooooooo-ee Conf .html
Pradeep Dasigi, Kyle Lo, Iz Beltagy, Arman Cohan, Noah A. Smith, and Matt Gardner. A dataset of information-seeking questions and answers anchored in research papers, 2021.
Stefan Elfwing, Eiji Uchibe, and Kenji Doya. Sigmoid-weighted linear units for neural network function approximation in reinforcement learning, 2017.
Jun Gao, Di He, Xu Tan, Tao Qin, Liwei Wang, and Tie- Liu. Representation degeneration problem training natural language generation models. arXiv preprint arXiv: /90. 2019.
Leo Gao, Jonathan Tow, Stella Biderman. Sid Black, Anthony DiPofi, Charles Foster, Laurence Golding, Jeffrey Hsu, Kyle McDonell, Niklas Muennighoff, Jason Phang, Laria Reynolds, Eric Tang, Anish Thite, Ben Wang, Kevin Wang, and Andy Zou. A framework for few-shot language model evaluation, September 2021a. URL https://: org/ wttrg//...28//eeeee 5281/zenodo. 5371628
Luyu Gao, Yunyi Zhang, Jiawei Han, and Jamie Callan, Scaling deep contrastive learning batch size under memory limited setup. arXiv preprint arXiv:2101.066 0069083,,, 2021b.
Xinyang Geng. Easylm: A simple and scalable training framework for large language models, 2023. URL https://gtttttbbb com/young-- haallyouungggggggggee /EasyLM
Xinyang Geng and Hao Liu. Openllama: An open reproduction of Ilama, May 2023. URL https plammpaaaanaamaaaaaaaaa
Tygithu..cmm comm
Mandy Guo, Joshua Ainslic, David Uthus, Santiago Ontanon, Jianmo Ni. Yun-Hsuan Sung, and Yinfei Yang. Longt5: Efficient text-to-text transformer for long sequences. arXiv preprint arXiv:2112.07916, 2021.
Yaru Hao, Yutao Sun, Li Dong, Zhixiong Han, Yuxian Gu, and Furu Wei. Structured prompting: Scaling in-context learning to 1, 000 examples. CoRR, abs/2212.06713 13, 2022. doi: 10.48550/ arXiw.2212.06713. URL URL httpp:tts:///0 org/10 448550/ariiv 2212. 06713
Adi Haviv, Ori Ram, Ofir Press, Peter Izsak, and Omer Levy. Transformer language models without positional encodings still learn positional information, 2022.
Alex Henry, Prudhvi Raj Dachapally, Shubham Shantaram Pawar, and Yuxuan Chen. Query-key normalization for transformers. In Trevor Cohn, Yulan He, and Yang Liu, editors, Findings of the Association for Computational Linguistics: EMNLP 2020, Online Event, 16-20 November 2020, volume EMNLP 2020 of Findings of ACL, pages 4246 4253. Association for Computational Linguistics, 2020. doi: 10. 086533v1/202D1 18653/v1/2020 findings- Sectangssemnlpp mmmlp.379 379
URL https://doi. org/10
Eduard Hovy, Laurie Gerber, Ulf Hermjakob, Chin- Lin, and Deepak Ravichandran. Toward semantics-bart based answer pinpointing. Proceedings of the First International Conference on Human Language Technology Research, 2001. |
https::///// wwwww/ aclweb. org/ anthology/
H01- 1069
Nihal Jain, Dejiao Zhang, Wasi Uddin Ahmad, Zijian Wang, Feng Nan, Xiaopeng Li, Ming Tan, Ramesh Nallapati, Baishakhi Ray, Parminder Bhatia, Xiaofei Ma, and Bing Xiang. Contraclm: Contrastive learning for causal language model, 2023.
12

---

Albert Qiaochu Jiang, Wenda Li, Szymon Tworkowski, Konrad Czechowski, Tomasz Odrzygóždž, Piotr Mitos, Yuhuai Wu, and Mateja Jamnik. Thor: Wielding hammers to integrate language models and automated theorem provers. In Alice H. Oh, Alekh Agarwal, Danielle Belgrave, and Kyunghyun Cho, editors, Advances in Neural Information Processing Systems, 2022. URL. htttss//opperreiiet wtwpaarrvwiiieeeeeornniieeeeee ttlp:oorrrmnidddltyt----------
Jeff Johnson, Matthijs Douze, and Hervé Jégou. Billion-scale similarity search with gpus, 2017.
Jean Kaddour, Oscar Key, Piotr Nawrot, Pasquale Minervini, and Matt J. Kusner. No train no gain: Revisiting efficient training algorithms for transformer-based language models, 2023.
kaiokendev. Things im learning while training superhot. https://kaaiokendev...c htwp:/ddwww.ittu....o.... Coneneingg-cntext---e
2023.
Urvashi Khandelwal, Omer Levy, Dan Jurafsky, Luke Zettlemoyer, and Mike Lewis. Generalization through memorization: Nearest neighbor language models. arXiv preprint arXiv: 1rXivv/911.000 1.00/72,
2019.
Denis Kocetkov, Raymond Li, Loubna Ben Allal, Jia Li, Chenghao Mou, Carlos Muñoz Ferrandis, Yacine Jernite, Margaret Mitchell, Scan Hughes, Thomas Wolf, Dzmitry Bahdanau. Leandro von Werra, and Harm de Vries. The stack: tb of permissively licensed source code. Preprint, 2022.
Taku Kudo and John Richardson. SentencePiece: A simple and language independent subword tokenizer and detokenizer for neural text processing. In Proceedings of the 2018 Conference on Empirical Methods in Natural Language Processing: System Demonstrations, pages 66-71, Brussels, Belgium, November 2018, Association for Computational Linguistics. doi: 10.18653/v1/ D18-2012. / oT1 2012
Aitor Lewkowycz, Anders Johan Andreassen, David Dohan, Ethan Dyer, Henryk Michalewski, Vinay Venkatesh Ramasesh, Ambrose Slone, Cem Anil, Imanol Schlag, Theo Gutman-Solo, Yuhuai Wu, Behnam Neyshabur, Guy Gur-Ari, and Vedant Misra. Solving quantitative reasoning problems with language models. Alice H. Oh, Alekh Agarwal, Danielle Belgrave, and Kyunghyun Cho. editors, Advances in Neural Information Processing Systems, 2022. URL / het:/forum?id:1FITZER...
Xin Li and Dan Roth. Learning question classifiers. In COLING 2002: The 19th International Conference Computational Linguistics, 2002. URL btttts::/www..ccweeb. anthooogy 602-1150)
Yujia Li, David H. Choi, Junyoung Chung, Nate Kushman, Julian Schrittwieser, Rémi Leblond, Tom Eccles, James Keeling, Felix Gimeno, Agustin Dal Lago, Thomas Hubert, Peter Choy, Cyprien de Masson d' Autume, Igor Babuschkin, Xinyun Chen, Po-Sen Huang, Johannes Welbl, Sven Gowal, Alexey Cherepanov, James Molloy, Daniel Mankowitz, Esme Sutherland Robson, Pushmeet Kohli, Nando de Freitas, Koray Kavukcuoglu, and Oriol Vinyals. Competition-level code generation with alphacode. CORR, abss/2233078/.. 2022. doi: 10.485500/rXiv.2200 wtt.:/xi..223.77788
Erik Lindgren, Sashank Reddi, Ruiqi Guo, and Sanjiv Kumar. Efficient training of retrieval models using negative cache. In M. Ranzato, A. Beygelzimer, Y. Dauphin, P.S. Liang, and J. Wort- man Vaughan, editors, Advances in Neural Information Processing Systems, volume 34, pages 4134-4146. Curran Associates, Inc., 2021. URL https::/pproceeddings nww.ips. cc/paper. .ddf
files/ f1t1://rrrr2222//222226666666660
Dacheng Li Rulin Shao Anze Xie Ying Sheng Lianmin Zheng Joseph F. Gonzalez lon Stoica Xuezhe Ma and Hao Zhang. How long can open-source llms truly promise context length?, June 2023.
URL htttps:/immss:
arg/blogg 20-2000-22992 longcha))
Maciej Mikula, Szymon Antoniak, Szymon Tworkowski, Albert Qiaochu Jiang, Jin Peng Zhou, Christian Szegedy, Lukasz Kuciński, Piotr Miloš, and Yuhuai Wu. Magnushammer: A transformer- based approach to premise selection, 2023.
Amirkeivan Mohtashami and Martin Jaggi. Landmark attention: Random-access infinite context length for transformers. CoRR, abs/2305. 16300, 2023. doi: 10.485550aaxXxv222555.60000 URL.
https: / //oii org/10 10. 48550/ aavvv 2305. 16300
13

---

MosaicML. Introducing 3pt-30b: Raising the bar for open-source foundation models, 2023. URL WWW mmmaiccilm......
www.mooaaccll..co//boo////g/// wwtp://g/ooooooee/////eeo mat00g/mpp33000 Accessed: 2023-06-22.
Piotr Nawrot, Szymon Tworkowski, Michal Tyrolski, Lukasz Kaiser, Yuhuai Wu, Christian Szegedy, and Henryk Michalewski. Hierarchical transformers are more efficient language models. CoRR, abs/21 137111 2021. URL htt.:hhtt/////t http://aawaw.../aaa//////1
wrgg/bs/////.....// 13711
Piotr Nawrot, Jan Chorowski, Adrian Lańcucci, and Edoardo M. Ponti. Efficient transformers with dynamic token pooling, 2023.
Erik Nijkamp, Hiroaki Hayashi, Tian Xie, Congying Xia, Bo Pang, Rui Meng, Wojciech Kryscinski, Lifu Tu, Meghana, Chen Xing, Jesse Vig, Lidiya Murakhovs' ka, Chien-Sheng Wu, Silvio Savarese, Yingbo Zhou, Shafiq Rayhan Joty, and Caiming Xiong. Long sequence modeling with xgen: A 7b Ilm trained on 8k input sequence length. Salesforce Al Research Blog, 2023. URL https www.:/oom/xxgi/e
//blog. //blog. www.sssseooreaaeeeeeerrereeeeeee
Stanislas Polu and Πya Sutskever. Generative language modeling for automated theorem proving. CoRR, abs/2009.03320 0039999900009..............,,3,3 2020. UtLptttss//aaxxi org/ abs/ 2009. 03393
Alec Radford, Jeffrey Wu, Rewon Child, David Luan, Dario Amodei, Ilya Sutskever, et al. Language models are unsupervised multitask learners. OpenAl blog, 1(8):9, 2019.
Alec Radford, Jong Wook Kim, Chris Hallacy, Aditya Ramesh, Gabriel Goh, Sandhini Agarwal, Girish Sastry, Amanda Askell, Pamela Mishkin, Jack Clark, et al. Learning transferable visual models from natural language supervision. In International conference on muchine learning, pages 9748-8763. PMLR, 2021.
Jack W Rae, Anna Potapenko, Siddhant M Jayakumar, Chloe Hillier, and Timothy P Lillicrap. Compressive transformers for long-range sequence modelling. arXiv preprint, 2019. URL
https://arxiv. org/ abt/ 05507
Colin Raffel, Noam Shazeer, Adam Roberts, Katherine Lee, Sharan Narang, Michael Matena, Yanqi Zhou, Wei Li, and Peter J. Liu. Exploring the limits of transfer learning with a unified text-to-text transformer. arXiv e-prints, 2019a.
Colin Raffel, Noam Shazeer, Adam Roberts, Katherine Lee, Sharan Narang, Michael Matena, Yanqi Zhou, Wei Li, and Peter J. Liu. Exploring the limits of transfer learning with a unified text-to-text transformer. CoRR, October 2019b.
Nir Ratner, Yoav Levine, Yonatan Belinkov, Ori Ram, Inbal Magar, Omri Abend, Ehud Karpas, Amnon Shashua, Kevin Leyton-Brown, and Yoav Shoham. Parallel context windows for large language models, 2023.
Uri Shaham, Elad Segal, Maor Ivgi, Avia Efrat, Ori Yoran, Adi Haviv, Ankit Gupta, Wenhan Xiong, Mor Geva, Jonathan Berant, and Omer Levy. Scrolls: Standardized comparison over long language sequences, 2022.
Jianlin Su, Yu Lu, Shengfeng Pan, Bo Wen, and Yunfeng Liu. Roformer: Enhanced transformer with rotary position embedding. CoRR, abs/2104.09864, 2021. /Rtpss:s:///r org/ abs/ 0184.
Simeng Sun, Kalpesh Krishna, Andrew Mattarella-Micke, and Mohit lyyer. Do long-range language models actually use long-range context?, 2021.
TogetherComputer. Redpajama: An open source recipe reproduce Ilama training dataset, 2023.
URL http:httt:///tthiiee hotp.//////oggeeeercopuutteee//e Ratp:e/teeerccoppptee///eeaaaa ReddajedPPajjajj ama- -
Hugo Touvron, Thibaut Lavril, Gautier Izacard, Xavier Martinet, Marie-Anne Lachaux, Timothée Lacroix, Baptiste Rozière, Naman Goyal, Eric Hambro, Faisal Azhar, et al. Llama: Open and efficient foundation language models. arXiv preprint arXiv:2302.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N. Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. CoRR, hbss/1066.03722 2017. URI.
http:// /tatpi// www.xii...rr/.ee. hts///b////// 1706. 03762
14

---

Yuhuai Wu, Markus Norman Rabe, DeLesley Hutchins, and Christian Szegedy. Memorizing trans- formers. In The Tenth International Conference on Learning Representations, ICLR 2022, Virtual Event, April 25-29, 2022. OpenReview.net, 2022. URL Tr TrjbxzRenf-
Manzil Zaheer, Guru Guruganesh, Kumar Avinava Dubey, Joshua Ainslie, Chris Alberti, Santiago Ontanon, Philip Pham, Anirudh Ravula, Qifan Wang. Li Yang, et al. Big bird: Transformers for longer sequences. Advances in neural information processing systems, 33:17283-17297, 2020.
Biao Zhang and Rico Sennrich. Root mean square layer normalization, 2019.
Zexuan Zhong, Tao Lei, and Danqi Chen. Training language models with memory augmentation. In Yoav Goldberg, Zornitsa Kozareva, and Yue Zhang, editors, Proceedings of the 2022 Conference on Empirical Methods in Natural Language Processing, EMNLP 2022, Abu Dhabi, United Arab Emirates, December 7-11, 2022, pages 5657-5667 5673. Association for Computational Linguisties, org/ 2022, emnlp- 1alppmmann... 382
2022.
15

---

Broader Impact
Recent rapid developments in language models have brought lot of new capabilities, At the same, these raised concerns about the social impact and very animated discussions the community. Our work develops a generie technique, which in principle, could be applied to virtually any language model and thus, by extending their capabilities, exacerbate threats. We note, however, that FoT does not create any new threats. Thus, we refer to the existing body of knowledge on the broader impact of language models, see e.g. Borgeaud et al. 2002]..
A LONGLLAMA
A.1 Architecture
Ope [Geng and Liu 2023 is an open-source reproduction of LLaMA [Touvron et al. 2023 uses a decoder-only architecture with rotary positional embeddings, and a few changes including pre-normalization with RMSNorm (Zhang and Sennrich, 2019], and SiLU activation [Elfwing et al. A SentencePiece [Kudo and Richardson 2018 Tokenizer with 32k vocabulary size is used.
2017
A.2 Extending context length with FoT
Positional encodings To achieve backward compatibility with the original LLaMA, we retain positional encodings in the local context. The tokens outside the local context are assigned the same position as the first token in the local context.
Dense attention to longer context To make the implementation simpler and less dependent on external software, we resign from using kNN lookup and perform attention over the whole memory. We have found only marginal performance differences between those two approaches to memory attention.
Crossbatch details For the 3B LONGLLAMA model, we set L = (6, 12, 18} as the memory layers. We vary the number of additional contexts d E (0, 2) across elements of the batch by dividing batch entries into four segments of equal size. Elements from the first segment only see local context (d 0). Elements from the second segment see two additional contexts (d 2), one from the same document (positive) and one from different one (negative). Elements from the third segment see
three additional contexts, two positives, and one negative. The last segment consists of elements exposed to three additional contexts coming from the same document. We abbreviate this setup as (0, 0), (1, 1), (2, 1), (3, 0).
For the 7B LONGLLAMA model, we set L = {8, 16, 24} as the memory layers. Here we divide batch entries into four segments and use the following setup: (0, 0), (1,2), (2, (3, 4).
Hyperparameters We follow the choices of Ope with respect to most of the hyperparam- eters, including using the same optimizer. During fine-tuning, we use batch size of 256K tokens and constant learning rate of 2e-5, which is lower than the learning rate at the end of Lae training (3c-5 after 1T tokens), and weight decay of 0.01.
A.3 LLaMA fine-tuning dataset
We use mixture based on RedPajama [TogetherComputer 2023 and The Stack (ooeekov et al. 2022] with the following proportions of each subset:
All subsets apart from python are taken directly from RedPajama. For the python subset, we gather Python source code from The Stack and, to obtain long documents for training, concatenate files that in the same subdirectory in random order, using similar procedure for the GitHub dataset in Section|H] H Additionally, we filter out short documents for some subsets of the original RedPajama, namely shorter than the Min. doc. length column indicates.
In case one document is too short to span across several contexts for crossbatch, then we concatenate it with the next document from the dataset.
16

---

python Table Proportions of RedPajama subsets for the LONGLLAMA fine-tuning mixture. For python subset, data from The Stack is used (see text for details). We only train on documents with length being least Min. doc. length, if specified (otherwise we train on all documents from that subset). The horizontal line separates long-context and short-context subsets.
Subset
Sampling proportion (%)
Min. doc. length
arxiv
25
4096
book
10
common_crme _rawl
29
c4
1024
github
2048
stackexchange
1024
wikipedia
1024
A.4 Language Model Evaluation Harness
To ensure that the performance of LoNGLLANE LLAMAs has not degraded in short context scenarios, evaluate our models on the Language Model Evaluation Harness benchmark [Gao et al.] [222a]] Table 5 compares results with OpenLLaMA [Geng and Liu 2023). Similarly to the authors of OpenLLaMA, we omit CB and WSC tasks.
Table 5: Model comparison across different tasks/metrics Language Model Evaluation Harness. The LONGLLAMA models were evaluated without context extension (i.e. as standard OpenLLaMA models). Results indicate that LONGLLAMAs maintain good performance in short context scenarios.
OaenLLaMA
LongLLaMA
Ope
LongLLaMA
Task/Metric
3B
3B
7B
7B
anli rl/acc anli r2/ace anli r3/acc are challenge/ace
0.33
0,32
0.33
0.35
0.32
0.33
0.36
0.37
0.35
0.35
0.38
0.36
0.34
0.34
0.37
0.37
are_challeegg/acc_nome
0.37
0.37
0.38
0.38
arc_ arc_easy/acc
0.69
0.68
0.72
0.70
vvc_easy/acc_nore
0.65
0.63
0.68
0.66
boolq/ace
0.68
0.68
0.71
0.71
hellaswag/ace
0.49
0,48
0.53
0.52
hellaswag/acc_nnn norm
0,67
0.65
0.72
0.71
openbookqa/acc
0.27
0.28
0.30
0.30
openbookqa/acc_noom norm
0.40
0.38
0.40
0.41
piqa/acc
0.75
0.73
0.76
0.75
piqa/ace_non _orm
0.76
0.75
0.77
0.76
record/em
0.88
0.87
0.89
0.89
record/f1
0.89
0.87
0.90
0.90
rte/acc
0.58
0.60
0.60
0.59
truthfulqa_mn 1m/maamcccl
0.22
0.24
0.23
0.24
nuthfulqa_me/mel
0.35
0.38
0.35
0.35
wic/ace
0,48
0,50
0,51
0.50
winogrande/ace
0.62
0.60
0.67
0.67
Average score
0.53
0.53
0.55
0.55
A.5
Question answering over research papers
We evaluate the context utilization of our model on the validation set of Qasper [Dasigi et al. 2021
from SCROLLS (Shaham et al. 2022). Details are in the Table 6
17

---

Table Zero-shot performance with different context lengths on the validation subset Qasper Dasigi 2021].
We use the implementation from Language Model Evaluation Hamess [Gao et al. 202 In the Harness implementation, yes/no questions are evaluated separately from open questions. Observe that LONGLLAMA 3B benefits from the extended context.
Context length
OpenI 3B
LONGLLAMA 3B
LLaMA 7B
LongChat 7B
2K
18.7
18.7
18.7
19.4
4K
20.7
21.2
6K
23.2
25.0
8K
26.6
28.8
B Architecture
This section describes the architecture and crossbatch details for non- non-LLaMA-baan models presented in this paper. The main differences are that for LLaMA-based models (LONGLLAMA) we maintain the positional encodings (with slight modification detailed in A.2 do not introduce the attention temperature parameter, and replace kNN with full dense attention.
B.1 Transformer models
For non -LaMA- -LaMee models we use the transformer architecture introduced in [Vaswani et al. 2017] with few standard changes. First, we use only the decoder without the encoder part. Secondly, we perform layer normalization before the input of both the attention and feed-forward modules. Additionally, we use Rotary Position Embedding [Su et al. 202T]. normalize keys and queries (Henry et al. 2020]. and introduce learnable temperature parameter for each attention head.
The hyperparameters for each model size be found in Appendix E For training the models on PG-19, we use the standard T5 tokenizer with 32k vocabulary [Raffel et al. [019b]. The larger models in Section [5.2 are trained with custom SentencePiece tokenizer Kudo and Richardson 2018 with 64k vocabulary size.
B.2 Memory attention layer
Memory attention layer & one of the transformer layers, which has access the additional context M. The memory stores (key, value) pairs. For each query in f, we retrieve the k most matching entries from M and them compute the attention value. More precisely, algorithm to pick Mtop == {(ee), valuei),... (keyk: valuex)} C M such that {gg ) are the top inner products in M. These are merged with the part of the local context before denoted as Czy and used to compute the attention value using the standard Transformer formula:
use the kNN J:=1....
=
Σ
s(key) v.
(key,e) McpPUCCo
where s(key) is the softmax score for key. This softmax is calculated as follows:
softmar ([ key)
5
MepEMo
where is a temperature parameter. In this approach, do not distinguish between the local context and the memory.
Another way of integrating Mrop is via gating. In this approach, we separately compute the attention value UM for Mtop and for the local context ve (using the standard Transformer formula). Then we use a gating mechanism to combine them:
= UM g + vC (1 (1- 9), g = obby):
where a the sigmoid function and by is a trainable bias. The gating approach was proposed in Wu al. 2022), formula [Wu et al. 2022 (2)).
18

---

We found our approach, i.c. using T to be equally effective, see Figure At the same time, ① is simpler and does not require additional parameters. Thus, we use it in our experiments.
For kNN lookup, we use the exact KNN search implemented FAISS [Jonnson et al. 2017]. The memory attention layer does with (key, value) pairs processed by l beforehand. In the single-doc setting, the memory is erased after each document.
positional encodings. The memory is populated incrementally
We do not use the r parameter for LONGLLAMAs as their architecture does not normalize keys and queries. For LONGLLAMAs, also replace the kNN search with dense attention and retain positional encodings (see Appendix(A22
B.3 Crossbatch training procedure
In FoT we choose a subset L of the attention layers for later augmentation with the memory of (key, value) pairs. Let an attention layer from During the training we expose this layer to mixture of (key, value) pairs from the current local context, Ceurr, and the previous local context and d contexts from other documents, Previ see also Figure for an illustration. this by modifying the input pipeline so that each batch index corresponds to different document (the batch index occupied by each document is fixed from the moment we load the document till we finish processing it).
achieve
More specifically, we embed the previous and the current local context for each document in the batch. Then we use Crey as source of the previous local context for a given document and d 1 contexts from other documents. For each element of the batch, the choices of those d additional contexts are fixed. We disable positional encoding in É, as we envision it to handle global information.
To be more precise, for each document ò within the batch and query q from the layer we create the set pº consisting of (key, value) pairs from the previous local context of document 5 along with pairs from 1 contexts gathered from Cprev. The attention value for q is given by
Σ
s(key) .,
where C6 consists of (key, value) pairs that preceded q in its local context and s(key) is the softmax score for key. We use softmax with learnable temperature
softmar ([
key)
Note that the only difference between and (2) is the source of the additional (key, value) pairs: p². This, in particular, implies that all the operations with respect to the previous context differentiable.
The number of different documents is equal to by (the batch size, i.e. each document has a separate index in the batch). Assume that document 8 has index i. We include into pd all tokens from Cpren with the batch indices {ú, (i + 1) mod bg...., - 1) mod bg}.
19

---

Single-doc Perplexity PG-19
25
20
15
10
No Gating
Sigmoid Gating
0
100k
200k
300k
400k
Step
Figure 4: Perplexity (on the test set) during training on PG-19. We train two Memorizing Transformer models [Wu al. 2022), one with original gating and one without (i.e. memory attention shared with local attention as described in Tp). We use the single-doc setting with 16k memory.
B.4 Qualitative analysis
Table provides brief qualitative analysis of FOT. It shows that the model can handle distractions and retrieve the parts of the character name from multi-document memory the PG-19 task dataset and appropriate definitions from the large dictionary (dictionary lookup task).
B.5
Memorizing Transformer
The Focused Transformer shares many similarities with the Memorizing Transformer [wu al. 2022 In this section, we summarize the differences between those two models.
Training The key difference between these two methods lies the training procedure. Our method uses crossbatch, see Section B.3 which, in nutshell, is the standard transformer training objective, but we additionally attend tokens from the previous context window, both from the same and different documents, see Appendix for for details. The Memorizing Transformer trains on tokens retrieved from the same document (it was envisioned for single-doc).
This has a few important consequences:
FoT does not use memory during training, while MT does. This may result in faster training; moreover, FoT always uses the most up-to-date values, while MT uses the values from memory, which may be outdated.
FoT differentiable through all (key, value) pairs, while MT does not differentiate through the retrieved tokens. We argue that this is key for joint training of well-structured key, value, and query embeddings and, consequently, good model performance.
FoT does require long documents in the training set, while MT does in order to capture long dependencies in memory. This is practically important, as many popular datasets consist of short documents.
We speculate that there may be benefits in blending these two approaches. One can, for example, argue that MT is better providing "hard' negatives for the model. We provide proof-of-concept experiment in Appendix C and leave this for future work.
Inference Both models use very similar memory attention layer. The difference how the retrieved (key, value) pairs are integrated. FoT treats the retrieved information in the same way as the local context. MT uses gating mechanism. Details are provided in Section B.2
20

---

Table Example elements retrieved by search along with scores on PG-19 dictionary lookup task. The first column shows the token associated with particular query (bolded) along with fragments of its context. The second column shows tokens associated with the keys retrieved by kNN for this query. The kNN shows what fraction of the attention mass dedicated retrieved keys corresponds Focus score calculated by taking
particular key. The following keys, calculating attention
key along 32 preceding weights them, and checking what fraction of attention the retrieved key gets. In the PG-19 setting the model was equipped with the memory of size 1096 spanning across parts of 8 documents. In the dictionary lookup setting the model was provided with memory of size 16M..
kNN
Focus
Text
kNN Results
Score Score
PG-19
Then we're here with the closed carriage at !en-! [They go together into the library.
0.69
0.99
DARBEY. [To SHE BA
Oh! [He sinks on to the settee with vacant stare, his arms hanging helplessly. DAR- BEY. [To SHE BA .J There his career a burden to him!
S HE BA takes the gold-rimmed pince-nez which hangs upon THE DEAN'S waistcoat and places before eyes
0.23
0.92
Papa! SHE BA Papsey! [THE DEAN rouses himself, discovers his children and removes his hat.
0.025
0.92
Then we're with closed carriage at ten-! [They go together into the library. DARBE [To SHEBA.]
0.71
0.99
THE DEAN gives DAR BE Y severe look, and with an important cough walks into the Library. The men and the girls speak in undertones.
Oh! [He sinks on to the settce with a vacant stare, his arms hanging helplessly. DARBE Y [To SHEBA.J There- now his career burden to him!
0.19
0.99
Oh, Salome! Papa! Papa! TARVER. The Dean? DARBE The Dean!
0.08
0.99
Dictionary Lookup Task
<q> 14 23
40
56
40
05
0.96
0.99
<q> 30 07 <リ>
36
17
26 63
<k> 30 55 07 23 <v>
36
26
0.88
1.0
<<> 10 41 26 39 <ソ>
48
38 11 24
<k> 10 41 26 39
<V>
48 38
11
24
0.87
0.99
C
Ablations
In this section, focus two key properties of crossbatch training procedure: differentiability and the inclusion negatives. We also discuss the relation to Memorizing Transformer in terms of the training protocol and memory integration. We refer to Appendix B.5 for a detailed technical description of differences between FoT and Memorizing Transformer.
C.1 Impact of differentiable keys and values
compare FOT to Memorizing Transformer, which uses a non-differentiable memory of keys and values during training. In the multi-doc experiment presented in Figure both MT and FoT are trained with local context of 512. We observe that FOT significantly better when the context expanded during inference, which confirms that differentiable keys and values are beneficial.
21

---

Table 8: Perplexity on PG-19 in the single-doc doc setting for various local context lengths during training. these experiments, we used the same context length both during training and evaluation.
Context Length
FoT d=1
MT
512
14.18
14.68
1024
14.17
14.46
2048
14.11
14.43
Mult-doo Perplexity on PG1D
Multi doc Perplexity PG19
17.5
-bul ded
-ingledds MT
15.2
-=FoT 082004
15
16.5
Transfomerr (no mem)
15.5
14.0
15
14.4
14.5
Singla-doc (Singla-doc avai)
0
100k
200K
300%
4C0K
500K
100k
200k
500k
400K
SD0K
Figure 5: Perplexity on PG-19 in the multi-doc set- ting. Both FoT and Multi-doc MT were trained with local context of size 512. During training, Multi-doc MT utilized memory of size 4096 shared across doc- uments. Diffcrentiability of keys and values results in better perplexity.
Figure 6: Importance of negatives in the multi- document setting, We compare FoT trained with d = to the one that started with d 2 and later switched
64. For additional comparison, we show the performance of MT trained with the memory of size 16K.
We also check whether differentiable keys and values can improve the performance in the single-doc setting. For this, compare FoT with d = 1 to MT with memory consisting of the previous local context. Table confirms that differentiable keys and values can also help this scenario.
C.2 Importance of negatives
We reaffirm the importance of negatives multi-document setting. previous experiments in Figure we already observed that increasing the number of negatives (i.c., increasing d) results in more attention mass being dedicated relevant tokens. In Figure we additionally show that the lack of negatives in training (d = 1) results in significant deterioration in model perplexity when the context length grows. This confirms that both using negatives and differentiability are important for FoT to work well.
C.3 Relation to Memorizing Transformer
Memorizing Transformer Wu et al. 2022 closely related to our method. The two key differences 1) the training protocol and 2) how the memory is integrated into the model. In this section, provide additional insights into these differences.
Training protocol In the previous sections, we have discussed the benefits of the crossbatch training, namely using the contrastive-inspiee objective and backpropagating through the previous context. potential advantage of the MT approach is that it is exposed to the whole memory during training (instead of just the previous context). We performed a groof-of-concey experiment combining the two approaches to explore this further. Namely, trained the model for 499k steps using crossbatch and fine-tuned it with the MT objective for steps. Interestingly, we observed a significant improvement compared to the MT training with the same step budget, see Figure 7 We believe there is further room to explore various training protocols combining the best of both worlds.
Memory integration FoT uses simple memory integration approach where the (key, value) pairs retrieved by KNN lookup are treated the same way as the local context. contrast, MT uses gating mechanism, weighted average of the memory, and local values; see details in Appendix B.2 We evaluated both approaches and found no difference in performance between these two memory integration methods. However, decided to approach because it does not require
22

---

Single-doc Perplexity on PG19
15.2
15
-Siigl--in
14.8
Single-dio MT 6Bk
Foo de2 fine-uuned 65k
14.6
14.
14.
14
20k
40k
50k
60k
Tokens Context
Figure 7: Single-doc eval of FoT finetuned for 1 steps differentiable memory. It achieves lower perplexity than MT, which has this memory for the whole training (500k steps).
any architectural changes (and thus makes fine-tuning existing models easy). For these reasons, we recommend using it. We speculate that the reason why the gating is not needed in FoT is another benefit the fact that the crossbatch training backpropagates through the (key, value) pairs from the previous context Cpree in contrast to MT that cannot backpropagate there and needs to rely local context when computing gradients for keys and values. Another reason might be the fact that Oprea is embedded for each batch, and thus staleness (see
[Wu et al. 2022 Section 3.2]) is avoided.
D Additional figures
16
Multidoc Perplexity on PG19
Single
Transformer mam)
FoT d-4
FQT das
d=2-66
14.5
Singio-dac (Singlo-doc aval).
14
100k
200k
300k
400k
500K
Figure 8: Perplexity in the multi-doc setting. FoT was trained with local context of size 512 and different d. FOT 2->64 started with d = 2 and then switched to = 64. Single-doc MT was trained with memory size of 16K. As inerease the memory size, the number of distractions (irrelevant keys) increases, making the task harder. Single-doc MT evaluated the single-doc setting soft lower bound since lacks distractions.
23.8
de--doc Perplesityoo on
de-
del, We-2
del, Wel
23.2
28
Singledoc
22.8
22.6
22.4
10k
20k
30k
40k
50k
80K
Tokens Context
Figure 9: Zero-shot performance on PG19 of FoT pretrained on C4. Model fine-tuned with the crossbatch dimension outperforms the with 1. Using the double (w 2) training context of 2048 beneficial.
23

---

E Hyperparameters
Table 9 shows hyperparameters used in experiments. We used context length 512 unless stated otherwise. Appendix F. Section 5.3 Section 5.5 we use the total batch size of 32K tokens. In Section 5. and Section 5.4 the total batch size 128K tokens.
Table 9: Hyperparameters for different model sizes. The batch size is given in number of tokens. For the 37M model local context length was 256 for FoT and 512 for the baseline.
Hyperparameter
Value
Common
#Layers
12
Index of memory attention layer (l)
8
Optimizer
AdaFactor
Learning schedule
Inverse Square Root
Warmup steps
1000
B1
0.9
Model-specific
#Params
37M
184M
Max learning rate
0.02
0.01
Min learning rate
0.01
0.0005
Embedding dim
512
1024
Head dim
64
128
#Heads
8
FeedForward dim
2048
4096
Local context length
256/512
512
Batch size
32K/64K
32K
#Number of training steps
5k
500%:
For the experiments described in Section 5.3 and Section rameter sweeps:
5.5
we performed the following hyperpa-
Learning rate: (1e-2, 3e-3, 1e-1e-1- 1e-3), chosen: 1c-2,
Batch size: {8K, 16K, 32K), chosen: 32K.
For the dictionary lookup task (Appendix F we checked the following hyperparameters:
Learning rate: {4e-2, 2e-2, 1e-1, 5e 3, 3e 3--3),
3}, chosen: 2e-2.
Number of dictionary tokens training step: {16K, 32K). chosen: 32K. Note that during the training number of document tokens dedicated the dictionary is the same the number of tokens dedicated to questions.
For most of the other hyperparameter choices, we followed (Wu et al. 2022 to provide a fair comparison.
E.1 Schedule of d
5.3 5.4 models with d
In Sections
and 5.5 for models with 21, 8} we used constant schedule, and for
64
trained with d = 2 for 450k steps and switched to d 64 for the final 50k

---

steps. In Appendix F we trained with d = until the model reached 98% accuracy and then we switched to d 128. For the 184M model in Section 5.2 we randomly sampled from {2, 128) in each training step.
Dictionary lookup task
1
0.8
Accuracy
0.6
FoT Transformer
0.4
0.2
256
16K 256K Tokens
16M
Figure 10: Accuracy vs number of dictionary tokens in a dictionary look-up task. The task format is as follows: <k>ki <v>ei randomly selected keys. Accuracy determined by measuring the predicted values after <q> tokens. Models were trained examples containing 512 tokens and evaluated with extended context length. FoT demonstrates high accuracy even when the memory size is large. The bascline transformer fails already for 16K tokens. Error bars represent the minimum and maximum on 10 seeds.
<apk:<>>n <y>u..... where dictionary is provided, followed by queries
We propose a dictionary lookup task test whether the model trained using our erossbatch method can utilize large memory database. Documents in this task are split into two parts. The first part defines keys and their associated values using the records the format:
<k>, k1, k2, ks, k4. <v>, 01, 02 U3 U4+
where <k> is special token that denotes the beginning of the defining sequence,
The second part consists queries about the values associated with the previously defined keys. The queries are in the following format:
<q>, k1, k2, ks, k4, <v>, U1, 02 U3, U4.
where <q> is a special token that denotes the beginning of the query. We mask the loss so that for such a question, only v1, 02 V3, U4 included. We a vocabulary of 64 tokens, keys and values are described using 4 tokens.
documents comprising 512 tokens. The first half of each document consists of definitions, whereas the second one consists of questions. For FoT, we use a local context of 256, thus the model needs to the memory attention layer answer the questions correctly. We start with d = and increase to d = 128 as soon as the model is able to reach 98% training accuracy. During the inference, use k: 32 (the number of keys retrieved by kNN). As baseline, standard transformer model trained with the context length of 512. In evaluation, we test different local context lengths, which quickly leads very poor results.
During training,
In evaluation, we use longer documents but make only the last 256 tokens correspond to questions. That is, as the context gets bigger (the token axis Figure [O], the number of definitions increases, but the number of queries remains the same.
G FoT fine-tuning
For comparison in Table 3 our model is pre-trained for 100k steps with total batch size of 128 (128K tokens per step, with 1024 local context). Then we fine-tune both FoT and baselines for additional 10k steps with the same batch size. When fine-tuning FOT, we randomly sample from (2, 128) in each training step to prevent the model from overfitting to large additional context length during training.
25

---

H
Datasets
Section outlines essential details concerning the PG-19 and arXiv datasets employed in this study. Now, we will present details about the remaining datasets:
GitHub We obtained a large corpus of permissively licensed Github repositories using BigQuery. By filtering for specific file extensions (C, C++, Java, Python, Go, and TypeScript), we captured individual source code files that are often short but have numerous dependencies and cross-references within the repository. To preserve the structure while shuffling the files and subdirectories a random order, we concatenated all the files within each repository, treating subdirectories unit, similarly to Wu et al. [2022).
Isabelle The Isabelle corpus comprises formal mathematical proofs in the form of theories written in a formal language. We combined theories from The Archive of Formal Proofs (from October 2021) and the Isabelle standard library to create corpus of theories licensed open source. Each theory focuses on topics like foundational logic, advanced analysis, algebra, or cryptography and consists of multiple files containing proofs. Similar the GitHub corpus, the files within each theory are concatenated into single document. However, unlike the Github corpus, we arrange the files based on their import dependencies, ensuring that later files can utilize sub-theorems proven in earlier files.
I Hardware and technical details
We used TPU virtual machines from the Google Cloud Platform (GCP). Each TPU virtual machine has TPUv2 TPUv3 cores totaling 64GB 128GB of device memory, 96 CPU cores, and over 300GB of RAM. In larger-scale experiments (Section 5.2) we used machines with 32 TPUv3 cores. For training the LONGLLAMA checkpoints, TPUv3- 128 pod provided by the TPU Research Cloud was used, which we gratefully acknowledge.
J Randomness
To evaluate the significance of our results, we conducted multiple runs for selected experiments in our study. In Figure I0 we calculate error bars, showing the minimum and maximum value over 10 runs of the same experiment. For the arXiv baseline experiment in Appendix [K] we performed three runs with different random seeds and calculated their standard deviation, which is equal to 0.002 perplexity. However, due to resource constraints, we were unable to conduct multiple runs for all experiments. Our preliminary findings indicate that the observed variance was minimal compared to the impact observed from other factors under investigation.
For the calculation of test perplexities, we used 1M tokens.
K
Additional experimental results
This section presents additional empirical results, providing detailed comparison of FOT with the Memorizing Transformer [Wu et al., [202) baseline. Both models are trained for the same number of 500k steps with local context of 2K and evaluated on the arXiv dataset the single-document setup, following [Wu et al. 2022]. In particular, we study how models trained with given context length perform when evaluated with different context lengths. These experiments differ from those in Section as the models were both trained and evaluated on the same dataset (arXiv), unlike the C4 training and zero-shot evaluation done in Section [5.2]
The MT baseline in Table 10 with a memory length of 2K struggles to utilize additional context beyond 32K tokens effectively. The model trained with 8K memory performs significantly better when evaluated with longer contexts, showing further perplexity gains at 64K tokens. We observe diminishing returns when scaling up the training memory length 16K tokens and beyond.
Using the same setup, we study the performance of FoT while varying d and w configurations, similarly to Section 5.4 see Table II Parameter values 1 and w 2 correspond to additional
https:/wwww.isaaaapoom
26

---

context lengths of 2K and 4K, respectively. In an additional context length, FoT outperforms the MT baseline, which shows the importance of trainable keys and values (see also Section C.Ip. Moreover, we confirm the findings from Section 5.4 that works significantly better than d in all settings. Our best configuration achieves 2.148 perplexity with 4 K additional context during training, compared to 2. 164 of MT with 16K additional context.
comparison to MT with 2K
Table 10:
Memorizing Transformer: arXiv perplexity values for different training memory lengths and
evaluation context sizes
Eval Context
Training Memory Length
2k
8k
16k
32k
4K
2.309
2.334
2.348
2.365
8K
2.242
2.244
2.252
2.265
16K
2.215
2.206
2.206
2.215
32K
2.199
2.178
2.177
2.181
64K
2.195
2.169
2.166
2.168
128K
2.195
2.168
2.164
2.166
Table 11: FoT: arXiv perplexity values different parameter combinations evaluation context sizes
Evaluation Context
Parameter Combinations (w, d), Training Memory Length
(1, I), 2048
(1, 2), 2048
(2, 1), 4096
(2, 2), 4096
4K
2.292
2.299
2.305
2.309
8K
2.229
2.224
2.214
2.217
16K
2.206
2.194
2.178
2.178
32K
2.192
2.176
2.159
2.156
64K
2.187
2.171
2.152
2.149
128K
2.187
2.171
2.152
2.148
Code
In Listing 2 we show the FoTs attention code (i.e., the code for the memory attention layers and
crossbatch training), see Section 3 Appendix B.2 Appendix B33 We note that the changes to the code are small; they are localized to the memory layer (the other layers follow the standard transformer protocol) and do not require any new trainable parameters.
27

---

Listing 2: Memory attention: Let be memory attention layer. During the training, make & attend to the (key, value) pairs from the local context, previous local context, and 1 contexts coming from other documents. During the inference, queries from l attend to the local context and nearest neighbors retrieved from memory. For simplicity, we provide the code for one head and assume that the crossbatch dimension d is equal to the batch size.
def mem aattn_layee (Q1, K1, v1, C1, Km, Vm Kp, Vp, attn_acf, mode): "." Attention mechanism for crossbatch and memory attention
Arga:
Q1, K1, v1: tensors of shape [batch, ctx_len, dim]
with queries, keys and values from the local context
Km, Vm:
tensors of shape [batch, ctx_len, k, dim]
with most matching memory keys for along with associated values
each query from tensors of shape [batch, ctx_len, dim] with keys and values from the previous context scale factor used before softmax either training inference
Kp, Vp:
attn_scf
mode:
Returns:
y: a vector with shape [batch, ctx_len, dim]
# attention to the local context
atcal_attention janseinsum ("bqd, bkd ->bqk", Q1, K1) local_ local_attention apply_causal mask (local_ attention)
attn_acf
if mode "train":
In train mode, we additionally use previous context and batch-1 contexts from other documents. prev_attention jup.einsum ceinsum""bqd, ("bqd, ckd ,>bqck", Q1, Kp)
shape prev_attention.
adiiiiond_ateenion prev_attention reshape (shape [: (-1,)) elif mode "inference":
# In the inference mode, we additionally use nearest neighbors # retrieved from memory. We retrieve k (key, value)
pairs for each query.
memory_attention additional__ttention nemory_attention else:
enppeiinum ("bqd, bqnd bqnd bqn", Q1, Km)
raise Exception ("Only train and inference modes are supported")
additional attn_scf
# We merge the raw attention scores and calculate the softmax attennin jnp. concatenate [[ocal_attente)
additional_atentioonl, axis=-1) combined_weight. = jax nn. softmax Cccominnd__ateentiin axis=-1)
ctx_len Q1. shape [1]
veights weiihte__wights(.... additional.
:ctx_len]
sembined_weighted (..., ctx_len:)
y jnp. einsum (" (bbqk, bkd -> b local_ weights, v1)
if mode "train";
prev_weights shape prev_weights.... shaghe
additional_weighe
prev_veights prev_veights. reshape (esapee (-1, ctx_len))
y jnp. einsum ("bqck, ckd bqd", prev_weights, Vp)
else:
semory_weights = additional_
+= jnp einsum bebgnbbqnd -> bqd", memory_veights, Vm)
return y
28


Reviews:
[{"review_id": "vvVABLNnKR", "reviewer": "Reviewer_Sux9", "summary": "The paper proposes a key challenge (distraction issue) when extending the attention layer to external (key, value) pairs, either from previous context of the same document or from other documents. This paper proposes the Focused Transformer (FoT), a model that can utilize a long context by retrieving kNN (key, value) pairs from memory. The architecture of FoT is similar to Memorizing Transformer (how retrieved items are integrated is different), but the training technique is different. This paper proposes to use crossbatch to train the model, which simply includes both previous context of the same document and also contexts from different documents. So the model can learn how to distinguish between useful information and distraction.", "questions": "* L144: why is there a positive and negative? Did you actually distinguish positive/negative in the contrastive loss?\n* Table 2: why do you use token-level accuracy here but use perplexity elsewhere.\n* How do you define external memory? Anything that is out of the original input can be called external.", "limitations": "The authors have well discussed the limitations of the research.", "rating": 5, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "* The paper focuses on an important and potentially impactful problem that extends the context window size of transformer models.\n* The proposed model and the training method is straightforward and easy to implement.\n* The experiments are fairly solid and have supported the points that the paper has raised.\n* I also like the synthetic dictionary task as an evaluation task to test if a language model has the ability to attend to desirable context and gather necessary information.", "weaknesses": "1. One main weakness of this paper is that it is not clear to me if the proposed model is designed to incorporate long context or to incorporate external memory which may come from a large corpus. \nAs the paper includes both single-doc and multi-doc experiments, I assume it is the latter case. Based on this assumption, the paper identifies “distraction issue”, which essentially means the model attends to more other documents when considering more documents during inference.\n\nI did not get the point here. If we allow the memory to contain items from other documents, we actually expect the model to extract useful information instead of treating them as “distraction”. When showing the problem of “distraction issue” (Fig 3), we only consider 64 documents, which is far from the real case where we want to use external corpus. Indeed, prior works have shown that using more external information can help the model to achieve better results instead hurting it, e.g., (Khandelwal et al., 2019), (Zhong et al., 2022).\n\nWhile if the paper is mainly considering only the long-context cases, we also should question if the “distraction issue” exists during inference or not, because we can always control the model to only attend to a single document.\n\nSo, I don’t think the paper is well-motivated or well-positioned. The authors are encouraged to clearly state what testing situations they are addressing and re-evaluate the proposed issue in that situation.\n\n2. The paper can be presented better. For example, it is never clear to me what “external memory” exactly means in the paper. Why “previous local context of the given document” is called positive and “contexts from unrelated documents” is called negative, given there are really no distinctions in the training objective (correct me if it is not the case)?\n\n3. The proposed CrossBatch technique is based on including both previous contexts from the same document and contexts from different documents in the same training batch. This method is very similar to the data batching method proposed in (Zhong et al., 2022). The authors are encouraged to discuss the differences.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their constructive feedback. \n\nWe admit deficiencies in clarity raised by the reviewer. We note that we focus on *the long-context capabilities*; see also the general response. In our experiments, we tested FoT in both single-doc and multi-doc scenarios to assess its potential usefulness. We found that FoT improves perplexity on single, long documents (see Section 4.5), and we believe this makes it applicable to generic long-context language modeling, which is strongly confirmed by our new experiments with large models. At the time of writing, we decided to keep some multi-doc experiments, e.g., to illustrate the distraction issue, which already impairs the model’s perplexity significantly at a relatively small scale (64 documents, see Figures 3,7). However, in retrospect, we recognize this might be confusing. To amend this, we apply the steps described in the general answer; in particular, we indicate the long context focus more explicitly. \n\nWe also note that there are practical applications where the multi-document setting is well-motivated, particularly repository-level code generation. We hope that our method could be scaled up to open possibilities for handling the entire repositories of code in context (possibly ~1M tokens in large codebases), which we plan to attempt in future work. \n\n\n**Regarding the 'external memory'**. By this, we understand anything outside the local context window, i.e., anything that is accessed additionally by memory attention layers. This clarification has been added to the paper. To make this name less confusing, we changed it to 'additional context'. \n\n**Regarding 'positive and negative examples'**. \n\nOur method is inspired by contrastive learning in the way how data is presented to the model. We assume the model is presented with distractions (possibly irrelevant documents) in the training phase. The previous local context (from the current document) is mixed with contexts from other documents in the batch. Intuitively, this 'forces' the model to learn to differentiate between the positives (tokens from the current document, which are likely to be useful) from the negatives (tokens from other documents, which are unlikely to be useful). We note that this is not standard contrastive learning, as we do not have a separate contrastive loss. We only use the standard language modeling loss. We have added a clarification to the paper.\n \n**Regarding Table 2**. We agree that it is not the best way to compare the models, but we were constrained by pre-trained models (as different tokenizers are used, comparing perplexity is not informative). We only aim to show that we get better accuracy with more context available for a given single model, in contrast to comparing token-level accuracies between models with different tokenizers, which is inconclusive. A comment has been added to the caption.\n\n\n\n**Regarding [1, 2]** \nWe first mention that our focus is different. As now clarified, we aim for a long context, while these papers are focused on retrieval from a large knowledge database. We have added a clarification to the related work section. On the technical level, [1] combines two probability distributions to predict the next token: one given by the model logits and the other created from retrieved pairs of (embedding, next token). Meanwhile, we extend the model context in a subset of attention layers, potentially allowing for reasoning within this extended context.\n\nWe thank the Reviewer for raising the topic of the usefulness of other documents in the batch. It was observed that nearest neighbors language models (kNN-LM) display almost linear perplexity gains wrt. datastore size[3]. Due to practical limitations, we embed tokens on the order of magnitude of ~100K per training batch, and documents in the batch are randomly sampled from a large corpus, which means it is unlikely that they are related to each other. Therefore, we should not expect significant perplexity gains for kNN-LM in that setting either, as the training bach comprises approximately 0.1% of the datastore. Empirically, we show that extending the model's context length with attention instead of using kNN-LM leads to perplexity increase due to the aforementioned distraction issue. To the best of our knowledge, the distraction (perplexity increase) resulting from increasing attention context length hasn't been studied before.\n\nWe agree with the Reviewer that TRIME [2] proposes a very similar objective inspired by contrastive learning, which is already mentioned in the related work section. The main difference is architectural: instead of attending to additional tokens in the memory layer, like in FoT, they combine probability distributions of the dense model and the retrieval database in the final layer, like [1]. Moreover, [2] focuses on retrieval from large databases, whereas our experiments mostly focus on long context. We have included this discussion in the related work section of the updated paper.\n\n\n\n**Regarding the distraction issue at the inference time**, giving the model multiple unrelated documents is an extreme case. The distraction issue could possibly occur in single-doc scenarios for long documents consisting of several chapters. Please note that despite alleviating the distraction issue, FoT allows training long-context models using short-context data and improves performance in single-doc cases (see Section 4.5).\n\n\n[1] Khandelwal et al., Generalization through Memorization: Nearest Neighbor Language Models, 2019 \n\n[2] Zhong et al., Training Language Models with Memory Augmentation, 2022\n\n[3] Xu et al., Why do Nearest Neighbor Language Models Work?, 2023\n\nIf our responses have adequately addressed your concerns, we kindly request your support and considerating of improving your score. If you have any further concerns or additional points to raise, we are eager to address them. Your feedback is valuable in enhancing the quality and impact of our research."}, {"review_id": "u8EZU2qxxy", "reviewer": "Reviewer_v42b", "summary": "This work presents a modified method to train Transformers with memory. The memory is based on k-Nearest Neighbors (with exact match). The work is based on the Memorizing Transformer paper from Y. Wu et al. 2022.\nThe model assumes a transformer model with a memory attention (an attention layer that has inputs from both current sequence and memory output). The (k-NN) memory returns the top-k matches for a query, and these are included in the input to the memory attention. One difference is the removal of a gating mechanism that mixes between memory and local context, becoming all input to an attention layer. \nIn addition, contrastive learning is proposed as a method to enhance the model capability to differentiate relevant keys from irrelevant ones in the memory layer. This is achieved by including negative samples from other documents, applying a forward pass up to the memory attention layer, and using these key,value pairs as simulated responses from the k-NN memory. These samples are mixed with the previous context block from the current sample. This method is differentiable and allows to backprop through these positive and negative samples, requiring only a modification to the data pipeline. The authors suggest that the contrastive learning is needed do reduce a \"distracting issue\" caused in a multi-document setup.\n\nEventually, the authors show the value of their method with a set of experiments that compare mainly against the memorizing transformer. The results show that memory can extend the effective context length for language modeling at inference time, they can solve a copy task using memory, a pre-trained model can be fine-tuned to introduce the memory attention layer, improves performance in single and multiple document scenarios, and an ablation study with the importance of the negative samples and the differentiable keys.\n\n", "questions": "I would appreciate it if the authors could answer the following questions:\n\n* Line 188: Why is layer 8 picked as the memory layer? How would someone choose the \"right\" layer for connecting the memory?\n* What is the improvement of an existing language model (like OPT, LLaMA, etc.) when fine-tuned to introduce the memory attention layer? \n* Why is your model not reaching perfect accuracy in the synthetic task? (Shallow) transformers are able to obtain and copy information. Is the same key related to multiple values in different documents in the dataset?\n* What is stored in memory during the evaluation in Section 4.3? What is the data in the memory?\n* What is the distance measure for k-NN? \n* How is the number of neighbors $k$ chosen for the memory retrieval?\n\n", "limitations": "The authors address limitations of their work in Section 5.\nScaling (approximate) k-NNs and memory has been addressed for quite some time. See the NeurIPS 2021 competition on Billion Scale Approximate Nearest Neighbor Search.\n\n ", "rating": 6, "confidence": 4, "soundness": 2, "presentation": 2, "contribution": 2, "strengths": "The paper mixes the idea of negative sampling and a twist to the memorizing transformer model with a k-NN memory. The topic is of significance to this community. \n\n* At inference time this model can increase the context length beyond the training length. Inheriting the original memorizing transformer property.\n* It is possible to fine-tune from a pre-trained model with no memory to a model with external memory.\n* The experiments show the value of the negative samples. The training methodology behind the idea is relatively simple.\n", "weaknesses": "The contribution of this work has limited novelty: both negative sampling (and contrastive learning [2]) and memorizing transformers have been proposed in the past for language modeling.\n\nThe clarity and explanations in the work could be improved. The work continuously omits important details or descriptions, sending the reader in most cases to the appendix. The references could make a better exploration of contrastive learning methods applied to language modeling (e.g., [2]). Please, see questions for additional details that may need clarifications. \n\nRetrieval-augmented language models are able to utilize memory to gain information from multiple documents. The \"distracting issue\" described in the paper suggests that the memory attention mechanism gets distracted by keys from different documents. However, this seems to contradict the findings in previous work (retrieval methods) that earn performance using multiple documents. \n\nThe experiments are mainly limited to the memorizing transformer and the vanilla transformer. However, retrieval language models solve the same multi-document problem. Also the dataset chosen follow the memorizing transformers, failing to compare with existing benchmarks in the long context use case, like SCROLLS [1].\n\n[1] Shaham et al., SCROLLS: Standardized CompaRison Over Long Language Sequences, 2022\n[2] Jain et al., CONTRACLM: Contrastive Learning For Causal Language Model, 2022", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the Reviewer for their thoughtful feedback. We acknowledge some deficiencies in the presentation. We focus on the long-context capabilities; the appropriate clarification is described in the general answer. In more detail, we aim for a single-stage method that can incorporate a large number of tokens directly in the model context (kNN attention can be used to approximate full attention). We observed that increasing context length naively gives worse results which is also confirmed e.g., in [1]. This is not in contradiction with the fact that the retrieval methods benefit from additional documents. The difference lies in the fact that they typically use a two-stage approach, with the retrieval part doing the hard job of extracting only a relatively small amount of tokens, which are efficiently processed within the standard context length [2]. We include this clarification in the paper. We also make a number of smaller adjustments to the paper, which hopefully make the paper easier to follow. If the Reviewer sees any specific issue, we'd be happy to address it.\n\nWe acknowledge some issues with clarity. In the method description, we outline the general idea and admit that it might be hard to infer details from it. We think presenting the details in the text would be quite cumbersome. To amend the situation, we plan to include a shortened version of the pseudocode from the Appendix in the main body. The pseudocode has been found helpful by Rev #fPf8. Thus, we hope it will satisfactorily complement the description. If you see any other parts which require clarification, please let us know.\n\n\nBelow we address the questions:\n1.  As noted in the general response, due to the limited computational resources, we could not perform full hyperparameter sweeps. In particular, for the memory layer, we have followed the choice of Memorizing Transformers (MT) [3].\n2.  Regarding the improvements in existing language models and benchmarking on additional long-context tasks, as noted in the general response, we present 3B and 7B models based on OpenLLaMA along with results on Qasper (SCROLLS benchmark), TREC, and WebQS where we show improved performance when the model is provided with additional context.\n3. Regarding the performance on the synthetic task. Please note that the model is trained in a much shorter context than it is evaluated, which makes it out of distribution.\n4. Regarding the memory content in Section 4.3 during evaluation - this is a single-doc memory; that is, in the additional context, we only store keys and values belonging to the currently processed document. \n5. The distance measure for kNN is inner product.\n6. We have tested values of $k\\in \\{32, 64, 128\\}$ and observed small differences in performance. We add this information to Appendix.\n\nThank you for pointing out [4]; we add the following description to the related work section:\n> CONTRACLM [4] 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 memory attention mechanism to handle longer context lengths. Nonetheless, exploring other contrastive learning objectives could be beneficial for further improving the memory key structure in future work.\n\n\n[1] Nelson F. Liu et al., Lost in the Middle: How Language Models Use Long Contexts, 2023\n\n[2] Sebastian Borgeaud et al., Improving language models by retrieving from trillions of tokens, 2021\n\n[3] Yuhuai Wu, Markus N. Rabe, DeLesley Hutchins, Christian Szegedy. Memorizing Transformers, 2022\n\n[4] Jain et al., CONTRACLM: Contrastive Learning For Causal Language Model, 2022\n\nWe again thank the Reviewer for raising important issues. We hope that our answers are satisfactory. If not, we’d be happy to provide more details. Otherwise, we’d appreciate it if the Reviewer reconsidered the final score of our submission."}, {"review_id": "CKMuUPjFjm", "reviewer": "Reviewer_fPf8", "summary": "This paper proposes the Focused Transformer (FoT), which modifies one layer of the Transformer model's Attention layer to Memory Attention, thus enabling the model to learn almost infinite length context without being limited by the constraints of local attention. Specifically, the paper also proposes a cross-batch training method that can adapt the model to any length context at a relatively low fine-tuning cost. Experimental results demonstrate the superiority of FoT on Long Context tasks.", "questions": "1. I don't quite understand why, during the cross-batch training stage, we need to distinguish between C_prev and C_curr? Isn't it better to treat all tokens across batches as negative learning directly?\n\n2. In the inference stage, how are the K and V to be input into Memory Attention calculated after kNN?\n\n3. After cross-batch training, will the model overfit long Context? Is its performance on short Context still good?", "limitations": "The authors have clearly and sufficiently explained the limitations of their research", "rating": 8, "confidence": 5, "soundness": 4, "presentation": 3, "contribution": 4, "strengths": "1. By only retaining one layer of Global Attention (Memory Attention), both training and inference are efficient;\n\n2. During the inference stage, it can flexibly follow the results of kNN retrieval;\n\n3. The cross-batch training method is very effective, and the training cost is also relatively low;\n\n4. The increase factor of the Context length for LLM reached 2^15!", "weaknesses": "1. With all due respect, I believe that the main text of the paper does not clearly introduce FoT and cross-batch training. I didn't understand it until I read the code in the appendix;\n\n2. The inference stage needs to be combined with kNN, which is inconsistent with the training mode and may affect the upper limit of the model's ability;\n\n3. There is a lack of comparison with other similar schemes for increasing Context length (such as Parallel Context Window).", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for the encouraging review.\n\nThe description of the method outlines the general idea, and we admit that it might be hard to infer details from it. We think presenting the details in the text would be quite cumbersome; thus, we plan to include the shortened versions of the code in the main body of the paper. We hope this will be satisfactory.\n\n **Regarding the kNN**, we consider kNN to be an approximation of the full dense attention. With such a perspective, there is no inconsistency. However, in practice, the approximation errors may impact the performance. We did not observe this in our experimental regime. We leave proper studies of this to future work. We also note that our fine-tuned versions of OpenLLaMA models use full dense attention instead of kNN, which we find performant and efficient enough. We also note that using kNN opens the possibility of using fast approximate indices (e.g., implemented in Faiss), which might be necessary for scaling the method. We have added this to future work.\n\n **Regarding the comparison with Parallel Context Window**, we add the following text to related work.\n> Parallel Context Window 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. Our method also allows us to improve the structure of $(key, value)$ space of existing models.\n\nRegarding the questions:\n1. We observed that it is important to have at least one positive example that brings additional related information to memory layers (for example, previous local context window $C_{prev}$). Otherwise, the model may learn to ignore memory layers.\n2. For each query in a memory layer, we take $k$ most matching keys from memory and add them to the attention for this query. That is, each query will attend only to all keys that precede it in the local context and its own $k$ most matching keys from memory. In the non-kNN approach, each query attends to the whole memory and all keys that precede it in the local context. To calculate $k$ most matching keys, we use the inner product. Note that in models presented in the paper, we remove positional encodings in memory layers. \n3. We have managed to fine-tune OpenLLaMA models so that they maintain the performance of the base models on short-context Language Model Evaluation Harness tasks and show improvements on long-context ones. For details, please refer to the table from the general response. \n    \nWe again thank you for an encouraging review. Should you have any further questions or concerns, we'd be happy to answer. We kindly ask to support our paper."}, {"review_id": "jZ8PqakQDG", "reviewer": "Reviewer_GWpw", "summary": "This paper proposes an improved training approach for memory-augmented Transformers on language modeling tasks. In particular, the authors identify the distraction issue in memory-augmented models, whereby the attention mechanism tends to focus on irrelevant contexts in the regime of long sequences. To address this concern, the authors propose cross-batch training, which is inspired by contrastive learning and exposes the attention mechanism to process both relevant and irrelevant documents. Extensive experiments are conducted to validate that the proposed cross-batch training can effectively shape the key-value space, mitigate the distraction issue, and successfully extend the model context size.", "questions": "1. How does the proposed technique compare to conventional packing techniques commonly used during pretraining (e.g., concatenating multiple documents together, possibly interleaved with [EOS] tokens)? This technique is mostly used to reduce padding but also generates a sequence with context from diverse documents.\n2. How well does the proposed model converge in comparison to vanilla transformers? Intuitively, as the attention mechanism in FoT is exposed to more irrelevant information, the learning process might slow down.\n3. How can one determine the optimal transformer layer for memory augmentation? Does allowing more transformer layers access to the memory aid in the retrieval of relevant information?", "limitations": "The authors have adequately addressed the limitations.", "rating": 8, "confidence": 4, "soundness": 4, "presentation": 4, "contribution": 3, "strengths": "- The paper is clearly written and well structured.\n- The identified distraction issue is compelling and makes a valuable contribution to the long-context research community, highlighting the need for more effective attention mechanisms or retrieving techniques.\n- The method is straightforward to use and demonstrates its significant effectiveness in identifying relevant information.", "weaknesses": "My main concern is about the increased training computational costs. The computational complexity for the memory-augmented attention layer increases from $O(bn^2)$ to $O(b d n^2)$, where $b$ represents the batch size, $d$ is the number of cross documents, and $n$ is the sequence length. Employing a large $d$ might introduce a significant computation overhead.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for a thoughtful review.\n\n**Regarding the computational cost**. We thank you for raising this important concern, which we have added to the limitation section. We also note that two factors mitigate this issue to some extent. As you noted, the increased cost occurs only in the memory layer. Second, FoT exhibits some context extrapolation, which might allow using smaller $n$ in training. Accordingly, we managed to fine-tune OpenLLaMA models using $d \\leq 8$\n    \nRegarding the questions:\n\n**Q1:**  We thank the Reviewer for proposing an interesting baseline. To answer the question, we fine-tune a vanilla OpenLLaMA model on sequences of length $4096$ (original seq. len $2048$), which we consider as just a standard “data packing\" baseline, and compare it to a FoT model trained on exactly the same data packed in the same way for 1B tokens. For clarity, we outline the following architectural differences between the baseline and FoT:\n* Additional context beyond $2048$ tokens is used in just a subset of layers\n* FoT does not use positional encodings in memory layers beyond its original context window ($2048$)\n\nThe results are as follows:\n    \n|Context/Setup | TREC: baseline $\\pm 1.0$  $~~$ | TREC: FoT $\\pm 1.0$  $~~$  | WebQS: baseline $\\pm 0.1$  $~~$  | WebQS: FoT $\\pm 0.1$  $~~$  |\n| - | - | - | - | - |\n| 2K | 52.8 | 55.6 | 20.7 | 20.8 |\n| 4K | 57.2 | 60.9 | 18.7 | 21.0 |\n\nWe observe accuracy improvements when more few-shot demonstrations are provided in the extended context (from 2K used by OpenLLaMA to 4K used in our fine-tuning). On TREC, the gains from additional context are significant for both models. Our method presents better data efficiency than the baseline. \n    \n**Q2:**  Starting with a large $d$ (crossbatch dimension) may slow down the process and result in the memory layer being ignored by the model. We have not seen any such problems when starting with a smaller value of $d\\leq8$. See the plot with the training loss comparison in the attached pdf.\n\n**Q3:** Due to the limited resources, we have followed the choice of Memorizing Transofrmer (MT) in picking the memory layer. We have also seen some additional gains from using multiple memory layers in our FoT fine-tuned OpenLLaMA models. \n\n\nIf your concerns have been sufficiently addressed in our responses, we humbly seek your support for the paper. Should you have any further concerns or additional points to raise, we are eager to address them."}]

Author rebuttal:
We would like to thank you for all your valuable feedback, both positive and negative, which we believe will help us to improve the quality of our work. 

We are delighted to note that the reviewers (Sux9, v42b, GWpw) prized the simplicity of our method and noted the potential impact (Sux9) of extending the context length (v42b, GWpw). Reviewer fPf8 noted the efficiency of FoT, and Sux9 prized the synthetic dictionary task. 

Reviewers Sux9 and v42b raised important concerns about the scope and contributions of the paper, which are addressed below. Moreover, we would like to advertise new important experiments.

### New experiments with large models
In the period between the submission and the rebuttal, we secured additional compute resources, which let us confirm that our method is useful for much larger models. We believe that this significantly strengthens the paper. Specifically, we fine-tuned $3B$ and $7B$ OpenLLama models with our FoT objective. The resulting models exhibit advancements in tasks requiring long context. Following that, we extend the contribution list accordingly.

Below we shortly summarise the properties of our models. We would be happy to provide more details here if needed. Otherwise, we present them in an additional section in the paper. Specifically, our new models: 
1. exhibit long-context capabilities on downstream tasks (see tables below), 
2. retain the performance of the original models on short-context tasks,
3. are more compute and memory efficient at inference, compared to vanilla Transformers with the same effective context length,
4. have some context extrapolation capabilities. We illustrate that our models manage a 256k context length for passkey retrieval task from [1]  even though being trained on 8k context. **(see pdf)**.

Ad 1. Our model exhibit performance gains from additional in-context few shot examples on TREC question classification [2, 3] and WebQS question answering [4]. What is more, it shows improvements in F1 score on Qasper (Question Answering over Scientific Research Papers) task [5], which is a part of SCROLLS [6].







| Context/Setup | TREC: FoT fine-tuned OpenLLaMA 3B | TREC: FoT fine-tuned OpenLLaMA 7B | WebQS: FoT fine-tuned OpenLLaMA 3B | WebQS: FoT fine-tuned OpenLLaMA 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                  |

For Qasper, we used the implementation from Language Model Evaluation Harness and observed that our model 3B model benefits from context increase. Below we provide zero-shot results. Note that LongChat 7B [7] was instruction fine-tuned.

|Context length | OpenLLaMA 3B | FoT fine-tuned OpenLLaMA 3B | LLaMA 7B | LongChat 7B |
| - | - | - | - | - |
| 2K | 18.7 | 18.7 | 18.7 | 19.4 |
| 4K | - | 20.7 | - | 21.2 |
| 6K | - | 23.2 | - | 25.0 |
| 8K | - | 26.6 | - | 28.8 |


Ad 2. Our fine-tuned OpenLLaMA models maintain the performance on the standard suite of short-context  tasks from Language Model Evaluation Harness (we use the same collection of tasks as OpenLLaMA and provide the average scores)

|Model | OpenLLaMA 3B | FoT fine-tuned OpenLLaMA 3B | OpenLLaMA 7B | FoT finetuned OpenLLaMA 7B|
| - | -| - | - | - |
|Average score | 0.53|	0.53 |	0.55 |	0.55|


### Scope and contributions of the paper

To clarify, *our paper focuses on the long-context capabilities*. We agree that the current writing is somewhat unclear. We have identified the following issues which might have caused the confusion:

- We now stress that handling large external databases was the initial motivation of FoT, which was later changed to long-context.
- We used the term 'external memory', which we now change to 'additional context'.
- Memorizing Transformer, on which we base our method, is framed as a retrieval method. We now explicitly state in the related work section that despite these similarities, our aim is different. Moreover, we amend the related work to include more long-context papers. 
- We include new long context tasks (see above). We keep the multi-doc experiments for illustrative purposes. However, we make explicit that the focus is on the long context. 

We thank the reviewers for pinpointing this clarity issue. We hope that the above changes will address the concerns. We would be happy to make further adjustments if the reviewers find it useful.


### Tuning and hyperparamters 

Reviewers (V42b, GWpw) raised questions about hyperparameters (e.g. the memory layers used). We note that some of the choices were educated guesses, as due to extreme computational cost, we could not perform a full hyperparameter search. For example, this was the case for the memory layer we based on the findings from Memorizing Transformer. This information is now added as a limitation.


[1] A. Mohtashami, et al. Landmark Attention: Random-Access Infinite Context Length for Transformers. 

[2] Li, Xin et al. Learning Question Classifiers.

[3] E. Hovy, et al. Toward semantics-based answer pinpointing.

[4] J. Berant, et al. Semantic Parsing on Freebase from Question-Answer Pairs.

[5] P. Dasigi, et al. A Dataset of Information-Seeking Questions and Answers Anchored in Research Papers.

[6] U. Shaham, et al. SCROLLS: Standardized CompaRison Over Long Language Sequences.

[7] D. Li*, et al. How Long Can Open-Source LLMs Truly Promise on Context Length?

Meta-review: The paper proposes the Focused Transformer (FoT) to addresses the limitations of local attention in Transformers by incorporating a Memory Attention layer, enabling it to learn from extensive context while mitigating distractions. This architecture, reminiscent of the Memorizing Transformer, utilizes a k-Nearest Neighbors (k-NN) memory retrieval mechanism to distinguish between relevant and irrelevant information. FoT's unique crossbatch training approach trains the model with context from both the same and different documents, enhancing its ability to discern valuable insights. Experimental results underscore FoT's success in handling long-context tasks and improving performance across various scenarios, positioning it as a promising solution for extending context and managing distractions through innovative attention mechanisms and training techniques.

All reviewers appreciate this work and the discussion between authors and reviewers has addressed most of reviewers' concerns. Therefore, I am happy to see this paper accepted at the conference.
Decision: Accept (poster)