Language model improves knowledge updates with mixed memory system

Mixture-of-Memories: Low-Overhead Language Modeling with Modular Knowledge Injection

Computation and Language

Summary

Language models have lots of facts stored inside them, but changing or controlling these facts can be tricky. The authors propose a system called Mixture-of-Memories that uses two different types of memory to better handle facts: one for quick internal updates and one for token-level knowledge. Their design also includes a smart router that decides how much to trust each memory when predicting words. This system makes it easier to update facts without needing to retrain the whole model and improves accuracy without slowing down responses.

What this means in practice

  • For ai engineers: Add or update factual knowledge in language models efficiently without retraining the whole system, improving accuracy in applications like question answering.
  • For chatbot developers: Enable chatbots to incorporate up-to-date facts during conversations with minimal delay and easy factual corrections.

Authors

Yifan Lu, Rubin Wei, Ziyun Zou, Junhong Liang, Haolin Yang, Xinlin Zhuang, Abhijit Das, Adinath Madhavrao Dukre, Shadab Khan, Zhouhan Lin, Imran Razzak

Abstract

Large language models store substantial factual knowledge in their parameters, yet this knowledge is difficult to control, update, and reliably elicit at inference time. Existing memory-augmented approaches typically rely on a single auxiliary mechanism, such as retrieval, parameter updates, or neural memory, and therefore apply the same form of correction to tokens with very different knowledge demands. We propose Mixture-of-Memories (MoM), a memory-augmented language modeling framework that separates routed neural memory from writable factual memory. MoM combines a frozen base LM with two neural memory experts: a latent memory that provides low-latency hidden-state corrections and a lexical memory that provides a token-space knowledge pathway. A lightweight Budget Router sets per-token mixture weights from uncertainty and disagreement features, replacing the fixed interpolation of prior memory modules. For post-hoc updates, an external explicit memory stays outside the router, because a router trained before facts are written learns to down-weight it; new facts are appended without retraining any component and used only when an answer-level gate accepts the retrieved fact. Experiments across language modeling, factual QA, update analysis, and efficiency show that MoM improves both predictive accuracy and updateability while retaining low inference overhead.