Research map / Architectures and efficiency / Attention and transformers
Memory Mosaics
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.
Introduces Memory Mosaics, networks of kernel-smoothing associative memories that match GPT2-scale transformers on BabiStories language modeling and outperform them on out-of-distribution and in-context learning tasks.
Problem
Transformers' compositional and in-context learning abilities are hard to interpret; the paper asks whether a more transparent architecture built from associative memories can match transformers on language modeling while offering better out-of-distribution and in-context learning.
Approach
Networks of associative memories (Nadaraya-Watson kernel regression) where values may peek one step into the future and keys are leaky averages; training is viewed as meta-learning that induces "predictive disentanglement", splitting the prediction task across independent memory units.
Key results
- On BabiStories, Memory Mosaics slightly outperform GPT2-style transformers at small depth, with training and validation losses becoming indistinguishable as depth increases.
- On RegBench in-context learning, Memory Mosaics outperform transformers, RNNs and state-space models across training set sizes spanning three orders of magnitude.
- On out-of-distribution Simple English Wikipedia, Memory Mosaics substantially outperform transformers after about 50 tokens, whereas the transformer plateaus after 100 to 150 tokens.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Use Memory Mosaic instead of RWKV | RegBench IID (100 training environments) · RegBench IID (100 training environments) | accuracy: 0.967 -> 0.959 | Hurt |
| Use Memory Mosaic instead of RWKV | RegBench IID (100 training environments) · RegBench IID (100 training environments) | TVD: 0.183 -> 0.417 | Hurt |
| Use Memory Mosaic instead of transformer | Language modeling on BabiStories, small-depth models · BabiStories | loss | Helped |
| Use Memory Mosaic instead of transformer | Language modeling on BabiStories, increasing depth · BabiStories | loss | No clear effect |
| Use Memory Mosaic instead of transformer | OOD prediction on Simple English Wikipedia after about 50 tokens · Simple English Wikipedia | per-token average loss | Helped |
| Use Memory Mosaic instead of previously tested architectures | In-context learning on held-out PFAs · RegBench OOD | prediction accuracy | Helped |
Surprising findings
- A single-layer Memory Mosaic can implement an induction head, which needs two attention layers in a transformer, and it works without position encoding or query/key distinction.
- Hyperparameters tuned for transformers transferred verbatim to Memory Mosaics and still matched them, which the authors note disadvantages Memory Mosaics.
Limitations
- Experiments are at medium scale (GPT2-small size) and the authors acknowledge larger-scale replication is needed.
- Hyperparameters were tuned for transformers and transferred to Memory Mosaics, possibly disadvantaging them.
- Predictive disentanglement is demonstrated only on a toy three-moon example and conjectured for transformers.
Experimental setup
Datasets: BabiStories TinyStories Simple English Wikipedia RegBench
Benchmarks: BabiStories RegBench Simple English Wikipedia 24 narrative consistency prompts
Metrics: validation loss prediction accuracy total variation distance per-token loss vs position
Baselines: GPT2-small transformer Mamba S4 RWKV linear transformer H3 GLA Hyena LSTM RetNet RoPE ALiBi
Builds on: transformers memory networks TinyStories RegBench
Closest papers by meaning
- Memory Mosaics at scale (NeurIPS 2025 · similarity 0.93)
- Scaling Laws for Associative Memories (ICLR 2024 · similarity 0.84)
- Understanding Factual Recall in Transformers via Associative Memories (ICLR 2025 · similarity 0.83)
- Memorization Capacity of Multi-Head Attention in Transformers (ICLR 2024 · similarity 0.82)
- Factual recall in linear associative memories: sharp asymptotics and mechanistic insights (NeurIPS 2026 · similarity 0.82)
- Geometric Factual Recall in Transformers (NeurIPS 2026 · similarity 0.82)
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.