MoME: Mixture-of-Memory Embeddings for Context-Aware Sparse Lookup
Muchen Li, Leonid Sigal, Renjie Liao
cs.CL · cs.AI
Quick verdict
MoME adds context-gated slot mixtures to token-indexed memory embeddings, beating VE/STEM/Bigram at sub-billion scale in controlled runs; gains are real but small and mostly sub-billion.
Summary
The paper proposes MoME, a context-aware conditional-memory mechanism for language models. Prior memory-embedding methods (Per-Layer Embedding, Value Embedding, STEM, Engram) index a learnable table deterministically by token identity or n-gram hash, so one token with several senses shares a single memory vector. MoME keeps the cheap token-indexed lookup but replaces each row with M slots and adds a learned gate over the hidden state that top-K selects and sigmoid-normalizes slot weights per value head, injecting the aggregated memory into the attention value stream through a residual gate. Optional kNN-based token grouping compresses rows and reinvests the saved capacity into slots. Controlled pretraining experiments cover a nanochat-style backbone (3e18 FLOPs, three seeds), Llama/MobileLLM 125M/350M, and Qwen3-0.6B at matched token budgets. MoME improves validation bpb and CORE over Value Embedding, Bigram, and STEM in most matched comparisons (e.g., Qwen3-0.6B CORE 0.2737 vs VE 0.2530), shows a better memory-size scaling trend than Bigram up to ~604M memory parameters, and can combine with Bigram or sparse-FFN MoE. Latency overhead falls from ~7.8% at 350M to ~1.6% at 8B scale. WiC-based routing analyses find sense-sensitive routing at 117/144 layer-head sites. Main limitations: small absolute gains (about 0.002-0.015 bpb), many key comparisons at a single seed, and no causal evidence that routing helps accuracy.
Review complete
Assessment at a glance
4
risk flags identified
● high 0
● medium 2
● low 2
5 strengths · 5 weaknesses · 5 author questions
AI text signal = Binoculars zero-shot detector score over 20 sampled passages (Fast-DetectGPT co-reported) — absence does not prove human authorship, and a flag does not imply misconduct.
Levels reflect verification complexity — not a quality score. Evolve never rates your paper.
Key claims under review
- Replacing each token-indexed memory row with a hidden-state-gated mixture of M slots improves validation bpb and CORE over Value Embedding, Bigram, and STEM in most iso-parameter and iso-FLOP comparisons at sub-billion scale.
- MoME exhibits a consistently better validation-bpb trajectory than matched-memory Bigram across the tested 75.5M-604M memory-size range on the nanochat d12 backbone.
- The context-aware routing is compatible with Bigram indexing and sparse-FFN MoE, and combining mechanisms improves over either alone at matched budgets.
- Inference latency overhead is modest and decreases with scale, from ~7.8% at 350M parameters to ~1.6% at 8B, due to parallel scheduling with the value projection.
- The learned slot router is sense-sensitive: same-sense occurrences route to the same slots more than different-sense occurrences, quantified on 670 filtered WiC pairs (117/144 sites positive for ΔJSD).
Strengths
Controlled, multi-backbone evaluation design
Table 2 (nanochat d12, 3e18 FLOPs, three seeds 42-44) holds memory parameters fixed at 151M/302M across Bigram, VEmbedding, and three MoME grouping factors; Table 3 repeats token-matched comparisons at 125M/350M/0.6B scales with wall-time columns, covering both iso-parameter and iso-FLOP angles.
Baselines and ablations are unusually thorough for a memory paper
Beyond VE/STEM/Bigram, Appendix A.8.3 (Table 16) adapts canonical Engram, JTok, JTok-M, and MoVE-Input at matched parameters; Table 18 isolates the context-conditioning contribution (hidden-state routing vs learned token-only vs fixed-random, with MoME losing 0.0054/0.0028 bpb and up to 0.0173 CORE without context); Table 19 adds dense width/depth-matched controls.
Routing interpretability is quantified, not just illustrated
Engineering-facing efficiency reporting
Honest limitations section
Weaknesses
Small absolute margins relative to run-to-run noise
Single-seed comparisons at the largest scales
Throughput regression vs the strongest bpb baseline
CORE metric moves are noisy and task-level results are mixed
Interpretability claims rest on curated examples plus a heavily filtered probe set
Full Review Report
$1.99 · one-timeUnlock the complete preliminary analysis Evolve produced for this paper.
Report preview
Scope & Claims
MoME (arXiv 2609.15126) targets a specific, well-scoped gap in the emerging memory-embedding literature: existing token-indexed memory tables (Per-Layer Embedding in Gemma 3n, Value Embedding, STEM, Engram) retrieve deterministically from the surface form, so a polysemous token collapses into one memory vector. The proposed fix keeps the cheap first-stage lookup (token ID or n-gram hash) and adds a second, context-aware stage: each row holds M slots; a linear gate on the hidden state selects top-K slots per value head with sigmoid-normalized weights; the aggregate is injected into the attention value stream via a residual gate γ=2σ(W h) (Sections 3.3-3.5). An optional kNN token-grouping map (Section 3.4) compresses rows by c_grp and reinvests capacity into slots. The claims are deliberately modest: competitive or better bpb/CORE at matched budgets, a better memory-scaling trend than Bigram, combinability with Bigram/MoE, low latency overhead, and sense-sensitive routing — all at sub-billion scale, as the limitations section itself states.
The remaining sections cover Methodology Assessment, Statistical Evidence, Results & Reproducibility, Limitations & Risks, and Verdict & Recommendations.
- All 5 strengths & 5 weaknesses, analyzed in detail
- 4 risk flags with severity, section & reviewer questions
- Methods & statistical evidence assessment
- 5 questions a peer reviewer would ask
- Full six-dimension evaluation report & recommendations
One-time purchase · this paper only
One-time purchase, this paper only. AI-generated preliminary analysis — not peer review.
Frequently asked questions
What is MoME (Mixture of Memory Embeddings)?
MoME is a modification to memory-augmented language models. Instead of giving each vocabulary token one fixed memory vector that is looked up by token ID, each row of the memory table holds M slots, and a learned gate reading the current hidden state picks which slots (top-K) to read at each position, weighting them and injecting the result into the attention value stream. This lets the same surface token retrieve different memory depending on context (e.g., 'python' as a language vs. a snake) while keeping the cheap deterministic lookup structure.
Does MoME beat Value Embedding and Mixture-of-Experts baselines?
In the paper's controlled pretraining runs (up to ~0.8B dense parameters, 3.3B-100B tokens), MoME generally achieves lower validation bits-per-byte and higher CORE scores than Value Embedding, Bigram, and STEM at matched memory parameters or matched tokens — e.g., on Qwen3-0.6B it reaches CORE 0.2737 vs 0.2530 for Value Embedding. Gains are small in absolute bpb terms (~0.002-0.015), and the largest comparisons are single-seed. It is not compared head-to-head with full MoE as a replacement; it is shown to combine with a sparse-FFN MoE.
How much does MoME slow down training and inference?
Training takes 1.04-1.08x the wall time of the no-memory base in the token-matched runs (Table 3), and full-sequence scoring throughput in the 100B-token run is ~9% lower than Value Embedding. Inference overhead per memory-augmented block falls with model size: about +7.8% for a 350M model, +5.5% at 1B, +2.35% at 4B, and +1.56% at 8B, because the memory branch is scheduled in parallel with the value projection (Tables 5 and 21).
Is the interpretability claim (polysemous tokens routing to different slots) statistically supported?
Partially. Beyond curated examples (bank, bug, drive, apple, python), the paper tests sense-sensitive routing on 670 strictly filtered Word-in-Context pairs across 144 layer-head sites: routing distributions diverge more for different-sense pairs at 117/144 sites, with the largest effect (ΔJSD 0.0928) having a bootstrap 95% CI of [0.049, 0.136]. Caveats acknowledged by the authors: no multiple-comparison correction across 144 sites, a noun-skewed filtered set, and no evidence the routing differences causally improve accuracy.
About this review
This is an AI-generated preliminary review by Evolve.Pub, produced from the openly available text of the paper. It is not peer review and does not represent the authors' views. The canonical record of the paper is at the publisher link above.
- Figures 2-9 are only available as embedded plot descriptions in the extracted text; trends were verified against the corresponding tables (17, 21, A.9) where possible.
- The extraction contains some garbled inline LaTeX/base64 artifacts in figures; no numeric values were taken from those regions.
- Appendix Table 15's full 22-task breakdown was included, but a few appendix subsections (A.9.4 gate heatmaps) are figure-only and could not be numerically audited.
- Code is claimed open-sourced ('The @ code and pretrained models are open-sourced'); the repository URL was not verifiable from the text.
Go deeper on this paper
Wrote a paper? Stress-test it before reviewers do.
Upload your manuscript for the same structured scrutiny — reasoning, statistics, evidence, figures, and reproducibility.