MoME: Mixture-of-Memory Embeddings for Context-Aware Sparse Lookup
Muchen Li, Leonid Sigal, Renjie Liao
University of British Columbia · Vector Institute for AI
Published September 14, 2026 · on the feed September 21, 2026 · ▲ 12 · arXiv 2609.15126
The 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.
The Brief
2:09A short podcast summary of the paper. Written and voiced by AI.
Transcript
MoME. Mixture-of-Memory Embeddings for Context-Aware Sparse Lookup. By Muchen Li and two others, from the University of British Columbia and the Vector Institute for A I. Published on the fourteenth of September. Its code repository is worked out: about sixty-five source files, with a readme and a licence.
Scaling large language models efficiently has pushed toward sparse capacity, first mixture-of-experts and more recently conditional memory: embedding tables indexed by token, which add cheap parametric lookups to the backbone. The trouble lies in how existing memory-embedding methods retrieve. They use a deterministic function of the surface form, so every contextual sense of a token collapses into one fixed entry. Python the programming language and python the animal receive the same row.
Mixture of Memory Embeddings, or MoME, makes that lookup context-aware. Each token's single memory row is replaced by a mixture of several slots, and a learned gate over the hidden state chooses which slots to read at each position. The same token can therefore reach different memory depending on where it appears.
Little else changes. The memory is still a table keyed by the token, and it still sits beside the backbone as a cheap source of extra parameters. What differs is that the table now offers several slots per token rather than one, and the choice among them is made by the gate, reading the model's hidden state at that position rather than the token's spelling alone.
The experiments are controlled pretraining runs across three backbones: nanochat, Llama-3 with MobileLLM, and Qwen3. Against Value Embedding, Bigram, and STEM baselines, MoME improves in settings matched for parameter count and for training compute. It shows a more promising trend as memory size grows, at sub-billion scale, and remains efficient in both training and inference. Qualitative analyses of routing on polysemous tokens further suggest that the learned mixture carries a degree of semantic interpretability, sending the same surface token to distinct memory slots under different senses.
The change is narrow and its motive is plain. A memory keyed only on spelling cannot tell senses apart, so MoME lets context choose which part of the memory a token reads, while keeping the lookup efficient.
Paper Quest · 7 chapters · about 40 min
Play this paper
Learn it by doing: make the calls the authors made, guess the result before you see it, and defend it in a viva graded on your own words. Chapter 1 is right here.
Chapter 1 🧩 · The problem
A cheap memory table with one blind spot: it cannot tell a snake from a language.
Your language model has a cheap add-on: a big lookup table where each token id fetches a learned row of extra parameters for the backbone. Only a few rows are touched per token, so it scales cheaply.
The road ahead
- 1 · The problem ~7 min A cheap lookup table gives a language model extra parameters, but it picks rows by spelling alone, so it cannot tell word senses apart.
- 2 · What came before ~3 min
- 3 · The key idea ~8 min Keep the token-keyed table, but give each token several memory slots and let a gate reading the hidden state pick which ones to use.
- 4 · Choose your own research ~3 min
- 5 · Build the method ~8 min The token picks its row of slots, a gate over the hidden state scores them, the top few are mixed, and the result joins the backbone.
- 6 · The experiment ~6 min You’ll guess: In matched pretraining comparisons (memory size, training compute or tokens), how did MoME compare with the Value Embedding, Bigram and STEM baselines?
- 7 · The viva ~5 min Three questions from the examiners, answered in your own words. Win it for the paper’s trophy.
Free, with Google or an email address. Your stars and progress are kept.