Self-attention Does Not Need $O(n^2)$ Memory

By Markus Norman Rabe, Charles Staats. 2021

205 citations recorded by the source corpus.

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.

Subfields: Artificial Intelligence, Hardware and Architecture, Efficiency & Compression

How researchers see this paper

The researchers whose trusted networks rank it highest, strongest advocate first. Every researcher gets a different order.

In Artificial Intelligence

In Hardware and Architecture

In Efficiency & Compression

Citations in the corpus

LLaMA: Open and Efficient Foundation Language Models (2023, methodology): “First, we use an efficient implementation of the causal multi-head attention operator, inspired by Rabe and Staats (2021) and Dao et al. (2022).”
FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness (2022, methodology): “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].”
Efficient Transformers: A Survey (2020, background): “As noted by (Rabe and Staats, 2021), for applications that require to flex on sequence length and memory needs time to time, it might be suffice to ‘just sequentially process it’ even if that might not be inherently as satisfying as finding a theoretical approximate.”
General-purpose, long-context autoregressive modeling with Perceiver AR (2022, methodology): “Using the chunked approach in (Kitaev et al., 2020; Rabe & Staats, 2021; Jumper et al., 2021) for the cross-attend layer allows scaling input length beyond even those limits without requiring any architecture .”
A Comprehensive Overview of Large Language Models (2023, unknown): “Here in masked language modeling, the attention is not bidirectional but unidirectional, attending either left-to-right or right-to-left context.”
OpenFold: retraining AlphaFold2 yields new insights into its learning mechanisms and capacity for generalization (2023, background): “…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…”
Baichuan 2: Open Large-scale Language Models (2023, methodology): “For the attention layer of Baichuan 2, we adopt the memory efficient attention (Rabe and Staats, 2021) implemented by xFormers2.”
Taming Throughput-Latency Tradeoff in LLM Inference with Sarathi-Serve (2024, methodology): “[ 38,39,60] have proposed various techniques to minimize the memory bottlenecks of self-attention with careful tiling and work partitioning.”
RWKV: Reinventing RNNs for the Transformer Era (2023, methodology): “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.”
Hungry Hungry Hippos: Towards Language Modeling with State Space Models (2022, background): “3There are several memory-efficient algorithms for attention [15, 52], though their time complexity is still quadratic in N , which is a lower-bound for attention [36].”
Ring Attention with Blockwise Transformers for Near-Infinite Context (2023, methodology): “It’s worth noting that on GPUs our approach Ring Attention can be also integrated with the more compute efficient Triton code [18] or CUDA code [9] of memory efficient attention [31], similarly on TPUs it is also compatible with Pallas [37].”
FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioning (2023, methodology): “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).”