Computational Complexity of Self-Attention in the Transformer Model

First, you are correct in your complexity calculations. So, what is the source of confusion?

When the original Attention paper was first introduced, it didn't require to calculate Q, V and K matrices, as the values were taken directly from the hidden states of the RNNs, and thus the complexity of Attention layer is O(n^2·d).

Now, to understand what Table 1 contains please keep in mind how most people scan papers: they read title, abstract, then look at figures and tables. Only then if the results were interesting, they read the paper more thoroughly. So, the main idea of the Attention is all you need paper was to replace the RNN layers completely with attention mechanism in seq2seq setting because RNNs were really slow to train. If you look at the Table 1 in this context, you see that it compares RNN, CNN and Attention and highlights the motivation for the paper: using Attention should have been beneficial over RNNs and CNNs. It should have been advantageous in 3 aspects: constant amount of calculation steps, constant amount of operations and lower computational complexity for usual Google setting, where n ~= 100 and d ~= 1000. But as any idea, it hit the hard wall of reality. And in reality in order for that great idea to work, they had to add positional encoding, reformulate the Attention and add multiple heads to it. The result is the Transformer architecture which while has the computational complexity of O(n^2·d + n·d^2) still is much faster then RNN (in a sense of wall clock time), and produces better results.

So the answer for your question is that attention layer the authors refer to in Table 1 is strictly the attention mechanism. It is not the complexity of the Transformer. They are very well aware about the complexity of their model (I quote):

Separable convolutions [6], however, decrease the complexity considerably, to O(k·n·d + n·d^2). Even with k = n, however, the complexity of a separable convolution is equal to the combination of a self-attention layer and a point-wise feed-forward layer, the approach we take in our model.

Strictly speaking, when considering the complexity of only the self-attention block (Fig 2 left, equation 1) the projection of x to q, k and v is not included in the self-attention. The complexities shown in table 1 are only for the very core of self-attention layer and thus are O(n^2 d).

one small caveat: the derivation seems correct to me but it'd be good to point out that there are 8 heads (denoted by h=8 in the original paper) for each mult-head attention module, and the attention matrices Q,K,V are shrunk by the same factor h. so the expressions are divided and multiplied by h for the second part (n*n*d) but the first part needs to be divided by h i believe.

I have the same question, and I think their claim is a little bit misleading. I will explain my understanding.

First, the Q,K,V 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.

This can be seen in Figure 2 of the transformer paper, 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. Q=K=V=X.

And, I found this slides from one of the author of the transformer paper, you can see clearly, O(n^2 d) is only for the dot-product attention, without the linear projection. While the complexity of multi-head attention is actually O(n^2 d+n d^2).

Also I don't think the argument of @igrinis is correct. Although it didn't require to calculate QKV in original attention paper, the complexity of alignment model(MLP) is actually O(d^2) for each pair of value, so total complexity of attention layer is O(n^2·d^2), even larger than the QKV attention.

You cannot compare this to a traditional RNN encoder-decoder, the architecture described in the paper is meant to improve upon the classical Attention Mechanism first established in this paper.

In its initial form, the attention mechanism was relying on a Neural network trained to retrieved the relevant hidden states of the encoder. Instead of relying on a fixed retrieval strategy (for instance: using the last hidden state) you allow the system some control over the process.

There is already a very good post on StackExchange explaining the differences in computational complexity here.

The paper you are describing is "replacing" this Neural Network with a dot product between two array, which less demanding computationally than having to train a Neural Network and relatively more efficient. But its not meant to be more efficient than a regular RRN-based auto-encoder without attention.

How is this any less demanding computationally?

In a traditionnal RNN / LSTM based auto encoder each time step is encoded into a vector h. The decoder usually (again, there's a lot of different architectures but that's the basic one) takes the last vector h as input produce the output sequence.

In this scenario there are no attention mechanism, this is straight forward reading the last encoded state. The problem with this architecture is that as your sequence gets longer, all the relevant information gets squeezed into the last encoded state h(t) and you loose relevant information.

Introducing attention mechanism

As described in the paper above, the original attention mechanism aims at circumventing this limitation by allowing the decoder to access not only the last encoded state but any of the previous encoded state and to combine them in order to improve on the prediction.

For each time steps, a probability vector alpha is computed by a neural network to choose encoded state to retreive :

If we restrict α to be an one-hot vector, this operation becomes the same as retrieving from a set of elements h with index α. With the restriction removed, the attention operation can be thought of as doing "proportional retrieval" according to the probability vector α

I won't copy paste the SE post , there you'll have the explanation on why the method of the dot product is computationally more efficient than the neural network.

Conclusion

The key take away is that you cannot compare this to a traditional RNN encoder-decoder because there are no attention mechanism in such network. It would be like comparing CNN with LSTM layer, these are just different architecture.