Mixture of memory embeddings improves context aware token lookups
MoME: Mixture-of-Memory Embeddings for Context-Aware Sparse Lookup
Computation and LanguageArtificial Intelligence
Summary
Large language models try to understand words by looking up stored information called embeddings. Usually, each word has one fixed entry, which makes it hard to capture different meanings of the same word in different contexts. The authors present Mixture of Memory Embeddings (MoME), which stores multiple entries per word and learns which to use based on the context. Experiments show MoME performs better than previous methods while using similar resources and can separate different meanings of words effectively.
What this means in practice
- •For ai model engineers: Build language models that better distinguish word meanings by dynamically selecting context-aware memory slots for embeddings.
- •For mobile app developers: Implement efficient small-scale language models that maintain rich contextual understanding using mixtures of memory embeddings.
Authors
Muchen Li, Leonid Sigal, Renjie Liao
Abstract
Scaling large language models efficiently has motivated sparse capacity mechanisms such as Mixture-of-Experts and, more recently, conditional memory: token-indexed embedding tables that augment the backbone with cheap parametric lookups. Existing memory-embedding methods retrieve via a deterministic function of the surface form, which collapses different contextual senses of the same token (e.g., python the language vs. the animal) into a single fixed entry. We introduce Mixture of Memory Embeddings (MoME), a context-aware memory mechanism that replaces each token's single memory row with a mixture of M slots and uses a learned gate over the hidden state to choose which slots to read at each position. In controlled pretraining experiments across nanochat, Llama-3/MobileLLM, and Qwen3 backbones, MoME improves over Value Embedding, Bigram, and STEM baselines in iso-parameter and iso-training-FLOP settings, shows a more promising memory-size scaling trend at sub-billion scale, and remains efficient in training and inference. Qualitative routing analyses on polysemous tokens further suggest that the learned mixture exhibits a degree of semantic interpretability, dispatching the same surface token to distinct memory slots under different senses.