SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Memory Retrieval Module. The fixed memory-size of cached memory bank in one GPU is 65536 key-value pairs of tokens. We enable each GPU to construct and update their own memory retrieval module for efficiency. For the implementation of the efficient token-to-chunk retrieval, we use the faiss [JDJ21] toolkit to construct an exact-search index on GPU to store the mean-pooled attention keys of text chunks and perform efficient retrieval. The faiss index maintains a fixed  $M / \text{csz}$  keys and provides the efficient exact search w.r.t. inner product. The retrieval takes about 15ms per 1k tokens, which is  $55\%$  timecost of backbone LLM forwarding pass. We can easily adapt the exact search index to approximate search index to gain more retrieval efficiency.
Baselines. In addition to the baseline of our pre-trained GPT-2* variant, we consider Memorizing Transformer (MemTRM) [WRHS22] and TRIME [ZLC22] as two memory-augmented baselines. The MemTRM model can be easily adapted to tune a pre-trained LLM to use external memory. We insert the KNN-augmented layer proposed by MemTRM as the same 18-th layer in the LLM decoder.
Dataset Splits	PG-22	ArXiv
S1	S2	S3	S4	S5
Len. Range	5K-10K	10K-100K	100K-500K	500K-1M	>1M	<60K
#Documents	500	100	30	8	1	100
Avg. #tokens	7.6K	47.6K	140K	640K	1.2M	15.4K
Table 1: Dataset Statistics of five splits of PG-22 based on length range and ArXiv.
To adapt TRIME for our experiments, we replace the batchfying function and loss function of training GPT-2* with those proposed by TRIME, which enables a memory-augmented adaptation tuning method for LLMs. The two reproduced baselines are trained for the same number of tokens under the same hyperparameter setting as LONGMEM.
3.2 Long-Context Language Modeling
The long-context language modeling can potentially benefit from the augmented memory of past long-contexts. The knowledge stored in retrieved attention key-values can provide valuable background and contextual information, helping models perform better in long-context language modeling. For instance, when trying to model a long-text book, acquiring knowledge from previous background and character relationships can be helpful in modeling the subsequent stories.
Evaluation Setting. We first compare LONGMEM and baselines on three long-context modeling datasets, Project Gutenberg 2020-2022, ArXiv, and ChapterBreak. The majority of included books or papers in these datasets have the length of at least 16k tokens. All listed datasets are evaluated in a zero-shot manner without any task-specific tuning. The detailed evaluation settings on the three datasets are as follows:
- Project Gutenberg 2020-2022 Language Modeling Dataset. We crawled and cleaned the books published between 2020 and 2022 under Project Gutenberg Library<sup>1</sup> to build up a completely new long-text modeling dataset, named PG-22. It is significantly different from our training subset PG-19 in terms of domains and writing styles, because books in PG-19 [RPJL19] are published before 1919. We provide different validation splits of PG-22 based on length range, and the data statistics are presented in Table 1.
- ArXiv Dataset. The ArXiv dataset includes papers in the areas of Math, Computer Science, and Physics. We select a validation split of ArXiv paper subset in the Pile corpus  $\left[\mathrm{GBB}^{+}20\right]$ . The ArXiv subset of Pile is excluded from our training and serves an out-of-distribution dataset. We report the token-level language modeling perplexity on the long-context language modeling benchmarks of PG-22 and ArXiv.
- ChapterBreak Benchmark. ChapterBreak [STI22] is a challenging suffix identification dataset that requires LLMs to distinguish the beginning of the ground-truth next chapter from a set of hard negative segments sampled from the same book, given the long context of previous chapters. ChapterBreak requires processing global long-context to comprehend and identify the correct suffix. [STI22] demonstrated that even state-of-the-art x-formers for long-text processing fail to effectively leverage long-range context to perform well on ChapterBreak. ChapterBreak has two subsets, the PG-19 subset and the Archive of Our Own (AO3) subset. As the PG-19 corpus has been included in the pre-training corpus of LONGMEM, it cannot be further used for evaluation. Thus, we select AO3 subset, which contains fan-fictions extracted from AO3. ChapterBreak provides 8 splits based on the prefix length from  $0.5\mathrm{k}$  to 8k tokens to fit the length limit of different models. The splits of 4k, 6k, and 8k prefix are selected for evaluation. For LLMs that cannot process over 4k tokens, we abandon the front prefix to fulfill the maximum input length of LLMs. For memory-augmented models (MemTRM and LONGMEM), we load the given  $4\mathrm{k}/6\mathrm{k}/8\mathrm{k}$  prefix contexts into the cached memory and then do the scoring. We use the perplexity as the scorer for each candidate suffix segment in a zero-shot manner. Then the suffix segment with lower perplexity is selected as the label. The suffix identification accuracy is used as the evaluation metric.
Results. The main results on evaluated long-context datasets are summarized in Table 2. The proposed LONGMEM model significantly outperforms all considered baselines on long-text language modeling
<sup>1</sup>https://www.gutenberg.org/
Model	In-Context
Len.	In-Memory
Len.	5K-10K	10K-100K	PG-22
100K-500K	500K-1M	>1M	ArXiv
GPT-2*	1k	N/A	22.78	24.39	24.12	24.97	18.07	11.05
MemTRM	1k	65K	21.77	23.56	23.23	24.16	17.39	10.81
TRIME	1k	65K	22.21	23.50	23.74	24.32	17.80	10.95
LONGMEM	1k	65K	21.29	23.01	22.55	23.35	16.71	10.05
Table 2: Evaluation results on long-context language modeling datasets. We report token-level perplexity (PPL) (lower the better) on all datasets.
Model	#Params	In-Context
Len.	In-Memory
Len.	ChapterBreaka03
ctx-4k	ctx-6k	ctx-8k
GPT-2-XL† [RWC+19]	1.5B	1K	N/A	24%	24%	24%
GPT-3† [BMR+20]	175B	2K	N/A	28%	28%	28%
LocalTRM† [RSVG21]	516M	8K	N/A	24%	24%	24%
RoutTRM† [RSVG21]	490M	8K	N/A	25%	24%	24%
Bigbird† [ZGD+20]	128M	4K	N/A	26%	26%	26%
GPT-2*	407M	1K	N/A	18.4%	18.4%	18.4%
MemTRM	407M	1K	∞	28.3%	28.7%	28.7%
LONGMEM	558M	1K	∞	37.7%	39.4%	40.5%
Table 3: Zero-shot Suffix Identification Accuracy on AO3 subset of ChapterBreak. Baselines marked with  $\dagger$  are directly cited from [STI22]. The MemTRM and LONGMEM loads the given  $4\mathrm{k} / 6\mathrm{k} / 8\mathrm{k}$  prefix contexts into cached memory, while the input length to local context is still 1k tokens.
datasets, with improvements of 1.38 to 1.62 perplexity on different length splits of  $PG-22$ , and 1.0 on ARXIV datasets. Surprisingly, the proposed method achieves the state-of-the-art performance of  $40.5\%$  accuracy on ChapterBreak $_{\mathrm{AO3}}$  suffix identification benchmark and outperforms both the strong long-context transformers and GPT-3 with 313x larger parameters. The substantial improvements on these datasets demonstrate that LONGMEM can comprehend past long-context in cached memory well for predicting future inputs.
3.3 Memory-Augmented In-Context Learning
LLMs have the emerging capability of in-context learning (ICL) via learning knowledge non-parametrically from few-shot demonstration examples in the local context. However, conventional in-context learning is heavily restricted by input context length, rendering it ineffective to absorb supervision from sufficient demonstration examples in the training set. With the proposed unlimited-length memory augmentation, LONGMEM can overcome the limitation of the number of demonstration examples in the local context and even attend on the whole training set by loading it into the cached memory. In this way, LONGMEM generalizes the conventional few-shot in-context learning to memory-augmented in-context learning with thousands of auxiliary demonstration examples.
Evaluation Setting. Here, we evaluate the in-context learning capability of baselines and the proposed LONGMEM model on five NLU datasets, SST-2 [SPW+13], MPQA [WWC05], MR [ABK+07], Subj [PL04] and SST-5 [SPW+13]. We evaluate models on two few-shot settings, 4-shot and 20-shot. The 4-shot case is the data-insufficient scenario, while the 20-shot demonstrations can almost fulfill the 1k input length and provide sufficient contextual self-supervisions. We transform the k-shot examples to semantically meaningful demonstration examples via fixed text template, i.e.,  $d_{i} = \text{"Review: } x_{i} \text{ Sentiment: } y_{i} \text{"}, \forall \{(x_{i}, y_{i})\}_{i=1}^{k} \in \mathcal{D}_{\text{train}}$  for sentiment analysis tasks. Additionally, we evaluate the 3-shot ICL on question-answering using SQuAD [RZLL16] under an open-ended generation setting. The details of all prompt templates are presented in Appendix D. Then we concatenate the demonstration examples with newlines to delimit them. The prediction label is directly generated using greedy decoding given the demonstration examples and test cases in context. The prediction accuracy is used as the evaluation metric. We report the mean and standard deviation of 6 runs with different random seeds to assess the randomness in selecting k-shot demonstration examples. As mentioned previously, the chunk size controls the granularity of retrieved text chunks. Since the considered NLU datasets require more fine-grained labels from cached memory, we perform
Model	In-Context
#Demons.	In-Memory
#Demons.	SST-2
ACC↑	MR
ACC↑	Subj
ACC↑	SST-5
ACC↑	MPQA
ACC↑	Avg.
Majority	N/A	N/A	50.9	50.0	50.0	20.0	50.0	44.2
GPT-2*	4	N/A	68.311.6	64.712.5	51.94.2	31.44.4	61.511.8	55.6
MemTRM	4	2000	67.512.4	64.611.3	53.26.0	29.64.4	63.012.1	55.6
TRIME	4	2000	69.514.5	63.89.8	51.51.5	31.86.7	63.612.9	56.0
LONGMEM	4	2000	71.814.0	65.111.0	53.83.7	36.06.8	65.412.8	58.4
GPT-2*	20	N/A	68.211.5	63.45.2	57.610.2	33.66.0	70.87.6	58.7
MemTRM	20	2000	65.19.6	65.19.3	58.210.6	31.96.3	72.77.4	58.6
TRIME	20	2000	74.313.9	71.52.5	57.511.4	33.04.6	69.87.8	61.1
LONGMEM	20	2000	78.014.1	78.63.3	65.68.5	36.57.5	74.67.3	66.7
a hyperparameter selection on the validation set of SST-2, and the best chunk-size 2 is used to report the results for MemTRM, TRIME and our model.
Results. The results on in-context learning are summarized in Table 5 and Table 4. LONGMEM achieves remarkable improvements on all NLU tasks under the 20-shot sufficient in-context setting, with +5.6 average scores increase over pretrained GPT-2*, MemTRM, and TRIME.
Meanwhile, LONGMEM also brings performance improvements on the 4-shot case. Additionally, LONGMEM improves the in-context learning capabilities of LLMs on open-ended generation tasks, with +4.5 EM score increase on SQuAD. The results indicate that having more demonstration examples loaded in cached memory can provide additional contextual cues to assist in-context learning. LONGMEM can utilize task-relevant knowledge from both local contextual demonstrations and in-memory augmented demonstrations, thereby achieving superior in-context learning capabilities.
Table 5: Accuracy [%] of 4-shot and 20-shot ICL on 5 NLU tasks (SST-2, mr, subj, SST-5, mpqa). We sample 2000 extra demonstration examples and load them into cached memory. The subscript is the standard deviation across 6 runs. Avg. refers to the average accuracy on 5 datasets.
Model	EM	F1
GPT-2*	22.282.3	30.782.0
MemTRM	22.843.5	32.652.8
LONGMEM	26.772.3	35.702.0
Table 4: Exact match (EM) and F1 scores of 3-shot (about 1k tokens) in-context learning on SQuAD. LONGMEM loads 200 extra demonstration examples into cached memory.
3.4 Ablation Studies
So far, we empirically verify the effectiveness and superiority of LONGMEM in utilizing cached memory for long-context modeling, long-context understanding, and many-shot in-context learning. Furthermore, we would like to investigate the extend to which the cached memory contributes to the long-context understanding capability of LONGMEM through an ablation study of removing memory augmentations. Besides, since the design of the cached memory bank involves several hyperparameters, such as memory size  $msz$  and chunk-size  $csz$ , we conduct a series of ablation studies to evaluate the effects of those choices.
Effects of Long-Term Memory Augmentation. To evaluate the effects and contributions of memory augmentations, we set the memory-size to 0 and maintain the SideNet parameters during inference. The results of LONGMEM without memory augmentation are shown in Table 6 of Appendix. As expected, without augmented long-term memory, the vanilla model with only backbone LLM and SideNet only gains 59.4 average scores on ICL NLU tasks, which is a 7.3 average accuracy decrease due to the removal of memory augmentation.
Effects of Chunk-Size. As analyzed before, the chunk-size  $csz$  controls the granularity of retrieval and thus it may make a difference to tasks with requirements of fine-grained retrieval. We perform an ablation study on the effects of various chunk-size choices  $csz \in \{2,4,8\}$  for in-context learning and the results are presented in 4(a). The chunk size of 2 yields the best performance on in-context learning tasks on five NLU datasets, which is consistent with the property of NLU tasks with the requirement of fine-grained retrieval and fusion towards classification label tokens.
Effects of Memory Size. The memory size (msz) controls the capacity of the memory bank. In general, the memory size should be compatible with the average length of documents or contexts,
(b)
Figure 4: (a) Accuracy on 5 NLU datasets given different chunk size during inference; (b)  $\Delta$ Perplexity on 4 splits of PG-22 given different memory size during inference, in which the perplexity when  $msz = 65k$  is used as baseline.
i.e., a set of books with average 16k tokens should deploy the memory size of 16k tokens in cached memory. The training  $msz$  of 65 tokens is excessive for downstream tasks such as ChapterBreak as the whole prefix context length does not exceed 65k tokens. Thus, we perform an ablation study on the effects of memory size  $msz \in \{8k, 16k, 32k, 65k\}$  during the inference stage on the PG-22 language modeling datasets and the results are shown in 4(b). To model the books with lengths of 8k-50k, the smaller memory size 16k which is consistent with the average length of target books yields the best perplexity.
Large Language Models. Large Language Models, i.e., GPT-3  $\left[\mathrm{BMR}^{+}20\right]$ , LLAMA  $\left[\mathrm{TMS}^{+}23\right]$ , GPT-4 [Ope23], significantly revolutionized NLP research and promoted the state-of-the-art of various language understanding, language generation  $\left[\mathrm{WZG}^{+}22\right]$ , and even vision-language tasks  $\left[\mathrm{WDC}^{+}22\right]$ . Additionally, enabled by multi-task instruction tuning  $\left[\mathrm{WBZ}^{+}21, \mathrm{OWJ}^{+}22\right]$ , LLMs exhibit "emergent abilities"  $\left[\mathrm{WTB}^{+}22\right]$  like mathematical reasoning  $\left[\mathrm{WWS}^{+}22\right]$ , code completion  $\left[\mathrm{CTJ}^{+}21\right]$ , etc.
x-formers. To enable transformers to attend on longer context, many variants of "x-formers" are proposed. Transformer-XL  $\left[\mathrm{DYY}^{+}19\right]$  proposes to cache attention keys and values of past segment and reuse them in recurrent manner. Recent seminal works of x-formers, including LinFormer [WLK  $^+$  20], LongFormer [BPC20], Routing Transformer [RSVG21], proposed various sparse attention mechanisms for decreasing  $O(n^{2})$  complexity to  $O(n\log n)$  or even  $O(n)$ . BigBird  $\left[\mathrm{ZGD}^{+}20\right]$  achieves a 4k sequence length via attending on a subset of context tokens. Although these x-formers achieve substantial efficiency improvements, such efficiency gains are not remarkable when modeling sequences that spans book-level length. Moreover, the largest sequence length of these methods is still upper-bounded by 16k tokens, making them invalid in modeling long-sequences at the book or wikipedia-page level (i.e., average 70k tokens for full-length books in PG19 dataset [RPJL19]).
Side-Tuning. The method of Side-Tuning  $\left[\mathrm{ZSZ}^{+}20, \mathrm{SCB22}\right]$  is a task-specific tuning method for pre-trained models via training a lightweight side-network that is fused with the fixed pre-trained network via summation. Our method inherits the idea of adopting a side-network but distinguishes the side-tuning method in terms of learning objective and cross-network fusion ways. LONGMEM proposes to augment LLMs with decoupled memory to retrain information from long past inputs without any task-specific tuning. The cross-network residual connections introduced here are novel and distinct from the vanilla summation used in Side-Tuning.
In this paper, we propose to augment LLMs with long-term memory for enabling them to memorize long-form context and gain long-form memory. The designed decoupled memory module can cache attention key and value pairs of past inputs for future retrieval and fusion. A decoupled residual SideNet is introduced as the memory retriever and reader, meanwhile the LLM itself is frozen and works as knowledge and memory encoder. Experiments on various long-contextual language modeling datasets demonstrate the effectiveness of our model over other memory-augmentation baselines. The proposed method can also enable in-context learning of LLMs to overcome the limited number of demonstration examples in context, which is constrained by the contextual length, via caching thousands of auxiliary demonstration examples in memory.
This work is done during the first author's internship at Microsoft Research. We would like to thank the anonymous reviewers for the helpful comments. We appreciate Yutao Sun and Yaru Hao for helpful suggestions on implementation and evaluation benchmarks. The first author was partly sponsored by the DARPA PTG program (HR001122C0009). Any opinions, findings, conclusions, or recommendations expressed in this paper are those of the authors and do not necessarily reflect the views of funding agencies.
[WZG+22] Weizhi Wang, Zhirui Zhang, Junliang Guo, Yinpei Dai, Boxing Chen, and Weihua Luo. Task-oriented dialogue system as natural language generation. In Proceedings of the 45th International ACM SIGIR Conference on Research and Development in Information Retrieval, pages 2698-2703, 2022.
[  $\mathrm{YD}\mathrm{Y}^{+}19]$  Zhilin Yang, Zihang Dai, Yiming Yang, Jaime G. Carbonell, Ruslan Salakhutdinov, and Quoc V. Le. XLNet: Generalized autoregressive pretraining for language understanding. In NeurIPS, 2019.
$\left[\mathrm{ZGD}^{+}20\right]$  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.
[ZLC22] Zexuan Zhong, Tao Lei, and Danqi Chen. Training language models with memory augmentation. arXiv preprint arXiv:2205.12674, 2022.
[  $\mathrm{ZSZ}^{+}20]$  Jeffrey O Zhang, Alexander Sax, Amir Zamir, Leonidas Guibas, and Jitendra Malik. Side-tuning: a baseline for network adaptation via additive side networks. In Computer Vision-ECCV 2020: 16th European Conference, Glasgow, UK, August 23-28, 2020, Proceedings, Part III 16, pages 698-714. Springer, 2020.
A Ablation Study on the Effect of Memory Augmentation
Model	In-Context
#Demons.	In-Memory
#Demons.	SST-2
ACC↑	MR
ACC↑	Subj
ACC↑	SST-5
ACC↑	MPQA
ACC↑	Avg.
Majority	N/A	N/A	50.9	50.0	50.0	20.0	50.0	44.2
GPT-2*	4	N/A	68.311.6	64.712.5	51.94.2	31.44.4	61.511.8	55.6
MemTRM	4	2000	67.512.4	64.611.3	53.26.0	29.64.4	63.012.1	55.6
TRIME	4	2000	69.514.5	63.89.8	51.51.5	31.86.7	63.612.9	56.0
LONGMEM	4	2000	71.814.0	65.111.0	53.83.7	36.06.8	65.412.8	58.4
w/o Memory	4	0	69.412.4	64.312.1	53.47.7	29.05.2	62.512.3	55.7
GPT-2*	20	N/A	68.211.5	63.45.2	57.610.2	33.66.0	70.87.6	58.7
MemTRM	20	2000	65.19.6	65.19.3	58.210.6	31.96.3	72.77.4	58.6
TRIME	20	2000	74.313.9	71.52.5	57.511.4	33.04.6	69.87.8	61.1
LONGMEM	20	2000	78.014.1	78.63.3	65.68.5	36.57.5	74.67.3	66.7
w/o Memory	20	0	70.012.8	70.86.2	52.94.6	30.96.4	72.57.5	59.4
B Inference Efficiency and GPU-Memory Efficiency
When the model is required to comprehend long sequences, the proposed method LONGMEM can load the out-of-boundary inputs into the cached memory as previous context. Thus, the memory usage and inference speed can be significantly improved compared with vanilla self-attention-based models. The detailed statistics in terms of the efficiency is presented in Table 7.
Table 6: Ablation study results on the effect of memory augmentation of 4-shot and 20-shot ICL on 5 NLU tasks (SST-2, mr, subj, SST-5, mpqa). We sample 2000 extra demonstration examples and load them into cached memory. The subscript is the standard deviation across 6 runs. Avg. refers to the average accuracy on 5 datasets. "w/o" is short for "without".
Model	In-Context
Len.	In-Memory
Len.	Inference Speed
(tokens/s)↑	GPU-Memory Usage
(MBs)↓
GPT-2*	4k	N/A	14666	20671
LONGMEM	1k	3k	22638	13335
GPT-2*	8k	N/A	8417	54195
LONGMEM	1k	7k	21343	13437
Table 7: The superiority of our method over fully dense self-attention (GPT-2*) in terms of inference speed and GPU-memory utilization.
The pre-training of reproduced GPT-2* iterates on 117B tokens in total, with 512 batch-size and 1024-token fixed segment-length. The Adam optimizer [KB15] is adopted in memory-augmented adaptation training. The pre-training and adaptation are trained on 16 32GB-Tesla-V100 GPUs. Other detailed training hyperparameters and settings are presented in Table 8.
Hyperparameter	LONGMEM
Reproduced GPT-2* Backbone LLM Hyperparameters
Parameters	407M
Precision	float16
Layers	24
Hidden dim.	1024
Attention heads	16
Head Dim	64
Vocab size	52k
Sequence length	1024
Position emb.	Alibi
Tied embedding	False
SideNet Hyperparameters
Parameters	151M
Precision	float16
Layers	12
Hidden dim.	1024
Attention heads	16
Head Dim	64
Sequence length	1024
Memory-Augmented Adaptation Hyperparameters
Global Batch Size	256
Learning rate	2.0e-4
Total tokens	26B
Warmup tokens	0
LR Decay style	polynomial
Adam (β1, β2)	(0.9, 0.98)
Adam eps	1e-06
Weight decay	0.01
Gradient clipping	2.0
D Prompting Templates
We present all hand-crafted in-context learning prompting templates and labels for 5 NLU datasets and Squad QA dataset in Tabel 9.
Table 8: Memory-Augmented Adaptation and Architectural Hyperparameters.
Task	Prompt	Labels
SST-2	Review: [Sentence] Sentiment: [Label]	{positive, negative}
MR	Review: [Sentence] Sentiment: [Label]	{positive, negative}
MPQA	Review: [Sentence] Sentiment: [Label]	{positive, negative}
SST-5	input: [Sentence] type: [Label]	{terrible,bad,okay,good,great}
Subj	input: [Sentence] type: [Label]	{objective, subjective}
Squad	Passage: [Passage]\nQuestion: [Question] Answer: [Answer]
Table 9: The hand-crafted prompts used to query the model predictions on the zero-shot evaluation of 5 NLU datasets and one question-answering dataset Squad.