6
$\begingroup$

In the attention mechanism, as described in the paper "Attention is all you need", normally the model learns weights $W_Q$ and $W_K$, i.e. the Queries and the Keys. The input of the attention layer is multiplied with both weights to get the Queries and Keys

\begin{equation} Q = W_Q^TX~\text{and}~K = W_K^TX \end{equation}

before computing the attention weights, called $a_w$ in this context, using the softmax function:

\begin{equation} a_w = \operatorname{softmax}\left(\frac{QK^T}{\sqrt{d_k}}\right) \end{equation}

where $\sqrt{d_k}$ is the square root of the dimension of the input. To my current understanding, and described in very newbie terms, the idea of using different weights for the Queries and Keys, is to project the input to different spaces, to potentially find useful connections between each element in the input. (A useful connection corresponds to a high softmax entry.)

My question is now, how important actually is the aspect of "different spaces"? What if we, in theory, without computational issues, use $4$ weight matrices? What if we only use $1$ weight matrix? That is, compute the keys and queries like so:

\begin{array}{cc} Q = W_Q^TX & K = W_Q^TX \end{array}

Wouldn't we still be able to find useful connections for each input?

$\endgroup$

2 Answers 2

8
$\begingroup$

The answer to your question is yes: you'd find useful connections still. Setting $W_Q \triangleq W_K$ is one of the suggestions for parameter efficiency from Kitaev et al. (2020 at ICLR). It has the effect of making the queries and keys identical.

The authors call this "shared-QK attention" and note in §5 that this parameter-tying approach is not worse than the standard attention mechanism:

A shared query-key space does not perform worse than regular attention; in fact, for enwik8 it appears to train slightly faster. In other words, we are not sacrificing accuracy by switching to shared-QK attention.

$\endgroup$
2
  • $\begingroup$ Is it then reasonable to say that if the data has abundant features, i.e. many potential useful connections between the different samples, a non shared key and query might be beneficial? Or why would we use non-shares weights at all? I have never seen it in implementations. $\endgroup$
    – kklaw
    Commented May 13, 2023 at 15:10
  • 1
    $\begingroup$ It appears in the Reformer implementation, which exists in many toolkits. The arguments for sharing or not sharing are the same as ever: amount of data, inductive biases, blah blah. Comments are not for extended discussion, so if you have other questions, it's best to post them as new ones and link back to this. $\endgroup$ Commented May 13, 2023 at 16:14
1
$\begingroup$

One reason to have different Query and Key vectors is to allow asymmetry. In the sentence "The movie was a box office bomb" the meaning of "movie" does not change because of the keys "bomb", but the word "bomb" is definitely interpreted differently because of the word "movie".

It also allows for the same word to be used differently as Query and Key. For instance, "Can you smell the smell of failure?". If you were doing a part of speech tagging task, it would be valuable to have different values for "smell" depending on whether it was being used as Query or a Key.

$\endgroup$

Your Answer

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

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