MEMORYLLM: Towards Self-Updatable Large Language Models

TL;DR

MemoryLLM tackles the challenge of updating LLM knowledge without full retraining by embedding a fixed-size memory pool in the latent space and enabling self-updates across layers. The approach uses a static φ backbone and a dynamic θ memory pool with a lightweight self-update mechanism to absorb new information while gradually forgetting older content. Through model editing benchmarks, long-context QA, and knowledge-retention tests, MemoryLLM demonstrates superior knowledge integration, strong retention, and robustness under nearly one million updates. The work provides a scalable, open-source pathway toward self-updating LLMs with practical implications for continual learning and deployment.

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.

Figures (9)

• Figure 1: The framework of MemoryLLM. (a) During generation, all memory tokens in the $l$-th layer of memory pool $\theta^l$ are attended by the hidden states $h_l$. (b) During self-update, The last $k$ memory tokens from $\theta_l$ are taken to be concatenated with the hidden states $h_l$ as the input to $\phi_l$. The output $h_{l+1}$ goes to the next layer. The last $K$ tokens of $h_{l+1}$ serve as the new memory tokens ${e_\theta^l}'$. We randomly drop $K$ tokens in $\theta_l$ and concatenate the left $\theta_l$ (denoted as $\theta^l(d)$) with ${e_\theta^l}'$ to obtain new memory $\theta_l'$.
• Figure 2: Training Process for new knowledge incorporation. During training, we randomly choose one of two shown processes to proceed with 50% probability each. The description pertains to the first layer, and the subsequent layers share an analogous procedure. After sampling $(x_1,x_2)$ from the dataset, we first perform self-update with $x_1$ as depicted in the left side of both processes. Subsequently, the modified memory ${e_\theta^1}'$ is employed to predict $x_2$. Of the two processes, the upper one maintains gradient flow throughout the entire process, optimizing the knowledge compression from $x_1$ to ${e_\theta^l}'$ ($l \in \{1,\cdots,L\}$). In contrast, the lower process executes the self-update without gradient. Both processes are designed to encourage the use of the knowledge in the memory pool for the prediction.
• Figure 3: Training process for continuous contexts understanding. We only draw two self-update steps here with $x_1, x_2$ though there should be $n-1$ self-updates in this training iteration. We show the procedure of $l$-th layer here. At the bottom of the figure, $h_1^n$ refers to the word embeddings of $x_n$, and $h_L^n$ is used for loss value calculation. Essentially we are compressing the knowledge from $x_1,\cdots,x_{n-1}$ into $\theta^{l^{n-1}}$ to predict $x_n$.
• 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$, we truncate $4096$ tokens from the given context as input to the model. The y-axis is the F1 score.
• 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.