A lot of recent research on Transformers has been devoted to reducing the cost of the self-attention mechanism:
$$\text{softmax}\left(\frac{Q K^T}{\sqrt{d}} \right)V,$$
As I understand it, the runtime, assuming $\{Q, K, V\}$ are each of shape $(n, d)$, is $O(n^2 d + n d^2)$. In general, the issue is the $n^2 d$ term, because the sequence length $n$ can be much bigger than the model dimension $d$. So far, so good.
But as far as I can tell, current research focuses on speedups for $Q K^T$, which is $O(n^2 d)$. There's less focus on computing $A V$, where $A = \text{softmax} \left(\frac{Q K^T}{\sqrt{d}} \right)$ -- which also has complexity $O(n^2 d)$.
Why is the first matrix product the limiting factor?
Examples of these faster Transformer architectures include Longformer, which approximates $QK^T$ as a low-rank-plus-banded matrix, Nystromformer, which approximates $\text{softmax}(QK^T)$ as a low-rank matrix with the Nystrom transformation, and Big Bird, which approximates it with a low-rank-plus-banded-plus-random matrix.
