Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

MemLLM: Finetuning LLMs to Use An Explicit Read-Write Memory

Ali Modarressi, Abdullatif Köksal, Ayyoob Imani, Mohsen Fayyaz, Hinrich Schütze

TL;DR

MemLLM introduces an explicit, structured read-write memory for LLMs and trains it via finetuning with a formal memory API. The approach uses a relational memory schema and retrieval by vector similarity to augment language modeling and enable robust knowledge editing. Experiments show improved perplexity, especially on target entities, and favorable performance on knowledge-editing tasks compared to baselines, with scalable memory size and interpretable results. The work provides a reusable training dataset and identifies limitations (restricted relation types) and avenues for future expansion and integration with broader grounding techniques.

Abstract

While current large language models (LLMs) perform well on many knowledge-related tasks, they are limited by relying on their parameters as an implicit storage mechanism. As a result, they struggle with memorizing rare events and with updating their memory as facts change over time. In addition, the uninterpretable nature of parametric memory makes it challenging to prevent hallucination. Model editing and augmenting LLMs with parameters specialized for memory are only partial solutions. In this paper, we introduce MemLLM, a novel method of enhancing LLMs by integrating a structured and explicit read-and-write memory module. MemLLM tackles the aforementioned challenges by enabling dynamic interaction with the memory and improving the LLM's capabilities in using stored knowledge. Our experiments indicate that MemLLM enhances the LLM's performance and interpretability, in language modeling in general and knowledge-intensive tasks in particular. We see MemLLM as an important step towards making LLMs more grounded and factual through memory augmentation. The project repository is publicly available at https://github.com/amodaresi/MemLLM

MemLLM: Finetuning LLMs to Use An Explicit Read-Write Memory

TL;DR

MemLLM introduces an explicit, structured read-write memory for LLMs and trains it via finetuning with a formal memory API. The approach uses a relational memory schema and retrieval by vector similarity to augment language modeling and enable robust knowledge editing. Experiments show improved perplexity, especially on target entities, and favorable performance on knowledge-editing tasks compared to baselines, with scalable memory size and interpretable results. The work provides a reusable training dataset and identifies limitations (restricted relation types) and avenues for future expansion and integration with broader grounding techniques.

Abstract

While current large language models (LLMs) perform well on many knowledge-related tasks, they are limited by relying on their parameters as an implicit storage mechanism. As a result, they struggle with memorizing rare events and with updating their memory as facts change over time. In addition, the uninterpretable nature of parametric memory makes it challenging to prevent hallucination. Model editing and augmenting LLMs with parameters specialized for memory are only partial solutions. In this paper, we introduce MemLLM, a novel method of enhancing LLMs by integrating a structured and explicit read-and-write memory module. MemLLM tackles the aforementioned challenges by enabling dynamic interaction with the memory and improving the LLM's capabilities in using stored knowledge. Our experiments indicate that MemLLM enhances the LLM's performance and interpretability, in language modeling in general and knowledge-intensive tasks in particular. We see MemLLM as an important step towards making LLMs more grounded and factual through memory augmentation. The project repository is publicly available at https://github.com/amodaresi/MemLLM
Paper Structure (19 sections, 1 equation, 5 figures, 4 tables, 1 algorithm)

This paper contains 19 sections, 1 equation, 5 figures, 4 tables, 1 algorithm.

Table of Contents

  1. Introduction
  2. Related work
  3. Methodology
  4. Memory structure
  5. Memory-API and Inference
  6. Memory-API and Inference
  7. Finetuning the LLM
  8. Experiments
  9. Setup
  10. Filtering Distant Supervision Relations
  11. Perplexity evaluation
  12. Knowledge Editing Evaluation
  13. Qualitative Analysis.
  14. Conclusion
  15. Memory-write Decoding Method
  16. ...and 4 more sections

Figures (5)

  • Figure 1: MemLLM memory schema. Each triple is stored in the "Triple Memory" using its subject entity ID, relation ID and object entity ID, along with a designated triple ID. Entities and relations are stored in separate tables, each containing their designated ID, name and corresponding vector embedding. Both tables are indexed based on their vector embeddings.
  • Figure 2: MemLLM inference with memory read and memory write
  • Figure 3: MemLLM training data pipeline.
  • Figure 4: Memory efficiency of storing structured triples vs. proposition-based storage. The y-axis represents the fraction of memory required compared to a RAG system that stores an embedding per fact.
  • Figure 5: The prompt for the distant supervision dataset filtering. This prompt includes the natural representation of the relation, the reasoning, and the final answer.