In the Attention is all you need paper, on the 4th page, we have equation 1, which describes the self-attention mechanism of the transformer architecture

$$ \text { Attention }(Q, K, V)=\operatorname{softmax}\left(\frac{Q K^{T}}{\sqrt{d_{k}}}\right) V $$

Everything is fine up to here.

Then they introduce the multi-head attention, which is described by the following equation.

$$ \begin{aligned} \text { MultiHead }(Q, K, V) &=\text { Concat}\left(\text {head}_{1}, \ldots, \text {head}_{\mathrm{h}}\right) W^{O} \\ \text { where head}_{\mathrm{i}} &=\text {Attention}\left(Q W_{i}^{Q}, K W_{i}^{K}, V W_{i}^{V}\right) \end{aligned} $$

Once the multi-head attention is motivated at the end of page 4, they state that for a single head (the $i$th head), the query $Q$ and key $K$ inputs are first linearly projected by $W_i^Q$ and $W_i^K$, then dot product is calculated, let's say $Q_i^p = Q W_i^Q$ and $K_i^p = K W_i^K$.

Therefore, the dot product of the projected query and key becomes the following from simple linear algebra.

$$Q_i^p {K_i^p}^\intercal = Q W_i^Q {W_i^K}^T K^T = Q W_i K^T,$$


$$W_i = W_i^Q {W_i^K}^T$$

Here, $W$ is the outer product of query projection by the key projection matrix. However, it is a matrix with shape $d_{model} \times d_{model}$. Why did the authors not define only a $W_i$ instead of $W_i^Q$ and $W_i^K$ pair which have $2 \times d_{model} \times d_{k}$ elements? In deep learning applications, I think it would be very inefficient.

Is there something that I am missing, like these 2 matrices $W_i^Q$ and $W_i^K$ should be separate because of this and that?


I'll use notation from the paper you cited, and any other readers should refer to the paper (widely available) for definitions of notation. The utility of using $W^Q$ and $W^K$, rather than $W$, lies in the fact that they allow us to add fewer parameters to our architecture. $W$ has dimension $d_{model} \times d_{model}$, which means that we are adding $d_{model}^2$ parameters to our architecture. $W^Q$ and $W^K$ each have dimension $d_{model} \times d_k$, and $d_k=\frac{d_{model}}{h}$. If we use these two matrices, we only add $2\frac{d_{model}^2}{h}$ parameters to our architecture, even though their multiplication (with the transpose) allows us to have the correct dimensions for matrix multiplication with $Q$ and $K$.

We do use $h$ attention heads, which then brings our number of parameters back up, but the multiple heads let the model attend to different pieces of information in our data.

  • $\begingroup$ I understand now, actually the resulting $W$ matrix is rank deficient when defined by inner product of $W^Q$ and $W^K$. This makes parameters fewer :) Thanks. $\endgroup$ Dec 17 '20 at 5:42
  • $\begingroup$ I just read a paper (arxiv.org/pdf/1906.07510.pdf) in which the authors use multi-head attention from Vaswani et al, but have $W^Q$ and $W^K$ be of dimension $d_{model} \times d_{model}$. I don't know if there is a reason for this, or if it's a typo. I plan to email them, and if there is a reason for it, I'll post it here. $\endgroup$
    – BioBroo
    Dec 17 '20 at 17:59

In practice, matrices $W^Q, W^K, W^V$ (each of size $d_{model}$ x $d_{model}$) are completely removed instead, and Transformer implementations just learn a single set of matrices $\{ W_i^{Q*}, W_i^{K*}, W_i^{V*} \}$ (each of size $d_{model}$ x $\frac{d_{model}}{h}$) for each head, where

$W_i^{Q*} = W^Q W_i^Q \\ W_i^{K*} = W^K W_i^K \\ W_i^{V*} = W^V W_i^V $

so that:

$Q_i(x) = x W_i^{Q*} = x W^Q W_i^{Q*} = Q W_i^Q\\ K_i(x) = x W_i^{K*} = x W^K W_i^{K*} = K W_i^K\\ V_i(x) = x W_i^{V*} = x W^V W_i^{V*} = V W_i^V\\ head_i(x) = softmax \left(\frac{Q_i(x) K_i(x)^T}{\sqrt{d_k}} \right) V_i(x)$.

I can confirm this with the original Transformer implementation in Tensor2Tensor, and also the BERT code that uses the encoding part of the Transformer.

  • $\begingroup$ I missed the part where $d_k=d_{model} / h$ . Thank you. $\endgroup$ Dec 17 '20 at 5:42
  • $\begingroup$ I just stumbled upon this Question after not understanding the original paper. $$ \text { where head }_{\mathrm{i}}=\text { Attention }\left(Q W_{i}^{Q}, K W_{i}^{K}, V W_{i}^{V}\right) $$I would really like to know if this redundant now and it can be$$ \text { head }_{\mathrm{i}}=\text { Attention }\left(Q_{i}, K_{i}, V_{i}\right)$$ where $$Q_{i}, K_{i}, V_{i} $$are simply generated from the input to the layer and the matrices $$ \left\{W_{i}^{Q }, W_{i}^{K }, W_{i}^{V }\right\}\in \mathbb{R} ^{D\times \dfrac{D}{h}} $$ respectively $\endgroup$
    – JimSi
    Jan 17 '21 at 11:19

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.