Trying to understand the dimensions of the Multihead Attention component in Transformer referring the following tutorial https://www.tensorflow.org/tutorials/text/transformer#setup

There are 2 unknown dimensions - depth and d_model which I dont understand.

For example, if I fix the dimensions of the Q,K,V as 64 and the number_of_attention_heads as 8, and input_embedding as 512 , can anyone please explain what is depth and d_model?

  • d_model is the dimensionality of the representations used as input to the multi-head attention, which is the same as the dimensionality of the output. In the case of normal transformers, d_model is the same size as the embedding size (i.e. 512). This naming convention comes from the original Transformer paper.

  • depth is d_model divided by the number of attention heads (i.e. 512 / 8 = 64). This is the dimensionality used for the individual attention heads. In the tutorial you linked, you can find this as self.depth = d_model // self.num_heads. Each attention head projects the original representation into a smaller representation of size depth, then computes the attention, and then all the attention head results are concatenated together, so that the final dimensionality is again d_model. You can find more details on the individual computations in this other answer.

Note that the implementation of the multi-head attention in the tutorial is not a straightforward implementation from the original paper but it is equivalent: in the original paper, there are different matrices $W_i^Q, W_i^K, W_i^V$ for each attention head $i$, while in the implementation of the tutorial there are combined matrices $W^Q, W^K, W^V$ that compute the projection for all attention heads, which is then split into the separate heads by means of the function split_heads.

  • $\begingroup$ Excellent. This is so helpful $\endgroup$ Apr 30 '21 at 19:22
  • $\begingroup$ Just to confirm, then the shape of Wq = Wk = Wv = (word_embedding_dim,depth) which is (512,64) in this case? $\endgroup$ Apr 30 '21 at 19:31
  • 1
    $\begingroup$ If there were separate $W_i^Q$, $W_i^K$ and $W_i^V$ for each head, then yes. However, in the tutorial they have combined project projections, so their shapes are $d_{model} \times d_{model}$, and after multiplying the input, the outputs are sliced into smaller matrices with shape $d_{model} \times depth$. $\endgroup$
    – noe
    Apr 30 '21 at 19:55
  • $\begingroup$ Thanks. I got it now. $\endgroup$ Apr 30 '21 at 23:29

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.