SlowGuess/NeurIPS_2023_OCR huggingface dataset:
We found our approach, i.e. using (1), to be equally effective, see Figure 4. At the same time, (1) 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 in FAISS [Johnson et al., 2017]. The memory attention layer does not use positional encodings. The memory is populated incrementally with (key, value) pairs processed by  $\ell$  beforehand. In the single-doc setting, the memory is erased after each document.
We do not use the  $\tau$  parameter for LONGLLAMAs as their architecture does not normalize keys and queries. For LONGLLAMAs, we also replace the kNN search with dense attention and retain positional encodings (see Appendix A.2).
B.3 Crossbatch training procedure
In FOT we choose a subset  $\mathcal{L}$  of the attention layers for later augmentation with the memory of (key, value) pairs. Let  $\ell$  an attention layer from  $\mathcal{L}$ . During the training we expose this layer to a mixture of (key, value) pairs from the current local context,  $C_{\mathrm{curr}}$ , and the previous local context and  $d - 1$  contexts from other documents,  $C_{\mathrm{prev}}$ ; see also Figure 2 for an illustration. We achieve this by modifying the input pipeline so that each batch index corresponds to a different document (the batch index occupied by each document is fixed from the moment we load the document till we finish processing it).
More specifically, we embed the previous and the current local context for each document in the batch. Then we use  $C_{\mathrm{prev}}$  as a 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  $\ell$ , as we envision it to handle global information.
To be more precise, for each document  $\delta$  within the batch and query  $q$  from the layer  $\ell$  we create the set  $p^{\delta}$  consisting of (key, value) pairs from the previous local context of document  $\delta$  along with pairs from  $d - 1$  contexts gathered from  $C_{\mathrm{prev}}$ . The attention value for  $q$  is given by
$$
v := \sum_ {(k e y, v) \in p ^ {\delta} \cup C _ {\text {c u r r}} ^ {\delta , <   q}} s (k e y) \cdot v, \tag {2}
$$
where  $C_{\mathrm{curr}}^{\delta, < q}$  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  $\tau$ :
$$
\operatorname {s o f t m a x} \left(\left[ \frac {\langle q , k e y \rangle}{\tau} \right] _ {k e y \in p ^ {\delta} \cup C _ {\mathrm {c u r r}} ^ {\delta , <   q}}\right).
$$
Note that the only difference between (1) and (2) is the source of the additional (key, value) pairs:  $p^{\delta}$ . This, in particular, implies that all the operations with respect to the previous context are differentiable.
The number of different documents is equal to  $b_{S}$  (the batch size, i.e. each document has a separate index in the batch). Assume that document  $\delta$  has index  $i$ . We include into  $p^{\delta}$  all tokens from  $C_{prev}$  with the batch indices in  $\{i, (i + 1) \mod b_s, \ldots, (i + d - 1) \mod b_s\}$ .
Figure 4: Perplexity (on the test set) during training on PG-19. We train two Memorizing Transformer models [Wu et al., 2022], one with original gating and one without (i.e. memory attention shared with local attention as described in (1)). We use the single-doc setting with 16k memory.
B.4 Qualitative analysis
Table 7 provides a brief qualitative analysis of FoT. It shows that the model can handle distractions and retrieve the parts of the character name from the multi-document memory in 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 et al., 2022]. In this section, we summarize the differences between those two models.
Training The key difference between these two methods lies in the training procedure. Our method uses crossbatch, see Section B.3, which, in a nutshell, is the standard transformer training objective, but we additionally attend to tokens from the previous context window, both from the same and different documents, see Appendix B.3 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 is 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 not 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 at providing 'hard' negatives for the model. We provide a proof-of-concept experiment in Appendix C.3, and leave this for future work.
Inference Both models use a very similar memory attention layer. The difference is how the retrieved (key, value) pairs are integrated. FoT treats the retrieved information in the same way as the local context. MT uses a gating mechanism. Details are provided in Section B.2.
Table 7: Example of elements retrieved by kNN search along with their scores on PG-19 and dictionary lookup task. The first column shows the token associated with a 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 score shows what fraction of the attention mass dedicated to retrieved keys corresponds to a particular key. The Focus score is calculated by taking the key along with 32 preceding and 32 following keys, calculating attention weights for 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 4096 spanning across parts of 8 documents. In the dictionary lookup setting the model was provided with memory of size  $16M$ .
Text	kNN Results	kNN Score	Focus Score
PG-19
S HE BA takes the gold-rimmed pince-nez which hangs upon THE DEAN'S waistcoat and places it before his eyes	Then if we're here with the closed carriage at ten-! [They go together into the library. DARBEY. [To SHE BA .]	0.69	0.99
Oh! [He sinks on to the settee with a vacant stare, his arms hanging helplessly. DAR- BEY. [To SHE BA .] There-now his career is a burden to him!	0.23	0.92
Papa! SHE BA . Papsey! [THE DEAN rouses himself, discovers his children and removes his hat.	0.025	0.92
THE DEAN gives DAR BE Y a severe look, and with an important cough walks into the Library. The men and the girls speak in undertones.	Then if we're here with the closed carriage at ten-! [They go together into the library. DARBE Y. [To SHEBA .]	0.71	0.99
Oh! [He sinks on to the settee with a vacant stare, his arms hanging helplessly. DARBE Y. [To SHEBA .] There-now his career is a burden to him!	0.19	0.99
Oh, Salome! Papa! Papa! TARVER. The Dean? DARBE Y. The Dean!	0.08	0.99
Dictionary Lookup Task
<q> 14 42 23 38 <v> 40 41 05 56	<k> 14 42 23 38 <v> 40 41 05 56	0.96	0.99
<q> 30 55 07 23 <v> 36 17 26 63	<k> 30 55 07 23 <v> 36 17 26 63	0.88	1.0
<q> 10 41 26 39 <v> 48 38 11 24	<k> 10 41 26 39 <v> 48 38 11 24	0.87	0.99
C Ablations
In this section, we focus on two key properties of crossbatch training procedure: differentiability and the inclusion of 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
We 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 5, both MT and FOT are trained with local context of 512. We observe that FOT is significantly better when the context is expanded during inference, which confirms that differentiable keys and values are beneficial.
Table 8: Perplexity on PG-19 in the single-doc setting for various local context lengths during training. In 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
Figure 5: Perplexity on PG-19 in the multi-doc setting. 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 8 documents. Differentiability of keys and values results in better perplexity.
Figure 6: Importance of negatives in the multi-document setting. We compare FoT trained with  $d = 1$  to the one that started with  $d = 2$  and later switched to  $d = 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, we compare FoT with  $d = 1$  to MT with memory consisting of the previous local context. Table 8 confirms that differentiable keys and values can also help in this scenario.
C.2 Importance of negatives
We reaffirm the importance of negatives in a multi-document setting. In previous experiments in Figure 3, we already observed that increasing the number of negatives (i.e., increasing  $d$ ) results in more attention mass being dedicated to relevant tokens. In Figure 6, we additionally show that the lack of negatives in training ( $d = 1$ ) results in a 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] is closely related to our method. The two key differences are 1) the training protocol and 2) how the memory is integrated into the model. In this section, we 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-inspired objective and backpropagating through the previous context. A 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 proof-of-concept experiment combining the two approaches to explore this further. Namely, we trained the model for 499k steps using crossbatch and fine-tuned it with the MT objective for 1k 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 a simple memory integration approach where the (key, value) pairs retrieved by kNN lookup are treated the same way as the local context. In contrast, MT uses a gating mechanism, a 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, we decided to use our approach because it does not require
Figure 7: Single-doc eval of FoT finetuned for 1k steps on non differentiable memory. It achieves lower perplexity than MT, which has access to 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 of the fact that the crossbatch training backpropagates through the (key, value) pairs from the previous context  $C_{prev}$  in contrast to MT that cannot backpropagate there and needs to rely on local context when computing gradients for keys and values. Another reason might be the fact that  $C_{prev}$  is embedded for each batch, and thus staleness (see [Wu et al., 2022, Section 3.2]) is avoided.
D Additional figures
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  $d = 64$ . Single-doc MT was trained with a memory size of  $16K$ . As we increase the memory size, the number of distractions (irrelevant keys) increases, making the task harder. Single-doc MT evaluated in the single-doc setting is a soft lower bound since it lacks distractions.
Figure 9: Zero-shot performance on PG19 of FoT pretrained on C4. Model fine-tuned with the crossbatch dimension  $d = 2$  outperforms the one with  $d = 1$ . Using the double  $(w = 2)$  training context of 2048 is beneficial.
E Hyperparameters
Table 9 shows hyperparameters used in our experiments. We used context length 512 unless stated otherwise. In Appendix F, Section 5.3, Section 5.5, we use the total batch size of  $32K$  tokens. In Section 5.2 and Section 5.4, the total batch size is  $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 rate schedule	Inverse Square Root
Warmup steps	1000
β1	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	8
FeedForward dim	2048	4096
Local context length	256/512	512
Batch size	32K/64K	32K
#Number of training steps	5k	500k
For the experiments described in Section 5.3 and Section 5.5 we performed the following hyperparameter sweeps:
- Learning rate:  $\{1\mathrm{e} - 2,5\mathrm{e} - 3,3\mathrm{e} - 3,1\mathrm{e} - 3\}$ , chosen:  $1\mathrm{e} - 2$
- Batch size:  $\{8K, 16K, 32K\}$ , chosen:  $32K$ .
For the dictionary lookup task (Appendix F) we checked the following hyperparameters:
- Learning rate:  $\{4\mathrm{e} - 2,2\mathrm{e} - 2,1\mathrm{e} - 2,5\mathrm{e} - 3,3\mathrm{e} - 3\}$ , chosen:  $2\mathrm{e} - 2$
- Number of dictionary tokens in training step:  $\{16K, 32K\}$ , chosen:  $32K$ . Note that during the training number of document tokens dedicated to the dictionary is the same as 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$
In Sections 5.3, 5.4 and 5.5 for models with  $d \in \{1,2,4,8\}$  we used constant schedule, and for models with  $d = 64$  we trained with  $d = 2$  for  $450k$  steps and switched to  $d = 64$  for the final  $50k$
steps. In Appendix F we trained with  $d = 1$  until the model reached  $98\%$  accuracy and then we switched to  $d = 128$ . For the  $184M$  model in Section 5.2, we randomly sampled  $d$  from  $\{2, 128\}$  in each training step.
F Dictionary lookup task
Figure 10: Accuracy vs number of dictionary tokens in a dictionary look-up task. The task format is as follows:  $\langle \mathbf{k} \rangle k_1 < \mathbf{v} \rangle v_1 < \mathbf{k} \rangle k_2 < \mathbf{v} \rangle v_2 \dots < \mathbf{k} \rangle k_n < \mathbf{v} \rangle v_n < \mathbf{q} \rangle k_i < \mathbf{v} \rangle v_i \dots$ , where a dictionary is provided, followed by queries on randomly selected keys. Accuracy is determined by measuring the predicted values  $v_i$  after  $\langle \mathbf{q} \rangle$  tokens. Models were trained on examples containing 512 tokens and evaluated with an extended context length. FoT demonstrates high accuracy even when the memory size is large. The baseline transformer fails already for  $16K$  tokens. Error bars represent the minimum and maximum on 10 seeds.
We propose a dictionary lookup task to test whether the model trained using our crossbatch method can utilize a large memory database. Documents in this task are split into two parts. The first part defines keys and their associated values using the records of the format:
$$
<   \mathrm {k} >, k _ {1}, k _ {2}, k _ {3}, k _ {4}, <   \mathrm {v} >, v _ {1}, v _ {2}, v _ {3}, v _ {4},
$$
where  $\langle \mathbf{k} \rangle$  is a special token that denotes the beginning of the defining sequence,
The second part consists of queries about the values associated with the previously defined keys. The queries are in the following format:
$$
<   \mathrm {q} >, k _ {1}, k _ {2}, k _ {3}, k _ {4}, <   \mathrm {v} >, v _ {1}, v _ {2}, v _ {3}, v _ {4},
$$
where  $\langle q \rangle$  is a special token that denotes the beginning of the query. We mask the loss so that for such a question, only  $v_{1}, v_{2}, v_{3}, v_{4}$  are included. We use a vocabulary of 64 tokens, keys and values are described using 4 tokens.
During training, we use 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 use the memory attention layer to answer the questions correctly. We start with  $d = 1$  and increase to  $d = 128$  as soon as the model is able to reach  $98\%$  training accuracy. During the inference, we use  $k = 32$  (the number of keys retrieved by kNN). As a baseline, we use a standard transformer model trained with the context length of 512. In evaluation, we test different local context lengths, which quickly leads to very poor results.
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 on Figure 10), 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 a 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  $d$  from  $\{2,128\}$  in each training step to prevent the model from overfitting to a large additional context length during training.
H Datasets
Section 5.1 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 in a random order, we concatenated all the files within each repository, treating subdirectories as a 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 a corpus of theories licensed as open source. Each theory focuses on topics like foundational logic, advanced analysis, algebra, or cryptography and consists of multiple files containing proofs. Similar to the GitHub corpus, the files within each theory are concatenated into a 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 8 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, a 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 10, 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 a detailed comparison of FoT with the Memorizing Transformer [Wu et al., 2022] baseline. Both models are trained for the same number of  $500k$  steps with local context of  $2K$  and evaluated on the arXiv dataset in the single-document setup, following [Wu et al., 2022]. In particular, we study how models trained with a given context length perform when evaluated with different context lengths. These experiments differ from those in Section 5.2, 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 to  $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 11. Parameter values  $w = 1$  and  $w = 2$  correspond to additional
<sup>3</sup>https://www.isa-afp.org
context lengths of  $2K$  and  $4K$ , respectively. In an apples-to-apples comparison to MT with  $2K$  additional context length, FoT outperforms the MT baseline, which shows the importance of trainable keys and values (see also Section C.1). Moreover, we confirm the findings from Section 5.4 that  $d = 2$  works significantly better than  $d = 1$  in all settings. Our best configuration achieves 2.148 perplexity with  $4K$  additional context during training, compared to 2.164 of MT with  $16K$  additional context.
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 for different parameter combinations and evaluation context sizes
Evaluation Context	Parameter Combinations (w, d), Training Memory Length
(1, 1), 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
L 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 B.3. 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.