SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Weizhi Wang†, Li Dong†, Hao Cheng†, Xiaodong Liu†, Xifeng Yan†, Jianfeng Gao‡, Furu Wei‡
†University of California, Santa Barbara ‡Microsoft Research weizhiwang@ucsb.edu, {lidong1, chehao, xiaodl}@microsoft.com
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. Our code is open-sourced at https://aka.ms/LongMem.
Large language models (LLMs) have revolutionized natural language processing with great successes in advancing the state-of-the-art on various understanding and generation tasks [DCLT19, RWC+19,  $\mathrm{LOG^{+}19}$ ,  $\mathrm{YDY^{+}19}$ ,  $\mathrm{BMR^{+}20}$ ,  $\mathrm{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 input context length. For instance, GPT-3  $\left[\mathrm{BMR}^{+}20\right]$  increases the input length from 1k tokens in GPT-2  $\left[\mathrm{RWC}^{+}19\right]$  to 2k tokens, thereby allowing for better capture of long-range dependencies. However, this approach typically incurs computation-intensive training from scratch and the in-context dense attention is still heavily constrained by the quadratic computation complexity of Transformer self-attention  $\left[\mathrm{VSP}^{+}17\right]$ . Another recent line of work  $\left[\mathrm{BPC}20, \mathrm{ZGD}^{+}20\right]$  instead focuses on developing in-context sparse attention to avoid the quadratic cost of self-attention, which
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- $K$  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-context sparse attention via dense attention over both in-context tokens and memorized tokens retrieved from a 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 staleness challenge during training due to its coupled memory design, which uses a single model for encoding memory and fusing memory for language modeling. In 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 a framework for Language Models Augmented with Long-Term Memory (LONGMEM). This framework enables language models to cache lengthy previous context or knowledge into a 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 a 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 memory-augmented layer of SideNet, the generated attention query of the current input is used to retrieve cached key-value pairs of previous contexts from the memory, and the corresponding memory augmentations are then fused into adaptable hidden states via a joint-attention mechanism. Furthermore, newly designed cross-network 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 is frozen during the efficient memory-augmented adaptation stage, LONGMEM can 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-augmented in-context learning with thousands of task-relevant
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-augmented in-context learning for natural language understanding (NLU) tasks. Experimental results demonstrate that our model consistently outperforms the strong baselines in 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\sim 1.62$  over different length splits of Gutenberg-2022 corpus. Notably, our model achieves state-of-the-art performance of  $40.5\%$  identification accuracy on ChapterBreak, a challenging long-context modeling benchmark, significantly surpassing existing strong x-former baselines. Finally, with 2k demonstration examples in memory, LONGMEM shows pronounced improvements in in-context learning on popular NLU tasks, compared with prominent memory-augmented and non-memory-augmented baselines.
To enable LLMs to 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 a 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-Term 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 on the Transformer architecture  $\left[\mathrm{VSP}^{+}17\right]$ . For LONGMEM, there are three key components: the frozen backbone LLM, SideNet, and Cache Memory Bank. As most existing pretrained LLMs can only take a 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 SideNet module can be viewed as an efficient adaption model that is trained to fuse the current input context and relevant cached previous contexts in the decoupled memory.
Formally, for a fixed-sized input text sequence  $\{\mathbf{x}_i\}_{i=1}^{|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  $\{\mathbf{x}_i\}_{i=1}^{|x|}$  into embedding space and outputs the initial hidden states,  $\mathbf{H}_{\mathrm{LLM}}^{0} \in \mathbb{R}^{|x| \times E}$ , 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,  $\mathbf{H}_{\mathrm{LLM}}^{l'} = f_{\theta_{\mathrm{LLM}}^{l'}}(\mathbf{H}_{\mathrm{LLM}}^{l' - 1}), \forall l' \in [1, L']$  and  $L'$  is 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 a head-wise vector queue  $\mathcal{Z}_k$ ,  $\mathcal{Z}_v \in \mathbb{R}^{H \times M \times d}$ , which maintains attention key-value pairs of latest  $M$  previous inputs  $\widetilde{\mathbf{K}}$ ,  $\widetilde{\mathbf{V}} \in \mathbb{R}^{H \times |x| \times d}$ , where  $H, d$  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  $\{\mathbf{H}_{\mathrm{LLM}}^{l^{\prime}}\}_{l^{\prime} = 1}^{L^{\prime}}$  and the past key-value pairs in the Cached Memory Bank for computing memory-augmented representations. Specifically, our SideNet of LONGMEM consists of  $(L - 1)$  normal Transformer decoder layers and one special memory-augmented 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 < L^{\prime}$ . Our SideNet encodes  $\mathbf{H}^0$  into memory-augmented contextual representation via  $(L - 1)$  normal Transformer decoder layers and a special memory-augmented layer.
The memory-augmented layer is an extension of the vanilla Transformer decoder layer that takes a memory-augmented input, including both top relevant key-value pairs in memory and the hidden states from the current input. Here, the cached key-value pairs are recalled using a token-based memory retrieval module (§2.3). For each current input token, the memory retrieval module  $s_{rt}(:)$  retrieves top-  $K$  relevant key-value pairs in the memory bank  $\{\widetilde{\pmb{k}}_{ij},\widetilde{\pmb{v}}_{ij}\}_{j = 1}^{K} = s_{rt}(\mathbf{x}_i)$ . Then SideNet computes the output using the memory-augmented input,  $\mathbf{H}_{\mathrm{Side}}^{m_s} = f_{\theta_{\mathrm{Mem}}}(\mathbf{H}_{\mathrm{Side}}^{m_s - 1},\{\{\widetilde{\pmb{k}}_{ij},\widetilde{\pmb{v}}_{ij}\}_{j = 1}^K\}_{i = 1}^{|x|})$ , where  $m_{s}$  is the layer index where we inject the memory-augmentation layer.
Finally, the token probability is computed using the last SideNet hidden states  $P(\mathbf{x}_i|\mathbf{x}_1,\dots ,\mathbf{x}_{i - 1}) =$  softmax(WH $^L$ ), 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], the training objective of LONGMEM is the standard left-to-right language modeling objective, which maximizes the likelihood of the next token based on the left context: max  $\sum_{x\in \mathcal{D}}\sum_{i = 1}^{\lvert\mathbf{x}\rvert}\log P(\mathbf{x}_i|\mathbf{x}_1,\dots ,\mathbf{x}_{i - 1})$ , where  $x$  is a randomly sampled sentence from the pre-training text corpus  $\mathcal{D}$ .
2.2 Residual SideNet
SideNet Architecture and Initialization. Here, we implement SideNet based on Transformer  $\left[\mathrm{VSP}^{+}17\right]$ . Specifically, the number of decoder layers  $L$  in SideNet is equal to the number of layers  $L^{\prime}$  in the backbone LLM divided by a reduction factor (a layer reduction factor of 2 is used throughout this work, i.e.,  $L^{\prime} = 2L$ ). The weights of each decoder layer in SideNet are initialized from the corresponding pre-trained decoder layer of the backbone LLM at the same depth:  $\Theta_{\mathrm{Side}}^{l} = \Theta_{\mathrm{LLM}}^{2l}$ . 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-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 our proposed cross-network residual connections to fuse representations from the backbone
LLM into SideNet. Specifically, we add the difference between output hidden states at  $2l$ -th and  $(2l - 2)$ -th layers of the backbone LLM as the residual connections to the output hidden states at  $l$ -th layer of SideNet. Then, the input to the next  $(l + 1)$ -th layer of SideNet is the sum of the original hidden state forwarded through the previous layer  $f_{\Theta_{\mathrm{side}}^l}(\mathbf{H}_{\mathrm{side}}^{l - 1})$  and the cross-network residual connection of the hidden state difference from the backbone LLM
$$
\mathbf {H} _ {\text {S i d e}} ^ {l} = f _ {\Theta_ {\text {S i d e}} ^ {l}} \left(\mathbf {H} _ {\text {S i d e}} ^ {l - 1}\right) + \left(\mathbf {H} _ {\text {L L M}} ^ {2 l} - \mathbf {H} _ {\text {L L M}} ^ {2 l - 2}\right), \forall l \in [ 1, L ], \tag {1}
$$
where  $\mathbf{H}^0$  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  $[\mathrm{VSP}^{+}17]$  will be performed as normal in  $f_{\Theta_{\mathrm{Side}}^l}(\mathbf{H}_{\mathrm{Side}}^{l - 1})$  and parallel to the proposed cross-network residual connections.
2.3 Memory Retrieval and Fusion
The long-term memory capability of LONGMEM is achieved via a memory-augmentation module for retrieval and fusion.
Token-to-Chunk Memory Retrieval. Instead of performing token-to-token retrieval, we focus on token-to-chunk retrieval for acceleration and integrity. A text-chunk refers to an n-gram structure of chunk-size  $csz$  number of contiguous tokens. The memory bank stores cached key-value pairs at the level of token chunks. We divide the memory bank into  $M / csz$  attention key-value paired chunks and use the mean-pooled vector on the chunk-size dimension to get the key vector for retrieval. Then we retrieve the top- $(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 a 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  $\{\widetilde{\mathbf{K}}_j, \widetilde{\mathbf{V}}_j\}_{j=1}^K$ . Adopting token-to-chunk 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  $[\mathrm{LGW}^{+}23]$  and  $[\mathrm{BMH}^{+}21]$ . The hyperparameter chunk-size  $csz$  controls the granularity of retrieved contexts, which can be empirically adjusted based on downstream tasks. For instance, in-context learning requires more fine-grained label tokens from demonstration examples cached in memory, where a smaller  $csz$  is helpful.
Memory Fusion. The memory fusion is performed within a special memory-augmented layer. As the conventional Transformer decoder layer uses the multi-head self-attention  $\left[\mathrm{VSP}^{+}17\right]$ , we follow [WRHS22] to extend it to a joint-attention 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  $\mathbf{H}^{l - 1}\in \mathbb{R}^{|x|\times d}$  and the corresponding retrieved attention key-value pairs are  $\{\widetilde{\mathbf{K}}_i,\widetilde{\mathbf{V}}_i\}_{i = 1}^{|x|}\in \mathbb{R}^{|x|\times \mathrm{K}\times \mathrm{d}}$ , the output hidden state for the  $l$ -th memory-augmented layer  $\mathbf{H}^l$  is computed as:
$$
\mathbf {A} = \operatorname {s o f t m a x} \left(\frac {\mathbf {Q} \mathbf {K} ^ {T}}{\sqrt {d}}\right) \mathbf {V}, \quad \mathbf {M} = \operatorname {C o n c a t} \left\{\operatorname {s o f t m a x} \left(\frac {\mathbf {Q} _ {i} \widetilde {\mathbf {K}} _ {i} ^ {T}}{\sqrt {d}}\right) \widetilde {\mathbf {V}} _ {i} \right\} _ {i = 1} ^ {| x |}, \tag {2}
$$
$$
\mathbf {H} ^ {l} = \operatorname {s i g m o i d} (g) \cdot \mathbf {A} + (1 - \operatorname {s i g m o i d} (g)) \cdot \mathbf {M}, \tag {3}
$$
where  $\mathbf{Q},\mathbf{K},\mathbf{V},\mathbf{A},\mathbf{M}\in \mathbb{R}^{|x|\times d}$ ,  $\mathrm{K}$  is the number of retrieved attention key-value pairs in cached memory for each token, and  $g$  is a trainable head-wise gating vector. The hidden state output from previous layer  $\mathbf{H}^{(l - 1)}$  is linearly projected into attention queries, keys, and values  $\mathbf{Q},\mathbf{K},\mathbf{V}$  separately via three matrices  $W^{Q},W^{K},W^{V}\in \mathbb{R}^{\mathrm{d}\times \mathrm{d}}$ . It is worth noting that the retrieved attention key-value pairs in cached memory are distinct to each token.
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; b) infinite-length in-context learning when loading a large number of demonstration examples into cached memory.
Figure 3: Batchfying the large text corpora into batches to ensure that each consecutive segments within each document is distributed in consecutive batches.
Batchfying the training corpora. The conventional batchfying process for large corpora truncates the whole corpora into consecutive fix-length text segments without padding and shuffles all segments to construct mini-batches  $\left[\mathrm{RWC}^{+}19\right]$ . 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 a document-level 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  $\left[\mathrm{GBB}^{+}20\right]$  as the training corpus, including BookCorpus2, Books3, OpenWebText2, Stack Exchange, Wikipedia, Gutenberg (PG-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  $\left[\mathrm{RWC}^{+}19\right]$  adopts absolute position embedding, which is found to perform poorly to enable LLM to learn long-distance dependencies  $\left[\mathrm{DYY}^{+}19\right]$ . The backbone LLM holds a  $L' = 24$ ,  $H = 16$ ,  $d = 64$  architecture. The SideNet holds a  $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  $csz$  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. The memory-augmentation 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.