Self-attention Does Not Need $O(n^2)$ Memory
This work addresses a critical bottleneck in processing long sequences for machine learning models, particularly on modern accelerators where memory is a limiting factor, by providing a memory-efficient alternative to standard self-attention.
The paper tackles the problem of high memory usage in self-attention mechanisms, which are often limited by O(n^2) memory requirements, by introducing an algorithm that reduces memory to O(1) for attention and O(log n) for self-attention, with a practical implementation achieving O(sqrt(n)) memory and reducing memory overhead by 59X for inference and 32X for differentiation at sequence length 16384.
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.