1
$\begingroup$

I was trying to understand transformer architecture from "Attention is all you need" paper. It says following regarding dimensions of different vectors:

  • The input consists of queries and keys of dimension $d_k$, and values of dimension $d_v$.
  • $MultiHead(Q, K, V ) = Concat(head_1, ..., head_h) W^O$ where $head_i = Attention(QW_i^Q,KW^K_i,VW^V_i)$
    where $W_i^Q\in\mathbb{R}^{d_{model}\times d_k}$, $W_i^K\in\mathbb{R}^{d_{model}\times d_k}$, $W_i^V\in\mathbb{R}^{d_{model}\times d_v}$, $W^O\in\mathbb{R}^{hd_{v}\times d_{model}}$
  • $h=8$ parallel attention layers or heads
  • $d_k=d_v=d_{model}/h=64$

From these I figured out dimensions of vectors at different position in the transformers model as follows (in red colored text):

enter image description here

I have following doubts:

Is dimension of $K$ (vector of multiple/all keys that current word-query needs to attend to) $=d_k$? Or $d_k$ is just for single key? If for single key, then what is the dimension for $K$? Same is the doubt with $Q$ and $d_q$. I feel $Q$ is set of all queries that can apply to single word. $K$ is the set of all keys that single word can attend to. If that is the case, then dimensions of $Q$ must be $d_q\times\text{number of queries to consider}$ and $K$ must be $d_k\times\text{number of keys to attend for each word-query}$ But then what is this number of queries and keys?

$\endgroup$
1
  • $\begingroup$ Please consider upvoting the answer if it is useful. Also, please consider accepting it if deemed correct, or alternatively, please indicate why you think it's not correct or not clear enough. $\endgroup$
    – noe
    Aug 28 at 22:43

1 Answer 1

1
$\begingroup$

Some of the dimensions of the diagram in the middle are somewhat misleading. Specifically, the sizes of $K$, $V$ and $Q$, which you set to $64 \times h$, which makes me think that they are split for each head before entering the multi-head attention. However, they are not split, they are simply projected with a different projection matrix for each head. Those projections map from the original $d_{model}$ dimensionality to $d_k, d_v, d_k$ dimensional spaces).

Is dimension of $K$ (vector of multiple/all keys that current word-query needs to attend to) $=d_k$? Or $d_k$ is just for single key? If for single key, then what is the dimension for $K$?

$d_k$ is the dimension of each of the key vectors (there are $n_k$ key vectors in a sequence $n_k$ tokens long, i.e. number of keys to consider). $K$ is actually a tensor of dimensions $n_k \times d_k$ (actually it has an extra dimension to computing in parallel multiple sequences in a minibatch: $b \times n_k \times d_k$, where $b$ is the batch size).

Same is the doubt with $Q$ and $d_q$. I feel $Q$ is set of all queries that can apply to single word. $K$ is the set of all keys that single word can attend to. If that is the case, then dimensions of $Q$ must be $d_q\times\text{number of queries to consider}$ and $K$ must be $d_k\times\text{number of keys to attend for each word-query}$ But then what is this number of queries and keys?

As with the keys, there is a query vector for each of the input tokens $n_q$ (i.e. number of queries to consider). Each of the query vectors is multiplied by each of the key vectors. The dimension of the query is $n_q \times d_k$ (ignoring the batch size) (note that there is no $d_q$ in the paper, $d_q$ is simply $d_k$). After multiplying all queries by all keys obtaining a tensor of dimensions $n_q \times n_k$, we take the softmax along the dimension of $n_k$ and use the resulting weights to weight-sum the value vectors obtaining a tensor of dimensions $d_v \times n_q$, that is, one vector for each query to consider.

$\endgroup$

Your Answer

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

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