Augmenting Language Models with Long-Term Memory
Authors: Weizhi Wang
Li Dong
Hao Cheng
Xiaodong Liu
Xifeng Yan
Jianfeng Gao
Furu Wei
Venue: NeurIPS 2023 poster (2023)
Keywords: large language models
long-term memory
long-text modeling and understanding
residual side-network
in-context learning
TL;DR: We propose to augment Large Language Models with Long-Term Memory for memorizing long-form previous contexts and meanwhile propose a residual SideNet as the memory retriever and reader.

Abstract: Existing large language models (LLMs) can only afford fix-sized inputs due to the input length limit, preventing them from utilizing rich long-context information from past inputs. To address this, we propose a framework, Language Models Augmented with Long-Term Memory (LongMem), which enables LLMs to memorize long history. We design a novel decoupled network architecture with the original backbone LLM frozen as a memory encoder and an adaptive residual side-network as a memory retriever and reader. Such a decoupled memory design can easily cache and update long-term past contexts for memory retrieval without suffering from memory staleness. Enhanced with memory-augmented adaptation training, LongMem can thus memorize long past context and use long-term memory for language modeling. The proposed memory retrieval module can handle unlimited-length context in its memory bank to benefit various downstream tasks. Typically, LongMem can enlarge the long-form memory to 65k tokens and thus cache many-shot extra demonstration examples as long-form memory for in-context learning. Experiments show that our method outperforms strong long-context models on ChapterBreak, a challenging long-context modeling benchmark, and achieves remarkable improvements on memory-augmented in-context learning over LLMs. The results demonstrate that the proposed method is effective in helping language models to memorize and utilize long-form contents.

Augmenting Language Models with Long-Term Memory
Weizhi Wang', Li Dong', Hao Cheng', Xiaodong Liu', Xifeng Yan Jianfeng Gao', Furu Wei University of California, Santa Barbara Miirosooft Research weizhiwang@ucsb. edu, {lidongl, chehao, xiaodl) [modl])@microooft.oo
Abstract
Existing large language models (LLMs) can only afford fix-sized inputs due to the input length limit, preventing them from utilizing rich long-context information from past inputs. To address this, we propose a framework, Language Models Augmented with Long-Term Memory (LONGMEM), which enables LLMs memorize long history. We design a novel decoupled network architecture with the original backbone LLM frozen memory encoder and adaptive residual side-network as a memory retriever and reader. Such a decoupled memory design easily cache and update long-term past contexts for memory retrieval without suffering from memory staleness. Enhanced with training, LONGMEM can thus memorize long past context and long-term memory for language modeling. The proposed memory retrieval module can handle unlimited-length context its memory bank to benefit various downstream tasks. Typically, LONGMEM can enlarge the long-form memory to 65k tokens and thus cache many-shot extra demonstration examples as long-form memory for in-context learning. Experiments show that our method outperforms strong long- context models on ChapterBreak, a challenging long-context modeling benchmark, and achieves remarkable improvements on nemory-augmentt over LLMss. The results demonstrate that the proposed method effective helping language models to memorize and utilize long-form contents. Our code is open-sourced at https: / akaa ms Longgg....
adaptation
in-context learning
Introduction
Large language models (LLMs) have revolutionized natural language processing with great successes in advancing the state-of-the-e on various understanding and generation tasks [DCLT19, RWC 19, LOG 19, YDY+ 19, BMR + 20, RSR+20) Most LLMs benefit from self-supervised training over large corpora via harvesting knowledge from fix-sized local context, showing emergent abilities, e.g., zero-shot prompting [RWC-19) in-context learning [BMR+ 20), and Chain-of-Thought (CoT) reasoning [WWS+ 22]. Nevertheless, the input length limit of existing LLMs prevents them from generalizing to real-world scenarios where the capability of processing long-form information beyond a fix-sized session is critical, e.g., long horizontal planning.
To address the length limit issue, the most straightforward method is to simply scale up the in- put context length. For instance, GPT-3 [BMR 20] increases the input length from Ik tokens in GPT-2 |RWC+19 19] to 2k tokens, thereby allowing for better capture of long-range dependen- cics. However, this approach typically incurs cumputation-intensiin training from scratch and the in-context dense attention is still heavily constrained by the quadratic computation complexity of Transformer self-attention [VSP+ 17]. Another recent line of work [BPC20, ZGD+200 +00 instead fo- cuses developing in-context sparse attention to avoid the quadratic cost of self eelf-attention, which
37th Conference on Neural Information Processing Systems (NcurIPS 2023).

---

Cached Memory Bank with Key, Value Pairs
Retrieved Attn Keys and Values
MongoMyrorry Retrieval
T
T
Search
Attn Keys and Values (Seg
Keys and Values (Seg B)
Keys and Values (Seg
...
Attention Query of Current Inputs
Memory Fusion
Large Language Model (Frozen)
Residual
Residual SideNet
Connections
(Trainable)
Long Sequence Inputs
Figure 1: Overview of the memory caching and retrieval flow of LONGMEM. The long text sequence is divided into fix-length segments with each previous segment processed through a frozen backbone LLM and the corresponding attention key and value vectors of m-th layer are cached into the memory bank. Given current inputs, the corresponding attention query vectors are used to retrieve the top- attention key-value pairs from previous segments stored in the memory bank, which will be then fused with local context for language modeling.
still largely requires training from scratch. In contrast, the prominent work, Memorizing Transformer (MemTRM) [WRHS22], approximates in in-context sparse attention via dense attention over both in- context tokens and memorized tokens retrieved from non-differentiable memory for Transformers. Thus, MemTRM scales up the resulting language model to handle up to 65k tokens and achieves substantial perplexity gains in modeling full-length books or long papers. However, MemTRM faces the memory stalenesy challenge during training due to its coupled memory design, which uses single model for encoding memory and fusing memory for language modeling. other words, as the model parameters are updated, cached older representations in memory may have distributional shifts from those from the latest model, thereby limiting the effectiveness of the memory augmentation.
In this paper, we present framework for Language Models Augmented with Long- Term Memory (LoNGMEM). This framework enables language models to cache lengthy previous context or knowledge into non-differentiable memory bank, and then utilize them via a decoupled memory module to mitigate the issue of memory staleness. To achieve decoupled memory, we design novel residual side-network (SideNet) in conjunction with a frozen backbone LLM. Paired attention keys and values of the previous context are extracted using the frozen backbone LLM, which are subsequently stored in the memory bank. In the eemory-aaggmente layer of SideNet, the generated attention query the current input is used to retrieve cached key-value pairs of previous contexts from the memory, and the corresponding memory augmentations then fused into adaptable hidden states via joint-attention mechanism. Furthermore, newly designed cross-nctwork residual connections between the SideNet and the frozen backbone LLM facilitate better knowledge transfer from the pretrained backbone LLM. Through continuous training of the residual SideNet to retrieve and fuse memory augmentations, the pre-trained LLM can be adapted to effectively leverage long- contextual memory for enhanced modeling. The detailed memory cache, retrieval and fusion process is illustrated in Figure 1.
Our decoupled memory design offers two key advantages. First, our proposed architecture effectively separates the process of encoding previous inputs into memory and the process of memory retrieval and fusion, thanks to the decoupled frozen backbone LLM and SideNet. In this way, the backbone LLM only works as the long-context knowledge encoder, while the residual SideNet works as the memory retriever and reader, which effectively resolves the issue of memory staleness. Second, directly adapting the entire LLM with memory augmentations is computationally inefficient and also prone to catastrophic forgetting. As the backbone LLM frozen during the efficient memory- augmented adaptation stage, LONGMEM not only tap into the pretrained knowledge but also avoid catastrophic forgetting.
LONGMEM is capable of taking various types of long-form text and knowledge into the memory bank based on downstream tasks. Here, we consider two representative cases, language modeling with full-length book contexts, and memory-augmentie in-context learning with thousands of task-relevant

---

Cached Bank
Ceche
fays and Valuea
Languoge
Memory Retrieval and Fusion
Decodar Layer
X
LLM Dccoder Layor
Hemmunn Layer
(V5)
Docoder Layer
Memory Fusion
ULM Decoder Layer
Salle Layar
b
Retiovall
Eribddding Layer
Lage
Cached Mensey Bank with Key, Vavue Pars
Redidall
Figure 2: Overview of LoNGMEM architecture. "MemAug" represents Memory-Augmented Layer.
demonstration examples. Specifically, we evaluate the effectiveness of the proposed LONGMEM on various long-text language modeling, and memory-augment in-context learning for natural language understanding (NLU) tasks. Experimental results demonstrate that our model consistently outperforms the strong baselines terms of long-text modeling and in-context learning abilities. Our method substantially improves LLM's long- context language modeling capabilities, with a reduction in perplexity of 1.38~~.62 ..22 over different length splits of Gutenberg-2022 corpus. Notably, our model achieves hate-of-theeart performance of 40.5% identification accuracy ChapterBreak, challenging long-context modeling benchmark, significantly surpassing existing strong x-former baselines. Finally, with 2k demonstration examples in memory, LoNGMEM shows pronounced improvements in-context learning on popular NLU tasks, compared with prominent memory- augmented and noo-memmry--aggmente baselines.
2 Methods
To enable LLMs harvest relevant information from the past long context in memory, we propose to augment the frozen backbone LLM with a decoupled memory module. To fuse the memory context information, we design novel lightweight residual SideNet, which can be continually trained in an efficient way. In the following, we first discuss the problem formulation of language modeling with memory augmentations. Then, we formally introduce our efficient residual SideNet for adapting the frozen pretrained LLM to jointly attend over local input context and retrieved memory context. Lastly, we provide our designed processes of how past memory is encoded, stored, recalled and fused for language modeling.
2.1 Language Models Augmented with Long-Teree Memory
Here, we focus on the high-level problem setup and defer more component details to later sections. Given its wide adoption for pretrained LLMs, our LONGMEM model is built the Transformer architecture [VSP+]] 17]. For LONGMEM, there are three key components: the frozen backbone LLM, SideNet, and Cache Memory Bank. As most existing pretrained LLMs can only take fix-sized input, only the input segment of a long sequence (e.g., a book) that can fit in the length limit is denoted as the current input as done for most existing autoregressive language models. Those previous segments that can not fit are denoted as previous inputs, which are used for memory augmentations. To tap into the learned knowledge of the pretrained LLM, both previous and current inputs are encoded using the frozen backbone LLM but different representations are extracted. For previous inputs, the key-value pairs from the Transformer self-attention at m-th layer are stored in Cache Memory Bank, whereas the hidden states from each LLM decoder layer for the current inputs are retained and transferred to SideNet. For each current input token, top relevant key-value vector pairs are retrieved as memory augmentations for language modeling. The SideNer module can be viewed as efficient adaption model that trained fuse the current input context and relevant cached previous contexts in the decoupled memory.
3

---

Formally, for a fix-sized input text sequence {x (the current input), LONGMEM first performs a forward pass using the backbone LLM (indicated in blue in Figure 2) without any gradient calculation. The embedding layer of the backbone LLM first encodes the input {x;
{
|sl i==1 into
embedding space and outputs the initial hidden states, HLLM E R where E. is the hidden dimension. Then each successive Transformer decoder layer of the frozen backbone LLM computes the new hidden states using the hidden states from the previous layer, HILM - Seti (H) V/' E [L, L'] and L' the total # layers for the backbone LLM. During the forward pass with the backbone LLM for all previous inputs, the key-value pairs used for self-attention at the m-th Transformer decoder layer are stored in Cached Memory Bank (highlighted in orange in upper-left of Figure 2).
These pairs are subsequently recalled as memory augmentations for future inputs.
Cached Memory Bank is head-wise vector queue Zg. Z, RRIxxxd, which maintains attention
key-value pairs of latest M previous inputs K, V E R ar) x|/dd where H, denotes the number of attention heads and per-head dimension respectively. After memory retrieval and fusion ($2.3), the memory bank removes the key-value pairs of the oldest sequences and appends the current sequences to the cached vector bank. This update mechanism ensures the language modeling causality at the sequences level and enables the memory bank to consistently maintain records of the most recent previous context for the current inputs.
After the forward pass with the backbone LLM,, the SideNet module then takes all current input hidden states from the backbone LLM [H] [H(LM) ) and the past key-value pairs the Cached Memory Bank for computing memory-augmentee representations. Specifically, our SideNet of LONGMEM consists of (L (L-1] 1) normal Transformer decoder layers and one special memory -uggneeed decoder layer. For efficient purposes, we mainly consider the case where #layers L of the SideNet is smaller than that of the backbone LLM, i.e., < L'. Our SideNet encodes H into aemory-augmented contextual representation via [L. 1) normal Transformer decoder layers and a special memory-augmenter layer.
The memory-augmented layer is an extension of the vanilla Transformer decoder layer that takes a memory- -augeetd input, including both top relevant key-value pairs memory and the hidden states from the current input. Here, the cached key-value pairs are recalled using token-based memory retrieval module ($2.3). For each current input token, the memory retrieval module syr(s) retrieves top-K relevant key-value pairs in the memory bank (kej, Tig 1K = s,y (x,). Then SideNet computes |/2 Side : [Kk,, ) )), where
the output using the uemory-augmentte input, H = fexmm (H m. m, is the layer index where we inject the
layer.
Finally, the token probability computed using the last SideNet hidden states P(x:|x1,'. . softmax (W where W is the frozen output embedding weight shared by both the backbone LLM and SideNet. We perform a memory- augmented adaptation training for LONGMEM to utilize the decoupled memory. Following the generative unsupervised pre-training [RNSS18] 18, the training objective of LONGMEM the standard left-to-right language modeling objective, which maximizes the likelihood of the next token based on the left context: max EnED Σi= log P(x:|x), "xx,)))) where x is randomly sampled sentence from the pre-training text corpus D.
2.2 Residual SideNet
SideNet Architecture and Initialization. Here, we implement SideNet based on Trans- former [VSP+17] Specifically, the number of decoder layers I in SideNet is equal to the number of layers L' in the backbone LLM divided by a reduction factor (a layer reduction factor of 2 is used throughout this work, i.e.. L' 2L). The weights of each decoder layer in SideNet are ini- tialized from the corresponding pre- trained decoder layer of the backbone LLM at the same depth: Side 0i As illustrated in Figure 2, the SideNet model takes the output of backbone LLM's embedding layer and reuses the language modeling head of the backbone LLM,, which remains frozen during the continual adaption stage. Throughout the memory- eegore-augmented adaptation stage, all other parameters of SideNet are updated based on the training signal. In this way, the lightweight SideNet adaptation achieves fast convergence with knowledge transferred from pre-trained parameters.
Cross-Network Residual Connections. To tap into knowledge from the pretrained backbone LLM,, we utilize proposed netss-n residual connections to fuse representations from the backbone

---

LLM into SideNet. Specifically, we add the difference between output hidden states at 21-th and (21 2)-th layers of the backbone LLM as the residual connections to the output hidden states at 1-th layer of SideNet. Then, the input to the next (7 + 1)-th layer of SideNet is the sum of the original hidden state forwarded through the previous layer feg.. (HH and the cross-network residual connection of the hidden state difference from the backbone LLM
H
feg. (Hi) + (HELM H VI [1,
(I)
where HD is the output of embedding layer. It is worth noting that the residual connections after the self-attention and feed-forward network of a decoder layer [VSP+ 17] will performed as normal in feg... (H)) and parallel to the proposed neoso- residual connections.
2.3 Memory Retrieval and Fusion
The long-term memory capability of LonGMEM is achieved via memory semmoryauummentttion module for retrieval and fusion.
torcnt Memory Retrieval. Instead of performing token-to-token retrieval, we focus on retrieval for acceleration and integrity. A text-chunk refers to an n-gram structure of chunk-size number of contiguous tokens. The memory bank stores cached key-value pairs the level of token chunks. We divide the memory bank into M/ csz attention key-value paired chunks and use the mean-pone pooled vector on the chunk-size dimension get the key vector for retrieval. Then we retrieve the top-(K' (K/csz) attention key-value chunks w.r.t the dot product between the attention query of the current input token and the mean-pooled attention key of candidate chunk. Finally, we squeeze the chunk-size dimension for retrieved key-value paired chunks and flatten them into K key-value pairs at token-level (K, V ッ Adopting loken-to-chung retrieval reduces the size of the retrieval index and accelerates the process. Meanwhile, the retrieval accuracy can be further improved, which is also observed in [L.GW+2) and [BMH-21]. The hyperparameter chunk-size caz controls the granularity of retrieved contexts, which can be empirically adjusted based on downstream tasks. For instance, in-context learning requires more examples cached in memory, where a smaller esz is helpful.
label tokens from demonstration
Memory Fusion. The memory fusion is performed within a special memory-augmenteed layer. As the conventional Transformer decoder layer uses the follow [WRHS22] to extend it to atttt-atention mechanism and propose a long-term memory fusion process to enable each token to attend on both local contexts and retrieved memory contexts. With the head-wise hidden state output from previous layer H' e Rlæiid and the corresponding retrieved attention key-value pairs are (K,, V.) i= E Rl the output hidden state for the I-th memory-augmented layer H' is computed as:
self-attention [VSP* 17], we
QK Vd
Q.K vd
A softmax
JV, M Concat (softmax
V
(2)
H sigmoid (g) A + sigmoid(g)) M,
where Q K, V. A, M E Rlalxd, K is the number of retrieved attention key-value pairs in cached memory for each token, and is a trainable head-wise gating vector. The hidden state output from previous layer (!-1) is linearly projected into attention queries, keys, and values Q, K, V separately via three matrices W W WV Rdxd It is worth noting that the retrieved attention key-value pairs in cached memory are distinct to each token.
3 Experiments
We evaluate our proposed LONGMEM model on different tasks that require long-context modeling: a) long-text language modeling and language understanding when loading the past long-context into cached memory; infinite-length in-context learning when loading a large number of demonstration examples into cached memory.
5

---

Long Decuments
Document Grouping
Truncateed Segments
Batchtying
Document1 1
9:1
Siz
Bi2
Sus
8GM
9GN
921
9zz
921
9an
Document2 2
Sat
Document
S41
9
San
9gs
B2s
Sax
Battch Batch Batch
Batch
Cached Memory
Update during 54h lteration
S
S14
Suz
Sys
Document
Fan
Syx
Fax
Sz1
9zz
9p
8z2
9z
Bqs
Syt
94
9yz
9,4
Sat
Sa1
5
5p
Su
Sas
Document
1..
U
Remove Oldest
Append Lastest
Batchfying the large text corpora into batches to ensure that each consecutive segments within each document distributed in consecutive batches.
Figure 3:
3.1 Training Setup
Batchfying the training corpora. The conventional batchyfing process for large corpora truncates the whole corpora into consecutive fix-length text segments without padding and shuffles all segments to construct mini-batches [RWC+ 19]. In contrast, LONGMEM must disable global shuffling and ensure the global causality at the segment level. Firstly, we divide all long documents in training corpora into batch-size number of document groups with equivalent length and then perform document-leve shuffling within each group. Then, we concatenate shuffled documents within one group and truncate them into ordered segments. In order to ensure that two consecutive segments of one long document are distributed in two consecutive input batches after batchfying, we select one segment from batch-size number of document groups with the same inner-group index. Thus a mini-batch with batch-size number of segments are constructed from exactly the batch-size number of document groups. In this way, as the training iteration steps, the cached attention key-value pairs in the memory bank are previous context of current inputs within the same document. The batchfying process is illustrated in Figure 3.
Training Corpus, Backbone LLM and Hyperparameter. We sample a subset of the Pile [GBB+20 as the training corpus, including BookCorpus2, Books3, WebTexe,, Stack Exchange, Wikipedia, Gutenberg (9)-19), NIH ExPorter, and Pile-CC datasets. We reproduce GPT-2* (407M-params) as the pre-trained backbone LLM with Alibi [PSL21] position embedding because original GPT- 2 [RWC+ 19] adopts absolute position embedding, which is found to perform poorly to enable LLM to learn long-distance dependencies [DYY+19] The backbone LLM holds a L,' 24, H 16, d = 64 architecture. The SideNet holds L 12, H 16, d 64 architecture. The training for memory- augmented adaptation iterates on 26B tokens, with a global 256 batch-size and 1024 sequence length. The chunk-size csa is 4 tokens and the memory size M is 65k key-value pairs of tokens. For each token, we retrieve K-64 attention key-value pairs for augmentation, which are K/csz=16 text chunks. layer is the 9-th layer of SideNet. The attention keys and values from 18-th layer of backbone LLM is cached into memory and used for future retrieval. Other training details are presented in Appendix C.
The mmeoor--aggmentoe
Memory Retrieval Module. The fixed memory-size of cached memory bank one GPU 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-chumi retricval, we use the faiss [JDJ21] toolkit to construct an exact- search index GPU to store the mean-poole attention keys of text chunks and perform efficient retrieval. The faiss index maintains a fixed M M/esz keys and provides the efficient exact search w.r.t. inner product. The retrieval takes about 15ms per Ik tokens, which is 55% timecost of backbone LLM forwarding pass. We can easily adapt the exact search index to approximate search index to gain retrieval efficiency.
Baselines. addition to the baseline our pre-trained GPT-2* variant, we consider Memorizing baselines. The TRMTRM model can be easily adapted to tune a pre-trained LLM to use external memory. We insert the KNN- KNN-augmented layer proposed by MemTRM the same 18-th layer in the LLM decoder.
Transformer (MenTRM) [WRHS22] and TRIME [ZLC22]] as two nemory-aaugmente

---

Dataset
PG-22
ArXiv
Splits
S1
$2
S3
S4
S5
Len. Range
5K-10K
10K-100k
100K-500K
500K-IM
>IM
<60k
#Documents
500
100
30
8
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 anemory-augmentt adaptation tuning method for LLMs. The two reproduced baselines are trained for the same number of tokens under the same hyperparameter setting 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-contex: 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 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! to build up 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 [GBB* 20]. The ArXiv subset of Pile is excluded from training and serves We report the token-level language modeling perplexity on the long-context language modeling benchmarks of PG-22 and ArXiv.
dataset.
ChapterBreak [STI22] is 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 (A03) subset. As the PG-19 corpus has been included in the pre- pre-training corpus of LONGMEM. it cannot be further used for evaluation. Thus, we select A03 subset, which contains fan-fictions extracted from AO3. ChapterBreak provides 8 splits based on the prefix length from 0.5k to 8k tokens to fit the length limit of different models. The splits of 4k, 6k, and 8k prefix are selected for evaluation. For L.MSs that cannot process over 4k tokens, we abandon the front prefix to fulfill the maximum input length of LLMs. For memory- -uggneeee models (MemTRM and LONGMEM), we load the given 4k/6k/8k prefix contexts into the cached memory and then do the scoring. we use the perplexity as the scorer for each candidate suffix segment in 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.
ChapterBreak Benchmark.
Results.
The main results evaluated long- long-context datasets summarized in Table 2. The proposed LONGMEM model significantly outperforms all considered baselines on long-text language modeling
hattss//www.uueeberg..rg

---

In-Context
In-Memory
PG-22
Model
ArXiv
Len.
Len.
5k-10k
10K-100K
100K-500K
300K-IM
>IM
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
Ik 65K
21.29
23.01
22.55
23.35
16.71
10.05
Table 2: Evaluation results on long-context language modeling datasets. perplexity (PPL) (lower the better) all datasets.
We report token-level
In-Context Len.
In-Memory Len.
ChapterBreak ctx-6k
Model
#Params
etx-4k
ctx-8k
GPT-2-XL [RWC*19]
IK
N/A
24%
24%
24%
GPT-3 1BMR*200
175B
2K
N/A
28%
28%
28%
LocalTRM [RSVG21]
516M
8K
N/A
24%
24%
24%
RourTRM [RSVG21]
490M
8K
N/A
25%
24%
24%
Bigbird' [2GD*20]
128M
4K
N/A
26%
26%
26%
GPT-2"
407M
IK
N/A
18.4%
18.4%
18.4%
McmTRM
407M
1K
DC
28.3%
28.7%
28.7%
LONGMEM
558M
1K
DC
37.7%
39.4% %
40.5%
Table 3: Zero-shot Suffix Identification Accuracy A03 subset of ChapterBreak. Baselines marked with i are directly cited from [STI22]. The MemTRM and LongMEM loads the given 4k/6k/k prefix contexts into cached memory, while the input length local context is still Ik 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-an performance of 40.5% accuracy on ChapterBreakAo: 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-conter 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 super- vision from sufficient demonstration examples in the training set. With the proposed unlimited-lenglh memory augmentation, LONGMEM can overcome the limitation of the number of demonstration examples the local context and even attend on the whole training set by loading 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] 13], MPQA [WWC05], MR [ABK+07 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 fulfill the 1k input length and provide sufficient contextual self-supervisioons. We transform the k-shot examples to semantically meaningful demonstration examples via fixed template, i.c.. d,="Review: x, Sentiment: y, (I,, }/ E Dirain for sentiment analysis tasks. Additionally, evaluate the 3-shot ICL on gresstoo-answerte 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- ane-grained labels from cached memory, perform
almost
using SQuAD [RZLL16] under an open-ended

---

In-Context #Demons.
In-Memory #Demons.
SST-2
MR
Subj
SST-5
MPQA
Model
Avg.
ACC1
ACC1
ACC1
ACC1
ACC1
Majority
N/A
50.9
50.0
50.0
20.0
50.0
44.2
GPT-2*
4
N/A
68.311.8
64.712.5
51.94.2
31.44.4
61.511.8
55.6
MemTRM
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.8p.8
51.51.5
31.86.7
63.612.9
56.0
LONGMEM
4
2000
71.8
65.111.0
53.83.7
36.0g.s
65,4
58.4
GPT-2*
20
N/A
68.211.5
63.4a.2
57.610.2
33.66.0
70.87.6
58.7
MemTRM
20
2000
65.19.G
65.19.a
58.210.6
31.96.3
72.77.4
58.6
TRIME
20
2000
74.313.9
71.52.5
57.5
33.04.6
69.87.8
61.1
LONGMEM
2000
78.0
78.63.3
65.68.5
36.57.5
74.67.3
66.7
Table 5: Accuracy [%] of 4-shot and 20-shot ICL on 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 runs. Avg. refers the average accuracy on 5 datasets.
a hyperparameter selection on the validation of SST-2, and the best chunk-size is used to report the results for MemTRM, TRIME and our model.
The results 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 improve- ments on the 4-shot case. Additionally, LONGMEM im- proves the in-context learning capabilities of LLMs on open-ended generation tasks, with +4.5 EM score in- crease 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 in-memory augmented demonstrations, thereby achieving superior in- context learming capabilities.
Results.
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 1 to- kens) in-context learning on SQuAD. LoNGMEM loads 200 extra demonstra- tion 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, would like investigate the extend 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 evaluate the effects of those choices.
Effects of Long- Ter-Term Memory Augmentation. To evaluate the effects and contributions of memory augmentations, set the memory-size to 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 L.MM and SideNet only gains 59.4 average scores on ICL NLU tasks, which is 7.3 average accuracy decrease due to the removal of memory augmentation.
Effeets of Chunk-Size. As analyzed before, the thunk-size csz controls the granularity of retrieval and thus it may make difference to tasks with requirements of fine-grained retrieval. We perform an ablation study on the effects of various chunk-size choices cs2 E {2, 4, 8} for in-context learning in-context
and the results are presented in 4(a). The chunk size of 2 yields the best performance 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 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,

---

60
Accurary
Good
mumbic
mardzc
me-11c
mec-th
1-
10-100K
0.1-0
Stllaor PG-22
(b)
Figure 4: (a) Accuracy on 5 NLU datasets given different chunk size during inference; (b) Apepllexity on splits of PG-22 given different memory size during inference, in which the perplexity when ms==65k is used as baseline.
ie., 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 E {8k, 16k, 32k, 65k) during the inference stage 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.
4 Related Work
Large Language Models. Large Language Models, i.e., GPT-3 [BMR= 20J, LLAMA [TMS+ GPT-4 [Ope23], significantly revolutionized NLP research and promoted the state-of-the-ard of various language understanding, language generation [WZG 22]. and even usion-language tasks [WDC* 22]. Additionally, enabled by multi-task instruction tuning [WBZ*2 21, OWJ 22], LLMs exhibit "emergent abilities" [WTB+t (TTB+222 222]
like mathematical reasoning [WWS+22) code completion [CTJ 21], etc.
x-formers. To enable transformers to attend on longer context, many variants of "x-formers" are pro- posed. Transformer-X X [DYY+19 19] proposes to cache attention keys and values of past segment and reuse them in recurrent manner. Recent seminal works of x x-formers, including LinFormer [WLK* 20], LongFormer [BPC20], Routing Transformer [RSVG21], proposed various sparse attention mecha- nisms for decreasing (nn)) complexity to (n log n.) or even Onn). BigBird 20G achieves 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 se- quences that spans book-level length. Moreover, the largest sequence length of these methods is still bounded by 16k tokens, making them invalid in modeling ong-sequences at the book or wikipedia page level (i.e., average 70k tokens for full-length books in PG19 dataset [RPJL19]).
Side- TuninT The method of Side- Side-Tuning [ZSZ 20, SCB22] is task-specific tuning method for are-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 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-netwond residual connections introduced here are novel and distinct from the vanilla summation used in Side-Tuning.
Conclusion
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- contextul language modeling datasets demonstrate effectiveness of our model over other memory-auge memory-augmentattin 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.
10

---

Acknowledgement
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 (HR001 122C0009.. 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.
References
[ABK*0T] Sören Auer, Christian Bizer, Georgi Kobilarov, Jens Lehmann, Richard Cyganiak, and Zachary Ives. Dbpedia: A mucleus for a web of open data, In The semantic web, pages 722-735. Springer, 2007.
[BMH*21] Sebastian Borgeaud, Arthur Mensch, Jordan Hoffmann, Trevor Cai, Eliza Rutherford, Katic Millican, George van den Driessche, Jean-Baptinn Lespiau, Bogdan Damoc, Aidan Clark, Diego de Las Casas, Aurelia Guy, Jacob Menick, Roman Ring, T. W. Hennigan, Saffron Huang, Lorenzo Maggiore, Chris Jones, Albin Cassirer, Andy Brock, Michela Paganini, Geoffrey Irving. Oriol Vinyals, Simon Osindero, Karen Simonyan, Jack W. Rae, Erich Elsen, and L. Sifre. Improving language models by retrieving from trillions of tokens. ArXiv, abs/21 12.0442,, 2021.
[BMR Tom B. Brown, Benjamin Mann, Nick Ryder, Melanic Subbiah, Jared Kaplan, Prafulla Dhariwal, Arvind Neclakantan, Pranav Shyam, Girish Sastry, Amanda Askell, Sandhini Agarwal, Ariel Herbert-Voss, Gretchen Krueger, T. J. Henighan, Rewon Child, Aditya Ramesh, Daniel M. Ziegler, Jeff 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. ArXiv, abs/2005. 14165, 2020.
[BPC20] Iz. Beltagy, Matthew Peters, and Arman Cohan. Longformer: The long-document transformer. arXiv preprint 200:v:2004.05 2020.
[CT Mark Chen, Jerry Tworek, Heewoo Jun, Qiming Yuan, Henrique Ponde de Oliveira Pinto, Jared Kaplan, Harri Edwards, Yuri Burda, Nicholas Joseph, Greg Brockman, et al. Evaluating large language models trained code. arXiv preprint arXiv:2107.03374 2021. 2021.
[DCLT19] Jacob Devlin, Ming- Wei Chang, Kenton Lee, and Kristina Toutanova. BERT: Pre-training of deep bidirectional transformers for language understanding. In NAACL, 2019.
[DYY*19] Zihang Dai, Zhilin Yang. Yiming Yang, Jaime Carbonell, Quoe V Le, and Ruslan Salakhutdi- nov. Transformer-xil xl: Attentive language models beyond fixed-length context. arXiv preprint arXiv:/901. 002860, 2019.
[GBB+20] Leo Gao, Stella Biderman, Sid Black, Laurence Golding, Travis Hoppe, Charles Foster, Jason Phang, Horace He, Anish Thite, Noa Nabeshima, et al. The pile: An 800gb dataset of diverse text for language modeling. arXiv preprint (0X.x:2111.00000 2020.
[JD221] Ieff Johnson, Matthijs Douze, and Hervé Jégou. Billion-scale similarity search with gpus. IEEE Transactions on Big Data, 7:535-547, 2021.
[KB15] Diederik P. Kingma and Jimmy Ba. Adam: A method for stochastic optimization. CoRR, abs/ 0bs/1412.6980, 2015.
[33 Rui Lv, Junliang Guo, Rui Wang, Xu Tan, Qi Liu, and Tao Qin. N-gram nearest neighbor machine translation. arXiv preprint Xrxiv:2301. /2866, 2023.
[LOG 19] Yinhan Liu, Myle Oit, Naman Goyal, Jingfei Du, Mandar Joshi, Danqi Chen, Omer Levy, Mike Lewis, Luke Zettlemoyer, and Veselin Stoyanov. RoBERTa: robustly optimized bert pretraining approach. ArXiv, abs//190.. 16692, 2019.
[Ope23] OpenAL. Gpt-4. https: //tp:///.pp/////////e/e wtt..//eeeeeeeeeeeeeeeeeeee gpt-4, 2023. Accessed on March 14, 2023.
[2W]*22] Long Ouyang, Jeffrey Wu, Xu Jiang, Diogo Almeida, Carroll Wainwright, Pamela Mishkin, Chong Zhang, Sandhini Agarwal, Katarina Slama, Alex Ray, et al. Training language models to follow instructions with human feedback. Advances in Neural Information Processing Systems. 01227730-27744 2022.

---

[PL04] Bo Pang and Lillian Lee. A sentimental education: Sentiment analysis using subjectivity summa-- rization based minimum cuts. arXiv preprint cx/0449058.. 2004.
[PSL21] Ofir Press, Noah Smith, and Mike Lewis. Train short, test long: Attention with linear biases enables input length extrapolation. arXiv preprint arXiv:2 12409, 2021.
[RNSS18] Alec Radford, Karthik Narasimhan, Tim Salimans, and Ilya Sutskever. Improving language understanding with unsupervised learning. 2018
[RPJL19] Jack W Rae, Anna Potapenko. Siddhant M Jayakumar, and Timothy P Lillicrap. Compressive transformers for long- range sequence modelling, arXiv preprint arXiv:1911.05507 2019.
[RSR 20] Colin Raffel, Noam M. Shazeer, Adam Roherts, Katherine Lee, Sharan Narang, Michael Matena, Yanqi Zhou, Wei Li, and Peter Liu. Exploring the limits of transfer learning with unified text-to-text transformer. ArXiv, abs/1910.10683. 2020.
[RSVG21] Aurko Roy, Mohammad Saffar, Ashish Vaswani, and David Grangier. Efficient content- content-based sparse attention with routing transformers. Transactions of the Association for Computational Linguistics, 9:53-68, 2021.
[RWC 19] Alec Radford, Jeff Wu, Rewon Child, David Luan, Dario Amodei, and Ilya Sutskever. Language models are unsupervised multitask learners. 2019.
[RZLL16] Pranav Rajpurkar, Jian Zhang, Konstantin Lopyrev, and Percy Liang. SQuAD: 100,000+ Questions for Machine Comprehension of Text. arXiv e-prints, page arXiv: 11..116066.55 05250, 2016.
[SCB22] Yi-Lin Sung, Jaemin Cho, and Mohit Bansal. Lst: Ladder side-tuning for parameter and memory efficient transfer learning. arXiv preprint arXiv:2206. 06522, 2022.
[3) Richard Socher, Alex Perelygin, Jean Wu, Jason Chuang, Christopher D Manning, Andrew Y Ng, and Christopher Potts. Recursive deep models for semantic compositionality over sentiment treebank. In Proceedings of the 2013 conference on empirical methods in natural language processing, pages 1631-1642, 2013.
[ST122] Simeng Sun, Katherine Thai, and Mohit lyyer. Chapterbreak: A challenge dataset for Long-range language models. arXiv preprint arXiv:2204. 10878, 2022.
[TMS+23] Hugo Touvron, Louis Martin, Kevin Stone, Peter Albert. Amjad Almahairi, Yasmine Babaei. Nikolay Bashlykov, Soumya Batra, Prajjwal Bhargava, Shruti Bhosale, et al. Llama Open foundation and fine-tuned chat models. arXiv preprint arXiv:2307.09228 09988, 2023.
[VSP* 17] Ashish Vaswani, Noam M. Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N. Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. In NIPS, 2017.
[WBZ+21] Jason Wei, Maarten Bosma, Vincent Zhao, Kelvin Guu, Adams Wei Yu, Brian Lester, Nan Du, Andrew M Dai, and Quoe Le. Finetuned language models are zero-shot learners, International Conference on Learning Representations, 2021.
[WDC+22] 22) Weizhi Wang, Li Dong, Hao Cheng, Haoyu Song, Xiaodong Liu, Xifeng Yan, Jianfeng Gao, and Furu Wei. Visually-augmented language modeling. arXiv preprint arXiv:2205.10178 2022.
[WLK+20] Sinong Wang, Belinda Z Li, Madian Khabsa, Han Fang, and Hao Ma. Linformer: Self-attention with linear complexity. arXiv preprint arXiv:2006.04760 2020.
[WRHS22] Yuhuai Wu, Markus N. Rabe, DeLesley Hutchins, and Christian Szegedy. Memorizing trans- formers. ArXiv, abs/2203.08913, 2022.
[WTB+22] Jason Wei, Yi Tay, Rishi Bommasani, Colin Raffel, Barret Zoph, Sebastian Borgeaud, Dani Yogatama, Maarten Bosma, Denny 7hou, Donald Metzler, et al. Emergent abilities of large language models. arXiv preprint arXiv:2206. 076822 2022.
[WWC05] Janyce Wiebe, Theresa Wilson, and Claire Cardie. Annotating expressions of opinions and emotions in language. Language resources and evaluation, 39(2):165 210, 2005.
[WWS*22] 22] Jason Wei, Xuezhi Wang, Dale Schuurmans, Maarten Bosma, Ed Chi, Quoe Le, and Denny Zhou. Chain of thought prompting elicits reasoning in large language models. arXiv preprint arXiv:2201.11903 2022.
12

---

[WZG+22] Weizhi Wang, Zhirui Zhang, Junliang Guo, Yinpei Dai, Boxing Chen, and Weihua Luo. Task- oriented dialogue system natural language generation. In Proceedings the 45th International ACM SIGIR Conference on Research and Development Information Retrieval, pages 2698-2703, 2022.
[YDY* 19] Zhilin Yang, Zihang Dai, Yiming Yang. Jaime G. Carbonell, Ruslan Salakhutdinov, and Quoe Le. XLNet: Generalized autoregressive pretraining for language understanding. In NeurlPS, 2019.
[ZGD 20] Manzil Zaheer, Guru Guruganesh, Kumar Avinava Dubey, Joshua Ainslie, Chris Alberti, Santiago Ontanon, Philip Pham, Anirudh Ravula, Qifan Wang, Li Yang, al. Big bird: Transformers for longer sequences. Advances neural information processing systems, 33:17283- 17297, 2020.
[ZLC22]
Zexuan Zhong,
Danqi Chen. Training language models with memory augmentation. arXiv preprint arXiv:2005. 12674, 2022.
[0S Jeffrey Zhang, Alexander Sax, Amir Zamir, Leonidas Guibas, and Jitendra Malik. Side-tuning: baseline for network adaptation via additive side networks. In Computer Vision-ECCC 2020: 16th European Conference, Glasgow, UK, August 23 23-28, 2020, Proceedings, Part III 16, pages 698 714, Springer, 2020.
A
Ablation Study on the Effect of Memory Augmentation
In-Context #Demons.
In-Memory #Demons.
SST-2 ACC1
MR ACC1
Subj
SST-5
MPQA
Model
Avg.
ACC1
ACC1
ACC1
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.7
51.94.2
31,44.4
61.511.8
55.6
MenTRM
4
2000
67.512.4
64.611.3
53.26.0
29.64.4
63.012.1
55.6
TRIME
4
2000
69.5
63.8g.8
51.51.6
31.86.7
63.612.9
56.0
LONGMEM
4
2000
71.814.0
65.111.0
53.83.7
36.0a.8
65.412.8
58.4
w/o Memory
4
0
69.412.4
64.312.1
53.47.7
29.0s.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.315.9
71.52.5
57.5
33.04.6
69.87.8
61.1
LONGMEM
20
2000
78.014.1
78.63.3
65.68.D
36.57.o
74.67.3
66.7
w/o Memory
20
70.012.8
70.8p.2
52.94.6
30.96.4
72.57.5
59.4
Table 6: Ablation study results the effect of memory augmentation of 4-shot and 20-shot ICL NLU tasks (SST-2, mr. subj. SST-5, mpqa). We sample 2000 extra demonstration examples and load them into cached memory. The subscript the standard deviation across 6 runs. Avg. refers the average accuracy on 5 datasets. "w/o" is short for "without".
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-bar based models. The detailed statisties in terms of the efficiency is presented in Table
In-Context Len.
In-Memory Len.
Inference Speed (tokens/s))
GPU-Memory Usage
Model
(MBs).
GPT-2*
4k
N/A
14666
20671
LONGMEM
lk
3k
22638
13335
GPT-2*
8k
N/A
8417
54195 13437
LONGMEM
lk
7k
21343
Table 7: The superiority of our method over fully dense self-attention (GPT-2*) in terms of inference speed and GPU-memory utilization.
13

---

C
Training Details
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- aammeed-augmented adaptation training. The pre-training and adaptation are trained on 16 22GB-Tesle V100 GPUs. Other detailed training pyppeaaaramtess and settings are presented in Table 8.
Hyperparameter
LONGMEM
Reproduced GPT-22 Backbone LLM Hyperparameters
Parameters
407M
Precision
float16
Layers
Hidden dim.
1024
Attention heads
16
Head Dim
64
Vocab size
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
Sequence length
1024
Mlemory-Augmented Adaptation Hyperparameters
Global Batch Size
256
Learning rate
2.0c-4
Total tokens
26B
Warmup tokens
LR Decay style
polynomial
Adam (31, (32)
(0.9, 0.98)
Adam eps
le-06
Weight decay
0.01
Gradient clipping
Table 8: Aegore-Augmented Adaptation and Architectural Hyperparameters.
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.
Task
Prompt
Labels
SST-2
Review: [Sentence] Sentiment: [Label]
(positive, negative
MR
Review: [Sentence] Sentiment: [Tabel]]
(positive, negative)
MPQA
Review: [Sentence] Sentiment: [Label]
(positive, negative]
SST-5
input: [Sentenee] type: [Label]
(terrbblebball.keyggooo.grgr
Subj
input: [Sentence] type: [Label]
(objective. subjective]
Squad
Passage:
[Passage]] \ Question: [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.
14


Reviews:
[{"review_id": "yA9RnNZ9wS", "reviewer": "Reviewer_yoPY", "summary": "This paper proposes a novel approach in handling long input context using memory framework. To be specific, any inputs that are longer than the maximum input size are cached in memory network (with the outcome of forward pass of LLM), and those cached memories are retrieved when \"current\" input is processed. The LLM remains fixed throughout training, hence being efficient; the trainable parameters are the fusion network and SideNet layers. The authors illustrate the effectiveness of the proposed approach on long-context benchmarks.", "questions": "Please refer to the Weakness", "limitations": "Please refer to the Weakness", "rating": 7, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- This paper proposes to a novel approach in handling long-context problem that transformer-based language models suffer from. \n\n- The paper evaluates the proposed idea on several benchmarks, and illustrate significant gain, especially on NLU tasks.\n\n- The idea is sounding, and the proposed approach conducts thorough ablation studies on modules introduced. \n", "weaknesses": "- some parts of the paper need proper citations. i.e. Line 63-65: \"directly adapting the entire ~~~suffers from catastrophic forgetting.\" needs a proper citation for the claim.\n\n- There are baselines that could be compared against in long-context language modeling tasks. For example, in Table 2, the authors compared against MemTRM and GPT2. LocalTRM is one of the baselines that the authors can compare to in language modeling task as well, as LocalTRM has similar number of parameters and has the context length of 8K. \n\n- It would be interested to see the inference speed of the proposed method, compared to the other baselines. It would be insightful to check how much of benefit is gained at how much expense of inference speed and computation cost. \n\n- Much performance gains are seen over the NLU tasks, yet the gains are somewhat trivial on language modeling benchmarks (PPL).", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank you a lot for your detailed feedback and suggestions!\n\nQ1: Citation for the claim for catastrophic forgetting issue\n\nWe will add the paper G-MAP as a citation for this claim.\n \nWan, Zhongwei, Yichun Yin, Wei Zhang, Jiaxin Shi, Lifeng Shang, Guangyong Chen, Xin Jiang, and Qun Liu. \"G-MAP: General Memory-Augmented Pre-trained Language Model for Domain Tasks.\" In Proceedings of the 2022 Conference on Empirical Methods in Natural Language Processing, pp. 6585-6597. 2022.\n \nQ2: LocalTRM on Language Modeling tasks\n \nThe LocalTRM adopts the relative positional bias as the positional embedding method to achieve length extrapolation. When LocalTRM is evaluated on language modeling tasks, it will down-grade to our reproduced GPT-2* which is also pre-trained language model with alibi relative positional bias.\n \n \nQ3: Inference speed and computational cost\n \nWhen the model is required to comprehend long sequences, the proposed method DEMA 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.\nThe detailed statistics in terms of the efficiency is presented in the following Table.\n \n| Model | In-Context Length | In-Memory Length | Inference Speed (tokens/s) | GPU-Memory Usage (MBs) |\n|--------------------------------|---------------------|--------------------|--------------------------|---------------------------|\n| GPT-2*                     \t| 4k              \t| N/A            \t| 14666                \t| 20671                 \t|\n| DEMA                     \t| 1k              \t| 3k             \t| 22638                \t| 13335                 \t|\n| GPT-2*                     \t| 8k              \t| N/A          \t  | 8417                 \t| 54195                 \t|\n| DEMA                    \t| 1k              \t| 7k             \t| 21343                \t| 13437                 \t|\n \nQ4: Perplexity Gains\n \nWe believe that the gains achieved by the proposed method on language modeling benchmarks are neither trivial nor subtle. The 1.38∼1.62 perplexity decreases are remarkable performance improvements on language modeling benchmarks. Normally, over 1.0 perplexity improvement can be regarded as strong demonstrations for the effectiveness of the method, which are also recognized by other significant works like Memorizing Transformer, KNN-LM, etc.\n"}, {"review_id": "GfmCgMS6Tw", "reviewer": "Reviewer_k2cK", "summary": "To address the limitation of existing large language models (LLMs), the authors introduce a framework called Decoupled-Memory-Augmented LLMs (DEMA) that utilizes long-context information. DEMA utilizes a decoupled network architecture with a frozen LLM as a memory encoder and an adaptive residual side-network as a memory retriever and reader. Experimental results show that DEMA outperforms strong long-context models and achieves significant improvements in memory-augmented in-context learning over LLMs.", "questions": "In line 111, you mention ‘total # layers’. Can you explain the meaning of the \"#\" symbol in this context?", "limitations": "None", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1.The idea of this work is novel and the motivation is strong. It tries to store and utilize long-form previous context or knowledge to address the input length limit of LLMs.\n\n2.The main experiments seem solid as this paper verifies their method on multiple datasets and the experimental results validate the effectiveness in helping LLMs to memorize and utilize long-form information.\n\n3.This paper is well-written and easy to follow. Their main claim is clear and easy to understand.", "weaknesses": "1.The paper introduces a novel lightweight residual SideNet, but there is a lack of discussion on the contribution of the SideNet module to the DEMA model in the ablation experiments.\n\n2.In Section 3.4, the ablation experiments show that Figure 2(b) does not provide a clear visualization of the impact of different memory sizes on perplexity, especially when msz=32k (the orange line) and msz=16k (the green line) are very close to each other. This lack of distinction makes it less evident to draw a conclusive result that suggests better performance when msz=16k, as mentioned in the paper.\n\n3.The baselines compared in the paper lack recent advancements, such as GPT-2 in 2019, GPT-3 and Bigbird in 2020, LocalTRM and RoutTRM in 2021, and MemTRM in March 2022.\n\n4.The experiments include only main results (comparing with baselines) and the ablation studies. It would be better to conduct more detailed analyses or qualitative analysis to verify the proposed methods (e.g. whether the issues proposed in the motivation are fixed).\n\n5.Related work section should discuss other memory augmented methods, especially the memory augment text generation methods.\n\n6.The inputs section in Figure 1 is ambiguous and confusing. It is unclear whether the letters a, b, ..., h are meant to represent individual letters or tokens. If they represent individual letters, then combining them does not form a word. However, if they refer to tokens, it is advisable to use more precise symbols that represent tokens to avoid any ambiguity.\n\n7.Figure2(a): sst-5 should be changed to SST-5, consistent with Table 5.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank you a lot for your detailed feedback and suggestions!\n \nQ1: Contribution of SideNet module\n \nIf the SideNet module is removed from the architecture, the model will down-grade to the architecture similar to Memorizing Transformer, which is used as the significant baseline to be compared with. Therefore, the Memorizing Transformer can be regarded as the ablation model without SideNet and the performance gap can demonstrate the contribution of proposed SideNet module.\n \nQ2: Ablation study on memory size\n \nWe revised the y-axis of Figure 2(b) as the delta perplexity to show the ablation results and the new Figure is presented in the one-page rebuttal PDF document. We will revise such description on the selection of superiority of msz=16k as the hyperparameters of msz=16k and 32k present similar superiority over other hyperparameters.\n \nQ3: Other up-to-date baselines\n \nWe follow the suggestions from other reviewers to add another up-to-date strong memory-augmented competitor TRIME-long[1] as the baseline and the results on memory-augmented in-context learning are shown below:\n \n| Model |   In-context #Demons |  SST-2 | MR  |  Subj   |  SST-5  | MPQA | Avg |\n|   ----  |----  | ----  | ----  | ----  | ----  | ----  | ----  |\n| GPT-2*   | 4        \t| $68.3_{11.6}$  | $64.7_{12.5}$  | $51.9_{4.2}$  | $31.4_{4.4}$      \t| $61.5_{11.8}$      \t| 55.6      \t|\n| MemTRM   |    4            |   $67.5_{12.4}$      \t| $64.6_{11.3}$      \t| $53.2_{6.0}$      \t| $29.6_{4.4}$      \t| $63.0_{12.1}$      \t| 55.6      \t|\n| TRIME | 4        \t|\t$69.5_{14.5}$  | $63.8_{9.8}$ | $51.5_{1.5}$ | $31.8_{6.7}$ | $63.6_{12.9}$ |  56.0\n| DEMA   | 4        \t| $71.8_{14.0}$   | $65.1_{11.0}$ | $53.8_{3.7}$ | $36.0_{6.8}$ | $65.4_{12.8}$ | 58.4  |\n| GPT-2*   | 20       \t| $68.2_{11.5}$      \t| $63.4_{5.2}$       \t| $57.6_{10.2}$     \t| $33.6_{6.0}$      \t| $70.8_{7.6}$       \t| 58.7      \t|\n| MemTRM   | 20 | $65.1_{9.6}$       \t| $65.1_{9.3}$       \t| $58.2_{10.6}$     \t| $31.9_{6.3}$    | $72.7_{7.4}$ | 58.6 |\n| TRIME   | 20  |  $74.3_{13.9}$   |  $71.5_{2.5}$  |  $57.5_{11.4}$ |  $33.0_{4.6}$ | $ 69.8_{7.8}$ | 61.1  |\n| DEMA   | 20  |  $78.0_{14.1}$ | $78.6_{3.3}$  | $65.6_{8.5}$ | $36.5_{7.5}$ | $74.6_{7.3}$  | 66.7 |\n \n[1] Zhong, Zexuan, Tao Lei, and Danqi Chen. \"Training Language Models with Memory Augmentation.\" In Proceedings of the 2022 Conference on Empirical Methods in Natural Language Processing, pp. 5657-5673. 2022.\n \n \nQ4: Additional Analysis\n \nWe perform an analysis on evaluating the memory staleness issue claimed in our motivation. We take a document in the PG22 corpus as the input to Memorizing Transformer training checkpoints at 1000th steps and 1100th steps and then the average cosine similarity between the cached attention values are computed to evaluate the Memory Stableness. The cosine similarity for Memorizing Transformer is 0.81 and our proposed method is 1 because the memory encoder is frozen in our proposed model.\n \nQ6: Notation Issues\n \nThe characters in Figure 1 refer to different tokens in an input sequence. We will follow your advice to revise such character symbols to actual texts to avoid the ambiguity. Additionally, the SST-5 notation in FIgure 2(a) will also be revised accordingly. At last, in line 111, he meaning of the \"#\" symbol is “the number of” and we will replace the symbol with textual description.\n"}, {"review_id": "83iMWeSJAv", "reviewer": "Reviewer_yw5j", "summary": "This paper proposes a side network as a memory retriever and reader to help the LLMs model long context input. Specifically, the key-value pairs of the history inputs are saved in the cached memory bank. The side network takes in the current hidden states and the retrieved past key-value pairs to compute the memory-augmented representations. Their approach achieves better performance in perplexity and accuracy on multiple datasets (i.e., Project Gutenberg 2020-2022 Language Modeling Dataset, Arxiv Dataset, ChapterBreak Benchmark). ", "questions": "na", "limitations": "na", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. The side network can retrieve and fuse memory-augmented long-context, which enables the LLM to model context with ultimate length. \n2. Experiments demonstrate the effectiveness of the proposed approach on long context datasets like Arxiv. Besides, this paper shows that LLM’s ability to model long context better could help both in context learning, and zero-shot performance. ", "weaknesses": "1. The evaluation mainly focuses on perplexity score, and accuracy of suffix identification. To evaluate whether LLM models the long context well and thus benefit its’ generation ability, tasks like summarization/title generation should be included. \n2. Missing comparison with Efficient Long-Text Understanding with Short-Text Models, which could take maximum 16K tokens as input. ", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We would like to thank you for your time and constructive suggestions.\n \nQ1: Evaluation on language generation tasks\n \nThank you for the great suggestion on verifying the proposed method on language generation tasks like summarization. As the proposed DEMA model is a pre-trained language model augmented with long-term memory, the designed evaluations focus more on the emerging capabilities enabled by proposed memory-augmented pre-training, especially in-context learning or zero-shot learning. Therefore, we evaluate the proposed method on both natural language understanding and natural language generation tasks on memory-augmented in-context learning settings. The performance improvement on Squad in Table 4 can demonstrate the effectiveness of the method on language generation tasks. The document summarization or reading comprehension tasks usually require full-parameter model fine-tuning on the knowledge domain of the dataset and thus not a good fit for evaluating the capability of pre-trained language models. But we will later tune the proposed method on BookSum dataset for evaluating such domain-specific adaptation capability.\n \nQ2: Comparison with efficient long-text understanding with short-text models\n \nThe baselines of LocalTRM and RoutTRM presented in Table 2 can be regarded as this type of models. And they are compared with the proposed method on long-text understanding benchmarks of ChapterBreak.\n"}, {"review_id": "lVUV6LTZTn", "reviewer": "Reviewer_Mats", "summary": "The paper proposes DeMA as a learning-based adapter for existing Language Models (LMs) to increase long-term memory and performance on long inputs.", "questions": "- L205-206 “Then, we concatenate shuffled documents within one group and truncate them into ordered segments.” Is it not clear to me the procedure to batch the input. I also read the appendix, but it is still not clear. Could you please elaborate a little bit on this aspect?\n- Why did you test your method just on the AO3 split of Chapterbreak and not on the PJ19 split?", "limitations": " ", "rating": 5, "confidence": 4, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "The method proposed is clear and sound and it could potentially be useful to the community since using transformers with a long context is an important open problem.", "weaknesses": "Experiments that prove the advantage of DeMA are limited:\n\n- Most of the comparisons are limited to just a single memory-augmented model (MemTRM) and GPT-2 (that does not have any memory mechanism). Table 3 shows comparisons with several memory-augmented models but just on a single task and with old memory-augmented models.\n- The experiments do not show comparisons with several recent memory-augmented competitors:\n    - “Memformer: A Memory-Augmented Transformer for Sequence Modeling” (Wu et al.). Here a comparison on WikiText-103 might have provided additional evidence of the superiority of DeMA.\n    - “Unlimiformer: Long-Range Transformers with Unlimited Length Input” (Bertsch et al.)\n    - “Training Language Models with Memory Augmentation” (Zhong et al.)\n    - “RWKV: Reinventing RNNs for the Transformer Era” (Peng et al.)\n- I’m not convinced by the memory-augmented In-context learning experiments (Table 5). What is the purpose of this experiment? Is it really a fair “few-shot” comparison since DeMA is using additional 2000 demonstration examples loaded into cached memory? I think a comparison with other memory-augmented models will leave any doubt.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank you a lot for your detailed feedback and suggestions!\n \nQ1: Comparisons with suggested baselines\n \nThank you so much for providing a series of competitors for comparisons. As Unlimitedformer is released within 10 days to the NeurIPS 2023 paper submission deadline, normally according to paper submission guidelines and our practical difficulties, it will not be considered as a baseline to be compared with. RWKV is a great work but it is released after the deadline of paper submission and thus impossible to be mentioned in the paper. The first pre-print version of Memformer is released in 2020 and the successor work like Memorizing Transformer used as the baseline in our paper has demonstrated more superiority as memory-augmented architectures. Additionally, we appreciate the suggestion on evaluating the method on Wikitext-103 dataset but the dataset is included as the subset of the Pile corpus used as the pre-training corpus of our work and RWKV. This is also the reason why later works pre-trained on Pile corpus did not perform evaluations on Wikitext-103.\n \nThe TRIME is both a great competitor with our work and a well-recognized memory-augmented architecture. We reproduced TRIME as a memory-augmented adaptation method with the same hyperparameters of external memory as DEMA and evaluated the TRIME-long pre-trained model on the memory-augmented in-context learning tasks. The results are shown below. The TRIME still significantly lags behind the proposed method when the same number of memory-augmented demonstration examples are cached into the external memory of TRIME.\n \n| Model |   In-context #Demons |  SST-2 | MR  |  Subj   |  SST-5  | MPQA | Avg |\n|   ----  |----  | ----  | ----  | ----  | ----  | ----  | ----  |\n| GPT-2*   | 4        \t| $68.3_{11.6}$  | $64.7_{12.5}$  | $51.9_{4.2}$  | $31.4_{4.4}$      \t| $61.5_{11.8}$      \t| 55.6      \t|\n| MemTRM   |    4            |   $67.5_{12.4}$      \t| $64.6_{11.3}$      \t| $53.2_{6.0}$      \t| $29.6_{4.4}$      \t| $63.0_{12.1}$      \t| 55.6      \t|\n| TRIME | 4        \t|    $69.5_{14.5}$  | $63.8_{9.8}$ | $51.5_{1.5}$ | $31.8_{6.7}$ | $63.6_{12.9}$ |  56.0\n| DEMA   | 4        \t| $71.8_{14.0}$   | $65.1_{11.0}$ | $53.8_{3.7}$ | $36.0_{6.8}$ | $65.4_{12.8}$ | 58.4  |\n| GPT-2*   | 20       \t| $68.2_{11.5}$      \t| $63.4_{5.2}$       \t| $57.6_{10.2}$     \t| $33.6_{6.0}$      \t| $70.8_{7.6}$       \t| 58.7      \t|\n| MemTRM   | 20 | $65.1_{9.6}$       \t| $65.1_{9.3}$       \t| $58.2_{10.6}$     \t| $31.9_{6.3}$    | $72.7_{7.4}$ | 58.6 |\n| TRIME   | 20  |  $74.3_{13.9}$   |  $71.5_{2.5}$  |  $57.5_{11.4}$ |  $33.0_{4.6}$ | $ 69.8_{7.8}$ | 61.1  |\n| DEMA   | 20  |  $78.0_{14.1}$ | $78.6_{3.3}$  | $65.6_{8.5}$ | $36.5_{7.5}$ | $74.6_{7.3}$  | 66.7 |\n \n \n \n \nQ2: Experiments on memory-augmented in-context learning\n \nThe context length of language models is usually limited to 1k-2k tokens and thus the demonstration examples used in the input prompt are limited for language models to acquire enough task representations to perform well. With the capability to utilize the demonstration examples in cached memory bank, the language models trained with our method can take advantage of even 2000 demonstration examples to gain great task representation and achieve much better performance on downstream tasks via in-context learning rather than fine-tuning. This is very crucial to the large-scale application of language models. With the augmented 2000 demonstrations, the evaluation setting is more appropriate to be described as memory-augmented many-shot in-context learning. Another significant memory-augmented model, Memorizing Transformer is also compared with our method. The results listed in Table 3 can demonstrate the superiority of our method over other memory-augmented architectures.\n \nQ3: Explanation on document batchfying process\n \nNormally, the long pre-training corpus will be truncated into fix-length segments and all segments will be randomly shuffled during training as there is no correlations between the segments in different batches. However, our proposed method requires that the i-th segment in the batch is exactly the previous context of the i-th segment in the next batch. In this way, the contextual segments in previous batches are already cached in memory bank and will benefit the language modeling towards the consecutive segments in the same document.\n \nQ4: Evaluation on only AO3 split\n \nThe ChapterBreak benchmark involves two splits, the PG-19 and the AO3. The whole PG-19 corpus has been included in the pre-training corpus Pile of our model and thus cannot be deployed as the testset or validation set. The evaluation on AO3 split is appropriate to evaluate the generalization capability and zero-hot learning capability of the pre-trained language models.\n \n \n \n"}, {"review_id": "NDgUIHlslH", "reviewer": "Reviewer_UtJk", "summary": "This paper proposes an architectural extension to a frozen LLM that allows for incorporation of external memory caches (also constructed with the frozen LLM), allow the model to handle extremely long contexts.", "questions": "- I believe GPT-2* refers to the reproduced GPT-2, but I don't think this is mentioned in the text.\n- Perplexity should be a positive number. Are the authors referring to negative log loss instead?\n- I would like to see experiments on with inference DEMA but with no memory cache (if I understand the model architecture correctly, this would still involve the SideNet parameters, which would differ from GPT-2*). This should apply to all current experiments.", "limitations": "N/A", "rating": 7, "confidence": 4, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "The approach presented is sensible (although I am not sufficiently familiar with the latest long-context modeling literature to determine if it is novel), and the experiments are convincing. In particular, I think the few-shot learning results are very strong.", "weaknesses": "- The explanation of the model architecture is fairly unclear to me, and I found it difficult to match the description of the architecture and the presented diagram. I recommend the authors rewrite the explanation of the model and obtaining feedback from a fresh reader for suggestions on improving the clarity. Alternatively, presenting the architecture in terms of pseudo-code may also be clearer.\n- The evaluations presented are largely newly proposed (besides ChapterBreak). They appear to be sensibly constructed, but it does mean the results are slightly difficult to compare to or contextualize to other works. I would be interested in seeing some experiments on more established tasks (e.g. long document summarization).", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank you a lot for your detailed feedback and suggestions!\n\n\nQ1: Issues of presentation\n \nThank you for the kind suggestion on optimizing the presentation. We follow your advice and collect some reader feedback. Then we decompose the main figure into two figures for better presentation. The first figure describes the memory encoding, caching, and retrieval flow of the proposed method. And the second figure illustrates more on the architectural contributions. You can refer to these revised figures in the one-page rebuttal PDF document attached.\n \nQ2: Evaluation on summarization tasks\n \nThank you for the great suggestion on verifying the proposed method on language generation tasks like summarization. As the proposed DEMA model is a pre-trained language model augmented with long-term memory, the designed evaluations focus more on the emerging capabilities enabled by proposed memory-augmented pre-training, especially in-context learning or zero-shot learning. Therefore, we evaluate the proposed method on both natural language understanding and natural language generation tasks on memory-augmented in-context learning settings. The performance improvement on Squad in Table 4 can demonstrate the effectiveness of the method on language generation tasks. The document summarization or reading comprehension tasks usually require full-parameter model fine-tuning on the knowledge domain of the dataset and thus not a good fit for evaluating the capability of pre-trained language models. But we will later tune the proposed method on BookSum dataset for evaluating such domain-specific adaptation capability.\n \nQ3: Notation for GPT-2*\n \nWe will add the notation for GPT-2* as the reproduced GPT-2 in the captions of Table 1-4. The training details for GPT-2* has been presented in Line 215-218.\n \nQ4: Notation for perplexity score\n\nWe apologize that the expression in Line 74 may lead to misunderstanding. The minus symbol used here represents the proposed achieves a decrease on the perplexity score.\n \nQ5: The ablation study on proposed method without memory-augmentation\nWe implement such ablation study via setting the hyperparameter memory-size equal to 0. Intuitively, if there is no long-term memory augmentation, the proposed method may have a performance degradation. We present the results of ablation model DEMA w/o Memory on memory-augmented in-context learning in the following Table.\n \n| Model |   In-context #Demons |  SST-2 | MR  |  Subj   |  SST-5  | MPQA | Avg |\n|   ----  |----  | ----  | ----  | ----  | ----  | ----  | ----  |\n| GPT-2*   | 4        \t| $68.3_{11.6}$  | $64.7_{12.5}$  | $51.9_{4.2}$  | $31.4_{4.4}$      \t| $61.5_{11.8}$      \t|  55.6          |\n| MemTRM   |    4      |   $67.5_{12.4}$      \t| $64.6_{11.3}$      \t| $53.2_{6.0}$      \t| $29.6_{4.4}$      \t| $63.0_{12.1}$      \t|  55.6          |\n| DEMA   | 4        \t| $71.8_{14.0}$   | $65.1_{11.0}$ | $53.8_{3.7}$ | $36.0_{6.8}$ | $65.4_{12.8}$ | 58.4  |\n| DEMA  w/o Memory | 4   | $69.4_{12.4}$   | $64.3_{12.1}$ | $53.4_{7.7}$ | $29.0_{5.2}$ | $62.5_{12.3}$ | 55.7 |\n| GPT-2*   | 20  | $68.2_{11.5}$   | $63.4_{5.2}$\t| $57.6_{10.2}$     \t| $33.6_{6.0}$  | $70.8_{7.6}$  | 58.7      \t|\n| MemTRM   | 20 | $65.1_{9.6}$    |  $65.1_{9.3}$    | $58.2_{10.6}$     \t| $31.9_{6.3}$    | $72.7_{7.4}$ | 58.6 |\n| DEMA   | 20  |  $78.0_{14.1}$ | $78.6_{3.3}$  | $65.6_{8.5}$ | $36.5_{7.5}$ | $74.6_{7.3}$  | 66.7 |\n| DEMA w/o Memory |  20   |  $70.0_{12.8}$   |  $70.8_{6.2}$\t| $52.9_{4.6}$\t| $30.9_{6.4}$ | $72.5_{7.5}$ | 59.4 |\n \n\n"}]

Author rebuttal:
This is the one-page rebuttal PDF document including the revisions on the Figure 1 and Figure 2 in the original paper following the suggestions from Reviewer utjk and Reviewer k2ck. We thank a lot to the detailed feedback and suggestions from all reviewers.

Meta-review: The paper proposes a method to cache the hidden state of the transformer into a memory bank and use the memory to extend the context length of large language models. Reviewers agreed on the comprehensiveness of the experiment setting and the soundness of this method.
Decision: Accept (poster)