I am reading Self-attention Does Not Need $O(n^{2})$ Memory which proposes an algorithm that requires $O(1)$ memory for one query and $O(log n)$ memory for self-attention, in theory. In practice the authors claim it is $O(\sqrt{n})$. $n$ here is the sequence length.

I am having quite the difficulty understanding why the algorithm requires $O(log n)$ memory when the sequence of $k, v$ pairs is not ordered. This supposes the query is fixed (not yet self-attention).

Here is the standard attention algorithm:

$s_{i} = dot(q, k_{i})$, $s_{i} = \frac{e^{s_{i}}}{\sum_{j}^{}{e^{s_{j}}}}$, $attention(q,k,v)=\sum_{j}^{}{v_{i}s_{i}'}$

Here is the suggested algorithm:

$s_{i} = dot(q, k_{i})$, $s_{i} = e^{s_{i}}$, $attention(q,k,v)=\frac{\sum_{j}^{}{v_{i}s_{i}'}}{\sum_{j}^{}{s_{j}'}}$

For ordered sequences of $(k, v)$ pairs, and computing the terms sequentially, it is clear that the memory complexity is $O(1)$. But I don't understand why it is $O(log n)$ for unordered sequences.

From my understanding, if we have a sequence of length $n$, we store one scalar $s^{*} \in \mathbb{R}$, this will constitute the $\sum_{j}^{}{s_{j}'}$ in the attention computation. We also store a vector $v^{*}\in \mathbb{R}^{d}$ that will constitute the $\sum_{j}^{}{v_{i}s_{i}'}$ in the attention computation. I don't understand why these quantities will be affected by the order of the pairs, the summation is commutative so whatever the order in which we process the pairs we should get the same value at the end.

I'm sure I'm missing something but I can't quite figure it out. I do not think it is the case, but maybe the authors mean by "inputs are provided in a different order" that we don't receive pairs of keys and values but they come individually? I don't think it is the case because then wouldn't the most efficient way to keep track of which values' indices were used for which keys' indices is to keep two hash tables, which leads to $O(n)$?



You must log in to answer this question.