Understanding dimensions of vectors at various places in transformer architecture

Question

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):

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?

Answer 1

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.