When I read the paper, Reformer: The Efficient Transformer, I cannot get the same complexity of the memory-efficient method in Table 1 (p. 5), which summarizes time/memory complexity of scaled dot-product, memory efficient, and LSH attention.
The memory complexity of the memory-efficient method is as follow:
$$\max \left(b n_{h} l d_{k}, b n_{h} l^{2}\right)$$
$b$: batch size $l$: sequence length $n_h$: the number of attention head $d_k$: the dimension of query or key
To the best of my knowledge, the memory-efficient method will do a loop for each query, therefore, the whole attention matrix will not show up.
So, shouldn't the memory complexity be $\max(b n_h l d_k, b n_h l)=(b n_h l d_k)$ instead of $\max(b n_h l d_k,b n_h l^2)$?
