8
$\begingroup$

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?

Improve this question
$\endgroup$

1 Answer 1

Reset to default
8
$\begingroup$

  • 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.

Improve this answer
$\endgroup$
5
  • $\begingroup$ Excellent. This is so helpful $\endgroup$
    – data_person
    Commented Apr 30, 2021 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$
    – data_person
    Commented Apr 30, 2021 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
    Commented Apr 30, 2021 at 19:55
  • $\begingroup$ May I ask a few questions here? My impression is that d_model = 512 is the word-embedding dimension, meaning each token, say "king", is a 512-dim vector. The input would be a sequence of words, eg, "I am king of the jungle". A Transformer layer would have a width which is the number of tokens it can process per each pass. I wonder why is this "width" not specified? Or is it the same as the "depth" referred to above? $\endgroup$
    – Yan King Yin
    Commented Jul 27, 2022 at 22:06
  • 1
    $\begingroup$ @YanKingYin what you are talking about is the sequence length, and it is not related to $depth$. The sequence length is variable. In the Transformer model, the limit of sequence length is imposed by the configured positional encoding table size and the device memory (the consumed memory is quadratic in the sequence length). $\endgroup$
    – noe
    Commented Jul 27, 2022 at 22:10

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service and acknowledge you have read our privacy policy.

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