MoME: Mixture-of-Memory Embeddings for Context-Aware Sparse Lookup

AI summary

Overview

Research area: Natural Language Processing / efficient large language model architecture, specifically conditional memory and memory-embedding augmentation of Transformers.

Technical level: Advanced. The paper assumes familiarity with Mixture-of-Experts routing, Transformer attention value streams, and pretraining evaluation metrics such as bits-per-byte (bpb) and the CORE score.

Scope: The paper introduces Mixture of Memory Embeddings (MoME), a context-aware replacement for deterministic token-indexed memory lookups, and evaluates it across three backbone families against Value Embedding, Bigram, and STEM baselines. Authors are Muchen Li, Leonid Sigal, and Renjie Liao (University of British Columbia, Vector Institute for AI, Canada CIFAR AI Chair, NSERC CRC Chair); the paper is arXiv:2609.15126v1 [cs.CL], dated 14 Sep 2026, licensed CC BY 4.0.

What This Paper Is About

Sparse conditional memory—auxiliary token-indexed embedding tables that a Transformer can look up cheaply—is a growing alternative to dense scaling, but existing methods retrieve memory with a deterministic function of the surface form. That means a token like python (the language vs. the animal) or bank (a financial institution vs. a river edge) is forced into one fixed memory vector regardless of context. MoME's goal is to keep the cheap token-indexed lookup while letting the model's hidden state decide which of several stored memory slots to read.

Key Contributions

1. A context-aware memory module. MoME replaces each token's single memory row with a mixture of M slots and adds a learned gate over the hidden state that selects which slots to read at each position, bringing context awareness into the token-indexed lookup of memory-embedding methods.

2. Cross-architecture evidence. Controlled pretraining experiments on nanochat-style, Llama-3/MobileLLM-style, and Qwen3-style backbones show improvements over Value Embedding, Bigram, and STEM baselines in iso-parameter and iso-training-FLOP settings, with the mechanism transferring across model families.

3. Memory-size scaling and composability. MoME shows a more favorable observed memory-size scaling trend than Bigram over the tested memory range at sub-billion scale, and can be combined with Bigram's first-stage indexer under compound scaling for additional gains.

4. Routing interpretability. Qualitative and quantitative routing analyses on polysemous tokens indicate the learned mixture dispatches the same surface token to distinct memory slots under different senses, rather than committing one route per token. Code and pretrained models are open-sourced.

Main Findings

• Nanochat iso-FLOP results. On the 12-layer 135M nanochat-style backbone at a fixed 3×10^18 training FLOPs (~3.3B tokens), the Base model reaches train/val bpb of 0.8783/0.8785 and CORE 0.1447; Bigram (151M/286M parameters) reaches 0.8632/0.8636 and CORE 0.1533; VEmbedding (151M/286M) reaches 0.8631/0.8633 and CORE 0.1522. MoME at c_grp=1 (151M/286M) reaches 0.8618/0.8621 and CORE 0.1571; c_grp=2 reaches 0.8609/0.8611 and CORE 0.1583; c_grp=4 reaches 0.8608/0.8610 and CORE 0.1554.

• Doubling memory helps all variants. At a 302M memory budget, MoME reaches val bpb 0.8569 (c_grp=1, CORE 0.1635), 0.8561 (c_grp=2, CORE 0.1664), and 0.8565 (c_grp=4, CORE 0.1515).

• Grouping factor sweep. Sweeping c_grp in {1, 2, 4} at fixed parameter count, c_grp=2 gives the highest CORE at both memory budgets and the lowest train and validation bpb at 302M; at 151M, c_grp=4 gives slightly lower train and validation bpb.

• Throughput overhead is small. MoME stays within <2% of the iso-parameter Bigram throughput. Throughput figures are 1.15M tok/s for Base, 1.09M for Bigram, 1.10M for VEmbedding, and 1.07M for the 151M-memory MoME runs.

• Llama/MobileLLM 125M. Base val bpb 0.8852 / CORE 0.1382; STEM (755M memory) 0.8633 / 0.1483; VEmbedding (95M) 0.8620 / 0.1530; MoME-A1/3 (95M) 0.8623 / 0.1587; MoME-A2/6 (189M) 0.8560 / 0.1686.

• Llama/MobileLLM 350M. Base 0.7750 / 0.1988; STEM (1342M) 0.7697 / 0.2179; VEmbedding (168M) 0.7651 / 0.2214; MoME-A2/5 (168M) 0.7647 / 0.2202; MoME-A2/10 (336M) 0.7610 / 0.2286.

• Qwen3 0.6B. Base 0.7594 / 0.2491; STEM (1409M) 0.7564 / 0.2556; VEmbedding (470M) 0.7427 / 0.2530; MoME-A2/8 (470M) 0.7461 / 0.2737. The paper states MoME improves both validation bpb and CORE over STEM across all three scales while using fewer memory parameters.

• Training wall time. MoME requires 1.04–1.08× the training wall time of the no-memory base across the Llama/MobileLLM and Qwen3 families.

• Memory-size scaling. Fixing the backbone to nanochat d12 under an iso-FLOP setting, MoME achieves lower validation bpb than the matched-memory Bigram row at every tested memory size, and shows substantially lower run-to-run variance in both training and validation bpb curves.

• Compound scaling with Bigram. Using the bigram indexer as MoME's first-stage indexer and sharing one memory table across augmented layers, the combination improves over Bigram at the same parameter budget. The best values at each memory budget in the compound scaling table are 0.8576/0.8582 at 302M (12V×6) and 0.8540/0.8546 at 604M (12V×12), versus Bigram's 0.8632/0.8636 at 151M.

• Inference latency. Measured per-step latency increase relative to Base: Llama/MobileLLM 350M 7.82% (value injection) vs. 8.05% (hidden-state injection); 1B 5.47% vs. 5.37%; Qwen3 0.6B 7.90% vs. 6.44%; 4B 2.35% vs. 4.60%; 8B 1.56% vs. 2.23%. In the qwen3_4b two-memory-layer benchmark, value-stream injection adds 0.509 ms, while the hidden-state-injection variant adds 0.996 ms.

• 100B-token ClimbMix training (nanochat d24). MoME-A2/12 (0.78B network, 0.61B memory) achieves ClimbMix 0.6504, FineWeb-Edu 0.7548, enwik9 0.8850, Shakespeare 1.4789, and CORE-22 0.3687, versus Dense (0.78B) at 0.6639 / 0.7677 / 0.9037 / 1.4993 / 0.3441 and VEmbedding (0.78B network, 0.60B memory) at 0.6463 / 0.7549 / 0.8860 / 1.4823 / 0.3484. MoME obtains lower bpb than Dense and VEmbedding on all three out-of-domain evaluations and a higher CORE-22, while its in-domain ClimbMix bpb is slightly higher than VEmbedding's. External references: Qwen3-0.6B (36T tokens) CORE-22 0.3753, Llama-3.2-1B (≤9T tokens) CORE-22 0.3629. Scoring throughput is 58.4k tok/s for MoME, 64.0k for VEmbedding, 43.8k for Qwen3-0.6B, and 51.7k for Llama-3.2-1B.

• Semantic routing. On a Qwen backbone at Layer 7 / Head 7, same-sense prompts for bank, bug, and drive route through the same memory slot while the changed-sense prompt jumps to a different slot. The paper notes that n-gram-based routing has too short a window to disambiguate these senses, while MoME conditions on the full prefix.

• Quantitative routing analysis. Using the Word-in-Context (WiC) dataset, 670 strictly filtered pairs are evaluated at all 144 memory layer–head sites in the d24 checkpoint, using Jensen–Shannon divergence and the chance-corrected top-2 decision overlap Õ of Olson et al. [44].

• Depth profile of memory use. Measured on the Table 6 model, both the gated memory residual magnitude and the mean injection gate are strongest near the beginning and especially the end of the model, with substantially lower utilization through the middle layers.

Methodology in Plain English

Existing memory-embedding methods store one vector per token (or per fixed n-gram hash) and always fetch that same vector. MoME keeps the lookup table indexed by token, but gives each row several "slots" — a tensor of shape N × M × d_value, where N is the number of addressable rows, M is the slots per row, and d_value is the per-head value dimension at the injection point. At each position, a small learned gate reads the current hidden state (which already encodes context from earlier layers) and produces scores over the M slots; a TopK operation picks K of them, and the selected memory vectors are combined into one aggregated vector. Weighting differs by K: a sigmoid-norm gate over the selected slots when K>1, and softmax when K=1.

That aggregated memory is added to the attention value stream through a separate per-head gate, scaled by 2·sigmoid(...) so that a zero-initialized logit yields a neutral residual scale of 1. Crucially, the memory branch runs in parallel with the standard value projection, so the only required join is the gated residual addition — a design choice the authors show reduces latency relative to injecting into hidden states. All value heads share a single memory table (capping per-layer memory parameters at N·M·d_value rather than N·M·H·d_value), but each head selects its own slot subset and weights.

Optionally, the token indexer f can group semantically similar tokens into shared rows using offline kNN matching in a lightly pretrained token-embedding space (reference taken from a nanochat-d12 base run of about 1B tokens and about 1e18 FLOPs). Grouping shrinks the row dimension by the factor c_grp and the freed capacity is offloaded into additional slots (M scaled by c_grp). All experiments use FineWeb-Edu with 524,288-token batches, except the 100B-token runs and d24 controls, which use ClimbMix with 1,048,576-token batches; all use 2048-token sequences and a shared 32k BPE tokenizer, with Muon for matrix-shaped parameters and AdamW for embeddings, unembeddings, scalars, value-memory tables, and other non-matrix parameters.

Why This Matters

Impact on research. The paper isolates a limitation shared by a whole family of memory-embedding methods — Per-Layer Embedding, Value Embedding, STEM, Engram — namely that addressing is a fixed function of the surface form even when fusion is context-dependent. It shows that making the lookup itself context-aware is compatible with cheap, sparse retrieval, and that the resulting router produces semantically structured behavior rather than arbitrary routing. It also demonstrates compatibility with existing first-stage indexers such as Bigram, suggesting memory mechanisms can be scaled compositionally rather than treated as mutually exclusive.

Real-world applications (implied by the paper's framing of retrieval-heavy regularities such as named entities and formulaic phrases):

• On-device or edge language models, where the latency measurements (down to 1.56% overhead at the 8B Llama/MobileLLM scale) and throughput figures matter.
• Retrieval of static, local patterns such as entities and formulas that Transformers otherwise spend early-layer capacity reconstructing.
• Serving systems that need to scale memory capacity with modest additional compute rather than scaling dense weights.
• Any polysemy-sensitive downstream use — the routing analysis directly targets words whose meaning changes with context.

Industry relevance. The paper emphasizes efficiency throughout: the memory branch adds negligible runtime overhead within the nanochat iso-FLOP study, keeps training wall time at 1.04–1.08× the no-memory base, and open-sources code and pretrained models. The observed latency advantage from value-stream injection over hidden-state injection, and the shrinking relative overhead as backbones grow, are directly relevant to deployment decisions.

Future Directions

• Scaling beyond sub-billion scale. The memory-size scaling study is run on nanochat d12 in a low-compute regime, and the 100B-token run uses a d24 backbone. Whether the favorable scaling trend versus Bigram persists at larger backbones and memory budgets is not established.

• Extending and automating token grouping. The grouping function relies on offline kNN matching in a lightly pretrained embedding space with c_grp in {2, 4} as the main configuration; the paper does not report learning the grouping end-to-end or exploring grouping at larger factors.

• Combining memory mechanisms further. MoME is shown to work with the Bigram indexer, and the best compound-scaling variant scales both the Bigram table and MoME slots. The space of other first-stage indexers, and the optimal allocation between table size and slot count, remains open.

• Making the routing analysis causal and broader. The routing analyses cover curated polysemous tokens plus 670 WiC pairs, with the WiC evaluation at 144 layer–head sites in one checkpoint. Whether the sense-sensitive routing produces measurable downstream gains on word-sense tasks is not reported.

Target Audience

Researchers and engineers working on efficient LLM pretraining, sparse capacity mechanisms, and conditional memory architectures — particularly those already familiar with Mixture-of-Experts routing, value-embedding augmentation, or n-gram memory designs. The paper is also relevant to practitioners evaluating whether to add memory lookup tables to a production backbone, since it reports parameter counts, throughput, training wall time, and inference latency overheads alongside quality metrics. Readers without a background in Transformer internals will find the architecture sections dense, though the motivation and routing-interpretability sections are accessible.

Authors’ 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.