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.


@Jindřich is exactly right. this is not an answer by itself, but rather some pointers to the implementation in the annotated transformer of every item he mentioned:

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

  • torch.from_numpy(np.random.randint(1, V, size=(batch, l))): $b \times l$
  • self.lut = nn.Embedding(vocab, d_model); self.lut(x): $b \times l \times d_m$

The output of the dense layer has the same shape $b \times l \times d_m = b \times l \times hd_k$:

  • self.linears = clones(nn.Linear(d_model, d_model), 4): these are $W^Q,W^K,W^V,W^O$ respectively, and their output is $b \times l \times d_m$
  • l(x).view(nbatches, -1, self.h, self.d_k): converts the output to $b \times l \times h \times d_k$

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

  • l(x).view(nbatches, -1, self.h, self.d_k).transpose(1, 2): converts the output to $b \times h \times l \times d_k$, done for K, Q and V.

Now, if you permute the dimensions...

  • scores = torch.matmul(query, key.transpose(-2, -1)): $[b \times h \times l \times d_k]\times[b \times h \times d_k \times l] = [b \times h \times l \times l]$
  • scores = scores.masked_fill(mask == 0, -1e9): mask is $b \times 1 \times 1 \times l$ for encoder layers, and $b \times 1 \times l \times l$ for decoder layers (actually $l-1$ but lets ignore that). Both are broadcastable to the scores tensor, the encoder mask is the same for all heads and for all positions, while the decoder mask is different for every position, as it hides the subsequent positions.

Then, if you do softmax in the last dimension...

  • p_attn = F.softmax(scores, dim = -1): $b \times h \times l \times l$, but now each vector sums to 1

By doing the batch matrix multiplication ... you get the weighted average

  • torch.matmul(p_attn, value): $[b \times h \times l \times l] \times [b \times h \times l \times d_k] = [b\times h\times l \times d_k]$, which are the weighted averages we wanted

So finally, the "concatenation" of all heads is in fact just another tensor reshaping

  • x.transpose(1, 2).contiguous(): $[b\times l\times h \times d_k]$
  • x.transpose(1, 2).contiguous().view(nbatches, -1, self.h * self.d_k): $[b\times l\times d_m]$

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service, privacy policy and cookie policy

Not the answer you're looking for? Browse other questions tagged or ask your own question.