So, I have a doubt in Attention is all you need:

The implementation of transformers on tensorflow's official documentation says:

Each multi-head attention block gets three inputs; Q (query), K (key), V (value). These are put through linear (Dense) layers and split up into multiple heads.

However, The paper mentions:

Instead of performing a single attention function with dmodel-dimensional keys, values and queries, we found it beneficial to linearly project the queries, keys and values h times with different, learned linear projections to dk, dk and dv dimensions, respectively. On each of these projected versions of queries, keys and values we then perform the attention function in parallel, yielding dv-dimensional output values.

There is no mention of splitting Q,K and V to obtain heads. Instead, the paper says that they are passed through 'h' different dense layers to convert d-model dimensional vectors to 'h' different dk,dk and dv dimensional vectors respectively. So basically the pseudocode, from what I understood, should look something like this:

Q,K & V are d-model dimensional vectors.

for i in range(h):
Qi = Dense(dk)(Q)
Ki = Dense(dk)(K)
Vi = Dense(dv)(V)
Ai = Attention(Qi, Ki, Vi)

A0, A1, A2, ..., Ah are then concatenated.


Is this right? or am I missing something here?

In principle, the pseudocode is correct, but it is not how it is implemented. The projection and dot-product attention can be done efficiently using matrix multiplication only for all heads simultaneously.

Instead of looping over the heads, you can use only one dense layer for queries, one for keys and one for values. E.g., for the keys, you would do a dense layer of dimension $$h \cdot d_k$$ and then reshape it.

Assuming you have batch size $$b$$, sequence length $$l$$ and model dimension $$d_m$$:

• The input of the dense layers is of shape $$b \times l \times d_m$$.

• The output of the dense layer is has shape $$b \times l \times hd_k$$ (or $$d_v$$ respectively).

• Then you can reshape the queries and keys to have shape $$b \times l \times h \times d_k$$.

Now, if you permute the dimensions, such that the queries have shape $$b \times h \times l \times d_k$$ and the keys have shape $$b \times h \times d_k \times l$$, you can do the batch matrix multiplication in the last two dimensions and you end with attention energies of shape $$b \times h \times l \times l$$. Then, if you do softmax in the last dimension, you get the attention distribution for each head and each query.

By doing the batch matrix multiplication with projected values that now with transposed dimensions to $$b \times h \times l \times d_v$$, you get the weighted average. So finally, the "concatenation" of all heads is in fact just another tensor reshaping.