Research map / Architectures and efficiency / Attention and transformers
On the Optimal Memorization Capacity of Transformers
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 Transformers memorize N sequences of length n with Õ(√N) parameters in next-token prediction (optimal up to log factors), and Õ(√(nN)) in seq-to-seq, with a matching lower bound only for hardmax Transformers.
Problem
How parameter-efficient Transformers are at memorizing N labeled input sequences of length n, i.e., their memorization capacity in terms of parameters, in next-token and sequence-to-sequence settings.
Approach
Builds a memorization construction in which a single uniform-averaging self-attention layer, combined with feed-forward networks that map an averaged sequence-id-like sum to labels, achieves Õ(√N) parameters independent of n; lower bounds are derived via VC/polynomial-partition arguments on parameters, and extended to hardmax seq-to-seq Transformers.
Key results
- Next-token prediction: Transformer memorizes N sequences with Õ(√N) parameters for d=Õ(√N), matching Ω(√N) lower bound up to log factors.
- Seq-to-seq: Õ(√(nN)) parameters suffice, and Ω(√(nN)/log(nN)) is necessary for hardmax Transformers.
- A single uniform-averaging attention layer already achieves optimal memorization capacity in next-token prediction.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Use parameter sharing with one uniform-attention layer | Transformer for next-token prediction · Next-token prediction | parameters required for memorization: d \times n -> \tilde{O}(\sqrt{N}) | Helped |
| Use a single self-attention layer as an averaging operation | Transformer for next-token prediction · Next-token prediction | parameters required for memorization: \tilde{O}(\sqrt{N}) | Helped |
| Memorize token-level labels rather than one label per sequence | Transformer · Sequence-to-sequence prediction | parameters required for memorization: \tilde{O}(\sqrt{N}) -> \tilde{O}(\sqrt{nN}) | Hurt |
| Use hardmax attention for the sequence-to-sequence lower-bound analysis | Transformer · Sequence-to-sequence prediction | parameters required for memorization (lower bound): \Omega\left(\sqrt{\frac{nN}{\log (nN)}}\right) | Hurt |
Surprising findings
- Simple uniform averaging attention suffices for optimal next-token memorization, so self-attention's advantage may lie in optimization or generalization.
- Memorization parameter count is nearly independent of input length n in next-token prediction.
Limitations
- Seq-to-seq lower bound applies only to hardmax Transformers; softmax case is open.
- Upper bound depends on embedding dimension d (Õ(d+√N) with linear in/out maps) and bit complexity grows with depth parameters.
- Lower bounds assume token-wise separated inputs and sometimes shattering over all label assignments.
Experimental setup
Datasets: two real-world datasets randomly generated dataset
Metrics: number of parameters bit complexity
Baselines: Kim et al. 2023 Transformer memorization Kajitsuka and Sato 2024 Vardi et al. 2021
Builds on: Vardi et al. memorization of ReLU networks Kim et al. Transformer memorization Bartlett VC dimension of ReLU networks Yun et al. Transformers are universal approximators
Closest papers by meaning
- Memorization Capacity of Multi-Head Attention in Transformers (ICLR 2024 · similarity 0.89)
- When Do Transformers Outperform Feedforward and Recurrent Networks? A Statistical Perspective (NeurIPS 2025 · similarity 0.88)
- kNN Attention Demystified: A Theoretical Exploration for Scalable Transformers (ICLR 2025 · similarity 0.88)
- Counting in Small Transformers: The Delicate Interplay between Attention and Feed-Forward Layers (ICML 2025 · similarity 0.88)
- Are Transformers with One Layer Self-Attention Using Low-Rank Weight Matrices Universal Approximators? (ICLR 2024 · similarity 0.87)
- Analyzing the Power of Chain of Thought through Memorization Capabilities (NeurIPS 2025 · similarity 0.87)
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.