Source-linked AI summary
Self-attention Does Not Need $O(n^2)$ Memory
Markus N. Rabe, Charles Staats
TL;DR
Self-attention is commonly described as requiring quadratic memory, which matters because modern accelerators are often memory constrained. The paper introduces exact incremental algorithms and a practical chunked implementation that reduce memory while retaining quadratic self-attention time. The implementation is numerically stable and can scale to sequence length 1M, though runtime measurements are approximate and the time complexity remains quadratic.
Problem
Self-attention has O(n^2) time and space complexity, while modern accelerators are often memory constrained and compute is relatively cheap.
Method
The paper presents exact O(1)-memory attention and O(log n)-memory self-attention algorithms, plus a numerically stable chunked implementation using O(√n) memory.
Results
The memory-efficient implementation removes the memory bottleneck of self-attention and scales to sequence length 1M while retaining quadratic time complexity.
Takeaways & Limitations
The algorithm can serve as a drop-in replacement to save memory and may support longer dense-attention datasets or revised architecture choices.
Takeaways & Limitations
The time complexity remains O(n^2) for self-attention, and isolated runtime measurements fluctuate and may not represent performance inside larger architectures.
Abstract
from arXiv · showhide
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.
1 Introduction
The paper argues that self-attention’s quadratic memory cost is a practical bottleneck on modern accelerators and presents exact, memory-efficient alternatives. Its implementation reduces memory while retaining quadratic self-attention time.
- Problem: O(n^2) space arises because self-attention computes and stores O(n) scores for each of n queries.Single-query attention uses O(n) time and memory; applying separate queries at every sequence position yields quadratic overall complexity.
- Motivation: Modern accelerators are often memory constrained while compute is relatively cheap, making self-attention’s space complexity a particular concern.This concern has motivated alternatives with more favorable complexity classes.
- Contribution: O(1) memory for attention and O(log n) memory for self-attention are achieved by the paper’s new algorithms.The basic algorithm is simple, while numerical feasibility requires an additional trick described later.
- Novelty: The memory-efficient attention algorithm is exact rather than approximate, computing the same function as standard attention.This supports using it as a drop-in replacement to save memory.
- Scope: The method still requires O(n^2) time for self-attention, so efficient long-context attention mechanisms remain alternatives to dense attention.The paper presents memory reduction, not a reduction in dense self-attention’s time complexity.
2 Algorithm
The algorithm computes attention incrementally by accumulating weighted values and normalization terms, then dividing once at the end. Sequential processing yields constant memory for one query and logarithmic memory for self-attention, with input-order assumptions affecting the bound.
- Reformulation: The division by the sum of exponentiated scores can be moved to the end using the distributive law.This reformulation enables incremental accumulation instead of storing every score.
- Related work: The displayed reformulation was later identified as a rediscovery of Jang et al.’s “lazy softmax” method.The paper distinguishes its work by discussing memory-complexity implications and additional innovations.
- Single-query attention: Constant memory suffices for single-query attention by maintaining one vector and one scalar while processing key-value pairs sequentially.The final attention result is obtained by dividing the accumulated vector by the accumulated scalar.
- Memory bound: O(log n) memory is required when inputs arrive in an order that necessitates storing an index into the sequence.With the specified query-then-pairs order, this additional index is not needed.
- Self-attention: Self-attention is computed by processing all queries sequentially, requiring one additional query index and producing O(n) outputs excluded from the space-complexity count.The resulting working-memory complexity is O(log n).
3 Numerical Stability
The incremental algorithm is numerically unstable because exponentiated scores can overflow, while maximum-score subtraction must be applied during accumulation. The paper resolves this by tracking the running maximum and renormalizing accumulated sums.
- Numerical problem: Scores ≥89 produce inf under exponentiation for bfloat16 and float32, causing numerical instability in attention.Standard softmax avoids this by subtracting the maximum score without changing the result.
- Numerical problem: Incremental accumulation cannot simply delay maximum-score subtraction because scores must be exponentiated before entering the cumulative sums.The maximum may also depend on a later score in the sequence.
- Stabilization: A running maximum and renormalized accumulators make the incremental algorithm numerically feasible.The algorithm updates the maximum, weighted-value sum, and score sum as each key-value pair is processed, then divides the final accumulators.
4 An Implementation For TPUs
The TPU-oriented implementation uses chunked attention to balance memory use, computational efficiency, and implementation simplicity. It processes query, key, and value chunks sequentially while writing each result directly to the output.
- Implementation: The implementation supports multiple attention heads and memory-efficient differentiation in JAX.The stated implementation targets efficient execution on TPUs while retaining memory-efficient differentiation.
- Design goal: The implementation balances simplicity, computational efficiency, and memory requirements rather than optimizing exclusively for minimum memory.The implementation is presented in Figure 1 and is suited for TPU execution.
- Attention computation: The attention code computes query-key dot products, stabilizes scores with a maximum score, and combines weighted values.The supplied implementation excerpts show einsum-based score computation, maximum-score normalization, and weighted value aggregation.
- Chunked computation: Chunked attention processes query chunks in an outer loop and key-value chunks sequentially within each query chunk.Each key-value chunk is summarized independently before the summaries are rescaled and combined.
- Memory-runtime trade-off: The implementation exposes query and key chunk sizes because memory and runtime depend on their hardware-specific choice.A constant query chunk size and key-value chunk size of √n are optimal for memory consumption, while runtime depends on hardware.
5 Empirical Analysis
The experiments evaluate memory, runtime, numerical agreement, differentiation, and translation quality for memory-efficient attention. The implementation scales to sequence length 1M, preserves results closely, and can outperform query chunking under memory constraints.
- Inference: Sequence length 1M is supported while self-attention retains quadratic time complexity.The implementation multiplies over more than 1 trillion query-key combinations at this sequence length.
- Inference: Runtime measurements are intended to show roughly similar performance, but fluctuate across runs and may differ inside larger architectures.The relative compute speed is the median over 100 runs; a small Transformer instead showed about a 4% increase in steps/sec.
- Inference: Memory-efficient attention is numerically close to standard attention, within 1.8 × 10−7 maximal absolute difference for tested inputs.The comparison covers cases where standard attention does not exceed 16GB device memory and uses sequence length 2^14.
- Differentiation: Checkpointing preserves the memory advantage during differentiation, although recomputation reduces relative compute speed.The algorithm summarizes attention chunks sequentially and never forms the full attention matrix.
- Training: After 100K WMT en-de training steps, evaluation accuracy was 62.69 with memory-efficient attention versus 62.59 with standard attention.The two Transformer implementations behaved almost identically throughout training, although the learning rate was lowered to 0.005.
- Comparison to Query Chunking: Memory-efficient attention can outperform query chunking in memory-constrained scenarios as sequence length increases.Chunking keys as well as queries avoids progressively shrinking query chunks to ≤64, which eventually slows query chunking significantly.
6 Related Work
Related work had explored delaying softmax division and reducing attention memory, but the authors distinguish their work through memory-complexity analysis, numerical-stability treatment, differentiation, and implementation.
- Prior attention algorithms: Jang et al. (2019) previously observed that softmax division can be delayed until the end of attention, but did not discuss memory complexity.The paper also says that work did not address numerical stability or backpropagation and had no known public implementation.
- Memory-efficient implementations: Dao et al. (2022) provided a CUDA implementation and reported speedups from reduced memory requirements on GPUs.This paper attributes its lack of similar gains to standard TPU self-attention already balancing FLOPs and memory bandwidth.
7 Conclusion
The paper presents a simple trick that dramatically reduces the memory requirement of attention and self-attention, a possibility the authors say the community appears to have overlooked.
- Conclusion: The proposed trick reduces attention and self-attention memory requirements dramatically without changing the paper’s central attention computation.The authors hope this raises awareness that attention is not intrinsically memory-hungry and may prompt revisiting neural- and hardware-architecture choices.