MemLLM: Finetuning LLMs to Use An Explicit Read-Write Memory (2404.11672v3)

Abstract: While current 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

• The paper introduces a method that integrates a structured explicit memory module into LLMs to improve retrieval of infrequent and time-sensitive information.
• It leverages relational triples and vector similarity indexing to manage dynamic memory-read and memory-write operations efficiently.
• Evaluation on the DOCRED dataset shows reduced perplexity and balanced recall-precision, enhancing factual adherence and reducing hallucinations.

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

The paper "MemLLM: Finetuning LLMs to Use An Explicit Read-Write Memory" (2404.11672) introduces a method for augmenting LLMs with a structured explicit memory module. This approach aims to address limitations in current LLMs related to the handling of infrequent knowledge and temporal degradation due to the reliance on implicit parametric storage.

Motivation and Background

While LLMs are proficient in numerous knowledge-intensive tasks, they often fall short when tasked with recalling infrequent or temporally sensitive information. Prior methods such as parametric memory pools, model editing, and retrieval-augmented generation (RAG) have highlighted some challenges. However, these methods encounter limitations, such as parametric distortions, interpretability issues, and inefficient retrieval processes. In contrast, the proposed MAuLLM method integrates a structured read-write memory component allowing efficient and dynamic interaction, thus aiming to enhance both performance and interpretability.

Methodology

The core innovation of MAuLLM is the integration of a structured memory component which the LLM can interact with explicitly. This memory is formatted as relational triples, which facilitates the LLM’s understanding and manipulation of stored information. The memory operations are controlled via a dedicated API enabling both memory-read and memory-write commands during inference.

Memory Structure

The structured memory is organized as triples: $\langle e_s, t, e_o \rangle$, where $e_s$ is the subject, $t$ is the relation, and $e_o$ is the object. The system relies on efficient retrieval techniques using vector similarities to locate relevant memory entries, leveraging known methods like Hierarchical Navigable Small World graph-based indexing.

Figure 1

Figure 1: The prompt for the distant supervision dataset filtering. This prompt includes the natural representation of the relation, the reasoning, and the final answer.

Memory-read and Memory-write Operations

• Memory Write: The LLM assesses each sentence in the text for potential relational information to store. Each sentence undergoes processing to extract relevant triples, which are then written to the memory using a defined API command.
• Memory Read: When generating text, the model can invoke memory-read operations to query relevant stored information. These queries are dynamically generated in response to the context, and the retrieved information aids in more factually accurate text generation.

Evaluation

The performance of MAuLLM was assessed using the DOCRED dataset, known for its comprehensive relation annotations derived from Wikipedia. In experiments, MAuLLM demonstrated improved performance over models without explicit memory features, particularly in handling entity-related perplexities. Notably, using augmented memories correctly improved target entity prediction, indicating better factual adherence and reduced hallucination risks.

Performance Analysis:

• Perplexity Reduction: MAuLLM achieved significant reductions in perplexity on entity targets, showcasing the READ/WRITE memory's efficacy in improving LLM accuracy.
• Recall and Precision: By refining the distant supervision input and filtering strategies, the model achieved a balanced gain in recall without a loss in precision, allowing for richer relational context processing.

Conclusions and Future Work

MAuLLM represents a robust step towards more capable and interpretable LLMs by incorporating a structured memory mechanism. This explicit memory format not only allows for scalable knowledge management but also enhances the model's adaptability across tasks requiring reliable knowledge retention over time.

Future research directions could involve adapting this framework to more diverse knowledge-domain applications or exploring alternate relational schemas that support complex or nested knowledge structures, extending the current methodology to broader and more dynamic real-world contexts.