$\text{Memory}^3$: Language Modeling with Explicit Memory (2407.01178v1)

Abstract: The training and inference of LLMs are together a costly process that transports knowledge from raw data to meaningful computation. Inspired by the memory hierarchy of the human brain, we reduce this cost by equipping LLMs with explicit memory, a memory format cheaper than model parameters and text retrieval-augmented generation (RAG). Conceptually, with most of its knowledge externalized to explicit memories, the LLM can enjoy a smaller parameter size, training cost, and inference cost, all proportional to the amount of remaining "abstract knowledge". As a preliminary proof of concept, we train from scratch a 2.4B LLM, which achieves better performance than much larger LLMs as well as RAG models, and maintains higher decoding speed than RAG. The model is named $\text{Memory}^3$, since explicit memory is the third form of memory in LLMs after implicit memory (model parameters) and working memory (context key-values). We introduce a memory circuitry theory to support the externalization of knowledge, and present novel techniques including a memory sparsification mechanism that makes storage tractable and a two-stage pretraining scheme that facilitates memory formation.

Overview of "Memory$^3$: LLMing with Explicit Memory"

The paper introduces the Memory$^3$ model, a novel approach to enhance the efficiency of LLMs by incorporating explicit memory. Inspired by the human brain's memory hierarchy, this model seeks to reduce substantial costs associated with training and inference in LLMs by externalizing specific knowledge into an explicit memory format. This memory format is presented as a cost-effective alternative to both model parameters and text retrieval-augmented generation (RAG).

Key Concepts and Methodology

The Memory$^3$ model focuses on separating knowledge into three distinct forms: implicit memory (model parameters), explicit memory, and external information. The goal is to optimize the storage and retrieval of knowledge by assigning it to the most efficient memory format based on usage frequency.

1. Memory Hierarchy for LLMs:

• Model Parameters: Store frequently used abstract knowledge.
• Explicit Memory: Suitable for moderate usage due to its moderate write and read costs.
• External Information: (RAG) Used for rare knowledge retrieval, minimizing write costs but increasing read costs.

2. Explicit Memory Design:

• Prior to inference, LLMs convert reference texts to explicit memories, reducing the computational burden during live operations.
• These memories are stored separately and retrieved as necessary, enhancing efficiency compared to traditional methods like RAG which often require real-time text processing.

3. Two-Stage Pretraining Approach:

• Warmup Stage: Initial model training without explicit memory to facilitate basic comprehension capabilities.
• Continual Train Stage: Introduces explicit memory, leveraging preprocessed references to build a more refined model.

Strong Numerical Results

The Memory$^3$ model, with 2.4B parameters, achieves superior performance compared to larger LLMs and RAG models. The explicit memory mechanism enables a smaller model to surpass state-of-the-art models in benchmark tasks and maintain higher decoding speeds, indicative of more efficient knowledge management.

Implications and Future Directions

Practical Implications:

• Reduced Training and Inference Costs: By externalizing specific knowledge, Memory$^3$ decreases the necessity for massive parameter sizes, leading to a more cost-effective training and inference process.
• Application Versatility: Facilitates quick adaptation to specialized tasks by simply updating the explicit memory bank, avoiding extensive retraining.

Theoretical Implications:

• Cognitive Alignment: The memory structure draws parallels to human cognitive processes, potentially guiding future developments in AI that mimic human-like reasoning and memory management.
• Enhanced Understanding: Provides insights into knowledge distribution and storage strategies within neural architectures.

Speculative Future Developments:

• Infinite Context Handling: Further exploration may lead to LLMs capable of handling longer contexts more efficiently, utilizing explicit memory to extend operational scopes.
• Improved Memory Consolidation Techniques: Developing methods to transition explicit memories into more permanent forms could enhance adaptability.
• Fascilitating Human-Like Reasoning: The anthropomorphic design of explicit memory might enable new reasoning capabilities that align more closely with human problem-solving.

Overall, the Memory$^3$ model represents a significant advancement in the efficient management of knowledge within LLMs, combining theoretical insights with practical benefits to push the boundaries of what is possible in AI development.