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

Markus N. Rabe, Charles Staats

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. Th…