Research map / Architectures and efficiency / Attention and transformers
Memory Mosaics at scale
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.
Scales Memory Mosaics (associative-memory networks) to llama-8B scale (9.9B params, 1T tokens) with three architectural changes (adaptive bandwidth, gated time-variant key extractor, 3-level memory), showing better new-knowledge storage and in-context learning than matched transformers, including transformers trained on 8x more data.
Problem
Whether the compositional and in-context learning advantages of Memory Mosaics, shown only at GPT-2 scale on small synthetic data, persist when scaled to LLM sizes and real-world data.
Approach
Replace transformer attention with associative-memory blocks using symmetric L2-normalized keys, add adaptive kernel bandwidth, a gated time-variant key extractor, and a 3-level memory split into short-term, long-term, and persistent memory; evaluate on persistent-knowledge, new-knowledge (RULER multi-document QA), and in-context classification dimensions.
Key results
- On RULER question-answer tasks at 32k task length after 32k fine-tuning, memory mosaics v2 large reaches 53.4% vs 41.1% for a transformer trained on 1T tokens.
- A transformer trained on 8T tokens reaches 46.9% on the same task, still below memory mosaics v2 trained on 1T tokens (53.4%).
- On 19 common benchmarks, memory mosaics v2 large averages 52.2 equal to transformer large 52.2; on 13 persistent-knowledge benchmarks it averages 56.8 vs 57.1.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Use Memory Mosaics v2 instead of a transformer | Large model · RULER multi-unrelated-documents QA, task-length 32k | score: 41.1 vs 53.4; 12.3% | Helped |
| Use Memory Mosaics v2 instead of a transformer | Large model · Persistent-knowledge benchmarks (13-task avg) | average score: 57.1 vs 56.8 | No clear effect |
| Train the transformer on 8T rather than 1T tokens | Transformer large · RULER multi-unrelated-documents QA, task-length 32k | score: 41.1 -> 46.9 | Helped |
| Train the transformer on 8T rather than 1T tokens | Transformer · Anonymous-label in-context classification | accuracy | Hurt |
| Remove long-term memory after training (ablation) | Memory Mosaics v2 large · Persistent-knowledge benchmarks (13-task avg) | average score: 56.8 vs 56.6 | No clear effect |
| Add long-term memory (ablation) | Memory Mosaics v2 large · Persistent-knowledge benchmarks (6-task avg) | average score: 34.9 -> 42.1 | Helped |
| Add long-term memory (ablation) | Memory Mosaics v2 large · RULER QA, task-length 32k | score: 20.2 -> 53.4 | Helped |
| Use stochastic long-term memory size during training (ablation) | Memory Mosaics v2 small · RULER QA, task-length 32k, no fine-tuning | score: 15.9 -> 31.7 (+15.8) | Helped |
| Augment the transformer with long-short-term attention (ablation) | Transformer · In-context classification | accuracy | No clear effect |
| Use Memory Mosaics v2 rather than a transformer | Large model · Semantic-label in-context classification | accuracy: more than 10% | Helped |
| Provide more demonstration shots | Memory Mosaics v2 · In-context classification | accuracy | Helped |
| Fine-tune for one mini-batch | Memory Mosaics v2 · RULER QA, task-length 32k | accuracy: 22% | Helped |
Surprising findings
- Augmenting a transformer with long-short term attention does not improve in-context learning, suggesting the gains come from the architecture as a whole.
- A single fine-tuning mini-batch on memory mosaics v2 yields a 22% accuracy improvement on the 32k RULER task, while transformers need 800 mini-batches to lag behind.
Limitations
- Hyperparameters were tuned for transformers and transferred to memory mosaics v2 without search.
- Performance degrades at 64k task length (46.4%) and the model does not address long-context compute cost.
- Several baselines (RNNs, state-space models) were excluded by design rather than empirically compared on the new-knowledge tasks.
Experimental setup
Datasets: diverse pretraining datamix (unnamed) RULER Banking77 Tacred Goemotion
Benchmarks: obqa arc easy winogrande arc challenge piqa boolq hellaswag nq siqa tqa gsm8k mmlu
Metrics: accuracy
Baselines: transformer large transformer small rwkv-v5-7b mamba-2.8b-slimpj llama2-7b mistral-7b-v0.2-base qwen-1.5-7b-base
Builds on: Memory Mosaics transformer llama architecture RULER
Closest papers by meaning
- Memory Mosaics (ICLR 2025 · similarity 0.93)
- Scaling Laws for Associative Memories (ICLR 2024 · similarity 0.86)
- Memory Layers at Scale (ICML 2025 · similarity 0.86)
- Beyond Attention: Breaking the Limits of Transformer Context Length with Recurrent Memory (AAAI 2024 · similarity 0.86)
- Memory Grafting: Scaling Language Model Pre-training via Offline Conditional Memory (NeurIPS 2026 · similarity 0.86)
- Language Model Memory and Memory Models for Language (NeurIPS 2026 · similarity 0.85)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.