# When calculating self-attention for Transformer ML architectures, why do we need both a key and a query weight matrix?

I'm trying to understand the math behind Transformers, specifically self-attention. This link, and many others, gives the formula to compute the output vectors from the input embeddings as:

$$Q=XW_Q,\;\;\;K=XW_K,\;\;\;V=XW_V$$ $$Attention(Q,K,V)=softmax(\frac{QK^T}{\sqrt d_k})V$$

But this eventually becomes

$$Attention(Q,K,V)=softmax(X\frac{W_QW_K^T}{\sqrt d_k}X^T)V$$

If $$W_Q$$ and $$W_K$$ are only ever used in the form $$\frac{W_QW_K^T}{\sqrt d_k}$$, why do we initialize both matrices at all? why not just define and initialize a single matrix $$W_{QK}$$, skip the matrix multiplication, and get rid of the redundant weights?

The weight matrices are $$n$$ by $$m$$ with $$n >> m$$. So $$W_Q W_K^T$$ is not just any matrix, it's $$n$$ by $$n$$ but with rank only $$m$$ -- there are fewer parameters, and computing $$QK^T$$ is much faster than $$X W' X^T$$ for some full rank $$W'$$
• Ah - that makes total sense! So even though $XW'X^T$ is one less matrix multiplication than $XW_QWK^TX^T$: 1) $W$ has $n^2$ weights while $W_Q$ and $W_K$ total have $2nm$ weights - so we would be adding unnecessary extra weights when $n>2m$, which will almost always be the case. 2): The number of operations to multiply an $i\times j$ matrix and a $j\times k$ matrix is approximately $ijk$. So $XW_QWK^TX^T$ takes around $2pnm + p^2m$ operations, which will almost always be less than the $pn^2+p^2n$ operations that $XW'X^T$ would take. Mar 26, 2021 at 14:27
• I do not see why the weight matrices would not be square matrices. Take a look here for example, in the MultiHeadedAttention class where the line self.linears = clones(nn.Linear(d_model, d_model), 4) initializes these matrices, if I am not mistaken. Jul 13, 2021 at 18:32
• @TimoDenk For 3 of 4 of those clones layers it would probably be clearer to write nn.Linear(d_model, self.d_k * h) to communicate the weights being instantiated are really h = 8 independent rectangular weight matrices of the expected shape $d_{model} \times d_k$ (not square). For a longer explanation see davidvandebunte.gitlab.io/executable-notes/notes/se/…. Aug 21, 2022 at 20:14