Research map / Architectures and efficiency / Attention and transformers
Understanding Factual Recall in Transformers via Associative Memories
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.
Shows theoretically and empirically that linear/MLP associative memories and a one-layer transformer (attention and/or MLP) can store factual associations at near-linear-in-parameter capacity, analyzes gradient-flow hallucination dynamics, and gives information-theoretic lower bounds.
Problem
How transformers optimally encode factual associations in their weights, i.e., what storage capacity shallow transformers can achieve for factual recall, and how training dynamics unfold.
Approach
Model factual recall via associative memories built from outer products of random sphere embeddings; show that a one-layer transformer can store SR facts using either attention value matrices or the MLP (trade-off), with near-optimal capacity matched by information-theoretic lower bounds.
Key results
- Linear associative memory stores Theta~(d^2) associations and MLP associative memory stores Theta~(md) associations, empirically matching theory.
- A one-layer attention-only or attention+MLP transformer achieves 100% accuracy on the synthetic factual recall task when attention or MLP parameters scale linearly (up to logs) with the number of facts SR.
- Information-theoretic lower bounds show Omega(SR log D) bits are needed for zero loss, so the constructions are near-optimal up to log factors.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Use attention-only storage with sufficient embedding dimension and attention capacity | One-layer transformer · Synthetic factual recall | accuracy: 100\% | Helped |
| Use an MLP associative memory with sufficient width and parameter count | One-layer transformer · Synthetic factual recall | accuracy: 100\% | Helped |
| Increase associative-memory parameters with random sphere embeddings | Linear associative memory · Discrete-vocabulary association storage | storage capacity: d^2 \gtrsim N\poly\log N | Helped |
| Increase MLP associative-memory parameters with random sphere embeddings | MLP associative memory · Discrete-vocabulary association storage | storage capacity: md \gtrsim N\poly\log N | Helped |
| Scale total model parameters with the number of facts | One-layer transformer · Synthetic factual recall | number of facts stored | Helped |
| Trade MLP parameters for self-attention parameters while retaining sufficient attention capacity | One-layer transformer · Synthetic factual recall, S = 32, R = 32, D = 8 | accuracy | No clear effect |
| Train the linear-attention model through the intermediate relation-only stage | Linear attention model with orthogonal embeddings · Synthetic factual recall, S = 16, R = 4, D = 8 | loss: \log\abs{\gA} to \log D to zero loss | Mixed |
| Train softmax attention with random embeddings | Single softmax attention head · Synthetic factual recall | loss | Mixed |
Surprising findings
- Gradient flow on a linear attention model exhibits an intermediate 'hallucination' plateau where the model predicts answers from the relation token alone.
Limitations
- Results are for synthetic one-to-one factual recall with shallow (one-layer) transformers and random embeddings.
- Optimization analysis uses linear attention with orthogonal embeddings, a simplification of the softmax model.
- Constructions require polylog factors and the learned models are not shown to coincide with these constructions in practice.
Experimental setup
Datasets: synthetic factual recall task (subject, relation, noise tokens, answers)
Benchmarks: synthetic factual recall task
Metrics: accuracy cross-entropy loss maximum number of facts stored at 99% accuracy
Baselines: one-hot embedding associative memory
Builds on: associative memories / Hopfield networks Allen-Zhu and Li factual knowledge capacity neural tangent kernel
Closest papers by meaning
- Memorization Capacity of Multi-Head Attention in Transformers (ICLR 2024 · similarity 0.89)
- Do LLMs dream of elephants (when told not to)? Latent concept association and associative memory in transformers (NeurIPS 2024 · similarity 0.89)
- Geometric Factual Recall in Transformers (NeurIPS 2026 · similarity 0.89)
- An Evolved Universal Transformer Memory (ICLR 2025 · similarity 0.89)
- Scaling Laws for Associative Memories (ICLR 2024 · similarity 0.89)
- Factual recall in linear associative memories: sharp asymptotics and mechanistic insights (NeurIPS 2026 · similarity 0.88)
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.