Self-attention Does Not Need $O(n^2)$ Memory
Open accessAbstract
We present a very simple algorithm for attention that requires $O(1)$ memory with respect to sequence length and an extension to self-attention that requires $O(\log n)$ memory. This is in contrast with the frequently stated belief that self-attention requires $O(n^2)$ memory. While the time complexity is still $O(n^2)$, device memory rather than compute capability is often the limiting factor on modern accelerators. Thus, reducing the memory requirements of attention allows processing of longer sequences than might otherwise be feasible. We provide a practical implementation for accelerators that requires $O(\sqrt{n})$ memory, is numerically stable, and is within a few percent of the runtime of the standard implementation of attention. We also demonstrate how to differentiate the function while remaining memory-efficient. For sequence length 16384, the memory overhead of self-attention is reduced by 59X for inference and by 32X for differentiation.
The claims and statistics of this paper haven’t been checked yet.
Run the check — free, no account needed. It reads the open-access text and checks the sources it cites.
Integrity signals
Registry
Cites
Statistics
Claims the paper makes
The claims and statistics of this paper haven’t been checked yet.
Run the check — free, no account needed. It reads the open-access text and checks the sources it cites.
Claims other papers make about this one
Each is checked against this paper’s own text: does it say what the citing paper says it does?
Naive memory-efficient attention still incurs quadratic HBM accesses, resulting in slower execution speed.
“This technique [60] has been used in the literature [51, 66] to show that attention computation does not need quadratic extra memory (though the number of HBM accesses is still quadratic, resulting in slow run-time).”
From FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness (2022) · cites this paper as [66]
The Introduction (p. 1) says the memory-efficient algorithm avoids quadratic memory use but still has quadratic time complexity. The Inference section (p. 4) reports roughly similar runtime to standard attention; it does not establish quadratic HBM accesses or attribute slow runtime to them. (LLM, conf=0.86)
Experimental — accuracy not yet measuredFlashAttention has smaller total memory requirement compared to Rabe and Staats [66] because it incrementally updates the output after processing each block, needing only one copy instead of K copies for K blocks.
“FlashAttention instead incrementally updates the output (Algorithm 1 line 12) after processing each block, so only one copy of the output is needed (instead of 𝐾 copies for 𝐾 blocks). This means that FlashAttention has smaller total memory requirement compared to Rabe and Staats [66].”
From FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness (2022) · cites this paper as Rabe and Staats [66]
Section 4 (p. 4) says the practical implementation keeps one summary per key/value chunk before rescaling them, producing √n summaries and O(√n) memory overhead. That supports the comparison with storing K intermediate results, but the paper does not establish the claimed smaller total memory requirement for FlashAttention. It also describes an O(log n)-memory self-attention algorithm in the Abstract (p. 1), so the comparison cannot be applied to all of its methods. (LLM, conf=0.84)
Experimental — accuracy not yet measuredFlashAttention simplifies the backward pass analytically, only recomputes the attention matrix and does not recompute the temporary output of each block, reducing memory requirement for the backward pass and yielding speedup compared to Rabe and Staats [66].
“FlashAttention instead simplifies the backward pass analytically (Appendices B.2 and B.4). It only recomputes the attention matrix and does not recompute the temporary output of each block. This reduces the memory requirement for the backward pass and yields speedup.”
From FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness (2022) · cites this paper as Rabe and Staats [66]
Split into 2 atoms — 0 confirmed, 2 partially supported, 0 disputed, 0 misattributed, 0 unverifiable. Expand the claim to see per-atom verdicts.
Experimental — accuracy not yet measuredBoth FlashAttention and Rabe and Staats [66] offer substantial memory saving in terms of total memory required.
“In terms of total memory required, both methods offer substantial memory saving.”
From FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness (2022) · cites this paper as Rabe and Staats [66]
The backward pass of attention can be done without quadratic extra memory by applying gradient checkpointing to the memory-efficient forward pass.
“Rabe and Staats [66] suggests that the backward pass can be done without quadratic extra memory by applying gradient checkpointing to the memory-efficient forward pass.”
From FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness (2022) · cites this paper as Rabe and Staats [66]
FlashAttention is faster than standard attention (2-4×) while Rabe and Staats [66] is around the same speed or slightly slower than standard attention.
“As a result, FlashAttention is faster than standard attention (2-4 ) while Rabe and Staats [66] is around the same speed or slightly slower than standard attention.”
From FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness (2022) · cites this paper as Rabe and Staats [66]
Rabe and Staats [66] focuses on reducing total memory footprint while FlashAttention focuses on reducing memory accesses.
“The first major difference is that Rabe and Staats [66] focuses on the reducing the total memory footprint (maximum amount of GPU memory required) while FlashAttention focuses on reducing memory accesses (the number of memory reads/writes).”
From FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness (2022) · cites this paper as Rabe and Staats [66]
Self-check
Percentages, sums and reported differences re-computed from the paper’s own numbers.
The self-check needs the paper’s full text.
Cites
Rethinking Attention with Performers
Choromanski, K., Likhosherstov, V., Dohan, D., Song, X., Gane, A., Sarlos, T., Hawkins, P., Davis, J., Mohiuddin, A., Kaiser, L., et al. (2020) · arXiv · doi:10.48550/arxiv.2009.14794 · Cited by 695 · see citing works
Flagged first, then by citations
Research works printed but not linked: 18
- Dzmitry Bahdanau, Kyunghyun Cho, and Yoshua Bengio. Neural machine translation by jointly learning to align and translate. In Yoshua Bengio and Yann LeCun (eds.), ICLR , 2015.
- Tianqi Chen, Bing Xu, Chiyuan Zhang, and Carlos Guestrin. Training deep nets with sublinear memory cost. CoRR , abs/1604.06174, 2016. · arXiv
- Rewon Child, Scott Gray, Alec Radford, and Ilya Sutskever. Generating long sequences with sparse transformers. CoRR , abs/1904.10509, 2019. URL http://arxiv.org/abs/1904.10509 . · arXiv
- Tri Dao, Daniel Y. Fu, Stefano Ermon, Atri Rudra, and Christopher Ré. Flashattention: Fast and memory-efficient exact attention with io-awareness, 2022. URL https://arxiv.org/abs/2205.14135 . · arXiv
- Jacob Devlin, Ming-Wei Chang, Kenton Lee, and Kristina Toutanova. BERT: pre-training of deep bidirectional transformers for language understanding. In Jill Burstein, Christy Doran, and Thamar Solorio (eds.), Proceedings of the 2019 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, NAACL-HLT 2019, Minneapolis, MN, USA, June 2-7, 2019, Volume 1 (Long and Short Papers) , pp. 4171-4186. Association for Computational Linguistics, 2019. doi: 10.18653/v1/n19-1423. URL https://doi.org/10.18653/v1/n19-1423 . · DOI · Paper page
- Curtis Hawthorne, Andrew Jaegle, C˘ at˘ alina Cangea, Sebastian Borgeaud, Charlie Nash, Mateusz Malinowski, Sander Dieleman, Oriol Vinyals, Matthew Botvinick, Ian Simon, Hannah Sheahan, Neil Zeghidour, Jean-Baptiste Alayrac, João Carreira, and Jesse Engel. General-purpose, long-context autoregressive modeling with perceiver AR, 2022. URL https://arxiv.org/abs/2202.07765 . · arXiv
- Hanhwi Jang, Joonsung Kim, Jae-Eon Jo, Jaewon Lee, and Jangwoo Kim. Mnnfast: A fast and scalable system architecture for memory-augmented neural networks. In 2019 ACM/IEEE 46th Annual International Symposium on Computer Architecture (ISCA) , pp. 250-263, 2019.
- Nikita Kitaev, Łukasz Kaiser, and Anselm Levskaya. Reformer: The efficient transformer. arXiv preprint arXiv:2001.04451 , 2020. · arXiv
- Xuezhe Ma, Xiang Kong, Sinong Wang, Chunting Zhou, Jonathan May, Hao Ma, and Luke Zettlemoyer. Luna: Linear unified nested attention. CoRR , abs/2106.01540, 2021. URL https://arxiv.org/abs/2106.01540 . · arXiv
- Jiezhong Qiu, Hao Ma, Omer Levy, Wen-tau Yih, Sinong Wang, and Jie Tang. Blockwise self-attention for long document understanding. In Findings of the Association for Computational Linguistics: EMNLP 2020 , pp. 25552565, 2020.
- Hongyu Ren, Hanjun Dai, Zihang Dai, Mengjiao Yang, Jure Leskovec, Dale Schuurmans, and Bo Dai. Combiner: Full attention transformer with sparse computation cost. arXiv preprint arXiv:2107.05768 , 2021. · arXiv
- Aurko Roy, Mohammad Saffar, Ashish Vaswani, and David Grangier. Efficient content-based sparse attention with routing transformers. Trans. Assoc. Comput. Linguistics , 9:53-68, 2021. URL https://transacl.org/ojs/index.php/tacl/article/view/2405 .
- Zhuoran Shen, Mingyuan Zhang, Haiyu Zhao, Shuai Yi, and Hongsheng Li. Efficient attention: Attention with linear complexities. In IEEE Winter Conference on Applications of Computer Vision, WACV 2021, Waikoloa,
- HI, USA, January 3-8, 2021 , pp. 3530-3538. IEEE, 2021. doi: 10.1109/WACV48630.2021.00357. URL https://doi.org/10.1109/WACV48630.2021.00357 . · DOI
- Yi Tay, Mostafa Dehghani, Samira Abnar, Yikang Shen, Dara Bahri, Philip Pham, Jinfeng Rao, Liu Yang, Sebastian Ruder, and Donald Metzler. Long range arena : A benchmark for efficient transformers. In 9th International Conference on Learning Representations, ICLR 2021, Virtual Event, Austria, May 3-7, 2021 . OpenReview.net, 2021. URL https://openreview.net/forum?id=qVyeW-grC2k .
- Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Łukasz Kaiser, and Illia Polosukhin. Attention is all you need. In Advances in neural information processing systems , pp. 5998-6008, 2017.
- Sinong Wang, Belinda Z Li, Madian Khabsa, Han Fang, and Hao Ma. Linformer: Self-attention with linear complexity. arXiv preprint arXiv:2006.04768 , 2020. · arXiv
- Manzil Zaheer, Guru Guruganesh, Kumar Avinava Dubey, Joshua Ainslie, Chris Alberti, Santiago Ontañón, Philip Pham, Anirudh Ravula, Qifan Wang, Li Yang, and Amr Ahmed. Big bird: Transformers for longer sequences. In Hugo Larochelle, Marc'Aurelio Ranzato, Raia Hadsell, Maria-Florina Balcan, and Hsuan-Tien Lin (eds.), Advances in Neural Information Processing Systems 33: Annual Conference on Neural Information Processing Systems 2020, NeurIPS 2020, December 6-12, 2020, virtual , 2020. URL https://proceedings.neurips.cc/paper/2020/hash/c8512d142a2d849725f31a9a7a361ab9-Abstract.html .
Cited by
Sorted by citations
FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness
Dao, T., Fu, D. Y., Ermon, S., Rudra, A., & Ré, C. (2022) · arXiv · doi:10.48550/arxiv.2205.14135 · Cited by 1,214 · see citing works
“Conceptually, both FlashAttention and Rabe and Staats [66] operate on blocks of the attention matrix using the well-established technique of tiling (or softmax scaling) [51, 60].”
BackgroundMethodResultInfluentialAhdritz, G., Bouatta, N., Floristean, C., Kadyan, S., Xia, Q., Gerecke, W., O’Donnell, T. J., Berenberg, D., Fisk, I., Zanichelli, N., et al. (2024) · Nature Methods · doi:10.1038/s41592-024-02272-z · Cited by 439 · see citing works
“…of neural network attention mechanisms (Vaswani et al. 2017) with favorable properties for unusually short and long sequences (Rabe and Staats 2021, Dao et al. 2022), module refactoring for lower memory usage, optional approximations of certain computations that reduce the memory…”
BackgroundFlashAttention-2: Faster Attention with Better Parallelism and Work Partitioning
Dao, T. (2023) · arXiv · doi:10.48550/arxiv.2307.08691 · Cited by 143 · see citing works
“As the softmax couples entire rows or blocks of row, online softmax [11, 13] can split the attention computation into blocks, and rescale the output of each block to finally get the right result (with no approximation).”
MethodCaution Ahead: Numerical Reasoning and Look‐Ahead Bias in AI Models
LEVY, B. (2026) · Journal of Accounting Research · doi:10.1111/1475-679x.70058 · Cited by 20 · see citing works
RWKV: Reinventing RNNs for the Transformer Era
Peng, B., Alcaide, E., Anthony, Q., Albalak, A., Arcadinho, S., Biderman, S., Cao, H., Cheng, X., Chung, M., Grella, M., et al. (2023) · arXiv · doi:10.48550/arxiv.2305.13048 · Cited by 17 · see citing works
“Some recent works like FlashAttention (Dao et al., 2022a) and others (Rabe and Staats, 2022; Jang et al., 2019) share similarities with RWKV’s chunked computation scheme.”
Method
- Authors
- Rabe, M. N., & Staats, C.
- Affiliations
- Google (United States) (US)
- ORCID
- Rabe, Markus N.
- Venue
- arXiv
- Publisher
- arXiv
- Published
- 2021
- ISSN
- 2331-8422 (ISSN-L)
- arXiv
- 2112.05682
- OpenAlex
- W4226001818
- Cited by
- 17
- References
- 1
- Citations by year
- 2026: 2 · 2025: 3 · 2024: 5 · 2023: 6 · 2022: 1
- Type
- posted content
- Language
- en
- Open access
- green
- arxiv.org
- Locations
- arXiv (Cornell University) (submittedVersion)
- arXiv (Cornell University)
- Topics
- Physical Sciences › Computer Science › Artificial Intelligence › Topic Modeling
- Physical Sciences › Computer Science › Hardware and Architecture › Parallel Computing and Optimization Techniques
- Physical Sciences › Computer Science › Computer Vision and Pattern Recognition › Advanced Neural Network Applications
- Keywords
- self-attention, space complexity, sequence length, memory-efficient computation, numerical stability, long sequence processing
About these checks
Registry, retraction, Cites and statistics results are computed, not judged, with one exception: a printed reference that no index lists is read once by a model before it counts as a citation error. Claims are extracted by a model and kept only when their quote appears verbatim in the text; each cited source is then checked, with a model’s help, for whether it says what the claim relies on. Those verdicts can be wrong — read the claim’s quote and the cited source before relying on one.
Statistics and arithmetic are recomputed, not judged.
Confidence bands are shown only for fields calibrated on our annotated test set. None is calibrated yet, so no band is shown.
What counts as a claim ›How we measure confidence ›How we check this ›
Are you an author, or do you see a problem with this page? Request a review
