Attention: Isn't it redundant to apply a linear layer to both the keys and values?

Question

Transformer attention is calculated $Attention(X) =X W^V\times \text{columnwise-softmax} (Att(X))$ where the attention attention matrix is $$Att(X) = Q \times K = {X W}^Q \times ({X W}^K)^T = {X W}^Q (W^K)^T X^T $$

But then couldn't one of $W^Q$ and $W^K$ be absorbed into the other? So one of them is redundant say $W^K$ and so attention only needs the parameters from $W^Q$ and $W^V$.

$$Attention(X) = {XW}^V \times \text{columnwise-softmax}({XW}^Q X^T).$$

Diagrams I use for thinking.

Answer 1

In multi-head attention the Keys, Values and Queries are chunked along the "channel" dimension, you need to apply a linear layer before you do this.

What I wanted to do wouldn't work because you would be chunking $X$. The chunks of $Q,K,V$ depend on all the channels of $X$, where as the chunks of $X$ do not.

There is code here to see exactly what is happening with multi-head attention:

q = rearrange(self.w_qs(q), 'b l (head k) -> head b l k', head=self.n_head)
k = rearrange(self.w_ks(k), 'b t (head k) -> head b t k', head=self.n_head)
v = rearrange(self.w_vs(v), 'b t (head v) -> head b t v', head=self.n_head)