Memory$^3$: Language Modeling with Explicit Memory
The training and inference of large language models (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.4 B 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 ${\rm 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.
Authors
- Hao Wu (ORCID: https://orcid.org/0000-0002-1951-6651)
- Wenqiang Wei (ORCID: https://orcid.org/0000-0003-2078-9056)
- Bo Tang (ORCID: https://orcid.org/0000-0001-7129-0250)
- Chenyang Xi (ORCID: https://orcid.org/0000-0002-8211-2430)
- E Weinan
- Zehao Lin (ORCID: https://orcid.org/0000-0002-4726-7867)
- Hongkang Yang (ORCID: https://orcid.org/0000-0003-4938-6445)
- Feiyu Xiong (ORCID: https://orcid.org/0000-0002-1456-2202)
- Kai Chen (ORCID: https://orcid.org/0000-0002-5624-2987)
- Wenjin Wang (ORCID: https://orcid.org/0000-0003-3942-5308)
- Yu Yu (ORCID: https://orcid.org/0000-0002-3109-7859)
- Jinbo Wang (ORCID: https://orcid.org/0000-0002-3527-1551)
- Shichao Song (ORCID: https://orcid.org/0009-0005-4575-9592)
- Linpeng Tang null
- Zhiyu Li (ORCID: https://orcid.org/0000-0003-3771-0199)
- Zeyun Tang
Publication Details
- Journal
- Journal of Machine Learning
- Published
- 2024-09-12
- DOI
- https://doi.org/10.4208/jml.240708
- Citations
- 4
- Primary Topic
- Topic Modeling
- Type
- article
- Field-Weighted Citation Impact
- 1.39