# Self-attention Does Not Need O(n2) Memory

arXiv: https://arxiv.org/abs/2112.05682 · HTML: https://mlsanity.com/paper/2112.05682/

## Key facts

- Cited by 154 papers in the map (data through 14 Oct 2026).
- Builds on 15 of its 22 references that are in the map.
- First posted on arXiv 10 Dec 2021.

## Builds on (most cited first, 15 of 15)

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding](https://mlsanity.com/paper/1810.04805/) (2018)
- [Neural Machine Translation by Jointly Learning to Align and Translate](https://mlsanity.com/paper/1409.0473/) (2014)
- [Big Bird: Transformers for Longer Sequences](https://mlsanity.com/paper/2007.14062/) (2020)
- [Reformer: The Efficient Transformer](https://mlsanity.com/paper/2001.04451/) (2020)
- [Rethinking Attention with Performers](https://mlsanity.com/paper/2009.14794/) (2020)
- [Generating Long Sequences with Sparse Transformers](https://mlsanity.com/paper/1904.10509/) (2019)
- [Linformer: Self-Attention with Linear Complexity](https://mlsanity.com/paper/2006.04768/) (2020)
- [Training Deep Nets with Sublinear Memory Cost](https://mlsanity.com/paper/1604.06174/) (2016)
- [Long Range Arena: A Benchmark for Efficient Transformers](https://mlsanity.com/paper/2011.04006/) (2020)
- [Efficient Content-Based Sparse Attention with Routing Transformers](https://mlsanity.com/paper/2003.05997/) (2020)
- [Efficient Attention: Attention with Linear Complexities](https://mlsanity.com/paper/1812.01243/) (2018)
- [Blockwise Self-Attention for Long Document Understanding](https://mlsanity.com/paper/1911.02972/) (2019)
- [Luna: Linear Unified Nested Attention](https://mlsanity.com/paper/2106.01540/) (2021)
- [Combiner: Full Attention Transformer with Sparse Computation Cost](https://mlsanity.com/paper/2107.05768/) (2021)

## Built on by (most cited first, 30 of 154)

- [GPT-4 Technical Report](https://mlsanity.com/paper/2303.08774/) (2023)
- [LLaMA: Open and Efficient Foundation Language Models](https://mlsanity.com/paper/2302.13971/) (2023)
- [FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness](https://mlsanity.com/paper/2205.14135/) (2022)
- [FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioning](https://mlsanity.com/paper/2307.08691/) (2023)
- [RWKV: Reinventing RNNs for the Transformer Era](https://mlsanity.com/paper/2305.13048/) (2023)
- [Baichuan 2: Open Large-scale Language Models](https://mlsanity.com/paper/2309.10305/) (2023)
- [A Comprehensive Overview of Large Language Models](https://mlsanity.com/paper/2307.06435/) (2023)
- [LightGlue: Local Feature Matching at Light Speed](https://mlsanity.com/paper/2306.13643/) (2023)
- [The Falcon Series of Open Language Models](https://mlsanity.com/paper/2311.16867/) (2023)
- [Taming Throughput-Latency Tradeoff in LLM Inference with Sarathi-Serve](https://mlsanity.com/paper/2403.02310/) (2024)
- [Hungry Hungry Hippos: Towards Language Modeling with State Space Models](https://mlsanity.com/paper/2212.14052/) (2022)
- [Ring Attention with Blockwise Transformers for Near-Infinite Context](https://mlsanity.com/paper/2310.01889/) (2023)
- [mGTE: Generalized Long-Context Text Representation and Reranking Models for Multilingual Text Retrieval](https://mlsanity.com/paper/2407.19669/) (2024)
- [Patch n' Pack: NaViT, a Vision Transformer for any Aspect Ratio and Resolution](https://mlsanity.com/paper/2307.06304/) (2023)
- [World Model on Million-Length Video And Language With Blockwise RingAttention](https://mlsanity.com/paper/2402.08268/) (2024)
- [SARATHI: Efficient LLM Inference by Piggybacking Decodes with Chunked Prefills](https://mlsanity.com/paper/2308.16369/) (2023)
- [Actions Speak Louder than Words: Trillion-Parameter Sequential Transducers for Generative Recommendations](https://mlsanity.com/paper/2402.17152/) (2024)
- [Flex Attention: A Programming Model for Generating Optimized Attention Kernels](https://mlsanity.com/paper/2412.05496/) (2024)
- [From Decoding to Meta-Generation: Inference-time Algorithms for Large Language Models](https://mlsanity.com/paper/2406.16838/) (2024)
- [ONE-PEACE: Exploring One General Representation Model Toward Unlimited Modalities](https://mlsanity.com/paper/2305.11172/) (2023)
- [Towards Efficient Generative Large Language Model Serving: A Survey from Algorithms to Systems](https://mlsanity.com/paper/2312.15234/) (2023)
- [Hydragen: High-Throughput LLM Inference with Shared Prefixes](https://mlsanity.com/paper/2402.05099/) (2024)
- [Advancing Transformer Architecture in Long-Context Large Language Models: A Comprehensive Survey](https://mlsanity.com/paper/2311.12351/) (2023)
- [Tracking Meets LoRA: Faster Training, Larger Model, Stronger Performance](https://mlsanity.com/paper/2403.05231/) (2024)
- [General-purpose, long-context autoregressive modeling with Perceiver AR](https://mlsanity.com/paper/2202.07765/) (2022)
- [SiLK -- Simple Learned Keypoints](https://mlsanity.com/paper/2304.06194/) (2023)
- [Striped Attention: Faster Ring Attention for Causal Transformers](https://mlsanity.com/paper/2311.09431/) (2023)
- [Geometric Algebra Transformer](https://mlsanity.com/paper/2305.18415/) (2023)
- [Post-Training Sparse Attention with Double Sparsity](https://mlsanity.com/paper/2408.07092/) (2024)
- [Replacing softmax with ReLU in Vision Transformers](https://mlsanity.com/paper/2309.08586/) (2023)

## Abstract (arXiv)

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(n2) memory. While the time complexity is still O(n2), 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(√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.