Research map / Architectures and efficiency / Attention and transformers

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

NeurIPS 2026 · Architectures and efficiency › Attention and transformers · approach: Sparse attention for long-context LLMs · 0 citations (OpenAlex)

Read the paperCheck your idea against 108,967 papers

AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.

Proposes MoME, a context-aware conditional memory that replaces each token's single memory row with M slots chosen by a learned gate over the hidden state, and tests it in controlled pretraining against Value Embedding, Bigram, and STEM baselines.

Problem

Token-indexed memory embeddings retrieve a deterministic, context-blind vector per token, so different senses of a polysemous token (e.g., "python") share one memory entry.

Approach

Each token row holds M memory slots and a learned per-head gate over the hidden state selects top-K slots and mixing weights, injected into the attention value stream via a gated residual; optional kNN token grouping shares rows among similar tokens to free capacity for more slots.

Key results

• At iso-151M memory on nanochat d12 (3e18 FLOPs), MoME (c_grp=2) reaches val bpb 0.8611 and CORE 0.1583 vs Bigram 0.8636/0.1533 and VEmbedding 0.8633/0.1522; at 302M memory the best MoME variant reaches val bpb 0.8561 and CORE 0.1664.
• On Llama/MobileLLM 125M, MoME-A2/6 (189M) reaches val bpb 0.8560 and CORE 0.1686 vs STEM 0.8633/0.1483 and VEmbedding 0.8620/0.1530; on Qwen3 0.6B, MoME-A2/8 reaches CORE 0.2737 vs VEmbedding 0.2530.
• In the 100B-token ClimbMix d24 run, MoME-A2/12 gets CORE-22 0.3687 vs Dense 0.3441 and VEmbedding 0.3484, but its in-domain ClimbMix bpb (0.6504) is slightly worse than VEmbedding (0.6463).

Surprising findings

• In the 100B-token run, MoME has slightly higher in-domain ClimbMix bpb than Value Embedding despite winning on out-of-domain bpb and CORE-22.
• Combining MoE-FFN with MoME gives the highest CORE-22 but a slightly worse validation bpb than MoE-FFN alone.

Limitations

• Computational scale is limited; the 100B-token scale check uses a single seed and a ~0.8B dense backbone.
• Routing interpretability (CORE activation and WiC analyses) is descriptive; causal interventions are not shown to improve downstream accuracy or robustness.
• Many comparisons are at small scale with mixed token budgets and some gains are small (e.g., CORE differences within error bars at nanochat scale).

Experimental setup

Datasets: FineWeb-Edu ClimbMix WiC enwik9 Shakespeare CORE-22 evaluation suite

Benchmarks: CORE-22 ARC-Easy ARC-Challenge BoolQ PIQA HellaSwag OpenBookQA WinoGrande BigBench tasks ClimbMix FineWeb-Edu enwik9

Metrics: validation bpb train bpb CORE raw accuracy training wall time inference latency JSD top-2 decision overlap

Baselines: Base Value Embedding Bigram STEM Dense Qwen3-0.6B Llama-3.2-1B MoE-FFN

Builds on: Value Embedding STEM Engram Mixture-of-Experts Muon

Closest papers by meaning

• Scaling Embedding Layers in Language Models (NeurIPS 2025 · similarity 0.88)
• Language Model Memory and Memory Models for Language (NeurIPS 2026 · similarity 0.87)
• Frozen Memory Is Not Enough: Rethinking External Memory as Extraction (NeurIPS 2026 · similarity 0.86)
• Data Engineering for Scaling Language Models to 128K Context (ICML 2024 · similarity 0.86)
• Memory Layers at Scale (ICML 2025 · similarity 0.86)
• Mixture-of-Recursions: Learning Dynamic Recursive Depths for Adaptive Token-Level Computation (NeurIPS 2025 · similarity 0.86)