MoME: Mixture-of-Memory Embeddings for Context-Aware Sparse Lookup
Muchen Li, Leonid Sigal, Renjie Liao
cs.CL, cs.AI
2026-09-14
MoME replaces each token memory row with M gated slots and beats Value Embedding, Bigram, and STEM on nanochat, Llama/MobileLLM, and Qwen3 at matched budgets.
Conditional memory is a second sparse-capacity axis beside Mixture-of-Experts: a cheap token-indexed table that stores static patterns such as names and formulaic phrases, so the backbone spends fewer early layers reconstructing a lookup table. Per-Layer Embedding in Gemma 3, Value Embedding, STEM, and Engram all take this route. Retrieval is still a deterministic function of the surface form. The programming-language python and the animal python share one row.
Capacity is allocated per token, not per meaning. cats and cat each occupy a row; bank has to hold a financial institution and a river edge in one vector. The missing primitive is a lookup that stays token-indexed and still chooses a slot from context.
MoME replaces each token row with M memory slots. Stage-one indexing still uses the token id, with an optional merge of near-duplicate tokens into one row. Stage two gates on the hidden state entering the block. Each value head independently selects K slots with TopK, aggregates them with a sigmoid-normalized mixture (softmax when K=1), and adds the result to the attention value stream through a residual gate. Default K=2, written MoME-A2/M.
Injection into the value stream lets the memory branch run beside the standard value projection; the join is one gated add. On a Qwen3-4B two-memory-layer microbenchmark that path adds 0.509 ms, while hidden-state-to-hidden-state injection adds 0.996 ms. Heads share one table to cap parameters; the gate stays per-head, so heads can read different slots. Optional kNN grouping shrinks the row count by factor c and spends the saved budget on more slots per row.
Because the first-stage indexer is swappable, MoME can sit on a Bigram hash as well. Under a matched budget the combination beats Bigram alone; at larger budgets, growing both the hash table and the slot count is the stronger setting.
Memory-augmented blocks alternate with ordinary transformer blocks.
On a 135M nanochat-style backbone at 3e18 FLOPs (about 3.3B tokens) and a 151M memory budget, MoME with grouping factor 2 records validation bpb 0.8611 and CORE 0.1583, against Value Embedding 0.8633/0.1522 and Bigram 0.8636/0.1533. Doubling memory to 302M lifts that grouping's CORE to 0.1664. Training throughput stays within 2% of iso-parameter Bigram.
The same recipe transfers at a fixed token budget. Llama/MobileLLM 125M, 55B tokens: MoME-A2/6 CORE 0.1686 versus Value Embedding 0.1530, STEM 0.1483, and the no-memory base 0.1382. At 350M and 20B tokens, MoME-A2/10 reaches CORE 0.2286 versus 0.2214 for Value Embedding. Qwen3 0.6B, 20B tokens: MoME-A2/8 CORE 0.2737 versus STEM 0.2556 and Value Embedding 0.2530; BoolQ is 63.39 for MoME and 47.55 for Value Embedding. Training wall time is 1.04-1.08 times the dense baseline.
After 100B ClimbMix tokens on nanochat d24 (0.78B backbone, 0.61B memory), MoME-A2/12 scores CORE-22 0.3687 versus 0.3484 for Value Embedding and 0.3441 for Dense. In-domain ClimbMix bpb is 0.6504, a shade worse than Value Embedding's 0.6463; FineWeb-Edu, enwik9, and Shakespeare are slightly better. Qwen3-0.6B reaches CORE-22 0.3753 after 36T tokens.
On WiC, 117 of 144 layer-head sites show higher routing-distribution divergence for different-sense pairs than for same-sense pairs, and 110 of 144 show the matching pattern in discrete slot overlap. Qualitative prompts for bank, bug, and drive reuse a slot under the same sense and jump when the sense changes.
If conditional memory only hashes surface forms, growing the table mainly enlarges the collision among senses. MoME shows that context-aware slots can sit on top of token indexing, and can share a first-stage indexer with Bigram. For a pretraining recipe that already injects a table into the value stream, the patch is small: extra slots and a gate, with almost no throughput hit.
The gain is real and modest. Matched-parameter validation bpb improves by about 0.002 over Value Embedding; CORE moves more. The 100B run still sits at 1.4B total parameters and a single seed.
Scale is the authors' main caveat: the 100B check uses an 0.8B dense backbone, 1.4B total parameters, and one seed. CORE activation maps and WiC routing are descriptive. There is no causal intervention that proves slot switching itself raises downstream accuracy or out-of-domain robustness. Grouping was swept only on nanochat; at 302M, grouping factor 4 drops CORE to 0.1515, so aggressive merges hurt. The STEM baseline follows the original 1/2-layer insertion and a much larger memory table, so the comparison is not an iso-parameter win over STEM.