SlowGuess/NeurIPS_2023_OCR huggingface dataset:
In this section, we 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 it 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  $1B$  tokens using the  $4K$  context length. We evaluate both models on a number of few-shot downstream tasks in the setting described in Section 4.4.
In most cases, see Table 2, we 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 1).
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 is not the case for the LONGLLAMA models by evaluating them using the LM Evaluation Harness library [Gao et al., 2021a]. On most tasks, the performance is kept intact; see Appendix A.4 for details. This also confirms that LONGLLAMAs could be used as a drop-in replacement of LLaMA models as they are compatible with the original LLaMA inference code.
5 Analysis of FoT
In this section, we perform extensive experiments on 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
an existing, pre-trained model? (3) How effectively can it handle distractions, and how does this capability translate to enhanced performance in long-context language modeling tasks? Moreover, we provide ablation studies of our method and additional analysis.
Architecture 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]; we pick  $\ell = 8$  as the memory attention layer. We tune  $k = 128$ , the number of top keys retrieved by kNN. In most experiments, we start training with a small crossbatch dimension  $d \leq 8$  and switch to  $d \geq 64$  after some training. For more details about the architecture and hyperparameters, see Appendix B and Appendix E.
Evaluation We distinguish two evaluation settings: single-document (abbreviated to single-doc) 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), arXiv (mathematical papers), GitHub (code), and Isabelle (formal proofs). PG-19 [Rae et al., 2019] is a large dataset of English-language books published prior to 1919, sourced from the Project Gutenberg archive. This dataset is a well-established benchmark for evaluating long-context language models [Sun et al., 2021]. The arXiv dataset contains LATEX 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 a book in PG19. For details on the remaining datasets, refer to Appendix H.
5.2 FOT fine-tuning and context length extrapolation
FoT is a 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, as we already demonstrated in Section 4. In this section, we deepen this analysis (on a smaller model) by studying perplexity improvements on various datasets.
As a 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-tuning and pre-training is the C4 dataset Raffel et al. [2019a] (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 a 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 3, 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 et al., 2022] fine-tuned with the local context of  $1K$  and memory size of  $16K$ , and Transformer-XL [Dai et al., 2019] fine-tuned with both local context and window length of  $1K$ . To ensure a 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-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 on 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 a synthetic dictionary lookup task. In this task, the model is first provided with  $k_{i} : v_{i}$  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
Table 3: Perplexity for different context lengths after fine-tuning a standard transformer model. The model is fine-tuned using the FoT objective (i.e., crossbatch) on C4 and evaluated zero-shot varying the context size. Transformer-XL [Dai et al., 2019] and Memorizing Transformer [Wu et al., 2022] fine-tuned in the same setting are used as baselines.
Method	Context Length	GitHub	Isabelle	arXiv	PG-19
FoT	2K	6.72	5.63	8.17	23.74
4K	5.88	4.93	7.44	23.25
16K	5.43	4.51	6.94	22.85
64K	5.32	4.44	6.81	22.65
Transformer-XL	2K	6.85	5.76	8.21	23.57
Memorizing Transformer	2K	8.10	7.34	9.39	24.03
4K	7.55	6.93	8.95	23.62
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 in the multi-document setting helps in language modeling. We pick the PG-19 dataset [Rae et al., 2019] 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 to 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 of 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 in 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 0.18 when scaling to  $>500k$  tokens. Importantly, the perplexity of FoT is close to this of Memorizing Transformer with the single-doc memory, which we treat as a soft lower bound since it is not exposed to distractions from unrelated books.
5.4 Context length extrapolation in single-doc
The original motivation behind FoT is to 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  $w = 1$  (the standard setting for FoT, that corresponds to the total differentiable context being  $2 \cdot 1024$ ) and  $w = 2$  (corresponding to the total differentiable context  $3 \cdot 1024$ ).
We observe, see Figure 9, improvements when context grows, even far beyond the training context length, which reaffirms the hypothesis that FoT helps with extrapolation to longer contexts. Moreover,  $d = 2$  is significantly better than  $d = 1$ . When comparing  $d = 1$  and  $w = 2$  to  $d = 2$  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 on 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 in the choice of memory layers.
Our research opens a 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 our experiments, we use the exact kNN search, which is not scalable to large memory. Using approximate kNN search will require a lot of engineering effort, as well as careful evaluation of the impact of the approximation on 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 a single TPUv3/TPUv2 machine, see also Appendix I. In future work, we want to further increase  $d$  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 a subset of 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 is mitigated. We expect that other CL methods could be beneficial, for example, hard negative mining to utilize a larger memory during training (see [Lindgren et al., 2021]). We leave this for future work.
Combining with other methods Developing 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 L.
# keys from other contexts will be encoded as if they
# were at the beginning of the local context
pkey_fst = pos_encode_as_first(xk=key)
# local context keys encoded in the standard way
pqquery, pkey = pos_encode(xq=query, xk=key)
# for each element of the batch we calculate indices of
# the batch that will be used in cross-batch
cross_batch_REL_ids = jnp.arange(0, -num attentions, -1)
        .reshape(1, -1)
batch_ids = jnp.arange(0, batch_size).reshape(-1, 1)
cross_batch_selector = cross_batch.rel_ids + batch_ids
# here we want other contexts
cross_batch_keys = pkey_fst[cross_batch_selector[:, 1:]]
# here we concatenate local context with other contexts
attention_keys =
    jnp concatenate([pkey[:, None], cross_batch_keys], axis=1)
cb_attnweights =
    jnp.einsum("bqhdt,bckhd->bhqck",
        pqquery, attention_keys, precision=precision)
We gratefully acknowledge the TPU Research Cloud program, which was instrumental to our research by providing significant computational resources. Parts of the project were realized using the resources of Poznański Centre 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 2019/35/O/ST6/03464. Henryk Michalewski was supported by the Polish National Science Center grant UMO-2018/29/B/ST6/02959.
Yuhuai Wu, Markus Norman Rabe, DeLesley Hutchins, and Christian Szegedy. Memorizing transformers. In The Tenth International Conference on Learning Representations, ICLR 2022, Virtual Event, April 25-29, 2022. OpenReview.net, 2022. URL https://openreview.net/forum?id=TrjbxzRcnf-.
Manzil Zaheer, Guru Guruganesh, Kumar Avinava Dubey, Joshua Ainslie, Chris Alberti, Santiago Ontonon, 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-5673. Association for Computational Linguistics, 2022. URL https://aclanthology.org/2022.emnlp-main.382.
Recent rapid developments in language models have brought a lot of new capabilities. At the same, these raised concerns about the social impact and very animated discussions in the community. Our work develops a generic 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. [2022].
A LONGLLAMA
OpenLLaMA [Geng and Liu, 2023] is an open-source reproduction of LLaMA [Touvron et al., 2023]. It 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., 2017]. A SentencePiece [Kudo and Richardson, 2018] tokenizer with 32k vocabulary size is used.
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  $\mathcal{L} = \{6,12,18\}$  as the memory layers. We vary the number of additional contexts  $d\in \{0,2,3\}$  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 a 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  $\frac{1}{4}(0,0), \frac{1}{4}(1,1), \frac{1}{4}(2,1), \frac{1}{4}(3,0)$ .
For the 7B LONGLLAMA model, we set  $\mathcal{L} = \{8,16,24\}$  as the memory layers. Here we divide batch entries into four segments and use the following setup:  $\frac{1}{4}(0,0), \frac{1}{4}(1,2), \frac{1}{4}(2,5), \frac{1}{4}(3,4)$ .
Hyperparameters We follow the choices of OpenLLaMA with respect to most of the hyperparameters, including using the same optimizer. During fine-tuning, we use a batch size of  $256K$  tokens and constant learning rate of 2e-5, which is lower than the learning rate at the end of OpenLLaMA training (3e-5 after 1T tokens), and weight decay of 0.01.
A.3 LLaMA fine-tuning dataset
We use a mixture based on RedPajama [TogetherComputer, 2023] and The Stack [Kocetkov 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 are in the same subdirectory in random order, using a similar procedure as for the GitHub dataset in Section 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.
Table 4: 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 at 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	-
python	25	4096
book	10	-
common_crawl	29	-
c4	5	1024
github	2	2048
stackexchange	2	1024
wikipedia	2	1024
A.4 Language Model Evaluation Harness
To ensure that the performance of LONGLLAMAs has not degraded in short context scenarios, we evaluate our models on the Language Model Evaluation Harness benchmark [Gao et al., 2021a]. Table 5 compares our 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 on 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.
Task/Metric	OpenLLaMA 3B	LongLLaMA 3B	OpenLLaMA 7B	LongLLaMA 7B
anli_r1/acc	0.33	0.32	0.33	0.35
anli_r2/acc	0.32	0.33	0.36	0.37
anli_r3/acc	0.35	0.35	0.38	0.36
arc_challenge/acc	0.34	0.34	0.37	0.37
arc_challenge/acc_norm	0.37	0.37	0.38	0.38
arc_easy/acc	0.69	0.68	0.72	0.70
arc_easy/acc_norm	0.65	0.63	0.68	0.66
boolq/acc	0.68	0.68	0.71	0.71
hellaswag/acc	0.49	0.48	0.53	0.52
hellaswag/acc_norm	0.67	0.65	0.72	0.71
openbookqa/acc	0.27	0.28	0.30	0.30
openbookqa/acc_norm	0.40	0.38	0.40	0.41
piqa/acc	0.75	0.73	0.76	0.75
piqa/acc_norm	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_mc/mc1	0.22	0.24	0.23	0.24
truthfulqa_mc/mc2	0.35	0.38	0.35	0.35
wic/acc	0.48	0.50	0.51	0.50
winogrande/acc	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.
Table 6: Zero-shot performance with different context lengths on the validation subset of Qasper [Dasigi et al., 2021]. We use the implementation from Language Model Evaluation Harness [Gao et al., 2021a]. In the Harness implementation, yes/no questions are evaluated separately from open questions. Observe that LONGLLAMA 3B benefits from the extended context.
Context length	OpenLLaMA 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-LLaMA-based models presented in this paper. The main differences are that for LLaMA-based models (LONGLLAMA) we maintain the positional encodings (with a 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-LLaMA-based models we use the transformer architecture introduced in [Vaswani et al., 2017] with a 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., 2021], normalize keys and queries [Henry et al., 2020], and introduce a learnable temperature parameter for each attention head.
The hyperparameters for each model size can be found in Appendix E. For training the models on PG-19, we use the standard T5 tokenizer with 32k vocabulary [Raffel et al., 2019b]. The larger models in Section 5.2 are trained with a custom SentencePiece tokenizer [Kudo and Richardson, 2018] with 64k vocabulary size.
B.2 Memory attention layer
Memory attention layer  $\ell$  is one of the transformer layers, which has access to the additional context  $M$ . The memory stores  $(key, value)$  pairs. For each query  $q$  in  $\ell$ , we retrieve the  $k$  most matching entries from  $M$  and use them to compute the attention value. More precisely, we use the kNN algorithm to pick  $M_{top} := \{(key_1, value_1), \ldots, (key_k, value_k)\} \subset M$  such that  $\{\langle q, key_i \rangle\}_{i=1,\ldots,k}$  are the top  $k$  inner products in  $M$ . These are merged with the part of the local context before  $q$  denoted as  $C_{<q}$  and used to compute the attention value using the standard Transformer formula:
$$
v := \sum_ {(k e y, v) \in M _ {t o p} \cup C _ {<   q}} s (k e y) \cdot v, \tag {1}
$$
where  $s(key)$  is the softmax score for key. This softmax is calculated as follows:
$$
\operatorname {s o f t m a x} \left(\left[ \frac {\langle q , k e y \rangle}{\tau} \right] _ {k e y \in M _ {t o p} \cup C _ {<   q}}\right),
$$
where  $\tau$  is a temperature parameter. In this approach, we do not distinguish between the local context and the memory.
Another way of integrating  $M_{top}$  is via gating. In this approach, we separately compute the attention value  $v_{M}$  for  $M_{top}$  and for the local context  $v_{C}$  (using the standard Transformer formula). Then we use a gating mechanism to combine them:
$$
v := v _ {M} \cdot g + v _ {C} \cdot (1 - g), \quad g = \sigma (b _ {g}),
$$
where  $\sigma$  is the sigmoid function and  $b_{g}$  is a trainable bias. The gating approach was proposed in [Wu et al., 2022], see formula [Wu et al., 2022, (2)].