Memorization and Knowledge Injection in Gated LLMs

TL;DR

MEGa presents a biologically inspired approach to lifelong memory in LLMs by embedding memories directly into model weights via gated LoRA adapters. A query-driven gating mechanism selects relevant memories at inference, enabling both memory recall and knowledge-based QA while preserving prior general knowledge. Across Fictional Character and Wikipedia 2024 Events, MEGa outperforms standard continual-learning baselines on recall and QA, with iRAG further enhancing QA performance and compositional reasoning. The work offers a scalable, memory-centric alternative to external retrieval systems, with implications for robust lifelong learning and personalized AI agents.

Abstract

Large Language Models (LLMs) currently struggle to sequentially add new memories and integrate new knowledge. These limitations contrast with the human ability to continuously learn from new experiences and acquire knowledge throughout life. Most existing approaches add memories either through large context windows or external memory buffers (e.g., Retrieval-Augmented Generation), and studies on knowledge injection rarely test scenarios resembling everyday life events. In this work, we introduce a continual learning framework, Memory Embedded in Gated LLMs (MEGa), which injects event memories directly into the weights of LLMs. Each memory is stored in a dedicated set of gated low-rank weights. During inference, a gating mechanism activates relevant memory weights by matching query embeddings to stored memory embeddings. This enables the model to both recall entire memories and answer related questions. On two datasets - fictional characters and Wikipedia events - MEGa outperforms baseline approaches in mitigating catastrophic forgetting. Our model draws inspiration from the complementary memory system of the human brain.

Figures (9)

• Figure 1: Illustration of our fine-tuning and inference pipeline. We store the embedding of each sample as its context key and initialize a LoRA adapter for each sample (i.e., raw text paragraphs). LoRA adapters are fine-tuned individually using a constructed user prompt. During inference, the base model processes the query to generate its embedding, which is then used to compute the LoRA gating weights. Questions can be augmented with additional instructions depending on the task. For the recall task, the model is prompted to reconstruct the entire relevant story. For the QA task, the model directly answers the question, with the instruction “Answer should be no more than one sentence” appended to ensure concise responses. In the internal RAG (iRAG) setup, the model first reconstructs the relevant story and then answers the question based on the reconstructed story. All inputs are formatted according to the Llama-3 instruct format.
• Figure 2: Recall forgetting curve comparing MEGa and other approaches. The model is prompted to recall the fine-tuned sample associated with a given question. Cosine similarity is computed between the recalled text and the original fine-tuning sample. The “Random pairs” baseline represents the average cosine similarity between embeddings of two randomly selected samples (0.465 for the Fictional Character dataset and 0.226 for the Wiki Events dataset).
• Figure 3: QA accuracy forgetting curve comparing MEGa and other approaches. The curve shows the QA accuracy for questions related to the first sample in the sequence, measured as the model is sequentially trained on additional samples. As more samples are introduced, accuracy on the first sample declines. For the Fictional Character dataset, the base model’s QA accuracy is 0.13%, overlapping with the x-axis.
• Figure 4: MMLU accuracy during continual fine-tuning.
• Figure 5: We found that the effectiveness of injecting new knowledge into instruct model depends on the "fine-tuning prompt". The category "fact" represents the prompts (n=15) that indicate the finetuning sample is a fact, for example "Please tell me a story that you memorized:"; the category "generic" represents the prompts (n=7) that do not imply the factuality of the sample, for example "Generate:"; the category "fake" represents the prompts (n=10) that indicate the sample is fictional, for example "Please make up a fake story:". We also have categories "random strings" (n=24), "spaces" (n=4), and "special characters" (n=10). The accuracy is the GPT judge's results on 750 questions related to 250 samples in each dataset. Fine-tuning is done per sample (no continual learning). The error bar represents the standard deviation among prompts.