<p>I have the same question, and I think their claim is a little bit misleading. I will explain my understanding.</p>
<p>First, the <code>Q</code>,<code>K</code>,<code>V</code> here is only symbols for distinguish in the formula of attention calculation, not means they have to be different. For vanilla self-attention, they are actually all equal to input X, without any linear projection/transforming.</p>
<p>This can be seen in Figure 2 of the <a href="https://papers.nips.cc/paper/2017/file/3f5ee243547dee91fbd053c1c4a845aa-Paper.pdf" rel="nofollow noreferrer">transformer paper</a>, the linear layer is not included in the dot-product attention layer(left), and in multi-head attention, there is addition linear layer before the dot-product attention(right). So their complexity result is for vanilla self-attention, without any linear projection, i.e. <code>Q</code>=<code>K</code>=<code>V</code>=<code>X</code>.</p>
<p>And, I found this <a href="https://nlp.stanford.edu/seminar/details/lkaiser.pdf" rel="nofollow noreferrer">slides</a> from one of the author of the transformer paper, you can see clearly, <code>O(n^2 d)</code> is only for the dot-product attention, without the linear projection. While the complexity of multi-head attention is actually <code>O(n^2 d+n d^2)</code>.</p>
<p>Also I don't think the argument of <a href="https://stackoverflow.com/a/65794564/4129549">@igrinis</a> is correct. Although it didn't require to calculate QKV in original <a href="https://arxiv.org/pdf/1409.0473.pdf" rel="nofollow noreferrer">attention paper</a>, the complexity of alignment model(MLP) is actually <code>O(d^2)</code> for each pair of value, so total complexity of attention layer is <code>O(n^2·d^2)</code>, even larger than the QKV attention.</p>
Computational Complexity of Self-Attention in the Transformer Model
<p>I recently went through the <a href="https://papers.nips.cc/paper/2017/file/3f5ee243547dee91fbd053c1c4a845aa-Paper.pdf" rel="noreferrer">Transformer</a> paper from Google Research describing how self-attention layers could completely replace traditional RNN-based sequence encoding layers for machine translation. In Table 1 of the paper, the authors compare the computational complexities of different sequence encoding layers, and state (later on) that self-attention layers are faster than RNN layers when the sequence length <code>n</code> is smaller than the dimensionality of the vector representations <code>d</code>.</p>
<p>However, the self-attention layer seems to have an inferior complexity than claimed if my understanding of the computations is correct. Let <code>X</code> be the input to a self-attention layer. Then, <code>X</code> will have shape <code>(n, d)</code> since there are <code>n</code> word-vectors (corresponding to rows) each of dimension <code>d</code>. Computing the output of self-attention requires the following steps (consider single-headed self-attention for simplicity):</p>
<ol>
<li>Linearly transforming the rows of <code>X</code> to compute the query <code>Q</code>, key <code>K</code>, and value <code>V</code> matrices, each of which has shape <code>(n, d)</code>. This is accomplished by post-multiplying <code>X</code> with 3 learned matrices of shape <code>(d, d)</code>, amounting to a computational complexity of <code>O(n d^2)</code>.</li>
<li>Computing the layer output, specified in Equation 1 of the paper as <code>SoftMax(Q Kt / sqrt(d)) V</code>, where the softmax is computed over each row. Computing <code>Q Kt</code> has complexity <code>O(n^2 d)</code>, and post-multiplying the resultant with <code>V</code> has complexity <code>O(n^2 d)</code> as well.</li>
</ol>
<p>Therefore, the total complexity of the layer is <code>O(n^2 d + n d^2)</code>, which is worse than that of a traditional RNN layer. I obtained the same result for multi-headed attention too, on considering the appropriate intermediate representation dimensionalities (<code>dk</code>, <code>dv</code>) and finally multiplying by the number of heads <code>h</code>.</p>
<p>Why have the authors ignored the cost of computing the Query, Key, and Value matrices while reporting total computational complexity?</p>
<p>I understand that the proposed layer is fully parallelizable across the <code>n</code> positions, but I believe that Table 1 does not take this into account anyway.</p>