MemoryLLM: Towards Self-Updatable Large Language Models 1 Introduction 2 Preliminaries 2.1 Problem Statement 2.2 Sketch of MemoryLLM 3 MemoryLLM 3.1 Structure Design 3.1.1 Memory Pool 3.1.2 Self-Update Process 3.1.3 Analysis of Forgetting 3.2 Training Strategy 3.2.1 New knowledge incorporation 3.2.2 Enhancing continuous contexts understanding 3.2.3 Mitigating forgetting problems 3.3 Model Instantiation 3.4 Discussions 4 Experiments 4.1 Evaluation Protocols 4.2 Implementation Details 4.3 Model Editing 4.3.1 Experimental Setup 4.3.2 Overall Performance Comparison 4.4 Long Context Evaluation 4.4.1 Experimental Setup 4.4.2 Overall Performance Comparison 4.4.3 Comparison with RAG methods 4.5 Knowledge Retention Experiments 4.5.1 Experimental Setup 4.6 Model Integrity Analysis 4.7 Ablation Study 4.7.1 Ablation Study of different K 𝐾 K italic_K and N 𝑁 N italic_N 4.7.2 Ablation Study of the Model Structures 5 Related Work 5.1 Memory based methods 5.2 Downsteam Tasks 6 Conclusion and Future Work A Details in Methodology A.1 Self-Update Process A.2 Training Strategy for New Knowledge Incorporation B Implementation Details B.1 Details for Mitigating Forgetting Problems C Additional Experiments C.1 Baselines for Model Editing MemoryLLM : Towards Self-Updatable Large Language Models Yu Wang * * * Work done during the internship at Amazon. Yifan Gao Xiusi Chen Haoming Jiang Shiyang Li Jingfeng Yang Qingyu Yin Zheng Li Xian Li Bing Yin Jingbo Shang Julian McAuley Abstract Existing Large Language Models (LLMs) usually remain static after deployment, which might make it hard to inject new knowledge into the model. We aim to build models containing a considerable portion of self-updatable parameters, enabling the model to integrate new knowledge effectively and efficiently. To this end, we introduce MemoryLLM , a model that comprises a transformer and a fixed-size memory pool within the latent space of the transformer. MemoryLLM can self-update with text knowledge and memorize the knowledge injected earlier. Our evaluations demonstrate the ability of MemoryLLM to effectively incorporate new knowledge, as evidenced by its performance on model editing benchmarks. Meanwhile, the model exhibits long-term information retention capacity, which is validated through our custom-designed evaluations and long-context benchmarks. MemoryLLM also shows operational integrity without any sign of performance degradation even after nearly a million memory updates. Our code and model are open-sourced at https://github.com/wangyu-ustc/MemoryLLM . memory, large language model 1 Introduction Despite the impressive performance LLMs demonstrate, a pivotal issue persists: How should we update the model with the latest knowledge? Previous solutions can be broadly categorized into three classes: (1) Retrieval-Based Methods: These methods rely on information retrieval in a knowledge base  (Khandelwal et al., 2019 ; Zhong et al., 2023 ) . They can yield strong results, but face challenges when redundancy in the knowledge base presents and suffer the logistical issue of managing an ever-expanding repository of knowledge. In multi-modality scenarios, retrieval-based methods might require enormous storage space to store all image data (24 images per second for humans) for retrieval purposes. (2) Model Editing: This class of methods involves making targeted edits to the model to adapt to new facts while preserving other desired capabilities  (Yao et al., 2023 ) . Existing methods primarily focus on fact-based editing, which is typically limited to single sentences. This limitation becomes more severe when one attempts to inject new knowledge in the form of longer and more complicated contexts. (3) Long Context Methods: Another alternative solution is to incorporate all the knowledge into the model’s context, which essentially makes the context into a knowledge base. This differs from retrieval-based methods in that the context directly informs the inference of the model. Methods in this category involve reducing the complexity of attention operations  (Child et al., 2019 ; Beltagy et al., 2020 ; Wang et al., 2020 ) , and modifying positional embeddings  (Press et al., 2021 ; Sun et al., 2023 ) to handle longer contexts. However, as complex reasoning tasks are thirsty for massive up-to-date knowledge, the inevitable context overload with long context methods becomes infeasible, as long as the context length is finite. In response to the challenges identified above, we introduce MemoryLLM , a model that embeds a substantial, fixed-size memory pool within its latent space, which serves as the self-updatable parameters. Specifically, we build the memory pool as hidden vectors within each layer of the transformer. At each layer, the memory pool contains memory tokens representing compressed knowledge. This design results in a memory pool that is less redundant than traditional knowledge bases in retrieval-based methods or contexts in long-context methods. To update the memory pool, we devise a self-update mechanism to propagate the new knowledge to every layer of the memory. During self-update, MemoryLLM only updates a proportion of memory in each layer to absorb the incoming knowledge. This allows previously stored knowledge to slowly phase out. These designs ensure MemoryLLM remains up-to-date while the old knowledge is slowly forgotten. After curated training, we update MemoryLLM nearly a million times without observing any performance deterioration. The evaluation of MemoryLLM focuses on several key aspects: (1) Integration of New Knowledge : The model’s performance is assessed with model editing benchmarks and QA tasks (long context QA benchmarks), where MemoryLLM demonstrates substantial improvements over existing methods. (2) Knowledge Retention Ability : MemoryLLM is evaluated on long context benchmarks and our knowledge retention experiments, showcasing its ability to recall knowledge. (3) Robustness : To test the integrity of the model, we subject MemoryLLM to almost a million update steps. The results show that our model is functioning properly even after extreme updates. In summary, our contributions are as follows: • We introduce MemoryLLM , which features an integrated memory pool within the latent space of an LLM. This memory pool is designed to manage new knowledge integration and encourage minimal information forgetting while being fixed-sized to circumvent the issue of uncontrolled growth. • We augment a 7B parameter model with an extensive memory pool comprising 1B parameters. • MemoryLLM demonstrates strong performance across various benchmarks, including model editing, long-context evaluation, and our knowledge retention experiments, showcasing its versatility and effectiveness in diverse applications. 2 Preliminaries 2.1 Problem Statement The primary challenge addressed in this paper is: How should we design a large language model that is capable of efficiently integrating new knowledge while minimizing the degradation of previously learned knowledge? To make the challenge more specific, we outline several essential properties that we hope to integrate into the new model: (1) Effiency : The process of knowledge injection into the model should be streamlined, potentially eliminating the need for back-propagation for efficiency. (2) Efficacy : It is crucial to ensure that the knowledge is effectively injected into the model, guaranteeing its impact on the model’s performance. (3) Knowledge Retention : Our model has a fixed-sized memory pool, implying a constant memorization capacity. This necessitates a mechanism for gradually phasing out older knowledge. (4) Integrity : The model must maintain full functionality regardless of the number of updates made to the memory pool. (5) Non-redundancy : We aim for more compact storage of knowledge, reducing redundancy, and optimizing memory usage. 2.2 Sketch of MemoryLLM To address the above challenges, our rough idea is to design a model denoted as ℳ θ , ϕ subscript ℳ 𝜃 italic-ϕ \mathcal{M}_{\theta,\phi} caligraphic_M start_POSTSUBSCRIPT italic_θ , italic_ϕ end_POSTSUBSCRIPT consisting of two sets of parameters: ϕ italic-ϕ \phi italic_ϕ and θ 𝜃 \theta italic_θ . Once we obtain the model, the ϕ italic-ϕ \phi italic_ϕ parameters should be static, while θ 𝜃 \theta italic_θ dynamically evolves when encountering new knowledge. This aligns with the intuition that some knowledge within an LLM should never change (persistent truths, encoded by ϕ italic-ϕ \phi italic_ϕ ) and some knowledge is being updated continuously (fresh information, modeled by θ 𝜃 \theta italic_θ ). Specifically, we use an existing large language model (Llama2) to model ϕ italic-ϕ \phi italic_ϕ , while θ 𝜃 \theta italic_θ is modeled by the memory pool with the detailed structure in Section 3.1.1 . Here we need to design the self-updating mechanism of θ 𝜃 \theta italic_θ that is pivotal to this process. Denoting the new knowledge as x 𝑥 x italic_x , a text paragraph, the self-updating process refers to updating θ 𝜃 \theta italic_θ in a way that does not compromise the general capabilities of the model while injecting the latest knowledge x 𝑥 x italic_x into the memory pool θ 𝜃 \theta italic_θ to obtain a new memory pool θ ′ superscript 𝜃 ′ \theta^{\prime} italic_θ start_POSTSUPERSCRIPT ′ end_POSTSUPERSCRIPT : θ ′ = U ⁢ ( θ , x ) superscript 𝜃 ′ 𝑈 𝜃 𝑥 \theta^{\prime}=U(\theta,x) italic_θ start_POSTSUPERSCRIPT ′ end_POSTSUPERSCRIPT = italic_U ( italic_θ , italic_x ) (1) Here U 𝑈 U italic_U is the update function which takes the memory pool θ 𝜃 \theta italic_θ and the new knowledge x 𝑥 x italic_x as input and outputs the new memory pool θ ′ superscript 𝜃 ′ \theta^{\prime} italic_θ start_POSTSUPERSCRIPT ′ end_POSTSUPERSCRIPT . Extending this process to multistep updating, consider a scenario with a never-ending context or a series of conversation histories, represented as ( x 1 , ⋯ , x n ) subscript 𝑥 1 ⋯ subscript 𝑥 𝑛 (x_{1},\cdots,x_{n}) ( italic_x start_POSTSUBSCRIPT 1 end_POSTSUBSCRIPT , ⋯ , italic_x start_POSTSUBSCRIPT italic_n end_POSTSUBSCRIPT ) , where x i , i ∈ { 1 , ⋯ , n } subscript 𝑥 𝑖 𝑖 1 ⋯ 𝑛 x_{i},i\in\{1,\cdots,n\} italic_x start_POSTSUBSCRIPT italic_i end_POSTSUBSCRIPT , italic_i ∈ { 1 , ⋯ , italic_n } is a text paragraph. The model requires the integration of all these contexts, which can be accomplished using the update function I 𝐼 I italic_I defined in Eq.( 1 ): θ n = U ( ⋯ ( U ( θ , x 1 ) , x n ) . \theta_{n}=U(\cdots(U(\theta,x_{1}),x_{n}). italic_θ start_POSTSUBSCRIPT italic_n end_POSTSUBSCRIPT = italic_U ( ⋯ ( italic_U ( italic_θ , italic_x start_POSTSUBSCRIPT 1 end_POSTSUBSCRIPT ) , italic_x start_POSTSUBSCRIPT italic_n end_POSTSUBSCRIPT ) . (2) We define the process self-updating as modifying the parameters θ 𝜃 \theta italic_θ with newly encountered knowledge x 𝑥 x italic_x , essentially enabling the model to read and assimilate knowledge. This design presents two primary challenges: (1) Parameter and Interaction Design : We need to determine the structure for θ 𝜃 \theta italic_θ and how it should interact with ϕ italic-ϕ \phi italic_ϕ , The goal is to allow the LLM to effectively use the knowledge from the θ 𝜃 \theta italic_θ in the generation process. (2) Update function design : It is crucial to design the update function U 𝑈 U italic_U such that θ 𝜃 \theta italic_θ can be updated without disturbing the old knowledge and undermining the overall capabilities of the model. 3 MemoryLLM 3.1 Structure Design (a) Generation (b) Self-Update Figure 1: The framework of MemoryLLM . (a) During generation, all memory tokens in the l 𝑙 l italic_l -th layer of memory pool θ l superscript 𝜃 𝑙 \theta^{l} italic_θ start_POSTSUPERSCRIPT italic_l end_POSTSUPERSCRIPT are attended by the hidden states h l subscript ℎ 𝑙 h_{l} italic_h start_POSTSUBSCRIPT italic_l end_POSTSUBSCRIPT . (b) During self-update, The last k 𝑘 k italic_k memory tokens from θ l subscript 𝜃 𝑙 \theta_{l} italic_θ start_POSTSUBSCRIPT italic_l end_POSTSUBSCRIPT are taken to be concatenated with the hidden states h l subscript ℎ 𝑙 h_{l} italic_h start_POSTSUBSCRIPT italic_l end_POSTSUBSCRIPT as the input to ϕ l subscript italic-ϕ 𝑙 \phi_{l} italic_ϕ start_POSTSUBSCRIPT italic_l end_POSTSUBSCRIPT . The output h l + 1 subscript ℎ 𝑙 1 h_{l+1} italic_h start_POSTSUBSCRIPT italic_l + 1 end_POSTSUBSCRIPT goes to the next layer. The last K 𝐾 K italic_K tokens of h l + 1 subscript ℎ 𝑙 1 h_{l+1} italic_h start_POSTSUBSCRIPT italic_l + 1 end_POSTSUBSCRIPT serve as the new memory tokens e θ l ′ superscript superscript subscript 𝑒 𝜃 𝑙 ′ {e_{\theta}^{l}}^{\prime} italic_e start_POSTSUBSCRIPT italic_θ end_POSTSUBSCRIPT start_POSTSUPERSCRIPT italic_l end_POSTSUPERSCRIPT start_POSTSUPERSCRIPT ′ end_POSTSUPERSCRIPT . We randomly drop K 𝐾 K italic_K tokens in θ l subscript 𝜃 𝑙 \theta_{l} italic_θ start_POSTSUBSCRIPT italic_l end_POSTSUBSCRIPT and concatenate the left θ l subscript 𝜃 𝑙 \theta_{l} italic_θ start_POSTSUBSCRIPT italic_l end_POSTSUBSCRIPT (denoted as θ l ⁢ ( d ) superscript 𝜃 𝑙 𝑑 \theta^{l}(d) italic_θ start_POSTSUPERSCRIPT italic_l end_POSTSUPERSCRIPT ( italic_d ) ) with e θ l ′ superscript superscript subscript 𝑒 𝜃 𝑙 ′ {e_{\theta}^{l}}^{\prime} italic_e start_POSTSUBSCRIPT italic_θ end_POSTSUBSCRIPT start_POSTSUPERSCRIPT italic_l end_POSTSUPERSCRIPT start_POSTSUPERSCRIPT ′ end_POSTSUPERSCRIPT to obtain new memory θ l ′ superscript subscript 𝜃 𝑙 ′ \theta_{l}^{\prime} italic_θ start_POSTSUBSCRIPT italic_l end_POSTSUBSCRIPT start_POSTSUPERSCRIPT ′ end_POSTSUPERSCRIPT . 3.1.1 Memory Pool We choose to instantiate ϕ italic-ϕ \phi italic_ϕ with an off-the-shelf LLM, specifically Llama2  (Touvron et al., 2023 ) . ϕ italic-ϕ \phi italic_ϕ consists of multiple transformer layers, denoted as ϕ = { ϕ l } l = 1 L italic-ϕ superscript subscript subscript italic-ϕ 𝑙 𝑙 1 𝐿 \phi=\{\phi_{l}\}_{l=1}^{L} italic_ϕ = { italic_ϕ start_POSTSUBSCRIPT italic_l end_POSTSUBSCRIPT } start_POSTSUBSCRIPT italic_l = 1 end_POSTSUBSCRIPT start_POSTSUPERSCRIPT italic_L end_POSTSUPERSCRIPT , where L 𝐿 L italic_L represents the total number of layers. To facilitate the transformer ϕ italic-ϕ \phi italic_ϕ to understand the memory pool θ 𝜃 \theta italic_θ , we conceptualize θ 𝜃 \theta italic_θ as hidden vectors within each transformer layer, symbolized as θ = { θ l } l = 1 L 𝜃 superscript subscript subscript 𝜃 𝑙 𝑙 1 𝐿 \theta=\{\theta_{l}\}_{l=1}^{L} italic_θ = { italic_θ start_POSTSUBSCRIPT italic_l end_POSTSUBSCRIPT } start_POSTSUBSCRIPT italic_l = 1 end_POSTSUBSCRIPT start_POSTSUPERSCRIPT italic_L end_POSTSUPERSCRIPT . Each θ l subscript 𝜃 𝑙 \theta_{l} italic_θ start_POSTSUBSCRIPT italic_l end_POSTSUBSCRIPT is of dimension N × d 𝑁 𝑑 N\times d italic_N × italic_d , corresponding to N 𝑁 N italic_N hidden states and the word embedding dimension d 𝑑 d italic_d in ϕ italic-ϕ \phi italic_ϕ . We term θ l subscript 𝜃 𝑙 \theta_{l} italic_θ start_POSTSUBSCRIPT italic_l end_POSTSUBSCRIPT memory tokens . The memory tokens serve as the representation of previous knowledge that the model has seen in a more compressed manner. We intend to maximize the memory size, so we assign the memory pool to every layer to significantly enlarge the memory pool. During the generation phase, all memory tokens are used, as illustrated in Figure 1(a) . The attention map is designed to enable every token in x 𝑥 x italic_x to attend to all memory tokens. If x 𝑥 x italic_x comprises n x subscript 𝑛 𝑥 n_{x} italic_n start_POSTSUBSCRIPT italic_x end_POSTSUBSCRIPT tokens, the attention map assumes a shape of n x × ( n x + N ) subscript 𝑛 𝑥 subscript 𝑛 𝑥 𝑁 n_{x}\times(n_{x}+N) italic_n start_POSTSUBSCRIPT italic_x end_POSTSUBSCRIPT × ( italic_n start_POSTSUBSCRIPT italic_x end_POSTSUBSCRIPT + italic_N ) , yielding a linear complexity w.r.t. the size of the memory pool. 3.1.2 Self-Update Process Figure 1(b) illustrates the self-update process. The goal of self-update is to ensure that MemoryLLM can always digest the latest knowledge and memorize the previously learned knowledge at its best. We discuss the self-update process in this subsection and prove in section  3.1.3 that MemoryLLM only forgets stale knowledge at an exponential decay rate with a theoretical guarantee. When introducing new knowledge x c subscript 𝑥 𝑐 x_{c} italic_x start_POSTSUBSCRIPT italic_c end_POSTSUBSCRIPT (in the following, we denote the new knowledge as context x c subscript 𝑥 𝑐 x_{c} italic_x start_POSTSUBSCRIPT italic_c end_POSTSUBSCRIPT to distinguish it from x 𝑥 x italic_x in the last section), the model must integrate x c subscript 𝑥 𝑐 x_{c} italic_x start_POSTSUBSCRIPT italic_c end_POSTSUBSCRIPT into θ 𝜃 \theta italic_θ as per Eq.( 1 ). To avoid additional modules and complexities, we use the transformer ϕ italic-ϕ \phi italic_ϕ for the update. Ideally, the input to ϕ l subscript italic-ϕ 𝑙 \phi_{l} italic_ϕ start_POSTSUBSCRIPT italic_l end_POSTSUBSCRIPT should be the memory pool θ l subscript 𝜃 𝑙 \theta_{l} italic_θ start_POSTSUBSCRIPT italic_l end_POSTSUBSCRIPT and the hidden states h l subscript ℎ 𝑙 h_{l} italic_h start_POSTSUBSCRIPT italic_l end_POSTSUBSCRIPT (where h 1 subscript ℎ 1 h_{1} italic_h start_POSTSUBSCRIPT 1 end_POSTSUBSCRIPT are the word embeddings of tokenized x c subscript 𝑥 𝑐 x_{c} italic_x start_POSTSUBSCRIPT italic_c end_POSTSUBSCRIPT ). We find it essential to maintain the gradient flow from both the self-update and the generation to achieve better performance (see Section   3.2.1 ). However, it is much more costly to feed the entire pool θ l subscript 𝜃 𝑙 \theta_{l} italic_θ start_POSTSUBSCRIPT italic_l end_POSTSUBSCRIPT to ϕ l subscript italic-ϕ 𝑙 \phi_{l} italic_ϕ start_POSTSUBSCRIPT italic_l end_POSTSUBSCRIPT during self-update. To solve this problem, we extract the last K 𝐾 K italic_K tokens of θ l subscript 𝜃 𝑙 \theta_{l} italic_θ start_POSTSUBSCRIPT italic_l end_POSTSUBSCRIPT where K << N much-less-than 𝐾 𝑁 K< K subscript 𝑛 subscript 𝑥 𝑐 𝐾 n_{x_{c}}>K italic_n start_POSTSUBSCRIPT italic_x start_POSTSUBSCRIPT italic_c end_POSTSUBSCRIPT end_POSTSUBSCRIPT > italic_K . The case when k > n x c 𝑘 subscript ... [OUTPUT TRUNCATED - 25580 chars omitted out of 75580 total] ... dge retention ability of MemoryLLM , where the knowledge from multiple updates earlier could boost performance. The performance of MemoryLLM , when the context length is less than 4k, is not the same as that of Llama2-7B, which can be attributed to the subset we used for training MemoryLLM , as we do not need to use the entire dataset for pertaining Llama2-7B for our model and a subset would inevitably have distribution shift from the original dataset. 4.4.3 Comparison with RAG methods In this section, we aim to explore the role of RAG methods in QA tasks which we argue is orthogonal to MemoryLLM . The primary goal of MemoryLLM is to achieve self-updatable LLM where the memory module serves as the parameters that could keep updating along the inference process, whereas RAG methods aim to retrieve the most relevant piece of information from the history. Intuitively, RAG is used to conduct coarse retrieval from millions of documents, while MemoryLLM can process the retrieved documents. We use BM25 retriever to extract 4k tokens from the whole context and use MemoryLLM to process these 4k tokens to generate the answer. The results are shown in Table 2 . Here MemoryLLM-7b-16k corresponds to the results in Figure 4 , and MemoryLLM-7b-all-BM25 means retrieving 4k tokens from the whole given context and using MemoryLLM to process the retrieved 4k tokens. The results show that using the BM25 retriever could enhance the model performance on certain datasets while not universally beneficial. Table 2: The performance comparison on long context QA benchmarks of our model with and without BM25 retriever. MemoryLLM -16k MemoryLLM -all-BM25 narrativeqa 20.64 15.60 qasper 19.57 20.30 multifieldqa_en 29.56 33.08 hotpotqa 34.03 32.27 2wikimqa 27.22 24.17 musique 13.47 15.36 Figure 4: Experimental Results on LongBench . The x-axis is the maximum context length for the QA task. For instance, with a maximum length of 4096 4096 4096 4096 , we truncate 4096 4096 4096 4096 tokens from the given context as input to the model. The y-axis is the F1 score. (a) SQuAD (b) NaturalQA Figure 5: Performance Comparison on SQuAD and NaturalQA. The x-axis shows the number of updates we perform on the model, where the context that contains the knowledge to answer the question is injected in Step 1. The y-axis reveals the accuracy of the model’s prediction after a certain number of updates. The accuracy is higher than the borderline indicating that the knowledge is not completely forgotten, while we wish the model to be more aligned with the exponential decay, i.e., the theoretical upper bound. (a) SQuAD (b) NaturalQA Figure 6: Model Integrity Check with SQuAD and NaturalQA. We plot accuracy along the updating process as well as the exponential moving average as the Smoothed (99.99%) value. We do not observe any decrease over 650k updates. (a) NaturalQA Acc Percentage (b) SQuAD Acc Percentage (c) NaturalQA Accuracy (d) SQuAD Accuracy Figure 7: Ablation Study with our knowledge retention experiments on NaturalQA and SQuAD. All models are trained with the same setting, 30 × 256 30 256 30\times 256 30 × 256 is our main model. The relevant knowledge for answering the question is injected in step 1, and the x-axis means the number of updates (steps) performed. The top figures show the ratio of the accuracy at each step compared with the accuracy at step 1 for better visualization of knowledge retention. 4.5 Knowledge Retention Experiments 4.5.1 Experimental Setup The datasets are prepared as below: SQuAD : Formatted as (context, question, answer) , where context and question are sentences, answer refers to the first answer in the list of ground-truth acceptable answers. Then we extract all the samples with answer shorter or equal to 3 3 3 3 tokens. The model generates 10 new tokens from the prompt “Question: Question Answer:”. Correct predictions cover the 3-token answer within the 10 generated tokens. A total of 2 , 250 2 250 2,250 2 , 250 samples are used for the accuracy calculation. NaturalQA : Formatted as (context, question, answer) , using the long answer as the context and the short answer as the ground truth. Samples with answers of 4 tokens or less are selected. Like SQuAD, the model generates 10 new tokens, and the correct predictions cover the 4-token answer. This yields 1,004 samples for analysis. The results are shown in Figure 5 . We assess MemoryLLM ’s forgetting rate, comparing it against a baseline (accuracy without context injected into the memory) and a theoretical upper bound. Denote the accuracy at step 1 as a u subscript 𝑎 𝑢 a_{u} italic_a start_POSTSUBSCRIPT italic_u end_POSTSUBSCRIPT , and the borderline accuracy as a b subscript 𝑎 𝑏 a_{b} italic_a start_POSTSUBSCRIPT italic_b end_POSTSUBSCRIPT . Then at step t 𝑡 t italic_t , we calculate the point on the curve with the following equation: a t = ( a u − a b ) ∗ ( N − K N ) t − 1 subscript 𝑎 𝑡 subscript 𝑎 𝑢 subscript 𝑎 𝑏 superscript 𝑁 𝐾 𝑁 𝑡 1 a_{t}=(a_{u}-a_{b})*\Big{(}\frac{N-K}{N}\Big{)}^{t-1} italic_a start_POSTSUBSCRIPT italic_t end_POSTSUBSCRIPT = ( italic_a start_POSTSUBSCRIPT italic_u end_POSTSUBSCRIPT - italic_a start_POSTSUBSCRIPT italic_b end_POSTSUBSCRIPT ) ∗ ( divide start_ARG italic_N - italic_K end_ARG start_ARG italic_N end_ARG ) start_POSTSUPERSCRIPT italic_t - 1 end_POSTSUPERSCRIPT (5) In our instantiation, N = 7 , 680 𝑁 7 680 N=7,680 italic_N = 7 , 680 and K = 256 𝐾 256 K=256 italic_K = 256 . Our findings indicate that the model retains knowledge even after 20 updates. However, it falls short of the exponential decay curve representing the upper bound. This gap can be attributed to the fact that even if the knowledge is partially corrupted after 20 steps of updating, it might be hard for the model to reveal the exact answer. The performance exceeding the upper bound at step 2 on the dataset SQuAD might be due to (a) the variation of inference and (b) dropping a small part of the memory may not affect the model predicting the words, while the exponential curve would drop. 4.6 Model Integrity Analysis To illustrate the integrity of our model, We update our model with NaturalQA and SQuAD mentioned in Section 4.5 . Each time we go through the whole dataset, we shuffle the dataset and inject it into the memory again. In this way, we can simulate infinite updates. Then during the updating process, we track if our model could answer the question related to the most recent context, obtaining a long binary array that indicates whether our model successfully answers the question related to the most recently injected context. With this binary array, we calculate the average accuracy of the last 2 , 250 2 250 2,250 2 , 250 samples for SQuAD and the last 1 , 004 1 004 1,004 1 , 004 samples for NaturalQA. The results are shown in Figure 6 . We continue running for up to 650 , 000 650 000 650,000 650 , 000 steps for 3 3 3 3 days. As shown in the figure, there is no sign of decreasing in accuracy even after the 650 , 000 650 000 650,000 650 , 000 steps, demonstrating the integrity of our model. From this observation, we argue that our model could be potentially updated for arbitrarily many times without affecting the functioning ability. 4.7 Ablation Study 4.7.1 Ablation Study of different K 𝐾 K italic_K and N 𝑁 N italic_N In this section, we study the effects of different K 𝐾 K italic_K and N 𝑁 N italic_N in Eq.( 3 ) with our knowledge-retention experiments. Our primary goal is to explore the forgetting ratio when the model has different memory sizes ( N 𝑁 N italic_N ) and numbers of tokens to store the new knowledge ( K 𝐾 K italic_K ). We vary N 𝑁 N italic_N to be { 10 × 256 , 20 × 256 , 30 × 256 } 10 256 20 256 30 256 \{10\times 256,20\times 256,30\times 256\} { 10 × 256 , 20 × 256 , 30 × 256 } , and K 𝐾 K italic_K to be { 256 , 512 } 256 512 \{256,512\} { 256 , 512 } . We also tried K = 128 𝐾 128 K=128 italic_K = 128 , but to find that the accuracy is much worse than the other settings (the step 1 accuracy of NQA and SQuAD under the setting K = 128 𝐾 128 K=128 italic_K = 128 are 0.34 and 0.25, respectively) , we omit this setting here. The results shown in Figure 7 reveal that (1) When K 𝐾 K italic_K is fixed, with greater N 𝑁 N italic_N , the forgetting ratio is smaller; (2) When N 𝑁 N italic_N is fixed, with smaller K 𝐾 K italic_K ( 10 × 512 10 512 10\times 512 10 × 512 vs. 20 × 256 20 256 20\times 256 20 × 256 , the latter yields better knowledge-retention ability), the forgetting ratio becomes smaller. These experiments support our intuition and show that with the improvement of N / K 𝑁 𝐾 N/K italic_N / italic_K , we can enable better knowledge-retention ability. 4.7.2 Ablation Study of the Model Structures In our main experiments, we train the model with the memory tokens augmented in every layer. To study the necessity of this design, we tried the following several settings: (1) Augment only one layer in the model with memory tokens; (2) Augment the last half of the layers in the model with the memory tokens (this design is inspired by Figure 6(a) in Fang et al. ( 2024 ) ). Then we find that design (1) leads to almost zero improvements with the context compared to the performance without the context, which means augmenting only one layer is almost useless. For design (2), We record the accuracies after injecting the context for one step: NaturalQA: 0.39, SQuAD: 0.22. For reference, the accuracy of NaturalQA and SQuAD in Figure 7 at step 1 is 0.46 and 0.39 respectively. This shows that having the memory tokens in both the first half and the second half layers is necessary for better performance. 5 Related Work 5.1 Memory based methods Previous memory-based methods share certain similarities with MemoryLLM . Among these methods, some use an external encoder to inject knowledge into the memory pool, such as the Memory Network  (Weston et al., 2014 ) , which focuses on rectifying the forgetting problems in RNNs. Follow-up work Sukhbaatar et al. ( 2015 ) computes the weighted sum of the entire memory pool as the representative vector of the memory. Others use the language model itself as the encoder to update the memory. Memory Transformer  (Burtsev & Sapunov, 2020 ) and RMT  (Bulatov et al., 2022 ) propose to add memory tokens when reading the contexts, where the memory pool is up to 20 20 20 20 tokens. EMMA  (Moro et al., 2023 ) has a slightly larger memory pool, which is the size of the chunk when injecting the contexts into the memory. These fixed-sized memory pools show promising results, although performance is limited by the size of the memory pool. This also shows the challenges of expanding memory and incorporating information without disturbing the original capability of the model. Other memory-based methods integrate the memory pool with unfixed size, where different forgetting mechanisms are adopted to handle the ever-growing problem. In this case, the memory pool would be in the form of (1) hidden states, such as (Adel, 2022 ) and MemoryBank  (Zhong et al., 2023 ) ; (2) key-value pairs, represented by KNN-LM  (Khandelwal et al., 2019 ) , LONGMEM  (Wang et al., 2023 ) . (3) vectors in hidden space. This involves the image captioning task  (Cornia et al., 2020 ) and Memformer  (Wu et al., 2022 ) . (4) raw texts. RET-LLM  (Modarressi et al., 2023 ) proposes to save the knowledge with triplets into the memory and then use API query to retrieve related information in the memory given the context. These methods have a more flexible memory pool. However, the memory pool might be redundant in terms of the stored knowledge. 5.2 Downsteam Tasks As MemoryLLM has a large memory pool that can be used to store knowledge, it could be used for downstream tasks such as model editing and long context tasks. For model editing tasks  (Yao et al., 2023 ) , MEND  (Mitchell et al., 2022 ) and ROME  (Meng et al., 2022 ) propose to modify the parameters of the LLM with the new given fact. During inference, MEND needs back-propagation and ROME requires the optimization for new MLP weights, while MemoryLLM , regarding the memory pool as part of the model parameters, could directly update the memory pool to store new facts. IKE  (Zheng et al., 2023 ) proposes to simply put the new facts in context, which is straightforward and intuitively similar to MemoryLLM in terms of this task. However, IKE would encounter the same problem as long context methods, i.e., the ever-growing contexts. For Long context tasks, representative methods can be categorized as follows: (1) Efficient Attention such as Longformer  (Beltagy et al., 2020 ) , Linformer  (Wang et al., 2020 ) , LongNet  (Jiayu et al., 2023 ) , (2) Positional Encoding like Alibi  (Press et al., 2021 ) , Positional Interpolation  (Chen et al., 2023a ) and Extrapolation  (Sun et al., 2023 ) , (3) Finetuning with longer context  (Xiong et al., 2023 ; Tworkowski et al., 2023 ) , (4) Memory-based methods  (Wang et al., 2023 ; Bulatov et al., 2022 ; Wu et al., 2022 ) . Among all these categories, MemoryLLM could fit into the fourth category where long contexts are absorbed into the memory, which is used for future prediction. 6 Conclusion and Future Work In this paper, we propose MemoryLLM , a language model consisting of a transformer and a huge memory pool within the latent space of the transformer, which serves as the self-updatable parameters of the model. MemoryLLM can perform self-updates on the memory with new knowledge, enabling effective knowledge incorporation and slow forgetting of previous knowledge. Comparisons against baselines for model editing and long context, together with a dedicated customized evaluation for knowledge retention analysis, demonstrate the superiority of MemoryLLM in effectively absorbing new knowledge and knowledge retention ability. In the future, it is of interest to extend the memory size as well as increase the compression rate, i.e., using fewer memory tokens during self-update to store the new knowledge. In addition, we aim to extend MemoryLLM to be multimodal, as the memory tokens of MemoryLLM may be suitable for storing multimodal knowledge. Impact Statement This paper presents work that aims to advance the field of Natural Language Processing, specifically the Large Language Models. There are many potential societal consequences of our work associated with LLMs, such as AI safety and reliability. Beyond LLMs, we feel no other consequences must be highlighted here. References Adel (2022) Adel, A. A. Global memory transformer for processing long documents. CoRR , abs/2212.01650, 2022. Bai et al. (2023) Bai, Y., Lv, X., Zhang, J., Lyu, H., Tang, J., Huang, Z., Du, Z., Liu, X., Zeng, A., Hou, L., et al. Longbench: A bilingual, multitask benchmark for long context understanding. arXiv preprint arXiv:2308.14508 , 2023. Beltagy et al. (2020) Beltagy, I., Peters, M. E., and Cohan, A. Longformer: The long-document transformer. CoRR , abs/2004.05150, 2020. URL https://arxiv.org/abs/2004.05150 . Bulatov et al. (2022) Bulatov, A., Kuratov, Y., and Burtsev, M. S. Recurrent memory transformer. In NeurIPS , 2022. Burtsev & Sapunov (2020) Burtsev, M. S. and Sapunov, G. V. Memory transformer. CoRR , abs/2006.11527, 2020. URL https://arxiv.org/abs/2006.11527 . Cao et al. (2021) Cao, N. D., Aziz, W., and Titov, I. Editing factual knowledge in language models. In EMNLP (1) , pp.  6491–6506. Association for Computational Linguistics, 2021. Chen et al. (2023a) Chen, S., Wong, S., Chen, L., and Tian, Y. Extending context window of large language models via positional interpolation. arXiv preprint arXiv:2306.15595 , 2023a. Chen et al. (2023b) Chen, Y., Qian, S., Tang, H., Lai, X., Liu, Z., Han, S., and Jia, J. Longlora: Efficient fine-tuning of long-context large language models. arXiv preprint arXiv:2309.12307 , 2023b. Child et al. (2019) Child, R., Gray, S., Radford, A., and Sutskever, I. Generating long sequences with sparse transformers. CoRR , abs/1904.10509, 2019. URL http://arxiv.org/abs/1904.10509 . Computer (2023) Computer, T. Redpajama: an open dataset for training large language models, 2023. URL https://github.com/togethercomputer/RedPajama-Data . Cornia et al. (2020) Cornia, M., Stefanini, M., Baraldi, L., and Cucchiara, R. Meshed-memory transformer for image captioning. In CVPR , pp.  10575–10584. Computer Vision Foundation / IEEE, 2020. Fang et al. (2024) Fang, J., Tang, L., Bi, H., Qin, Y., Sun, S., Li, Z., Li, H., Li, Y., Cong, X., Yan, Y., et al. Unimem: Towards a unified view of long-context large language models. arXiv preprint arXiv:2402.03009 , 2024. Geng & Liu (2023) Geng, X. and Liu, H. Openllama: An open reproduction of llama, May 2023. URL https://github.com/openlm-research/open_llama . Jiayu et al. (2023) Jiayu, D., Shuming, M., Li, D., Xingxing, Z., Shaohan, H., Wenhui, W., and Wei†, F. Longnet: Scaling transformers to 1,000,000,000 tokens. arXiv preprint arXiv:2307.02486 , 2023. Khandelwal et al. (2019) Khandelwal, U., Levy, O., Jurafsky, D., Zettlemoyer, L., and Lewis, M. Generalization through memorization: Nearest neighbor language models. arXiv preprint arXiv:1911.00172 , 2019. Levy et al. (2017) Levy, O., Seo, M., Choi, E., and Zettlemoyer, L. Zero-shot relation extraction via reading comprehension. In Levy, R. and Specia, L. (eds.), Proceedings of the 21st Conference on Computational Natural Language Learning (CoNLL 2017), Vancouver, Canada, August 3-4, 2017 , pp.  333–342. Association for Computational Linguistics, 2017. doi: 10.18653/V1/K17-1034 . URL https://doi.org/10.18653/v1/K17-1034 . Meng et al. (2022) Meng, K., Bau, D., Andonian, A., and Belinkov, Y. Locating and editing factual associations in gpt. Advances in Neural Information Processing Systems , 35:17359–17372, 2022. Mitchell et al. (2022) Mitchell, E., Lin, C., Bosselut, A., Finn, C., and Manning, C. D. Fast model editing at scale. In ICLR . OpenReview.net, 2022. Modarressi et al. (2023) Modarressi, A., Imani, A., Fayyaz, M., and Schütze, H. Ret-llm: Towards a general read-write memory for large language models. arXiv preprint arXiv:2305.14322 , 2023. Moro et al. (2023) Moro, G., Ragazzi, L., Valgimigli, L., Frisoni, G., Sartori, C., and Marfia, G. Efficient memory-enhanced transformer for long-document summarization in low-resource regimes. Sensors , 23(7):3542, 2023. Press et al. (2021) Press, O., Smith, N. A., and Lewis, M. Train short, test long: Attention with linear biases enables input length extrapolation. arXiv preprint arXiv:2108.12409 , 2021. Raffel et al. (2020) Raffel, C., Shazeer, N., Roberts, A., Lee, K., Narang, S., Matena, M., Zhou, Y., Li, W., and Liu, P. J. Exploring the limits of transfer learning with a unified text-to-text transformer. J. Mach. Learn. Res. , 21:140:1–140:67, 2020. Sukhbaatar et al. (2015) Sukhbaatar, S., Weston, J., Fergus, R., et al. End-to-end memory networks. Advances in neural information processing systems , 28, 2015. Sun et al. (2023) Sun, Y., Dong, L., Patra, B., Ma, S., Huang, S., Benhaim, A., Chaudhary, V., Song, X., and Wei, F. A length-extrapolatable transformer. In ACL (1) , pp.  14590–14604. Association for Computational Linguistics, 2023. Touvron et al. (2023) Touvron, H., Martin, L., Stone, K., Albert, P., Almahairi, A., Babaei, Y., Bashlykov, N., Batra, S., Bhargava, P., Bhosale, S., Bikel, D., Blecher, L., Canton-Ferrer, C., Chen, M., Cucurull, G., Esiobu, D., Fernandes, J., Fu, J., Fu, W., Fuller, B., Gao, C., Goswami, V., Goyal, N., Hartshorn, A., Hosseini, S., Hou, R., Inan, H., Kardas, M., Kerkez, V., Khabsa, M., Kloumann, I., Korenev, A., Koura, P. S., Lachaux, M., Lavril, T., Lee, J., Liskovich, D., Lu, Y., Mao, Y., Martinet, X., Mihaylov, T., Mishra, P., Molybog, I., Nie, Y., Poulton, A., Reizenstein, J., Rungta, R., Saladi, K., Schelten, A., Silva, R., Smith, E. M., Subramanian, R., Tan, X. E., Tang, B., Taylor, R., Williams, A., Kuan, J. X., Xu, P., Yan, Z., Zarov, I., Zhang, Y., Fan, A., Kambadur, M., Narang, S., Rodriguez, A., Stojnic, R., Edunov, S., and Scialom, T. Llama 2: Open foundation and fine-tuned chat models. CoRR , abs/2307.09288, 2023. Tworkowski et al. (2023) Tworkowski, S., Staniszewski, K., Pacek, M., Wu, Y., Michalewski, H., and Miłoś, P. Focused transformer: Contrastive training for context scaling. arXiv preprint arXiv:2307.03170 , 2023. Wang et al. (2020) Wang, S., Li, B. Z., Khabsa, M., Fang, H., and Ma, H. Linformer: Self-attention with linear complexity. arXiv preprint arXiv:2006.04768 , 2020. Wang et al. (2023) Wang, W., Dong, L., Cheng, H., Liu, X., Yan, X., Gao, J., and Wei, F. Augmenting language models with long-term memory. arXiv preprint arXiv:2306.07174 , 2023. Weston et al. (2014) Weston, J., Chopra, S., and Bordes, A. Memory networks. arXiv preprint arXiv:1410.3916 , 2014. Wu et al. (2022) Wu, Q., Lan, Z., Qian, K., Gu, J., Geramifard, A., and Yu, Z. Memformer: A memory-augmented transformer for sequence modeling. In AACL/IJCNLP (Findings) , pp.  308–318. Association for Computational Linguistics, 2022. Xiong et al. (2023) Xiong, W., Liu, J., Molybog, I., Zhang, H., Bhargava, P., Hou, R., Martin, L., Rungta, R., Sankararaman, K. A., Oguz, B., et al. Effective long-context scaling of foundation models. arXiv preprint arXiv:2309.16039 , 2023. Yao et al. (2023) Yao, Y., Wang, P., Tian, B., Cheng, S., Li, Z., Deng, S., Chen, H., and Zhang, N. Editing large language models: Problems, methods, and opportunities. CoRR , abs/2305.13172, 2023. Zheng et al. (2023) Zheng, C., Li, L., Dong, Q., Fan, Y., Wu, Z., Xu, J., and Chang, B. Can we edit factual knowledge by in-context learning? arXiv preprint arXiv:2305.12740 , 2023. Zhong et al. (2023) Zhong, W., Guo, L., Gao, Q., and Wang, Y. Memorybank: Enhancing large language models with long-term memory. arXiv preprint arXiv:2305.10250 , 2023. Zhu et al. (2020) Zhu, C., Rawat, A. S., Zaheer, M., Bhojanapalli, S., Li, D., Yu, F. X., and Kumar, S. Modifying memories in transformer models. CoRR , abs/2012.00363, 2020. Appendix A Details in Methodology A.1 Self-Update Process In Section 3.1.2 , we illustrate the self-update process with the scenario of the input context x c subscript 𝑥 𝑐 x_{c} italic_x start_POSTSUBSCRIPT italic_c end_POSTSUBSCRIPT having more than K 𝐾 K italic_K tokens. As for the case when x c subscript 𝑥 𝑐 x_{c} italic_x start_POSTSUBSCRIPT italic_c end_POSTSUBSCRIPT has less than K 𝐾 K italic_K tokens, we draw the process in Figure 8 . As shown in this figure, we input e θ l superscript subscript 𝑒 𝜃 𝑙 e_{\theta}^{l} italic_e start_POSTSUBSCRIPT italic_θ end_POSTSUBSCRIPT start_POSTSUPERSCRIPT italic_l end_POSTSUPERSCRIPT and h l subscript ℎ 𝑙 h_{l} italic_h start_POSTSUBSCRIPT italic_l end_POSTSUBSCRIPT into the transformer layer ϕ l subscript italic-ϕ 𝑙 \phi_{l} italic_ϕ start_POSTSUBSCRIPT italic_l end_POSTSUBSCRIPT to obtain e θ l ′ superscript superscript subscript 𝑒 𝜃 𝑙 ′ {e_{\theta}^{l}}^{\prime} italic_e start_POSTSUBSCRIPT italic_θ end_POSTSUBSCRIPT start_POSTSUPERSCRIPT italic_l end_POSTSUPERSCRIPT start_POSTSUPERSCRIPT ′ end_POSTSUPERSCRIPT where the last n x c subscript 𝑛 subscript 𝑥 𝑐 n_{x_{c}} italic_n start_POSTSUBSCRIPT italic_x start_POSTSUBSCRIPT italic_c end_POSTSUBSCRIPT end_POSTSUBSCRIPT tokens are passed into the next layer. Figure 8: Self-Update process when the number of tokens is smaller than the number of memory tokens needed. A.2 Training Strategy for New Knowledge Incorporation Figure 9: Ideal Training Routine for Latest Knowledge Incorporation As shown in Figure 9 , compared with Section 3.2.1 , the ideal case is to perform the whole process, i.e., self-update and the prediction on x 2 subscript 𝑥 2 x_{2} italic_x start_POSTSUBSCRIPT 2 end_POSTSUBSCRIPT , with gradient flow, so that the cross-entropy loss could be backpropagated to x 1 subscript 𝑥 1 x_{1} italic_x start_POSTSUBSCRIPT 1 end_POSTSUBSCRIPT . However, this would induce unaffordable memory consumption, thus we decompose this process into two processes in Figure 2 . Appendix B Implementation Details B.1 Details for Mitigating Forgetting Problems As mentioned in Section 3.2.3 , we need to sample one main document d 𝑑 d italic_d = { x 1 , ⋯ , x n } subscript 𝑥 1 ⋯ subscript 𝑥 𝑛 \{x_{1},\cdots,x_{n}\} { italic_x start_POSTSUBSCRIPT 1 end_POSTSUBSCRIPT , ⋯ , italic_x start_POSTSUBSCRIPT italic_n end_POSTSUBSCRIPT } and multiple side documents and inject all the side documents into the memory after the injection of { x 1 , ⋯ , x n − 1 } subscript 𝑥 1 ⋯ subscript 𝑥 𝑛 1 \{x_{1},\cdots,x_{n-1}\} { italic_x start_POSTSUBSCRIPT 1 end_POSTSUBSCRIPT , ⋯ , italic_x start_POSTSUBSCRIPT italic_n - 1 end_POSTSUBSCRIPT } , then we calculate the loss on x n subscript 𝑥 𝑛 x_{n} italic_x start_POSTSUBSCRIPT italic_n end_POSTSUBSCRIPT to update the model. However, instead of sampling multiple documents at each step, we develop a more efficient strategy during training. We provide the pseudo-code in Algorithm 1 . Algorithm 1 Training Strategy for Mitigating Forgetting Problems 0:   Training data 𝒟 𝒟 \mathcal{D} caligraphic_D ; 1:   Initialize the indicator r 0 = 1 subscript 𝑟 0 1 r_{0}=1 italic_r start_POSTSUBSCRIPT 0 end_POSTSUBSCRIPT = 1 , l = 0 𝑙 0 l=0 italic_l = 0 ; 2:   Initialize the cache x c ⁢ a ⁢ c ⁢ h ⁢ e subscript 𝑥 𝑐 𝑎 𝑐 ℎ 𝑒 x_{cache} italic_x start_POSTSUBSCRIPT italic_c italic_a italic_c italic_h italic_e end_POSTSUBSCRIPT = None; 3:    for   d ∈ 𝒟 𝑑 𝒟 d\in\mathcal{D} italic_d ∈ caligraphic_D   do 4:       n = 𝑛 absent n= italic_n = the number of contexts in d 𝑑 d italic_d ; 5:       { x 1 , ⋯ , x n } = d subscript 𝑥 1 ⋯ subscript 𝑥 𝑛 𝑑 \{x_{1},\cdots,x_{n}\}=d { italic_x start_POSTSUBSCRIPT 1 end_POSTSUBSCRIPT , ⋯ , italic_x start_POSTSUBSCRIPT italic_n end_POSTSUBSCRIPT } = italic_d ; 6:       if   r 0 = = 1 r_{0}==1 italic_r start_POSTSUBSCRIPT 0 end_POSTSUBSCRIPT = = 1 or l = = 0 l==0 italic_l = = 0   then 7:          r 𝑟 r italic_r = 0; 8:       else 9:          r 𝑟 r italic_r = Random(0, 2); 10:       end   if 11:       if   r = = 0 r==0 italic_r = = 0 and r 0 = = 0 r_{0}==0 italic_r start_POSTSUBSCRIPT 0 end_POSTSUBSCRIPT = = 0   then 12:         Inject { x 1 , ⋯ , x n − 1 } subscript 𝑥 1 ⋯ subscript 𝑥 𝑛 1 \{x_{1},\cdots,x_{n-1}\} { italic_x start_POSTSUBSCRIPT 1 end_POSTSUBSCRIPT , ⋯ , italic_x start_POSTSUBSCRIPT italic_n - 1 end_POSTSUBSCRIPT } into the memory pool; 13:         Calculate the cross-entropy loss on x n subscript 𝑥 𝑛 x_{n} italic_x start_POSTSUBSCRIPT italic_n end_POSTSUBSCRIPT and update the model; 14:          l + = n limit-from 𝑙 𝑛 l+=n italic_l + = italic_n ; 15:       else   if   r = = 0 r==0 italic_r = = 0 and r 0 = = 1 r_{0}==1 italic_r start_POSTSUBSCRIPT 0 end_POSTSUBSCRIPT = = 1   then 16:         Inject { x 1 , ⋯ , x n − 1 } subscript 𝑥 1 ⋯ subscript 𝑥 𝑛 1 \{x_{1},\cdots,x_{n-1}\} { italic_x start_POSTSUBSCRIPT 1 end_POSTSUBSCRIPT , ⋯ , italic_x start_POSTSUBSCRIPT italic_n - 1 end_POSTSUBSCRIPT } into the memory pool; 17:         Calculate the cross-entropy loss on x n subscript 𝑥 𝑛 x_{n} italic_x start_POSTSUBSCRIPT italic_n end_POSTSUBSCRIPT and update the model; 18:          x c ⁢ a ⁢ c ⁢ h ⁢ e = x n subscript 𝑥 𝑐 𝑎 𝑐 ℎ 𝑒 subscript 𝑥 𝑛 x_{cache}=x_{n} italic_x start_POSTSUBSCRIPT italic_c italic_a italic_c italic_h italic_e end_POSTSUBSCRIPT = italic_x start_POSTSUBSCRIPT italic_n end_POSTSUBSCRIPT ; 19:          l + = n limit-from 𝑙 𝑛 l+=n italic_l + = italic_n ; 20:       else   if   r = = 1 r==1 italic_r = = 1   then 21:         Calculate the cross-entropy loss on x c ⁢ a ⁢ c ⁢ h ⁢ e subscript 𝑥 𝑐 𝑎 𝑐 ℎ 𝑒 x_{cache} italic_x start_POSTSUBSCRIPT italic_c italic_a italic_c italic_h italic_e end_POSTSUBSCRIPT and update the model; 22:          l = 0 𝑙 0 l=0 italic_l = 0 ; 23:       end   if 24:       r 0 = r subscript 𝑟 0 𝑟 r_{0}=r italic_r start_POSTSUBSCRIPT 0 end_POSTSUBSCRIPT = italic_r ; 25:    end   for Note that at every step, we inject the knowledge into the memory pool, thus after a random number of steps, the useful knowledge for predicting x c ⁢ a ⁢ c ⁢ h ⁢ e subscript 𝑥 𝑐 𝑎 𝑐 ℎ 𝑒 x_{cache} italic_x start_POSTSUBSCRIPT italic_c italic_a italic_c italic_h italic_e end_POSTSUBSCRIPT must be somewhere in the memory pool, we need to encourage the model to extract the relevant knowledge. If the model could extract the knowledge from the memory that was injected long ago, we could mitigate the forgetting problems. Appendix C Additional Experiments C.1 Baselines for Model Editing We introduce the details of the baselines for the model editing experiments here: FT (Finetuning): which applies Adam with early stopping at one layer to finetune the model on the given fact. FT-L (Constrained Finetuning)  (Zhu et al., 2020 ) : a parameter-space L ∞ subscript 𝐿 L_{\infty} italic_L start_POSTSUBSCRIPT ∞ end_POSTSUBSCRIPT norm constraint is imposed on the weight changes. IKE (In-context knowledge editing)  (Zheng et al., 2023 ) : The facts used to edit the model are saved in the contexts, which are inputted into the model during inference. This method is only implemented on CounterFactual so we compare our model with it on the CounterFactual benchmark. ROME (Rank-One Model Editing)  (Meng et al., 2022 ) : After identifying that MLPs in LLMs are the major modules for saving knowledge, ROME proposes to alter the MLP matrix by regarding the matrix as a key-value store and then insert a new key-value pair into the matrix, obtaining a new one that contains the injected information. Generated on Sun May 26 23:05:34 2024 by L a T e XML